Update app_quant_latent.py

app_quant_latent.py

@@ -555,90 +555,7 @@ def safe_get_latents(pipe, height, width, generator, device, LOGS):
 # Main generation function (kept exactly as your logic)
 # --------------------------
 @spaces.GPU
-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))
-    latent_gallery = []
-    final_gallery = []
-
-    try:
-        # ==========================================================
-        # ADVANCED LATENT MODE (hack using latent tensor)
-        # ==========================================================
-        try:
-            # Get initial latent tensor
-            latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
-
-            # 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), :, :]
-
-                # Convert step latent to PIL
-                try:
-                    latent_img = latent_to_image(step_latent)
-                except Exception:
-                    latent_img = placeholder
-
-                latent_gallery.append(latent_img)
-
-                # Stream intermediate latent preview to Gradio
-                yield (
-                    None,           # final image not ready
-                    latent_gallery, # gallery updates live
-                    "\n".join(LOGS),
-                )
-
-            # 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),
-            )
-
-        # ==========================================================
-        # 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(f"❌ Total failure: {e}")
-        placeholder_img = placeholder
-        yield (
-            placeholder_img,
-            [placeholder_img],
-            "\n".join(LOGS),
-        )
-
+def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0): LOGS = [] 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 = [] try: # --- 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"] latents = pipe.scheduler.step(noise_pred, t, latents)["prev_sample"] # 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"⚠️ 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 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
 
 @spaces.GPU
 def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):