Spaces:
Paused
Paused
File size: 4,277 Bytes
5efc5e5 | 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 | """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
@app.route("/")
def index():
return send_from_directory(APP_DIR, "index.html")
@app.route("/style.css")
def style():
return send_from_directory(APP_DIR, "style.css")
@app.route("/lliquid.gif")
def logo():
return send_from_directory(APP_DIR, "lliquid.gif")
@app.route("/api/ready")
def ready():
return {"ready": MODEL is not None}
@app.route("/api/generate")
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
@stream_with_context
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)
|