Merge branch 'main' into copilot/add-unload-model-option

Author: lstein

invokeai/app/invocations/fields.py

@@ -243,6 +243,12 @@ class BoardField(BaseModel):
     board_id: str = Field(description="The id of the board")
 
 
+class StylePresetField(BaseModel):
+    """A style preset primitive field"""
+
+    style_preset_id: str = Field(description="The id of the style preset")
+
+
 class DenoiseMaskField(BaseModel):
     """An inpaint mask field"""
 

invokeai/app/invocations/prompt_template.py

@@ -0,0 +1,57 @@
+from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
+from invokeai.app.invocations.fields import InputField, OutputField, StylePresetField, UIComponent
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+@invocation_output("prompt_template_output")
+class PromptTemplateOutput(BaseInvocationOutput):
+    """Output for the Prompt Template node"""
+
+    positive_prompt: str = OutputField(description="The positive prompt with the template applied")
+    negative_prompt: str = OutputField(description="The negative prompt with the template applied")
+
+
+@invocation(
+    "prompt_template",
+    title="Prompt Template",
+    tags=["prompt", "template", "style", "preset"],
+    category="prompt",
+    version="1.0.0",
+)
+class PromptTemplateInvocation(BaseInvocation):
+    """Applies a Style Preset template to positive and negative prompts.
+
+    Select a Style Preset and provide positive/negative prompts. The node replaces
+    {prompt} placeholders in the template with your input prompts.
+    """
+
+    style_preset: StylePresetField = InputField(
+        description="The Style Preset to use as a template",
+    )
+    positive_prompt: str = InputField(
+        default="",
+        description="The positive prompt to insert into the template's {prompt} placeholder",
+        ui_component=UIComponent.Textarea,
+    )
+    negative_prompt: str = InputField(
+        default="",
+        description="The negative prompt to insert into the template's {prompt} placeholder",
+        ui_component=UIComponent.Textarea,
+    )
+
+    def invoke(self, context: InvocationContext) -> PromptTemplateOutput:
+        # Fetch the style preset from the database
+        style_preset = context._services.style_preset_records.get(self.style_preset.style_preset_id)
+
+        # Get the template prompts
+        positive_template = style_preset.preset_data.positive_prompt
+        negative_template = style_preset.preset_data.negative_prompt
+
+        # Replace {prompt} placeholder with the input prompts
+        rendered_positive = positive_template.replace("{prompt}", self.positive_prompt)
+        rendered_negative = negative_template.replace("{prompt}", self.negative_prompt)
+
+        return PromptTemplateOutput(
+            positive_prompt=rendered_positive,
+            negative_prompt=rendered_negative,
+        )

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/public/locales/en.json

@@ -1985,6 +1985,7 @@
         "allLoaded": "All Workflows Loaded",
         "searchPlaceholder": "Search by name, description or tags",
         "filterByTags": "Filter by Tags",
+        "tags": "Tags",
         "yourWorkflows": "Your Workflows",
         "recentlyOpened": "Recently Opened",
         "noRecentWorkflows": "No Recent Workflows",
@@ -2676,7 +2677,9 @@
         "useForTemplate": "Use For Prompt Template",
         "viewList": "View Template List",
         "viewModeTooltip": "This is how your prompt will look with your currently selected template. To edit your prompt, click anywhere in the text box.",
-        "togglePromptPreviews": "Toggle Prompt Previews"
+        "togglePromptPreviews": "Toggle Prompt Previews",
+        "selectPreset": "Select Style Preset",
+        "noMatchingPresets": "No matching presets"
     },
 
     "ui": {

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/InputFieldRenderer.tsx

@@ -55,6 +55,8 @@ import {
   isStringFieldInputTemplate,
   isStringGeneratorFieldInputInstance,
   isStringGeneratorFieldInputTemplate,
+  isStylePresetFieldInputInstance,
+  isStylePresetFieldInputTemplate,
 } from 'features/nodes/types/field';
 import type { NodeFieldElement } from 'features/nodes/types/workflow';
 import { memo } from 'react';
@@ -67,6 +69,7 @@ import ColorFieldInputComponent from './inputs/ColorFieldInputComponent';
 import EnumFieldInputComponent from './inputs/EnumFieldInputComponent';
 import ImageFieldInputComponent from './inputs/ImageFieldInputComponent';
 import SchedulerFieldInputComponent from './inputs/SchedulerFieldInputComponent';
+import StylePresetFieldInputComponent from './inputs/StylePresetFieldInputComponent';
 
 type Props = {
   nodeId: string;
@@ -206,6 +209,13 @@ export const InputFieldRenderer = memo(({ nodeId, fieldName, settings }: Props)
     return <BoardFieldInputComponent nodeId={nodeId} field={field} fieldTemplate={template} />;
   }
 
+  if (isStylePresetFieldInputTemplate(template)) {
+    if (!isStylePresetFieldInputInstance(field)) {
+      return null;
+    }
+    return <StylePresetFieldInputComponent nodeId={nodeId} field={field} fieldTemplate={template} />;
+  }
+
   if (isModelIdentifierFieldInputTemplate(template)) {
     if (!isModelIdentifierFieldInputInstance(field)) {
       return null;

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/StylePresetFieldInputComponent.tsx

@@ -0,0 +1,73 @@
+import type { ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library';
+import { Combobox } from '@invoke-ai/ui-library';
+import { useAppDispatch } from 'app/store/storeHooks';
+import { fieldStylePresetValueChanged } from 'features/nodes/store/nodesSlice';
+import { NO_DRAG_CLASS, NO_WHEEL_CLASS } from 'features/nodes/types/constants';
+import type { StylePresetFieldInputInstance, StylePresetFieldInputTemplate } from 'features/nodes/types/field';
+import { memo, useCallback, useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useListStylePresetsQuery } from 'services/api/endpoints/stylePresets';
+
+import type { FieldComponentProps } from './types';
+
+const StylePresetFieldInputComponent = (
+  props: FieldComponentProps<StylePresetFieldInputInstance, StylePresetFieldInputTemplate>
+) => {
+  const { nodeId, field } = props;
+  const dispatch = useAppDispatch();
+  const { t } = useTranslation();
+  const { data: stylePresets, isLoading } = useListStylePresetsQuery();
+
+  const options = useMemo<ComboboxOption[]>(() => {
+    const _options: ComboboxOption[] = [];
+    if (stylePresets) {
+      for (const preset of stylePresets) {
+        _options.push({
+          label: preset.name,
+          value: preset.id,
+        });
+      }
+    }
+    return _options;
+  }, [stylePresets]);
+
+  const onChange = useCallback<ComboboxOnChange>(
+    (v) => {
+      if (!v) {
+        return;
+      }
+
+      dispatch(
+        fieldStylePresetValueChanged({
+          nodeId,
+          fieldName: field.name,
+          value: { style_preset_id: v.value },
+        })
+      );
+    },
+    [dispatch, field.name, nodeId]
+  );
+
+  const value = useMemo(() => {
+    const _value = field.value;
+    if (!_value) {
+      return null;
+    }
+    return options.find((o) => o.value === _value.style_preset_id) ?? null;
+  }, [field.value, options]);
+
+  const noOptionsMessage = useCallback(() => t('stylePresets.noMatchingPresets'), [t]);
+
+  return (
+    <Combobox
+      className={`${NO_WHEEL_CLASS} ${NO_DRAG_CLASS}`}
+      value={value}
+      options={options}
+      onChange={onChange}
+      placeholder={isLoading ? t('common.loading') : t('stylePresets.selectPreset')}
+      noOptionsMessage={noOptionsMessage}
+    />
+  );
+};
+
+export default memo(StylePresetFieldInputComponent);