Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -17,7 +17,6 @@
 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
@@ -154,58 +153,9 @@ def __call__(self, x):
 
     return base_out + lora_out
 
-
-# -----------------------------------------------------------------------------
-# Helper: The "Discovery" Logic (Graph Transformation)
-# -----------------------------------------------------------------------------
-
-def inject_lora(
-    model: nnx.Module, 
-    rank: int, 
-    rngs: nnx.Rngs, 
-    scale: float = 1.0,
-    target_linear: bool = True, 
-    target_conv: bool = False
-):
-    """
-    Traverses the NNX model graph and replaces target layers with LoRA wrappers.
-    This modifies the 'model' object in-place.
-    """
-    for path, module in nnx.iter_graph(model):
-        # If path is like ('block', 'linear'), attr_name is 'linear'
-        # and parent is model.block
-        attr_name = path[-1]
-        if len(path) == 1:
-          parent = model
-        else:
-          parent = model
-          for key in path[:-1]:
-            parent = getattr(parent, key)
-
-        # Handle Linear Layers
-        if target_linear and isinstance(module, nnx.Linear):
-            # Do not wrap if it's already wrapped (sanity check)
-            if isinstance(parent, BaseLoRALayer): 
-                continue
-
-            print(f"Injecting LoRA (Linear) at {'.'.join([str(p) for p in path])}")
-            wrapper = LoRALinear(base_layer=module, rank=rank, scale=scale, rngs=rngs)
-            setattr(parent, attr_name, wrapper)
-
-        # Handle Conv Layers
-        elif target_conv and isinstance(module, nnx.Conv):
-            if isinstance(parent, BaseLoRALayer):
-                continue
-
-            print(f"Injecting LoRA (Conv) at {'.'.join([str(p) for p in path])}")
-            wrapper = LoRAConv(base_layer=module, rank=rank, scale=scale, rngs=rngs)
-            setattr(parent, attr_name, wrapper)
-
-    return model
-
 def _to_jax_array(v):
     if isinstance(v, torch.Tensor):
-        return dlpack.from_dlpack(torch.utils.dlpack.to_dlpack(v))
+        return dlpack.from_dlpack(v)
     return jnp.array(v)
 
 def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=None):