# app.py — FLUX.1-schnell (1 step, 256x256), async jobs + gallery + download/clear # - POST /generate -> returns job_id immediately # - GET /status/{job_id} -> ready / failed / image_url # - GET / -> gallery of session images (with prompts) # - GET /download-all -> zip of all images + prompts # - POST /clear-all -> wipe /data/images # # Storage-friendly: saves JPEG (not PNG), trims image folder by count/size. import os import io import uuid import threading from typing import Optional, List, Tuple, Dict from fastapi import FastAPI, HTTPException, Request from fastapi.responses import HTMLResponse, StreamingResponse, RedirectResponse from fastapi.staticfiles import StaticFiles from pydantic import BaseModel # ---------- Perf & cache envs ---------- os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("MKL_NUM_THREADS", "1") os.environ.setdefault("HF_HOME", "/data/hf") os.environ.setdefault("TRANSFORMERS_CACHE", "/data/hf/transformers") os.makedirs(os.environ["HF_HOME"], exist_ok=True) os.makedirs(os.environ["TRANSFORMERS_CACHE"], exist_ok=True) # ---------- Token (your exact snippet) ---------- HF_TOKEN = ( os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN") or (open("/run/secrets/HF_TOKEN").read().strip() if os.path.exists("/run/secrets/HF_TOKEN") else None) ) # ---------- Storage ---------- IMAGE_DIR = "/data/images" os.makedirs(IMAGE_DIR, exist_ok=True) # ---------- App ---------- app = FastAPI(title="flux-schnell-async", version="0.2.0") app.mount("/images", StaticFiles(directory=IMAGE_DIR), name="images") # ---------- Schemas ---------- class GenerateRequest(BaseModel): prompt: str class GenerateResponse(BaseModel): job_id: str class StatusResponse(BaseModel): job_id: str ready: bool image_url: Optional[str] = None failed: Optional[bool] = None error: Optional[str] = None # ---------- Model ---------- import torch # noqa from PIL import Image # noqa from diffusers import FluxPipeline # noqa DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 PIPE = None # Hard-coded generation params W = 256 H = 256 STEPS = 1 GUIDANCE = 0.0 MODEL_ID = "black-forest-labs/FLUX.1-schnell" # ---------- In-memory job store ---------- # JOBS[job_id] = { # "status": "pending" | "done" | "failed", # "prompt": str, # "image_name": Optional[str], # JPG filename when done # "error": Optional[str] # } JOBS: Dict[str, Dict] = {} JOBS_LOCK = threading.Lock() # ---------- Helpers ---------- def _list_images_with_prompts() -> List[Tuple[str, float, str]]: items = [] try: for f in os.listdir(IMAGE_DIR): if not f.lower().endswith((".jpg", ".jpeg")): continue path = os.path.join(IMAGE_DIR, f) mtime = os.path.getmtime(path) stem, _ = os.path.splitext(f) prompt_text = "" ptxt = os.path.join(IMAGE_DIR, f"{stem}.txt") if os.path.exists(ptxt): try: with open(ptxt, "r", encoding="utf-8") as pf: prompt_text = pf.read().strip() except Exception: pass items.append((f, mtime, prompt_text)) items.sort(key=lambda t: t[1], reverse=True) except Exception: pass return items def _dir_size_bytes(path: str) -> int: total = 0 for root, _, files in os.walk(path): for name in files: try: total += os.path.getsize(os.path.join(root, name)) except OSError: pass return total def _trim_images(max_keep=300, max_bytes=2_000_000_000): """ Keep at most `max_keep` files (images+txt) and at most `max_bytes` bytes (~2GB) in IMAGE_DIR. Newest are kept; oldest are deleted first. """ files = [] for f in os.listdir(IMAGE_DIR): if f.lower().endswith((".jpg", ".jpeg", ".txt")): p = os.path.join(IMAGE_DIR, f) try: files.append((p, os.path.getmtime(p))) except OSError: pass files.sort(key=lambda t: t[1], reverse=True) # newest first # by count for p, _ in files[max_keep:]: try: os.remove(p) except OSError: pass # by size files = sorted([(os.path.join(IMAGE_DIR, f), os.path.getmtime(os.path.join(IMAGE_DIR, f))) for f in os.listdir(IMAGE_DIR) if f.lower().endswith((".jpg", ".jpeg", ".txt"))], key=lambda t: t[1], reverse=True) while _dir_size_bytes(IMAGE_DIR) > max_bytes and files: p, _ = files.pop() # oldest try: os.remove(p) except OSError: pass def _render_gallery(base_url: str, files: List[Tuple[str, float, str]], limit: int = 100) -> str: items = files[:limit] cards = [] for name, _, prompt in items: url = f"{base_url}/images/{name}" safe_prompt = (prompt or "[no prompt saved]").replace("<", "<").replace(">", ">") cards.append( f""" {name}
{safe_prompt}
{name}
 """ ) grid = "\n".join(cards) or "

No images yet. POST to /generate to create one.

" return f""" FLUX.1-schnell — Gallery

FLUX.1-schnell — Session Gallery

{len(items)} image(s)
{grid}

Create a job:

curl -X POST '{base_url}/generate' -H 'Content-Type: application/json' -d '{{"prompt":"a serene landscape"}}'

Poll status:

curl '{base_url}/status/<job_id>'
""" # ---------- Startup: load pipeline ---------- @app.on_event("startup") def _load_model(): global PIPE kwargs = dict(dtype=DTYPE, use_safetensors=True, low_cpu_mem_usage=True) if HF_TOKEN: kwargs["token"] = HF_TOKEN PIPE = FluxPipeline.from_pretrained(MODEL_ID, **kwargs) if DEVICE == "cuda": PIPE = PIPE.to("cuda") else: try: PIPE.enable_attention_slicing() PIPE.enable_vae_slicing() PIPE.enable_vae_tiling() except Exception: pass # ---------- Background worker ---------- def _run_job(job_id: str, prompt: str): try: # HARD-CODED params: 256x256, 1 step, guidance 0.0 result = PIPE( prompt=prompt, width=W, height=H, num_inference_steps=STEPS, guidance_scale=GUIDANCE, ) if not result.images: raise RuntimeError("pipeline returned no images") image = result.images[0].convert("RGB") # Save JPEG (smaller than PNG) jpg_name = f"{job_id}.jpg" txt_name = f"{job_id}.txt" jpg_path = os.path.join(IMAGE_DIR, jpg_name) image.save(jpg_path, format="JPEG", quality=90, optimize=True, progressive=True) try: with open(os.path.join(IMAGE_DIR, txt_name), "w", encoding="utf-8") as f: f.write(prompt) except Exception: pass # Trim folder to stay under caps _trim_images(max_keep=300, max_bytes=2_000_000_000) # ~2 GB with JOBS_LOCK: JOBS[job_id]["status"] = "done" JOBS[job_id]["image_name"] = jpg_name except Exception as e: with JOBS_LOCK: JOBS[job_id]["status"] = "failed" JOBS[job_id]["error"] = str(e) # ---------- Routes ---------- @app.get("/", response_class=HTMLResponse) def index(request: Request, limit: int = 100): base = str(request.base_url).rstrip("/") files = _list_images_with_prompts() html = _render_gallery(base, files, limit=limit) return HTMLResponse(content=html, status_code=200) @app.get("/download-all") def download_all(): import zipfile buf = io.BytesIO() with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as z: for f in os.listdir(IMAGE_DIR): if f.lower().endswith((".jpg", ".jpeg", ".txt")): z.write(os.path.join(IMAGE_DIR, f), arcname=f) buf.seek(0) return StreamingResponse( buf, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=images.zip"}, ) @app.post("/clear-all") def clear_all(): deleted = 0 for f in os.listdir(IMAGE_DIR): if f.lower().endswith((".jpg", ".jpeg", ".txt")): try: os.remove(os.path.join(IMAGE_DIR, f)) deleted += 1 except Exception: pass return RedirectResponse("/", status_code=303) @app.post("/generate", response_model=GenerateResponse) def generate(req: GenerateRequest): # Create job and return immediately (no delay) job_id = str(uuid.uuid4()) with JOBS_LOCK: JOBS[job_id] = {"status": "pending", "prompt": req.prompt, "image_name": None, "error": None} t = threading.Thread(target=_run_job, args=(job_id, req.prompt), daemon=True) t.start() return GenerateResponse(job_id=job_id) @app.get("/status/{job_id}", response_model=StatusResponse) def status(job_id: str, request: Request): with JOBS_LOCK: job = JOBS.get(job_id) if not job: # Unknown job: report not ready (per your spec) return StatusResponse(job_id=job_id, ready=False) if job["status"] == "pending": return StatusResponse(job_id=job_id, ready=False) if job["status"] == "failed": return StatusResponse(job_id=job_id, ready=False, failed=True, error=job.get("error") or "job failed") # done base = str(request.base_url).rstrip("/") image_url = f"{base}/images/{job['image_name']}" return StatusResponse(job_id=job_id, ready=True, image_url=image_url)