InvokeAI/invokeai/app/util/ti_utils.py

import re
from typing import List, Tuple

import invokeai.backend.util.logging as logger
from invokeai.app.services.model_records import UnknownModelException
from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.model_manager.config import BaseModelType, ModelType
from invokeai.backend.textual_inversion import TextualInversionModelRaw


def extract_ti_triggers_from_prompt(prompt: str) -> List[str]:
    ti_triggers: List[str] = []
    for trigger in re.findall(r"<[a-zA-Z0-9., _-]+>", prompt):
        ti_triggers.append(str(trigger))
    return ti_triggers


def generate_ti_list(
    prompt: str, base: BaseModelType, context: InvocationContext
) -> List[Tuple[str, TextualInversionModelRaw]]:
    ti_list: List[Tuple[str, TextualInversionModelRaw]] = []
    for trigger in extract_ti_triggers_from_prompt(prompt):
        name_or_key = trigger[1:-1]
        try:
            loaded_model = context.models.load(name_or_key)
            model = loaded_model.model
            assert isinstance(model, TextualInversionModelRaw)
            assert loaded_model.config.base == base
            ti_list.append((name_or_key, model))
        except UnknownModelException:
            try:
                loaded_model = context.models.load_by_attrs(
                    name=name_or_key, base=base, type=ModelType.TextualInversion
                )
                model = loaded_model.model
                assert isinstance(model, TextualInversionModelRaw)
                assert loaded_model.config.base == base
                ti_list.append((name_or_key, model))
            except UnknownModelException:
                pass
        except ValueError:
            logger.warning(f'trigger: "{trigger}" more than one similarly-named textual inversion models')
        except AssertionError:
            logger.warning(f'trigger: "{trigger}" not a valid textual inversion model for this graph')
        except Exception:
            logger.warning(f'Failed to load TI model for trigger: "{trigger}"')
    return ti_list