app.py

"""Z-Anime 6B Image Generation (CPU/GPU) via sd-cli binary

CLI:  python app.py "prompt" --seed 42 --cfg 1.0
GUI:  python app.py --gradio
"""

import os, sys, time, subprocess, tempfile, threading, argparse

# ---------------------------------------------------------------------------
# Paths — auto-detect local vs Docker
# ---------------------------------------------------------------------------
_LOCAL_MODELS = os.path.join(os.path.dirname(__file__), "models")
_DOCKER_MODELS = "/app/models"
MODELS_DIR = _LOCAL_MODELS if os.path.isdir(_LOCAL_MODELS) else _DOCKER_MODELS

_LOCAL_SDCLI = os.path.join(os.path.dirname(__file__), "..", "sd-cpp-tools", "build", "bin", "Release", "sd-cli.exe")
_DOCKER_SDCLI = "/app/sd-cli"
SD_CLI = _LOCAL_SDCLI if os.path.isfile(_LOCAL_SDCLI) else _DOCKER_SDCLI

DIFFUSION = os.path.join(MODELS_DIR, "z-anime-distill-4step-q5_0.gguf")
LLM = os.path.join(MODELS_DIR, "qwen3_4b_iq4xs.gguf")
VAE = os.path.join(MODELS_DIR, "ae.safetensors")

RESOLUTIONS = ["512x512", "768x512", "512x768"]
STEPS = 4
TIMEOUT = 10800

_active_proc = None
_proc_lock = threading.Lock()

# ---------------------------------------------------------------------------
# Core generation (shared by CLI and GUI)
# ---------------------------------------------------------------------------

def generate_image(prompt, negative_prompt="", resolution="512x512",
                   seed=-1, cfg=1.0, output_path=None):
    """Generate an anime image using Z-Anime 6B model.

    Args:
        prompt: Text description of the image to generate.
        negative_prompt: Things to avoid in the generated image.
        resolution: Image resolution (512x512, 768x512, or 512x768).
        seed: Random seed (-1 for random).
        cfg: CFG scale (1.0 recommended for distill, higher = slower).
        output_path: Where to save the image (auto if None).

    Returns:
        tuple: (output_path, status_message)
    """
    global _active_proc

    if not prompt or not prompt.strip():
        raise ValueError("Please enter a prompt.")

    prompt = prompt.strip()[:500]
    w, h = (int(x) for x in resolution.split("x"))
    seed = int(seed or -1)
    cfg = float(cfg)

    if output_path is None:
        f = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        output_path = f.name
        f.close()

    cmd = [
        SD_CLI,
        "--diffusion-model", DIFFUSION,
        "--llm", LLM,
        "--vae", VAE,
        "-p", prompt,
        "-n", negative_prompt or "",
        "-W", str(w),
        "-H", str(h),
        "--steps", str(STEPS),
        "--cfg-scale", str(cfg),
        "--sampling-method", "euler_a",
        "-o", output_path,
        "--diffusion-fa",
        "--diffusion-conv-direct",
        "--vae-tiling",
        "--vae-conv-direct",
        "--tensor-type-rules", "^vae=f32",
        "-v",
    ]
    if seed >= 0:
        cmd += ["-s", str(seed)]

    print(f"[gen] {w}x{h} steps={STEPS} cfg={cfg} seed={seed} prompt={prompt[:80]}")
    t0 = time.time()

    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    with _proc_lock:
        _active_proc = proc

    try:
        stdout, stderr = proc.communicate(timeout=TIMEOUT)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait()
        with _proc_lock:
            _active_proc = None
        raise RuntimeError(f"Generation timed out ({TIMEOUT // 60} min limit)")

    elapsed = time.time() - t0
    with _proc_lock:
        _active_proc = None

    if proc.returncode != 0:
        err = stderr.decode(errors="replace")[-500:] if stderr else "Unknown error"
        if proc.returncode == -9:
            raise RuntimeError("Out of memory (killed by OS). Try 512x512.")
        raise RuntimeError(f"sd-cli failed (code {proc.returncode}): {err}")

    if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
        raise RuntimeError("No output image generated")

    status = f"Generated in {elapsed:.1f}s ({w}x{h}, {STEPS} steps, cfg {cfg})"
    print(f"[gen] {status}")
    return output_path, status


# ---------------------------------------------------------------------------
# CLI mode
# ---------------------------------------------------------------------------

def cli_main():
    parser = argparse.ArgumentParser(description="Z-Anime 6B Image Generation")
    parser.add_argument("prompt", help="Text prompt for image generation")
    parser.add_argument("-n", "--negative", default="lowres, bad anatomy, bad hands, text, error, worst quality, blurry",
                        help="Negative prompt")
    parser.add_argument("-r", "--resolution", default="512x512", choices=RESOLUTIONS)
    parser.add_argument("-s", "--seed", type=int, default=-1, help="Random seed (-1=random)")
    parser.add_argument("-c", "--cfg", type=float, default=1.0, help="CFG scale (1.0 recommended)")
    parser.add_argument("-o", "--output", default=None, help="Output file path")
    args = parser.parse_args()

    if args.output is None:
        args.output = f"z-anime_seed{args.seed}_cfg{args.cfg}.png"

    try:
        path, status = generate_image(
            prompt=args.prompt,
            negative_prompt=args.negative,
            resolution=args.resolution,
            seed=args.seed,
            cfg=args.cfg,
            output_path=args.output,
        )
        print(f"  Output: {path}")
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


# ---------------------------------------------------------------------------
# Gradio GUI mode
# ---------------------------------------------------------------------------

def gradio_main():
    import mmap
    from PIL import Image
    import gradio as gr

    # Warm up page cache
    print("[init] Preloading models into page cache...")
    t0 = time.time()
    for model_path in [DIFFUSION, LLM, VAE]:
        if os.path.exists(model_path):
            sz = os.path.getsize(model_path)
            with open(model_path, "rb") as f:
                mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
                mm.read()
                mm.close()
            print(f"  {os.path.basename(model_path)}: {sz / 1e9:.2f} GB cached")
    print(f"[init] Page cache warm in {time.time() - t0:.1f}s")

    def gui_generate(prompt, negative_prompt, resolution, cfg, seed):
        try:
            path, status = generate_image(prompt, negative_prompt, resolution,
                                          int(seed or -1), cfg=float(cfg or 1.0))
            return Image.open(path), status
        except Exception as e:
            raise gr.Error(str(e))

    with gr.Blocks(title="Z-Anime (CPU)") as demo:
        gr.Markdown(
            "**[Z-Anime 6B](https://huggingface.co/SeeSee21/Z-Anime)** S3-DiT Q5_0 GGUF "
            "(distill 4-step) via [sd.cpp](https://github.com/leejet/stable-diffusion.cpp) | "
            "~30 min at 512x512 on free CPU"
        )
        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(label="Prompt", lines=3,
                    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.")
                neg_input = gr.Textbox(label="Negative Prompt", lines=2,
                    value="worst quality, low quality, lowres, blurry, bad anatomy, deformed hands, extra fingers, missing fingers, watermark, signature, text, error, censored")
                with gr.Row():
                    res_input = gr.Dropdown(choices=RESOLUTIONS, value="512x512", label="Resolution")
                    cfg_input = gr.Slider(minimum=1.0, maximum=1.5, value=1.0, step=0.1, label="CFG (1.0 best, max 1.5)")
                    seed_input = gr.Number(value=-1, label="Seed (-1=random)", precision=0)
                gen_btn = gr.Button("Generate (4 steps)", variant="primary", size="lg")
            with gr.Column():
                output_img = gr.Image(type="pil", label="Output")
                status_box = gr.Textbox(label="Status", interactive=False)

        gen_btn.click(fn=gui_generate,
            inputs=[prompt_input, neg_input, res_input, cfg_input, seed_input],
            outputs=[output_img, status_box],
            concurrency_limit=1,
            api_name="generate")

        gr.Examples(
            examples=[
                ["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.",
                 "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",
                 "512x512", 1.0, -1],
            ],
            inputs=[prompt_input, neg_input, res_input, cfg_input, seed_input],
            outputs=[output_img, status_box],
            fn=gui_generate,
            cache_examples=True,
            cache_mode="lazy",
        )

        def _on_unload():
            with _proc_lock:
                proc = _active_proc
            if proc and proc.poll() is None:
                print("[cleanup] User disconnected, killing sd-cli process")
                proc.kill()

        demo.unload(_on_unload)

    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True,
                theme="NoCrypt/miku", mcp_server=True)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] != "--gradio":
        cli_main()
    else:
        gradio_main()