Update app.py

app.py

@@ -7,7 +7,6 @@ import torch
 import gradio as gr
 from threading import Lock
 from contextlib import contextmanager
-from huggingface_hub import snapshot_download
 
 # --- LOGGING FOR UI ---
 LOG_BUFFER = []
@@ -21,10 +20,9 @@ def log(message):
             LOG_BUFFER.pop(0)
     return "\n".join(LOG_BUFFER)
 
-# 🚀 Initialization
 _initial_logs = log("🚀 Initializing Ultimate Z-Image Turbo CPU Edition...")
 
-# --- ENVIRONMENT SETUP ---
+# CPU THREAD OPTIMIZATION
 CPU_THREADS = min(8, os.cpu_count() or 1)
 os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
 os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
@@ -36,9 +34,6 @@ os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 os.environ["TRANSFORMERS_CACHE"] = "./hf_cache"
 os.environ["HF_DATASETS_CACHE"] = "./hf_cache"
-os.environ["HF_HUB_OFFLINE"] = "1"
-os.environ["TRANSFORMERS_OFFLINE"] = "1"
-os.environ["HF_DATASETS_OFFLINE"] = "1"
 
 torch.set_num_threads(CPU_THREADS)
 torch.set_grad_enabled(False)
@@ -65,23 +60,6 @@ pipe = None
 _pipe_lock = Lock()
 _generation_lock = Lock()
 
-# --- Pre-download full snapshot once ---
-MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
-MODEL_LOCAL = os.path.join(CACHE_DIR, "Z-Image-Turbo-snapshot")
-os.makedirs(MODEL_LOCAL, exist_ok=True)
-
-if not os.listdir(MODEL_LOCAL):
-    log("📥 Downloading full model snapshot, please wait...")
-    snapshot_download(
-        repo_id=MODEL_ID,
-        cache_dir=MODEL_LOCAL,
-        local_dir=MODEL_LOCAL,
-        local_dir_use_symlinks=False
-    )
-    log(f"📦 Model snapshot cached at: {MODEL_LOCAL}")
-else:
-    log(f"📦 Model snapshot already exists at: {MODEL_LOCAL}")
-
 @contextmanager
 def managed_memory():
     try:
@@ -102,9 +80,9 @@ def load_pipeline():
         start_load = time.time()
 
         pipe = ZImagePipeline.from_pretrained(
-            MODEL_LOCAL,
+            "Tongyi-MAI/Z-Image-Turbo",
             torch_dtype=DTYPE,
-            local_files_only=True,
+            cache_dir=CACHE_DIR,
             low_cpu_mem_usage=True
         )
 
@@ -153,14 +131,15 @@ def generate(prompt, quality_mode, seed, progress=gr.Progress()):
             generator = torch.Generator("cpu").manual_seed(seed)
             start_time = time.time()
 
-            def progress_callback(step, *_):
+            def diffusers_progress_callback(pipeline, step_index, timestep, callback_kwargs):
                 elapsed = time.time() - start_time
-                avg = elapsed / (step + 1) if step >= 0 else 0
-                remaining = avg * (steps - step - 1)
+                avg = elapsed / (step_index + 1) if step_index >= 0 else 0
+                remaining = avg * (steps - step_index - 1)
                 progress(
-                    (step + 1) / steps,
-                    desc=f"Step {step+1}/{steps} | ETA {remaining:.1f}s"
+                    (step_index + 1) / steps,
+                    desc=f"Step {step_index+1}/{steps} | ETA {remaining:.1f}s"
                 )
+                return callback_kwargs
 
             result = pipe(
                 prompt=prompt,
@@ -170,8 +149,8 @@ def generate(prompt, quality_mode, seed, progress=gr.Progress()):
                 num_inference_steps=steps,
                 guidance_scale=0.0,
                 generator=generator,
-                callback=progress_callback,
-                callback_steps=1,
+                callback_on_step_end=diffusers_progress_callback,
+                callback_on_step_end_tensor_inputs=["latents"],
                 output_type="pil"
             )
 
@@ -184,7 +163,6 @@ def generate(prompt, quality_mode, seed, progress=gr.Progress()):
 
             return image, seed
 
-# --- GRADIO UI ---
 with gr.Blocks(title="🚀 Z-Image Turbo Pro Max + Live Logs") as demo:
     gr.Markdown("## GPU‑FREE CPU Turbo — Live Logs Below")
 
@@ -215,7 +193,7 @@ with gr.Blocks(title="🚀 Z-Image Turbo Pro Max + Live Logs") as demo:
 
     def wrapped_generate(prompt, quality_mode, seed):
         image, used_seed = generate(prompt, quality_mode, seed)
-        logs = log(f"🧠 Latest status: Finished generation.")
+        logs = log("🧠 Latest status: Finished generation.")
         return image, used_seed, logs
 
     generate_btn.click(