app.py

from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from huggingface_hub import hf_hub_download   # <-- جديد
import torch
from PIL import Image
import io
from fastapi.responses import StreamingResponse

app = FastAPI()

# --- ريبو الملفات على Hugging Face ---
MODEL_REPO = "ebraam1/interior-sd-models"
BASE_MODEL_FILE = "Interior.safetensors"
LORA_FILE = "Interior_lora.safetensors"

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# --- تحميل الملفات من الـ Hub ---
print("Downloading base model...")
base_path = hf_hub_download(repo_id=MODEL_REPO, filename=BASE_MODEL_FILE)

print("Downloading LoRA...")
lora_path = hf_hub_download(repo_id=MODEL_REPO, filename=LORA_FILE)

# --- تحميل البايبلاينز من الملف المحلي (نفس كود كولاب) ---
print("Loading base model...")
txt2img = StableDiffusionPipeline.from_single_file(
    base_path,                # <-- المسار المحلي بعد التحميل
    torch_dtype=dtype,
    safety_checker=None
).to(device)

img2img = StableDiffusionImg2ImgPipeline.from_single_file(
    base_path,
    torch_dtype=dtype,
    safety_checker=None
).to(device)

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)

print("LoRA loaded 🔥")

class Prompt(BaseModel):
    prompt: str

def to_bytes(img):
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return buf

@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")

@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]
    return StreamingResponse(to_bytes(image), media_type="image/png")