Update app_quant_latent.py

app_quant_latent.py

@@ -247,22 +247,76 @@ log_system_stats("AFTER PIPELINE BUILD")
 
 
 from PIL import Image
+import torch
+
+# --------------------------
+# Helper: Safe latent extractor
+# --------------------------
+def safe_get_latents(pipe, height, width, generator, device, LOGS):
+    """
+    Attempts multiple ways to get latents.
+    Returns a valid tensor even if pipeline hides UNet.
+    """
+    # Try official prepare_latents
+    try:
+        if hasattr(pipe, "unet") and hasattr(pipe.unet, "in_channels"):
+            num_channels = pipe.unet.in_channels
+            latents = pipe.prepare_latents(
+                batch_size=1,
+                num_channels=num_channels,
+                height=height,
+                width=width,
+                dtype=torch.float32,
+                device=device,
+                generator=generator
+            )
+            LOGS.append("✅ Latents extracted using official prepare_latents.")
+            return latents
+    except Exception as e:
+        LOGS.append(f"⚠️ Official latent extraction failed: {e}")
+
+    # Try hidden internal attribute
+    try:
+        if hasattr(pipe, "_default_latents"):
+            LOGS.append("⚠️ Using hidden _default_latents.")
+            return pipe._default_latents
+    except:
+        pass
+
+    # Fallback: raw Gaussian tensor
+    try:
+        LOGS.append("⚠️ Using raw Gaussian latents fallback.")
+        return torch.randn(
+            (1, 4, height // 8, width // 8),
+            generator=generator,
+            device=device,
+            dtype=torch.float32
+        )
+    except Exception as e:
+        LOGS.append(f"⚠️ Gaussian fallback failed: {e}")
+
+    LOGS.append("❗ Using CPU hard fallback latents.")
+    return torch.randn((1, 4, height // 8, width // 8))
+
+
+# --------------------------
+# Main generation function
+# --------------------------
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
-
     LOGS = []
     latents = None
     image = None
     gallery = []
 
-    # safe placeholder image
+    # placeholder image if all fails
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
 
     try:
         generator = torch.Generator(device).manual_seed(int(seed))
 
         # -------------------------------
-        # Try advanced latent extractor
+        # Try advanced latent extraction
         # -------------------------------
         try:
             latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
@@ -285,9 +339,6 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
             LOGS.append(f"⚠️ Latent mode failed: {e}")
             LOGS.append("🔁 Switching to standard pipeline...")
 
-            # -------------------------------
-            # Standard generation fallback
-            # -------------------------------
             try:
                 output = pipe(
                     prompt=prompt,
@@ -306,13 +357,11 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
                 image = placeholder
                 gallery = [image]
 
-        # final safe return
         return image, gallery, LOGS
 
     except Exception as e:
         LOGS.append(f"❌ Total failure: {e}")
         return placeholder, [placeholder], LOGS
-
 @spaces.GPU
 def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
     """