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
@@ -13,28 +13,19 @@ MODEL_PATH = "Interior.safetensors"
 LORA_PATH = "Interior_lora.safetensors"
 
 # ========================
-#  CPU Optimization
+# ⚡ 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
+torch.set_num_threads(os.cpu_count())
 
-print("Loading base model...")
+print("Loading model...")
 
 # ========================
-# TXT2IMG PIPELINE
+# SINGLE PIPELINE (IMPORTANT FIX)
 # ========================
-txt2img = StableDiffusionPipeline.from_single_file(
-    MODEL_PATH,
-    torch_dtype=dtype,
-    safety_checker=None
-).to(device)
-
-# ========================
-# IMG2IMG PIPELINE
-# ========================
-img2img = StableDiffusionImg2ImgPipeline.from_single_file(
+pipe = StableDiffusionPipeline.from_single_file(
     MODEL_PATH,
     torch_dtype=dtype,
     safety_checker=None
@@ -42,22 +33,16 @@ img2img = StableDiffusionImg2ImgPipeline.from_single_file(
 
 print("Loading LoRA...")
 
-txt2img.load_lora_weights(LORA_PATH)
-img2img.load_lora_weights(LORA_PATH)
-
-txt2img.fuse_lora(lora_scale=0.8)
-img2img.fuse_lora(lora_scale=0.8)
+pipe.load_lora_weights(LORA_PATH)
+pipe.fuse_lora(lora_scale=0.7)
 
 # ========================
-#  SPEED BOOSTS
+# SPEED BOOSTS
 # ========================
-txt2img.enable_attention_slicing()
-txt2img.enable_vae_slicing()
-
-img2img.enable_attention_slicing()
-img2img.enable_vae_slicing()
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
 
-print("LoRA loaded 🔥")
+print("Model ready 🔥")
 
 
 # ========================
@@ -86,24 +71,24 @@ def home():
 
 
 # ========================
-# TXT2IMG ENDPOINT (FAST MODE)
+# TXT2IMG (FAST MODE)
 # ========================
 @app.post("/txt2img")
 def generate(data: Prompt):
 
-    image = txt2img(
+    image = pipe(
         data.prompt,
-        num_inference_steps=10,   # ⚡ أسرع حاجة
+        num_inference_steps=6,   # ⚡ أسرع بكتير
         guidance_scale=5,
-        height=384,
-        width=384
+        height=256,
+        width=256
     ).images[0]
 
     return StreamingResponse(to_bytes(image), media_type="image/png")
 
 
 # ========================
-# IMG2IMG ENDPOINT (FAST MODE)
+# IMG2IMG (FAST MODE)
 # ========================
 @app.post("/img2img")
 async def img2img_api(
@@ -112,13 +97,13 @@ async def img2img_api(
 ):
 
     img = Image.open(io.BytesIO(await file.read())).convert("RGB")
-    img = img.resize((384, 384))  # ⚡ تسريع مهم جدًا
+    img = img.resize((256, 256))  # ⚡ أسرع بشكل واضح
 
-    image = img2img(
+    image = pipe(
         prompt=prompt,
         image=img,
         strength=0.6,
-        num_inference_steps=10,
+        num_inference_steps=6,
         guidance_scale=5
     ).images[0]