Spaces:
Paused
Paused
| """Masked-diffusion chat — server-side CPU inference. | |
| The model (LFM2.5-Encoder-350M-Diffusion, f32 safetensors) is loaded once and runs | |
| on the server's CPU. Nothing is downloaded to the visitor's browser: the page streams | |
| denoising frames from /api/generate as newline-delimited JSON and animates them, | |
| exactly like the old in-browser WebGPU build did. | |
| """ | |
| import json | |
| import os | |
| import threading | |
| import torch | |
| from flask import Flask, Response, request, send_from_directory, stream_with_context | |
| from tokenizers import Tokenizer | |
| import lfm2_diffusion as L | |
| APP_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| # Match threads to the cores actually granted to the container (cgroup quota), | |
| # not the host's core count — oversubscription badly slows CPU matmuls. | |
| _CPUS = L.effective_cpus() | |
| torch.set_num_threads(_CPUS) | |
| try: | |
| torch.set_num_interop_threads(1) | |
| except RuntimeError: | |
| pass | |
| app = Flask(__name__) | |
| # ---- load model + tokenizer once, at import (per gunicorn worker) ----------- | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| print(f"[boot] pulling {L.MODEL_ID} from the Hub, loading on {_CPUS} CPU thread(s)…", flush=True) | |
| from huggingface_hub import hf_hub_download | |
| TOK = Tokenizer.from_file(hf_hub_download(L.MODEL_ID, "tokenizer.json", token=HF_TOKEN)) | |
| MODEL = L.Lfm2Diffusion.from_hub(token=HF_TOKEN) | |
| print("[boot] calibrating + quantizing (smoothed int8)…", flush=True) | |
| MODEL.quantize_int8([TOK.encode(L.render_chat_prompt(p)).ids for p in L.CALIB_PROMPTS]) | |
| _GEN_LOCK = threading.Lock() | |
| # warm up the graph so the first real request is not the slow one | |
| list(L.diffuse(MODEL, TOK.encode(L.render_chat_prompt("Hi")).ids, | |
| max_new=8, steps=4, block_size=8)) | |
| print("[boot] model ready", flush=True) | |
| def _clampi(v, lo, hi, default): | |
| try: | |
| return max(lo, min(hi, int(float(v)))) | |
| except (TypeError, ValueError): | |
| return default | |
| def index(): | |
| return send_from_directory(APP_DIR, "index.html") | |
| def style(): | |
| return send_from_directory(APP_DIR, "style.css") | |
| def logo(): | |
| return send_from_directory(APP_DIR, "lliquid.gif") | |
| def ready(): | |
| return {"ready": MODEL is not None} | |
| def generate(): | |
| prompt = (request.args.get("prompt") or "").strip() | |
| if not prompt: | |
| return {"error": "empty prompt"}, 400 | |
| opts = dict( | |
| max_new=_clampi(request.args.get("mnt"), 16, 192, 64), | |
| steps=_clampi(request.args.get("steps"), 8, 128, 32), | |
| block_size=_clampi(request.args.get("blk"), 8, 64, 16), | |
| temperature=_clampi(request.args.get("temp"), 0, 100, 0) / 100.0, | |
| tau=_clampi(request.args.get("tau"), 50, 100, 90) / 100.0, | |
| ) | |
| prompt_ids = TOK.encode(L.render_chat_prompt(prompt)).ids | |
| def frames(): | |
| # Serialize generation: the CPU model handles one request at a time. | |
| if not _GEN_LOCK.acquire(blocking=False): | |
| yield json.dumps({"error": "The model is busy with another request. Try again in a moment."}) + "\n" | |
| return | |
| try: | |
| last = None | |
| for fr in L.diffuse(MODEL, prompt_ids, **opts): | |
| last = fr["gen_ids"] | |
| tokens = [None if t == L.MASK_ID | |
| else TOK.decode([t], skip_special_tokens=True) | |
| for t in fr["gen_ids"]] | |
| yield json.dumps({ | |
| "step": fr["step"], | |
| "tokens": tokens, | |
| "revealed": sorted(fr["revealed"]), | |
| }) + "\n" | |
| gen = [t for t in (last or []) if t != L.MASK_ID] | |
| text = L.trim(TOK.decode(gen, skip_special_tokens=True)) | |
| yield json.dumps({"done": True, "text": text or "(empty)"}) + "\n" | |
| except Exception as exc: # surface errors to the UI | |
| yield json.dumps({"error": str(exc)}) + "\n" | |
| finally: | |
| _GEN_LOCK.release() | |
| return Response(frames(), mimetype="application/x-ndjson", headers={ | |
| "Cache-Control": "no-cache, no-transform", | |
| "X-Accel-Buffering": "no", | |
| }) | |
| if __name__ == "__main__": | |
| app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), threaded=True) | |