Add support for AI toolkit lora format

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -121,19 +121,22 @@ def parse_lora_dict(state_dict, dtype):
       continue
 
     # Standard LoRA
-    m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
+    m = re.match(r"^(.*?)\.(lora_down|lora_up)\.weight$", k)
     if not m:
-      m = re.match(r"^(.*?)\.lora\.(down|up)\.weight$", k)
-    if not m:
-      m = re.match(r"^(.*?)\.(lora_down|lora_up)\.weight$", k)
+      m = re.match(r"^(.*?)\.(lora_A|lora_B)\.weight$", k)
 
     if m:
       key_base, weight_type = m.group(1), m.group(2).replace("lora_", "")
+      if weight_type == "A":
+        weight_type = "down"
+      elif weight_type == "B":
+        weight_type = "up"
       if key_base not in lora_params:
         lora_params[key_base] = {}
       lora_params[key_base][weight_type] = _to_jax_array(v, dtype=dtype)
     else:
       # Fallback for exact matches of diffs if regex failed above
+      max_logging.log(f"Key {k} did not match any LoRA pattern.")
       pass
 
   return lora_params