LoRA support for Modulation layer in WAN2.2

Author: Perseus14

src/maxdiffusion/loaders/wan_lora_nnx_loader.py

@@ -19,7 +19,7 @@
 from .lora_pipeline import StableDiffusionLoraLoaderMixin
 from ..models import lora_nnx
 from .. import max_logging
-from . import lora_conversion_utils, preprocess_wan_lora_dict
+from . import lora_conversion_utils
 
 
 class Wan2_1NNXLoraLoader(LoRABaseMixin):
@@ -53,7 +53,7 @@ def translate_fn(nnx_path_str):
     if hasattr(pipeline, "transformer") and transformer_weight_name:
       max_logging.log(f"Merging LoRA into transformer with rank={rank}")
       h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=transformer_weight_name, **kwargs)
-      h_state_dict = preprocess_wan_lora_dict(h_state_dict)
+      h_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(h_state_dict)
       merge_fn(pipeline.transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype)
     else:
       max_logging.log("transformer not found or no weight name provided for LoRA.")
@@ -94,7 +94,7 @@ def translate_fn(nnx_path_str: str):
     if hasattr(pipeline, "high_noise_transformer") and high_noise_weight_name:
       max_logging.log(f"Merging LoRA into high_noise_transformer with rank={rank}")
       h_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=high_noise_weight_name, **kwargs)
-      h_state_dict = preprocess_wan_lora_dict(h_state_dict)
+      h_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(h_state_dict)
       merge_fn(pipeline.high_noise_transformer, h_state_dict, rank, scale, translate_fn, dtype=dtype)
     else:
       max_logging.log("high_noise_transformer not found or no weight name provided for LoRA.")
@@ -103,7 +103,7 @@ def translate_fn(nnx_path_str: str):
     if hasattr(pipeline, "low_noise_transformer") and low_noise_weight_name:
       max_logging.log(f"Merging LoRA into low_noise_transformer with rank={rank}")
       l_state_dict, _ = lora_loader.lora_state_dict(lora_model_path, weight_name=low_noise_weight_name, **kwargs)
-      l_state_dict = preprocess_wan_lora_dict(l_state_dict)
+      l_state_dict = lora_conversion_utils.preprocess_wan_lora_dict(l_state_dict)
       merge_fn(pipeline.low_noise_transformer, l_state_dict, rank, scale, translate_fn, dtype=dtype)
     else:
       max_logging.log("low_noise_transformer not found or no weight name provided for LoRA.")