prompt_cache.py

"""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