Merge branch 'main' into feat/z-image-regional-guidance

Author: lstein

invokeai/backend/model_manager/configs/lora.py

@@ -150,11 +150,16 @@ def _validate_base(cls, mod: ModelOnDisk) -> None:
 
     @classmethod
     def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None:
-        # First rule out ControlLoRA and Diffusers LoRA
+        # First rule out ControlLoRA
         flux_format = _get_flux_lora_format(mod)
         if flux_format in [FluxLoRAFormat.Control]:
             raise NotAMatchError("model looks like Control LoRA")
 
+        # If it's a recognized Flux LoRA format (Kohya, Diffusers, OneTrainer, AIToolkit, XLabs, etc.),
+        # it's valid and we skip the heuristic check
+        if flux_format is not None:
+            return
+
         # Note: Existence of these key prefixes/suffixes does not guarantee that this is a LoRA.
         # Some main models have these keys, likely due to the creator merging in a LoRA.
         has_key_with_lora_prefix = state_dict_has_any_keys_starting_with(

invokeai/backend/model_manager/load/model_loaders/lora.py

@@ -41,6 +41,10 @@
     is_state_dict_likely_in_flux_onetrainer_format,
     lora_model_from_flux_onetrainer_state_dict,
 )
+from invokeai.backend.patches.lora_conversions.flux_xlabs_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_xlabs_format,
+    lora_model_from_flux_xlabs_state_dict,
+)
 from invokeai.backend.patches.lora_conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict
 from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format
 from invokeai.backend.patches.lora_conversions.z_image_lora_conversion_utils import lora_model_from_z_image_state_dict
@@ -118,6 +122,8 @@ def _load_model(
                     model = lora_model_from_flux_control_state_dict(state_dict=state_dict)
                 elif is_state_dict_likely_in_flux_aitoolkit_format(state_dict=state_dict):
                     model = lora_model_from_flux_aitoolkit_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_in_flux_xlabs_format(state_dict=state_dict):
+                    model = lora_model_from_flux_xlabs_state_dict(state_dict=state_dict)
                 else:
                     raise ValueError("LoRA model is in unsupported FLUX format")
             else:

invokeai/backend/model_manager/taxonomy.py

@@ -171,6 +171,7 @@ class FluxLoRAFormat(str, Enum):
     OneTrainer = "flux.onetrainer"
     Control = "flux.control"
     AIToolkit = "flux.aitoolkit"
+    XLabs = "flux.xlabs"
 
 
 AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, FluxVariantType]

invokeai/backend/patches/lora_conversions/flux_xlabs_lora_conversion_utils.py

@@ -0,0 +1,92 @@
+import re
+from typing import Any, Dict
+
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict
+from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+
+# A regex pattern that matches all of the transformer keys in the xlabs FLUX LoRA format.
+# Example keys:
+#   double_blocks.0.processor.qkv_lora1.down.weight
+#   double_blocks.0.processor.qkv_lora1.up.weight
+#   double_blocks.0.processor.proj_lora1.down.weight
+#   double_blocks.0.processor.proj_lora1.up.weight
+#   double_blocks.0.processor.qkv_lora2.down.weight
+#   double_blocks.0.processor.proj_lora2.up.weight
+FLUX_XLABS_KEY_REGEX = r"double_blocks\.(\d+)\.processor\.(qkv|proj)_lora([12])\.(down|up)\.weight"
+
+
+def is_state_dict_likely_in_flux_xlabs_format(state_dict: dict[str | int, Any]) -> bool:
+    """Checks if the provided state dict is likely in the xlabs FLUX LoRA format.
+
+    The xlabs format is characterized by keys matching the pattern:
+    double_blocks.{block_idx}.processor.{qkv|proj}_lora{1|2}.{down|up}.weight
+
+    Where:
+    - lora1 corresponds to the image attention stream (img_attn)
+    - lora2 corresponds to the text attention stream (txt_attn)
+    """
+    if not state_dict:
+        return False
+
+    # Check that all keys match the xlabs pattern
+    for key in state_dict.keys():
+        if not isinstance(key, str):
+            continue
+        if not re.match(FLUX_XLABS_KEY_REGEX, key):
+            return False
+
+    # Ensure we have at least some valid keys
+    return any(isinstance(k, str) and re.match(FLUX_XLABS_KEY_REGEX, k) for k in state_dict.keys())
+
+
+def lora_model_from_flux_xlabs_state_dict(state_dict: Dict[str, torch.Tensor]) -> ModelPatchRaw:
+    """Converts an xlabs FLUX LoRA state dict to the InvokeAI ModelPatchRaw format.
+
+    The xlabs format uses:
+    - lora1 for image attention stream (img_attn)
+    - lora2 for text attention stream (txt_attn)
+    - qkv for query/key/value projection
+    - proj for output projection
+
+    Key mapping:
+    - double_blocks.X.processor.qkv_lora1 -> double_blocks.X.img_attn.qkv
+    - double_blocks.X.processor.proj_lora1 -> double_blocks.X.img_attn.proj
+    - double_blocks.X.processor.qkv_lora2 -> double_blocks.X.txt_attn.qkv
+    - double_blocks.X.processor.proj_lora2 -> double_blocks.X.txt_attn.proj
+    """
+    # Group keys by layer (without the .down.weight/.up.weight suffix)
+    grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {}
+
+    for key, value in state_dict.items():
+        match = re.match(FLUX_XLABS_KEY_REGEX, key)
+        if not match:
+            raise ValueError(f"Key '{key}' does not match the expected pattern for xlabs FLUX LoRA weights.")
+
+        block_idx = match.group(1)
+        component = match.group(2)  # qkv or proj
+        lora_stream = match.group(3)  # 1 or 2
+        direction = match.group(4)  # down or up
+
+        # Map lora1 -> img_attn, lora2 -> txt_attn
+        attn_type = "img_attn" if lora_stream == "1" else "txt_attn"
+
+        # Create the InvokeAI-style layer key
+        layer_key = f"double_blocks.{block_idx}.{attn_type}.{component}"
+
+        if layer_key not in grouped_state_dict:
+            grouped_state_dict[layer_key] = {}
+
+        # Map down/up to lora_down/lora_up
+        param_name = f"lora_{direction}.weight"
+        grouped_state_dict[layer_key][param_name] = value
+
+    # Create LoRA layers
+    layers: dict[str, BaseLayerPatch] = {}
+    for layer_key, layer_state_dict in grouped_state_dict.items():
+        layers[FLUX_LORA_TRANSFORMER_PREFIX + layer_key] = any_lora_layer_from_state_dict(layer_state_dict)
+
+    return ModelPatchRaw(layers=layers)

invokeai/backend/patches/lora_conversions/formats.py

@@ -14,6 +14,9 @@
 from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_utils import (
     is_state_dict_likely_in_flux_onetrainer_format,
 )
+from invokeai.backend.patches.lora_conversions.flux_xlabs_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_xlabs_format,
+)
 
 
 def flux_format_from_state_dict(
@@ -30,5 +33,7 @@ def flux_format_from_state_dict(
         return FluxLoRAFormat.Control
     elif is_state_dict_likely_in_flux_aitoolkit_format(state_dict, metadata):
         return FluxLoRAFormat.AIToolkit
+    elif is_state_dict_likely_in_flux_xlabs_format(state_dict):
+        return FluxLoRAFormat.XLabs
     else:
         return None

invokeai/frontend/web/src/common/util/promptAST.test.ts

@@ -7,72 +7,76 @@ describe('promptAST', () => {
     it('should tokenize basic text', () => {
       const tokens = tokenize('a cat');
       expect(tokens).toEqual([
-        { type: 'word', value: 'a' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'word', value: 'cat' },
+        { type: 'word', value: 'a', start: 0, end: 1 },
+        { type: 'whitespace', value: ' ', start: 1, end: 2 },
+        { type: 'word', value: 'cat', start: 2, end: 5 },
       ]);
     });
 
     it('should tokenize groups with parentheses', () => {
       const tokens = tokenize('(a cat)');
       expect(tokens).toEqual([
-        { type: 'lparen' },
-        { type: 'word', value: 'a' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'word', value: 'cat' },
-        { type: 'rparen' },
+        { type: 'lparen', start: 0, end: 1 },
+        { type: 'word', value: 'a', start: 1, end: 2 },
+        { type: 'whitespace', value: ' ', start: 2, end: 3 },
+        { type: 'word', value: 'cat', start: 3, end: 6 },
+        { type: 'rparen', start: 6, end: 7 },
       ]);
     });
 
     it('should tokenize escaped parentheses', () => {
       const tokens = tokenize('\\(medium\\)');
       expect(tokens).toEqual([
-        { type: 'escaped_paren', value: '(' },
-        { type: 'word', value: 'medium' },
-        { type: 'escaped_paren', value: ')' },
+        { type: 'escaped_paren', value: '(', start: 0, end: 2 },
+        { type: 'word', value: 'medium', start: 2, end: 8 },
+        { type: 'escaped_paren', value: ')', start: 8, end: 10 },
       ]);
     });
 
     it('should tokenize mixed escaped and unescaped parentheses', () => {
       const tokens = tokenize('colored pencil \\(medium\\) (enhanced)');
       expect(tokens).toEqual([
-        { type: 'word', value: 'colored' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'word', value: 'pencil' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'escaped_paren', value: '(' },
-        { type: 'word', value: 'medium' },
-        { type: 'escaped_paren', value: ')' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'lparen' },
-        { type: 'word', value: 'enhanced' },
-        { type: 'rparen' },
+        { type: 'word', value: 'colored', start: 0, end: 7 },
+        { type: 'whitespace', value: ' ', start: 7, end: 8 },
+        { type: 'word', value: 'pencil', start: 8, end: 14 },
+        { type: 'whitespace', value: ' ', start: 14, end: 15 },
+        { type: 'escaped_paren', value: '(', start: 15, end: 17 },
+        { type: 'word', value: 'medium', start: 17, end: 23 },
+        { type: 'escaped_paren', value: ')', start: 23, end: 25 },
+        { type: 'whitespace', value: ' ', start: 25, end: 26 },
+        { type: 'lparen', start: 26, end: 27 },
+        { type: 'word', value: 'enhanced', start: 27, end: 35 },
+        { type: 'rparen', start: 35, end: 36 },
       ]);
     });
 
     it('should tokenize groups with weights', () => {
       const tokens = tokenize('(a cat)1.2');
       expect(tokens).toEqual([
-        { type: 'lparen' },
-        { type: 'word', value: 'a' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'word', value: 'cat' },
-        { type: 'rparen' },
-        { type: 'weight', value: 1.2 },
+        { type: 'lparen', start: 0, end: 1 },
+        { type: 'word', value: 'a', start: 1, end: 2 },
+        { type: 'whitespace', value: ' ', start: 2, end: 3 },
+        { type: 'word', value: 'cat', start: 3, end: 6 },
+        { type: 'rparen', start: 6, end: 7 },
+        { type: 'weight', value: 1.2, start: 7, end: 10 },
       ]);
     });
 
     it('should tokenize words with weights', () => {
       const tokens = tokenize('cat+');
       expect(tokens).toEqual([
-        { type: 'word', value: 'cat' },
-        { type: 'weight', value: '+' },
+        { type: 'word', value: 'cat', start: 0, end: 3 },
+        { type: 'weight', value: '+', start: 3, end: 4 },
       ]);
     });
 
     it('should tokenize embeddings', () => {
       const tokens = tokenize('<embedding_name>');
-      expect(tokens).toEqual([{ type: 'lembed' }, { type: 'word', value: 'embedding_name' }, { type: 'rembed' }]);
+      expect(tokens).toEqual([
+        { type: 'lembed', start: 0, end: 1 },
+        { type: 'word', value: 'embedding_name', start: 1, end: 15 },
+        { type: 'rembed', start: 15, end: 16 },
+      ]);
     });
   });
 
@@ -81,9 +85,9 @@ describe('promptAST', () => {
       const tokens = tokenize('a cat');
       const ast = parseTokens(tokens);
       expect(ast).toEqual([
-        { type: 'word', text: 'a' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'word', text: 'cat' },
+        { type: 'word', text: 'a', range: { start: 0, end: 1 }, attention: undefined },
+        { type: 'whitespace', value: ' ', range: { start: 1, end: 2 } },
+        { type: 'word', text: 'cat', range: { start: 2, end: 5 }, attention: undefined },
       ]);
     });
 
@@ -93,10 +97,12 @@ describe('promptAST', () => {
       expect(ast).toEqual([
         {
           type: 'group',
+          range: { start: 0, end: 7 },
+          attention: undefined,
           children: [
-            { type: 'word', text: 'a' },
-            { type: 'whitespace', value: ' ' },
-            { type: 'word', text: 'cat' },
+            { type: 'word', text: 'a', range: { start: 1, end: 2 }, attention: undefined },
+            { type: 'whitespace', value: ' ', range: { start: 2, end: 3 } },
+            { type: 'word', text: 'cat', range: { start: 3, end: 6 }, attention: undefined },
           ],
         },
       ]);
@@ -106,27 +112,29 @@ describe('promptAST', () => {
       const tokens = tokenize('\\(medium\\)');
       const ast = parseTokens(tokens);
       expect(ast).toEqual([
-        { type: 'escaped_paren', value: '(' },
-        { type: 'word', text: 'medium' },
-        { type: 'escaped_paren', value: ')' },
+        { type: 'escaped_paren', value: '(', range: { start: 0, end: 2 } },
+        { type: 'word', text: 'medium', range: { start: 2, end: 8 }, attention: undefined },
+        { type: 'escaped_paren', value: ')', range: { start: 8, end: 10 } },
       ]);
     });
 
     it('should parse mixed escaped and unescaped parentheses', () => {
       const tokens = tokenize('colored pencil \\(medium\\) (enhanced)');
       const ast = parseTokens(tokens);
       expect(ast).toEqual([
-        { type: 'word', text: 'colored' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'word', text: 'pencil' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'escaped_paren', value: '(' },
-        { type: 'word', text: 'medium' },
-        { type: 'escaped_paren', value: ')' },
-        { type: 'whitespace', value: ' ' },
+        { type: 'word', text: 'colored', range: { start: 0, end: 7 }, attention: undefined },
+        { type: 'whitespace', value: ' ', range: { start: 7, end: 8 } },
+        { type: 'word', text: 'pencil', range: { start: 8, end: 14 }, attention: undefined },
+        { type: 'whitespace', value: ' ', range: { start: 14, end: 15 } },
+        { type: 'escaped_paren', value: '(', range: { start: 15, end: 17 } },
+        { type: 'word', text: 'medium', range: { start: 17, end: 23 }, attention: undefined },
+        { type: 'escaped_paren', value: ')', range: { start: 23, end: 25 } },
+        { type: 'whitespace', value: ' ', range: { start: 25, end: 26 } },
         {
           type: 'group',
-          children: [{ type: 'word', text: 'enhanced' }],
+          range: { start: 26, end: 36 },
+          attention: undefined,
+          children: [{ type: 'word', text: 'enhanced', range: { start: 27, end: 35 }, attention: undefined }],
         },
       ]);
     });
@@ -138,10 +146,11 @@ describe('promptAST', () => {
         {
           type: 'group',
           attention: 1.2,
+          range: { start: 0, end: 10 },
           children: [
-            { type: 'word', text: 'a' },
-            { type: 'whitespace', value: ' ' },
-            { type: 'word', text: 'cat' },
+            { type: 'word', text: 'a', range: { start: 1, end: 2 }, attention: undefined },
+            { type: 'whitespace', value: ' ', range: { start: 2, end: 3 } },
+            { type: 'word', text: 'cat', range: { start: 3, end: 6 }, attention: undefined },
           ],
         },
       ]);
@@ -150,13 +159,13 @@ describe('promptAST', () => {
     it('should parse words with attention', () => {
       const tokens = tokenize('cat+');
       const ast = parseTokens(tokens);
-      expect(ast).toEqual([{ type: 'word', text: 'cat', attention: '+' }]);
+      expect(ast).toEqual([{ type: 'word', text: 'cat', attention: '+', range: { start: 0, end: 4 } }]);
     });
 
     it('should parse embeddings', () => {
       const tokens = tokenize('<embedding_name>');
       const ast = parseTokens(tokens);
-      expect(ast).toEqual([{ type: 'embedding', value: 'embedding_name' }]);
+      expect(ast).toEqual([{ type: 'embedding', value: 'embedding_name', range: { start: 0, end: 16 } }]);
     });
   });
 
@@ -243,19 +252,20 @@ describe('promptAST', () => {
 
       // Should have escaped parens as nodes and a group with attention
       expect(ast).toEqual([
-        { type: 'word', text: 'portrait' },
-        { type: 'whitespace', value: ' ' },
-        { type: 'escaped_paren', value: '(' },
-        { type: 'word', text: 'realistic' },
-        { type: 'escaped_paren', value: ')' },
-        { type: 'whitespace', value: ' ' },
+        { type: 'word', text: 'portrait', range: { start: 0, end: 8 }, attention: undefined },
+        { type: 'whitespace', value: ' ', range: { start: 8, end: 9 } },
+        { type: 'escaped_paren', value: '(', range: { start: 9, end: 11 } },
+        { type: 'word', text: 'realistic', range: { start: 11, end: 20 }, attention: undefined },
+        { type: 'escaped_paren', value: ')', range: { start: 20, end: 22 } },
+        { type: 'whitespace', value: ' ', range: { start: 22, end: 23 } },
         {
           type: 'group',
           attention: 1.2,
+          range: { start: 23, end: 40 },
           children: [
-            { type: 'word', text: 'high' },
-            { type: 'whitespace', value: ' ' },
-            { type: 'word', text: 'quality' },
+            { type: 'word', text: 'high', range: { start: 24, end: 28 }, attention: undefined },
+            { type: 'whitespace', value: ' ', range: { start: 28, end: 29 } },
+            { type: 'word', text: 'quality', range: { start: 29, end: 36 }, attention: undefined },
           ],
         },
       ]);