Update app.py

app.py

@@ -4,6 +4,7 @@ from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 import torch
 from PIL import Image
 import io
+import os
 from fastapi.responses import StreamingResponse
 
 app = FastAPI()
@@ -11,25 +12,34 @@ app = FastAPI()
 MODEL_PATH = "Interior.safetensors"
 LORA_PATH = "Interior_lora.safetensors"
 
+# ========================
+#  CPU Optimization
+# ========================
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
+torch.set_num_threads(os.cpu_count())  # تسريع CPU execution
 
 print("Loading base model...")
 
+# ========================
+# TXT2IMG PIPELINE
+# ========================
 txt2img = StableDiffusionPipeline.from_single_file(
     MODEL_PATH,
     torch_dtype=dtype,
     safety_checker=None
 ).to(device)
 
+# ========================
+# IMG2IMG PIPELINE
+# ========================
 img2img = StableDiffusionImg2ImgPipeline.from_single_file(
     MODEL_PATH,
     torch_dtype=dtype,
     safety_checker=None
 ).to(device)
 
-
 print("Loading LoRA...")
 
 txt2img.load_lora_weights(LORA_PATH)
@@ -38,13 +48,28 @@ img2img.load_lora_weights(LORA_PATH)
 txt2img.fuse_lora(lora_scale=0.8)
 img2img.fuse_lora(lora_scale=0.8)
 
+# ========================
+#  SPEED BOOSTS
+# ========================
+txt2img.enable_attention_slicing()
+txt2img.enable_vae_slicing()
+
+img2img.enable_attention_slicing()
+img2img.enable_vae_slicing()
+
 print("LoRA loaded 🔥")
 
 
+# ========================
+# REQUEST MODEL
+# ========================
 class Prompt(BaseModel):
     prompt: str
 
 
+# ========================
+# IMAGE UTILS
+# ========================
 def to_bytes(img):
     buf = io.BytesIO()
     img.save(buf, format="PNG")
@@ -52,19 +77,49 @@ def to_bytes(img):
     return buf
 
 
+# ========================
+# HEALTH CHECK
+# ========================
 @app.get("/")
 def home():
     return {"status": "API is running 🚀"}
 
 
+# ========================
+# TXT2IMG ENDPOINT (FAST MODE)
+# ========================
 @app.post("/txt2img")
 def generate(data: Prompt):
-    image = txt2img(data.prompt).images[0]
+
+    image = txt2img(
+        data.prompt,
+        num_inference_steps=10,   # ⚡ أسرع حاجة
+        guidance_scale=5,
+        height=384,
+        width=384
+    ).images[0]
+
     return StreamingResponse(to_bytes(image), media_type="image/png")
 
 
+# ========================
+# IMG2IMG ENDPOINT (FAST MODE)
+# ========================
 @app.post("/img2img")
-async def img2img_api(file: UploadFile = File(...), prompt: str = ""):
-    img = Image.open(io.BytesIO(await file.read())).convert("RGB").resize((512, 512))
-    image = img2img(prompt=prompt, image=img).images[0]
+async def img2img_api(
+    file: UploadFile = File(...),
+    prompt: str = ""
+):
+
+    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    img = img.resize((384, 384))  # ⚡ تسريع مهم جدًا
+
+    image = img2img(
+        prompt=prompt,
+        image=img,
+        strength=0.6,
+        num_inference_steps=10,
+        guidance_scale=5
+    ).images[0]
+
     return StreamingResponse(to_bytes(image), media_type="image/png")