Upload app.py

app.py

@@ -0,0 +1,68 @@
+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_PATH = "Interior.safetensors"
+LORA_PATH = "Interior_lora.safetensors"
+
+print("Loading base model...")
+
+txt2img = 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,
+    safety_checker=None
+).to("cpu")
+
+print("Loading LoRA...")
+
+txt2img.load_lora_weights(LORA_PATH)
+img2img.load_lora_weights(LORA_PATH)
+
+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.post("/txt2img")
+def generate(data: Prompt):
+
+    image = txt2img(
+        data.prompt,
+        cross_attention_kwargs={"scale": 0.8}
+    ).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,
+        cross_attention_kwargs={"scale": 0.8}
+    ).images[0]
+
+    return StreamingResponse(to_bytes(image), media_type="image/png")