Hugging Face's logo

Agnes-AI
/
Agnes-Video-2.0-MoE

audio-video-generation
multimodal
sparse-moe
video-generation
Model card Files
xet
 Community
Agnes-Video-2.0-MoE / inference_code /infer.sh
Agnes-AI's picture
Agnes-AI
Add TIRx inference kernels and B200 denoising timings
dceb455 verified
Raw History Blame Contribute Delete
2.47 kB
#!/usr/bin/env bash
# Set OMNIVAE_CODE_ROOT before running.
#
# Standard inference:
# bash infer.sh --config-load-path configs/config_template.json --generate \
# --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
#
# TIRx kernels on the unpruned SparseMoE checkpoint:
# TIRX=1 bash infer.sh --config-load-path configs/config_template.json --generate \
# --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
#
# TIRx kernels on an elastic EENP checkpoint (SwiGLU7, QK, pre-norm, gather/scatter):
# TIRX_ELASTIC=1 bash infer.sh --config-load-path configs/config_template.json --generate \
# --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
#
# Elastic width without TIRx kernels:
# ELASTIC=1 bash infer.sh --config-load-path configs/config_template.json --generate \
# --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
set -euo pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export PYTHONPATH="${REPO}:${OMNIVAE_CODE_ROOT:-}${PYTHONPATH:+:${PYTHONPATH}}"
export OMNIVAE_CODE_ROOT="${OMNIVAE_CODE_ROOT:-}"
export OMNIVAE_AUDIO_CKPT_DIR="${OMNIVAE_AUDIO_CKPT_DIR:-}"
export DAVINCI_DISABLE_MAGI_COMPILE="${DAVINCI_DISABLE_MAGI_COMPILE:-1}"
export MAGI_DISABLE_COMPILE="${MAGI_DISABLE_COMPILE:-1}"
export TORCHDYNAMO_DISABLE="${TORCHDYNAMO_DISABLE:-1}"
export PYTHONUNBUFFERED=1
cd "$REPO"
if [[ "${TIRX_ELASTIC:-0}" == "1" ]]; then
export MAGI_TIRX_ELASTIC_IO="${MAGI_TIRX_ELASTIC_IO:-1}"
export MAGI_TIRX_QK_ROPE="${MAGI_TIRX_QK_ROPE:-1}"
export MAGI_TIRX_PRENORM="${MAGI_TIRX_PRENORM:-1}"
export MAGI_TIRX_SWIGLU7="${MAGI_TIRX_SWIGLU7:-1}"
export MAGI_TIRX_PREWARM="${MAGI_TIRX_PREWARM:-1}"
exec python infer_tirx_elastic_io.py "$@"
fi
if [[ "${TIRX:-0}" == "1" ]]; then
export MAGI_TIRX_SCATTER="${MAGI_TIRX_SCATTER:-1}"
export MAGI_TIRX_GATHER="${MAGI_TIRX_GATHER:-1}"
export MAGI_TIRX_DOWN="${MAGI_TIRX_DOWN:-1}"
export MAGI_TIRX_DOWN_PREWARM="${MAGI_TIRX_DOWN_PREWARM:-1}"
export MAGI_TIRX_QK_ROPE="${MAGI_TIRX_QK_ROPE:-1}"
export MAGI_TIRX_PRENORM="${MAGI_TIRX_PRENORM:-1}"
export MAGI_TIRX_SWIGLU7="${MAGI_TIRX_SWIGLU7:-1}"
export MAGI_TIRX_PREWARM="${MAGI_TIRX_PREWARM:-1}"
exec python infer_tirx_moe_scatter.py "$@"
fi
if [[ "${ELASTIC:-0}" == "1" ]]; then
exec python crown/infer_entry.py "$@"
fi
exec python inference/test_infer_seedance.py "$@"