Wan2.1-T2V-14B-NABLA-0.5-STA-11-5-5

Instala las dependencias y carga el modelo con Diffusers:

pip install -U diffusers transformers accelerate

import torch
from diffusers import DiffusionPipeline

# Cambia a "mps" para dispositivos Apple
pipe = DiffusionPipeline.from_pretrained(
    "ai-forever/Wan2.1-T2V-14B-NABLA-0.5-STA-11-5-5",
    dtype=torch.bfloat16,
    device_map="cuda"
)

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

Para generar vídeo con NABLA y STA mediante el repositorio Wan2.1-NABLA:

python generate.py \
  --task t2v-14B \
  --size 1280*720 \
  --ckpt_dir ./Wan2.1-T2V-14B-NABLA-0.5-STA-11-5-5 \
  --prompt {prompt} \
  --frame_num {frame_num}

También puede aplicarse NABLA junto con STA en modo zero-shot sobre el modelo Wan2.1 original:

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}

Se necesita PyTorch 2.7 o posterior para disponer de FlexAttention.