Update app_quant_latent.py

app_quant_latent.py

@@ -474,50 +474,21 @@ import torch
 # Helper: Safe latent extractor
 # --------------------------
 def safe_get_latents(pipe, height, width, generator, device, LOGS):
-    """
-    Attempts multiple ways to get latents.
-    Returns a valid tensor even if pipeline hides UNet.
-    """
-    # Try official prepare_latents
-    try:
-        if hasattr(pipe, "unet") and hasattr(pipe.unet, "in_channels"):
-            num_channels = pipe.unet.in_channels
-            latents = pipe.prepare_latents(
-                batch_size=1,
-                num_channels=num_channels,
-                height=height,
-                width=width,
-                dtype=torch.float32,
-                device=device,
-                generator=generator
-            )
-            LOGS.append("✅ Latents extracted using official prepare_latents.")
-            return latents
-    except Exception as e:
-        LOGS.append(f"⚠️ Official latent extraction failed: {e}")
-
-    # Try hidden internal attribute
     try:
-        if hasattr(pipe, "_default_latents"):
-            LOGS.append("⚠️ Using hidden _default_latents.")
-            return pipe._default_latents
-    except:
-        pass
-
-    # Fallback: raw Gaussian tensor
-    try:
-        LOGS.append("⚠️ Using raw Gaussian latents fallback.")
-        return torch.randn(
-            (1, 4, height // 8, width // 8),
-            generator=generator,
+        latents = pipe.prepare_latents(
+            batch_size=1,
+            num_channels=getattr(pipe.unet, "in_channels", 4),
+            height=height,
+            width=width,
+            dtype=torch.float32,
             device=device,
-            dtype=torch.float32
+            generator=generator
         )
+        LOGS.append(f"🔹 Latents shape: {latents.shape}, dtype: {latents.dtype}, device: {latents.device}")
+        return latents
     except Exception as e:
-        LOGS.append(f"⚠️ Gaussian fallback failed: {e}")
-
-    LOGS.append("❗ Using CPU hard fallback latents.")
-    return torch.randn((1, 4, height // 8, width // 8))
+        LOGS.append(f"⚠️ Latent extraction failed: {e}")
+        return torch.randn((1, 4, height // 8, width // 8), generator=generator, device=device)
 
 
 # --------------------------
@@ -706,11 +677,11 @@ 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:
+  with gr.Blocks(title="Z-Image- experiment - dont run")as demo:
 
-    gr.Markdown("## 🔥 Z-Image Turbo — Output & Logs Viewer")
+     gr.Markdown("## 🔥 Z-Image Turbo — Output & Logs Viewer")
 
-    with gr.Tabs():
+     with gr.Tabs():
         with gr.TabItem("Output"):
             with gr.Row():
                 with gr.Column(scale=1):
@@ -728,16 +699,16 @@ def generate_image_backup(prompt, height, width, steps, seed, guidance_scale=0.0
         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):
+     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(
+     generate_btn.click(
         fn=run_and_update_logs,
         inputs=[prompt, height, width, steps, seed],
         outputs=[final_image, latent_gallery, logs_box]
-    )
+      )