From 83948cafd1b79a4cd6fb5f71038b6d28daf04868 Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Thu, 6 Mar 2025 07:32:04 -0800
Subject: [PATCH] WAN 2.1 support

---
 comfy/app/frontend_management.py | 2 +-
 comfy/model_management.py        | 2 +-
 comfy/nodes/base_nodes.py        | 2 +-
 comfy/nodes/package_typing.py    | 2 +-
 comfy/sd.py                      | 1 +
 comfy_extras/nodes/nodes_wan.py  | 4 ++--
 6 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/comfy/app/frontend_management.py b/comfy/app/frontend_management.py
index 2e8f70700..e1805ede1 100644
--- a/comfy/app/frontend_management.py
+++ b/comfy/app/frontend_management.py
@@ -17,7 +17,7 @@ import importlib.resources
 from typing_extensions import NotRequired
 
 from ..cli_args import DEFAULT_VERSION_STRING
-from ..cmd.folder_paths import add_model_folder_path
+from ..cmd.folder_paths import add_model_folder_path  # pylint: disable=import-error
 
 REQUEST_TIMEOUT = 10  # seconds
 
diff --git a/comfy/model_management.py b/comfy/model_management.py
index 9b4526ec8..bd2dd0809 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -120,7 +120,7 @@ except:
     npu_available = False
 
 try:
-    import torch_mlu  # noqa: F401
+    import torch_mlu  # pylint: disable=import-error, noqa: F401
 
     _ = torch.mlu.device_count()
     mlu_available = torch.mlu.is_available()
diff --git a/comfy/nodes/base_nodes.py b/comfy/nodes/base_nodes.py
index 93cc25be3..0b557c685 100644
--- a/comfy/nodes/base_nodes.py
+++ b/comfy/nodes/base_nodes.py
@@ -977,7 +977,7 @@ class CLIPLoader:
         elif type == "lumina2":
             clip_type = sd.CLIPType.LUMINA2
         elif type == "wan":
-            clip_type = comfy.sd.CLIPType.WAN
+            clip_type = sd.CLIPType.WAN
         else:
             logging.warning(f"Unknown clip type argument passed: {type} for model {clip_name}")
 
diff --git a/comfy/nodes/package_typing.py b/comfy/nodes/package_typing.py
index f449f65ec..bb0ce8431 100644
--- a/comfy/nodes/package_typing.py
+++ b/comfy/nodes/package_typing.py
@@ -101,7 +101,7 @@ class SaveNodeResultT(TypedDict, total=True):
     type: Literal["output", "input", "temp"]
 
 
-SaveNodeResult = SaveNodeResultT | FileLocator
+SaveNodeResult = SaveNodeResultT
 
 
 class UIImagesImagesResult(TypedDict, total=True):
diff --git a/comfy/sd.py b/comfy/sd.py
index 9e2dcb52f..5c24ae96e 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+import json
 import logging
 import math
 import os.path
diff --git a/comfy_extras/nodes/nodes_wan.py b/comfy_extras/nodes/nodes_wan.py
index dc30eb546..e4b368b91 100644
--- a/comfy_extras/nodes/nodes_wan.py
+++ b/comfy_extras/nodes/nodes_wan.py
@@ -1,5 +1,5 @@
-import nodes
-import node_helpers
+from comfy.nodes import base_nodes as nodes
+from comfy import node_helpers
 import torch
 import comfy.model_management
 import comfy.utils