diffusers/main/ru/optimization/cache.md

Кэширование

Кэширование повторно использует выходные данные промежуточных слоёв между шагами денойзинга, чтобы ускорить инференс. Этот метод требует больше памяти и не нуждается в обучении. Включите метод с конфигурацией на трансформере, который поддерживает кэширование.

Выберите метод кэширования

Выберите метод в зависимости от того, сколько параметров конфигурации вы будете задавать, и от требуемой точности. Каждый метод по-разному балансирует между скоростью, потреблением памяти и степенью соответствия результатов запуску без кэша. Перед тем как полагаться на какой-либо метод, сравните результаты с кэшем и без него на собственных запросах.

Метод	Когда использовать	Примечания
Кэш текстовых KV	Изображение NucleusMoE, повторное использование проекций K/V текста между шагами	В настоящее время хуки блокируют образы NucleusMoE
SeaCache	Генерация видео Cosmos 3	Настройки часто не переносятся между моделями
FirstBlockCache	Нужен один основной параметр для управления соотношением скорости и качества в зарегистрированном трансформере	Один параметр threshold управляет частотой пропуска шагов
MagCache	Есть ли у вас коэффициенты амплитуд для вашей контрольной точки и планировщика, или сначала будет выполнена калибровка	Соотношения зависят от контрольной точки и планировщика
TaylorSeer	Хотите предсказывать поздние активации по ранним шагам	Более высокое значение max_order требует больше памяти
PAB	Видео, готовность настраивать повторное использование внимания (диапазоны пропусков блоков и временных шагов для каждого типа внимания)	Диапазоны пропуска требуют настройки для каждой модели
FasterCache	Как PAB, плюс опциональное пропускание ветвей CFG	Экспериментальный

Pyramid Attention Broadcast

Pyramid Attention Broadcast (PAB) приближает внимание на шагах денойзинга, повторно используя выходы внимания для некоторых блоков и временных шагов вместо пересчёта на каждом шаге. Конфигурация разделяет типы внимания (пространственное, временное, кросс-внимание), если модель их поддерживает. Не каждая модель видео предоставляет все три типа, поэтому задавайте только диапазоны, соответствующие имеющимся у вас блокам.

Каждый тип использует *_attention_block_skip_range (как часто пересчитывать или переиспользовать значения внутри окна) и *_attention_timestep_skip_range (какие шаги денойзинга можно пропускать). Необходимо передать current_timestep_callback, чтобы хук мог читать текущий шаг пайплайна. Более широкие или агрессивные пропуски обычно означают большую скорость и больший риск потери качества.

Передайте PyramidAttentionBroadcastConfig, чтобы включить его.

import torch
from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

config = PyramidAttentionBroadcastConfig(
    spatial_attention_block_skip_range=2,
    spatial_attention_timestep_skip_range=(100, 800),
    current_timestep_callback=lambda: pipe.current_timestep,
)
pipe.transformer.enable_cache(config)

FasterCache

FasterCache кэширует и повторно использует признаки внимания по аналогии с PAB. Кроме того, он может пропускать необусловленную ветвь при классификаторно-свободном управлении и оценивать её по обусловленной ветви, если последовательные латенты достаточно избыточны.

Передайте FasterCacheConfig, чтобы включить его. Как и в случае с PAB, задайте *_attention_block_skip_range и *_attention_timestep_skip_range для имеющихся у вас типов внимания, а также опции пропуска ветвей CFG, если они вам нужны.

import torch
from diffusers import CogVideoXPipeline, FasterCacheConfig

pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

config = FasterCacheConfig(
    spatial_attention_block_skip_range=2,
    spatial_attention_timestep_skip_range=(-1, 681),
    current_timestep_callback=lambda: pipe.current_timestep,
    attention_weight_callback=lambda _: 0.3,
    unconditional_batch_skip_range=5,
    unconditional_batch_timestep_skip_range=(-1, 781),
    tensor_format="BFCHW",
)
pipe.transformer.enable_cache(config)

SeaCache

SeaCache сравнивает индикаторы Spectral Evolution Aware (SEA) между последовательными шагами шумоподавления. Если накопленное изменение остаётся ниже порога, метод пропускает стек блоков трансформера и предсказывает выход на основе кэшированных остатков. Метод является приближённым и предназначен для генерации видео.

Встроенные адаптеры для SeaCache включают:

• Cosmos 3 — это основная оптимизированная и протестированная интеграция.
• Wan T2V использует универсальный путь с повторяющимися блоками в качестве демонстрации того, как передавать необработанные визуальные латенты в SeaCache. Те же параметры кэша могут не переноситься на Wan или другие варианты Wan.

Включите SeaCache для трансформера. Цикл денойзинга Cosmos 3 добавляет шаг планировщика, сигму и номер шага к каждому cache_context, поэтому дополнительные параметры не требуются.

from diffusers import Cosmos3OmniPipeline, SeaCacheConfig

pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano")
pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))

SeaCache может изменять выходные данные. Вызывайте pipe.transformer.disable_cache(), если вам необходимо, чтобы на каждом шаге выполнялся полный трансформер. Тот же вызов для включения работает с Cosmos3OmniPipeline, Cosmos3OmniModularPipeline и Cosmos3DistilledModularPipeline.

Чтобы использовать SeaCache с другим видео-трансформером, следуйте инструкции Добавление кэширования в новую модель. SeaCache также требует step_index, sigma и num_inference_steps в каждом cache_context, а также raw_vision_callback, если встроенный адаптер отсутствует. Настраивайте параметры для каждой модели и планировщика.

FirstBlockCache

FirstBlockCacheConfig проверяет, насколько сильно ранние слои денойзера изменяются от одного временного шага к другому. Если изменение незначительно, модель пропускает дорогостоящие последующие слои и повторно использует предыдущий вывод.

Включите его через enable_cache, чтобы disable_cache и is_cache_enabled оставались синхронизированными. Значение по умолчанию для threshold — 0.05. Более высокое значение, например 0.2, приводит к более частому пропуску для дополнительной скорости, но качество генерации может снизиться.

import torch
from diffusers import DiffusionPipeline, FirstBlockCacheConfig

pipe = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16
)
pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=0.2))

Кэш TaylorSeer

TaylorSeer Cache ускоряет инференс диффузионных моделей с помощью разложений в ряд Тейлора по шагам денойзинга. Он предсказывает активации для более поздних шагов на основе более ранних и повторно использует эти предсказания на протяжении нескольких шагов, что снижает объём полных вычислений трансформера.

• cache_interval: количество шагов, в течение которых переиспользуются кэшированные результаты, прежде чем выполняется полный прямой проход
• disable_cache_before_step: Начальные шаги, использующие полные вычисления для сбора данных для приближений
• max_order: Более высокие порядки разложения Тейлора могут быть точнее, но требуют больше памяти. Оставьте это значение равным 1, если у вас нет веских причин его менять.

import torch
from diffusers import FluxPipeline, TaylorSeerCacheConfig

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

config = TaylorSeerCacheConfig(
    cache_interval=5,
    max_order=1,
    disable_cache_before_step=10,
    taylor_factors_dtype=torch.bfloat16,
)
pipe.transformer.enable_cache(config)

MagCache

MagCache пропускает блоки трансформера на основе величины обновления остаточного соединения. Величины обновления предсказуемо затухают в процессе денойзинга, а MagCache отслеживает бюджет ошибки на основе заранее вычисленных отношений величин (mag_ratios), чтобы определить, когда повторное использование безопасно. Эти отношения зависят от конкретной контрольной точки и планировщика. Отношения, полученные для большого числа шагов, можно интерполировать для меньшего числа шагов.

Встроенные FLUX_MAG_RATIOS были измерены для FLUX.1. Другим моделям, включая Flux2 и Flux2 Klein, требуется собственный прогон калибровки.

MagCache следует двум шагам:

1. Калибровка: выполните один проход инференса с calibrate=True. Хук измеряет амплитуды остаточных сигналов и выводит рассчитанные соотношения.
2. Вывод: отключите кэш калибровки, затем передайте эти соотношения в MagCacheConfig для ускорения.

Classifier-free guidance может влиять на калибровку. Пайплайны, использующие true CFG с последовательными контекстами, например Flux, когда true_cfg_scale > 1, входят в cache_context("cond") и cache_context("uncond") по отдельности. Калибровка может выводить один массив на каждый контекст, но в большинстве случаев следует использовать условный массив. Пайплайны, которые батчат CFG, конкатенируя условные и безусловные входные данные (например, CogVideoX), формируют единый совместный массив, который можно использовать напрямую.

import torch
from diffusers import FluxPipeline, MagCacheConfig
from diffusers.hooks.mag_cache import FLUX_MAG_RATIOS

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"

# 1. Calibration Step
# Run full inference to measure model behavior.
calib_config = MagCacheConfig(calibrate=True, num_inference_steps=4)
pipe.transformer.enable_cache(calib_config)

# Run a prompt to trigger calibration
pipe("A cat playing chess", num_inference_steps=4)
# Prints: [MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):

# 2. Inference Step
# Disable calibration hooks before enabling MagCache for inference.
pipe.transformer.disable_cache()

# Apply ratios from calibration, or use the Flux defaults:
# mag_ratios=FLUX_MAG_RATIOS
mag_config = MagCacheConfig(
    mag_ratios=[1.0, 1.37, 0.97, 0.87],
    num_inference_steps=4
)

pipe.transformer.enable_cache(mag_config)

image = pipe("A cat playing chess", num_inference_steps=4).images[0]

Кэш текста KV

TextKVCacheConfig вычисляет проекции ключей и значений для текста один раз и переиспользует их на каждом шаге денойзинга, вместо того чтобы пересчитывать их заново. apply_text_kv_cache() на данный момент подключается только к NucleusMoEImageTransformerBlock. Включите его с помощью enable_cache.

import torch
from diffusers import NucleusMoEImagePipeline, TextKVCacheConfig

pipe = NucleusMoEImagePipeline.from_pretrained(
    "NucleusAI/NucleusMoE-Image", dtype=torch.bfloat16
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

pipe.transformer.enable_cache(TextKVCacheConfig())

image = pipe("A cat holding a sign that says hello world", num_inference_steps=50).images[0]

Добавление кэширования в новую модель

Методы кэширования подключаются к модели через хуки, которые зависят от класса модели, её трансформерных блоков и цикла денуазинга пайплайна. Настройте все три компонента перед вызовом enable_cache для новой модели.

1. Наследуйте CacheMixin в классе трансформера. Он добавляет enable_cache, disable_cache и cache_context.

2. Зарегистрируйте выходную раскладку блока трансформера в _register_transformer_blocks_metadata в hooks/_helpers.py. Методы, которые пропускают или повторно используют целые блоки, такие как FirstBlockCache, MagCache и SeaCache, используют её для поиска скрытых состояний в выходах блока.

   TransformerBlockRegistry.register(
       model_class=MyTransformerBlock,
       metadata=TransformerBlockMetadata(
           return_hidden_states_index=0,
           return_encoder_hidden_states_index=None,
       ),
   )
   

3. Оберните каждый вызов трансформера в цикле денойзинга пайплайна с помощью cache_context. Используйте разные имена, например "cond" и "uncond", если classifier-free guidance выполняет условную и безусловную ветви как отдельные вызовы, чтобы каждая ветвь сохраняла собственное состояние кэша.

   with self.transformer.cache_context("cond"):
       noise_pred = self.transformer(hidden_states=latents, timestep=timestep, ...)[0]
   

PAB и FasterCache также требуют current_timestep_callback, чтобы хуки могли читать текущий шаг. Проверьте качество вывода на новой модели, так как настройки кэша, настроенные для одной модели, часто не переносятся на другую.