Update app_quant_latent.py

app_quant_latent.py

@@ -579,6 +579,7 @@ def upload_latents_to_hf(latent_dict, filename="latents.pt"):
         os.remove(local_path)
         raise e
 
+
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
@@ -589,6 +590,90 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     latent_gallery = []
     final_gallery = []
 
+    # --- Try generating latent previews ---
+    try:
+        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
+
+        # Decode latent tensor to PIL for preview with robust fallbacks
+        latent_img = placeholder
+        try:
+            with torch.no_grad():
+                # 1️⃣ Try normal VAE decode if available
+                if hasattr(pipe, "vae") and hasattr(pipe.vae, "decode"):
+                    try:
+                        latent_img_tensor = pipe.vae.decode(latents).sample  # [1,3,H,W]
+                        latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
+                        latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
+                        latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
+                    except Exception as e1:
+                        LOGS.append(f"⚠️ VAE decode failed: {e1}")
+
+                # 2️⃣ Collapse first 3 channels if decode failed
+                if latent_img is placeholder and latents.shape[1] >= 3:
+                    ch = latents[0, :3, :, :]
+                    ch = (ch - ch.min()) / (ch.max() - ch.min() + 1e-8)
+                    latent_img = Image.fromarray((ch.permute(1, 2, 0).cpu().numpy() * 255).astype('uint8'))
+
+                # 3️⃣ Collapse all channels to mean -> replicate to RGB
+                if latent_img is placeholder:
+                    mean_ch = latents[0].mean(dim=0, keepdim=True)  # [1,H,W]
+                    mean_ch = (mean_ch - mean_ch.min()) / (mean_ch.max() - mean_ch.min() + 1e-8)
+                    latent_img = Image.fromarray(
+                        torch.cat([mean_ch]*3, dim=0).permute(1,2,0).cpu().numpy().astype('uint8')
+                    )
+
+        except Exception as e:
+            LOGS.append(f"⚠️ Latent to image conversion failed: {e}")
+            latent_img = placeholder
+
+        latent_gallery.append(latent_img)
+        yield None, latent_gallery, LOGS  # show preview immediately
+
+        # Save latents to HF for later testing
+        latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
+        try:
+            hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
+            LOGS.append(f"🔹 Latents uploaded: {hf_url}")
+        except Exception as e:
+            LOGS.append(f"⚠️ Failed to upload latents: {e}")
+
+    except Exception as e:
+        LOGS.append(f"⚠️ Latent generation failed: {e}")
+        latent_gallery.append(placeholder)
+        yield None, latent_gallery, LOGS
+
+    # --- Final image: completely untouched, uses standard pipeline ---
+    try:
+        output = pipe(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        final_img = output.images[0]
+        final_gallery.append(final_img)
+        latent_gallery.append(final_img)  # fallback preview if needed
+        LOGS.append("✅ Standard pipeline succeeded.")
+        yield final_img, latent_gallery, LOGS
+
+    except Exception as e2:
+        LOGS.append(f"❌ Standard pipeline failed: {e2}")
+        final_gallery.append(placeholder)
+        latent_gallery.append(placeholder)
+        yield placeholder, latent_gallery, LOGS
+# this version generate well for final and gives a tensor back for latent
+@spaces.GPU
+def generate_image_workswell(prompt, height, width, steps, seed, guidance_scale=0.0):
+    LOGS = []
+    device = "cuda"
+    generator = torch.Generator(device).manual_seed(int(seed))
+
+    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
+    latent_gallery = []
+    final_gallery = []
+
     # --- Try generating latent previews ---
     try:
         latents = safe_get_latents(pipe, height, width, generator, device, LOGS)