WanVideoDecode cleanup

Author: pollockjj

wanvideo.py

@@ -261,7 +261,6 @@ def loadmodel(self, model_name, precision, device=None, quantization="disabled")
 
         return text_encoder, device
 
-
 class WanVideoTextEncodeCached:
     @classmethod
     def INPUT_TYPES(s):
@@ -409,7 +408,6 @@ def loadmodel(self, model_name, load_device=None, precision="fp16", parallel=Fal
 
         return vae_model, load_device
 
-
 class WanVideoBlockSwap:
     @classmethod
     def INPUT_TYPES(s):
@@ -665,24 +663,15 @@ def VALIDATE_INPUTS(s, tile_x, tile_y, tile_stride_x, tile_stride_y):
     def decode(self, vae, load_device, samples, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y, normalization="default"):
         from . import set_current_device
 
+        original_decode = NODE_CLASS_MAPPINGS["WanVideoDecode"]()
+        decode_module = inspect.getmodule(original_decode)
+        original_module_device = decode_module.device
+
         set_current_device(load_device)      
         compute_device_to_be_patched = mm.get_torch_device()
-        
-        logger.info(f"[MultiGPU WanVideoWrapper][WanVideoDecodeMultiGPU] load device: {load_device}")
-
-        original_loader = NODE_CLASS_MAPPINGS["WanVideoDecode"]()
-        loader_module = inspect.getmodule(original_loader)
-
-        original_module_device = loader_module.device
-
-        loader_module.device = compute_device_to_be_patched
-
-        result = original_loader.decode(vae[0], samples, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y, normalization)
-
-        loader_module.device = original_module_device
-
-        decode = result[0]
-
-        return (decode,)
-
+        decode_module.device = compute_device_to_be_patched
 
+        try:
+            return original_decode.decode(vae[0], samples, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y, normalization)
+        finally:
+            decode_module.device = original_module_device