InvokeAI 2.2.4.5__py3-none-any.whl

ldm/invoke/concepts_lib.py

@@ -0,0 +1,169 @@
+"""
+Query and install embeddings from the HuggingFace SD Concepts Library
+at https://huggingface.co/sd-concepts-library.
+
+The interface is through the Concepts() object.
+"""
+import os
+import re
+import traceback
+from typing import Callable
+from urllib import request, error as ul_error
+from huggingface_hub import HfFolder, hf_hub_url, ModelSearchArguments, ModelFilter, HfApi
+from ldm.invoke.globals import Globals
+
+class Concepts(object):
+    def __init__(self, root=None):
+        '''
+        Initialize the Concepts object. May optionally pass a root directory.
+        '''
+        self.root = root or Globals.root
+        self.hf_api = HfApi()
+        self.concept_list = None
+        self.concepts_loaded = dict()
+        self.triggers = dict()            # concept name to trigger phrase
+        self.concept_names = dict()       # trigger phrase to concept name
+        self.match_trigger = re.compile('(<[\w\- >]+>)') # trigger is slightly less restrictive than HF concept name
+        self.match_concept = re.compile('<([\w\-]+)>') # HF concept name can only contain A-Za-z0-9_-
+
+    def list_concepts(self)->list:
+        '''
+        Return a list of all the concepts by name, without the 'sd-concepts-library' part.
+        '''
+        if self.concept_list is not None:
+            return self.concept_list
+        try:
+            models = self.hf_api.list_models(filter=ModelFilter(model_name='sd-concepts-library/'))
+            self.concept_list = [a.id.split('/')[1] for a in models]
+        except Exception as e:
+            print(f' ** WARNING: Hugging Face textual inversion concepts libraries could not be loaded. The error was {str(e)}.')
+            print(' ** You may load .bin and .pt file(s) manually using the --embedding_directory argument.')
+        return self.concept_list
+
+    def get_concept_model_path(self, concept_name:str)->str:
+        '''
+        Returns the path to the 'learned_embeds.bin' file in
+        the named concept. Returns None if invalid or cannot
+        be downloaded.
+        '''
+        if not concept_name in self.list_concepts():
+            print(f'This concept is not known to the Hugging Face library. Generation will continue without the concept.')
+            return None
+        return self.get_concept_file(concept_name.lower(),'learned_embeds.bin')
+
+    def concept_to_trigger(self, concept_name:str)->str:
+        '''
+        Given a concept name returns its trigger by looking in the
+        "token_identifier.txt" file.
+        '''
+        if concept_name in self.triggers:
+            return self.triggers[concept_name]
+        file = self.get_concept_file(concept_name, 'token_identifier.txt', local_only=True)
+        if not file:
+            return None
+        with open(file,'r') as f:
+            trigger = f.readline()
+            trigger = trigger.strip()
+        self.triggers[concept_name] = trigger
+        self.concept_names[trigger] = concept_name
+        return trigger
+
+    def trigger_to_concept(self, trigger:str)->str:
+        '''
+        Given a trigger phrase, maps it to the concept library name.
+        Only works if concept_to_trigger() has previously been called
+        on this library. There needs to be a persistent database for
+        this.
+        '''
+        concept = self.concept_names.get(trigger,None)
+        return f'<{concept}>' if concept else f'{trigger}'
+
+    def replace_triggers_with_concepts(self, prompt:str)->str:
+        '''
+        Given a prompt string that contains <trigger> tags, replace these
+        tags with the concept name. The reason for this is so that the
+        concept names get stored in the prompt metadata. There is no
+        controlling of colliding triggers in the SD library, so it is
+        better to store the concept name (unique) than the concept trigger
+        (not necessarily unique!)
+        '''
+        if not prompt:
+            return prompt
+        triggers = self.match_trigger.findall(prompt)
+        if not triggers:
+            return prompt
+
+        def do_replace(match)->str:
+            return self.trigger_to_concept(match.group(1)) or f'<{match.group(1)}>'
+        return self.match_trigger.sub(do_replace, prompt)
+
+    def replace_concepts_with_triggers(self, prompt:str, load_concepts_callback: Callable[[list], any])->str:
+        '''
+        Given a prompt string that contains `<concept_name>` tags, replace
+        these tags with the appropriate trigger.
+
+        If any `<concept_name>` tags are found, `load_concepts_callback()` is called with a list
+        of `concepts_name` strings.
+        '''
+        concepts = self.match_concept.findall(prompt)
+        if not concepts:
+            return prompt
+        load_concepts_callback(concepts)
+
+        def do_replace(match)->str:
+            return self.concept_to_trigger(match.group(1)) or f'<{match.group(1)}>'
+        return self.match_concept.sub(do_replace, prompt)
+
+    def get_concept_file(self, concept_name:str, file_name:str='learned_embeds.bin' , local_only:bool=False)->str:
+        if not self.concept_is_downloaded(concept_name) and not local_only:
+            self.download_concept(concept_name)
+        path = os.path.join(self._concept_path(concept_name), file_name)
+        return path if os.path.exists(path) else None
+
+    def concept_is_downloaded(self, concept_name)->bool:
+        concept_directory = self._concept_path(concept_name)
+        return os.path.exists(concept_directory)
+
+    def download_concept(self,concept_name)->bool:
+        repo_id = self._concept_id(concept_name)
+        dest = self._concept_path(concept_name)
+
+        access_token = HfFolder.get_token()
+        header = [("Authorization", f'Bearer {access_token}')] if access_token else []
+        opener = request.build_opener()
+        opener.addheaders = header
+        request.install_opener(opener)
+
+        os.makedirs(dest, exist_ok=True)
+        succeeded = True
+
+        bytes = 0
+        def tally_download_size(chunk, size, total):
+            nonlocal bytes
+            if chunk==0:
+                bytes += total
+
+        print(f'>> Downloading {repo_id}...',end='')
+        try:
+            for file in ('README.md','learned_embeds.bin','token_identifier.txt','type_of_concept.txt'):
+                url = hf_hub_url(repo_id, file)
+                request.urlretrieve(url, os.path.join(dest,file),reporthook=tally_download_size)
+        except ul_error.HTTPError as e:
+            if e.code==404:
+                print(f'This concept is not known to the Hugging Face library. Generation will continue without the concept.')
+            else:
+                print(f'Failed to download {concept_name}/{file} ({str(e)}. Generation will continue without the concept.)')
+            os.rmdir(dest)
+            return False
+        except ul_error.URLError as e:
+            print(f'ERROR: {str(e)}. This may reflect a network issue. Generation will continue without the concept.')
+            os.rmdir(dest)
+            return False
+        print('...{:.2f}Kb'.format(bytes/1024))
+        return succeeded
+
+    def _concept_id(self, concept_name:str)->str:
+        return f'sd-concepts-library/{concept_name}'
+
+    def _concept_path(self, concept_name:str)->str:
+        return os.path.join(self.root,'models','sd-concepts-library',concept_name)

ldm/invoke/conditioning.py

@@ -0,0 +1,289 @@
+'''
+This module handles the generation of the conditioning tensors.
+
+Useful function exports:
+
+get_uc_and_c_and_ec()           get the conditioned and unconditioned latent, and edited conditioning if we're doing cross-attention control
+
+'''
+import re
+from typing import Union
+
+import torch
+
+from .prompt_parser import PromptParser, Blend, FlattenedPrompt, \
+    CrossAttentionControlledFragment, CrossAttentionControlSubstitute, Fragment
+from ..models.diffusion import cross_attention_control
+from ..models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent
+from ..modules.encoders.modules import WeightedFrozenCLIPEmbedder
+
+
+def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_legacy_blend=False):
+    prompt, negative_prompt = get_prompt_structure(prompt_string,
+                                                   skip_normalize_legacy_blend=skip_normalize_legacy_blend)
+    conditioning = _get_conditioning_for_prompt(prompt, negative_prompt, model, log_tokens)
+
+    return conditioning
+
+
+def get_prompt_structure(prompt_string, skip_normalize_legacy_blend: bool = False) -> (
+Union[FlattenedPrompt, Blend], FlattenedPrompt):
+    """
+    parse the passed-in prompt string and return tuple (positive_prompt, negative_prompt)
+    """
+    prompt, negative_prompt = _parse_prompt_string(prompt_string,
+                                                   skip_normalize_legacy_blend=skip_normalize_legacy_blend)
+    return prompt, negative_prompt
+
+
+def get_tokens_for_prompt(model, parsed_prompt: FlattenedPrompt, truncate_if_too_long=True) -> [str]:
+    text_fragments = [x.text if type(x) is Fragment else
+                      (" ".join([f.text for f in x.original]) if type(x) is CrossAttentionControlSubstitute else
+                       str(x))
+                      for x in parsed_prompt.children]
+    text = " ".join(text_fragments)
+    tokens = model.cond_stage_model.tokenizer.tokenize(text)
+    if truncate_if_too_long:
+        max_tokens_length = model.cond_stage_model.max_length - 2 # typically 75
+        tokens = tokens[0:max_tokens_length]
+    return tokens
+
+
+def _parse_prompt_string(prompt_string_uncleaned, skip_normalize_legacy_blend=False) -> Union[FlattenedPrompt, Blend]:
+    # Extract Unconditioned Words From Prompt
+    unconditioned_words = ''
+    unconditional_regex = r'\[(.*?)\]'
+    unconditionals = re.findall(unconditional_regex, prompt_string_uncleaned)
+
+    if len(unconditionals) > 0:
+        unconditioned_words = ' '.join(unconditionals)
+
+        # Remove Unconditioned Words From Prompt
+        unconditional_regex_compile = re.compile(unconditional_regex)
+        clean_prompt = unconditional_regex_compile.sub(' ', prompt_string_uncleaned)
+        prompt_string_cleaned = re.sub(' +', ' ', clean_prompt)
+    else:
+        prompt_string_cleaned = prompt_string_uncleaned
+
+    pp = PromptParser()
+
+    parsed_prompt: Union[FlattenedPrompt, Blend] = None
+    legacy_blend: Blend = pp.parse_legacy_blend(prompt_string_cleaned, skip_normalize_legacy_blend)
+    if legacy_blend is not None:
+        parsed_prompt = legacy_blend
+    else:
+        # we don't support conjunctions for now
+        parsed_prompt = pp.parse_conjunction(prompt_string_cleaned).prompts[0]
+
+    parsed_negative_prompt: FlattenedPrompt = pp.parse_conjunction(unconditioned_words).prompts[0]
+    return parsed_prompt, parsed_negative_prompt
+
+
+def _get_conditioning_for_prompt(parsed_prompt: Union[Blend, FlattenedPrompt], parsed_negative_prompt: FlattenedPrompt,
+                                 model, log_tokens=False) \
+    -> tuple[torch.Tensor, torch.Tensor, InvokeAIDiffuserComponent.ExtraConditioningInfo]:
+    """
+    Process prompt structure and tokens, and return (conditioning, unconditioning, extra_conditioning_info)
+    """
+
+    if log_tokens:
+        print(f">> Parsed prompt to {parsed_prompt}")
+        print(f">> Parsed negative prompt to {parsed_negative_prompt}")
+
+    conditioning = None
+    cac_args: cross_attention_control.Arguments = None
+
+    if type(parsed_prompt) is Blend:
+        conditioning = _get_conditioning_for_blend(model, parsed_prompt, log_tokens)
+    elif type(parsed_prompt) is FlattenedPrompt:
+        if parsed_prompt.wants_cross_attention_control:
+            conditioning, cac_args = _get_conditioning_for_cross_attention_control(model, parsed_prompt, log_tokens)
+
+        else:
+            conditioning, _ = _get_embeddings_and_tokens_for_prompt(model,
+                                                                    parsed_prompt,
+                                                                    log_tokens=log_tokens,
+                                                                    log_display_label="(prompt)")
+    else:
+        raise ValueError(f"parsed_prompt is '{type(parsed_prompt)}' which is not a supported prompt type")
+
+    unconditioning, _ = _get_embeddings_and_tokens_for_prompt(model,
+                                                              parsed_negative_prompt,
+                                                              log_tokens=log_tokens,
+                                                              log_display_label="(unconditioning)")
+    if isinstance(conditioning, dict):
+        # hybrid conditioning is in play
+        unconditioning, conditioning = _flatten_hybrid_conditioning(unconditioning, conditioning)
+        if cac_args is not None:
+            print(
+                ">> Hybrid conditioning cannot currently be combined with cross attention control. Cross attention control will be ignored.")
+            cac_args = None
+
+    if type(parsed_prompt) is Blend:
+        blend: Blend = parsed_prompt
+        all_token_sequences = [get_tokens_for_prompt(model, p) for p in blend.prompts]
+        longest_token_sequence = max(all_token_sequences, key=lambda t: len(t))
+        eos_token_index = len(longest_token_sequence)+1
+    else:
+        tokens = get_tokens_for_prompt(model, parsed_prompt)
+        eos_token_index = len(tokens)+1
+    return (
+        unconditioning, conditioning, InvokeAIDiffuserComponent.ExtraConditioningInfo(
+            tokens_count_including_eos_bos=eos_token_index + 1,
+            cross_attention_control_args=cac_args
+        )
+    )
+
+
+def _get_conditioning_for_cross_attention_control(model, prompt: FlattenedPrompt, log_tokens: bool = True):
+    original_prompt = FlattenedPrompt()
+    edited_prompt = FlattenedPrompt()
+    # for name, a0, a1, b0, b1 in edit_opcodes: only name == 'equal' is currently parsed
+    original_token_count = 0
+    edited_token_count = 0
+    edit_options = []
+    edit_opcodes = []
+    # beginning of sequence
+    edit_opcodes.append(
+        ('equal', original_token_count, original_token_count + 1, edited_token_count, edited_token_count + 1))
+    edit_options.append(None)
+    original_token_count += 1
+    edited_token_count += 1
+    for fragment in prompt.children:
+        if type(fragment) is CrossAttentionControlSubstitute:
+            original_prompt.append(fragment.original)
+            edited_prompt.append(fragment.edited)
+
+            to_replace_token_count = _get_tokens_length(model, fragment.original)
+            replacement_token_count = _get_tokens_length(model, fragment.edited)
+            edit_opcodes.append(('replace',
+                                 original_token_count, original_token_count + to_replace_token_count,
+                                 edited_token_count, edited_token_count + replacement_token_count
+                                 ))
+            original_token_count += to_replace_token_count
+            edited_token_count += replacement_token_count
+            edit_options.append(fragment.options)
+        # elif type(fragment) is CrossAttentionControlAppend:
+        #    edited_prompt.append(fragment.fragment)
+        else:
+            # regular fragment
+            original_prompt.append(fragment)
+            edited_prompt.append(fragment)
+
+            count = _get_tokens_length(model, [fragment])
+            edit_opcodes.append(('equal', original_token_count, original_token_count + count, edited_token_count,
+                                 edited_token_count + count))
+            edit_options.append(None)
+            original_token_count += count
+            edited_token_count += count
+    # end of sequence
+    edit_opcodes.append(
+        ('equal', original_token_count, original_token_count + 1, edited_token_count, edited_token_count + 1))
+    edit_options.append(None)
+    original_token_count += 1
+    edited_token_count += 1
+    original_embeddings, original_tokens = _get_embeddings_and_tokens_for_prompt(model,
+                                                                                 original_prompt,
+                                                                                 log_tokens=log_tokens,
+                                                                                 log_display_label="(.swap originals)")
+    # naïvely building a single edited_embeddings like this disregards the effects of changing the absolute location of
+    # subsequent tokens when there is >1 edit and earlier edits change the total token count.
+    # eg "a cat.swap(smiling dog, s_start=0.5) eating a hotdog.swap(pizza)" - when the 'pizza' edit is active but the
+    # 'cat' edit is not, the 'pizza' feature vector will nevertheless be affected by the introduction of the extra
+    # token 'smiling' in the inactive 'cat' edit.
+    # todo: build multiple edited_embeddings, one for each edit, and pass just the edited fragments through to the CrossAttentionControl functions
+    edited_embeddings, edited_tokens = _get_embeddings_and_tokens_for_prompt(model,
+                                                                             edited_prompt,
+                                                                             log_tokens=log_tokens,
+                                                                             log_display_label="(.swap replacements)")
+    conditioning = original_embeddings
+    edited_conditioning = edited_embeddings
+    # print('>> got edit_opcodes', edit_opcodes, 'options', edit_options)
+    cac_args = cross_attention_control.Arguments(
+        edited_conditioning=edited_conditioning,
+        edit_opcodes=edit_opcodes,
+        edit_options=edit_options
+    )
+    return conditioning, cac_args
+
+
+def _get_conditioning_for_blend(model, blend: Blend, log_tokens: bool = False):
+    embeddings_to_blend = None
+    for i, flattened_prompt in enumerate(blend.prompts):
+        this_embedding, _ = _get_embeddings_and_tokens_for_prompt(model,
+                                                                  flattened_prompt,
+                                                                  log_tokens=log_tokens,
+                                                                  log_display_label=f"(blend part {i + 1}, weight={blend.weights[i]})")
+        embeddings_to_blend = this_embedding if embeddings_to_blend is None else torch.cat(
+            (embeddings_to_blend, this_embedding))
+    conditioning = WeightedFrozenCLIPEmbedder.apply_embedding_weights(embeddings_to_blend.unsqueeze(0),
+                                                                      blend.weights,
+                                                                      normalize=blend.normalize_weights)
+    return conditioning
+
+
+def _get_embeddings_and_tokens_for_prompt(model, flattened_prompt: FlattenedPrompt, log_tokens: bool = False,
+                                          log_display_label: str = None):
+    if type(flattened_prompt) is not FlattenedPrompt:
+        raise Exception(f"embeddings can only be made from FlattenedPrompts, got {type(flattened_prompt)} instead")
+    fragments = [x.text for x in flattened_prompt.children]
+    weights = [x.weight for x in flattened_prompt.children]
+    embeddings, tokens = model.get_learned_conditioning([fragments], return_tokens=True, fragment_weights=[weights])
+    if log_tokens:
+        text = " ".join(fragments)
+        log_tokenization(text, model, display_label=log_display_label)
+
+    return embeddings, tokens
+
+
+def _get_tokens_length(model, fragments: list[Fragment]):
+    fragment_texts = [x.text for x in fragments]
+    tokens = model.cond_stage_model.get_tokens(fragment_texts, include_start_and_end_markers=False)
+    return sum([len(x) for x in tokens])
+
+
+def _flatten_hybrid_conditioning(uncond, cond):
+    '''
+    This handles the choice between a conditional conditioning
+    that is a tensor (used by cross attention) vs one that has additional
+    dimensions as well, as used by 'hybrid'
+    '''
+    assert isinstance(uncond, dict)
+    assert isinstance(cond, dict)
+    cond_flattened = dict()
+    for k in cond:
+        if isinstance(cond[k], list):
+            cond_flattened[k] = [
+                torch.cat([uncond[k][i], cond[k][i]])
+                for i in range(len(cond[k]))
+            ]
+        else:
+            cond_flattened[k] = torch.cat([uncond[k], cond[k]])
+    return uncond, cond_flattened
+
+
+def log_tokenization(text, model, display_label=None):
+    """ shows how the prompt is tokenized
+    # usually tokens have '</w>' to indicate end-of-word,
+    # but for readability it has been replaced with ' '
+    """
+
+    tokens = model.cond_stage_model.tokenizer.tokenize(text)
+    tokenized = ""
+    discarded = ""
+    usedTokens = 0
+    totalTokens = len(tokens)
+    for i in range(0, totalTokens):
+        token = tokens[i].replace('</w>', ' ')
+        # alternate color
+        s = (usedTokens % 6) + 1
+        if i < model.cond_stage_model.max_length:
+            tokenized = tokenized + f"\x1b[0;3{s};40m{token}"
+            usedTokens += 1
+        else:  # over max token length
+            discarded = discarded + f"\x1b[0;3{s};40m{token}"
+    print(f"\n>> Tokens {display_label or ''} ({usedTokens}):\n{tokenized}\x1b[0m")
+    if discarded != "":
+        print(
+            f">> Tokens Discarded ({totalTokens - usedTokens}):\n{discarded}\x1b[0m"
+        )

ldm/invoke/devices.py

@@ -0,0 +1,27 @@
+import torch
+from torch import autocast
+from contextlib import nullcontext
+
+def choose_torch_device() -> str:
+    '''Convenience routine for guessing which GPU device to run model on'''
+    if torch.cuda.is_available():
+        return 'cuda'
+    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        return 'mps'
+    return 'cpu'
+
+def choose_precision(device) -> str:
+    '''Returns an appropriate precision for the given torch device'''
+    if device.type == 'cuda':
+        device_name = torch.cuda.get_device_name(device)
+        if not ('GeForce GTX 1660' in device_name or 'GeForce GTX 1650' in device_name):
+            return 'float16'
+    return 'float32'
+
+def choose_autocast(precision):
+    '''Returns an autocast context or nullcontext for the given precision string'''
+    # float16 currently requires autocast to avoid errors like:
+    # 'expected scalar type Half but found Float'
+    if precision == 'autocast' or precision == 'float16':
+        return autocast
+    return nullcontext

ldm/invoke/generator/__init__.py

@@ -0,0 +1,4 @@
+'''
+Initialization file for the ldm.invoke.generator package
+'''
+from .base import Generator