Update app_quant_latent.py

app_quant_latent.py

@@ -691,7 +691,7 @@ def generate_image_all_latents(prompt, height, width, steps, seed, guidance_scal
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
-    device = "cuda"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     generator = torch.Generator(device).manual_seed(int(seed))
 
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
@@ -711,33 +711,29 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
         for i, alpha in enumerate(preview_steps):
             try:
                 with torch.no_grad():
-                    # Simulate progression
+                    # --- simulate progression like Z-Image Turbo ---
                     preview_latent = latents * alpha + torch.randn_like(latents) * (1 - alpha)
+
                     # 🛠 FIX: move to same device as VAE and match dtype
                     preview_latent = preview_latent.to(pipe.vae.device).to(pipe.vae.dtype)
 
                     # Decode latent
-                    latent_img_tensor = pipe.vae.decode(preview_latent).sample  # [1,3,H,W]
-                    latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
+                    decoded = pipe.vae.decode(preview_latent).sample  # [1,3,H,W]
+                    decoded = (decoded / 2 + 0.5).clamp(0, 1)
 
                     # Convert to PIL
-                    latent_img_tensor = latent_img_tensor[0].permute(1, 2, 0).cpu().numpy()  # HWC
-                    latent_img = Image.fromarray((latent_img_tensor * 255).astype("uint8"))
+                    decoded = decoded[0].permute(1, 2, 0).cpu().numpy()  # HWC
+                    latent_img = Image.fromarray((decoded * 255).astype("uint8"))
 
             except Exception as e:
                 LOGS.append(f"⚠️ Latent preview decode failed: {e}")
                 latent_img = placeholder
 
             latent_gallery.append(latent_img)
-            yield None, latent_gallery, LOGS
-
-        # Upload latents (optional)
-        # latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
-        # try:
-        #     hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
-        #     LOGS.append(f"🔹 Latents uploaded: {hf_url}")
-        # except Exception as e:
-        #     LOGS.append(f"⚠️ Failed to upload latents: {e}")
+            yield None, latent_gallery, LOGS  # update Gradio with intermediate preview
+
+        # Optionally, you can store/upload last few latents here for later
+        # last_latents = latents[-4:].cpu()
 
     except Exception as e:
         LOGS.append(f"⚠️ Latent generation failed: {e}")