Update app.py

app.py

@@ -39,18 +39,29 @@ pipe.to("cuda")
 def pipeline_debug_info(pipe):
     info = []
 
-    # Transformer config
+    # Show basic pipeline internals
+    info.append("=== PIPELINE DEBUG INFO ===")
+
+    # Transformer info (Z-Image uses DiT/Transformer backbone)
     try:
+        info.append(f"Transformer type: {pipe.transformer.__class__.__name__}")
+        info.append(f"Transformer sample_size: {pipe.transformer.config.sample_size}")
         info.append(f"Transformer attention backend: {pipe.transformer.config.attn_implementation}")
     except:
-        info.append("No transformer.attn_implementation")
+        info.append("Transformer diagnostics failed")
 
-    # Processor class
+    # VAE diagnostics
     try:
-        proc = pipe.transformer.blocks[0].attn.processor
-        info.append(f"Processor type: {type(proc)}")
-    except Exception as e:
-        info.append(f"Processor error: {e}")
+        info.append(f"VAE latent channels: {pipe.vae.config.latent_channels}")
+        info.append(f"VAE scaling factor: {pipe.vae.config.scaling_factor}")
+    except:
+        info.append("VAE diagnostics failed")
+
+    # Attn backend globally
+    try:
+        info.append(f"Pipeline attention backend: {pipe.config.attn_implementation}")
+    except:
+        pass
 
     return "\n".join(info)
 
@@ -79,7 +90,7 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
     num_images = min(max(1, int(num_images)), 3)
 
     # Debug pipe info
-    log(pipeline_debug_info())
+    log(pipeline_debug_info(pipe))
 
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
@@ -96,18 +107,20 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         output_type="pil",
     )
 
-    # Tensor diagnostics (shapes only)
+    # Correct latent diagnostics (Z-Image uses VAE + Transformer)
     try:
-        latent_shape = pipe.unet.config.sample_size
-        log(f"UNet latent resolution (approx): {latent_shape}")
-    except:
-        pass
+        log(f"VAE latent channels: {pipe.vae.config.latent_channels}")
+        log(f"VAE scaling factor: {pipe.vae.config.scaling_factor}")
+        log(f"Transformer latent size: {pipe.transformer.config.sample_size}")
+    except Exception as e:
+        log(f"Latent diagnostics error: {e}")
 
     log("Pipeline finished.")
     log("Returning images...")
 
     return result.images, seed, log_buffer.getvalue()
 
+
 # ------------------------
 # GRADIO UI
 # ------------------------