From e32afb355415210eac19e2a65fae55deb7b51d4d Mon Sep 17 00:00:00 2001
From: jax <jaxinformationtransformer@gmail.com>
Date: Sat, 8 Mar 2025 22:51:22 +0800
Subject: [PATCH] init Wan Node

---
 __init__.py    |  4 +++-
 nodes/nodes.py | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 9582ded..8d045b8 100644
--- a/__init__.py
+++ b/__init__.py
@@ -9,7 +9,7 @@ from .nodes.save_file import ComfyOnlineSaveFile
 from .nodes.LoadEmbedding import EmbeddingLoader
 from .nodes.AudioSave import SaveAudioAsWav
 from .nodes.TextSave import TextSave
-
+from .nodes.nodes import LoadWanVideoLoraFromCivitAIWithDownloader
 from .nodes.LoadCheckpointFromCivitAI import LoadCheckpointFromCivitAIWithDownloader
 
 NODE_CLASS_MAPPINGS = { 
@@ -17,6 +17,7 @@ NODE_CLASS_MAPPINGS = {
     "LoadLoraFromComfyOnline":LoadLoraFromComfyOnlineWithDownloader,
     "LoadHunyuanLoraFromCivitAI":LoadHunyuanLoraFromCivitAIWithDownloader,
     "LoadHunyuanLoraFromComfyOnline":LoadHunyuanLoraFromComfyOnlineWithDownloader,
+    "LoadWanVideoLoraFromCivitAI":LoadWanVideoLoraFromCivitAIWithDownloader,
     "LoadHunyuanLoraFromHF":LoadHunyuanLoraFromHFWithDownloader,
     "LoadLoraFromHF":LoadLoraFromHFWithDownloader,
     "LoadLoraFromHFWithToken": LoadLoraFromHFWithToken,
@@ -36,6 +37,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "LoadHunyuanLoraFromComfyOnline":" Load Lora From ComfyOnline",
     "LoadHunyuanLoraFromCivitAI": "Load HunyuanLora From CivitAI",
     "LoadHunyuanLoraFromHF": "Load HunyuanLora From HF",
+    "LoadWanVideoLoraFromCivitAI": "Load WanVideoLora From CivitAI",
     'ComfyOnlineUploadAnything': "ComfyOnlineUploadAnything",
     "ComfyOnlineSaveFile": "ComfyOnlineSaveFile",
     "EmbeddingLoader": "Load Embedding",
diff --git a/nodes/nodes.py b/nodes/nodes.py
index 72721a1..750735b 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -14,6 +14,66 @@ import os
 folder_paths.add_model_folder_path("tmp_hunyuan_loras", os.path.join(folder_paths.models_dir, "tmp_hunyuan_loras"))
 
 
+class LoadWanVideoLoraFromCivitAIWithDownloader:
+    def __init__(self):
+        self.loaded_lora = None
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "civitai_model_id": ("STRING", {"default": "", "tooltip": "The ID of the model to download from CivitAI."}),
+                "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.0001, "tooltip": "LORA strength, set to 0.0 to unmerge the LORA"}),
+            },
+            "optional": {
+                "prev_lora":("HYVIDLORA", {"default": None, "tooltip": "For loading multiple LoRAs"}),
+                "blocks":("SELECTEDBLOCKS", ),
+            }
+        }
+
+    RETURN_TYPES = ("WANVIDLORA",)
+    RETURN_NAMES = ("lora", )
+    FUNCTION = "getlorapath"
+    CATEGORY = "WanVideoWrapper"
+    DESCRIPTION = "Select a LoRA model from civitAI"
+
+    def getlorapath(self, civitai_model_id, strength, blocks=None, prev_lora=None, fuse_lora=False):
+        # 获取 CivitAI Token
+        civitai_token_id = os.getenv("CIVITAI_TOKEN", "").strip()
+        if not civitai_token_id:
+            raise RuntimeError("CIVITAI_TOKEN environment variable is not set or empty.")
+        # 目标存储路径为 loras 目录
+        loras_dir = folder_paths.get_folder_paths("tmp_wanvideo_loras")[0]
+
+        # 下载文件到 loras 目录
+        lora_filename = f"tmp_civit_{civitai_model_id or 'downloaded_lora'}.safetensors"  # 生成临时文件名
+        lora_path = os.path.join(loras_dir, lora_filename)
+        
+        self.download_from_civitai(civitai_model_id, civitai_token_id, lora_path)
+
+        loras_list = []
+        lora = {
+            "path": lora_path,
+            "strength": strength,
+            "name": lora_filename,
+            "fuse_lora": fuse_lora,
+            "blocks": blocks
+        }
+        if prev_lora is not None:
+            loras_list.extend(prev_lora)
+
+        loras_list.append(lora)
+        return (loras_list,)
+
+    def download_from_civitai(self, model_id, token_id, lora_path):
+        print("Downloading LoRA from CivitAI")
+        print(f"\tModel ID: {model_id}")
+        print(f"\tToken ID: {token_id}")
+        print(f"\tSave path: {lora_path}")
+        # 实现下载逻辑
+        download_cai(model_id, token_id, lora_path)
+
+
 class LoadHunyuanLoraFromCivitAIWithDownloader:
     def __init__(self):
         self.loaded_lora = None