Update app_quant_latent.py

app_quant_latent.py

@@ -560,70 +560,66 @@ 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_gallery = []
+    final_image = None
 
     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 latent to preview image
+                # Convert to preview
                 try:
-                    latent_img = latent_to_image(latents)  # returns single PIL image
+                    latent_img = latent_to_image(latents)
                 except Exception:
-                    latent_img = placeholder
+                    latent_img = Image.new("RGB", (width, height), "white")
 
                 latent_gallery.append(latent_img)
 
-                # Yield intermediate update: final gallery empty for now
-                yield latent_gallery, None, LOGS  # latent updates, final image empty
+                # ---- LIVE UPDATE ----
+                yield latent_gallery, None, "\n".join(LOGS)
 
-            # decode final image after all timesteps
-            final_img = pipe.decode_latents(latents)[0]
-            final_gallery.append(final_img)
+            # Final decode
+            final = pipe.decode_latents(latents)[0]
+            final_image = final
             LOGS.append("✅ Advanced latent pipeline succeeded.")
-            yield latent_gallery, final_img, LOGS 
+
+            # ---- FINAL OUTPUT ----
+            yield latent_gallery, final_image, "\n".join(LOGS)
 
         except Exception as e:
             LOGS.append(f"⚠️ Advanced latent mode failed: {e}")
             LOGS.append("🔁 Switching to standard pipeline...")
 
-            # 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
+            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)
 
     except Exception as e:
         LOGS.append(f"❌ Total failure: {e}")
-        final_gallery.append(placeholder)
-        latent_gallery.append(placeholder)
-        yield placeholder, latent_gallery, final_gallery, LOGS
+        placeholder = Image.new("RGB", (width, height), "white")
+        yield [placeholder], placeholder, "\n".join(LOGS)