WeReCooking/Z-Anime-CPU
7
1"""Z-Anime 6B Image Generation (CPU/GPU) via sd-cli binary2 3CLI: python app.py "prompt" --seed 42 --cfg 1.04GUI: python app.py --gradio5"""6 7import os, sys, time, subprocess, tempfile, threading, argparse8 9# ---------------------------------------------------------------------------10# Paths — auto-detect local vs Docker11# ---------------------------------------------------------------------------12_LOCAL_MODELS = os.path.join(os.path.dirname(__file__), "models")13_DOCKER_MODELS = "/app/models"14MODELS_DIR = _LOCAL_MODELS if os.path.isdir(_LOCAL_MODELS) else _DOCKER_MODELS15 16_LOCAL_SDCLI = os.path.join(os.path.dirname(__file__), "..", "sd-cpp-tools", "build", "bin", "Release", "sd-cli.exe")17_DOCKER_SDCLI = "/app/sd-cli"18SD_CLI = _LOCAL_SDCLI if os.path.isfile(_LOCAL_SDCLI) else _DOCKER_SDCLI19 20DIFFUSION = os.path.join(MODELS_DIR, "z-anime-distill-4step-q5_0.gguf")21LLM = os.path.join(MODELS_DIR, "qwen3_4b_iq4xs.gguf")22VAE = os.path.join(MODELS_DIR, "ae.safetensors")23 24RESOLUTIONS = ["512x512", "768x512", "512x768"]25STEPS = 426TIMEOUT = 1080027 28_active_proc = None29_proc_lock = threading.Lock()30 31# ---------------------------------------------------------------------------32# Core generation (shared by CLI and GUI)33# ---------------------------------------------------------------------------34 35def generate_image(prompt, negative_prompt="", resolution="512x512",36 seed=-1, cfg=1.0, output_path=None):37 """Generate an anime image using Z-Anime 6B model.38 39 Args:40 prompt: Text description of the image to generate.41 negative_prompt: Things to avoid in the generated image.42 resolution: Image resolution (512x512, 768x512, or 512x768).43 seed: Random seed (-1 for random).44 cfg: CFG scale (1.0 recommended for distill, higher = slower).45 output_path: Where to save the image (auto if None).46 47 Returns:48 tuple: (output_path, status_message)49 """50 global _active_proc51 52 if not prompt or not prompt.strip():53 raise ValueError("Please enter a prompt.")54 55 prompt = prompt.strip()[:500]56 w, h = (int(x) for x in resolution.split("x"))57 seed = int(seed or -1)58 cfg = float(cfg)59 60 if output_path is None:61 f = tempfile.NamedTemporaryFile(suffix=".png", delete=False)62 output_path = f.name63 f.close()64 65 cmd = [66 SD_CLI,67 "--diffusion-model", DIFFUSION,68 "--llm", LLM,69 "--vae", VAE,70 "-p", prompt,71 "-n", negative_prompt or "",72 "-W", str(w),73 "-H", str(h),74 "--steps", str(STEPS),75 "--cfg-scale", str(cfg),76 "--sampling-method", "euler_a",77 "-o", output_path,78 "--diffusion-fa",79 "--diffusion-conv-direct",80 "--vae-tiling",81 "--vae-conv-direct",82 "--tensor-type-rules", "^vae=f32",83 "-v",84 ]85 if seed >= 0:86 cmd += ["-s", str(seed)]87 88 print(f"[gen] {w}x{h} steps={STEPS} cfg={cfg} seed={seed} prompt={prompt[:80]}")89 t0 = time.time()90 91 proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)92 with _proc_lock:93 _active_proc = proc94 95 try:96 stdout, stderr = proc.communicate(timeout=TIMEOUT)97 except subprocess.TimeoutExpired:98 proc.kill()99 proc.wait()100 with _proc_lock:101 _active_proc = None102 raise RuntimeError(f"Generation timed out ({TIMEOUT // 60} min limit)")103 104 elapsed = time.time() - t0105 with _proc_lock:106 _active_proc = None107 108 if proc.returncode != 0:109 err = stderr.decode(errors="replace")[-500:] if stderr else "Unknown error"110 if proc.returncode == -9:111 raise RuntimeError("Out of memory (killed by OS). Try 512x512.")112 raise RuntimeError(f"sd-cli failed (code {proc.returncode}): {err}")113 114 if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:115 raise RuntimeError("No output image generated")116 117 status = f"Generated in {elapsed:.1f}s ({w}x{h}, {STEPS} steps, cfg {cfg})"118 print(f"[gen] {status}")119 return output_path, status120 121 122# ---------------------------------------------------------------------------123# CLI mode124# ---------------------------------------------------------------------------125 126def cli_main():127 parser = argparse.ArgumentParser(description="Z-Anime 6B Image Generation")128 parser.add_argument("prompt", help="Text prompt for image generation")129 parser.add_argument("-n", "--negative", default="lowres, bad anatomy, bad hands, text, error, worst quality, blurry",130 help="Negative prompt")131 parser.add_argument("-r", "--resolution", default="512x512", choices=RESOLUTIONS)132 parser.add_argument("-s", "--seed", type=int, default=-1, help="Random seed (-1=random)")133 parser.add_argument("-c", "--cfg", type=float, default=1.0, help="CFG scale (1.0 recommended)")134 parser.add_argument("-o", "--output", default=None, help="Output file path")135 args = parser.parse_args()136 137 if args.output is None:138 args.output = f"z-anime_seed{args.seed}_cfg{args.cfg}.png"139 140 try:141 path, status = generate_image(142 prompt=args.prompt,143 negative_prompt=args.negative,144 resolution=args.resolution,145 seed=args.seed,146 cfg=args.cfg,147 output_path=args.output,148 )149 print(f" Output: {path}")150 except Exception as e:151 print(f"ERROR: {e}", file=sys.stderr)152 sys.exit(1)153 154 155# ---------------------------------------------------------------------------156# Gradio GUI mode157# ---------------------------------------------------------------------------158 159def gradio_main():160 import mmap161 from PIL import Image162 import gradio as gr163 164 # Warm up page cache165 print("[init] Preloading models into page cache...")166 t0 = time.time()167 for model_path in [DIFFUSION, LLM, VAE]:168 if os.path.exists(model_path):169 sz = os.path.getsize(model_path)170 with open(model_path, "rb") as f:171 mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)172 mm.read()173 mm.close()174 print(f" {os.path.basename(model_path)}: {sz / 1e9:.2f} GB cached")175 print(f"[init] Page cache warm in {time.time() - t0:.1f}s")176 177 def gui_generate(prompt, negative_prompt, resolution, cfg, seed):178 try:179 path, status = generate_image(prompt, negative_prompt, resolution,180 int(seed or -1), cfg=float(cfg or 1.0))181 return Image.open(path), status182 except Exception as e:183 raise gr.Error(str(e))184 185 with gr.Blocks(title="Z-Anime (CPU)") as demo:186 gr.Markdown(187 "**[Z-Anime 6B](https://huggingface.co/SeeSee21/Z-Anime)** S3-DiT Q5_0 GGUF "188 "(distill 4-step) via [sd.cpp](https://github.com/leejet/stable-diffusion.cpp) | "189 "~30 min at 512x512 on free CPU"190 )191 with gr.Row():192 with gr.Column():193 prompt_input = gr.Textbox(label="Prompt", lines=3,194 placeholder="An anime girl with long silver hair and sharp blue eyes, wearing ornate fantasy armor with glowing runes. She stands on a cliff overlooking a vast kingdom at sunset, wind catching her cape. Dramatic cinematic lighting, beautiful background art, detailed shading, professional anime illustration.")195 neg_input = gr.Textbox(label="Negative Prompt", lines=2,196 value="worst quality, low quality, lowres, blurry, bad anatomy, deformed hands, extra fingers, missing fingers, watermark, signature, text, error, censored")197 with gr.Row():198 res_input = gr.Dropdown(choices=RESOLUTIONS, value="512x512", label="Resolution")199 cfg_input = gr.Slider(minimum=1.0, maximum=1.5, value=1.0, step=0.1, label="CFG (1.0 best, max 1.5)")200 seed_input = gr.Number(value=-1, label="Seed (-1=random)", precision=0)201 gen_btn = gr.Button("Generate (4 steps)", variant="primary", size="lg")202 with gr.Column():203 output_img = gr.Image(type="pil", label="Output")204 status_box = gr.Textbox(label="Status", interactive=False)205 206 gen_btn.click(fn=gui_generate,207 inputs=[prompt_input, neg_input, res_input, cfg_input, seed_input],208 outputs=[output_img, status_box],209 concurrency_limit=1,210 api_name="generate")211 212 gr.Examples(213 examples=[214 ["An anime girl with long silver hair and sharp blue eyes, wearing ornate fantasy armor with glowing runes. She stands on a cliff overlooking a vast kingdom at sunset, wind catching her cape. Dramatic cinematic lighting, beautiful background art, detailed shading, professional anime illustration.",215 "worst quality, low quality, lowres, blurry, bad anatomy, deformed hands, extra fingers, fused fingers, missing fingers, bad proportions, wrong proportions, extra limbs, broken limbs, duplicate body parts, asymmetrical eyes, distorted face, warped features, poorly drawn face, mutated, extra eyes, cropped head, cut-off body, bad framing, jpeg artifacts, compression artifacts, watermark, logo, signature, text, error, noisy, oversmoothed, muddy colors, background clutter, censored, 3d, chibi, character doll, sepia, high contrast",216 "512x512", 1.0, -1],217 ],218 inputs=[prompt_input, neg_input, res_input, cfg_input, seed_input],219 outputs=[output_img, status_box],220 fn=gui_generate,221 cache_examples=True,222 cache_mode="lazy",223 )224 225 def _on_unload():226 with _proc_lock:227 proc = _active_proc228 if proc and proc.poll() is None:229 print("[cleanup] User disconnected, killing sd-cli process")230 proc.kill()231 232 demo.unload(_on_unload)233 234 demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True,235 theme="NoCrypt/miku", mcp_server=True)236 237 238# ---------------------------------------------------------------------------239# Entry point240# ---------------------------------------------------------------------------241 242if __name__ == "__main__":243 if len(sys.argv) > 1 and sys.argv[1] != "--gradio":244 cli_main()245 else:246 gradio_main()247