Enable VAE tiling+slicing only when resolution >= 768 (inside @spaces.GPU); reset after the call

app.py

@@ -241,6 +241,16 @@ def generate(
     num_inference_steps = int(num_inference_steps)
     H = W = int(resolution)
 
+    # Heavy resolutions push the VAE decode over the 48GB MIG budget — slice + tile.
+    vae_tiling_enabled = False
+    if H >= 768:
+        try:
+            pipeline.vae.enable_tiling()
+            pipeline.vae.enable_slicing()
+            vae_tiling_enabled = True
+        except Exception as _e:
+            print(f"[pid] VAE tiling/slicing not available: {_e}", flush=True)
+
     # initial: show the live preview, hide the final slider
     yield gr.update(visible=True, value=None, label="Generating Z-Image…"), gr.update(visible=False, value=None), gr.update(value=seed)
 
@@ -303,6 +313,9 @@ def generate(
         (baseline_01[0].clamp(0, 1).permute(1, 2, 0).float().cpu().numpy() * 255).astype(np.uint8)
     )
 
+    # Free Z-Image VAE intermediates before PiD takes over the GPU
+    torch.cuda.empty_cache()
+
     # ---- PiD upscaling on the final latent, streaming the 4 internal steps ----
     final_sigma = float(pipeline.scheduler.sigmas[-1].item())
     pid_img = None
@@ -314,6 +327,13 @@ def generate(
             gr.update(),
         )
 
+    if vae_tiling_enabled:
+        try:
+            pipeline.vae.disable_tiling()
+            pipeline.vae.disable_slicing()
+        except Exception:
+            pass
+
     # ---- Done: hide live preview, show the A/B slider ----
     yield (
         gr.update(visible=False, value=None),