File size: 5,011 Bytes
736cf48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
136
137
138
139
140
141
142
143
144
145
146
"""Z-Anime 6B Image Generation (CPU) via sd-cli binary"""

import os, time, subprocess, tempfile, threading
from PIL import Image
import gradio as gr

# ---------------------------------------------------------------------------
# Model paths (downloaded at build time)
# ---------------------------------------------------------------------------
DIFFUSION = "/app/models/z-anime-8steps-q5_0.gguf"
LLM = "/app/models/qwen3_4b_q8_0.gguf"
VAE = "/app/models/ae.safetensors"

RESOLUTIONS = ["512x512", "768x512", "512x768"]
STEPS = 8
CFG = 1.0
TIMEOUT = 10800

_active_proc = None
_proc_lock = threading.Lock()

# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def generate(prompt, negative_prompt, resolution, seed):
    global _active_proc

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

    prompt = prompt.strip()[:500]
    w, h = (int(x) for x in resolution.split("x"))
    seed = int(seed or -1) if seed is not None else -1

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
        output_path = f.name

    cmd = [
        "/app/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",
        "--schedule", "beta",
        "-o", output_path,
        "--diffusion-fa",
        "--vae-tiling",
        "-v",
    ]
    if seed >= 0:
        cmd += ["-s", str(seed)]

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

    try:
        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()
            raise

        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 gr.Error("Out of memory (killed by OS). Try 512x512.")
            raise gr.Error(f"sd-cli failed (code {proc.returncode}): {err}")

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

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

    except subprocess.TimeoutExpired:
        with _proc_lock:
            _active_proc = None
        raise gr.Error(f"Generation timed out ({TIMEOUT//60} min limit)")
    except gr.Error:
        with _proc_lock:
            _active_proc = None
        raise
    except Exception as e:
        with _proc_lock:
            _active_proc = None
        raise gr.Error(f"Error: {e}")

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
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 8-step) via [sd.cpp](https://github.com/leejet/stable-diffusion.cpp) | "
        "Free CPU inference"
    )
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", lines=3,
                placeholder="anime girl with silver hair, fantasy armor, dramatic lighting")
            neg_input = gr.Textbox(label="Negative Prompt", lines=2,
                value="lowres, bad anatomy, bad hands, text, error, worst quality, blurry")
            with gr.Row():
                res_input = gr.Dropdown(choices=RESOLUTIONS, value="512x512",
                    label="Resolution")
                seed_input = gr.Number(value=-1, label="Seed (-1=random)", precision=0)
            gen_btn = gr.Button("Generate (8 steps, CFG 1)", 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=generate,
        inputs=[prompt_input, neg_input, res_input, seed_input],
        outputs=[output_img, status_box],
        concurrency_limit=1)

    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)