Create app.py

app.py

@@ -0,0 +1,190 @@
+import os
+import time
+import random
+import gc
+import torch
+import gradio as gr
+
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, AutoModel
+from diffusers import (
+    ZImagePipeline,
+    ZImageTransformer2DModel,
+    GGUFQuantizationConfig,
+    AutoencoderKL,
+    FlowMatchEulerDiscreteScheduler
+)
+
+# =========================
+# FORCE CPU ENV
+# =========================
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+cpu_cores = os.cpu_count() or 1
+
+torch.set_num_threads(cpu_cores)
+torch.set_num_interop_threads(cpu_cores)
+
+os.environ["OMP_NUM_THREADS"] = str(cpu_cores)
+os.environ["MKL_NUM_THREADS"] = str(cpu_cores)
+
+torch.backends.mkldnn.enabled = True
+torch.backends.quantized.engine = "fbgemm"
+torch.backends.cudnn.enabled = False
+torch.set_float32_matmul_precision("high")
+
+dtype = torch.float32
+device = torch.device("cpu")
+
+# =========================
+# MODEL CONFIG
+# =========================
+BASE_MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
+TEXT_ENCODER_ID = "Qwen/Qwen3-4B"
+GGUF_REPO_ID = "unsloth/Z-Image-Turbo-GGUF"
+GGUF_FILENAME = "z-image-turbo-Q2_K.gguf"
+CACHE_DIR = "models"
+
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+def download_if_needed(repo_id, filename):
+    local_path = os.path.join(CACHE_DIR, filename)
+    if os.path.exists(local_path):
+        print("Model cached locally.")
+        return local_path
+
+    print("Downloading model (first run)...")
+    path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        cache_dir=CACHE_DIR,
+        resume_download=True
+    )
+    print("Download finished.")
+    return path
+
+# =========================
+# LOAD PIPELINE CPU ONLY
+# =========================
+def load_pipeline():
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        BASE_MODEL_ID,
+        subfolder="scheduler",
+        cache_dir=CACHE_DIR
+    )
+
+    vae = AutoencoderKL.from_pretrained(
+        BASE_MODEL_ID,
+        subfolder="vae",
+        torch_dtype=dtype,
+        cache_dir=CACHE_DIR
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        TEXT_ENCODER_ID,
+        cache_dir=CACHE_DIR
+    )
+
+    text_encoder = AutoModel.from_pretrained(
+        TEXT_ENCODER_ID,
+        torch_dtype=dtype,
+        cache_dir=CACHE_DIR
+    ).to(device)
+
+    gguf_path = download_if_needed(GGUF_REPO_ID, GGUF_FILENAME)
+
+    transformer = ZImageTransformer2DModel.from_single_file(
+        gguf_path,
+        quantization_config=GGUFQuantizationConfig(compute_dtype=dtype),
+        torch_dtype=dtype
+    ).to(device)
+
+    pipe = ZImagePipeline(
+        vae=vae.to(device),
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        transformer=transformer,
+        scheduler=scheduler
+    ).to(device)
+
+    pipe.unet.to(memory_format=torch.channels_last)
+    pipe.text_encoder.to(memory_format=torch.channels_last)
+
+    pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
+    pipe.text_encoder = torch.compile(pipe.text_encoder, mode="max-autotune", fullgraph=True)
+
+    return pipe
+
+pipe = load_pipeline()
+
+# Warmup compile
+with torch.inference_mode():
+    _ = pipe(
+        prompt="warmup",
+        width=256,
+        height=256,
+        num_inference_steps=1,
+        guidance_scale=1.0
+    )
+
+# =========================
+# GENERATION WITH PROGRESS
+# =========================
+def generate(prompt, seed, progress=gr.Progress()):
+    if not prompt:
+        raise gr.Error("Prompt required")
+
+    if seed < 0:
+        seed = random.randint(0, 2**31 - 1)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    total_steps = 4
+    start_time = time.time()
+
+    def step_callback(step, timestep, latents):
+        elapsed = time.time() - start_time
+        done = step + 1
+        avg = elapsed / done
+        eta = avg * (total_steps - done)
+        progress(done / total_steps, desc=f"Step {done}/{total_steps} | ETA {eta:.1f}s")
+
+    with torch.inference_mode():
+        gc.disable()
+        try:
+            image = pipe(
+                prompt=prompt,
+                width=256,
+                height=256,
+                num_inference_steps=total_steps,
+                guidance_scale=1.0,
+                generator=generator,
+                callback=step_callback,
+                callback_steps=1
+            ).images[0]
+        finally:
+            gc.enable()
+
+    return image, seed
+
+# =========================
+# UI + QUEUE
+# =========================
+with gr.Blocks(title="Z-Image Turbo Q2_K CPU MAX") as demo:
+    gr.Markdown("# Z-Image Turbo Q2_K — FULL CPU MAX MODE")
+
+    prompt = gr.Textbox(label="Prompt", lines=3)
+    seed = gr.Number(label="Seed (-1 random)", value=-1, precision=0)
+    btn = gr.Button("Generate")
+
+    image_out = gr.Image()
+    seed_out = gr.Number(interactive=False)
+
+    btn.click(generate, inputs=[prompt, seed], outputs=[image_out, seed_out])
+
+demo.queue(max_size=10, concurrency_count=1)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)