From e1157f343b8630059e137bb71278ebe57f766575 Mon Sep 17 00:00:00 2001 From: Billy Date: Mon, 23 Jun 2025 13:51:16 +1000 Subject: [PATCH] Support for Flux and SDXL --- invokeai/backend/model_manager/config.py | 6 +----- invokeai/backend/model_manager/load/model_loaders/lora.py | 5 ++--- invokeai/backend/model_manager/omi.py | 8 +++----- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/invokeai/backend/model_manager/config.py b/invokeai/backend/model_manager/config.py index b995815d1f..e1ebb1856e 100644 --- a/invokeai/backend/model_manager/config.py +++ b/invokeai/backend/model_manager/config.py @@ -356,15 +356,11 @@ class LoRAOmiConfig(LoRAConfigBase, ModelConfigBase): base_str = base_str.lower() if "stable-diffusion-v1" in base_str: - base = BaseModelType.StableDiffusion1 - elif "stable-diffusion-v3" in base_str: - base = BaseModelType.StableDiffusion3 - elif base_str == "stable-diffusion-xl-v1-base": base = BaseModelType.StableDiffusionXL elif "flux" in base_str: base = BaseModelType.Flux else: - raise InvalidModelConfigException(f"Unrecognised base architecture for OMI LoRA: {base_str}") + raise InvalidModelConfigException(f"Unrecognised/unsupported base architecture for OMI LoRA: {base_str}") return {"base": base} diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index be6da03e24..3c67986bee 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -41,8 +41,6 @@ from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import @ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.LoRA, format=ModelFormat.OMI) -@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.LoRA, format=ModelFormat.OMI) -@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion3, type=ModelType.LoRA, format=ModelFormat.OMI) @ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.LoRA, format=ModelFormat.OMI) @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LoRA, format=ModelFormat.Diffusers) @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LoRA, format=ModelFormat.LyCORIS) @@ -78,7 +76,8 @@ class LoRALoader(ModelLoader): else: state_dict = torch.load(model_path, map_location="cpu") - if config.format == ModelFormat.OMI: + # At the time of writing, we support the OMI standard for base models Flux and SDXL + if config.format == ModelFormat.OMI and self._model_base in [BaseModelType.StableDiffusionXL, BaseModelType.Flux]: state_dict = convert_from_omi(state_dict, config.base) # type: ignore # Apply state_dict key conversions, if necessary. diff --git a/invokeai/backend/model_manager/omi.py b/invokeai/backend/model_manager/omi.py index 94b339672c..a94c5b9ff2 100644 --- a/invokeai/backend/model_manager/omi.py +++ b/invokeai/backend/model_manager/omi.py @@ -12,9 +12,7 @@ def convert_from_omi(weights_sd: StateDict, base: BaseModelType): keyset = { BaseModelType.Flux: convert_flux_lora_key_sets(), BaseModelType.StableDiffusionXL: convert_sdxl_lora_key_sets(), - BaseModelType.StableDiffusion1: convert_sd_lora_key_sets(), - BaseModelType.StableDiffusion3: convert_sd3_lora_key_sets(), }[base] - - target = "diffusers" # alternatively, "legacy_diffusers" - return lora_util.__convert(weights_sd, keyset, "omi", target) # type: ignore + source = "omi" + target = "legacy_diffusers" + return lora_util.__convert(weights_sd, keyset, source, target) # type: ignore