Update app.py

app.py

@@ -1,20 +1,19 @@
-import spaces  # must come first before any torch import
+import spaces  # must come first before any torch or CUDA import
 import io
 import base64
+import torch
 import asyncio
 from fastapi import FastAPI, Request, Form
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image
-import torch
 from diffusers import StableDiffusionPipeline
-from peft import PeftModel  # required for LoRA
+from safetensors.torch import load_file
 
 # ======================
 # App Setup
 # ======================
 app = FastAPI(title="Ghibli + Anime Image Generator (ZeroGPU CPU mode)")
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -27,49 +26,42 @@ app.add_middleware(
 # ======================
 device = "cpu"
 dtype = torch.float32
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
 
 print("🔹 Loading base SDXL model (CPU)...")
-base_model = "stabilityai/stable-diffusion-xl-base-1.0"
-
-# Load once and reuse
 pipe_ghibli = StableDiffusionPipeline.from_pretrained(
-    base_model,
+    BASE_MODEL,
     torch_dtype=dtype,
-    safety_checker=None,  # You can re-enable this if you prefer
 )
 pipe_ghibli.to(device)
 
 pipe_anime = StableDiffusionPipeline.from_pretrained(
-    base_model,
+    BASE_MODEL,
     torch_dtype=dtype,
-    safety_checker=None,
 )
 pipe_anime.to(device)
 
-# --- Apply LoRA weights manually via PEFT ---
-try:
-    print("🎨 Loading Ghibli LoRA...")
-    PeftModel.from_pretrained(pipe_ghibli.unet, "./studioghibli_flux_r32-v2.safetensors")
-    print("✅ Ghibli LoRA loaded.")
-except Exception as e:
-    print(f"⚠️ Ghibli LoRA load failed: {e}")
-
-try:
-    print("🎨 Loading Anime LoRA...")
-    PeftModel.from_pretrained(pipe_anime.unet, "./Flux_1_Dev_LoRA_AestheticAnime.safetensors")
-    print("✅ Anime LoRA loaded.")
-except Exception as e:
-    print(f"⚠️ Anime LoRA load failed: {e}")
+# --- Load LoRA weights directly ---
+def apply_lora(pipe, path):
+    try:
+        print(f"🎨 Applying LoRA: {path}")
+        lora_weights = load_file(path)
+        missing, unexpected = pipe.unet.load_state_dict(lora_weights, strict=False)
+        print(f"✅ LoRA loaded successfully. Missing: {len(missing)}, Unexpected: {len(unexpected)}")
+    except Exception as e:
+        print(f"⚠️ Failed to load LoRA ({path}): {e}")
+
+apply_lora(pipe_ghibli, "./studioghibli_flux_r32-v2.safetensors")
+apply_lora(pipe_anime, "./Flux_1_Dev_LoRA_AestheticAnime.safetensors")
 
 # ======================
 # Helper functions
 # ======================
-
 async def generate_image(prompt: str, style: str, seed: int):
     generator = torch.Generator(device=device).manual_seed(seed)
     pipe = pipe_anime if style.lower() == "anime" else pipe_ghibli
-
     print(f"🎨 Generating {style} image for: {prompt}")
+
     image = pipe(
         prompt=f"{prompt}, {style} style, cinematic lighting, high quality",
         num_inference_steps=30,
@@ -94,9 +86,8 @@ def image_to_base64(img: Image.Image) -> str:
 
 
 # ======================
-# API Routes
+# Routes
 # ======================
-
 @app.get("/", response_class=HTMLResponse)
 def home():
     return """
@@ -114,7 +105,7 @@ def home():
 
 
 @app.post("/api/generate")
-async def api_generate(request: Request, prompt: str = Form(...), style: str = Form(...), seed: int = Form(42)):
+async def api_generate(prompt: str = Form(...), style: str = Form("Ghibli"), seed: int = Form(42)):
     try:
         print(f"📩 Received: {prompt} | Style={style} | Seed={seed}")
         image = await generate_image(prompt, style, seed)
@@ -126,9 +117,6 @@ async def api_generate(request: Request, prompt: str = Form(...), style: str = F
         return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
 
 
-# ======================
-# ZeroGPU keep-alive
-# ======================
 @spaces.GPU
 def keep_alive():
     """Required for ZeroGPU Spaces — keeps container active."""
@@ -136,9 +124,6 @@ def keep_alive():
     return "OK"
 
 
-# ======================
-# Run the app
-# ======================
 if __name__ == "__main__":
     import uvicorn
     print("🚀 Launching FastAPI on port 7860 (ZeroGPU mode)")