Vace training && Comfyui update && Fix flux readme && Update MV2V demo

comfyui/cogvideox_fun/nodes.py

@@ -198,13 +198,18 @@ def INPUT_TYPES(s):
     CATEGORY = "CogVideoXFUNWrapper"
 
     def load_lora(self, cogvideoxfun_model, lora_name, strength_model, lora_cache):
+        new_funmodels = dict(cogvideoxfun_model)  
+
         if lora_name is not None:
-            cogvideoxfun_model['lora_cache'] = lora_cache
-            cogvideoxfun_model['loras'] = cogvideoxfun_model.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)]
-            cogvideoxfun_model['strength_model'] = cogvideoxfun_model.get("strength_model", []) + [strength_model]
-            return (cogvideoxfun_model,)
-        else:
-            return (cogvideoxfun_model,)
+            lora_path = folder_paths.get_full_path("loras", lora_name)
+            if lora_path is None:
+                raise FileNotFoundError(f"LoRA 文件未找到: {lora_name}")
+
+            new_funmodels['lora_cache'] = lora_cache
+            new_funmodels['loras'] = cogvideoxfun_model.get("loras", []) + [lora_path]
+            new_funmodels['strength_model'] = cogvideoxfun_model.get("strength_model", []) + [strength_model]
+
+        return (new_funmodels,)
 
 class CogVideoXFunT2VSampler:
     @classmethod

comfyui/comfyui_nodes.py

@@ -1,4 +1,5 @@
 import json
+import os
 
 import cv2
 import numpy as np
@@ -83,20 +84,28 @@ def INPUT_TYPES(s):
     def compile(self, cache_size_limit, funmodels):
         torch._dynamo.config.cache_size_limit = cache_size_limit
         if hasattr(funmodels["pipeline"].transformer, "blocks"):
-            for i in range(len(funmodels["pipeline"].transformer.blocks)):
-                funmodels["pipeline"].transformer.blocks[i] = torch.compile(funmodels["pipeline"].transformer.blocks[i])
+            for i, block in enumerate(funmodels["pipeline"].transformer.blocks):
+                if hasattr(block, "_orig_mod"):
+                    block = block._orig_mod
+                funmodels["pipeline"].transformer.blocks[i] = torch.compile(block)
 
             if hasattr(funmodels["pipeline"], "transformer_2") and funmodels["pipeline"].transformer_2 is not None:
-                for i in range(len(funmodels["pipeline"].transformer_2.blocks)):
-                    funmodels["pipeline"].transformer_2.blocks[i] = torch.compile(funmodels["pipeline"].transformer_2.blocks[i])
+                for i, block in enumerate(funmodels["pipeline"].transformer_2.blocks):
+                    if hasattr(block, "_orig_mod"):
+                        block = block._orig_mod
+                    funmodels["pipeline"].transformer_2.blocks[i] = torch.compile(block)
 
         elif hasattr(funmodels["pipeline"].transformer, "transformer_blocks"):
-            for i in range(len(funmodels["pipeline"].transformer.transformer_blocks)):
-                funmodels["pipeline"].transformer.transformer_blocks[i] = torch.compile(funmodels["pipeline"].transformer.transformer_blocks[i])
-
+            for i, block in enumerate(funmodels["pipeline"].transformer.transformer_blocks):
+                if hasattr(block, "_orig_mod"):
+                    block = block._orig_mod
+                funmodels["pipeline"].transformer.transformer_blocks[i] = torch.compile(block)
+
             if hasattr(funmodels["pipeline"], "transformer_2") and funmodels["pipeline"].transformer_2 is not None:
-                for i in range(len(funmodels["pipeline"].transformer_2.transformer_blocks)):
-                    funmodels["pipeline"].transformer_2.transformer_blocks[i] = torch.compile(funmodels["pipeline"].transformer_2.transformer_blocks[i])
+                for i, block in enumerate(funmodels["pipeline"].transformer_2.transformer_blocks):
+                    if hasattr(block, "_orig_mod"):
+                        block = block._orig_mod
+                    funmodels["pipeline"].transformer_2.transformer_blocks[i] = torch.compile(block)
 
         else:
             funmodels["pipeline"].transformer.forward = torch.compile(funmodels["pipeline"].transformer.forward)
@@ -106,6 +115,31 @@ def compile(self, cache_size_limit, funmodels):
 
         print("Add Compile")
         return (funmodels,)
+
+class FunAttention:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "attention_type": (
+                    ["flash", "sage", "torch"],
+                    {"default": "flash"},
+                ),
+                "funmodels": ("FunModels",)
+            }
+        }
+    RETURN_TYPES = ("FunModels",)
+    RETURN_NAMES = ("funmodels",)
+    FUNCTION = "funattention"
+    CATEGORY = "CogVideoXFUNWrapper"
+
+    def funattention(self, attention_type, funmodels):
+        os.environ['VIDEOX_ATTENTION_TYPE'] = {
+            "flash": "FLASH_ATTENTION",
+            "sage": "SAGE_ATTENTION",
+            "torch": "TORCH_SCALED_DOT"
+        }[attention_type]
+        return (funmodels,)
 
 class LoadConfig:
     @classmethod
@@ -376,6 +410,7 @@ def run(self,camera_pose,fx,fy,cx,cy):
     "FunTextBox": FunTextBox,
     "FunRiflex": FunRiflex,
     "FunCompile": FunCompile,
+    "FunAttention": FunAttention,
 
     "LoadCogVideoXFunModel": LoadCogVideoXFunModel,
     "LoadCogVideoXFunLora": LoadCogVideoXFunLora,
@@ -436,6 +471,7 @@ def run(self,camera_pose,fx,fy,cx,cy):
     "FunTextBox": "FunTextBox",
     "FunRiflex": "FunRiflex",
     "FunCompile": "FunCompile",
+    "FunAttention": "FunAttention",
 
     "LoadWanClipEncoderModel": "Load Wan ClipEncoder Model",
     "LoadWanTextEncoderModel": "Load Wan TextEncoder Model",

comfyui/wan2_1/nodes.py

@@ -612,13 +612,17 @@ def INPUT_TYPES(s):
     CATEGORY = "CogVideoXFUNWrapper"
 
     def load_lora(self, funmodels, lora_name, strength_model, lora_cache):
+        new_funmodels = dict(funmodels)  
+
         if lora_name is not None:
-            funmodels['lora_cache'] = lora_cache
-            funmodels['loras'] = funmodels.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)]
-            funmodels['strength_model'] = funmodels.get("strength_model", []) + [strength_model]
-            return (funmodels,)
-        else:
-            return (funmodels,)
+            lora_path = folder_paths.get_full_path("loras", lora_name)
+
+            new_funmodels['lora_cache'] = lora_cache
+            new_funmodels['loras'] = funmodels.get("loras", []) + [lora_path]
+            new_funmodels['strength_model'] = funmodels.get("strength_model", []) + [strength_model]
+
+        return (new_funmodels,)
+
 
 class WanT2VSampler:
     @classmethod

comfyui/wan2_1_fun/nodes.py

@@ -238,13 +238,16 @@ def INPUT_TYPES(s):
     CATEGORY = "CogVideoXFUNWrapper"
 
     def load_lora(self, funmodels, lora_name, strength_model, lora_cache):
+        new_funmodels = dict(funmodels)  
+
         if lora_name is not None:
-            funmodels['lora_cache'] = lora_cache
-            funmodels['loras'] = funmodels.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)]
-            funmodels['strength_model'] = funmodels.get("strength_model", []) + [strength_model]
-            return (funmodels,)
-        else:
-            return (funmodels,)
+            lora_path = folder_paths.get_full_path("loras", lora_name)
+
+            new_funmodels['lora_cache'] = lora_cache
+            new_funmodels['loras'] = funmodels.get("loras", []) + [lora_path]
+            new_funmodels['strength_model'] = funmodels.get("strength_model", []) + [strength_model]
+
+        return (new_funmodels,)
 
 class WanFunT2VSampler:
     @classmethod

comfyui/wan2_2/nodes.py

@@ -463,14 +463,16 @@ def INPUT_TYPES(s):
     CATEGORY = "CogVideoXFUNWrapper"
 
     def load_lora(self, funmodels, lora_name, lora_high_name, strength_model, lora_cache):
+        new_funmodels = dict(funmodels)
         if lora_name is not None:
-            funmodels['lora_cache'] = lora_cache
-            funmodels['loras'] = funmodels.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)]
-            funmodels['loras_high'] = funmodels.get("loras_high", []) + [folder_paths.get_full_path("loras", lora_high_name)]
-            funmodels['strength_model'] = funmodels.get("strength_model", []) + [strength_model]
-            return (funmodels,)
-        else:
-            return (funmodels,)
+            loras = list(new_funmodels.get("loras", [])) + [folder_paths.get_full_path("loras", lora_name)]
+            loras_high = list(new_funmodels.get("loras_high", [])) + [folder_paths.get_full_path("loras", lora_high_name)]
+            strength_models = list(new_funmodels.get("strength_model", [])) + [strength_model]
+            new_funmodels['loras'] = loras
+            new_funmodels['loras_high'] = loras_high
+            new_funmodels['strength_model'] = strength_models
+            new_funmodels['lora_cache'] = lora_cache
+        return (new_funmodels,)
 
 class Wan2_2T2VSampler:
     @classmethod

comfyui/wan2_2_fun/nodes.py

@@ -258,14 +258,16 @@ def INPUT_TYPES(s):
     CATEGORY = "CogVideoXFUNWrapper"
 
     def load_lora(self, funmodels, lora_name, lora_high_name, strength_model, lora_cache):
+        new_funmodels = dict(funmodels)
         if lora_name is not None:
-            funmodels['lora_cache'] = lora_cache
-            funmodels['loras'] = funmodels.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)]
-            funmodels['loras_high'] = funmodels.get("loras_high", []) + [folder_paths.get_full_path("loras", lora_high_name)]
-            funmodels['strength_model'] = funmodels.get("strength_model", []) + [strength_model]
-            return (funmodels,)
-        else:
-            return (funmodels,)
+            loras = list(new_funmodels.get("loras", [])) + [folder_paths.get_full_path("loras", lora_name)]
+            loras_high = list(new_funmodels.get("loras_high", [])) + [folder_paths.get_full_path("loras", lora_high_name)]
+            strength_models = list(new_funmodels.get("strength_model", [])) + [strength_model]
+            new_funmodels['loras'] = loras
+            new_funmodels['loras_high'] = loras_high
+            new_funmodels['strength_model'] = strength_models
+            new_funmodels['lora_cache'] = lora_cache
+        return (new_funmodels,)
 
 class Wan2_2FunT2VSampler:
     @classmethod

examples/cogvideox_fun/predict_i2v.py

@@ -202,7 +202,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 if partial_video_length is not None:
     partial_video_length = int((partial_video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -292,7 +292,7 @@
         ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/cogvideox_fun/predict_t2v.py

@@ -194,7 +194,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -232,7 +232,7 @@
         ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/cogvideox_fun/predict_v2v.py

@@ -201,7 +201,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
 latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
@@ -227,7 +227,7 @@
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/cogvideox_fun/predict_v2v_control.py

@@ -188,7 +188,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
 latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
@@ -212,7 +212,7 @@
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/flux/predict_t2i.py

@@ -184,7 +184,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     sample = pipeline(
@@ -198,7 +198,7 @@
     ).images
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/phantom/predict_s2v.py

@@ -244,7 +244,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -272,7 +272,7 @@
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/qwenimage/predict_t2i.py

@@ -176,7 +176,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     sample = pipeline(
@@ -190,7 +190,7 @@
     ).images
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/wan2.1/predict_i2v.py

@@ -245,7 +245,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -273,7 +273,7 @@
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/wan2.1/predict_t2v.py

@@ -232,7 +232,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -254,7 +254,7 @@
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):

examples/wan2.1_fun/predict_i2v.py

@@ -246,7 +246,7 @@
 generator = torch.Generator(device=device).manual_seed(seed)
 
 if lora_path is not None:
-    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 with torch.no_grad():
     video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
@@ -274,7 +274,7 @@
     ).videos
 
 if lora_path is not None:
-    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+    pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
 
 def save_results():
     if not os.path.exists(save_path):