Spaces:
Running
Running
File size: 4,540 Bytes
07d8ca5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
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=4,
guidance_scale=3,
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_credentials=True,
allow_methods=["*"],
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 form = document.getElementById("f");
const promptInput = document.getElementById("prompt");
const seedInput = document.getElementById("seed");
const resultDiv = document.getElementById("out");
form.addEventListener("submit", async (e) => {
e.preventDefault();
const prompt = promptInput.value.trim();
if (!prompt) {
resultDiv.innerHTML = "<p style='color:red'>❌ Please enter a prompt</p>";
return;
}
resultDiv.innerHTML = "<p>⏳ Generating...</p>";
const payload = {
prompt: prompt,
seed: parseInt(seedInput.value || 42)
};
const res = await fetch("/api/generate", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload)
});
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>
"""
@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)
|