Update app_quant_latent.py

app_quant_latent.py

@@ -558,93 +558,71 @@ def safe_get_latents(pipe, height, width, generator, device, LOGS):
 def generate_image(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))
+    generator = torch.Generator(device).manual_seed(int(seed))
     latent_gallery = []
-    final_gallery = []
-
-    try:
-        # ==========================================================
-        # ADVANCED LATENT MODE
-        # ==========================================================
-        try:
-            latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
-            for i, t in enumerate(pipe.scheduler.timesteps):
-                with torch.no_grad():
-                    noise_pred = pipe.unet(
-                        latents, t,
-                        encoder_hidden_states=pipe.get_text_embeddings(prompt)
-                    )["sample"]
+    # Gradio-safe streaming buffer
+    stream_buffer = {"updated": False}
 
-                    latents = pipe.scheduler.step(noise_pred, t, latents)["prev_sample"]
+    # --------------------------
+    # CALLBACK CALLED EACH STEP
+    # --------------------------
+    def latent_callback(step, timestep, latents):
+        try:
+            img = latent_to_image(latents)     # convert latent → PIL
+        except Exception:
+            img = Image.new("RGB", (width, height), "white")
 
-                # convert latent → image
-                try:
-                    latent_img = latent_to_image(latents)
-                except Exception:
-                    latent_img = placeholder
+        latent_gallery.append(img)
+        stream_buffer["updated"] = True
 
-                latent_gallery.append(latent_img)
+    # --------------------------
+    # MAIN GENERATOR LOOP
+    # --------------------------
+    try:
+        # start generation
+        gen = pipe.generate(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            callback=latent_callback,
+            callback_steps=1,
+            stream=True,     # << IMPORTANT for manual iteration
+        )
 
-                # 🔥 STREAM update
+        # stream steps as they happen
+        for out in gen:
+            if stream_buffer["updated"]:
+                stream_buffer["updated"] = False
                 yield (
-                    None,               # final_image
-                    latent_gallery,     # latent gallery list
-                    "\n".join(LOGS),    # logs
+                    None,                # final image not ready
+                    latent_gallery,      # growing thumbnails
+                    "\n".join(LOGS),     # logs
                 )
 
-            # ---------------------
-            # FINAL decode
-            # ---------------------
-            final_img = pipe.decode_latents(latents)[0]
-            final_gallery.append(final_img)
-            LOGS.append("✅ Advanced latent pipeline succeeded.")
-
-            yield (
-                final_img,
-                latent_gallery,
-                "\n".join(LOGS),
-            )
-
-        # ==========================================================
-        # FALLBACK STANDARD PIPELINE
-        # ==========================================================
-        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_img = output.images[0]
-            latent_gallery.append(final_img)
-            LOGS.append("✅ Standard pipeline succeeded.")
+        # final result
+        final_img = out.images[0]
 
-            yield (
-                final_img,
-                latent_gallery,
-                "\n".join(LOGS),
-            )
+        yield (
+            final_img,
+            latent_gallery,
+            "\n".join(LOGS),
+        )
 
     except Exception as e:
-        LOGS.append(f"❌ Total failure: {e}")
-        placeholder_img = placeholder
+        LOGS.append(str(e))
+        ph = Image.new("RGB", (width, height), "white")
         yield (
-            placeholder_img,
-            [placeholder_img],
+            ph,
+            [ph],
             "\n".join(LOGS),
         )
 
 
-
 @spaces.GPU
 def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
     """