InvokeAI/invokeai/app/invocations/util/choose_model.py
2023-05-11 00:19:20 -04:00

15 lines
758 B
Python

from invokeai.backend.model_management.model_manager_service import ModelManagerService, SDModelType
def choose_model(model_manager: ModelManagerService, model_name: str, model_type: SDModelType=SDModelType.diffusers):
"""Returns the default model if the `model_name` not a valid model, else returns the selected model."""
logger = model_manager.logger
if model_name and not model_manager.valid_model(model_name, model_type):
default_model_name = model_manager.default_model()
logger.warning(f"\'{model_name}\' is not a valid model name. Using default model \'{default_model_name}\' instead.")
model = model_manager.get_model()
else:
model = model_manager.get_model(model_name, model_type)
return model