Update app.py

app.py

@@ -1,97 +1,198 @@
 import os
-import gc
 import sys
 import time
+import json
+import gc
 import random
 import torch
 import gradio as gr
-from threading import Lock, Event
+import requests
+from threading import Lock, Event, Thread
 from contextlib import contextmanager
-from huggingface_hub import snapshot_download, LocalEntryNotFoundError
+from urllib.parse import urlparse
 
-# ----------- LOGGING -----------
+from huggingface_hub import hf_hub_download, hf_hub_url
+from huggingface_hub.utils import RepositoryNotFoundError
 
+# ===== LOGGING =====
 LOG_BUFFER = []
 LOG_LOCK = Lock()
-
 def log(msg):
     with LOG_LOCK:
-        timestamp = time.strftime('%H:%M:%S')
-        LOG_BUFFER.append(f"{timestamp} | {msg}")
+        t = time.strftime("%H:%M:%S")
+        entry = f"{t} | {msg}"
+        LOG_BUFFER.append(entry)
         if len(LOG_BUFFER) > 500:
             LOG_BUFFER.pop(0)
-        print(msg)
     return "\n".join(LOG_BUFFER)
 
-# ----------- ENV CONFIG -----------
-
-CPU_THREADS = min(8, os.cpu_count() or 1)
-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"] = ""
+# ===== ENV SETUP =====
 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"
+os.makedirs("./hf_cache", exist_ok=True)
 
 torch.set_grad_enabled(False)
-torch.set_num_threads(CPU_THREADS)
-torch.backends.mkldnn.enabled = True
+torch.set_num_threads(min(8, os.cpu_count() or 1))
 torch.set_float32_matmul_precision("medium")
 
 DEVICE = "cpu"
 DTYPE = torch.float32
-os.makedirs("./hf_cache", exist_ok=True)
 
 try:
-    from diffusers import ZImagePipeline
-    log("Imported diffusers successfully.")
+    from diffusers import ZImagePipeline, GGUFQuantizationConfig, ZImageTransformer2DModel
+    log("Loaded diffusers modules")
 except ImportError as e:
-    log(f"Import diffusers failed: {e}")
+    log(f"Import error: {e}")
     sys.exit(1)
 
-pipe_cache = {}
-pipe_lock = Lock()
-generation_lock = Lock()
+# ===== DOWNLOAD CONTEXT =====
 interrupt_event = Event()
+pipe_cache = {}
+download_lock = Lock()
 
-# ----------- SNAPSHOT WITH RETRY -----------
-
+# ===== MODEL LIST =====
 MODEL_SPECS = {
-    "Z-Image Turbo": "Tongyi-MAI/Z-Image-Turbo",
-    # Optionally add quantized variants here
-    # "Z-Image Turbo GGUF": "unsloth/Z-Image-Turbo-GGUF",
+    "Turbo Full": "Tongyi-MAI/Z-Image-Turbo",
+    "Turbo Q2_K 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 -----------
+# ===== DOWNLOAD HELPERS =====
+def list_repo_files(repo_id):
+    """
+    Returns a list of (filename, size) tuples by doing a dry run
+    (no actual data downloaded).
+    """
+    try:
+        infos = hf_hub_download(repo_id, dry_run=True)
+        return [(info.rfilename, info.size_in_bytes) for info in infos]
+    except Exception as e:
+        log(f"List failed: {e}")
+        return []
+
+def download_file_chunked(repo_id, filename, target_dir, progress_updater):
+    """
+    Download a single file by streaming signed URL chunks.
+    Supports resume by checking existing file size.
+    """
+    local_path = os.path.join(target_dir, filename)
+    tmp_path = local_path + ".part"
+
+    os.makedirs(os.path.dirname(local_path), exist_ok=True)
+    already = 0
+    if os.path.exists(tmp_path):
+        already = os.path.getsize(tmp_path)
+
+    # Get a fresh signed URL from HF for that file
+    try:
+        url = hf_hub_url(repo_id, filename)
+    except RepositoryNotFoundError:
+        # fallback to normal
+        url = hf_hub_download(repo_id, filename=filename)
+
+    headers = {}
+    if already > 0:
+        headers["Range"] = f"bytes={already}-"
+
+    with requests.get(url, headers=headers, stream=True, timeout=10) as r:
+        total = int(r.headers.get("Content-Length", 0)) + already
+        with open(tmp_path, "ab") as f:
+            downloaded = already
+            for chunk in r.iter_content(chunk_size=1024*256):
+                if interrupt_event.is_set():
+                    return False
+                if not chunk:
+                    continue
+                f.write(chunk)
+                downloaded += len(chunk)
+                progress_updater(downloaded / total)
+
+    os.rename(tmp_path, local_path)
+    return True
+
+def parallel_download_repo(repo_id, progress: gr.Progress):
+    """
+    Download all files in the repo in parallel with per-file progress.
+    """
+    base_dir = os.path.join("./hf_cache", repo_id.replace("/", "_"))
+    files = list_repo_files(repo_id)
+    if not files:
+        progress(1.0, desc="No files to download")
+        return
+
+    total_bytes = sum(sz for _, sz in files)
+    downloaded_bytes = 0
+
+    def file_thread(filename, size):
+        nonlocal downloaded_bytes
+        success = download_file_chunked(
+            repo_id, filename, base_dir,
+            lambda frac: progress((downloaded_bytes + frac * size) / total_bytes,
+                                  desc=f"{filename} {frac*100:.1f}%")
+        )
+        if success:
+            with download_lock:
+                downloaded_bytes += size
+
+    threads = []
+    for fname, size in files:
+        if interrupt_event.is_set():
+            break
+        # skip if fully cached already
+        local_full = os.path.join(base_dir, fname)
+        if os.path.exists(local_full) and os.path.getsize(local_full) == size:
+            downloaded_bytes += size
+            continue
+        t = Thread(target=file_thread, args=(fname, size))
+        t.start()
+        threads.append(t)
+
+    for t in threads:
+        t.join()
+
+# ===== PIPELINE LOADER =====
+def load_pipeline(model_key):
+    """
+    Load the HF pipeline, using quantized GGUF if selected.
+    """
+    if model_key in pipe_cache:
+        return pipe_cache[model_key]
+
+    repo = MODEL_SPECS[model_key]
+    repo_cache = os.path.join("./hf_cache", repo.replace("/", "_"))
+
+    # ensure cache
+    if not os.path.isdir(repo_cache) or not os.listdir(repo_cache):
+        raise gr.Error("Model not downloaded; press Preload first")
+
+    # load model
+    if "GGUF" in model_key:
+        # pick .gguf file
+        files = [f for f in os.listdir(repo_cache) if f.endswith(".gguf")]
+        if not files:
+            raise gr.Error("Quantized file not found")
+        gguf = os.path.join(repo_cache, files[0])
+        transformer = ZImageTransformer2DModel.from_single_file(
+            gguf,
+            quantization_config=GGUFQuantizationConfig(compute_dtype=DTYPE),
+            torch_dtype=DTYPE
+        )
+        pipe = ZImagePipeline.from_pretrained(
+            "Tongyi-MAI/Z-Image-Turbo",
+            transformer=transformer,
+            torch_dtype=DTYPE,
+            cache_dir="./hf_cache"
+        )
+    else:
+        pipe = ZImagePipeline.from_pretrained(repo_cache, torch_dtype=DTYPE, local_files_only=True)
+
+    pipe.to(DEVICE)
+    pipe.vae.eval()
+    pipe.text_encoder.eval()
+    pipe.transformer.eval()
+    pipe_cache[model_key] = pipe
+    return pipe
 
 @contextmanager
 def managed_memory():
@@ -103,127 +204,97 @@ def managed_memory():
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-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")
-            log("Transformer compiled.")
-        except Exception as e:
-            log(f"Transformer compile skipped: {e}")
-        pipe_cache[model_name] = pipe
-        return pipe
-
-# ----------- GENERATION LOGIC -----------
-
+# ===== GENERATION =====
 @torch.inference_mode()
 @torch.no_grad()
-def generate(prompt, quality_mode, seed, model_name):
+def generate(prompt, quality_mode, seed, model_key):
     if not prompt.strip():
-        raise gr.Error("Prompt cannot be empty!")
+        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),
+        "ultra_fast": (1,256),
+        "fast": (1,256),
+        "balanced": (2,256),
+        "quality": (4,384),
+        "ultra_quality": (4,512),
     }
-    steps, size = PRESETS.get(quality_mode, (1, 256))
+    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()
+    seed = int(seed) if seed>=0 else random.randint(0,2**31-1)
+    log(f"Gen: {prompt[:30]} | {quality_mode} | {model_key} | seed={seed}")
 
-        preview_images = []
+    with managed_memory():
+        pipe = load_pipeline(model_key)
+        gen = torch.Generator("cpu").manual_seed(seed)
+        previews=[]
+        start = time.time()
 
-        def progress_cb(pipeline, step_idx, timestep, cbk):
+        def cb(ppl, step, timestep, cbk):
             if interrupt_event.is_set():
-                raise KeyboardInterrupt("Generation interrupted")
-            if step_idx % 2 == 0:  # preview every 2 steps
+                ppl._interrupt=True
+            if step % 2 == 0:
                 try:
-                    preview_images.append(pipeline.image_from_latents(pipeline.latents))
-                except Exception:
+                    previews.append(ppl.image_from_latents(cbk["latents"]))
+                except:
                     pass
             return cbk
 
-        try:
-            result = pipe(
-                prompt=prompt,
-                negative_prompt=None,
-                width=width,
-                height=height,
-                num_inference_steps=steps,
-                guidance_scale=0.0,
-                generator=generator,
-                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
-
-        del result
-        gc.collect()
-
-        preview_images.append(final_image)
-        return final_image, seed, preview_images
-
-# ----------- GRADIO UI -----------
-
-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")
+        result = pipe(
+            prompt=prompt,
+            negative_prompt=None,
+            width=width,
+            height=height,
+            num_inference_steps=steps,
+            guidance_scale=0.0,
+            generator=gen,
+            callback_on_step_end=cb,
+            callback_on_step_end_tensor_inputs=["latents"],
+            output_type="pil"
+        )
+        final = result.images[0]
+        previews.append(final)
+        log(f"Generated in {time.time()-start:.1f}s")
+    return final, seed, previews
+
+# ===== GRADIO UI =====
+with gr.Blocks(title="Z‑Image Turbo CPU Downloader + UI A‑Progress") as demo:
+    gr.Markdown("## True parallel download UI + chunked progress")
 
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label="Prompt", lines=4)
-            quality_mode = gr.Radio(
-                choices=["ultra_fast","fast","balanced","quality","ultra_quality"],
-                value="fast",
-                label="Quality Mode"
-            )
-            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")
+            quality = gr.Radio(["ultra_fast","fast","balanced","quality","ultra_quality"], value="fast")
+            seed = gr.Number(value=-1, precision=0, label="Seed")
+            model_select = gr.Dropdown(list(MODEL_SPECS.keys()), value=list(MODEL_SPECS.keys())[0], label="Model")
+            preload = gr.Button("PRELOAD MODELS")
             gen_btn = gr.Button("GENERATE")
-            interrupt_btn = gr.Button("STOP")
-
+            stop_btn = gr.Button("STOP")
         with gr.Column():
-            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)
+            out_image = gr.Image(label="Final")
+            used_seed = gr.Number(label="Seed Used")
+            preview = gr.Gallery(label="Preview Frames")
+            logs = gr.Textbox(label="Logs", lines=25)
+
+    def do_preload(progress=gr.Progress()):
+        interrupt_event.clear()
+        for key, repo in MODEL_SPECS.items():
+            parallel_download_repo(repo, progress)
+        return log("📦 Preload finished")
 
-    def on_generate(prompt, quality_mode, seed, model_choice):
+    def do_gen(prompt, quality, seed, model_key):
         interrupt_event.clear()
-        final_img, used_seed, previews = generate(prompt, quality_mode, seed, model_choice)
-        return final_img, used_seed, previews, log("Generation done.")
+        img, used, previews = generate(prompt, quality, seed, model_key)
+        return img, used, previews, log("🧠 Generation done")
 
-    def on_interrupt():
+    def do_stop():
         interrupt_event.set()
-        return log("📌 Interrupt requested")
+        return log("🔴 Interrupt set")
 
-    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)
+    preload.click(do_preload, outputs=logs)
+    gen_btn.click(do_gen, inputs=[prompt,quality,seed,model_select],
+                  outputs=[out_image,used_seed,preview,logs])
+    stop_btn.click(do_stop, outputs=logs)
 
 demo.queue()
 demo.launch(server_name="0.0.0.0", server_port=7860)