app.py

# 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("<", "&lt;").replace(">", "&gt;")
        cards.append(
            f"""
            <a class="card" href="{url}" target="_blank" rel="noopener">
              <img src="{url}" loading="lazy" alt="{name}" />
              <div class="caption"><b>{safe_prompt}</b></div>
              <div class="caption">{name}</div>
            </a>
            """
        )
    grid = "\n".join(cards) or "<p>No images yet. POST to <code>/generate</code> to create one.</p>"
    return f"""
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1" />
<title>FLUX.1-schnell — Gallery</title>
<style>
  :root {{ --bg:#0b0f14; --fg:#e6edf3; --muted:#93a1b0; --card:#121820; --accent:#4ea1ff; }}
  body {{ margin:0; font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Arial, Noto Sans, "Apple Color Emoji","Segoe UI Emoji"; background:var(--bg); color:var(--fg); }}
  header {{ padding:16px 20px; border-bottom:1px solid #1b2530; display:flex; gap:10px; align-items:center; flex-wrap:wrap; }}
  h1 {{ font-size:18px; margin:0; }}
  .meta {{ font-size:12px; color:var(--muted); margin-left:auto; }}
  main {{ padding:16px; }}
  .actions {{ display:flex; gap:10px; flex-wrap:wrap; }}
  .btn {{ appearance:none; border:1px solid #1b2530; background:var(--card); color:var(--fg); padding:8px 12px; border-radius:10px; cursor:pointer; text-decoration:none; display:inline-block; }}
  .btn:hover {{ border-color:var(--accent); }}
  .grid {{ display:grid; grid-template-columns:repeat(auto-fill,minmax(220px,1fr)); gap:14px; margin-top:14px; }}
  .card {{ display:block; background:var(--card); border:1px solid #1b2530; border-radius:12px; overflow:hidden; text-decoration:none; color:inherit; }}
  .card:hover {{ outline:2px solid var(--accent); }}
  img {{ display:block; width:100%; height:220px; object-fit:cover; background:#0a0a0a; }}
  .caption {{ padding:8px 10px; font-size:12px; color:var(--muted); white-space:nowrap; overflow:hidden; text-overflow:ellipsis; }}
  .howto {{ margin-top:16px; font-size:13px; color:var(--muted); line-height:1.5; }}
  code {{ background:#0f1620; padding:2px 6px; border-radius:6px; }}
</style>
</head>
<body>
<header>
  <h1>FLUX.1-schnell — Session Gallery</h1>
  <div class="meta">{len(items)} image(s)</div>
  <div class="actions">
    <form method="get" action="{base_url}/download-all">
      <button class="btn" type="submit">⬇️ Download all (.zip)</button>
    </form>
    <form method="post" action="{base_url}/clear-all" onsubmit="return confirm('Delete ALL saved images and prompts?')">
      <button class="btn" type="submit">🗑️ Clear folder</button>
    </form>
  </div>
</header>
<main>
  <div class="grid">
    {grid}
  </div>
  <div class="howto">
    <p>Create a job:</p>
    <pre><code>curl -X POST '{base_url}/generate' -H 'Content-Type: application/json' -d '{{"prompt":"a serene landscape"}}'</code></pre>
    <p>Poll status:</p>
    <pre><code>curl '{base_url}/status/&lt;job_id&gt;'</code></pre>
  </div>
</main>
</body>
</html>
"""

# ---------- 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)