File size: 1,915 Bytes
371818e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """Bounded reuse of exact encoder outputs for repeated text-only prompts."""
from collections import OrderedDict
import torch
def cache_prompt_embeddings(pipe, capacity=8):
"""Use only for an immutable inference pipeline; clear after changing the encoder.
Image edits and caller-supplied embeddings always use the normal encoder path.
Cached tensors stay on the original device. No embedding precision is changed.
"""
if hasattr(pipe, '_image21_prompt_cache'):
return pipe
cache = OrderedDict()
original = pipe.encode_prompt
pipe._image21_prompt_cache = cache
pipe._image21_prompt_cache_stats = {'hits': 0, 'misses': 0}
def encode(*args, **kwargs):
prompt = kwargs.get('prompt')
safe = (not args and isinstance(prompt, str) and kwargs.get('image') is None
and kwargs.get('prompt_embeds') is None
and kwargs.get('prompt_embeds_mask') is None
and kwargs.get('image_pad_mask') is None
and kwargs.get('num_images_per_prompt', 1) == 1
and not torch.is_grad_enabled() and not torch.is_autocast_enabled()
and set(kwargs) <= {'prompt','image','device','num_images_per_prompt',
'prompt_embeds','prompt_embeds_mask','image_pad_mask'})
if not safe:
return original(*args, **kwargs)
key = (prompt, str(kwargs.get('device') or pipe._execution_device), str(pipe.text_encoder.dtype))
if key in cache:
cache.move_to_end(key)
pipe._image21_prompt_cache_stats['hits'] += 1
return cache[key]
result = original(*args, **kwargs)
pipe._image21_prompt_cache_stats['misses'] += 1
cache[key] = result
if len(cache) > capacity:
cache.popitem(last=False)
return result
pipe.encode_prompt = encode
return pipe
|