Hugging Face's logo

Diffusers documentation

AnyFlowFARTransformer3DModel

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→