Update app_quant_latent.py

app_quant_latent.py

@@ -5,60 +5,81 @@ import sys
 import platform
 import diffusers
 import transformers
+import psutil
 import os
-import torchvision.transforms as T
+import time
 
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from diffusers import ZImagePipeline, AutoModel
 from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 
 # ============================================================
-
 # LOGGING BUFFER
-
 # ============================================================
-
 LOGS = ""
 def log(msg):
-   global LOGS
-   print(msg)
-   LOGS += msg + "\n"
-   return msg
+    global LOGS
+    print(msg)
+    LOGS += msg + "\n"
+    return msg
 
+
+# ============================================================
+# SYSTEM METRICS — LIVE GPU + CPU MONITORING
 # ============================================================
+def log_system_stats(tag=""):
+    try:
+        log(f"\n===== 🔥 SYSTEM STATS {tag} =====")
 
-# ENVIRONMENT INFO
+        # ============= GPU STATS =============
+        if torch.cuda.is_available():
+            allocated = torch.cuda.memory_allocated(0) / 1e9
+            reserved = torch.cuda.memory_reserved(0) / 1e9
+            total = torch.cuda.get_device_properties(0).total_memory / 1e9
+            free = total - allocated
+
+            log(f"💠 GPU Total     : {total:.2f} GB")
+            log(f"💠 GPU Allocated : {allocated:.2f} GB")
+            log(f"💠 GPU Reserved  : {reserved:.2f} GB")
+            log(f"💠 GPU Free      : {free:.2f} GB")
+
+        # ============= CPU STATS ============
+        cpu = psutil.cpu_percent()
+        ram_used = psutil.virtual_memory().used / 1e9
+        ram_total = psutil.virtual_memory().total / 1e9
+
+        log(f"🧠 CPU Usage     : {cpu}%")
+        log(f"🧠 RAM Used      : {ram_used:.2f} GB / {ram_total:.2f} GB")
+
+    except Exception as e:
+        log(f"⚠️ Failed to log system stats: {e}")
 
-# ============================================================
 
+# ============================================================
+# ENVIRONMENT INFO
+# ============================================================
 log("===================================================")
-log("🔍 Z-IMAGE-TURBO DEBUGGING + ROBUST TRANSFORMER INSPECTION")
+log("🔍 Z-IMAGE-TURBO DEBUGGING + LIVE METRIC LOGGER")
 log("===================================================\n")
 
-log(f"📌 PYTHON VERSION       : {sys.version.replace(chr(10), ' ')}")
+log(f"📌 PYTHON VERSION       : {sys.version.replace(chr(10),' ')}")
 log(f"📌 PLATFORM             : {platform.platform()}")
-log(f"📌 TORCH VERSION        : {torch.**version**}")
-log(f"📌 TRANSFORMERS VERSION : {transformers.**version**}")
-log(f"📌 DIFFUSERS VERSION    : {diffusers.**version**}")
+log(f"📌 TORCH VERSION        : {torch.__version__}")
+log(f"📌 TRANSFORMERS VERSION : {transformers.__version__}")
+log(f"📌 DIFFUSERS VERSION    : {diffusers.__version__}")
 log(f"📌 CUDA AVAILABLE       : {torch.cuda.is_available()}")
 
-if torch.cuda.is_available():
-log(f"📌 GPU NAME             : {torch.cuda.get_device_name(0)}")
-log(f"📌 GPU CAPABILITY       : {torch.cuda.get_device_capability(0)}")
-log(f"📌 GPU MEMORY (TOTAL)   : {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB")
-log(f"📌 FLASH ATTENTION      : {torch.backends.cuda.flash_sdp_enabled()}")
-else:
-raise RuntimeError("❌ CUDA is REQUIRED but not available.")
+log_system_stats("AT STARTUP")
+
+if not torch.cuda.is_available():
+    raise RuntimeError("❌ CUDA Required")
 
 device = "cuda"
 gpu_id = 0
 
 # ============================================================
-
 # MODEL SETTINGS
-
 # ============================================================
-
 model_cache = "./weights/"
 model_id = "Tongyi-MAI/Z-Image-Turbo"
 torch_dtype = torch.bfloat16
@@ -72,232 +93,286 @@ log(f"Model Cache Directory : {model_cache}")
 log(f"torch_dtype           : {torch_dtype}")
 log(f"USE_CPU_OFFLOAD       : {USE_CPU_OFFLOAD}")
 
-# ============================================================
+log_system_stats("BEFORE TRANSFORMER LOAD")
 
-# ROBUST TRANSFORMER INSPECTION FUNCTION
 
 # ============================================================
+# FUNCTION TO CONVERT LATENTS TO IMAGE
+# ============================================================
+def latent_to_image(latent):
+    try:
+        img_tensor = pipe.vae.decode(latent)
+        img_tensor = (img_tensor / 2 + 0.5).clamp(0, 1)
+        pil_img = T.ToPILImage()(img_tensor[0])
+        return pil_img
+    except Exception as e:
+        log(f"⚠️ Failed to decode latent: {e}")
+        return None
 
-def inspect_transformer(model, model_name="Transformer"):
-log(f"\n🔍 {model_name} Architecture Details:")
-try:
-block_attrs = ["transformer_blocks", "blocks", "layers", "encoder_blocks", "model"]
-blocks = None
-for attr in block_attrs:
-blocks = getattr(model, attr, None)
-if blocks is not None:
-break
-
-```
-    if blocks is None:
-        log(f"⚠️ Could not find transformer blocks in {model_name}, skipping detailed block info")
-    else:
-        try:
-            log(f"Number of Transformer Modules : {len(blocks)}")
-            for i, block in enumerate(blocks):
-                log(f"  Block {i}: {block.__class__.__name__}")
-                attn_type = getattr(block, "attn", None)
-                if attn_type:
-                    log(f"    Attention: {attn_type.__class__.__name__}")
-                    flash_enabled = getattr(attn_type, "flash", None)
-                    log(f"    FlashAttention Enabled? : {flash_enabled}")
-        except Exception as e:
-            log(f"⚠️ Error inspecting blocks: {e}")
-
-    config = getattr(model, "config", None)
-    if config:
-        log(f"Hidden size: {getattr(config, 'hidden_size', 'N/A')}")
-        log(f"Number of attention heads: {getattr(config, 'num_attention_heads', 'N/A')}")
-        log(f"Number of layers: {getattr(config, 'num_hidden_layers', 'N/A')}")
-        log(f"Intermediate size: {getattr(config, 'intermediate_size', 'N/A')}")
-    else:
-        log(f"⚠️ No config attribute found in {model_name}")
-except Exception as e:
-    log(f"⚠️ Failed to inspect {model_name}: {e}")
-```
 
+
+# ============================================================
+# SAFE TRANSFORMER INSPECTION
 # ============================================================
+def inspect_transformer(model, name):
+    log(f"\n🔍 Inspecting {name}")
+    try:
+        candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
+        blocks = None
 
-# LOAD TRANSFORMER BLOCK
+        for attr in candidates:
+            if hasattr(model, attr):
+                blocks = getattr(model, attr)
+                break
+
+        if blocks is None:
+            log(f"⚠️ No block structure found in {name}")
+            return
+
+        if hasattr(blocks, "__len__"):
+            log(f"Total Blocks = {len(blocks)}")
+        else:
+            log("⚠️ Blocks exist but are not iterable")
+
+        for i in range(min(10, len(blocks) if hasattr(blocks, "__len__") else 0)):
+            log(f"Block {i} = {blocks[i].__class__.__name__}")
+
+    except Exception as e:
+        log(f"⚠️ Transformer inspect error: {e}")
 
-# ============================================================
 
+# ============================================================
+# LOAD TRANSFORMER — WITH LIVE STATS
+# ============================================================
 log("\n===================================================")
 log("🔧 LOADING TRANSFORMER BLOCK")
 log("===================================================")
 
-quantization_config = DiffusersBitsAndBytesConfig(
-load_in_4bit=True,
-bnb_4bit_quant_type="nf4",
-bnb_4bit_compute_dtype=torch_dtype,
-bnb_4bit_use_double_quant=True,
-llm_int8_skip_modules=["transformer_blocks.0.img_mod"],
-)
-log("4-bit Quantization Config (Transformer):")
-log(str(quantization_config))
-
-transformer = AutoModel.from_pretrained(
-model_id,
-cache_dir=model_cache,
-subfolder="transformer",
-quantization_config=quantization_config,
-torch_dtype=torch_dtype,
-device_map=device,
-)
-log("✅ Transformer block loaded successfully.")
-inspect_transformer(transformer, "Transformer")
+log("📌 Logging memory before load:")
+log_system_stats("START TRANSFORMER LOAD")
+
+try:
+    quant_cfg = DiffusersBitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch_dtype,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    transformer = AutoModel.from_pretrained(
+        model_id,
+        cache_dir=model_cache,
+        subfolder="transformer",
+        quantization_config=quant_cfg,
+        torch_dtype=torch_dtype,
+        device_map=device,
+    )
+    log("✅ Transformer loaded successfully.")
 
-if USE_CPU_OFFLOAD:
-transformer = transformer.to("cpu")
+except Exception as e:
+    log(f"❌ Transformer load failed: {e}")
+    transformer = None
 
-# ============================================================
+log_system_stats("AFTER TRANSFORMER LOAD")
 
-# LOAD TEXT ENCODER
+if transformer:
+    inspect_transformer(transformer, "Transformer")
 
-# ============================================================
 
+# ============================================================
+# LOAD TEXT ENCODER
+# ============================================================
 log("\n===================================================")
 log("🔧 LOADING TEXT ENCODER")
 log("===================================================")
 
-quantization_config = TransformersBitsAndBytesConfig(
-load_in_4bit=True,
-bnb_4bit_quant_type="nf4",
-bnb_4bit_compute_dtype=torch_dtype,
-bnb_4bit_use_double_quant=True,
-)
-log("4-bit Quantization Config (Text Encoder):")
-log(str(quantization_config))
-
-text_encoder = AutoModel.from_pretrained(
-model_id,
-cache_dir=model_cache,
-subfolder="text_encoder",
-quantization_config=quantization_config,
-torch_dtype=torch_dtype,
-device_map=device,
-)
-log("✅ Text encoder loaded successfully.")
-inspect_transformer(text_encoder, "Text Encoder")
+log_system_stats("START TEXT ENCODER LOAD")
 
-if USE_CPU_OFFLOAD:
-text_encoder = text_encoder.to("cpu")
+try:
+    quant_cfg2 = TransformersBitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch_dtype,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    text_encoder = AutoModel.from_pretrained(
+        model_id,
+        cache_dir=model_cache,
+        subfolder="text_encoder",
+        quantization_config=quant_cfg2,
+        torch_dtype=torch_dtype,
+        device_map=device,
+    )
+    log("✅ Text encoder loaded successfully.")
 
-# ============================================================
+except Exception as e:
+    log(f"❌ Text encoder load failed: {e}")
+    text_encoder = None
 
-# BUILD PIPELINE
+log_system_stats("AFTER TEXT ENCODER LOAD")
 
-# ============================================================
+if text_encoder:
+    inspect_transformer(text_encoder, "Text Encoder")
 
+
+# ============================================================
+# BUILD PIPELINE
+# ============================================================
 log("\n===================================================")
-log("🔧 BUILDING Z-IMAGE-TURBO PIPELINE")
+log("🔧 BUILDING PIPELINE")
 log("===================================================")
 
-pipe = ZImagePipeline.from_pretrained(
-model_id,
-transformer=transformer,
-text_encoder=text_encoder,
-torch_dtype=torch_dtype,
-)
+log_system_stats("START PIPELINE BUILD")
 
-if USE_CPU_OFFLOAD:
-pipe.enable_model_cpu_offload(gpu_id=gpu_id)
-log("⚙ CPU OFFLOAD ENABLED")
-else:
-pipe.to(device)
-log("⚙ Pipeline moved to GPU")
+try:
+    pipe = ZImagePipeline.from_pretrained(
+        model_id,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        torch_dtype=torch_dtype,
+    )
+    pipe.to(device)
+    log("✅ Pipeline built successfully.")
 
-log("✅ Pipeline ready.")
+except Exception as e:
+    log(f"❌ Pipeline build failed: {e}")
+    pipe = None
 
-# ============================================================
+log_system_stats("AFTER PIPELINE BUILD")
 
-# FUNCTION TO CONVERT LATENTS TO IMAGE
 
 # ============================================================
+# INFERENCE
+# ============================================================
+@spaces.GPU
+def generate_image(prompt, height, width, steps, seed):
+global LOGS
+LOGS = ""  # reset logs
+
+
+log("===================================================")
+log("🎨 RUNNING INFERENCE")
+log("===================================================")
+log_system_stats("BEFORE INFERENCE")
 
-def latent_to_image(latent):
 try:
-img_tensor = pipe.vae.decode(latent)
-img_tensor = (img_tensor / 2 + 0.5).clamp(0, 1)
-pil_img = T.ToPILImage()(img_tensor[0])
-return pil_img
-except Exception as e:
-log(f"⚠️ Failed to decode latent: {e}")
-return None
+    generator = torch.Generator(device).manual_seed(seed)
+    latent_history = []
+
+    # Callback to save latents and GPU info
+    def save_latents(step, timestep, latents):
+        latent_history.append(latents.detach().clone())
+        gpu_mem = torch.cuda.memory_allocated(0)/1e9
+        log(f"Step {step} - GPU Memory Used: {gpu_mem:.2f} GB")
+
+    # Step 3: Loop over pipeline for step-wise generation
+    for step, img in pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=0.0,
+        generator=generator,
+        callback=save_latents,
+        callback_steps=1
+    ).iter():
+        # Optionally: yield intermediate images or just store latents
+        current_latent = latent_history[-1] if latent_history else None
+        # You can process current_latent here if needed
+
+    log("✅ Inference finished.")
+    log_system_stats("AFTER INFERENCE")
+
+    # Return final image + logs
+    return img, LOGS
 
-# ============================================================
+except Exception as e:
+    log(f"❌ Inference error: {e}")
+    return None, LOGS
 
-# REAL-TIME INFERENCE FUNCTION
 
-# ============================================================
 
 @spaces.GPU
-def generate_image_realtime(prompt, height, width, steps, seed):
+def generate_image(prompt, height, width, steps, seed):
 global LOGS
-LOGS = ""
+LOGS = ""  # reset logs
+
+
 log("===================================================")
-log("🎨 RUNNING REAL-TIME INFERENCE")
+log("🎨 RUNNING INFERENCE")
 log("===================================================")
-log(f"Prompt     : {prompt}")
-log(f"Resolution : {width} x {height}")
-log(f"Steps      : {steps}")
-log(f"Seed       : {seed}")
-
-```
-generator = torch.Generator(device).manual_seed(seed)
-latent_history = []
-
-# Define callback to save latents and GPU info
-def save_latents(step, timestep, latents):
-    latent_history.append(latents.detach().clone())
-    gpu_mem = torch.cuda.memory_allocated(0)/1e9
-    log(f"Step {step} - GPU Memory Used: {gpu_mem:.2f} GB")
-
-# Yield images step-by-step
-for step, img in pipe(
-    prompt=prompt,
-    height=height,
-    width=width,
-    num_inference_steps=steps,
-    guidance_scale=0.0,
-    generator=generator,
-    callback=save_latents,
-    callback_steps=1
-).iter():
-    current_latent = latent_history[-1] if latent_history else None
-    latent_images = [latent_to_image(l) for l in latent_history if l is not None]
-    yield img, latent_images, LOGS
-```
+log_system_stats("BEFORE INFERENCE")
 
-# ============================================================
+try:
+    generator = torch.Generator(device).manual_seed(seed)
+    latent_history = []
+
+    # Callback to save latents and GPU info
+    def save_latents(step, timestep, latents):
+        latent_history.append(latents.detach().clone())
+        gpu_mem = torch.cuda.memory_allocated(0)/1e9
+        log(f"Step {step} - GPU Memory Used: {gpu_mem:.2f} GB")
+
+    # Step-wise loop just for latent capture
+    for step, _ in pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=0.0,
+        generator=generator,
+        callback=save_latents,
+        callback_steps=1
+    ).iter():
+        pass  # only capturing latents, ignoring intermediate images
+
+    # Original final image generation
+    output = pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=0.0,
+        generator=generator,
+    )
+
+    log("✅ Inference finished.")
+    log_system_stats("AFTER INFERENCE")
+
+    return output.images[0], latent_history, LOGS
+
+except Exception as e:
+    log(f"❌ Inference error: {e}")
+    return None, None, LOGS
 
-# GRADIO UI
 
+# ============================================================
+# UI
 # ============================================================
 
 with gr.Blocks(title="Z-Image-Turbo Generator") as demo:
-gr.Markdown("# **🚀 Z-Image-Turbo — Real-Time Latent & Transformer Logs**")
+gr.Markdown("# **🚀 Z-Image-Turbo — Final Image & Latents**")
 
-```
 with gr.Row():
     with gr.Column(scale=1):
         prompt = gr.Textbox(label="Prompt", value="Realistic mid-aged male image")
         height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
         width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
-        steps = gr.Slider(1, 16, value=9, step=1, label="Inference Steps")
-        seed = gr.Slider(0, 999999, value=42, step=1, label="Seed")
-        btn = gr.Button("Generate", variant="primary")
+        steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
+        seed = gr.Number(value=42, label="Seed")
+        run_btn = gr.Button("Generate Image")
 
     with gr.Column(scale=1):
-        output_image = gr.Image(label="Final Output Image")
-        latent_gallery = gr.Gallery(label="Latent Evolution", elem_id="latent_gallery").style(grid=[2], height="auto")
-        logs_panel = gr.Textbox(label="📜 Transformer & GPU Logs", lines=25, interactive=False)
+        final_image = gr.Image(label="Final Image")
+        latent_gallery = gr.Gallery(label="Latent Steps").style(grid=[4], height="256px")
+        logs_box = gr.Textbox(label="Logs", lines=15)
 
-btn.click(
-    generate_image_realtime,
+run_btn.click(
+    generate_image,
     inputs=[prompt, height, width, steps, seed],
-    outputs=[output_image, latent_gallery, logs_panel],
+    outputs=[final_image, latent_gallery, logs_box]
 )
-```
 
-demo.launch()
+
+demo.launch()