Update app_quant.py

app_quant.py

@@ -3,10 +3,9 @@ import spaces
 import gradio as gr
 import sys
 import platform
-import os
 import diffusers
 import transformers
-
+import os
 
 from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
 from diffusers import ZImagePipeline, AutoModel
@@ -16,7 +15,6 @@ from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
 # LOGGING BUFFER
 # ============================================================
 LOGS = ""
-
 def log(msg):
     global LOGS
     print(msg)
@@ -27,7 +25,7 @@ def log(msg):
 # ENVIRONMENT INFO
 # ============================================================
 log("===================================================")
-log("🔍 Z-IMAGE-TURBO DEBUGGING + DETAILED TRANSFORMER INSIGHTS")
+log("🔍 Z-IMAGE-TURBO DEBUGGING + ROBUST TRANSFORMER INSPECTION")
 log("===================================================\n")
 
 log(f"📌 PYTHON VERSION       : {sys.version.replace(chr(10), ' ')}")
@@ -35,7 +33,6 @@ 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():
@@ -65,6 +62,45 @@ log(f"Model Cache Directory : {model_cache}")
 log(f"torch_dtype           : {torch_dtype}")
 log(f"USE_CPU_OFFLOAD       : {USE_CPU_OFFLOAD}")
 
+# ============================================================
+# ROBUST TRANSFORMER INSPECTION FUNCTION
+# ============================================================
+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}")
+
 # ============================================================
 # LOAD TRANSFORMER BLOCK
 # ============================================================
@@ -75,11 +111,10 @@ log("===================================================")
 quantization_config = DiffusersBitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
+    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))
 
@@ -92,25 +127,7 @@ transformer = AutoModel.from_pretrained(
     device_map=device,
 )
 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}")
+inspect_transformer(transformer, "Transformer")
 
 if USE_CPU_OFFLOAD:
     transformer = transformer.to("cpu")
@@ -125,10 +142,9 @@ log("===================================================")
 quantization_config = TransformersBitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_compute_dtype=torch_dtype,
     bnb_4bit_use_double_quant=True,
 )
-
 log("4-bit Quantization Config (Text Encoder):")
 log(str(quantization_config))
 
@@ -141,23 +157,7 @@ text_encoder = AutoModel.from_pretrained(
     device_map=device,
 )
 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}")
+inspect_transformer(text_encoder, "Text Encoder")
 
 if USE_CPU_OFFLOAD:
     text_encoder = text_encoder.to("cpu")
@@ -191,12 +191,10 @@ log("✅ Pipeline ready.")
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed):
     global LOGS
-    LOGS = ""  # Reset logs for this run
-
+    LOGS = ""  # reset logs
     log("===================================================")
     log("🎨 RUNNING INFERENCE")
     log("===================================================")
-
     log(f"Prompt     : {prompt}")
     log(f"Resolution : {width} x {height}")
     log(f"Steps      : {steps}")
@@ -212,7 +210,6 @@ def generate_image(prompt, height, width, steps, seed):
         guidance_scale=0.0,
         generator=generator,
     )
-
     log("✅ Inference Finished")
     return out.images[0], LOGS