Hugging Face's logo

Diffusers documentation

AnyFlowFARTransformer3DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.43.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

AnyFlowFARTransformer3DModel

Причинный (FAR) 3D-трансформер, используемый в AnyFlowFARPipeline — вариант FAR для AnyFlow. См. страницу AnyFlowFARPipeline для информации о статье, авторах и опубликованных контрольных точках. Он расширяет базовую модель Wan2.1 версии v0.35.1 тремя дополнениями:

  1. Причинная блок-маска FAR с использованием torch.nn.attention.flex_attention, поддерживающая поштучную авторегрессионную генерацию, как описано в FAR.
  2. Встраивание патчей сжатых кадров (far_patch_embedding) для контекстных (уже сгенерированных) кадров, инициализируемое при создании из полноразрешающего patch_embedding методом трилинейной интерполяции.
  3. Встраивание потока с двумя временными шагами (аналогично AnyFlowTransformer3DModel) — каждый вызов forward использует в качестве условия и исходный временной шаг t, и целевой временной шаг r.

Расписание по умолчанию (chunk_partition) хранится в конфигурации модели; опубликованные чекпоинты NVIDIA AnyFlow-FAR используют [1, 3, 3, 3, 3, 3, 3, 2] для стандартной настройки на 81 кадр. Метод forward принимает переопределение chunk_partition для каждого вызова, поэтому тот же чекпоинт также обрабатывает другие конфигурации num_frames без дообучения.

from diffusers import AnyFlowFARTransformer3DModel

# Causal AnyFlow checkpoint (FAR):
transformer = AnyFlowFARTransformer3DModel.from_pretrained(
    "nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers", subfolder="transformer"
)

AnyFlowFARTransformer3DModel

класс diffusers.AnyFlowFARTransformer3DModel

< >

( patch_size: typing.Tuple[int] = (1, 2, 2)compressed_patch_size: typing.Tuple[int] = (1, 4, 4)full_chunk_limit: int = 3num_attention_heads: int = 40attention_head_dim: int = 128in_channels: int = 16out_channels: int = 16text_dim: int = 4096freq_dim: int = 256ffn_dim: int = 13824num_layers: int = 40cross_attn_norm: bool = Trueeps: float = 1e-06image_dim: typing.Optional[int] = Nonerope_max_seq_len: int = 1024gate_value: float = 0.25deltatime_type: str = 'r'chunk_partition: typing.Tuple[int, ...] = (1, 3, 3, 3, 3, 3, 3, 2) )

Параметры

  • patch_size (Tuple[int], по умолчанию (1, 2, 2)) — размеры 3D-патчей для фрагментов полного разрешения.
  • compressed_patch_size (Tuple[int], по умолчанию (1, 4, 4)) — Более крупные размеры патчей для FAR-сжатых (контекстных) фрагментов.
  • full_chunk_limit (int, по умолчанию 3) — Максимальное количество чанков полного разрешения, прежде чем более ранние чанки будут понижены до сжатого контекста FAR. В выпущенных контрольных точках используется значение 3.
  • num_attention_heads (int, по умолчанию 40) — Количество головок внимания.
  • attention_head_dim (int, по умолчанию 128) — Количество каналов в каждой голове.
  • in_channels (int, по умолчанию 16) — Количество каналов во входном латентном представлении.
  • out_channels (int, по умолчанию 16) — Количество каналов в выходном латентном представлении.
  • text_dim (int, по умолчанию 4096) — Размерность входных векторов для текстовых эмбеддингов (UMT5).
  • freq_dim (int, по умолчанию 256) — Размерность для синусоидных временных эмбеддингов.
  • ffn_dim (int, по умолчанию 13824) — Промежуточное измерение в полносвязной сети.
  • num_layers (int, по умолчанию 40) — Количество блоков трансформера.
  • cross_attn_norm (bool, по умолчанию True) — Включает нормализацию кросс-внимания.
  • eps (float, по умолчанию 1e-6) — Значение эпсилона для слоёв нормализации.
  • image_dim (Optional[int], необязательный, по умолчанию None) — Размерность эмбеддинга изображения для I2V-условия.
  • rope_max_seq_len (int, по умолчанию 1024) — Максимальная длина последовательности, используемая для предварительного вычисления частот вращательных позиционных кодировок.
  • gate_value (float, по умолчанию 0.25) — Коэффициент смешивания между эмбеддингами исходного временного шага и дельта-шага.
  • deltatime_type (str, по умолчанию ‘r’) — Либо "r" (delta — целевой шаг), либо "t-r" (delta — абсолютный интервал).
  • chunk_partition (Tuple[int, …], по умолчанию (1, 3, 3, 3, 3, 3, 3, 2)) — Количество кадров в каждом чанке по умолчанию, используемое конвейером. Опубликованные чекпоинты NVIDIA AnyFlow-FAR рассчитаны на num_frames=81 (21 латентный кадр при временном шаге VAE, равном 4) с разбиением 1 + 3*6 + 2. Для другого num_frames необходимо передать соответствующее значение chunk_partition в AnyFlowFARPipeline.call() (и аналогично в forward).

Причинный (FAR) 3D-трансформер для выборки карт потоков AnyFlow с поштучной авто регрессивной генерацией.

Расширяет базовую модель Wan2.1 версии v0.35.1:

  • Причинная блок-маска FAR с помощью torch.nn.attention.flex_attention, поддерживающая поштучную авторегрессионную генерацию (FAR).
  • Встраивание патчей сжатых кадров far_patch_embedding для контекстных (уже сгенерированных) кадров, инициализируемое из patch_embedding с помощью трилинейной интерполяции, чтобы только что сконструированная модель уже находилась в разумной стартовой точке даже до дообучения с помощью LoRA.
  • Встраивание потока с двойным шагом времени для выборки с любым шагом (аналогично AnyFlowTransformer3DModel).

Для простого двунаправленного T2V используйте AnyFlowTransformer3DModel — этот вариант пропускает причинное маскирование FAR и far_patch_embedding и примерно на 5–10% меньше по размеру.

build_attention_mask

< >

( chunk_partition: typing.List[int]height: intwidth: inthas_clean_context: bool = Falsedevice: typing.Optional[torch.device] = Nonemode: str = 'train' ) → ~torch.nn.attention.flex_attention.BlockMask

Параметры

  • chunk_partition — количество кадров на каждый чанк (сумма должна равняться общему числу латентных кадров).
  • height, width — пространственные размеры латентного представления.
  • has_clean_context — True, если clean_hidden_states будут передаваться через forward (обучение V2V/I2V); только наличие этого флага влияет на структуру маски.
  • device — устройство для результирующего BlockMask. Маска не перемещается автоматически с помощью device_map="auto"; создавайте её на том же устройстве, где будут находиться входные данные трансформера.
  • mode — "train" (соответствует _forward_train) или "cache" (соответствует _forward_cache). Авто регрессионный путь _forward_inference обращается через KV-кэш и не имеет здесь режима.

Возвращает

~torch.nn.attention.flex_attention.BlockMask

causal mask, охватывающая раскладку FAR, дополненная нулями до кратного 128 значения по размерности последовательности (требование к размеру блока BlockMask).

Исключения

ValueError

  • ValueError — если mode не является ни "train", ни "cache".

Предварительно создайте каузальную ~torch.nn.attention.flex_attention.BlockMask вне forward.

Передайте результат через аргумент attention_mask метода forward, чтобы сделать весь трансформер совместимым с torch.compile(fullgraph=True). Без заранее созданной маски forward переходит на внутренний путь её построения — этот путь использует flex_attention.create_block_mask(_compile=False) и нарушает граф компиляции.

forward

< >

( hidden_states: Tensortimestep: Tensorr_timestep: Tensorencoder_hidden_states: Tensorchunk_partition: typing.List[int]encoder_hidden_states_image: typing.Optional[torch.Tensor] = Noneclean_hidden_states: typing.Optional[torch.Tensor] = Noneclean_timestep: typing.Optional[torch.Tensor] = Nonekv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = Nonekv_cache_flag: typing.Optional[typing.Dict[str, typing.Any]] = Noneattention_mask: typing.Optional[torch.nn.attention.flex_attention.BlockMask] = Noneattention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = Nonereturn_dict: bool = True ) → [~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] or tuple

Параметры

  • hidden_states (torch.Tensor) — Латентный вход с формой (B, F, C, H, W).
  • timestep (torch.Tensor) — Исходный (более шумный) шаг потока t.
  • r_timestep (torch.Tensor) — Целевой (более чистый) шаг временной шкалы flow-map r.
  • encoder_hidden_states (torch.Tensor) — векторные представления текста UMT5.
  • chunk_partition (List[int]) — Количество кадров в каждом чанке; сумма должна совпадать с числом латентных кадров в hidden_states.
  • encoder_hidden_states_image (torch.Tensor, optional) — вложение изображения для I2V; при наличии конкатенируется перед текстовыми токенами.
  • clean_hidden_states (torch.Tensor, optional) — Чистые (без шума) кадры, используемые в качестве условия для тренировочного rollout.
  • clean_timestep (torch.Tensor, optional) — Шаги времени для чистых кадров условий в тренировочном rollout.
  • kv_cache (List[Dict[str, torch.Tensor]], optional) — Кэш KV по блокам для авторегрессионного вывода. None выбирает режим обучения.
  • kv_cache_flag (Dict[str, Any], необязательный) — метаданные KV-кэша (например, флаг is_cache_step и количество токенов).
  • attention_mask (BlockMask, optional) — Готовая каузальная маска, как правило, создаваемая с помощью build_attention_mask. Используется в путях обучения и префилла KV-кэша; путь авторегрессионного вывода обращается к KV-кэшу и не использует полную маску. Если None, пути обучения / кэша создают маску внутренне; этот запасной вариант несовместим с компиляцией (базовая функция flex_attention.create_block_mask нарушает граф при fullgraph=True), поэтому передавайте готовую маску, если оборачиваете forward в torch.compile.
  • attention_kwargs (dict, optional) — Передаётся в обработчики внимания.
  • return_dict (bool, необязательный, по умолчанию True) — Если False, возвращает позиционные кортежи вместо объекта данных класса вывода.

Возвращает

[~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] или tuple

Если return_dict равно False, возвращается обычный tuple. В противном случае, при обучении с причинной раскруткой (kv_cache is None) возвращается [~models.transformer_2d.Transformer2DModelOutput], а при префилле кэша и авторегрессионном выводе — [AnyFlowFARTransformerOutput].

Прямой проход FAR с причинной зависимостью. Направляет выполнение на один из трёх внутренних путей:

  • kv_cache is None → rollout для обучения с каузальностью (возвращает Transformer2DModelOutput).
  • kv_cache is not None и kv_cache_flag["is_cache_step"] → cache-prefill (возвращает AnyFlowFARTransformerOutput с sample=None).
  • В противном случае → шаг авто-регрессивного вывода (возвращает AnyFlowFARTransformerOutput).

AnyFlowFARTransformerOutput

класс diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput

< >

( sample: typing.Optional[torch.Tensor] = Nonekv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None )

Параметры

  • sample (torch.Tensor или None) — Предсказанная цель денойзинга для авто-регрессивного фрагмента. None для пути предварительного заполнения кэша, который только записывает KV-кэш и не производит пригодного для использования сэмпла.
  • kv_cache (list[dict[str, torch.Tensor]], optional) — Состояние кэша KV для каждого блока, используемое последующими шагами авторегрессивного генерирования.

Dataclass для выходных данных каузальных forward-путей модели AnyFlowFARTransformer3DModel.

Update on GitHub

←AllegroTransformer3DModel AnyFlowTransformer3DModel→