Update app_quant.py

app_quant.py

@@ -5,7 +5,9 @@ import sys
 import platform
 import diffusers
 import transformers
+import psutil
 import os
+import time
 
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from diffusers import ZImagePipeline, AutoModel
@@ -21,27 +23,56 @@ def log(msg):
     LOGS += msg + "\n"
     return msg
 
+
+# ============================================================
+# SYSTEM METRICS — LIVE GPU + CPU MONITORING
+# ============================================================
+def log_system_stats(tag=""):
+    try:
+        log(f"\n===== 🔥 SYSTEM STATS {tag} =====")
+
+        # ============= 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"📌 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
@@ -62,75 +93,75 @@ 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
+# SAFE TRANSFORMER INSPECTION
 # ============================================================
-def inspect_transformer(model, model_name="Transformer"):
-    log(f"\n🔍 {model_name} Architecture Details:")
+def inspect_transformer(model, name):
+    log(f"\n🔍 Inspecting {name}")
     try:
-        block_attrs = ["transformer_blocks", "blocks", "layers", "encoder_blocks", "model"]
+        candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
         blocks = None
-        for attr in block_attrs:
-            blocks = getattr(model, attr, None)
-            if blocks is not None:
+
+        for attr in candidates:
+            if hasattr(model, attr):
+                blocks = getattr(model, attr)
                 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')}")
+            log(f"⚠️ No block structure found in {name}")
+            return
+
+        if hasattr(blocks, "__len__"):
+            log(f"Total Blocks = {len(blocks)}")
         else:
-            log(f"⚠️ No config attribute found in {model_name}")
+            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"⚠️ Failed to inspect {model_name}: {e}")
+        log(f"⚠️ Transformer inspect error: {e}")
+
 
 # ============================================================
-# LOAD TRANSFORMER BLOCK
+# 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")
-
-if USE_CPU_OFFLOAD:
-    transformer = transformer.to("cpu")
+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.")
+
+except Exception as e:
+    log(f"❌ Transformer load failed: {e}")
+    transformer = None
+
+log_system_stats("AFTER TRANSFORMER LOAD")
+
+if transformer:
+    inspect_transformer(transformer, "Transformer")
+
 
 # ============================================================
 # LOAD TEXT ENCODER
@@ -139,104 +170,118 @@ 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")
-
-if USE_CPU_OFFLOAD:
-    text_encoder = text_encoder.to("cpu")
+log_system_stats("START TEXT ENCODER LOAD")
+
+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
+
+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,
-)
-
-if USE_CPU_OFFLOAD:
-    pipe.enable_model_cpu_offload(gpu_id=gpu_id)
-    log("⚙ CPU OFFLOAD ENABLED")
-else:
+log_system_stats("START PIPELINE BUILD")
+
+try:
+    pipe = ZImagePipeline.from_pretrained(
+        model_id,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        torch_dtype=torch_dtype,
+    )
     pipe.to(device)
-    log("⚙ Pipeline moved to GPU")
+    log("✅ Pipeline built successfully.")
+
+except Exception as e:
+    log(f"❌ Pipeline build failed: {e}")
+    pipe = None
+
+log_system_stats("AFTER PIPELINE BUILD")
 
-log("✅ Pipeline ready.")
 
 # ============================================================
-# INFERENCE FUNCTION
+# INFERENCE
 # ============================================================
-@spaces.GPU
+@spaces.GGPU
 def generate_image(prompt, height, width, steps, seed):
     global LOGS
     LOGS = ""  # reset logs
+
     log("===================================================")
     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)
-
-    out = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        num_inference_steps=steps,
-        guidance_scale=0.0,
-        generator=generator,
-    )
-    log("✅ Inference Finished")
-    return out.images[0], LOGS
+    log_system_stats("BEFORE INFERENCE")
+
+    try:
+        generator = torch.Generator(device).manual_seed(seed)
+
+        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], LOGS
+
+    except Exception as e:
+        log(f"❌ Inference error: {e}")
+        return None, LOGS
+
 
 # ============================================================
-# GRADIO UI
+# UI
 # ============================================================
-with gr.Blocks(title="Z-Image-Turbo Generator") as demo:
-    gr.Markdown("# **Z-Image-Turbo — 4bit Quantized Image Generator**")
+with gr.Blocks(title="Z-Image Turbo Debugger") as demo:
+    gr.Markdown("## **Z-Image Turbo — Full Debug + Live GPU/CPU Monitor**")
 
     with gr.Row():
         with gr.Column(scale=1):
-            prompt = gr.Textbox(label="Prompt", value="Realistic mid-aged male image")
+            prompt = gr.Textbox(label="Prompt", value="Realistic male portrait")
             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")
+            steps = gr.Slider(1, 16, value=9, step=1, label="Steps")
             seed = gr.Slider(0, 999999, value=42, step=1, label="Seed")
-
-            btn = gr.Button("Generate", variant="primary")
+            btn = gr.Button("Generate")
 
         with gr.Column(scale=1):
-            output_image = gr.Image(label="Output Image")
-            logs_panel = gr.Textbox(label="📜 Transformer Logs", lines=25, interactive=False)
+            image_out = gr.Image(label="Output")
+            logs_panel = gr.Textbox(label="📜 Logs", lines=30)
 
-    btn.click(
-        generate_image,
-        inputs=[prompt, height, width, steps, seed],
-        outputs=[output_image, logs_panel],
-    )
+    btn.click(generate_image,
+              inputs=[prompt, height, width, steps, seed],
+              outputs=[image_out, logs_panel])
 
 demo.launch()