Update app.py

app.py

@@ -1,110 +1,126 @@
 from fastapi import FastAPI, UploadFile, File
 from pydantic import BaseModel
-from diffusers import StableDiffusionPipeline
-import torch
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+from huggingface_hub import hf_hub_download
+from fastapi.responses import StreamingResponse
 from PIL import Image
+import torch
+
 import io
-import os
-from fastapi.responses import StreamingResponse
+
+
 
 app = FastAPI()
 
-MODEL_PATH = "Interior.safetensors"
-LORA_PATH = "Interior_lora.safetensors"
+# =========================
+# تحميل الموديلات من HuggingFace
+# =========================
+MODEL_PATH = hf_hub_download(
+    repo_id="ebraam1/interior-sd-models",
+    filename="Interior.safetensors"
+)
+
+LORA_PATH = hf_hub_download(
+    repo_id="ebraam1/interior-sd-models",
+    filename="Interior_lora.safetensors"
+)
+
+print("Loading base model...")
+
+# =========================
+# Load Stable Diffusion
+# =========================
+txt2img = StableDiffusionPipeline.from_single_file(
+    MODEL_PATH,
+    torch_dtype=torch.float16,
+    safety_checker=None
+).to("cpu")   # غيّرها لـ "cuda" لو GPU متاح
+
+img2img = StableDiffusionImg2ImgPipeline.from_single_file(
+
 
-# ========================
-# ⚡ 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())
 
-print("Loading model...")
 
-# ========================
-# SINGLE PIPELINE (IMPORTANT FIX)
-# ========================
-pipe = StableDiffusionPipeline.from_single_file(
     MODEL_PATH,
-    torch_dtype=dtype,
+    torch_dtype=torch.float16,
     safety_checker=None
-).to(device)
+).to("cpu")
 
 print("Loading LoRA...")
 
-pipe.load_lora_weights(LORA_PATH)
-pipe.fuse_lora(lora_scale=0.7)
+txt2img.load_lora_weights(LORA_PATH)
+img2img.load_lora_weights(LORA_PATH)
+
+
+
+
+
 
-# ========================
-# SPEED BOOSTS
-# ========================
-pipe.enable_attention_slicing()
-pipe.enable_vae_slicing()
 
-print("Model ready 🔥")
+print("LoRA loaded 🔥")
 
+# =========================
+# API Schema
+# =========================
 
-# ========================
-# REQUEST MODEL
-# ========================
 class Prompt(BaseModel):
     prompt: str
 
+# =========================
+# Helper: PIL → Bytes
+# =========================
 
-# ========================
-# IMAGE UTILS
-# ========================
 def to_bytes(img):
     buf = io.BytesIO()
     img.save(buf, format="PNG")
     buf.seek(0)
     return buf
 
+# =========================
+# TEXT → IMAGE
+# =========================
+
+
+
+
+
+
 
-# ========================
-# HEALTH CHECK
-# ========================
-@app.get("/")
-def home():
-    return {"status": "API is running 🚀"}
 
 
-# ========================
-# TXT2IMG (FAST MODE)
-# ========================
 @app.post("/txt2img")
 def generate(data: Prompt):
 
-    image = pipe(
+    image = txt2img(
         data.prompt,
-        num_inference_steps=6,   # ⚡ أسرع بكتير
-        guidance_scale=5,
-        height=256,
-        width=256
+        cross_attention_kwargs={"scale": 0.8}
+
+
+
     ).images[0]
 
     return StreamingResponse(to_bytes(image), media_type="image/png")
 
+# =========================
+# IMAGE → IMAGE
+# =========================
 
-# ========================
-# IMG2IMG (FAST MODE)
-# ========================
 @app.post("/img2img")
-async def img2img_api(
-    file: UploadFile = File(...),
-    prompt: str = ""
-):
+async def img2img_api(file: UploadFile = File(...), prompt: str = ""):
+
+
+
 
-    img = Image.open(io.BytesIO(await file.read())).convert("RGB")
-    img = img.resize((256, 256))  # ⚡ أسرع بشكل واضح
+    img = Image.open(io.BytesIO(await file.read())).convert("RGB").resize((512, 512))
 
-    image = pipe(
+
+    image = img2img(
         prompt=prompt,
         image=img,
-        strength=0.6,
-        num_inference_steps=6,
-        guidance_scale=5
+        cross_attention_kwargs={"scale": 0.8}
+
+
     ).images[0]
 
-    return StreamingResponse(to_bytes(image), media_type="image/png")
+    return StreamingResponse(to_bytes(image), media_type="image/png")