Update app_quant_latent.py

app_quant_latent.py

@@ -464,6 +464,28 @@ except Exception as e:
 log_system_stats("AFTER PIPELINE BUILD")
 
 
+# -----------------------------
+# Monkey-patch prepare_latents
+# -----------------------------
+original_prepare_latents = pipe.prepare_latents
+
+def logged_prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+    # Call original method
+    result_latents = original_prepare_latents(
+        batch_size, num_channels_latents, height, width, dtype, device, generator, latents
+    )
+
+    # Save info for logging
+    log_msg = f"🔹 prepare_latents called | shape={result_latents.shape}, dtype={result_latents.dtype}, device={result_latents.device}"
+    if hasattr(self, "_latents_log"):
+        self._latents_log.append(log_msg)
+    else:
+        self._latents_log = [log_msg]
+
+    return result_latents
+
+# Apply patch
+pipe.prepare_latents = logged_prepare_latents.__get__(pipe)
 
 
 
@@ -526,6 +548,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
 
             image = output.images[0]
             gallery = [image]
+            LOGS.extend(getattr(pipe, "_latents_log", []))
             LOGS.append("✅ Advanced latent pipeline succeeded.")
 
         except Exception as e:
@@ -677,38 +700,39 @@ def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0
 #       outputs=[final_image, latent_gallery, logs_box]
 #      )
 
-with gr.Blocks(title="Z-Image- experiment - dont run")as demo:
-
-  gr.Markdown("## 🔥 Z-Image Turbo — Output & Logs Viewer")
-
-     with gr.Tabs():
-        with gr.TabItem("Output"):
+with gr.Blocks(title="Z-Image-Turbo") as demo:
+    with gr.Tabs():
+        with gr.TabItem("Image & Latents"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    prompt = gr.Textbox(label="Prompt", lines=2, value="boat in ocean")
+                    prompt = gr.Textbox(label="Prompt", value="boat in Ocean")
                     height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
                     width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
-                    steps = gr.Slider(1, 50, value=20, step=1, label="Steps")
+                    steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
                     seed = gr.Number(value=42, label="Seed")
-                    generate_btn = gr.Button("🚀 Generate Image")
+                    run_btn = gr.Button("Generate Image")
 
                 with gr.Column(scale=1):
                     final_image = gr.Image(label="Final Image")
-                    latent_gallery = gr.Gallery(label="Latent Steps", columns=4, height=256, preview=True)
+                    latent_gallery = gr.Gallery(
+                        label="Latent Steps", columns=4, height=256, preview=True
+                    )
+
+            run_btn.click(
+                generate_image,
+                inputs=[prompt, height, width, steps, seed],
+                outputs=[final_image, latent_gallery, None]
+            )
 
         with gr.TabItem("Logs"):
-            logs_box = gr.HTML(value=f"<pre'>{LOGS}</pre>", label="Full Logs")
-
-        def run_and_update_logs(prompt, height, width, steps, seed):
-          img, gallery, logs = generate_image(prompt, height, width, steps, seed)
-          combined_logs = "\n".join(logs) + "\n\n" + LOGS  # append global logs
-          return img, gallery, f"<pre id='logbox'>{combined_logs}</pre>"
-
-         generate_btn.click(
-             fn=run_and_update_logs,
-              inputs=[prompt, height, width, steps, seed],
-            outputs=[final_image, latent_gallery, logs_box]
-         )
+            logs_box = gr.Textbox(label="All Logs", lines=25)
+            
+            run_btn.click(
+                generate_image,
+                inputs=[prompt, height, width, steps, seed],
+                outputs=[None, None, logs_box]
+            )
+