Refactored based on review

Author: Perseus14

src/maxdiffusion/generate_wan.py

@@ -28,7 +28,7 @@
 from google.cloud import storage
 import flax
 from maxdiffusion.common_types import WAN2_1, WAN2_2
-from maxdiffusion.loaders.wan_lora_nnx_loader import Wan2_1NnxLoraLoader, Wan2_2NnxLoraLoader
+from maxdiffusion.loaders.wan_lora_nnx_loader import Wan2_1NNXLoraLoader, Wan2_2NNXLoraLoader
 
 
 def upload_video_to_gcs(output_dir: str, video_path: str):
@@ -195,37 +195,33 @@ def run(config, pipeline=None, filename_prefix=""):
   # If LoRA is specified, inject layers and load weights.
   if config.enable_lora and hasattr(config, "lora_config") and config.lora_config and config.lora_config["lora_model_name_or_path"]:
     if model_key == WAN2_1:
-      lora_loader = Wan2_1NnxLoraLoader()
+      lora_loader = Wan2_1NNXLoraLoader()
       lora_config = config.lora_config
-
-      if len(lora_config["lora_model_name_or_path"]) > 1:
-          max_logging.log("Found multiple LoRAs in config, but only loading the first one.")
-
-      pipeline = lora_loader.load_lora_weights(
-          pipeline,
-          lora_config["lora_model_name_or_path"][0],
-          transformer_weight_name=lora_config["weight_name"][0],
-          rank=lora_config["rank"][0],
-          scale=lora_config["scale"][0],
-          scan_layers=config.scan_layers,
-      )
+      for i in range(len(lora_config["lora_model_name_or_path"])):
+        pipeline = lora_loader.load_lora_weights(
+            pipeline,
+            lora_config["lora_model_name_or_path"][i],
+            transformer_weight_name=lora_config["weight_name"][i],
+            rank=lora_config["rank"][i],
+            scale=lora_config["scale"][i],
+            scan_layers=config.scan_layers,
+            dtype=config.weights_dtype,
+        )
 
     if model_key == WAN2_2:
-      lora_loader = Wan2_2NnxLoraLoader()
+      lora_loader = Wan2_2NNXLoraLoader()
       lora_config = config.lora_config
-
-      if len(lora_config["lora_model_name_or_path"]) > 1:
-          max_logging.log("Found multiple LoRAs in config, but only loading the first one.")
-
-      pipeline = lora_loader.load_lora_weights(
-          pipeline,
-          lora_config["lora_model_name_or_path"][0],
-          high_noise_weight_name=lora_config["high_noise_weight_name"][0],
-          low_noise_weight_name=lora_config["low_noise_weight_name"][0],
-          rank=lora_config["rank"][0],
-          scale=lora_config["scale"][0],
-          scan_layers=config.scan_layers,
-      )
+      for i in range(len(lora_config["lora_model_name_or_path"])):
+        pipeline = lora_loader.load_lora_weights(
+            pipeline,
+            lora_config["lora_model_name_or_path"][i],
+            high_noise_weight_name=lora_config["high_noise_weight_name"][i],
+            low_noise_weight_name=lora_config["low_noise_weight_name"][i],
+            rank=lora_config["rank"][i],
+            scale=lora_config["scale"][i],
+            scan_layers=config.scan_layers,
+            dtype=config.weights_dtype,
+        )
 
   s0 = time.perf_counter()
 

src/maxdiffusion/loaders/__init__.py

@@ -14,4 +14,4 @@
 
 from .lora_pipeline import StableDiffusionLoraLoaderMixin
 from .flux_lora_pipeline import FluxLoraLoaderMixin
-from .wan_lora_nnx_loader import Wan2_1NnxLoraLoader, Wan2_2NnxLoraLoader
+from .wan_lora_nnx_loader import Wan2_1NNXLoraLoader, Wan2_2NNXLoraLoader

src/maxdiffusion/loaders/wan_lora_nnx_loader.py

@@ -22,7 +22,7 @@
 from . import lora_conversion_utils
 
 
-class Wan2_1NnxLoraLoader(LoRABaseMixin):
+class Wan2_1NNXLoraLoader(LoRABaseMixin):
   """
   Handles loading LoRA weights into NNX-based WAN 2.1 model.
   Assumes WAN pipeline contains 'transformer'
@@ -37,6 +37,7 @@ def load_lora_weights(
       rank: int,
       scale: float = 1.0,
       scan_layers: bool = False,
+      dtype: str = "float32",
       **kwargs,
   ):
     """
@@ -53,14 +54,14 @@ 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)
-      merge_fn(pipeline.transformer, h_state_dict, rank, scale, translate_fn)
+      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.")
 
     return pipeline
 
 
-class Wan2_2NnxLoraLoader(LoRABaseMixin):
+class Wan2_2NNXLoraLoader(LoRABaseMixin):
   """
   Handles loading LoRA weights into NNX-based WAN 2.2 model.
   Assumes WAN pipeline contains 'high_noise_transformer' and 'low_noise_transformer'
@@ -76,6 +77,7 @@ def load_lora_weights(
       rank: int,
       scale: float = 1.0,
       scan_layers: bool = False,
+      dtype: str = "float32",
       **kwargs,
   ):
     """
@@ -92,15 +94,15 @@ 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)
-      merge_fn(pipeline.high_noise_transformer, h_state_dict, rank, scale, translate_fn)
+      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.")
 
     # Handle low noise model
     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)
-      merge_fn(pipeline.low_noise_transformer, l_state_dict, rank, scale, translate_fn)
+      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.")