Update app_quant_latent.py

app_quant_latent.py

@@ -513,22 +513,39 @@ import torch
 # Helper: Safe latent extractor
 # --------------------------
 def safe_get_latents(pipe, height, width, generator, device, LOGS):
+    """
+    Safely prepare latents for any ZImagePipeline variant.
+    Returns latents tensor, logs issues instead of failing.
+    """
     try:
+        # Determine number of channels
+        num_channels = 4  # default fallback
+        if hasattr(pipe, "unet") and hasattr(pipe.unet, "in_channels"):
+            num_channels = pipe.unet.in_channels
+        elif hasattr(pipe, "vae") and hasattr(pipe.vae, "latent_channels"):
+            num_channels = pipe.vae.latent_channels  # some pipelines define this
+        LOGS.append(f"🔹 Using num_channels={num_channels} for latents")
+
         latents = pipe.prepare_latents(
             batch_size=1,
-            num_channels=getattr(pipe.unet, "in_channels", 4),
+            num_channels_latents=num_channels,
             height=height,
             width=width,
             dtype=torch.float32,
             device=device,
-            generator=generator
+            generator=generator,
         )
+
         LOGS.append(f"🔹 Latents shape: {latents.shape}, dtype: {latents.dtype}, device: {latents.device}")
         return latents
     except Exception as e:
         LOGS.append(f"⚠️ Latent extraction failed: {e}")
-        return torch.randn((1, 4, height // 8, width // 8), generator=generator, device=device)
-
+        # fallback: guess a safe shape
+        fallback_channels = 16  # try standard default for ZImage pipelines
+        latents = torch.randn((1, fallback_channels, height // 8, width // 8),
+                              generator=generator, device=device)
+        LOGS.append(f"🔹 Using fallback random latents shape: {latents.shape}")
+        return latents
 
 # --------------------------
 # Main generation function (kept exactly as your logic)