Update app_quant_latent.py

app_quant_latent.py

@@ -589,75 +589,54 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     latent_gallery = []
     final_gallery = []
 
+    # --- Try generating latent previews ---
     try:
-        try:
-            latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
-
-            # --- Latent preview only ---
-            try:
-                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 latent preview immediately
-            yield None, latent_gallery, LOGS
-
-            # --- Final image decoding ---
-            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.")
-
-            # Save latents to dict and upload to Hugging Face
-            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}")
-
-            yield final_img, latent_gallery, LOGS
+        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
+        # Decode latent tensor to PIL for preview
+        try:
+            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)
+                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 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"⚠️ Advanced latent mode failed: {e}")
-            LOGS.append("🔁 Switching to 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
-                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
+            LOGS.append(f"⚠️ Failed to upload latents: {e}")
 
     except Exception as e:
-        LOGS.append(f"❌ Total failure: {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