From fe9fb7e2f0e16e66ca9ff9ff40c9ebb62e251c16 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 31 Mar 2025 01:25:37 +0300
Subject: [PATCH] Only use custom preview code for TAEW

---
 nodes.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 5a7c688..0454b97 100644
--- a/nodes.py
+++ b/nodes.py
@@ -28,6 +28,7 @@ import comfy.model_base
 import comfy.latent_formats
 from comfy.clip_vision import clip_preprocess, ClipVisionModel
 from comfy.sd import load_lora_for_models
+from comfy.cli_args import args, LatentPreviewMethod
 
 script_directory = os.path.dirname(os.path.abspath(__file__))
 
@@ -1944,7 +1945,10 @@ class WanVideoSampler:
            
         pbar = ProgressBar(steps)
 
-        from .latent_preview import prepare_callback
+        if args.preview_method in [LatentPreviewMethod.Auto, LatentPreviewMethod.Latent2RGB]: #default for latent2rgb
+            from latent_preview import prepare_callback
+        else:
+            from .latent_preview import prepare_callback #custom for tiny VAE previews
         callback = prepare_callback(patcher, steps)
 
         #blockswap init