Update app_quant_latent.py

app_quant_latent.py

@@ -560,66 +560,70 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     device = "cuda"
     generator = torch.Generator(device).manual_seed(int(seed))
 
+    # placeholders
+    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
     latent_gallery = []
-    final_image = None
+    final_gallery = []
 
     try:
-        # Try advanced latent mode
+        # --- Try advanced latent mode ---
         try:
             latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
             for i, t in enumerate(pipe.scheduler.timesteps):
+                # Step-wise denoising
                 with torch.no_grad():
-                    noise_pred = pipe.unet(
-                        latents, 
-                        t, 
-                        encoder_hidden_states=pipe.get_text_embeddings(prompt)
-                    )["sample"]
-
+                    noise_pred = pipe.unet(latents, t, encoder_hidden_states=pipe.get_text_embeddings(prompt))["sample"]
                     latents = pipe.scheduler.step(noise_pred, t, latents)["prev_sample"]
 
-                # Convert to preview
+                # Convert latent to preview image
                 try:
-                    latent_img = latent_to_image(latents)
+                    latent_img = latent_to_image(latents)  # returns single PIL image
                 except Exception:
-                    latent_img = Image.new("RGB", (width, height), "white")
+                    latent_img = placeholder
 
                 latent_gallery.append(latent_img)
 
-                # ---- LIVE UPDATE ----
-                yield latent_gallery, None, "\n".join(LOGS)
+                # Yield intermediate update: final gallery empty for now
+                yield None, latent_gallery, final_gallery, LOGS
 
-            # Final decode
-            final = pipe.decode_latents(latents)[0]
-            final_image = final
+            # decode final image after all timesteps
+            final_img = pipe.decode_latents(latents)[0]
+            final_gallery.append(final_img)
             LOGS.append("✅ Advanced latent pipeline succeeded.")
-
-            # ---- FINAL OUTPUT ----
-            yield latent_gallery, final_image, "\n".join(LOGS)
+            yield final_img, latent_gallery, final_gallery, LOGS
 
         except Exception as e:
             LOGS.append(f"⚠️ Advanced latent mode failed: {e}")
             LOGS.append("🔁 Switching to standard pipeline...")
 
-            output = pipe(
-                prompt=prompt,
-                height=height,
-                width=width,
-                num_inference_steps=steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-            )
-
-            final_image = output.images[0]
-            latent_gallery.append(final_image)
-            LOGS.append("✅ Standard pipeline succeeded.")
-
-            yield latent_gallery, final_image, "\n".join(LOGS)
+            # Standard pipeline fallback
+            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)  # optionally show in latent gallery as last step
+                LOGS.append("✅ Standard pipeline succeeded.")
+                yield final_img, latent_gallery, final_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, final_gallery, LOGS
 
     except Exception as e:
         LOGS.append(f"❌ Total failure: {e}")
-        placeholder = Image.new("RGB", (width, height), "white")
-        yield [placeholder], placeholder, "\n".join(LOGS)
+        final_gallery.append(placeholder)
+        latent_gallery.append(placeholder)
+        yield placeholder, latent_gallery, final_gallery, LOGS