Update app_quant_latent.py

app_quant_latent.py

@@ -581,62 +581,61 @@ def upload_latents_to_hf(latent_dict, filename="latents.pt"):
 
 
 
+import asyncio
+import torch
+from PIL import Image
+
+async def async_upload_latents(latent_dict, filename, LOGS):
+    try:
+        hf_url = await upload_latents_to_hf(latent_dict, filename=filename)  # assume this can be async
+        LOGS.append(f"🔹 All preview latents uploaded: {hf_url}")
+    except Exception as e:
+        LOGS.append(f"⚠️ Failed to upload all preview latents: {e}")
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
-    device = "cuda"
+    device = "cpu"  # CPU mode
     generator = torch.Generator(device).manual_seed(int(seed))
 
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
     latent_gallery = []
     final_gallery = []
 
-    all_latents = []  # store all preview latents
+    all_latents = []
 
     # --- Try generating latent previews ---
     try:
         latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
+        latents = latents.to("cpu")  # ensure CPU dtype
 
-        # Loop through timesteps for preview generation
         for i, t in enumerate(pipe.scheduler.timesteps):
             try:
-                # Convert latent tensor to PIL for preview
                 with torch.no_grad():
-                    # Some pipelines may require same dtype as bias
-                    latent_to_decode = latents.to(pipe.vae.dtype)
+                    latent_to_decode = latents.to(pipe.vae.dtype).cpu()
                     latent_img_tensor = pipe.vae.decode(latent_to_decode).sample  # [1,3,H,W]
                     latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
-                    latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
+                    latent_img_tensor = latent_img_tensor.permute(0, 2, 3, 1)[0]
                     latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype("uint8"))
             except Exception:
                 latent_img = placeholder
                 LOGS.append("⚠️ Latent preview decode failed.")
 
             latent_gallery.append(latent_img)
-            all_latents.append(latents.cpu().clone())  # save current latent
+            all_latents.append(latents.cpu().clone())
 
-            # Yield intermediate preview every few steps
-            if i % max(1, len(pipe.scheduler.timesteps) // 10) == 0:
+            if i % max(1, len(pipe.scheduler.timesteps)//10) == 0:
                 yield None, latent_gallery, LOGS
 
-        # Upload full series of latents
-        try:
-            latent_dict = {
-                "latents_series": all_latents,
-                "prompt": prompt,
-                "seed": seed
-            }
-            hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_series_{seed}.pt")
-            LOGS.append(f"🔹 All preview latents uploaded: {hf_url}")
-        except Exception as e:
-            LOGS.append(f"⚠️ Failed to upload all preview latents: {e}")
+        # --- Async upload of full latent series ---
+        latent_dict = {"latents_series": all_latents, "prompt": prompt, "seed": seed}
+        asyncio.create_task(async_upload_latents(latent_dict, f"latents_series_{seed}.pt", LOGS))
 
     except Exception as e:
         LOGS.append(f"⚠️ Latent generation failed: {e}")
         latent_gallery.append(placeholder)
         yield None, latent_gallery, LOGS
 
-    # --- Final image: completely untouched, uses standard pipeline ---
+    # --- Final image: unchanged, uses standard pipeline ---
     try:
         output = pipe(
             prompt=prompt,
@@ -648,7 +647,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
         )
         final_img = output.images[0]
         final_gallery.append(final_img)
-        latent_gallery.append(final_img)  # fallback preview if needed
+        latent_gallery.append(final_img)
         LOGS.append("✅ Standard pipeline succeeded.")
         yield final_img, latent_gallery, LOGS
 
@@ -743,6 +742,7 @@ def generate_image0(prompt, height, width, steps, seed, guidance_scale=0.0):
 
 
 with gr.Blocks(title="Z-Image-Turbo") as demo:
+    gr.Markdown("# 🎨 DO NOT RUN THIS ")
     with gr.Tabs():
         with gr.TabItem("Image & Latents"):
             with gr.Row():