Gibhili

Running

App Files Files

videopix commited on Nov 12

Commit

4cb20ce

verified ·

1 Parent(s): 3a9efeb

Update app.py

Browse files

Files changed (1) hide show

app.py +58 -138

app.py CHANGED Viewed

@@ -1,89 +1,55 @@
-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 FluxPipeline
 from PIL import Image
-import base64
-import io
-import asyncio
 from concurrent.futures import ThreadPoolExecutor
-# -----------------------------
-# Hugging Face Token Support
-# -----------------------------
-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 = {}
-executor = ThreadPoolExecutor(max_workers=2)  # reduce thread pressure for CPU
-semaphore = asyncio.Semaphore(1)  # avoids CPU overload
-# -----------------------------
-# Load FLUX Pipeline (Optimized)
-# -----------------------------
 def load_pipeline():
-    if "flux" in _cached_pipelines:
-        return _cached_pipelines["flux"]
-    print("🔹 Loading FLUX.1-schnell Model (Optimized for CPU)")
     pipe = FluxPipeline.from_pretrained(
         BASE_MODEL,
-        torch_dtype=torch.float16,          # <= Faster
         use_auth_token=HF_TOKEN,
-    )
-    pipe.to("cpu", dtype=torch.float16)     # <= Ensure CPU uses FP16
     pipe.enable_attention_slicing()
     pipe.enable_vae_tiling()
-    _cached_pipelines["flux"] = pipe
     return pipe
-# -----------------------------
-# Image Generation (Optimized)
-# -----------------------------
-def generate_image_sync(prompt: str, style: str = None, seed: int = 42):
     pipe = load_pipeline()
-    generator = torch.Generator(device="cpu").manual_seed(int(seed))
-    width = 1024   # <= Faster resolution
-    height = 576
     image = pipe(
         prompt=prompt,
-        width=width,
-        height=height,
-        num_inference_steps=8,   # <= Lower steps for speed
-        guidance_scale=2.8,      # <= Balanced for FLUX
-        generator=generator,
     ).images[0]
-    return image
-# -----------------------------
-# Async Wrapper
-# -----------------------------
-async def generate_image_async(prompt, style, seed):
     async with semaphore:
         loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(executor, generate_image_sync, prompt, style, seed)
-# -----------------------------
-# FastAPI App Setup
-# -----------------------------
-app = FastAPI(title="FLUX Image Generator API", version="2.2")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -92,103 +58,57 @@ app.add_middleware(
     allow_headers=["*"],
 )
 @app.get("/", response_class=HTMLResponse)
 def home():
     return """
-    <html>
-        <head>
-            <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: 90%; }
-            </style>
-        </head>
-        <body>
-            <h2>🎨 FLUX.1-schnell Image Generator</h2>
-            <form id="generateForm">
-                <input id="prompt" placeholder="Describe your image..." required><br>
-                <select id="style">
-                    <option value="default">Default (FLUX)</option>
-                </select><br>
-                <input id="seed" type="number" value="42"><br>
-                <button type="submit">Generate Image</button>
-            </form>
-            <div id="result"></div>
-            <script>
-                const form = document.getElementById("generateForm");
-                const resultDiv = document.getElementById("result");
-                form.addEventListener("submit", async (e) => {
-                    e.preventDefault();
-                    resultDiv.innerHTML = "<p>⏳ Generating image...</p>";
-                    const data = {
-                        prompt: document.getElementById("prompt").value,
-                        style: document.getElementById("style").value,
-                        seed: parseInt(document.getElementById("seed").value)
-                    };
-                    const res = await fetch("/api/generate", {
-                        method: "POST",
-                        headers: { "Content-Type": "application/json" },
-                        body: JSON.stringify(data)
-                    });
-                    const json = await res.json();
-                    if (json.status === "success") {
-                        resultDiv.innerHTML = `<img src="data:image/png;base64,${json.image_base64}"/><p>✅ Done!</p>`;
-                    } else {
-                        resultDiv.innerHTML = `<p style='color:red'>❌ ${json.message}</p>`;
-                    }
-                });
-            </script>
-        </body>
-    </html>
     """
-# -----------------------------
-# API Endpoint
-# -----------------------------
 @app.post("/api/generate")
 async def api_generate(request: Request):
     try:
         data = await request.json()
-        prompt = data.get("prompt", "").strip()
-        style = data.get("style", "default")
-        seed = data.get("seed", 42)
         if not prompt:
-            return JSONResponse({"status": "error", "message": "Prompt required"}, status_code=400)
     except Exception:
-        return JSONResponse({"status": "error", "message": "Invalid JSON"}, status_code=400)
     try:
-        image = await generate_image_async(prompt, style, seed)
-        buffer = io.BytesIO()
-        image.save(buffer, format="PNG")
-        img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        return JSONResponse({
-            "status": "success",
-            "prompt": prompt,
-            "style": style,
-            "image_base64": img_base64
-        })
     except Exception as e:
         print(f"❌ Error: {e}")
-        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)
-# -----------------------------
-# ZeroGPU Keep Alive
-# -----------------------------
 @spaces.GPU
-def keep_alive():
-    return "ZeroGPU Ready"
 if __name__ == "__main__":
     import uvicorn
-    print("🚀 Launching FastAPI (FLUX optimized for CPU)")
     keep_alive()
     uvicorn.run(app, host="0.0.0.0", port=7860)

+import os, io, base64, asyncio, torch, spaces
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse, JSONResponse
 from diffusers import FluxPipeline
 from PIL import Image
 from concurrent.futures import ThreadPoolExecutor
+HF_TOKEN = os.getenv("HF_TOKEN")
 BASE_MODEL = "black-forest-labs/FLUX.1-schnell"
+_cached = {}
+# moderate concurrency so CPU doesn’t choke
+executor = ThreadPoolExecutor(max_workers=3)
+semaphore = asyncio.Semaphore(3)
 def load_pipeline():
+    if "flux" in _cached:
+        return _cached["flux"]
+    print("🔹 Loading FLUX.1-schnell (fast mode)")
     pipe = FluxPipeline.from_pretrained(
         BASE_MODEL,
+        torch_dtype=torch.float16,
         use_auth_token=HF_TOKEN,
+    ).to("cpu", dtype=torch.float16)
     pipe.enable_attention_slicing()
     pipe.enable_vae_tiling()
+    _cached["flux"] = pipe
     return pipe
+def generate_image_sync(prompt: str, seed: int = 42):
     pipe = load_pipeline()
+    gen = torch.Generator(device="cpu").manual_seed(int(seed))
+    # smaller size and steps for speed
+    w, h = 768, 432
     image = pipe(
         prompt=prompt,
+        width=w,
+        height=h,
+        num_inference_steps=6,
+        guidance_scale=2.5,
+        generator=gen,
     ).images[0]
+    # slight upscale back to 960×540 to keep output clear
+    return image.resize((960, 540), Image.BICUBIC)
+async def generate_image_async(prompt, seed):
     async with semaphore:
         loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(executor, generate_image_sync, prompt, seed)
+app = FastAPI(title="FLUX Fast API", version="3.1")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
     allow_headers=["*"],
 )
 @app.get("/", response_class=HTMLResponse)
 def home():
     return """
+    <html><head><title>FLUX Fast</title>
+    <style>body{font-family:Arial;text-align:center;padding:2rem}
+    input,button{margin:.5rem;padding:.6rem;width:300px;border-radius:6px;border:1px solid #ccc}
+    button{background:#444;color:#fff}button:hover{background:#333}
+    img{margin-top:1rem;max-width:90%;border-radius:12px}</style></head>
+    <body><h2>🎨 FLUX Fast Generator</h2>
+    <form id='f'><input id='prompt' placeholder='Describe image...' required><br>
+    <input id='seed' type='number' value='42'><br>
+    <button>Generate</button></form><div id='out'></div>
+    <script>
+    const f=document.getElementById('f'),o=document.getElementById('out');
+    f.addEventListener('submit',async e=>{
+      e.preventDefault();o.innerHTML='⏳ Generating...';
+      const res=await fetch('/api/generate',{method:'POST',headers:{'Content-Type':'application/json'},
+        body:JSON.stringify({prompt:prompt.value,seed:+seed.value})});
+      const j=await res.json();
+      if(j.status==='success')o.innerHTML=`<img src="data:image/png;base64,${j.image_base64}"/><p>✅ Done!</p>`;
+      else o.innerHTML=`<p style='color:red'>❌ ${j.message}</p>`;
+    });
+    </script></body></html>
     """
 @app.post("/api/generate")
 async def api_generate(request: Request):
     try:
         data = await request.json()
+        prompt = str(data.get("prompt", "")).strip()
+        seed = int(data.get("seed", 42))
         if not prompt:
+            return JSONResponse({"status": "error", "message": "Prompt required"}, 400)
     except Exception:
+        return JSONResponse({"status": "error", "message": "Invalid JSON"}, 400)
     try:
+        image = await generate_image_async(prompt, seed)
+        buf = io.BytesIO()
+        image.save(buf, format="PNG")
+        img64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        return JSONResponse({"status": "success", "prompt": prompt, "image_base64": img64})
     except Exception as e:
         print(f"❌ Error: {e}")
+        return JSONResponse({"status": "error", "message": str(e)}, 500)
 @spaces.GPU
+def keep_alive(): return "ZeroGPU Ready"
 if __name__ == "__main__":
     import uvicorn
+    print("🚀 Launching Fast FLUX API")
     keep_alive()
     uvicorn.run(app, host="0.0.0.0", port=7860)