Кэширование

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

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

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

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, чтобы хуки могли читать текущий шаг. Проверьте качество вывода на новой модели, так как настройки кэша, настроенные для одной модели, часто не переносятся на другую.