mirror of
https://github.com/vladmandic/automatic
synced 2026-09-09 06:18:43 +02:00
CalamitousFelicitousness
c904585f87
0dbe27c37 keyed the prompt cache on the negative prompt and the cfg flag
and hashed both prompts to tuples for it. The positive prompt was hashed
before the list check that builds the key, so a list prompt, which is what
every pipeline passes, was never stored or found and the text encoder ran
on every generation. Only the negative prompt needs hashing.
115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
import os
|
|
import time
|
|
from modules import shared, errors, timer, sd_models
|
|
from modules.logger import log
|
|
from modules.attention import context as attention_context
|
|
|
|
|
|
class PromptCache:
|
|
def __init__(self):
|
|
self.cache = {}
|
|
self.id = None
|
|
self.max = 16
|
|
|
|
@staticmethod
|
|
def _hashable(val):
|
|
if isinstance(val, list):
|
|
return tuple(val)
|
|
return val
|
|
|
|
def get(self, prompt, negative_prompt=None, cfg_enabled=None):
|
|
if self.id != id(shared.sd_model):
|
|
self.cache.clear()
|
|
self.id = id(shared.sd_model)
|
|
log.debug(f'Encode: prompt cache activate id={self.id} depth={len(self.cache)}')
|
|
negative_prompt = self._hashable(negative_prompt)
|
|
if (isinstance(prompt, list) and len(prompt) == 1 and isinstance(prompt[0], str)):
|
|
cached = self.cache.get((prompt[0], negative_prompt, cfg_enabled), None)
|
|
elif isinstance(prompt, str):
|
|
cached = self.cache.get((prompt, negative_prompt, cfg_enabled), None)
|
|
else:
|
|
cached = None
|
|
if cached:
|
|
if isinstance(prompt, list):
|
|
log.debug(f'Encode: prompt={prompt} cache={len(self.cache)} hit')
|
|
else:
|
|
log.debug(f'Encode: prompt="{prompt}" cache={len(self.cache)} hit')
|
|
return cached
|
|
|
|
def set(self, prompt, encoded, negative_prompt=None, cfg_enabled=None):
|
|
if len(self.cache) >= self.max:
|
|
oldest_key = next(iter(self.cache))
|
|
del self.cache[oldest_key]
|
|
negative_prompt = self._hashable(negative_prompt)
|
|
if (isinstance(prompt, list) and len(prompt) == 1 and isinstance(prompt[0], str)):
|
|
self.cache[(prompt[0], negative_prompt, cfg_enabled)] = encoded
|
|
elif isinstance(prompt, str):
|
|
self.cache[(prompt, negative_prompt, cfg_enabled)] = encoded
|
|
|
|
|
|
prompt_cache = PromptCache()
|
|
|
|
|
|
def hijack_encode_prompt(*args, **kwargs):
|
|
jobid = shared.state.begin('TE Encode')
|
|
t0 = time.time()
|
|
if 'max_sequence_length' in kwargs and kwargs['max_sequence_length'] is not None:
|
|
kwargs['max_sequence_length'] = max(kwargs['max_sequence_length'], int(os.environ.get('MAX_SEQUENCE_LENGTH', 256)))
|
|
res = None
|
|
try:
|
|
args_copy = list(args)
|
|
patch_prompt = False
|
|
prompt = kwargs.get('prompt', None)
|
|
if prompt is None and len(args_copy) > 0:
|
|
prompt = args_copy[0]
|
|
patch_prompt = True
|
|
prompt = [p.strip(", \n") if isinstance(p, str) else p for p in prompt] if isinstance(prompt, list) else prompt
|
|
res = prompt
|
|
|
|
if hasattr(shared.sd_model, 'before_prompt_encode'):
|
|
log.debug('Encode: op=before')
|
|
res = shared.sd_model.before_prompt_encode(prompt)
|
|
if patch_prompt:
|
|
args_copy[0] = res
|
|
|
|
# cache key must include cfg-affecting kwargs since encode_prompt output (e.g. negative_prompt_embeds) depends on them
|
|
negative_prompt = kwargs.get('negative_prompt', None)
|
|
cfg_enabled = kwargs.get('do_classifier_free_guidance', None)
|
|
cached = prompt_cache.get(prompt, negative_prompt, cfg_enabled)
|
|
if cached is not None:
|
|
res = cached
|
|
else:
|
|
if isinstance(prompt, list):
|
|
log.debug(f'Encode: prompt={prompt} hijack=True')
|
|
else:
|
|
log.debug(f'Encode: prompt="{prompt}" hijack=True')
|
|
with attention_context.role('te'):
|
|
if hasattr(shared.sd_model, 'orig_encode_prompt'):
|
|
res = shared.sd_model.orig_encode_prompt(*args_copy, **kwargs)
|
|
else:
|
|
res = shared.sd_model.encode_prompt(*args_copy, **kwargs)
|
|
prompt_cache.set(prompt, res, negative_prompt, cfg_enabled)
|
|
|
|
if hasattr(shared.sd_model, 'after_prompt_encode'):
|
|
log.debug('Encode: op=after')
|
|
res = shared.sd_model.after_prompt_encode(res)
|
|
|
|
except Exception as e:
|
|
log.error(f'Encode prompt: {e}')
|
|
errors.display(e, 'Encode prompt')
|
|
t1 = time.time()
|
|
timer.process.add('te', t1-t0)
|
|
if t1 - t0 > 10:
|
|
log.warning(f'Encode: time={t1-t0:.3f} long encode prompt')
|
|
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
|
|
shared.state.end(jobid)
|
|
# from modules import memstats
|
|
# log.debug(f'Encode: memory={memstats.memory_stats()}')
|
|
return res
|
|
|
|
|
|
def init_hijack(pipe):
|
|
if (pipe is not None) and not hasattr(pipe, 'orig_encode_prompt') and hasattr(pipe, 'encode_prompt'):
|
|
pipe.orig_encode_prompt = pipe.encode_prompt
|
|
pipe.encode_prompt = hijack_encode_prompt
|