Update app_quant_latent.py

app_quant_latent.py

@@ -553,42 +553,46 @@ def safe_get_latents(pipe, height, width, generator, device, LOGS):
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
-    latents = None
-    image = None
-    gallery = []
+    device = "cuda"
+    generator = torch.Generator(device).manual_seed(int(seed))
 
-    # placeholder image if all fails
+    # placeholders
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
-    print(prompt)
+    latent_gallery = []
+    final_gallery = []
 
     try:
-        generator = torch.Generator(device).manual_seed(int(seed))
-
-        # -------------------------------
-        # Try advanced latent extraction
-        # -------------------------------
+        # --- Try advanced latent mode ---
         try:
             latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
-            output = pipe(
-                prompt=prompt,
-                height=height,
-                width=width,
-                num_inference_steps=steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                latents=latents
-            )
+            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"]
+                    latents = pipe.scheduler.step(noise_pred, t, latents)["prev_sample"]
 
-            image = output.images[0]
-            gallery = [image]
-            LOGS.extend(getattr(pipe, "_latents_log", []))
+                # Convert latent to preview image
+                try:
+                    latent_img = latent_to_image(latents, pipe.vae)[0]
+                except Exception:
+                    latent_img = placeholder
+                latent_gallery.append(latent_img)
+
+                # Yield intermediate update: latents updated, final gallery empty
+                yield None, latent_gallery, final_gallery, LOGS
+
+            # decode final image
+            final_img = pipe.decode_latents(latents)[0]
+            final_gallery.append(final_img)
             LOGS.append("✅ Advanced latent pipeline succeeded.")
+            yield final_img, latent_gallery, final_gallery, LOGS
 
         except Exception as e:
-            LOGS.append(f"⚠️ Latent mode failed: {e}")
+            LOGS.append(f"⚠️ Advanced latent mode failed: {e}")
             LOGS.append("🔁 Switching to standard pipeline...")
 
+            # Standard pipeline fallback
             try:
                 output = pipe(
                     prompt=prompt,
@@ -598,21 +602,23 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
                     guidance_scale=guidance_scale,
                     generator=generator,
                 )
-                image = output.images[0]
-                gallery = [image]
+                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}")
-                image = placeholder
-                gallery = [image]
-
-        return image, gallery, LOGS
+                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}")
-        return placeholder, [placeholder], LOGS
-
+        final_gallery.append(placeholder)
+        latent_gallery.append(placeholder)
+        yield placeholder, latent_gallery, final_gallery, LOGS
 @spaces.GPU
 def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
     """