Update app.py

app.py

@@ -1,48 +1,49 @@
-import io, base64
-from fastapi import FastAPI
-from pydantic import BaseModel
-from PIL import Image
-import torch
-
-from optimum.intel.openvino import OVStableDiffusionPipeline
-
-# مدل بهینه برای CPU (INT8). می‌تونی بعداً عوضش کنی.
-MODEL_ID = "OpenVINO/stable-diffusion-v1-5-int8-ov"
-
-print("Loading model ...")
-pipe = OVStableDiffusionPipeline.from_pretrained(MODEL_ID)
-pipe.reshape(512, 512)  # ثابت نگه داشتن ورودی => سریع‌تر روی CPU
-pipe.compile()
-print("Model loaded.")
-
-app = FastAPI(title="Txt2Img CPU API")
-
-class Req(BaseModel):
-    prompt: str
-    negative_prompt: str | None = None
-    steps: int = 20
-    guidance: float = 7.5
-    seed: int | None = None
-    width: int = 512
-    height: int = 512
-
-@app.get("/healthz")
-def health():
-    return {"ok": True}
-
-@app.post("/txt2img")
-def txt2img(r: Req):
-    g = torch.Generator(device="cpu")
-    if r.seed is not None:
-        g.manual_seed(r.seed)
-    image = pipe(
-        prompt=r.prompt,
-        negative_prompt=r.negative_prompt,
-        num_inference_steps=r.steps,
-        guidance_scale=r.guidance,
-        width=r.width, height=r.height,
-        generator=g
-    ).images[0]
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
-    return {"image_base64": base64.b64encode(buf.getvalue()).decode()}
+import io, base64, os, random
+from fastapi import FastAPI
+from pydantic import BaseModel
+from PIL import Image
+
+print(">>> importing optimum.intel.openvino ...")
+from optimum.intel.openvino import OVStableDiffusionPipeline
+print(">>> import OK")
+
+MODEL_ID = os.environ.get("MODEL_ID", "OpenVINO/stable-diffusion-v1-5-int8-ov")
+
+print("Loading model ...")
+pipe = OVStableDiffusionPipeline.from_pretrained(MODEL_ID)
+pipe.reshape(512, 512)  # برای CPU بهتر
+pipe.compile()
+print("Model loaded.")
+
+app = FastAPI(title="Txt2Img CPU API")
+
+class Req(BaseModel):
+    prompt: str
+    negative_prompt: str | None = None
+    steps: int = 20
+    guidance: float = 7.5
+    seed: int | None = None
+    width: int = 512
+    height: int = 512
+
+@app.get("/healthz")
+def health():
+    return {"ok": True}
+
+@app.post("/txt2img")
+def txt2img(r: Req):
+    # در OpenVINO نیازی به torch.Generator نیست؛ seed را روی random ست می‌کنیم (اختیاری)
+    if r.seed is not None:
+        random.seed(r.seed)
+
+    image = pipe(
+        prompt=r.prompt,
+        negative_prompt=r.negative_prompt,
+        num_inference_steps=r.steps,
+        guidance_scale=r.guidance,
+        width=r.width, height=r.height,
+    ).images[0]
+
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    return {"image_base64": base64.b64encode(buf.getvalue()).decode()}