add prompt caching and tokenizer info

modules/prompt_parser_diffusers.py

@@ -8,9 +8,14 @@ from transformers import PreTrainedTokenizer
 from modules import shared, prompt_parser, devices
 
 
+debug_enabled = os.environ.get('SD_PROMPT_DEBUG', None)
 debug = shared.log.trace if os.environ.get('SD_PROMPT_DEBUG', None) is not None else lambda *args, **kwargs: None
 debug('Trace: PROMPT')
 orig_encode_token_ids_to_embeddings = EmbeddingsProvider._encode_token_ids_to_embeddings # pylint: disable=protected-access
+token_dict = None
+token_type = None
+cache = {}
+cache_type = None
 
 
 def compel_hijack(self, token_ids: torch.Tensor,
@@ -97,10 +102,41 @@ def get_prompt_schedule(prompt, steps):
     return temp, len(schedule) > 1
 
 
+def get_tokens(msg, prompt):
+    global token_dict, token_type # pylint: disable=global-statement
+    if shared.backend != shared.Backend.DIFFUSERS:
+        return
+    if shared.sd_loaded and hasattr(shared.sd_model, 'tokenizer') and shared.sd_model.tokenizer is not None:
+        if token_dict is None or token_type != shared.sd_model_type:
+            token_type = shared.sd_model_type
+            fn = os.path.join(shared.sd_model.tokenizer.name_or_path, 'tokenizer', 'vocab.json')
+            token_dict = shared.readfile(fn, silent=True)
+            for k, v in shared.sd_model.tokenizer.added_tokens_decoder.items():
+                token_dict[str(v)] = k
+            shared.log.debug(f'Tokenizer: words={len(token_dict)} file="{fn}"')
+        has_bos_token = shared.sd_model.tokenizer.bos_token_id is not None
+        has_eos_token = shared.sd_model.tokenizer.eos_token_id is not None
+        ids = shared.sd_model.tokenizer(prompt)
+        ids = getattr(ids, 'input_ids', [])
+        tokens = []
+        for i in ids:
+            tokens.append(list(token_dict.keys())[list(token_dict.values()).index(i)])
+        token_count = len(ids) - int(has_bos_token) - int(has_eos_token)
+        shared.log.trace(f'Prompt tokenizer: type={msg} tokens={token_count} {tokens}')
+
+
 def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, clip_skip: typing.Optional[int] = None):
     if 'StableDiffusion' not in pipe.__class__.__name__ and 'DemoFusion' not in pipe.__class__.__name__:
         shared.log.warning(f"Prompt parser not supported: {pipe.__class__.__name__}")
         return
+    elif prompts == cache.get('prompts', None) and negative_prompts == cache.get('negative_prompts', None) and cache.get('model_type', None) == shared.sd_model_type:
+        p.prompt_embeds = cache.get('prompt_embeds', None)
+        p.positive_pooleds = cache.get('positive_pooleds', None)
+        p.negative_embeds = cache.get('negative_embeds', None)
+        p.negative_pooleds = cache.get('negative_pooleds', None)
+        p.scheduled_prompt = cache.get('scheduled_prompt', None)
+        debug("Prompt encode: cached")
+        return
     else:
         t0 = time.time()
         positive_schedule, scheduled = get_prompt_schedule(prompts[0], steps)
@@ -111,7 +147,6 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
         p.negative_embeds = []
         p.negative_pooleds = []
 
-        cache = {}
         for i in range(max(len(positive_schedule), len(negative_schedule))):
             positive_prompt = positive_schedule[i % len(positive_schedule)]
             negative_prompt = negative_schedule[i % len(negative_schedule)]
@@ -124,12 +159,23 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
             prompt_embed, positive_pooled, negative_embed, negative_pooled = results
             if prompt_embed is not None:
                 p.prompt_embeds.append(torch.cat([prompt_embed] * len(prompts), dim=0))
+                cache['prompt_embeds'] = p.prompt_embeds
             if negative_embed is not None:
                 p.negative_embeds.append(torch.cat([negative_embed] * len(negative_prompts), dim=0))
+                cache['negative_embeds'] = p.negative_embeds
             if positive_pooled is not None:
                 p.positive_pooleds.append(torch.cat([positive_pooled] * len(prompts), dim=0))
+                cache['positive_pooleds'] = p.positive_pooleds
             if negative_pooled is not None:
                 p.negative_pooleds.append(torch.cat([negative_pooled] * len(negative_prompts), dim=0))
+                cache['negative_pooleds'] = p.negative_pooleds
+
+        cache['prompts'] = prompts
+        cache['negative_prompts'] = negative_prompts
+        cache['model_type'] = shared.sd_model_type
+        if debug_enabled:
+            get_tokens('positive', prompts[0])
+            get_tokens('negative', negative_prompts[0])
         debug(f"Prompt encode: time={(time.time() - t0):.3f}")
         return