Fix

Author: Perseus14

src/maxdiffusion/configs/base_wan_lora_27b.yml

@@ -305,9 +305,9 @@ lightning_ckpt: ""
 lora_rank: 64
 # Values are lists to support multiple LoRA loading during inference in the future.
 lora_config: {
-  lora_model_name_or_path: ["lightx2v/Wan2.2-Lightning/Wan2.2-T2V-A14B-4steps-lora-250928"],
-  high_noise_weight_name: ["high_noise_model.safetensors"],
-  low_noise_weight_name: ["low_noise_model.safetensors"],
+  lora_model_name_or_path: ["lightx2v/Wan2.2-Lightning"],
+  high_noise_weight_name: ["Wan2.2-T2V-A14B-4steps-lora-250928/high_noise_model.safetensors"],
+  low_noise_weight_name: ["Wan2.2-T2V-A14B-4steps-lora-250928/low_noise_model.safetensors"],
   adapter_name: ["wan22-lightning-lora"],
   scale: [1.0],
   from_pt: []

src/maxdiffusion/generate_wan.py

@@ -166,7 +166,6 @@ def run(config, pipeline=None, filename_prefix=""):
         low_noise_weight_name=lora_config["low_noise_weight_name"][0],
         rank=config.lora_rank,
         scale=lora_config["scale"][0],
-        rng=jax.random.key(config.seed),
     )
 
   s0 = time.perf_counter()

src/maxdiffusion/loaders/wan_lora_nnx_loader.py

@@ -14,10 +14,8 @@
 
 """NNX-based LoRA loader for WAN models."""
 
-import re
 from flax import nnx
 import jax
-import jax.numpy as jnp
 from .lora_base import LoRABaseMixin
 from .lora_pipeline import StableDiffusionLoraLoaderMixin
 from ..models import lora_nnx
@@ -38,16 +36,11 @@ def load_lora_weights(
       low_noise_weight_name: str,
       rank: int,
       scale: float = 1.0,
-      rng: jax.Array = None,
       **kwargs,
   ):
     """
-    Injects LoRA layers into the pipeline and loads weights
-    from a checkpoint.
+    Merges LoRA weights into the pipeline from a checkpoint.
     """
-    if rng is None:
-      rng = jax.random.key(0)
-
     lora_loader = StableDiffusionLoraLoaderMixin()
 
     # Handle high noise model