Update app.py

app.py

@@ -1,10 +1,10 @@
 import spaces
 import torch
+import os
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse, JSONResponse
-from diffusers import StableDiffusionXLPipeline
-from safetensors.torch import load_file
+from diffusers import FluxPipeline
 from PIL import Image
 import base64
 import io
@@ -12,97 +12,80 @@ import asyncio
 from concurrent.futures import ThreadPoolExecutor
 
 # -----------------------------
-# Model Configuration
+# Hugging Face Token Support
 # -----------------------------
-BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
-LORA_MODELS = {
-    "Ghibli": "./studioghibli_flux_r32-v2.safetensors",
-    "GH1bli": "./gh1bli-style.safetensors"
-}
+HF_TOKEN = os.getenv("HF_TOKEN")  # Must be set on server / spaces
+
+# -----------------------------
+# Model (FLUX.1-schnell)
+# -----------------------------
+BASE_MODEL = "black-forest-labs/FLUX.1-schnell"
 
 _cached_pipelines = {}
-# Thread pool for parallel inference
-executor = ThreadPoolExecutor(max_workers=6)
-# Semaphore controls how many can run at once (to protect memory)
-semaphore = asyncio.Semaphore(6)
+executor = ThreadPoolExecutor(max_workers=4)
+semaphore = asyncio.Semaphore(4)
 
 # -----------------------------
-# Pipeline Loader
+# Load FLUX Pipeline
 # -----------------------------
-def load_pipeline(style="Ghibli"):
-    if style in _cached_pipelines:
-        return _cached_pipelines[style]
+def load_pipeline():
+    if "flux" in _cached_pipelines:
+        return _cached_pipelines["flux"]
 
-    print(f"🔹 Loading SDXL model for style: {style}")
-    device = "cpu"
-    dtype = torch.float32
+    print("🔹 Loading FLUX.1-schnell Model")
 
-    pipe = StableDiffusionXLPipeline.from_pretrained(
+    pipe = FluxPipeline.from_pretrained(
         BASE_MODEL,
-        dtype=dtype,
-        use_safetensors=True,
-    ).to(device)
+        torch_dtype=torch.float32,
+        use_auth_token=HF_TOKEN,  # <--- token applied here
+    ).to("cpu")
 
     pipe.enable_attention_slicing()
     pipe.enable_vae_tiling()
 
-    model_path = LORA_MODELS.get(style, LORA_MODELS["Ghibli"])
-    try:
-        print(f"🎨 Applying LoRA weights: {model_path}")
-        lora_weights = load_file(model_path)
-        pipe.unet.load_state_dict(lora_weights, strict=False)
-        print("✅ LoRA loaded successfully.")
-    except Exception as e:
-        print(f"⚠️ LoRA load failed: {e}")
-
-    _cached_pipelines[style] = pipe
+    _cached_pipelines["flux"] = pipe
     return pipe
 
 
 # -----------------------------
-# Synchronous Image Generation
+# Image Generation
 # -----------------------------
-def generate_image_sync(prompt: str, style: str = "Ghibli", seed: int = 42):
-    pipe = load_pipeline(style)
+def generate_image_sync(prompt: str, style: str = None, seed: int = 42):
+    pipe = load_pipeline()
     generator = torch.Generator(device="cpu").manual_seed(int(seed))
 
-    enhanced_prompt = (
-        f"{prompt}, Studio Ghibli style, full-frame composition, centered subject, "
-        "clean background, cinematic tone, detailed illustration, digital painting"
-    )
+    width = 1920
+    height = 1080
 
     image = pipe(
-        prompt=enhanced_prompt,
-        height=512,
-        width=512,
-        num_inference_steps=30,
-        guidance_scale=7.0,
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=20,
+        guidance_scale=3.5,
         generator=generator,
     ).images[0]
 
-    w, h = image.size
-    image = image.crop((5, 5, w - 5, h - 5)).resize((512, 512))
     return image
 
 
 # -----------------------------
-# Async Wrapper for Concurrency
+# Async Wrapper
 # -----------------------------
 async def generate_image_async(prompt, style, seed):
     async with semaphore:
-        loop = asyncio.get_event_loop()
+        loop = asyncio.get_running_loop()
         return await loop.run_in_executor(executor, generate_image_sync, prompt, style, seed)
 
 
 # -----------------------------
 # FastAPI App Setup
 # -----------------------------
-app = FastAPI(title="Studio Ghibli Generator API", version="2.0")
+app = FastAPI(title="FLUX Image Generator API", version="2.1")
 
-# Enable cross-app requests (CORS)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allow any domain (you can restrict if needed)
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -114,22 +97,21 @@ def home():
     return """
     <html>
         <head>
-            <title>Studio Ghibli Generator</title>
+            <title>FLUX Generator</title>
             <style>
                 body { font-family: Arial; text-align: center; padding: 2rem; background-color: #f9f9f9; }
                 input, select, button { margin: 0.5rem; padding: 0.6rem; width: 300px; border-radius: 6px; border: 1px solid #ccc; }
                 button { background-color: #444; color: white; cursor: pointer; }
                 button:hover { background-color: #333; }
-                img { margin-top: 1rem; border-radius: 12px; max-width: 512px; }
+                img { margin-top: 1rem; border-radius: 12px; max-width: 90%; }
             </style>
         </head>
         <body>
-            <h2>🎨 Studio Ghibli Generator</h2>
+            <h2>🎨 FLUX.1-schnell Image Generator</h2>
             <form id="generateForm">
-                <input id="prompt" placeholder="e.g. a boy riding a dragon" required><br>
+                <input id="prompt" placeholder="Describe your image..." required><br>
                 <select id="style">
-                    <option value="Ghibli">Ghibli</option>
-                    <option value="GH1bli">GH1bli</option>
+                    <option value="default">Default (FLUX)</option>
                 </select><br>
                 <input id="seed" type="number" value="42"><br>
                 <button type="submit">Generate Image</button>
@@ -172,7 +154,7 @@ async def api_generate(request: Request):
     try:
         data = await request.json()
         prompt = data.get("prompt", "").strip()
-        style = data.get("style", "Ghibli")
+        style = data.get("style", "default")
         seed = data.get("seed", 42)
         if not prompt:
             return JSONResponse({"status": "error", "message": "Prompt required"}, status_code=400)
@@ -197,7 +179,7 @@ async def api_generate(request: Request):
 
 
 # -----------------------------
-# ZeroGPU / Keep Alive Hook
+# ZeroGPU Keep Alive
 # -----------------------------
 @spaces.GPU
 def keep_alive():
@@ -206,6 +188,6 @@ def keep_alive():
 
 if __name__ == "__main__":
     import uvicorn
-    print("🚀 Launching FastAPI (Multi-Request / CPU / ZeroGPU Mode)")
+    print("🚀 Launching FastAPI with FLUX + HF Token")
     keep_alive()
     uvicorn.run(app, host="0.0.0.0", port=7860)