Update app.py

app.py

@@ -5,60 +5,93 @@ import time
 import random
 import torch
 import gradio as gr
-from threading import Lock
+from threading import Lock, Event
 from contextlib import contextmanager
+from huggingface_hub import snapshot_download, LocalEntryNotFoundError
+
+# ----------- LOGGING -----------
 
-# --- LOGGING FOR UI ---
 LOG_BUFFER = []
 LOG_LOCK = Lock()
 
-def log(message):
-    print(message)
+def log(msg):
     with LOG_LOCK:
-        LOG_BUFFER.append(f"{time.strftime('%H:%M:%S')} | {message}")
+        timestamp = time.strftime('%H:%M:%S')
+        LOG_BUFFER.append(f"{timestamp} | {msg}")
         if len(LOG_BUFFER) > 500:
             LOG_BUFFER.pop(0)
+        print(msg)
     return "\n".join(LOG_BUFFER)
 
-_initial_logs = log("🚀 Initializing Ultimate Z-Image Turbo CPU Edition...")
+# ----------- ENV CONFIG -----------
 
-# 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)
-os.environ["OPENBLAS_NUM_THREADS"] = str(CPU_THREADS)
-os.environ["VECLIB_MAXIMUM_THREADS"] = str(CPU_THREADS)
-os.environ["NUMEXPR_NUM_THREADS"] = str(CPU_THREADS)
+for var in ["OMP_NUM_THREADS","MKL_NUM_THREADS","OPENBLAS_NUM_THREADS","VECLIB_MAXIMUM_THREADS","NUMEXPR_NUM_THREADS"]:
+    os.environ[var] = str(CPU_THREADS)
+
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+os.environ["HF_HUB_OFFLINE"] = "1"
+os.environ["TRANSFORMERS_OFFLINE"] = "1"
 os.environ["TRANSFORMERS_CACHE"] = "./hf_cache"
 os.environ["HF_DATASETS_CACHE"] = "./hf_cache"
 
-torch.set_num_threads(CPU_THREADS)
 torch.set_grad_enabled(False)
+torch.set_num_threads(CPU_THREADS)
 torch.backends.mkldnn.enabled = True
-torch.backends.mkldnn.deterministic = False
-torch.set_flush_denormal(True)
 torch.set_float32_matmul_precision("medium")
 
 DEVICE = "cpu"
 DTYPE = torch.float32
-CACHE_DIR = "./hf_cache"
-os.makedirs(CACHE_DIR, exist_ok=True)
-
-log(f"⚡ CPU Threads: {CPU_THREADS}, Device: {DEVICE}, DType: {DTYPE}")
+os.makedirs("./hf_cache", exist_ok=True)
 
 try:
     from diffusers import ZImagePipeline
-    log("📦 diffusers imported successfully")
+    log("Imported diffusers successfully.")
 except ImportError as e:
-    log(f"❌ Import Error: {e}")
+    log(f"Import diffusers failed: {e}")
     sys.exit(1)
 
-pipe = None
-_pipe_lock = Lock()
-_generation_lock = Lock()
+pipe_cache = {}
+pipe_lock = Lock()
+generation_lock = Lock()
+interrupt_event = Event()
+
+# ----------- SNAPSHOT WITH RETRY -----------
+
+MODEL_SPECS = {
+    "Z-Image Turbo": "Tongyi-MAI/Z-Image-Turbo",
+    # Optionally add quantized variants here
+    # "Z-Image Turbo GGUF": "unsloth/Z-Image-Turbo-GGUF",
+}
+
+def download_snapshot_with_retry(repo_id, local_path, retries=3):
+    attempt = 1
+    while attempt <= retries:
+        log(f"Snapshot attempt {attempt}/{retries} for {repo_id}...")
+        try:
+            # snapshot_download respects HF cache and will skip downloads if cached
+            path = snapshot_download(repo_id=repo_id, local_dir=local_path, local_dir_use_symlinks=False)
+            log(f"Snapshot fully downloaded: {path}")
+            return path
+        except Exception as e:
+            log(f"⚠️ snapshot_download failed: {e}")
+            attempt += 1
+            time.sleep(2)
+    raise RuntimeError(f"Failed to download snapshot of {repo_id} after {retries} attempts")
+
+# Ensure snapshot is present
+for model_name, repo_id in MODEL_SPECS.items():
+    local_dir = os.path.join("./hf_cache", f"{model_name}_snapshot")
+    if not os.path.isdir(local_dir) or not os.listdir(local_dir):
+        log(f"📥 No snapshot for {model_name}, starting download...")
+        try:
+            download_snapshot_with_retry(repo_id, local_dir, retries=3)
+        except RuntimeError as err:
+            log(f"❌ Snapshot download error: {err}")
+
+# ----------- PIPELINE LOADING -----------
 
 @contextmanager
 def managed_memory():
@@ -70,77 +103,68 @@ def managed_memory():
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-def load_pipeline():
-    global pipe
-    with _pipe_lock:
-        if pipe is not None:
-            return pipe
-
-        log("📦 Loading Z-Image Turbo pipeline...")
-        start_load = time.time()
-
-        pipe = ZImagePipeline.from_pretrained(
-            "Tongyi-MAI/Z-Image-Turbo",
-            torch_dtype=DTYPE,
-            cache_dir=CACHE_DIR,
-            low_cpu_mem_usage=True
-        )
-
-        pipe = pipe.to(DEVICE)
+def load_pipeline(model_name):
+    if model_name in pipe_cache:
+        return pipe_cache[model_name]
+    with pipe_lock:
+        log(f"Loading {model_name} pipeline.")
+        repo_dir = os.path.join("./hf_cache", f"{model_name}_snapshot")
+        try:
+            pipe = ZImagePipeline.from_pretrained(repo_dir, torch_dtype=DTYPE, local_files_only=True, low_cpu_mem_usage=True)
+        except LocalEntryNotFoundError:
+            log(f"Incomplete local snapshot for {model_name}, retrying online load.")
+            pipe = ZImagePipeline.from_pretrained(MODEL_SPECS[model_name], torch_dtype=DTYPE, cache_dir="./hf_cache", low_cpu_mem_usage=True)
+        pipe.to(DEVICE)
         pipe.vae.eval()
         pipe.text_encoder.eval()
         pipe.transformer.eval()
-
         try:
-            pipe.transformer = torch.compile(
-                pipe.transformer,
-                mode="reduce-overhead",
-                fullgraph=False,
-                dynamic=False
-            )
-            log("✅ Transformer compiled successfully!")
-        except Exception as compile_error:
-            log(f"⚠️ torch.compile() failed: {compile_error}")
-
-        load_time = time.time() - start_load
-        log(f"✅ Pipeline loaded in {load_time:.2f}s")
+            pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")
+            log("Transformer compiled.")
+        except Exception as e:
+            log(f"Transformer compile skipped: {e}")
+        pipe_cache[model_name] = pipe
         return pipe
 
+# ----------- GENERATION LOGIC -----------
+
 @torch.inference_mode()
 @torch.no_grad()
-def generate(prompt, quality_mode, seed, progress=gr.Progress()):
+def generate(prompt, quality_mode, seed, model_name):
     if not prompt.strip():
-        raise gr.Error("🎯 Prompt cannot be empty!")
-
-    quality_settings = {
-        "ultra_fast": {"steps": 1, "width": 256, "height": 256},
-        "fast": {"steps": 1, "width": 256, "height": 256},
-        "balanced": {"steps": 2, "width": 256, "height": 256},
-        "quality": {"steps": 4, "width": 384, "height": 384},
-        "ultra_quality": {"steps": 4, "width": 512, "height": 512}
+        raise gr.Error("Prompt cannot be empty!")
+
+    PRESETS = {
+        "ultra_fast": (1, 256),
+        "fast": (1, 256),
+        "balanced": (2, 256),
+        "quality": (4, 384),
+        "ultra_quality": (4, 512),
     }
-    settings = quality_settings.get(quality_mode, quality_settings["fast"])
-    steps, width, height = settings["steps"], settings["width"], settings["height"]
-
-    seed = int(seed) if seed >= 0 else random.randint(0, 2**31 - 1)
-    log(f"🎨 Generating: '{prompt[:50]}...' | Mode: {quality_mode} | {width}x{height} | Seed: {seed}")
-
-    with managed_memory():
-        with _generation_lock:
-            pipe = load_pipeline()
-            generator = torch.Generator("cpu").manual_seed(seed)
-            start_time = time.time()
-
-            def diffusers_progress_callback(pipeline, step_index, timestep, callback_kwargs):
-                elapsed = time.time() - start_time
-                avg = elapsed / (step_index + 1) if step_index >= 0 else 0
-                remaining = avg * (steps - step_index - 1)
-                progress(
-                    (step_index + 1) / steps,
-                    desc=f"Step {step_index+1}/{steps} | ETA {remaining:.1f}s"
-                )
-                return callback_kwargs
+    steps, size = PRESETS.get(quality_mode, (1, 256))
+    width = height = size
+
+    seed = int(seed) if seed >= 0 else random.randint(0, (2**31)-1)
+    log(f"Generating: '{prompt[:40]}...' | {quality_mode} | {width}x{height} | seed={seed}")
+
+    with managed_memory(), generation_lock:
+        pipe = load_pipeline(model_name)
+        generator = torch.Generator("cpu").manual_seed(seed)
+        start_time = time.time()
 
+        preview_images = []
+
+        def progress_cb(pipeline, step_idx, timestep, cbk):
+            if interrupt_event.is_set():
+                raise KeyboardInterrupt("Generation interrupted")
+            if step_idx % 2 == 0:  # preview every 2 steps
+                try:
+                    preview_images.append(pipeline.image_from_latents(pipeline.latents))
+                except Exception:
+                    pass
+            return cbk
+
+        try:
             result = pipe(
                 prompt=prompt,
                 negative_prompt=None,
@@ -149,61 +173,57 @@ def generate(prompt, quality_mode, seed, progress=gr.Progress()):
                 num_inference_steps=steps,
                 guidance_scale=0.0,
                 generator=generator,
-                callback_on_step_end=diffusers_progress_callback,
+                callback_on_step_end=progress_cb,
                 callback_on_step_end_tensor_inputs=["latents"],
                 output_type="pil"
             )
+            final_image = result.images[0]
+            log(f"Done in {time.time()-start_time:.1f}s")
+        except KeyboardInterrupt:
+            log("⚠️ Generation interrupted.")
+            return None, seed, preview_images
 
-            image = result.images[0]
-            elapsed = time.time() - start_time
-            log(f"✅ Generated in {elapsed:.2f}s | Seed: {seed}")
+        del result
+        gc.collect()
 
-            del result
-            gc.collect()
+        preview_images.append(final_image)
+        return final_image, seed, preview_images
 
-            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")
+with gr.Blocks(title="🤩✨ Z‑Image Turbo CPU Ultimate + Retry + Preview + Interrupt") as demo:
+    gr.Markdown("## Full feature CPU image generator — true snapshot retry + preview frames")
 
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label="Prompt", lines=4)
             quality_mode = gr.Radio(
-                choices=[
-                    ("Ultra Fast", "ultra_fast"),
-                    ("Fast", "fast"),
-                    ("Balanced", "balanced"),
-                    ("Quality", "quality"),
-                    ("Ultra Quality", "ultra_quality")
-                ],
+                choices=["ultra_fast","fast","balanced","quality","ultra_quality"],
                 value="fast",
                 label="Quality Mode"
             )
-            seed = gr.Number(value=-1, precision=0, label="Seed")
-            generate_btn = gr.Button("GENERATE")
+            seed = gr.Number(value=-1, precision=0, label="Seed (-1=random)")
+            model_choice = gr.Dropdown(list(MODEL_SPECS.keys()), value=list(MODEL_SPECS.keys())[0], label="Select model")
+            gen_btn = gr.Button("GENERATE")
+            interrupt_btn = gr.Button("STOP")
+
         with gr.Column():
-            output_image = gr.Image(label="Output")
-            used_seed = gr.Number(label="Seed Used", interactive=False)
-            log_output = gr.Textbox(
-                label="Live System Log",
-                lines=15,
-                interactive=False
-            )
+            out_img = gr.Image(label="Final Image")
+            out_seed = gr.Number(label="Seed Used", interactive=False)
+            preview_gallery = gr.Gallery(label="Preview frames")
+            log_output = gr.Textbox(label="Live System Log", lines=15, interactive=False)
 
-    def wrapped_generate(prompt, quality_mode, seed):
-        image, used_seed = generate(prompt, quality_mode, seed)
-        logs = log("🧠 Latest status: Finished generation.")
-        return image, used_seed, logs
+    def on_generate(prompt, quality_mode, seed, model_choice):
+        interrupt_event.clear()
+        final_img, used_seed, previews = generate(prompt, quality_mode, seed, model_choice)
+        return final_img, used_seed, previews, log("Generation done.")
 
-    generate_btn.click(
-        wrapped_generate,
-        inputs=[prompt, quality_mode, seed],
-        outputs=[output_image, used_seed, log_output],
-        concurrency_limit=1
-    )
+    def on_interrupt():
+        interrupt_event.set()
+        return log("📌 Interrupt requested")
 
-demo.queue(max_size=3)
+    gen_btn.click(on_generate, inputs=[prompt, quality_mode, seed, model_choice], outputs=[out_img, out_seed, preview_gallery, log_output])
+    interrupt_btn.click(on_interrupt, inputs=None, outputs=log_output)
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.queue()
+demo.launch(server_name="0.0.0.0", server_port=7860)