Update app_quant_latent.py

app_quant_latent.py

@@ -555,7 +555,86 @@ 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)) # 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
+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):