Update app_quant_latent.py

app_quant_latent.py

@@ -102,14 +102,18 @@ log_system_stats("BEFORE TRANSFORMER LOAD")
 # FUNCTION TO CONVERT LATENTS TO IMAGE
 # ============================================================
 def latent_to_image(latent):
+    """
+    Convert a latent tensor to a PIL image using pipe.vae
+    """
     try:
         img_tensor = pipe.vae.decode(latent)
         img_tensor = (img_tensor / 2 + 0.5).clamp(0, 1)
-        pil_img = T.ToPILImage()(img_tensor[0])
+        pil_img = T.ToPILImage()(img_tensor[0].cpu())  # <--- single image
         return pil_img
     except Exception as e:
         log(f"⚠️ Failed to decode latent: {e}")
-        return None
+        # fallback blank image
+        return Image.new("RGB", (latent.shape[-1]*8, latent.shape[-2]*8), color=(255,255,255))
 
 
 
@@ -574,15 +578,16 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
 
                 # Convert latent to preview image
                 try:
-                    latent_img = latent_to_image(latents, pipe.vae)[0]
+                    latent_img = latent_to_image(latents)  # returns single PIL image
                 except Exception:
                     latent_img = placeholder
+
                 latent_gallery.append(latent_img)
 
-                # Yield intermediate update: latents updated, final gallery empty
+                # Yield intermediate update: final gallery empty for now
                 yield None, latent_gallery, final_gallery, LOGS
 
-            # decode final image
+            # decode final image after all timesteps
             final_img = pipe.decode_latents(latents)[0]
             final_gallery.append(final_img)
             LOGS.append("✅ Advanced latent pipeline succeeded.")
@@ -619,6 +624,10 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
         final_gallery.append(placeholder)
         latent_gallery.append(placeholder)
         yield placeholder, latent_gallery, final_gallery, LOGS
+
+
+
+
 @spaces.GPU
 def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
     """