Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, UploadFile, File
 from pydantic import BaseModel
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionPipeline
 import torch
 from PIL import Image
 import io
@@ -11,27 +11,32 @@ app = FastAPI()
 MODEL_PATH = "Interior.safetensors"
 LORA_PATH = "Interior_lora.safetensors"
 
-print("Loading base model...")
+print("Loading model...")
 
-txt2img = StableDiffusionPipeline.from_single_file(
+# ========================
+# ⚡ SAFE CPU CONFIG
+# ========================
+pipe = StableDiffusionPipeline.from_single_file(
     MODEL_PATH,
-    torch_dtype=torch.float16,
-    safety_checker=None
-).to("cpu")  # هنرجعها GPU لو متاح لاحقًا
-
-img2img = StableDiffusionImg2ImgPipeline.from_single_file(
-    MODEL_PATH,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float32,   # ✔ مهم جدًا
     safety_checker=None
 ).to("cpu")
 
 print("Loading LoRA...")
 
-txt2img.load_lora_weights(LORA_PATH)
-img2img.load_lora_weights(LORA_PATH)
+pipe.load_lora_weights(LORA_PATH)
+pipe.fuse_lora(lora_scale=0.8)
 
-print("LoRA loaded 🔥")
+# ⚡ speed boost
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
 
+print("Model ready 🔥")
+
+
+# ========================
+# REQUEST MODEL
+# ========================
 class Prompt(BaseModel):
     prompt: str
 
@@ -43,26 +48,38 @@ def to_bytes(img):
     return buf
 
 
+# ========================
+# TXT2IMG
+# ========================
 @app.post("/txt2img")
 def generate(data: Prompt):
 
-    image = txt2img(
+    image = pipe(
         data.prompt,
-        cross_attention_kwargs={"scale": 0.8}
+        num_inference_steps=6,   # ⚡ سريع
+        guidance_scale=5,
+        height=256,
+        width=256
     ).images[0]
 
     return StreamingResponse(to_bytes(image), media_type="image/png")
 
 
+# ========================
+# IMG2IMG
+# ========================
 @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))
+    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    img = img.resize((256, 256))  # ⚡ أسرع
 
-    image = img2img(
+    image = pipe(
         prompt=prompt,
         image=img,
-        cross_attention_kwargs={"scale": 0.8}
+        strength=0.6,
+        num_inference_steps=6,
+        guidance_scale=5
     ).images[0]
 
-    return StreamingResponse(to_bytes(image), media_type="image/png")
+    return StreamingResponse(to_bytes(image), media_type="image/png")