Update app.py

app.py

@@ -11,27 +11,36 @@ app = FastAPI()
 MODEL_PATH = "Interior.safetensors"
 LORA_PATH = "Interior_lora.safetensors"
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+
 print("Loading base model...")
 
 txt2img = StableDiffusionPipeline.from_single_file(
     MODEL_PATH,
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     safety_checker=None
-).to("cpu")  # هنرجعها GPU لو متاح لاحقًا
+).to(device)
 
 img2img = StableDiffusionImg2ImgPipeline.from_single_file(
     MODEL_PATH,
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     safety_checker=None
-).to("cpu")
+).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
 
@@ -43,26 +52,19 @@ def to_bytes(img):
     return buf
 
 
-@app.post("/txt2img")
-def generate(data: Prompt):
+@app.get("/")
+def home():
+    return {"status": "API is running 🚀"}
 
-    image = txt2img(
-        data.prompt,
-        cross_attention_kwargs={"scale": 0.8}
-    ).images[0]
 
+@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,
-        cross_attention_kwargs={"scale": 0.8}
-    ).images[0]
-
+    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")