ai-forever/Wan2.1-T2V-14B-NABLA-0.6-STA-11-3-3

ai-forever

Texto a video

Modelo de generación de video a partir de texto basado en Wan 2.1 T2V 14B y ajustado con NABLA, un mecanismo de atención dispersa adaptativa a nivel de bloques, combinado con Sliding Tile Attention. Está diseñado para acelerar el entrenamiento y la inferencia de modelos Diffusion Transformer de video de alta resolución, manteniendo métricas de calidad comparables a la atención completa.

Como usar

Instalación y uso básico con Diffusers:
pip install -U diffusers transformers accelerate

import torch
from diffusers import DiffusionPipeline

# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained(
    "ai-forever/Wan2.1-T2V-14B-NABLA-0.6-STA-11-3-3",
    dtype=torch.bfloat16,
    device_map="cuda"
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]

Uso con los pesos ajustados NABLA 0.6 + STA 11-3-3:
python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B-NABLA-0.6-STA-11-3-3 --prompt {prompt} --frame_num {frame_num}

También puede usarse en modo zero-shot con NABLA, STA o una combinación de ambos:
python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt {prompt} --frame_num {frame_num} --sparse_algo nabla-{thr}

python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt {prompt} --frame_num {frame_num} --sparse_algo sta-{wT}-{wH}-{wW}

python generate.py --task t2v-14B --size 1280*720 --ckpt_dir ./Wan2.1-T2V-14B --prompt {prompt} --frame_num {frame_num} --sparse_algo nabla-{thr}_sta-{wT}-{wH}-{wW}

Requiere PyTorch 2.7 o superior para Flex Attention. NABLA admite resoluciones de video divisibles por 128; por ejemplo, 1280x720 se convierte a 1280x640.

Funcionalidades

Generación de video desde texto con arquitectura Diffusion Transformer basada en Wan 2.1 T2V 14B.
Usa NABLA con umbral 0.6 junto con STA configurado con ventanas wT=11, wH=3 y wW=3.
Atención dispersa dinámica que reduce el coste computacional mediante máscaras generadas a partir de Q/K comprimidos.
Compatible con PyTorch Flex Attention, sin kernels CUDA personalizados.
Orientado a video de alta resolución, incluyendo flujos 720p+ y contenido de mayor duración.
Reporta aceleraciones de hasta 2.7x en entrenamiento e inferencia según tamaño de video y configuración.
Licencia Apache 2.0.

Casos de uso

Generar videos desde prompts de texto con menor tiempo de inferencia que Wan 2.1 con atención completa.
Experimentar con atención dispersa adaptativa en modelos de video basados en DiT.
Crear pruebas y benchmarks de generación de video 720p+ con configuraciones NABLA y STA.
Investigar compromisos entre velocidad y calidad usando CLIP, VBench y evaluación humana.