diff --git a/omni/docker/Dockerfile-extension b/omni/docker/Dockerfile-extension index 4b82d7bb..461af191 100644 --- a/omni/docker/Dockerfile-extension +++ b/omni/docker/Dockerfile-extension @@ -23,6 +23,9 @@ COPY ./patches/extension_comfyui_reactor.patch /tmp/ COPY ./patches/extension_comfyui_rmbg.patch /tmp/ COPY ./patches/extension_SenseNova-U1.patch /tmp COPY ./patches/extension_comfyui-advancedliveportrait.patch /tmp +COPY ./patches/extension_comfyui_unload_model.patch /tmp +COPY ./patches/extension_comfyui_mmaudio.patch /tmp +COPY ./patches/extension_comfyui_slk_joy_caption_two.patch /tmp COPY ./patches/comfyui-pulid-flux-ll-add_intel_xpu_support_and_bug_fix.patch /tmp COPY ./patches/comfyui_wanblockswap-add_intel_xpu_support.patch /tmp @@ -141,6 +144,36 @@ RUN --mount=type=cache,target=/root/.cache/pip \ pip install -r requirements.txt && \ rm -f /tmp/extension_comfyui-advancedliveportrait.patch +RUN --mount=type=cache,target=/root/.cache/pip \ + cd /llm/ComfyUI/custom_nodes && \ + rm -rf ComfyUI-Unload-Model && \ + git clone --depth 1 https://github.com/SeanScripts/ComfyUI-Unload-Model.git && \ + cd ComfyUI-Unload-Model && \ + git fetch --depth 1 origin ac5ffb4ed05546545ce7cf38e7b69b5152714eed && \ + git checkout ac5ffb4ed05546545ce7cf38e7b69b5152714eed && \ + git apply /tmp/extension_comfyui_unload_model.patch && \ + rm -f /tmp/extension_comfyui_unload_model.patch + +RUN --mount=type=cache,target=/root/.cache/pip \ + cd /llm/ComfyUI/custom_nodes && \ + rm -rf ComfyUI-MMAudio && \ + git clone --depth 1 https://github.com/kijai/ComfyUI-MMAudio.git && \ + cd ComfyUI-MMAudio && \ + git fetch --depth 1 origin 8eaeb72edc3aaf2059b57f2d96a1f6f689f19ae2 && \ + git checkout 8eaeb72edc3aaf2059b57f2d96a1f6f689f19ae2 && \ + git apply /tmp/extension_comfyui_mmaudio.patch && \ + rm -f /tmp/extension_comfyui_mmaudio.patch + +RUN --mount=type=cache,target=/root/.cache/pip \ + cd /llm/ComfyUI/custom_nodes && \ + rm -rf ComfyUI_SLK_joy_caption_two && \ + git clone --depth 1 https://github.com/EvilBT/ComfyUI_SLK_joy_caption_two.git && \ + cd ComfyUI_SLK_joy_caption_two && \ + git fetch --depth 1 origin 667751cab945bd8e9fb0be4d557d47e36821350a && \ + git checkout 667751cab945bd8e9fb0be4d557d47e36821350a && \ + git apply /tmp/extension_comfyui_slk_joy_caption_two.patch && \ + rm -f /tmp/extension_comfyui_slk_joy_caption_two.patch + RUN --mount=type=cache,target=/root/.cache/pip \ cd /llm/ComfyUI/custom_nodes && \ rm -rf ComfyUI_PuLID_Flux_ll && \ @@ -163,6 +196,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \ git apply /tmp/comfyui_wanblockswap-add_intel_xpu_support.patch && \ rm -f /tmp/comfyui_wanblockswap-add_intel_xpu_support.patch + RUN rm -f /usr/lib/python3*/EXTERNALLY-MANAGED WORKDIR /llm/entrypoints diff --git a/omni/extension_comfyui_slk_joy_caption_two.patch b/omni/extension_comfyui_slk_joy_caption_two.patch new file mode 100644 index 00000000..a11840fc --- /dev/null +++ b/omni/extension_comfyui_slk_joy_caption_two.patch @@ -0,0 +1,374 @@ +--- a/joy_caption_two_node.py ++++ b/joy_caption_two_node.py +@@ -27,6 +27,56 @@ + + BASE_MODEL_PATH = Path(folder_paths.models_dir, "Joy_caption_two") + ++def _autocast_device_type(): ++ if torch.cuda.is_available(): ++ return "cuda" ++ xpu_backend = getattr(torch, "xpu", None) ++ if xpu_backend is not None and xpu_backend.is_available(): ++ return "xpu" ++ return "cpu" ++ ++ ++def _autocast_enabled(): ++ return _autocast_device_type() != "cpu" ++ ++ ++def _empty_cache(): ++ if torch.cuda.is_available(): ++ torch.cuda.empty_cache() ++ return ++ xpu_backend = getattr(torch, "xpu", None) ++ if xpu_backend is not None and xpu_backend.is_available(): ++ xpu_backend.empty_cache() ++ ++ ++def _safe_clear_llm(joy_two_pipeline, low_vram): ++ if joy_two_pipeline is not None and getattr(joy_two_pipeline, "llm", None) is not None: ++ joy_two_pipeline.llm.clear_gpu(low_vram) ++ ++ ++def _disable_incompatible_torchao_for_peft(): ++ try: ++ import importlib_metadata ++ import packaging.version ++ import peft.import_utils as peft_import_utils ++ import peft.tuners.lora.torchao as peft_torchao ++ ++ try: ++ torchao_version = packaging.version.parse(importlib_metadata.version("torchao")) ++ except Exception: ++ return ++ ++ if torchao_version < packaging.version.parse("0.16.0"): ++ def _always_false(): ++ return False ++ ++ peft_import_utils.is_torchao_available = _always_false ++ peft_torchao.is_torchao_available = _always_false ++ print(f"Disabled PEFT torchao path due to incompatible torchao version: {torchao_version}") ++ except Exception as e: ++ print(f"Torchao compatibility probe skipped: {e}") ++ ++ + def tensor2pil(t_image: torch.Tensor) -> Image: + return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +@@ -58,13 +108,11 @@ + clip_model.eval() + clip_model.requires_grad_(False) + self.model = clip_model +- +- self.patcher = ModelPatcher(self.model, load_device=self.load_device, offload_device=self.offload_device) +- + def encode_image(self, pixel_values): +- #print(f"{id(self)}之前 in JoyClipVisionModel: {next(self.model.parameters()).device}") # 打印模型参数的设备 +- load_models_gpu([self.patcher], force_full_load=True, force_patch_weights=True) +- #print(f"之后 in JoyClipVisionModel: {next(self.model.parameters()).device}") # 打印模型参数的设备 ++ # Avoid ModelPatcher for SigLIP vision model: it may try to set a read-only `device` property. ++ model_device = next(self.model.parameters()).device ++ if model_device != self.load_device: ++ self.model = self.model.to(self.load_device) + vision_outputs = self.model(pixel_values=pixel_values, output_hidden_states=True) + return vision_outputs + +@@ -139,8 +187,21 @@ + self.patcher = ModelPatcher(self.image_adapter, load_device=self.load_device, + offload_device=self.offload_device) + +- def embedded_image(self, hidden_states): ++ def embedded_image(self, vision_outputs): + load_models_gpu([self.patcher], force_full_load=True, force_patch_weights=True) ++ hidden_states = getattr(vision_outputs, "hidden_states", None) ++ if hidden_states is None: ++ last_hidden_state = getattr(vision_outputs, "last_hidden_state", None) ++ if last_hidden_state is not None: ++ # Keep at least two items so ImageAdapter can safely index [-2]. ++ hidden_states = (last_hidden_state, last_hidden_state) ++ ++ if hidden_states is None: ++ hidden_states = vision_outputs ++ ++ if isinstance(hidden_states, (tuple, list)) and len(hidden_states) == 1: ++ hidden_states = (hidden_states[0], hidden_states[0]) ++ + embedded_images = self.image_adapter(hidden_states) + embedded_images.to(self.load_device) + return embedded_images +@@ -161,40 +222,62 @@ + self.tokenizer = tokenizer + self.text_model = None + ++ def _load_text_model_once(self, text_model_path, llm_path): ++ try: ++ text_model = AutoModelForCausalLM.from_pretrained( ++ text_model_path, ++ local_files_only=True, ++ trust_remote_code=True, ++ torch_dtype=self.type, ++ ) ++ except Exception as e: ++ err_msg = str(e) ++ if "`weight` is not an nn.Module" in err_msg or "torchao" in err_msg: ++ print(f"Primary adapter load failed ({e}); falling back to base+PEFT path") ++ _disable_incompatible_torchao_for_peft() ++ from peft import PeftModel ++ ++ base_model = AutoModelForCausalLM.from_pretrained( ++ llm_path, ++ local_files_only=True, ++ trust_remote_code=True, ++ torch_dtype=self.type, ++ ) ++ text_model = PeftModel.from_pretrained( ++ base_model, ++ text_model_path, ++ local_files_only=True, ++ is_trainable=False, ++ ) ++ else: ++ raise ++ ++ text_model = text_model.to(self.load_device) ++ text_model.eval() ++ return text_model ++ + def load_llm_model(self): + if self.text_model is None or self.current_model_id != self.model_id: + print(f"Loading LLM: {self.model_id}") + LLM_PATH = download_hg_model(self.model_id, "LLM") + text_model_path = os.path.join(BASE_MODEL_PATH, "text_model") + modify_json_value(os.path.join(text_model_path, "adapter_config.json"), "base_model_name_or_path", +- LLM_PATH) +- max_retries = 5 # 设置最大重试次数 ++ LLM_PATH) ++ _disable_incompatible_torchao_for_peft() ++ max_retries = 5 # retry loading when free memory is too low + retries = 0 + while True: + free_vram = get_free_memory()/1024/1024 +- # print(f"现在的显存{retries}:{free_vram}") + if free_vram > 6400: +- text_model = AutoModelForCausalLM.from_pretrained(text_model_path, +- device_map=self.load_device, +- local_files_only=True, +- trust_remote_code=True, torch_dtype=self.type) +- text_model.eval() +- self.text_model = text_model ++ self.text_model = self._load_text_model_once(text_model_path, LLM_PATH) ++ break ++ clear_cache() ++ retries += 1 ++ if retries > max_retries: ++ self.text_model = self._load_text_model_once(text_model_path, LLM_PATH) + break +- else: +- clear_cache() +- retries += 1 +- if retries > max_retries: +- text_model = AutoModelForCausalLM.from_pretrained(text_model_path, +- device_map=self.load_device, +- local_files_only=True, +- trust_remote_code=True, +- torch_dtype=self.type) +- text_model.eval() +- self.text_model = text_model +- break +- time.sleep(1 + retries / 2) +- # print(f"现在呢:{get_free_memory()/1024/1024}") ++ time.sleep(1 + retries / 2) ++ + self.current_model_id = self.model_id + return self.text_model + +@@ -202,7 +285,7 @@ + del self.text_model + self.text_model = None + self.current_model_id = None +- torch.cuda.empty_cache() ++ _empty_cache() + import gc + gc.collect() + if low_vram: +@@ -294,7 +377,7 @@ + FUNCTION = "generate" + + def generate(self, joy_two_pipeline: JoyTwoPipeline, image, caption_type, caption_length, low_vram): +- torch.cuda.empty_cache() ++ _empty_cache() + + if joy_two_pipeline.clip_model is None: + joy_two_pipeline.parent.loadModels() +@@ -337,9 +420,9 @@ + + # Embed image + # This results in Batch x Image Tokens x Features +- with torch.amp.autocast_mode.autocast('cuda', enabled=True): ++ with torch.amp.autocast_mode.autocast(_autocast_device_type(), enabled=_autocast_enabled()): + vision_outputs = joy_two_pipeline.clip_model.encode_image(pixel_values) +- embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs.hidden_states) ++ embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs) + + if low_vram: + pixel_values.to(joy_two_pipeline.offload_device) +@@ -382,7 +465,7 @@ + + text_model = joy_two_pipeline.llm.load_llm_model() + # Embed the tokens +- convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) ++ convo_embeds = text_model.get_input_embeddings()( convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) + + # Construct the input + input_embeds = torch.cat([ +@@ -416,7 +499,7 @@ + + caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0] + +- joy_two_pipeline.llm.clear_gpu(low_vram) ++ _safe_clear_llm(joy_two_pipeline, low_vram) + + return (caption.strip(), ) + +@@ -449,7 +532,7 @@ + FUNCTION = "generate" + + def generate(self, joy_two_pipeline: JoyTwoPipeline, image, extra_options, caption_type, caption_length, name, custom_prompt, low_vram, top_p, temperature): +- torch.cuda.empty_cache() ++ _empty_cache() + + if joy_two_pipeline.clip_model == None: + joy_two_pipeline.parent.loadModels() +@@ -500,9 +583,9 @@ + + # Embed image + # This results in Batch x Image Tokens x Features +- with torch.amp.autocast_mode.autocast('cuda', enabled=True): ++ with torch.amp.autocast_mode.autocast(_autocast_device_type(), enabled=_autocast_enabled()): + vision_outputs = joy_two_pipeline.clip_model.encode_image(pixel_values) +- embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs.hidden_states) ++ embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs) + + if low_vram: + pixel_values.to(joy_two_pipeline.offload_device) +@@ -545,7 +628,7 @@ + + text_model = joy_two_pipeline.llm.load_llm_model() + # Embed the tokens +- convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) ++ convo_embeds = text_model.get_input_embeddings()( convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) + # Construct the input + input_embeds = torch.cat([ + convo_embeds[:, :preamble_len], # Part before the prompt +@@ -579,7 +662,7 @@ + + caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0] + +- joy_two_pipeline.llm.clear_gpu(low_vram) ++ _safe_clear_llm(joy_two_pipeline, low_vram) + + return (caption.strip(), ) + +@@ -608,16 +691,16 @@ + FUNCTION = "generate" + + def generate_caption(self, joy_two_pipeline: JoyTwoPipeline, image, prompt, low_vram=True): +- torch.cuda.empty_cache() ++ _empty_cache() + pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0 + pixel_values = TVF.normalize(pixel_values, [0.5], [0.5]) + pixel_values = pixel_values.to(joy_two_pipeline.load_device) + + # Embed image + # This results in Batch x Image Tokens x Features +- with torch.amp.autocast_mode.autocast('cuda', enabled=True): ++ with torch.amp.autocast_mode.autocast(_autocast_device_type(), enabled=_autocast_enabled()): + vision_outputs = joy_two_pipeline.clip_model.encode_image(pixel_values) +- embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs.hidden_states) ++ embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs) + + if low_vram: + pixel_values.to(joy_two_pipeline.offload_device) +@@ -659,7 +742,7 @@ + + text_model = joy_two_pipeline.llm.load_llm_model() + # Embed the tokens +- convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) ++ convo_embeds = text_model.get_input_embeddings()( convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) + # Construct the input + input_embeds = torch.cat([ + convo_embeds[:, :preamble_len], # Part before the prompt +@@ -695,7 +778,7 @@ + return caption.strip() + + def generate(self, joy_two_pipeline: JoyTwoPipeline, input_dir, output_dir, caption_type, caption_length, low_vram): +- torch.cuda.empty_cache() ++ _empty_cache() + + if joy_two_pipeline.clip_model == None: + joy_two_pipeline.parent.loadModels() +@@ -767,7 +850,7 @@ + # NOTE: I found the default processor for so400M to have worse results than just using PIL directly + # image = clip_processor(images=input_image, return_tensors='pt').pixel_values + +- joy_two_pipeline.llm.clear_gpu(low_vram) ++ _safe_clear_llm(joy_two_pipeline, low_vram) + + return (f"result: finished count: {finished_image_count}, error count: {error_image_count}", ) + +@@ -806,16 +889,16 @@ + FUNCTION = "generate" + + def generate_caption(self, joy_two_pipeline: JoyTwoPipeline, image, prompt, top_p=0.9, temperature=0.6, low_vram=True): +- torch.cuda.empty_cache() ++ _empty_cache() + pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0 + pixel_values = TVF.normalize(pixel_values, [0.5], [0.5]) + pixel_values = pixel_values.to(joy_two_pipeline.load_device) + + # Embed image + # This results in Batch x Image Tokens x Features +- with torch.amp.autocast_mode.autocast('cuda', enabled=True): ++ with torch.amp.autocast_mode.autocast(_autocast_device_type(), enabled=_autocast_enabled()): + vision_outputs = joy_two_pipeline.clip_model.encode_image(pixel_values) +- embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs.hidden_states) ++ embedded_images = joy_two_pipeline.image_adapter.embedded_image(vision_outputs) + + if low_vram: + pixel_values.to(joy_two_pipeline.offload_device) +@@ -857,7 +940,7 @@ + + text_model = joy_two_pipeline.llm.load_llm_model() + # Embed the tokens +- convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) ++ convo_embeds = text_model.get_input_embeddings()( convo_tokens.unsqueeze(0).to(joy_two_pipeline.load_device)) + # Construct the input + input_embeds = torch.cat([ + convo_embeds[:, :preamble_len], # Part before the prompt +@@ -896,7 +979,7 @@ + def generate(self, joy_two_pipeline: JoyTwoPipeline, input_dir, output_dir, rename, prefix_name, start_index, + extra_options, caption_type, caption_length, name, custom_prompt, low_vram, top_p, temperature, + prefix_caption, suffix_caption): +- torch.cuda.empty_cache() ++ _empty_cache() + + if joy_two_pipeline.clip_model == None: + joy_two_pipeline.parent.loadModels() +@@ -986,7 +1069,7 @@ + # NOTE: I found the default processor for so400M to have worse results than just using PIL directly + # image = clip_processor(images=input_image, return_tensors='pt').pixel_values + +- joy_two_pipeline.llm.clear_gpu(low_vram) ++ _safe_clear_llm(joy_two_pipeline, low_vram) + + return (f"result: finished count: {finished_image_count}, error count: {error_image_count}", ) + +@@ -1016,4 +1099,4 @@ + for selected, option in zip(options_selected, options): + if selected: + values.append(option) +- return (values, ) +\ No newline at end of file ++ return (values, ) diff --git a/omni/omni_xpu_kernel/omni_xpu_kernel/norm/__init__.py b/omni/omni_xpu_kernel/omni_xpu_kernel/norm/__init__.py index c6e0ff11..7b1cc1b4 100644 --- a/omni/omni_xpu_kernel/omni_xpu_kernel/norm/__init__.py +++ b/omni/omni_xpu_kernel/omni_xpu_kernel/norm/__init__.py @@ -76,7 +76,7 @@ def layer_norm( - hidden_size must be <= 8192 and divisible by 32 - Supports fp32, fp16, bf16 """ - return _get_native().layer_norm(input, weight, bias, eps) + return _get_native().layer_norm(input.contiguous(), weight, bias, eps) def fused_add_rms_norm( diff --git a/omni/patches/extension_comfyui_mmaudio.patch b/omni/patches/extension_comfyui_mmaudio.patch new file mode 100644 index 00000000..4fa7fe07 --- /dev/null +++ b/omni/patches/extension_comfyui_mmaudio.patch @@ -0,0 +1,80 @@ +diff --git a/mmaudio/ext/autoencoder/vae.py b/mmaudio/ext/autoencoder/vae.py +index efa3731..9fb69cb 100644 +--- a/mmaudio/ext/autoencoder/vae.py ++++ b/mmaudio/ext/autoencoder/vae.py +@@ -75,11 +75,11 @@ class VAE(nn.Module): + super().__init__() + + if data_dim == 80: +- self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32).cuda()) +- self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32).cuda()) ++ self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32)) ++ self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32)) + elif data_dim == 128: +- self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32).cuda()) +- self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32).cuda()) ++ self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32)) ++ self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32)) + + self.data_mean = self.data_mean.view(1, -1, 1) + self.data_std = self.data_std.view(1, -1, 1) +diff --git a/mmaudio/ext/rotary_embeddings.py b/mmaudio/ext/rotary_embeddings.py +index 1ea9d56..5a57aaa 100644 +--- a/mmaudio/ext/rotary_embeddings.py ++++ b/mmaudio/ext/rotary_embeddings.py +@@ -16,20 +16,18 @@ def compute_rope_rotations(length: int, + device: Union[torch.device, str] = 'cpu') -> Tensor: + assert dim % 2 == 0 + +- with torch.amp.autocast(device_type='cuda', enabled=False): +- pos = torch.arange(length, dtype=torch.float32, device=device) +- freqs = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) +- freqs *= freq_scaling ++ pos = torch.arange(length, dtype=torch.float32, device=device) ++ freqs = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) ++ freqs *= freq_scaling + +- rot = torch.einsum('..., f -> ... f', pos, freqs) +- rot = torch.stack([torch.cos(rot), -torch.sin(rot), torch.sin(rot), torch.cos(rot)], dim=-1) +- rot = rearrange(rot, 'n d (i j) -> 1 n d i j', i=2, j=2) +- return rot ++ rot = torch.einsum('..., f -> ... f', pos, freqs) ++ rot = torch.stack([torch.cos(rot), -torch.sin(rot), torch.sin(rot), torch.cos(rot)], dim=-1) ++ rot = rearrange(rot, 'n d (i j) -> 1 n d i j', i=2, j=2) ++ return rot + + + def apply_rope(x: Tensor, rot: Tensor) -> tuple[Tensor, Tensor]: +- with torch.amp.autocast(device_type='cuda', enabled=False): +- _x = x.float() +- _x = _x.view(*_x.shape[:-1], -1, 1, 2) +- x_out = rot[..., 0] * _x[..., 0] + rot[..., 1] * _x[..., 1] +- return x_out.reshape(*x.shape).to(dtype=x.dtype) ++ _x = x.float() ++ _x = _x.view(*_x.shape[:-1], -1, 1, 2) ++ x_out = rot[..., 0] * _x[..., 0] + rot[..., 1] * _x[..., 1] ++ return x_out.reshape(*x.shape).to(dtype=x.dtype) +diff --git a/mmaudio/model/embeddings.py b/mmaudio/model/embeddings.py +index 297feb4..cf5699f 100644 +--- a/mmaudio/model/embeddings.py ++++ b/mmaudio/model/embeddings.py +@@ -20,13 +20,12 @@ class TimestepEmbedder(nn.Module): + self.max_period = max_period + assert dim % 2 == 0, 'dim must be even.' + +- with torch.autocast('cuda', enabled=False): +- self.freqs = nn.Buffer( +- 1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) / +- frequency_embedding_size)), +- persistent=False) +- freq_scale = 10000 / max_period +- self.freqs = freq_scale * self.freqs ++ self.freqs = nn.Buffer( ++ 1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) / ++ frequency_embedding_size)), ++ persistent=False) ++ freq_scale = 10000 / max_period ++ self.freqs = freq_scale * self.freqs + + def timestep_embedding(self, t): + """ diff --git a/omni/patches/extension_comfyui_unload_model.patch b/omni/patches/extension_comfyui_unload_model.patch new file mode 100644 index 00000000..d81ee39d --- /dev/null +++ b/omni/patches/extension_comfyui_unload_model.patch @@ -0,0 +1,43 @@ +diff --git a/unloadModel.py b/unloadModel.py +index 021c2cb..38c92fb 100644 +--- a/unloadModel.py ++++ b/unloadModel.py +@@ -3,6 +3,18 @@ import gc + import torch + import time + ++def clear_torch_cache(): ++ """Clear backend cache for CUDA or XPU when available.""" ++ if torch.cuda.is_available(): ++ torch.cuda.empty_cache() ++ if hasattr(torch.cuda, "ipc_collect"): ++ torch.cuda.ipc_collect() ++ return ++ ++ if torch.xpu.is_available(): ++ torch.xpu.empty_cache() ++ ++ + # Note: This doesn't work with reroute for some reason? + class AnyType(str): + def __ne__(self, __value: object) -> bool: +@@ -47,8 +59,7 @@ class UnloadModelNode: + try: + print(" - Clearing Cache...") + gc.collect() +- torch.cuda.empty_cache() +- torch.cuda.ipc_collect() ++ clear_torch_cache() + except: + print(" - Unable to clear cache") + return (kwargs.get("value"),) +@@ -76,8 +87,7 @@ class UnloadAllModelsNode: + try: + print(" - Clearing Cache...") + gc.collect() +- torch.cuda.empty_cache() +- torch.cuda.ipc_collect() ++ clear_torch_cache() + except: + print(" - Unable to clear cache") + return (kwargs.get("value"),) diff --git a/omni/patches/extention_comfyui_wanvideowrapper.patch b/omni/patches/extention_comfyui_wanvideowrapper.patch index fcf867f7..a5208ad1 100644 --- a/omni/patches/extention_comfyui_wanvideowrapper.patch +++ b/omni/patches/extention_comfyui_wanvideowrapper.patch @@ -1,5 +1,5 @@ diff --git a/custom_linear.py b/custom_linear.py -index 9fe4762..0977968 100644 +index 9fe4762..0a85e3e 100644 --- a/custom_linear.py +++ b/custom_linear.py @@ -13,7 +13,7 @@ def apply_lora(weight: torch.Tensor, lora_diff_0: torch.Tensor, lora_diff_1: tor @@ -20,8 +20,134 @@ index 9fe4762..0977968 100644 @apply_single_lora.register_fake def _(weight, lora_diff, lora_strength): +@@ -31,6 +31,11 @@ def _(weight, lora_diff, lora_strength): + + @torch.library.custom_op("wanvideo::linear_forward", mutates_args=()) + def linear_forward(input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: ++ # Fallback alignment for mixed offload paths (cpu/xpu). ++ if weight.device != input.device: ++ weight = weight.to(device=input.device) ++ if bias is not None and (bias.device != input.device or bias.dtype != weight.dtype): ++ bias = bias.to(device=input.device, dtype=weight.dtype) + return torch.nn.functional.linear(input, weight, bias) + + @linear_forward.register_fake +@@ -221,6 +226,10 @@ class CustomLinear(nn.Linear): + + for idx, lora_diff_names in enumerate(self.lora_diffs): + lora_strength = self._get_lora_strength(idx) ++ if isinstance(lora_strength, torch.Tensor): ++ lora_strength = lora_strength.to(device=weight.device, dtype=weight.dtype) ++ else: ++ lora_strength = torch.tensor(lora_strength, device=weight.device, dtype=weight.dtype) + + if isinstance(lora_diff_names, tuple): + lora_diff_0 = getattr(self, lora_diff_names[0]) +@@ -239,7 +248,8 @@ class CustomLinear(nn.Linear): + def _prepare_weight(self, input): + """Prepare weight tensor - handles both regular and GGUF weights""" + if self.is_gguf: +- weight = dequantize_gguf_tensor(self.weight).to(self.compute_dtype) ++ target_dtype = self.compute_dtype if self.compute_dtype is not None else input.dtype ++ weight = dequantize_gguf_tensor(self.weight).to(device=input.device, dtype=target_dtype) + else: + weight = self.weight.to(input) + return weight +@@ -248,7 +258,7 @@ class CustomLinear(nn.Linear): + weight = self._prepare_weight(input) + + if self.bias is not None: +- bias = self.bias.to(input if not self.is_gguf else self.compute_dtype) ++ bias = self.bias.to(device=input.device, dtype=weight.dtype) + else: + bias = None + +diff --git a/nodes.py b/nodes.py +index 25d33cb..0b0faf6 100644 +--- a/nodes.py ++++ b/nodes.py +@@ -198,7 +198,7 @@ class WanVideoTextEncodeCached: + "negative_prompt": ("STRING", {"default": "", "multiline": True} ), + "quantization": (['disabled', 'fp8_e4m3fn'], {"default": 'disabled', "tooltip": "optional quantization method"}), + "use_disk_cache": ("BOOLEAN", {"default": True, "tooltip": "Cache the text embeddings to disk for faster re-use, under the custom_nodes/ComfyUI-WanVideoWrapper/text_embed_cache directory"}), +- "device": (["gpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}), ++ "device": (["gpu", "xpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}), + }, + "optional": { + "extender_args": ("WANVIDEOPROMPTEXTENDER_ARGS", {"tooltip": "Use this node to extend the prompt with additional text."}), +@@ -237,7 +237,7 @@ of the original Wan templates or a custom system prompt. + log.info("Using WanVideoPromptExtender to process prompts") + qwen, = QwenLoader().load( + extender_args["model"], +- load_device="main_device" if device == "gpu" else "cpu", ++ load_device="main_device" if device in ("gpu", "xpu") else "cpu", + precision=precision) + positive_prompt, = WanVideoPromptExtender().generate( + qwen=qwen, +@@ -294,7 +294,7 @@ class WanVideoTextEncode: + "force_offload": ("BOOLEAN", {"default": True}), + "model_to_offload": ("WANVIDEOMODEL", {"tooltip": "Model to move to offload_device before encoding"}), + "use_disk_cache": ("BOOLEAN", {"default": False, "tooltip": "Cache the text embeddings to disk for faster re-use, under the custom_nodes/ComfyUI-WanVideoWrapper/text_embed_cache directory"}), +- "device": (["gpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}), ++ "device": (["gpu", "xpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}), + } + } + +@@ -323,7 +323,7 @@ class WanVideoTextEncode: + if t5 is None: + raise ValueError("No cached text embeds found for prompts, please provide a T5 encoder.") + +- if model_to_offload is not None and device == "gpu": ++ if model_to_offload is not None and device in ("gpu", "xpu"): + try: + log.info(f"Moving video model to {offload_device}") + model_to_offload.model.to(offload_device) +@@ -358,6 +358,12 @@ class WanVideoTextEncode: + + if device == "gpu": + device_to = mm.get_torch_device() ++ elif device == "xpu": ++ if hasattr(torch, "xpu") and torch.xpu.is_available(): ++ device_to = torch.device("xpu") ++ else: ++ log.warning("XPU selected for text encoding but not available, falling back to CPU.") ++ device_to = torch.device("cpu") + else: + device_to = torch.device("cpu") + +@@ -464,7 +470,7 @@ class WanVideoTextEncodeSingle: + "force_offload": ("BOOLEAN", {"default": True}), + "model_to_offload": ("WANVIDEOMODEL", {"tooltip": "Model to move to offload_device before encoding"}), + "use_disk_cache": ("BOOLEAN", {"default": False, "tooltip": "Cache the text embeddings to disk for faster re-use, under the custom_nodes/ComfyUI-WanVideoWrapper/text_embed_cache directory"}), +- "device": (["gpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}), ++ "device": (["gpu", "xpu", "cpu"], {"default": "gpu", "tooltip": "Device to run the text encoding on."}), + } + } + +@@ -498,7 +504,7 @@ class WanVideoTextEncodeSingle: + + if encoded is None: + try: +- if model_to_offload is not None and device == "gpu": ++ if model_to_offload is not None and device in ("gpu", "xpu"): + log.info(f"Moving video model to {offload_device}") + model_to_offload.model.to(offload_device) + mm.soft_empty_cache() +@@ -510,6 +516,12 @@ class WanVideoTextEncodeSingle: + + if device == "gpu": + device_to = mm.get_torch_device() ++ elif device == "xpu": ++ if hasattr(torch, "xpu") and torch.xpu.is_available(): ++ device_to = torch.device("xpu") ++ else: ++ log.warning("XPU selected for text encoding but not available, falling back to CPU.") ++ device_to = torch.device("cpu") + else: + device_to = torch.device("cpu") + diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py -index ce101d8..15ba0b9 100644 +index ce101d8..c67b12f 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -302,7 +302,7 @@ class WanRMSNorm(nn.Module): @@ -33,7 +159,7 @@ index ce101d8..15ba0b9 100644 def _norm(self, x): return x * (torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)).to(x.dtype) -@@ -320,7 +320,7 @@ class WanRMSNorm(nn.Module): +@@ -320,33 +320,39 @@ class WanRMSNorm(nn.Module): chunk = x[:, start_idx:end_idx, :] norm_factor = torch.rsqrt(chunk.pow(2).mean(dim=-1, keepdim=True) + self.eps) @@ -42,27 +168,60 @@ index ce101d8..15ba0b9 100644 start_idx = end_idx -@@ -328,6 +328,8 @@ class WanRMSNorm(nn.Module): + return output class WanFusedRMSNorm(nn.RMSNorm): ++ def _forward_with_local_weight(self, x): ++ weight = self.weight ++ if weight is not None and weight.device != x.device: ++ weight = weight.to(device=x.device) ++ return F.rms_norm(x, self.normalized_shape, weight, self.eps) ++ def forward(self, x, num_chunks=1): -+ if self.weight is not None and self.weight.device != x.device: -+ self.weight.data = self.weight.data.to(x.device) use_chunked = num_chunks > 1 if use_chunked: return self.forward_chunked(x, num_chunks) -@@ -359,6 +361,10 @@ class WanLayerNorm(nn.LayerNorm): + else: +- return super().forward(x) ++ return self._forward_with_local_weight(x) + + def forward_chunked(self, x, num_chunks=4): + output = torch.empty_like(x) +- ++ + chunk_sizes = [x.shape[1] // num_chunks + (1 if i < x.shape[1] % num_chunks else 0) + for i in range(num_chunks)] +- ++ + start_idx = 0 + for size in chunk_sizes: + end_idx = start_idx + size + chunk = x[:, start_idx:end_idx, :] +- output[:, start_idx:end_idx, :] = super().forward(chunk) ++ output[:, start_idx:end_idx, :] = self._forward_with_local_weight(chunk) + start_idx = end_idx +- ++ + return output + + class WanLayerNorm(nn.LayerNorm): +@@ -359,7 +365,14 @@ class WanLayerNorm(nn.LayerNorm): Args: x(Tensor): Shape [B, L, C] """ -+ if self.weight is not None and self.weight.device != x.device: -+ self.weight.data = self.weight.data.to(x.device) -+ if self.bias is not None and self.bias.device != x.device: -+ self.bias.data = self.bias.data.to(x.device) - return super().forward(x) +- return super().forward(x) ++ weight = self.weight ++ bias = self.bias ++ if weight is not None and weight.device != x.device: ++ weight = weight.to(device=x.device) ++ if bias is not None and bias.device != x.device: ++ target_dtype = weight.dtype if weight is not None else bias.dtype ++ bias = bias.to(device=x.device, dtype=target_dtype) ++ return F.layer_norm(x, self.normalized_shape, weight, bias, self.eps) #region selfattn -@@ -968,9 +974,9 @@ class WanAttentionBlock(nn.Module): + class WanSelfAttention(nn.Module): +@@ -968,9 +981,9 @@ class WanAttentionBlock(nn.Module): if e.shape[-1] == 512: e = self.modulation(e) return e.unsqueeze(2).chunk(6, dim=-1) @@ -74,7 +233,7 @@ index ce101d8..15ba0b9 100644 return [ei.squeeze(1) for ei in e_mod.unbind(dim=1)] -@@ -1552,9 +1558,9 @@ class Head(nn.Module): +@@ -1552,9 +1565,9 @@ class Head(nn.Module): def get_mod(self, e): if e.dim() == 2: @@ -86,7 +245,7 @@ index ce101d8..15ba0b9 100644 return [ei.squeeze(1) for ei in e] def forward(self, x, e, e_tr=None, tr_start=0, tr_num=0, **kwargs): -@@ -3294,8 +3300,8 @@ class WanModel(torch.nn.Module): +@@ -3294,8 +3307,8 @@ class WanModel(torch.nn.Module): compute_end = time.perf_counter() compute_time = compute_end - compute_start to_cpu_transfer_start = time.perf_counter() @@ -97,3 +256,19 @@ index ce101d8..15ba0b9 100644 if self.block_swap_debug: to_cpu_transfer_end = time.perf_counter() to_cpu_transfer_time = to_cpu_transfer_end - to_cpu_transfer_start +diff --git a/wanvideo/modules/wananimate/face_blocks.py b/wanvideo/modules/wananimate/face_blocks.py +index e5399af..452ccab 100644 +--- a/wanvideo/modules/wananimate/face_blocks.py ++++ b/wanvideo/modules/wananimate/face_blocks.py +@@ -78,7 +78,10 @@ class RMSNorm(nn.Module): + def forward(self, x): + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): +- output = output * self.weight ++ weight = self.weight ++ if weight.device != output.device or weight.dtype != output.dtype: ++ weight = weight.to(device=output.device, dtype=output.dtype) ++ output = output * weight + return output + +