Update app_quant_latent.py

app_quant_latent.py

@@ -248,52 +248,6 @@ log_system_stats("AFTER PIPELINE BUILD")
 
 @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
-
-    # 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_history, LOGS
-
- except Exception as e:
-    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 = []
@@ -353,9 +307,10 @@ def generate_image(prompt, height, width, steps, seed):
 # ============================================================
 
 with gr.Blocks(title="Z-Image-Turbo Generator") as demo:
-   gr.Markdown("# **🚀 Z-Image-Turbo — Final Image & Latents**")
+  gr.Markdown("# **🚀 Z-Image-Turbo — Final Image & Latents**")
 
-   with gr.Row():
+
+  with gr.Row():
     with gr.Column(scale=1):
         prompt = gr.Textbox(label="Prompt", value="Realistic mid-aged male image")
         height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
@@ -366,11 +321,15 @@ 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", grid=[4], height=256)
+        latent_gallery = gr.Gallery(label="Latent Steps").style(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()