Update app_quant_latent.py

app_quant_latent.py

@@ -558,67 +558,84 @@ 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))
-    latent_gallery = []
 
-    # Gradio-safe streaming buffer
-    stream_buffer = {"updated": False}
+    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
+    latent_gallery = []
+    final_gallery = []
 
-    # --------------------------
-    # CALLBACK CALLED EACH STEP
-    # --------------------------
-    def latent_callback(step, timestep, latents):
+    try:
+        # ==========================================================
+        # ADVANCED LATENT MODE (hack using latent tensor)
+        # ==========================================================
         try:
-            img = latent_to_image(latents)     # convert latent → PIL
-        except Exception:
-            img = Image.new("RGB", (width, height), "white")
+            # Get initial latent tensor
+            latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
-        latent_gallery.append(img)
-        stream_buffer["updated"] = True
+            # Fake step-wise preview by slicing latent channels / noise
+            num_preview_steps = min(6, latents.shape[1])  # e.g. 6 slices
+            for i in range(num_preview_steps):
+                # Take a slice of latent channels to simulate intermediate step
+                step_latent = latents[:, : (i + 1), :, :]
 
-    # --------------------------
-    # 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
-        )
+                # Convert step latent to PIL
+                try:
+                    latent_img = latent_to_image(step_latent)
+                except Exception:
+                    latent_img = placeholder
 
-        # stream steps as they happen
-        for out in gen:
-            if stream_buffer["updated"]:
-                stream_buffer["updated"] = False
+                latent_gallery.append(latent_img)
+
+                # Stream intermediate latent preview to Gradio
                 yield (
-                    None,                # final image not ready
-                    latent_gallery,      # growing thumbnails
-                    "\n".join(LOGS),     # logs
+                    None,           # final image not ready
+                    latent_gallery, # gallery updates live
+                    "\n".join(LOGS),
                 )
 
-        # final result
-        final_img = out.images[0]
+            # Decode final latent tensor into 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,
-            "\n".join(LOGS),
-        )
+            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)  # last step in gallery
+            LOGS.append("✅ Standard pipeline succeeded.")
+
+            yield (
+                final_img,
+                latent_gallery,
+                "\n".join(LOGS),
+            )
 
     except Exception as e:
-        LOGS.append(str(e))
-        ph = Image.new("RGB", (width, height), "white")
+        LOGS.append(f"❌ Total failure: {e}")
+        placeholder_img = placeholder
         yield (
-            ph,
-            [ph],
+            placeholder_img,
+            [placeholder_img],
             "\n".join(LOGS),
         )