Update app_quant_latent.py

app_quant_latent.py

@@ -291,6 +291,62 @@ def generate_image(prompt, height, width, steps, seed):
     log(f"❌ Inference error: {e}")
     return None, None, LOGS
 
+@spaces.GPU
+def generate_image(prompt, height, width, steps, seed):
+
+ try:
+    generator = torch.Generator(device).manual_seed(seed)
+    latent_history = []
+
+    # Callback to save latents and GPU info
+    def save_latents(step, timestep, latents):
+        latent_history.append(latents.detach().clone())
+        gpu_mem = torch.cuda.memory_allocated(0)/1e9
+        log(f"Step {step} - GPU Memory Used: {gpu_mem:.2f} GB")
+
+    # Step-wise loop just for latent capture
+    for step, _ in pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=0.0,
+        generator=generator,
+        callback=save_latents,
+        callback_steps=1
+    ).iter():
+        pass  # only capturing latents, ignoring intermediate images
+
+    # Convert latents to PIL images for gallery
+    latent_images = []
+    for latent in latent_history:
+        try:
+            img_tensor = pipe.vae.decode(latent)
+            img_tensor = (img_tensor / 2 + 0.5).clamp(0, 1)
+            pil_img = T.ToPILImage()(img_tensor[0].cpu())
+            latent_images.append(pil_img)
+        except Exception as e:
+            log(f"⚠️ Failed to convert latent to image: {e}")
+
+    # Original final image generation
+    output = pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=0.0,
+        generator=generator,
+    )
+
+    log("✅ Inference finished.")
+    log_system_stats("AFTER INFERENCE")
+
+    return output.images[0], latent_images, LOGS
+
+ except Exception as e:
+    log(f"❌ Inference error: {e}")
+    return None, None, LOGS
+
 
 # ============================================================
 # UI
@@ -310,14 +366,11 @@ with gr.Blocks(title="Z-Image-Turbo Generator") as demo:
 
     with gr.Column(scale=1):
         final_image = gr.Image(label="Final Image")
-        latent_gallery = gr.Gallery(label="Latent Steps").style(grid=[4], height="256px")
+        latent_gallery = gr.Gallery(label="Latent Steps", grid=[4], height=256)
         logs_box = gr.Textbox(label="Logs", lines=15)
 
-    run_btn.click(
-      generate_image,
-       inputs=[prompt, height, width, steps, seed],
-        outputs=[final_image, latent_gallery, logs_box]
-     )
+  
+    run_btn.click( generate_image, inputs=[prompt, height, width, steps, seed], outputs=[final_image, latent_gallery, logs_box] )
 
 
 demo.launch()