Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -16,6 +16,9 @@
 
 from typing import Union, Tuple, Optional
 import re
+import torch
+import torch.utils.dlpack
+from jax import dlpack
 import jax.numpy as jnp
 from flax import nnx
 from .. import max_logging
@@ -200,6 +203,11 @@ def inject_lora(
 
     return model
 
+def _to_jax_array(v):
+    if isinstance(v, torch.Tensor):
+        return dlpack.from_dlpack(torch.utils.dlpack.to_dlpack(v))
+    return jnp.array(v)
+
 def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=None):
     """
     Merges weights from a Diffusers-formatted state dict directly
@@ -213,7 +221,7 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=N
             key_base = k[:-len(".alpha")]
             if key_base not in lora_params:
                 lora_params[key_base] = {}
-            lora_params[key_base]["alpha"] = jnp.array(v)
+            lora_params[key_base]["alpha"] = _to_jax_array(v)
             continue
 
         m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
@@ -232,7 +240,7 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=N
                     continue
         if key_base not in lora_params:
             lora_params[key_base] = {}
-        lora_params[key_base][weight_type] = jnp.array(v)
+        lora_params[key_base][weight_type] = _to_jax_array(v)
     max_logging.log(f"Parsed {len(lora_params)} unique LoRA module keys.")
 
     assigned_count = 0
@@ -289,7 +297,7 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
             key_base = k[:-len(".alpha")]
             if key_base not in lora_params:
                 lora_params[key_base] = {}
-            lora_params[key_base]["alpha"] = jnp.array(v)
+            lora_params[key_base]["alpha"] = _to_jax_array(v)
             continue
 
         m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
@@ -308,7 +316,7 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
                     continue
         if key_base not in lora_params:
             lora_params[key_base] = {}
-        lora_params[key_base][weight_type] = jnp.array(v)
+        lora_params[key_base][weight_type] = _to_jax_array(v)
     max_logging.log(f"Parsed {len(lora_params)} unique LoRA module keys for scanned merge.")
 
     assigned_count = 0