Update app_quant_latent.py

app_quant_latent.py

@@ -698,59 +698,51 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     latent_gallery = []
     final_gallery = []
 
-    last_latents = []  # store last 5 preview latents on CPU
+    last_latents = []
 
     try:
-        # --- Initial latents ---
+        # --- Initial latents (noise) ---
         latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
-        latents = latents.float().to(cpu_device)  # move to CPU
-
-        num_previews = min(10, steps)
-        preview_indices = torch.linspace(0, steps - 1, num_previews).long()
+        latents = latents.float().to(device)
+
+        # --- Run 1-2 diffusion steps to inject prompt info ---
+        with torch.no_grad():
+            partial_output = pipe(
+                prompt=prompt,
+                num_inference_steps=min(2, steps),
+                guidance_scale=guidance_scale,
+                generator=generator,
+                output_type="latent"
+            )
+        latents = partial_output  # latents now contain partial image info
 
-        for i, step_idx in enumerate(preview_indices):
+        # --- Last 5 previews: interpolate from partial to final latent ---
+        num_previews = 5
+        for i, alpha in enumerate(np.linspace(0.2, 1.0, num_previews)):
             try:
-                with torch.no_grad():
-                    # --- Z-Image Turbo-style denoise simulation ---
-                    t = 1.0 - (i / num_previews)  # linear decay [1.0 -> 0.0]
-                    noise_scale = t ** 0.5  # reduce noise over steps (sqrt for smoother)
-                    denoise_latent = latents * t + torch.randn_like(latents) * noise_scale
+                preview_latent = latents * alpha
+                preview_latent = preview_latent.to(pipe.vae.device).to(pipe.vae.dtype)
 
-                    # Move to VAE device & dtype
-                    denoise_latent = denoise_latent.to(pipe.vae.device).to(pipe.vae.dtype)
-
-                    # Decode latent to image
-                    decoded = pipe.vae.decode(denoise_latent, return_dict=False)[0]
-                    decoded = (decoded / 2 + 0.5).clamp(0, 1)
-                    decoded = decoded.cpu().permute(0, 2, 3, 1).float().numpy()
-                    decoded = (decoded * 255).round().astype("uint8")
-                    latent_img = Image.fromarray(decoded[0])
+                decoded = pipe.vae.decode(preview_latent, return_dict=False)[0]
+                decoded = (decoded / 2 + 0.5).clamp(0, 1)
+                decoded = decoded.cpu().permute(0, 2, 3, 1).float().numpy()
+                decoded = (decoded * 255).round().astype("uint8")
+                latent_img = Image.fromarray(decoded[0])
 
             except Exception as e:
                 LOGS.append(f"⚠️ Latent preview decode failed: {e}")
                 latent_img = placeholder
 
             latent_gallery.append(latent_img)
-
-            # Keep last 5 latents only
-            last_latents.append(denoise_latent.cpu().clone())
-            if len(last_latents) > 5:
-                last_latents.pop(0)
-
-            # Show only last 5 previews in UI
+            last_latents.append(preview_latent.cpu())
             yield None, latent_gallery[-5:], LOGS
 
-        # Optionally: upload last 5 latents
-        # latent_dict = {"latents": last_latents, "prompt": prompt, "seed": seed}
-        # hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_last5_{seed}.pt")
-        # LOGS.append(f"🔹 Last 5 latents uploaded: {hf_url}")
-
     except Exception as e:
         LOGS.append(f"⚠️ Latent generation failed: {e}")
         latent_gallery.append(placeholder)
         yield None, latent_gallery[-5:], LOGS
 
-    # --- Final image on GPU ---
+    # --- Final image: full pipeline ---
     try:
         output = pipe(
             prompt=prompt,
@@ -764,15 +756,13 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
         final_gallery.append(final_img)
         latent_gallery.append(final_img)
         LOGS.append("✅ Standard pipeline succeeded.")
-        yield final_img, latent_gallery[-5:] + [final_img], LOGS  # last 5 previews + final
+        yield final_img, latent_gallery[-5:] + [final_img], LOGS
 
     except Exception as e2:
         LOGS.append(f"❌ Standard pipeline failed: {e2}")
         final_gallery.append(placeholder)
         latent_gallery.append(placeholder)
         yield placeholder, latent_gallery[-5:] + [placeholder], LOGS
-
-
 # this is astable vesopn tha can gen final and a noise to latent
 @spaces.GPU
 def generate_image_verygood_realnoise(prompt, height, width, steps, seed, guidance_scale=0.0):