Update app_quant.py

app_quant.py

@@ -1,33 +1,77 @@
 import torch
 import spaces
 import gradio as gr
+import sys
+import platform
+import os
+import diffusers
+import transformers
+import peft
 
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from diffusers import ZImagePipeline, AutoModel
 from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 
 # ============================================================
-# Model Settings
+# LOGGING BUFFER
 # ============================================================
-model_cache = "./weights/"
-model_id = "Tongyi-MAI/Z-Image-Turbo"
-torch_dtype = torch.bfloat16
-USE_CPU_OFFLOAD = False
+LOGS = ""
+
+def log(msg):
+    global LOGS
+    print(msg)
+    LOGS += msg + "\n"
+    return msg
 
 # ============================================================
-# GPU Check
+# ENVIRONMENT INFO
 # ============================================================
+log("===================================================")
+log("🔍 Z-IMAGE-TURBO DEBUGGING + DETAILED TRANSFORMER INSIGHTS")
+log("===================================================\n")
+
+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"📌 PEFT VERSION         : {peft.__version__}")
+log(f"📌 CUDA AVAILABLE       : {torch.cuda.is_available()}")
+
 if torch.cuda.is_available():
-    print(f"INFO: CUDA available: {torch.cuda.get_device_name(0)} (count={torch.cuda.device_count()})")
-    device = "cuda:0"
-    gpu_id = 0
+    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("ERROR: CUDA not available. This program requires a CUDA-enabled GPU.")
+    raise RuntimeError("❌ CUDA is REQUIRED but not available.")
+
+device = "cuda"
+gpu_id = 0
+
+# ============================================================
+# MODEL SETTINGS
+# ============================================================
+model_cache = "./weights/"
+model_id = "Tongyi-MAI/Z-Image-Turbo"
+torch_dtype = torch.bfloat16
+USE_CPU_OFFLOAD = False
+
+log("\n===================================================")
+log("🧠 MODEL CONFIGURATION")
+log("===================================================")
+log(f"Model ID              : {model_id}")
+log(f"Model Cache Directory : {model_cache}")
+log(f"torch_dtype           : {torch_dtype}")
+log(f"USE_CPU_OFFLOAD       : {USE_CPU_OFFLOAD}")
 
 # ============================================================
-# Load Transformer
+# LOAD TRANSFORMER BLOCK
 # ============================================================
-print("INFO: Loading transformer block ...")
+log("\n===================================================")
+log("🔧 LOADING TRANSFORMER BLOCK")
+log("===================================================")
+
 quantization_config = DiffusersBitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -35,6 +79,10 @@ quantization_config = DiffusersBitsAndBytesConfig(
     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,
@@ -43,21 +91,47 @@ transformer = AutoModel.from_pretrained(
     torch_dtype=torch_dtype,
     device_map=device,
 )
-print("INFO: Transformer block loaded.")
+log("✅ Transformer block loaded successfully.")
+
+# ------------------------------------------------------------
+# TRANSFORMER INSIGHTS
+# ------------------------------------------------------------
+log("🔍 Transformer Architecture Details:")
+log(f"Number of Transformer Modules : {len(transformer.transformer_blocks)}")
+for i, block in enumerate(transformer.transformer_blocks):
+    log(f"  Block {i}: {block.__class__.__name__}")
+    # Log attention type if possible
+    attn_type = getattr(block, "attn", None)
+    if attn_type:
+        log(f"    Attention: {attn_type.__class__.__name__}")
+    # Check for FlashAttention usage if attribute exists
+    flash_enabled = getattr(attn_type, "flash", None)
+    log(f"    FlashAttention Enabled? : {flash_enabled}")
+log(f"Hidden size: {transformer.config.hidden_size}")
+log(f"Number of attention heads: {transformer.config.num_attention_heads}")
+log(f"Number of layers: {transformer.config.num_hidden_layers}")
+log(f"Intermediate size: {transformer.config.intermediate_size}")
 
 if USE_CPU_OFFLOAD:
     transformer = transformer.to("cpu")
 
 # ============================================================
-# Load Text Encoder
+# LOAD TEXT ENCODER
 # ============================================================
-print("INFO: Loading 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.bfloat16,
     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,
@@ -66,15 +140,35 @@ text_encoder = AutoModel.from_pretrained(
     torch_dtype=torch_dtype,
     device_map=device,
 )
-print("INFO: Text encoder loaded.")
+log("✅ Text encoder loaded successfully.")
+
+# ------------------------------------------------------------
+# TEXT ENCODER INSIGHTS
+# ------------------------------------------------------------
+log("🔍 Text Encoder Architecture Details:")
+log(f"Number of Transformer Modules : {len(text_encoder.transformer_blocks)}")
+for i, block in enumerate(text_encoder.transformer_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}")
+log(f"Hidden size: {text_encoder.config.hidden_size}")
+log(f"Number of attention heads: {text_encoder.config.num_attention_heads}")
+log(f"Number of layers: {text_encoder.config.num_hidden_layers}")
+log(f"Intermediate size: {text_encoder.config.intermediate_size}")
 
 if USE_CPU_OFFLOAD:
     text_encoder = text_encoder.to("cpu")
 
 # ============================================================
-# Build Pipeline
+# BUILD PIPELINE
 # ============================================================
-print("INFO: Building pipeline ...")
+log("\n===================================================")
+log("🔧 BUILDING Z-IMAGE-TURBO PIPELINE")
+log("===================================================")
+
 pipe = ZImagePipeline.from_pretrained(
     model_id,
     transformer=transformer,
@@ -84,19 +178,33 @@ pipe = ZImagePipeline.from_pretrained(
 
 if USE_CPU_OFFLOAD:
     pipe.enable_model_cpu_offload(gpu_id=gpu_id)
-    print("INFO: CPU offload active")
+    log("⚙ CPU OFFLOAD ENABLED")
 else:
     pipe.to(device)
-    print("INFO: Pipeline to GPU")
+    log("⚙ Pipeline moved to GPU")
+
+log("✅ Pipeline ready.")
 
 # ============================================================
-# Inference Function for Gradio
+# INFERENCE FUNCTION
 # ============================================================
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed):
+    global LOGS
+    LOGS = ""  # Reset logs for this run
+
+    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)
-    
-    output = pipe(
+
+    out = pipe(
         prompt=prompt,
         height=height,
         width=width,
@@ -105,14 +213,14 @@ def generate_image(prompt, height, width, steps, seed):
         generator=generator,
     )
 
-    return output.images[0]
-
+    log("✅ Inference Finished")
+    return out.images[0], LOGS
 
 # ============================================================
-# Gradio UI
+# GRADIO UI
 # ============================================================
 with gr.Blocks(title="Z-Image-Turbo Generator") as demo:
-    gr.Markdown("# **Z-Image-Turbo — 4bit Quantized Image Generator**")
+    gr.Markdown("# **🚀 Z-Image-Turbo — Transformer Deep Logs**")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -123,17 +231,15 @@ with gr.Blocks(title="Z-Image-Turbo Generator") as demo:
             seed = gr.Slider(0, 999999, value=42, step=1, label="Seed")
 
             btn = gr.Button("Generate", variant="primary")
-        
+
         with gr.Column(scale=1):
             output_image = gr.Image(label="Output Image")
+            logs_panel = gr.Textbox(label="📜 Transformer Logs", lines=25, interactive=False)
 
     btn.click(
         generate_image,
         inputs=[prompt, height, width, steps, seed],
-        outputs=[output_image],
+        outputs=[output_image, logs_panel],
     )
 
-# ============================================================
-# Launch
-# ============================================================
 demo.launch()