Diffusers documentation
WanAnimate2Transformer3DModel
WanAnimate2Transformer3DModel
Модель Diffusion Transformer для 3D-видеоподобных данных, используемая в Wan-Animate-2 от команды Alibaba Wan. Она анимирует изображение персонажа с помощью движения из управляющего видео посредством механизма контекстной ссылки: каждый сегмент сначала выполняет проход ссылки (kv_cache_mode="extract"), кэшируя справочные K/V для каждого слоя, после чего проходы денойзинга (kv_cache_mode="cached") совместно обращаются к токенам генерации и кэшированным справочным токенам через гибкую BlockMask.
Модель можно загрузить с помощью следующего фрагмента кода.
from diffusers import WanAnimate2Transformer3DModel
transformer = WanAnimate2Transformer3DModel.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Diffusers", subfolder="transformer", dtype=torch.bfloat16)WanAnimate2Transformer3DModel
класс diffusers.WanAnimate2Transformer3DModel
< source >( patch_size: tuple = (1, 2, 2)text_len: int = 512in_dim: int = 36dim: int = 5120ffn_dim: int = 13824freq_dim: int = 256text_dim: int = 4096out_dim: int = 16num_heads: int = 40num_layers: int = 40cross_attn_norm: bool = Trueeps: float = 1e-06use_img_emb: bool = Truerefer_offset_t: int = 1refer_offset_h: int = 0refer_offset_w: int = -1refer_stride: int = 1 )
Параметры
- patch_size (tuple[int], по умолчанию (1, 2, 2)) — размеры 3D-патчей для встраивания видео (t_patch, h_patch, w_patch).
- text_len (int, по умолчанию 512) — Фиксированная длина для текстовых эмбеддингов.
- in_dim (int, по умолчанию 36) — Количество каналов на входе (2 * latent_channels + 4 для маскирующего канала).
- dim (int, по умолчанию 5120) — Количество каналов в трансформере.
- ffn_dim (int, по умолчанию 13824) — Промежуточное измерение в полносвязной сети.
- freq_dim (int, по умолчанию 256) — Размерность для синусоидных временных эмбеддингов.
- text_dim (int, по умолчанию 4096) — Размерность входных векторов текстовых эмбеддингов.
- out_dim (int, по умолчанию 16) — Количество каналов на выходе.
- num_heads (int, по умолчанию 40) — Количество головок внимания.
- num_layers (int, по умолчанию 40) — Количество используемых блоков трансформера.
- cross_attn_norm (bool, по умолчанию True) — Включает нормализацию кросс-внимания.
- eps (float, по умолчанию 1e-6) — Значение эпсилона для слоёв нормализации.
- use_img_emb (bool, по умолчанию True) — Использовать ли CLIP image embedding.
- refer_offset_t (int, по умолчанию 1) — смещение RoPE для временного измерения эталонного сигнала.
- refer_offset_h (int, по умолчанию 0) — смещение RoPE для измерения высоты эталонного изображения.
- refer_offset_w (int, по умолчанию -1) — смещение RoPE для ширины опорного изображения. Значение -1 означает использование размера сетки генерации.
- refer_stride (int, по умолчанию 1) — Шаг применения RoPE для эталонного представления.
Модель Transformer для видео-подобных данных, используемая в модели Wan-Animate-2.
Wan-Animate-2 использует механизм внимания в контексте с кэшем KV: сначала кодируется видео-референс (kv_cache_mode="extract") для заполнения [WanAnimate2KVCache], затем на каждом шаге денойзинга (kv_cache_mode="cached") внимание совместно обрабатывает токены генерации и кэшированные K/V референса через flex BlockMask. Таким образом, само-внимание генерации выполняется только на бэкенде внимания flex; все остальные виды внимания в модели работают на любом бэкенде.
forward
< source >( hidden_states: list[torch.Tensor]timestep: Tensorencoder_hidden_states: list[torch.Tensor]condition_latents: list[torch.Tensor]kv_cache: WanAnimate2KVCachekv_cache_mode: strseq_len: intencoder_hidden_states_image: typing.Optional[torch.Tensor] = Noneoffset_grid_sizes: typing.Optional[torch.Tensor] = Nonereference_grid_sizes: typing.Optional[torch.Tensor] = Noneorigin_len: int | None = Noneorigin_area: list[int] | None = Noneis_uncondtion: bool = Falsereturn_dict: bool = True ) → Transformer2DModelOutput or tuple(list[torch.Tensor])
Параметры
- hidden_states (
list[torch.Tensor]) — Латенты для этого прохода — эталонные латенты, еслиkv_cache_mode="extract", или латенты шумной генерации, еслиkv_cache_mode="cached". - timestep (
torch.Tensor) — Шаг денойзинга. Игнорируется приkv_cache_mode="extract", который использует фиксированный шаг 1. - encoder_hidden_states (
list[torch.Tensor]) — Векторные представления текста для этого прохода. - condition_latents (
list[torch.Tensor]) — Условные латенты, конкатенированные сhidden_statesперед patch embedding. - kv_cache (
WanAnimate2KVCache) — записывается приkv_cache_mode="extract", читается при"cached". - kv_cache_mode (
str) —"extract"выполняет эталонный проход и заполняетkv_cache;"cached"выполняет шаг денойзинга по кэшированным эталонным токенам. - seq_len (
int) — Количество токенов, которое должен содержать каждый образец после встраивания патчей. - encoder_hidden_states_image (
torch.Tensor, необязательный) — встраивания изображений CLIP, используются, когда модель настроена сuse_img_emb. - offset_grid_sizes (
torch.Tensor, необязательный) — Сетка патчей для опорных латентов, используется для разрешения значенийrefer_offset_*, равных -1. Обязателен приkv_cache_mode="extract"; при"cached"та же сетка описывается параметромreference_grid_sizesи используется вместо него. - reference_grid_sizes (
torch.Tensor, необязательный) — Сетка патчей для опорных латентов, используется для опорных вращательных эмбеддингов. Обязателен приkv_cache_mode="cached". - origin_len (
int, необязательный) — Количество кадров полного видео, на основе которого строится маска блока в контексте. Обязателен приkv_cache_mode="cached". - origin_area (
list[int], необязательный) — Пространственный размер[height, width]полного видео, на основе которого строится маска блока в контексте. Обязателен приkv_cache_mode="cached". - is_uncondtion (
bool, необязательный) — Указывает, является ли это безусловной ветвью classifier-free guidance. - return_dict (
bool, необязательный, по умолчаниюTrue) — Определяет, следует ли возвращать~models.transformer_2d.Transformer2DModelOutputвместо обычного кортежа.
Возвращает
Transformer2DModelOutput или tuple(list[torch.Tensor])
Предсказанный образец на каждый входной латентный, без патчирования; обычный кортеж, если return_dict равно False.
Transformer2DModelOutput
класс diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Параметры
- sample (
torch.Tensorформы(batch_size, num_channels, height, width)или(batch size, num_vector_embeds - 1, num_latent_pixels), если Transformer2DModel дискретный) — Выходные скрытые состояния, обусловленные входомencoder_hidden_states. Если модель дискретная, возвращает распределения вероятностей для латентных пикселей без шума.
Вывод Transformer2DModel.