Update app.py

app.py

@@ -1,194 +1,209 @@
 import os
 import gc
+import sys
 import time
 import random
 import torch
 import gradio as gr
-
-# =====================================================
-# 🛠 CPU OPTIMIZED SETTINGS
-# =====================================================
-
+from threading import Lock
+from contextlib import contextmanager
+
+# --- LOGGING FOR UI ---
+LOG_BUFFER = []
+LOG_LOCK = Lock()
+
+def log(message):
+    print(message)
+    with LOG_LOCK:
+        LOG_BUFFER.append(f"{time.strftime('%H:%M:%S')} | {message}")
+        if len(LOG_BUFFER) > 500:
+            LOG_BUFFER.pop(0)
+    return "\n".join(LOG_BUFFER)
+
+# 🚀 Initialization log
+_initial_logs = log("🚀 Initializing Ultimate Z-Image Turbo CPU Edition...")
+
+# 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)
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
-
-CPU_THREADS = min(4, os.cpu_count() or 1)
-os.environ["OMP_NUM_THREADS"] = str(CPU_THREADS)
-os.environ["MKL_NUM_THREADS"] = str(CPU_THREADS)
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
 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.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)
 
-print("⚡ Z-Image Turbo CPU — Optimized (Latest Docs)")
+log(f"⚡ CPU Threads: {CPU_THREADS}, Device: {DEVICE}, DType: {DTYPE}")
 
 try:
-    from huggingface_hub import hf_hub_download
-    from diffusers import (
-        ZImagePipeline,
-        ZImageTransformer2DModel,
-        GGUFQuantizationConfig,
-        AutoencoderKL,
-        FlowMatchEulerDiscreteScheduler
-    )
-    from transformers import CLIPTokenizer, CLIPTextModel
+    from diffusers import ZImagePipeline
+    log("📦 diffusers imported successfully")
 except ImportError as e:
-    print(f"⚠️ Imports may not load: {e}")
+    log(f"❌ Import Error: {e}")
+    sys.exit(1)
 
 pipe = None
-_pipe_lock = False
+_pipe_lock = Lock()
+_generation_lock = Lock()
 
-def load_text_encoder_min():
-    tokenizer = CLIPTokenizer.from_pretrained(
-        "openai/clip-vit-base-patch32", cache_dir=CACHE_DIR
-    )
-    text_encoder = CLIPTextModel.from_pretrained(
-        "openai/clip-vit-base-patch32",
-        cache_dir=CACHE_DIR,
-        torch_dtype=DTYPE,
-        low_cpu_mem_usage=True
-    )
-    return tokenizer, text_encoder
+@contextmanager
+def managed_memory():
+    try:
+        yield
+    finally:
+        gc.collect()
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 def load_pipeline():
-    global pipe, _pipe_lock
-    if pipe is not None:
-        return pipe
-    if _pipe_lock:
-        raise gr.Error("Pipeline already loading…")
-    _pipe_lock = True
+    global pipe
+    with _pipe_lock:
+        if pipe is not None:
+            return pipe
 
-    try:
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            "Tongyi-MAI/Z-Image-Turbo",
-            subfolder="scheduler",
-            cache_dir=CACHE_DIR,
-            low_cpu_mem_usage=True
-        )
+        log("📦 Loading Z-Image Turbo pipeline...")
+        start_load = time.time()
 
-        vae = AutoencoderKL.from_pretrained(
+        pipe = ZImagePipeline.from_pretrained(
             "Tongyi-MAI/Z-Image-Turbo",
-            subfolder="vae",
-            cache_dir=CACHE_DIR,
             torch_dtype=DTYPE,
-            low_cpu_mem_usage=True
-        )
-
-        tokenizer, text_encoder = load_text_encoder_min()
-
-        gguf_path = hf_hub_download(
-            repo_id="unsloth/Z-Image-Turbo-GGUF",
-            filename="z-image-turbo-Q2_K.gguf",
             cache_dir=CACHE_DIR,
-            resume_download=True
-        )
-
-        transformer = ZImageTransformer2DModel.from_single_file(
-            gguf_path,
-            quantization_config=GGUFQuantizationConfig(compute_dtype=DTYPE),
-            torch_dtype=DTYPE,
             low_cpu_mem_usage=True
         )
 
-        pipe = ZImagePipeline(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler
-        ).to(DEVICE)
-
-        pipe.enable_attention_slicing()
-        pipe.enable_vae_slicing()
-        pipe.enable_vae_tiling()
-        pipe.set_progress_bar_config(disable=True)
-
+        pipe = 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")
         return pipe
 
-    except Exception as e:
-        raise gr.Error(f"Failed to load model: {e}")
-
-    finally:
-        _pipe_lock = False
-
 @torch.inference_mode()
-def generate(prompt, width, height, steps, seed, progress=gr.Progress()):
+@torch.no_grad()
+def generate(prompt, quality_mode, seed, progress=gr.Progress()):
     if not prompt.strip():
-        raise gr.Error("Prompt required")
-
-    width = (max(256, min(int(width), 512)) // 64) * 64
-    height = (max(256, min(int(height), 512)) // 64) * 64
-    steps = max(1, min(int(steps), 4))
-
-    if seed < 0 or seed == "":
-        seed = random.randint(0, 2**31 - 1)
-    else:
-        seed = int(seed)
-
-    gc.collect()
-
-    pipe = load_pipeline()
-    generator = torch.Generator(device="cpu").manual_seed(seed)
-
-    start = time.time()
-    def callback(step, *_):
-        elapsed = time.time() - start
-        avg = elapsed / (step + 1)
-        remaining = avg * (steps - step - 1)
-        progress((step+1)/steps, desc=f"Step {step+1}/{steps} | ETA {remaining:.1f}s")
-
-    result = pipe(
-        prompt=prompt,
-        negative_prompt=None,
-        width=width,
-        height=height,
-        num_inference_steps=steps,
-        guidance_scale=1.0,
-        generator=generator,
-        callback=callback,
-        callback_steps=1,
-        output_type="pil"
-    )
-
-    image = result.images[0]
-    del result
-    gc.collect()
-    return image, seed
-
-with gr.Blocks() as demo:
-    gr.Markdown("# ⚡ Z-Image Turbo — CPU Optimized")
-
-    prompt = gr.Textbox(label="Prompt")
-    width = gr.Slider(256, 512, 256, step=64, label="Width")
-    height = gr.Slider(256, 512, 256, step=64, label="Height")
-    steps = gr.Slider(1, 4, 2, step=1, label="Steps")
-    seed = gr.Number(value=-1, precision=0, label="Seed (-1=random)")
-
-    btn = gr.Button("🚀 Generate")
-    output = gr.Image(label="Output")
-    used_seed = gr.Number(label="Seed Used", interactive=False)
-
-    btn.click(
-        generate,
-        inputs=[prompt, width, height, steps, seed],
-        outputs=[output, used_seed],
+        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}
+    }
+    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 progress_callback(step, *_):
+                elapsed = time.time() - start_time
+                avg = elapsed / (step + 1) if step >= 0 else 0
+                remaining = avg * (steps - step - 1)
+                progress(
+                    (step + 1) / steps,
+                    desc=f"Step {step+1}/{steps} | ETA {remaining:.1f}s"
+                )
+
+            result = pipe(
+                prompt=prompt,
+                negative_prompt=None,
+                width=width,
+                height=height,
+                num_inference_steps=steps,
+                guidance_scale=0.0,
+                generator=generator,
+                callback=progress_callback,
+                callback_steps=1,
+                output_type="pil"
+            )
+
+            image = result.images[0]
+            elapsed = time.time() - start_time
+            log(f"✅ Generated in {elapsed:.2f}s | Seed: {seed}")
+
+            del result
+            gc.collect()
+
+            return image, seed
+
+with gr.Blocks(title="🚀 Z-Image Turbo Pro Max + Live Logs") as demo:
+    gr.Markdown("## GPU‑FREE CPU Turbo — Live Logs Below")
+
+    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")
+                ],
+                value="fast",
+                label="Quality Mode"
+            )
+            seed = gr.Number(value=-1, precision=0, label="Seed")
+            generate_btn = gr.Button("GENERATE")
+        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
+            )
+
+    def wrapped_generate(prompt, quality_mode, seed):
+        image, used_seed = generate(prompt, quality_mode, seed)
+        logs = log(f"🧠 Latest status: Finished generation.")
+        return image, used_seed, logs
+
+    generate_btn.click(
+        wrapped_generate,
+        inputs=[prompt, quality_mode, seed],
+        outputs=[output_image, used_seed, log_output],
         concurrency_limit=1
     )
 
-# Enable queue with up to 2 pending jobs
-demo.queue(max_size=2)  # queues events per current Gradio docs :contentReference[oaicite:1]{index=1}
+demo.queue(max_size=3)
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)