diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index e0dfd07dbe..13760dc11c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -81,7 +81,10 @@ class LoRALoader(ModelLoader): state_dict = torch.load(model_path, map_location="cpu") # 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]: + 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 a94c5b9ff2..e595e6d442 100644 --- a/invokeai/backend/model_manager/omi.py +++ b/invokeai/backend/model_manager/omi.py @@ -1,7 +1,5 @@ import omi_model_standards.convert.lora.convert_lora_util as lora_util from omi_model_standards.convert.lora.convert_flux_lora import convert_flux_lora_key_sets -from omi_model_standards.convert.lora.convert_sd3_lora import convert_sd3_lora_key_sets -from omi_model_standards.convert.lora.convert_sd_lora import convert_sd_lora_key_sets from omi_model_standards.convert.lora.convert_sdxl_lora import convert_sdxl_lora_key_sets from invokeai.backend.model_manager.model_on_disk import StateDict @@ -15,4 +13,4 @@ def convert_from_omi(weights_sd: StateDict, base: BaseModelType): }[base] source = "omi" target = "legacy_diffusers" - return lora_util.__convert(weights_sd, keyset, source, target) # type: ignore + return lora_util.__convert(weights_sd, keyset, source, target) # type: ignore