Update app_quant_latent.py

app_quant_latent.py

@@ -116,30 +116,242 @@ def latent_to_image(latent):
 # SAFE TRANSFORMER INSPECTION
 # ============================================================
 def inspect_transformer(model, name):
-    log(f"\n🔍 Inspecting {name}")
+    log(f"\n🔍🔍 FULL TRANSFORMER DEBUG DUMP: {name}")
+    log("=" * 80)
+
     try:
+        log(f"Model class          : {model.__class__.__name__}")
+        log(f"DType                : {getattr(model, 'dtype', 'unknown')}")
+        log(f"Device               : {next(model.parameters()).device}")
+        log(f"Requires Grad?       : {any(p.requires_grad for p in model.parameters())}")
+
+        # Check quantization
+        if hasattr(model, "is_loaded_in_4bit"):
+            log(f"4bit Quantization    : {model.is_loaded_in_4bit}")
+        if hasattr(model, "is_loaded_in_8bit"):
+            log(f"8bit Quantization    : {model.is_loaded_in_8bit}")
+
+        # Find blocks
         candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
         blocks = None
+        chosen_attr = None
 
         for attr in candidates:
             if hasattr(model, attr):
                 blocks = getattr(model, attr)
+                chosen_attr = attr
                 break
 
+        log(f"Block container attr : {chosen_attr}")
+
+        if blocks is None:
+            log("⚠️ No valid block container found.")
+            return
+
+        if not hasattr(blocks, "__len__"):
+            log("⚠️ Blocks exist but not iterable.")
+            return
+
+        total = len(blocks)
+        log(f"Total Blocks         : {total}")
+        log("-" * 80)
+
+        # Inspect first N blocks
+        N = min(20, total)
+        for i in range(N):
+            block = blocks[i]
+            log(f"\n🧩 Block [{i}/{total-1}]")
+            log(f"Class: {block.__class__.__name__}")
+
+            # Print submodules
+            for n, m in block.named_children():
+                log(f"  ├─ {n}: {m.__class__.__name__}")
+
+            # Print attention related
+            if hasattr(block, "attn"):
+                attn = block.attn
+                log(f"  ├─ Attention: {attn.__class__.__name__}")
+                log(f"  │   Heads     : {getattr(attn, 'num_heads', 'unknown')}")
+                log(f"  │   Dim       : {getattr(attn, 'hidden_size', 'unknown')}")
+                log(f"  │   Backend   : {getattr(attn, 'attention_backend', 'unknown')}")
+
+            # Device + dtype info
+            try:
+                dev = next(block.parameters()).device
+                log(f"  ├─ Device     : {dev}")
+            except StopIteration:
+                pass
+
+            try:
+                dt = next(block.parameters()).dtype
+                log(f"  ├─ DType      : {dt}")
+            except StopIteration:
+                pass
+
+        log("\n🔚 END TRANSFORMER DEBUG DUMP")
+        log("=" * 80)
+
+    except Exception as e:
+        log(f"❌ ERROR IN INSPECTOR: {e}")
+import torch
+import time
+
+# ---------- UTILITY ----------
+def pretty_header(title):
+    log("\n\n" + "=" * 80)
+    log(f"🎛️  {title}")
+    log("=" * 80 + "\n")
+
+
+# ---------- MEMORY ----------
+def get_vram(prefix=""):
+    try:
+        allocated = torch.cuda.memory_allocated() / 1024**2
+        reserved = torch.cuda.memory_reserved() / 1024**2
+        log(f"{prefix}Allocated VRAM : {allocated:.2f} MB")
+        log(f"{prefix}Reserved VRAM  : {reserved:.2f} MB")
+    except:
+        log(f"{prefix}VRAM: CUDA not available")
+
+
+# ---------- MODULE INSPECT ----------
+def inspect_module(name, module):
+    pretty_header(f"🔬 Inspecting {name}")
+
+    try:
+        log(f"📦 Class          : {module.__class__.__name__}")
+        log(f"🔢 DType          : {getattr(module, 'dtype', 'unknown')}")
+        log(f"💻 Device         : {next(module.parameters()).device}")
+        log(f"🧮 Params         : {sum(p.numel() for p in module.parameters()):,}")
+
+        # Quantization state
+        if hasattr(module, "is_loaded_in_4bit"):
+            log(f"⚙️  4-bit QLoRA     : {module.is_loaded_in_4bit}")
+        if hasattr(module, "is_loaded_in_8bit"):
+            log(f"⚙️  8-bit load      : {module.is_loaded_in_8bit}")
+
+        # Attention backend (DiT)
+        if hasattr(module, "set_attention_backend"):
+            try:
+                attn = getattr(module, "attention_backend", None)
+                log(f"🚀 Attention Backend: {attn}")
+            except:
+                pass
+
+        # Search for blocks
+        candidates = ["transformer_blocks", "blocks", "layers", "encoder", "model"]
+        blocks = None
+        chosen_attr = None
+
+        for attr in candidates:
+            if hasattr(module, attr):
+                blocks = getattr(module, attr)
+                chosen_attr = attr
+                break
+
+        log(f"\n📚 Block Container : {chosen_attr}")
+
         if blocks is None:
-            log(f"⚠️ No block structure found in {name}")
+            log("⚠️ No block structure found")
             return
 
-        if hasattr(blocks, "__len__"):
-            log(f"Total Blocks = {len(blocks)}")
-        else:
+        if not hasattr(blocks, "__len__"):
             log("⚠️ Blocks exist but are not iterable")
+            return
+
+        total = len(blocks)
+        log(f"🔢 Total Blocks   : {total}\n")
+
+        # Inspect first 15 blocks
+        N = min(15, total)
+
+        for i in range(N):
+            blk = blocks[i]
+            log(f"\n🧩 Block [{i}/{total-1}] — {blk.__class__.__name__}")
+
+            for n, m in blk.named_children():
+                log(f"   ├─ {n:<15} {m.__class__.__name__}")
+
+            # Attention details
+            if hasattr(blk, "attn"):
+                a = blk.attn
+                log(f"   ├─ Attention")
+                log(f"   │    Heads    : {getattr(a, 'num_heads', 'unknown')}")
+                log(f"   │    Dim      : {getattr(a, 'hidden_size', 'unknown')}")
+                log(f"   │    Backend  : {getattr(a, 'attention_backend', 'unknown')}")
+
+            # Device / dtype
+            try:
+                log(f"   ├─ Device     : {next(blk.parameters()).device}")
+                log(f"   ├─ DType      : {next(blk.parameters()).dtype}")
+            except StopIteration:
+                pass
+
+        get_vram("   ▶ ")
+
+    except Exception as e:
+        log(f"❌ Module inspect error: {e}")
+
+
+# ---------- LORA INSPECTION ----------
+def inspect_loras(pipe):
+    pretty_header("🧩 LoRA ADAPTERS")
+
+    try:
+        if not hasattr(pipe, "lora_state_dict") and not hasattr(pipe, "adapter_names"):
+            log("⚠️ No LoRA system detected.")
+            return
 
-        for i in range(min(10, len(blocks) if hasattr(blocks, "__len__") else 0)):
-            log(f"Block {i} = {blocks[i].__class__.__name__}")
+        if hasattr(pipe, "adapter_names"):
+            names = pipe.adapter_names
+            log(f"Available Adapters: {names}")
+
+        if hasattr(pipe, "active_adapters"):
+            log(f"Active Adapters   : {pipe.active_adapters}")
+
+        if hasattr(pipe, "lora_scale"):
+            log(f"LoRA Scale        : {pipe.lora_scale}")
+
+        # LoRA modules
+        if hasattr(pipe, "transformer") and hasattr(pipe.transformer, "modules"):
+            for name, module in pipe.transformer.named_modules():
+                if "lora" in name.lower():
+                    log(f"   🔧 LoRA Module: {name}  ({module.__class__.__name__})")
 
     except Exception as e:
-        log(f"⚠️ Transformer inspect error: {e}")
+        log(f"❌ LoRA inspect error: {e}")
+
+
+# ---------- PIPELINE INSPECTOR ----------
+def debug_pipeline(pipe):
+    pretty_header("🚀 FULL PIPELINE DEBUGGING")
+
+    try:
+        log(f"Pipeline Class       : {pipe.__class__.__name__}")
+        log(f"Attention Impl       : {getattr(pipe, 'attn_implementation', 'unknown')}")
+        log(f"Device               : {pipe.device}")
+    except:
+        pass
+
+    get_vram("▶ ")
+
+    # Inspect TRANSFORMER
+    if hasattr(pipe, "transformer"):
+        inspect_module("Transformer", pipe.transformer)
+
+    # Inspect TEXT ENCODER
+    if hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
+        inspect_module("Text Encoder", pipe.text_encoder)
+
+    # Inspect UNET (if ZImage pipeline has it)
+    if hasattr(pipe, "unet"):
+        inspect_module("UNet", pipe.unet)
+
+    # LoRA adapters
+    inspect_loras(pipe)
+
+    pretty_header("🎉 END DEBUG REPORT")
+
 
 
 # ============================================================
@@ -240,6 +452,7 @@ try:
     
     pipe.set_adapters(["lora",], adapter_weights=[1.])
     pipe.fuse_lora(adapter_names=["lora"], lora_scale=0.75)
+    debug_pipeline(pipe)
   #  pipe.unload_lora_weights()
     pipe.to("cuda")
     log("✅ Pipeline built successfully.")