.dstack.yml

type: task
name: wan22

repos:
  # Clones it to `/workflow` (the default working directory)
  - https://github.com/Wan-Video/Wan2.2.git

python: 3.12
nvcc: true

env:
  - PROMPT="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
  # Required for storing cache on a volume
  - UV_LINK_MODE=copy
commands:
  # Install flash-attn
  - |
    uv pip install torch
    uv pip install flash-attn --no-build-isolation 
  # Install dependencies
  - |
    uv pip install . decord librosa
    uv pip install "huggingface_hub[cli]"
    hf download Wan-AI/Wan2.2-T2V-A14B --local-dir /root/.cache/Wan2.2-T2V-A14B
  # Generate video
  - |
    if [ ${DSTACK_GPUS_NUM} -gt 1 ]; then
      torchrun \
        --nproc_per_node=${DSTACK_GPUS_NUM} \
        generate.py \
        --task t2v-A14B \
        --size 1280*720 \
        --ckpt_dir /root/.cache/Wan2.2-T2V-A14B \
        --dit_fsdp --t5_fsdp --ulysses_size ${DSTACK_GPUS_NUM} \
        --save_file ${DSTACK_RUN_NAME}.mp4 \
        --prompt "${PROMPT}"
    else
      python generate.py \
        --task t2v-A14B \
        --size 1280*720 \
        --ckpt_dir /root/.cache/Wan2.2-T2V-A14B \
        --offload_model True \
        --convert_model_dtype \
        --save_file ${DSTACK_RUN_NAME}.mp4 \
        --prompt "${PROMPT}"
    fi
  # Upload video
  - curl https://bashupload.com/ -T ./${DSTACK_RUN_NAME}.mp4

resources: 
  gpu:
    name: [H100, H200]
    count: 1..8
  disk: 300GB

# Change to on-demand for disabling spot
spot_policy: auto

volumes:
  # Cache pip packages and HF models
  - instance_path: /root/dstack-cache
    path: /root/.cache/
    optional: true