Update app.py

app.py

@@ -1,29 +1,26 @@
 import os
+import gc
 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
+# HARD CPU MODE
 # =========================
 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)
 
@@ -32,41 +29,33 @@ 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")
+dtype = torch.float16
 
 # =========================
 # 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)
+def download_gguf():
+    local_path = os.path.join(CACHE_DIR, GGUF_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,
+    return hf_hub_download(
+        repo_id=GGUF_REPO_ID,
+        filename=GGUF_FILENAME,
         cache_dir=CACHE_DIR,
         resume_download=True
     )
-    print("Download finished.")
-    return path
 
 # =========================
-# LOAD PIPELINE CPU ONLY
+# LOAD PIPELINE ULTRA LEAN
 # =========================
 def load_pipeline():
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
@@ -75,25 +64,15 @@ def load_pipeline():
         cache_dir=CACHE_DIR
     )
 
-    vae = AutoencoderKL.from_pretrained(
+    pipe = ZImagePipeline.from_pretrained(
         BASE_MODEL_ID,
-        subfolder="vae",
+        scheduler=scheduler,
         torch_dtype=dtype,
-        cache_dir=CACHE_DIR
+        cache_dir=CACHE_DIR,
+        low_cpu_mem_usage=True
     )
 
-    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)
+    gguf_path = download_gguf()
 
     transformer = ZImageTransformer2DModel.from_single_file(
         gguf_path,
@@ -101,36 +80,20 @@ def load_pipeline():
         torch_dtype=dtype
     ).to(device)
 
-    pipe = ZImagePipeline(
-        vae=vae.to(device),
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        transformer=transformer,
-        scheduler=scheduler
-    ).to(device)
+    pipe.transformer = transformer
 
-    pipe.unet.to(memory_format=torch.channels_last)
-    pipe.text_encoder.to(memory_format=torch.channels_last)
+    pipe.enable_attention_slicing()
+    pipe.enable_vae_slicing()
+    pipe.enable_sequential_cpu_offload()
 
-    pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True)
-    pipe.text_encoder = torch.compile(pipe.text_encoder, mode="max-autotune", fullgraph=True)
+    pipe = pipe.to(device)
 
     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
+# GENERATION (MIN RAM)
 # =========================
 def generate(prompt, seed, progress=gr.Progress()):
     if not prompt:
@@ -141,39 +104,40 @@ def generate(prompt, seed, progress=gr.Progress()):
 
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    total_steps = 4
-    start_time = time.time()
+    steps = 4
+    width = 256
+    height = 256
+
+    start = time.time()
 
-    def step_callback(step, timestep, latents):
-        elapsed = time.time() - start_time
+    def callback(step, timestep, latents):
         done = step + 1
+        elapsed = time.time() - start
         avg = elapsed / done
-        eta = avg * (total_steps - done)
-        progress(done / total_steps, desc=f"Step {done}/{total_steps} | ETA {eta:.1f}s")
+        eta = avg * (steps - done)
+        progress(done / steps, desc=f"Step {done}/{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()
+        gc.collect()
+        image = pipe(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_inference_steps=steps,
+            guidance_scale=1.0,
+            generator=generator,
+            callback=callback,
+            callback_steps=1
+        ).images[0]
+        gc.collect()
 
     return image, seed
 
 # =========================
-# UI + QUEUE
+# UI
 # =========================
-with gr.Blocks(title="Z-Image Turbo Q2_K CPU MAX") as demo:
-    gr.Markdown("# Z-Image Turbo Q2_K — FULL CPU MAX MODE")
+with gr.Blocks(title="Z-Image Turbo Ultra Lean CPU") as demo:
+    gr.Markdown("# Z-Image Turbo Q2_K — Ultra Lean 16GB CPU Mode")
 
     prompt = gr.Textbox(label="Prompt", lines=3)
     seed = gr.Number(label="Seed (-1 random)", value=-1, precision=0)
@@ -184,7 +148,8 @@ with gr.Blocks(title="Z-Image Turbo Q2_K CPU MAX") as demo:
 
     btn.click(generate, inputs=[prompt, seed], outputs=[image_out, seed_out])
 
-demo.queue(max_size=10, concurrency_count=1)
+demo.queue(max_size=5, concurrency_count=1)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+