| |
| """POCKET-Zimage-CPU — FastAPI demo (CPU only, no GPU). |
| Serves index.html and runs stable-diffusion.cpp on CPU. |
| """ |
| import os, re, json, time, base64, shutil, subprocess, threading, uuid, urllib.request |
| from pathlib import Path |
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import FileResponse, JSONResponse |
| from fastapi.staticfiles import StaticFiles |
| from pydantic import BaseModel, Field |
|
|
| ROOT = Path(__file__).parent |
| MODELS = ROOT / "models" |
| OUT = ROOT / "outputs" |
| MODELS.mkdir(exist_ok=True) |
| OUT.mkdir(exist_ok=True) |
|
|
| SD_BIN = os.environ.get("SD_BIN", str(ROOT / "bin" / "sd-cli")) |
| THREADS = int(os.environ.get("SD_THREADS", str(min(8, os.cpu_count() or 4)))) |
|
|
| FILES = { |
| "diffusion": ("z_image_turbo-Q4_0-pocket.gguf", |
| "https://huggingface.co/FINAL-Bench/POCKET-Zimage-CPU/resolve/main/z_image_turbo-Q4_0-pocket.gguf"), |
| "vae": ("ae.safetensors", |
| "https://huggingface.co/Tongyi-MAI/Z-Image-Turbo/resolve/main/vae/diffusion_pytorch_model.safetensors"), |
| "llm": ("Qwen3-4B-Instruct-2507-Q4_K_M.gguf", |
| "https://huggingface.co/lmstudio-community/Qwen3-4B-Instruct-2507-GGUF/resolve/main/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"), |
| } |
|
|
| _state = {"ready": False, "msg": "idle", "pct": 0} |
| _lock = threading.Lock() |
|
|
|
|
| HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") |
|
|
|
|
| def _fetch(name: str, url: str): |
| dst = MODELS / name |
| if dst.exists() and dst.stat().st_size > 10_000_000: |
| return |
| tmp = dst.with_suffix(dst.suffix + ".part") |
| hdr = {"User-Agent": "pocket-zimage/1.0"} |
| if HF_TOKEN and "huggingface.co" in url: |
| hdr["Authorization"] = "Bearer " + HF_TOKEN |
| req = urllib.request.Request(url, headers=hdr) |
| with urllib.request.urlopen(req) as r, open(tmp, "wb") as f: |
| total = int(r.headers.get("Content-Length") or 0) |
| done = 0 |
| while True: |
| chunk = r.read(1 << 20) |
| if not chunk: |
| break |
| f.write(chunk) |
| done += len(chunk) |
| if total: |
| _state["pct"] = int(done * 100 / total) |
| _state["msg"] = f"downloading {name} {_state['pct']}%" |
| tmp.rename(dst) |
|
|
|
|
| def _prepare(): |
| try: |
| for key, (name, url) in FILES.items(): |
| _state["msg"] = f"fetching {name}" |
| _fetch(name, url) |
| _state.update(ready=True, msg="ready", pct=100) |
| except Exception as e: |
| _state.update(ready=False, msg=f"error: {e}") |
|
|
|
|
| threading.Thread(target=_prepare, daemon=True).start() |
|
|
| app = FastAPI(title="POCKET-Zimage-CPU") |
|
|
|
|
| class GenReq(BaseModel): |
| prompt: str = Field(min_length=1, max_length=600) |
| steps: int = Field(default=3, ge=1, le=8) |
| size: int = Field(default=512) |
| seed: int = Field(default=-1) |
|
|
|
|
| @app.get("/api/status") |
| def status(): |
| return {**_state, "threads": THREADS, "gpu": False} |
|
|
|
|
| @app.post("/api/generate") |
| def generate(req: GenReq): |
| if not _state["ready"]: |
| raise HTTPException(503, _state["msg"]) |
| if req.size not in (384, 512, 768): |
| raise HTTPException(400, "size must be 384, 512 or 768") |
| if not _lock.acquire(blocking=False): |
| raise HTTPException(429, "busy - one image at a time on CPU") |
| try: |
| tag = uuid.uuid4().hex[:10] |
| dst = OUT / f"{tag}.png" |
| seed = req.seed if req.seed >= 0 else int(time.time()) % 2_000_000_000 |
| cmd = [ |
| SD_BIN, |
| "--diffusion-model", str(MODELS / FILES["diffusion"][0]), |
| "--vae", str(MODELS / FILES["vae"][0]), |
| "--llm", str(MODELS / FILES["llm"][0]), |
| "-p", req.prompt, |
| "--cfg-scale", "1.0", |
| "--steps", str(req.steps), |
| "--fa", "--vae-tiling", |
| "-t", str(THREADS), |
| "-H", str(req.size), "-W", str(req.size), |
| "--seed", str(seed), |
| "-o", str(dst), |
| ] |
| t0 = time.time() |
| p = subprocess.run(cmd, capture_output=True, text=True, timeout=900) |
| took = round(time.time() - t0, 1) |
| if not dst.exists(): |
| tail = (p.stderr or p.stdout or "")[-400:] |
| raise HTTPException(500, f"generation failed: {tail}") |
| b64 = base64.b64encode(dst.read_bytes()).decode() |
| dst.unlink(missing_ok=True) |
| return JSONResponse({"image": f"data:image/png;base64,{b64}", |
| "seconds": took, "seed": seed, |
| "steps": req.steps, "size": req.size}) |
| finally: |
| _lock.release() |
|
|
|
|
| @app.get("/") |
| def index(): |
| return FileResponse(ROOT / "index.html") |
|
|
|
|
| app.mount("/static", StaticFiles(directory=ROOT / "static"), name="static") |
|
|
