Create inference.py

inference.py

@@ -0,0 +1,63 @@
+import os
+
+# ===== CPU optimization =====
+CPU_THREADS = str(os.cpu_count() or 2)
+
+os.environ["OMP_NUM_THREADS"] = CPU_THREADS
+os.environ["OPENVINO_NUM_THREADS"] = CPU_THREADS
+os.environ["MKL_NUM_THREADS"] = CPU_THREADS
+os.environ["NUMEXPR_NUM_THREADS"] = CPU_THREADS
+os.environ["OMP_WAIT_POLICY"] = "PASSIVE"
+os.environ["KMP_BLOCKTIME"] = "0"
+
+import torch
+from optimum.intel import OVZImagePipeline
+
+# ===== Global =====
+pipe = None
+generator = torch.Generator("cpu")
+
+
+def load_model():
+    global pipe
+
+    if pipe is None:
+        pipe = OVZImagePipeline.from_pretrained(
+            "hsuwill000/Z-Image-Turbo-ov",
+            device="cpu"
+        )
+
+        # warmup
+        pipe(
+            prompt="warmup",
+            height=512,
+            width=512,
+            num_inference_steps=1,
+            guidance_scale=0.0,
+        )
+
+    return pipe
+
+
+def generate(
+    prompt: str,
+    height: int = 512,
+    width: int = 512,
+    steps: int = 9,
+    seed: int = -1,
+):
+    pipe = load_model()
+
+    if seed != -1:
+        generator.manual_seed(int(seed))
+
+    image = pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=0.0,
+        generator=generator,
+    ).images[0]
+
+    return image