From 76484a08070f541d85b44f03c5eeeb09c4670ce6 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 26 Feb 2025 02:45:40 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 91db929..ff45230 100644
--- a/nodes.py
+++ b/nodes.py
@@ -133,7 +133,7 @@ class WanVideoModelLoader:
                 "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}),
 
             "base_precision": (["fp32", "bf16", "fp16"], {"default": "bf16"}),
-            "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2', 'fp8_scaled', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}),
+            "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}),
             "load_device": (["main_device", "offload_device"], {"default": "main_device"}),
             },
             "optional": {
@@ -363,7 +363,7 @@ class WanVideoVAELoader:
         vae = WanVideoVAE(dtype=dtype)
         vae.load_state_dict(vae_sd)
         vae.eval()
-        vae.to(device = device, dtype = dtype)
+        vae.to(device = offload_device, dtype = dtype)
             
 
         return (vae,)