Refactored based on review

Author: Perseus14

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -276,8 +276,8 @@ profiler_steps: 10
 enable_jax_named_scopes: False
 
 # Generation parameters
-prompt: "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-prompt_2: "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+prompt: "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. Appearing behind him is a giant, translucent, pink spiritual manifestation (faxiang) that is synchronized with the man's action and pose." #"An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+prompt_2: "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. Appearing behind him is a giant, translucent, pink spiritual manifestation (faxiang) that is synchronized with the man's action and pose." #"An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
 negative_prompt: "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
 do_classifier_free_guidance: True
 height: 480
@@ -303,11 +303,11 @@ lightning_ckpt: ""
 enable_lora: False
 # Values are lists to support multiple LoRA loading during inference in the future.
 lora_config: {
-  rank: [64],
-  lora_model_name_or_path: ["lightx2v/Wan2.1-Distill-Loras"],
-  weight_name: ["wan2.1_i2v_lora_rank64_lightx2v_4step.safetensors"],
-  adapter_name: ["wan21-distill-lora-i2v"],
-  scale: [1.0],
+  rank: [64, 32],
+  lora_model_name_or_path: ["lightx2v/Wan2.1-Distill-Loras", "starsfriday/Wan2.1-Divine-Power-LoRA"],
+  weight_name: ["wan2.1_i2v_lora_rank64_lightx2v_4step.safetensors", "divine-power.safetensors"],
+  adapter_name: ["wan21-distill-lora-i2v", "divine-power-lora"],
+  scale: [1.0, 1.0],
   from_pt: []
 }
 # Ex with values:

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -277,8 +277,8 @@ profiler_steps: 10
 enable_jax_named_scopes: False
 
 # Generation parameters
-prompt: "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-prompt_2: "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+prompt: "orbit 180 around an astronaut on the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+prompt_2: "orbit 180 around an astronaut on the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
 negative_prompt: "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
 do_classifier_free_guidance: True
 height: 480
@@ -288,10 +288,10 @@ flow_shift: 3.0
 
 # Reference for below guidance scale and boundary values: https://github.com/Wan-Video/Wan2.2/blob/main/wan/configs/wan_t2v_A14B.py
 # guidance scale factor for low noise transformer
-guidance_scale_low: 3.0    
+guidance_scale_low: 3.0
 
 # guidance scale factor for high noise transformer
-guidance_scale_high: 4.0   
+guidance_scale_high: 4.0
 
 # The timestep threshold. If `t` is at or above this value,
 # the `high_noise_model` is considered as the required model.
@@ -315,12 +315,12 @@ lightning_ckpt: ""
 enable_lora: False
 # Values are lists to support multiple LoRA loading during inference in the future.
 lora_config: {
-  rank: [64],
-  lora_model_name_or_path: ["lightx2v/Wan2.2-Distill-Loras"],
-  high_noise_weight_name: ["wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors"],
-  low_noise_weight_name: [""], # Empty or "wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors"
-  adapter_name: ["wan22-lightning-lora"],
-  scale: [1.0],
+  rank: [64, 16],
+  lora_model_name_or_path: ["lightx2v/Wan2.2-Distill-Loras", "ostris/wan22_i2v_14b_orbit_shot_lora"],
+  high_noise_weight_name: ["wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors", "wan22_14b_i2v_orbit_high_noise.safetensors"],
+  low_noise_weight_name: ["wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors", "wan22_14b_i2v_orbit_low_noise.safetensors"], # Empty or "wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors"
+  adapter_name: ["wan22-distill-lora", "wan22-orbit-lora"],
+  scale: [1.0, 1.0],
   from_pt: []
 }
 # Ex with values:

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):
@@ -191,39 +191,40 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
     pipeline, _, _ = checkpoint_loader.load_checkpoint()
 
   # 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 (
+      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/lora_conversion_utils.py

@@ -610,6 +610,20 @@ def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
   return new_state_dict
 
 
+def preprocess_wan_lora_dict(state_dict):
+  """
+  Preprocesses WAN LoRA dict to convert diff_m to modulation.diff.
+  """
+  new_d = {}
+  for k, v in state_dict.items():
+    if k.endswith(".diff_m"):
+      new_k = k.removesuffix(".diff_m") + ".modulation.diff"
+      new_d[new_k] = v
+    else:
+      new_d[k] = v
+  return new_d
+
+
 def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
   """
   Translates WAN NNX path to Diffusers/LoRA keys.
@@ -637,6 +651,8 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
     return "diffusion_model.patch_embedding"
   if nnx_path_str == "proj_out":
     return "diffusion_model.head.head"
+  if nnx_path_str == "scale_shift_table":
+    return "diffusion_model.head.modulation"
   if nnx_path_str == "condition_embedder.time_proj":
     return "diffusion_model.time_projection.1"
 
@@ -667,7 +683,7 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
       "ffn.proj_out": "ffn.2",  # Down proj
       # Global Norms & Modulation
       "norm2.layer_norm": "norm3",
-      "scale_shift_table": "modulation",
+      "adaln_scale_shift_table": "modulation",
       "proj_out": "head.head",
   }
 

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,
   ):
     """
@@ -49,18 +50,18 @@ def load_lora_weights(
     def translate_fn(nnx_path_str):
       return lora_conversion_utils.translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=scan_layers)
 
-    # Handle high noise model
     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)
+      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.")
 
     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,17 @@ 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)
+      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.")
 
     # 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)
+      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.")