File size: 4,713 Bytes
4e09904 2f85585 4e09904 2f85585 4e09904 | 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | # -*- coding: utf-8 -*-
"""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: # noqa: BLE001
_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")
|