Update app.py

app.py

@@ -1,91 +1,70 @@
-from fastapi import FastAPI, UploadFile, File, Query
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
-from huggingface_hub import hf_hub_download
-from fastapi.responses import StreamingResponse
-from PIL import Image
+from fastapi import FastAPI, UploadFile, File
+from pydantic import BaseModel
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 import torch
+from PIL import Image
 import io
+from fastapi.responses import StreamingResponse
 
 app = FastAPI()
 
-MODEL_REPO = "ebraam1/interior-sd-models"
-BASE_MODEL = "Interior.safetensors"
-LORA_FILE = "Interior_lora.safetensors"
+MODEL_PATH = "Interior.safetensors"
+LORA_PATH = "Interior_lora.safetensors"
 
-# تحميل الموديل الأساسي
-print("Downloading base model...")
-base_path = hf_hub_download(repo_id=MODEL_REPO, filename=BASE_MODEL)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
-print("Loading base model on CPU...")
-pipe = StableDiffusionPipeline.from_single_file(
-    base_path,
-    torch_dtype=torch.float32,
-    safety_checker=None,
-    requires_safety_checker=False
-).to("cpu")
 
-# أفضل Scheduler للجودة العالية
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+print("Loading base model...")
+
+txt2img = StableDiffusionPipeline.from_single_file(
+    MODEL_PATH,
+    torch_dtype=dtype,
+    safety_checker=None
+).to(device)
+
+img2img = StableDiffusionImg2ImgPipeline.from_single_file(
+    MODEL_PATH,
+    torch_dtype=dtype,
+    safety_checker=None
+).to(device)
 
-# تحميل LoRA ودمجه
-print("Downloading LoRA...")
-lora_path = hf_hub_download(repo_id=MODEL_REPO, filename=LORA_FILE)
 
 print("Loading LoRA...")
-pipe.load_lora_weights(lora_path)
-pipe.fuse_lora(lora_scale=1.0)
 
-# تحسينات الذاكرة
-pipe.enable_attention_slicing()
-pipe.enable_vae_slicing()
+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)
 
-print("Model ready 🔥 (Max quality mode - Landscape)")
+print("LoRA loaded 🔥")
 
-def to_bytes(img: Image.Image) -> io.BytesIO:
+
+class Prompt(BaseModel):
+    prompt: str
+
+
+def to_bytes(img):
     buf = io.BytesIO()
     img.save(buf, format="PNG")
     buf.seek(0)
     return buf
 
-# ================== Text-to-Image (GET) ==================
-@app.get("/txt2img")
-def txt2img(
-    prompt: str = Query(..., description="الوصف"),
-    negative_prompt: str = Query("", description="العناصر اللي عايز تتجنبها"),
-    steps: int = Query(20, ge=1, le=30),
-    guidance: float = Query(9.0, ge=1.0, le=20.0),
-    height: int = Query(512, ge=256, le=768),   # أصبح أقل (عرضي)
-    width: int = Query(768, ge=256, le=768)     # أصبح أكبر
-):
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt or None,
-        num_inference_steps=steps,
-        guidance_scale=guidance,
-        height=height,
-        width=width
-    ).images[0]
+
+@app.get("/")
+def home():
+    return {"status": "API is running 🚀"}
+
+
+@app.post("/txt2img")
+def generate(data: Prompt):
+    image = txt2img(data.prompt).images[0]
     return StreamingResponse(to_bytes(image), media_type="image/png")
 
-# ================== Image-to-Image (POST) ==================
+
 @app.post("/img2img")
-async def img2img(
-    file: UploadFile = File(...),
-    prompt: str = Query(""),
-    negative_prompt: str = Query(""),
-    steps: int = Query(20, ge=1, le=30),
-    guidance: float = Query(9.0, ge=1.0, le=20.0),
-    strength: float = Query(0.6, ge=0.0, le=1.0),
-    height: int = Query(512, ge=256, le=768),
-    width: int = Query(768, ge=256, le=768)
-):
-    img = Image.open(io.BytesIO(await file.read())).convert("RGB").resize((width, height))
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt or None,
-        image=img,
-        strength=strength,
-        num_inference_steps=steps,
-        guidance_scale=guidance
-    ).images[0]
+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]
     return StreamingResponse(to_bytes(image), media_type="image/png")