Update app_quant_latent.py

app_quant_latent.py

@@ -593,18 +593,31 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
         try:
             latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
-            # --- FIX: Decode latent directly for preview (no unet) ---
+            # --- FIX: Convert latent tensor to PIL image ---
             try:
-                latent_img = latent_to_image(latents, pipe.vae)[0]  # rough preview
+                with torch.no_grad():
+                    latent_img_tensor = pipe.vae.decode(latents).sample  # [1, 3, H, W]
+                    latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)  # normalize
+                    latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]  # HWC
+                    latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
             except Exception:
                 latent_img = placeholder
+
             latent_gallery.append(latent_img)
 
             # Yield initial latent preview
             yield None, latent_gallery, LOGS
 
-            # Decode final image from latent
-            final_img = pipe.decode_latents(latents)[0]
+            # --- Final image from latent ---
+            try:
+                with torch.no_grad():
+                    final_img_tensor = pipe.vae.decode(latents).sample
+                    final_img_tensor = (final_img_tensor / 2 + 0.5).clamp(0, 1)
+                    final_img_tensor = final_img_tensor.cpu().permute(0, 2, 3, 1)[0]
+                    final_img = Image.fromarray((final_img_tensor.numpy() * 255).astype('uint8'))
+            except Exception:
+                final_img = placeholder
+
             final_gallery.append(final_img)
             LOGS.append("✅ Advanced latent pipeline succeeded.")
 
@@ -633,7 +646,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
                 )
                 final_img = output.images[0]
                 final_gallery.append(final_img)
-                latent_gallery.append(final_img)  # fallback latent preview
+                latent_gallery.append(final_img)  # fallback preview
                 LOGS.append("✅ Standard pipeline succeeded.")
                 yield final_img, latent_gallery, LOGS