nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Infer sta tea with torch.compile (#190)

Browse Source 
Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>
This commit is contained in:
Yongqi Chen
2025-02-18 11:29:36 -08:00
committed by GitHub
co-authored by Peiyuan Zhang
parent ac490399c6
commit 76a12cc8a1
7 changed files with 73 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+10 -5
View File
@@ -24,7 +24,7 @@ Dev in progress and highly experimental.
## Change Log
- ```2025/02/14```: Release the inference code and kernel for Sliding Tile Attention(https://hao-ai-lab.github.io/blogs/sta/).
- ```2025/02/14```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
- ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
- ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
- ```2024/12/17```: `FastVideo` v1.0 is released.
@@ -41,11 +41,16 @@ To try sliding tile attention (optional), please follow the instruction in [csrc
## 🚀 Inference
### Inference HunyuanVideo with Sliding Tile Attention
```bash
python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
# Inference with FA2 full attn
sh scripts/inference/inference_hunyuan_STA.sh
First, download the model:
```
python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
```
We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
```
scripts/inference/inference_hunyuan_STA.sh
```
### Video Demos using STA + Teacache
Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on a single H100.
### Inference FastHunyuan on single RTX4090
We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py
+12 -5
View File
@@ -869,13 +869,20 @@ class HunyuanVideoPipeline(DiffusionPipeline):
num_warmup_steps = len(
timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
def dict_to_3d_list(mask_strategy, t_max=50, l_max=60, h_max=24):
result = [[[None for _ in range(h_max)]
for _ in range(l_max)]
for _ in range(t_max)]
if mask_strategy is None:
return result
for key, value in mask_strategy.items():
t, l, h = map(int, key.split('_'))
result[t][l][h] = value
return result
mask_strategy = dict_to_3d_list(mask_strategy)
# if is_progress_bar:
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
mask_param = [
mask_strategy, i
] # if mask_strategy is None, STA will not be used
if self.interrupt:
continue
@@ -912,7 +919,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
encoder_hidden_states,
t_expand,
prompt_mask,
mask_param=mask_param,
mask_strategy=mask_strategy[i],
guidance=guidance_expand,
return_dict=False,
)[0]
fastvideo/models/hunyuan/inference.py
+2
View File
@@ -100,6 +100,8 @@ class Inference(object):
)
model = model.to(device)
model = Inference.load_state_dict(args, model, pretrained_model_path)
if args.enable_torch_compile:
model = torch.compile(model)
model.eval()
# ============================= Build extra models ========================
fastvideo/models/hunyuan/modules/attenion.py
+4 -5
View File
@@ -76,7 +76,7 @@ def parallel_attention(q,
img_q_len,
img_kv_len,
text_mask,
mask_param=None):
mask_strategy=None):
query, encoder_query = q
key, encoder_key = k
value, encoder_value = v
@@ -101,7 +101,7 @@ def parallel_attention(q,
sequence_length = query.size(1)
encoder_sequence_length = encoder_query.size(1)
if mask_param[0] is not None:
if mask_strategy is not None:
query = torch.cat([tile(query, nccl_info.sp_size), encoder_query],
dim=1).transpose(1, 2)
key = torch.cat([tile(key, nccl_info.sp_size), encoder_key],
@@ -110,9 +110,8 @@ def parallel_attention(q,
dim=1).transpose(1, 2)
head_num = query.size(1)
mask_strategy, time_step, layer_idx = mask_param
windows = [
mask_strategy[f"{time_step}_{layer_idx}_{head_idx}"]
mask_strategy[head_idx]
for head_idx in range(head_num)
]
@@ -135,7 +134,7 @@ def parallel_attention(q,
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
(sequence_length, encoder_sequence_length), dim=1)
if mask_param[0] is not None:
if mask_strategy is not None:
hidden_states = untile(hidden_states, nccl_info.sp_size)
if get_sequence_parallel_state():
fastvideo/models/hunyuan/modules/models.py
+7 -11
View File
@@ -138,7 +138,7 @@ class MMDoubleStreamBlock(nn.Module):
vec: torch.Tensor,
freqs_cis: tuple = None,
text_mask: torch.Tensor = None,
mask_param=None,
mask_strategy=None,
) -> Tuple[torch.Tensor, torch.Tensor]:
(
img_mod1_shift,
@@ -215,7 +215,7 @@ class MMDoubleStreamBlock(nn.Module):
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
text_mask=text_mask,
mask_param=mask_param,
mask_strategy=mask_strategy,
)
# attention computation end
@@ -319,7 +319,7 @@ class MMSingleStreamBlock(nn.Module):
txt_len: int,
freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
text_mask: torch.Tensor = None,
mask_param=None,
mask_strategy=None,
) -> torch.Tensor:
mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
@@ -365,7 +365,7 @@ class MMSingleStreamBlock(nn.Module):
img_q_len=img_q.shape[1],
img_kv_len=img_k.shape[1],
text_mask=text_mask,
mask_param=mask_param,
mask_strategy=mask_strategy,
)
# attention computation end
@@ -597,7 +597,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_attention_mask: torch.Tensor,
mask_param=None,
mask_strategy=None,
output_features=False,
output_features_stride=8,
attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -656,28 +656,24 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
# --------------------- Pass through DiT blocks ------------------------
for index, block in enumerate(self.double_blocks):
mask_param.append(index)
double_block_args = [
img, txt, vec, freqs_cis, text_mask, mask_param
img, txt, vec, freqs_cis, text_mask, mask_strategy[index]
]
img, txt = block(*double_block_args)
mask_param = mask_param[:-1]
# Merge txt and img to pass through single stream blocks.
x = torch.cat((img, txt), 1)
if output_features:
features_list = []
if len(self.single_blocks) > 0:
for index, block in enumerate(self.single_blocks):
mask_param.append(index + len(self.double_blocks))
single_block_args = [
x,
vec,
txt_seq_len,
(freqs_cos, freqs_sin),
text_mask,
mask_param,
mask_strategy[index + len(self.double_blocks)],
]
mask_param = mask_param[:-1]
x = block(*single_block_args)
if output_features and _ % output_features_stride == 0:
features_list.append(x[:, :img_seq_len, ...])
fastvideo/sample/sample_t2v_hunyuan_STA.py
+12 -14
View File
@@ -23,7 +23,7 @@ def teacache_forward(
encoder_hidden_states: torch.Tensor,
timestep: torch.LongTensor,
encoder_attention_mask: torch.Tensor,
mask_param=None,
mask_strategy=None,
output_features=False,
output_features_stride=8,
attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -122,12 +122,10 @@ def teacache_forward(
ori_img = img.clone()
# --------------------- Pass through DiT blocks ------------------------
for index, block in enumerate(self.double_blocks):
mask_param.append(index)
double_block_args = [
img, txt, vec, freqs_cis, text_mask, mask_param
img, txt, vec, freqs_cis, text_mask, mask_strategy[index]
]
img, txt = block(*double_block_args)
mask_param = mask_param[:-1]
# Merge txt and img to pass through single stream blocks.
x = torch.cat((img, txt), 1)
@@ -135,16 +133,14 @@ def teacache_forward(
features_list = []
if len(self.single_blocks) > 0:
for index, block in enumerate(self.single_blocks):
mask_param.append(index + len(self.double_blocks))
single_block_args = [
x,
vec,
txt_seq_len,
(freqs_cos, freqs_sin),
text_mask,
mask_param,
mask_strategy[index + len(self.double_blocks)],
]
mask_param = mask_param[:-1]
x = block(*single_block_args)
if output_features and _ % output_features_stride == 0:
features_list.append(x[:, :img_seq_len, ...])
@@ -154,28 +150,24 @@ def teacache_forward(
else:
# --------------------- Pass through DiT blocks ------------------------
for index, block in enumerate(self.double_blocks):
mask_param.append(index)
double_block_args = [
img, txt, vec, freqs_cis, text_mask, mask_param
img, txt, vec, freqs_cis, text_mask, mask_strategy[index]
]
img, txt = block(*double_block_args)
mask_param = mask_param[:-1]
# Merge txt and img to pass through single stream blocks.
x = torch.cat((img, txt), 1)
if output_features:
features_list = []
if len(self.single_blocks) > 0:
for index, block in enumerate(self.single_blocks):
mask_param.append(index + len(self.double_blocks))
single_block_args = [
x,
vec,
txt_seq_len,
(freqs_cos, freqs_sin),
text_mask,
mask_param,
mask_strategy[index + len(self.double_blocks)],
]
mask_param = mask_param[:-1]
x = block(*single_block_args)
if output_features and _ % output_features_stride == 0:
features_list.append(x[:, :img_seq_len, ...])
@@ -214,7 +206,8 @@ def main(args):
models_root_path = Path(args.model_path)
if not models_root_path.exists():
raise ValueError(f"`models_root` not exists: {models_root_path}")
if args.enable_teacache and args.enable_torch_compile:
parser.error("--enable_teacache and --enable_torch_compile cannot be used simultaneously. Please enable only one of these options.")
# Create save folder to save the samples
save_path = args.output_path
os.makedirs(os.path.dirname(save_path), exist_ok=True)
@@ -450,6 +443,11 @@ if __name__ == "__main__":
action="store_true",
help="Use teacache for speeding up inference",
)
parser.add_argument(
"--enable_torch_compile",
action="store_true",
help="Use torch.compile for speeding up STA inference without teacache",
)
parser.add_argument("--mask_strategy_file_path",
type=str,
default="assets/mask_strategy.json")
scripts/inference/inference_hunyuan_STA.sh
+26 -2
View File
@@ -1,9 +1,10 @@
#!/bin/bash
# Inference with STA + Teacache
num_gpus=1
mask_strategy_file_path=assets/mask_strategy.json
export MODEL_BASE=data/hunyuan
rel_l1_thresh=0.2
rel_l1_thresh=0.15
CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
fastvideo/sample/sample_t2v_hunyuan_STA.py \
--height 768 \
@@ -22,4 +23,27 @@ CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_p
--rel_l1_thresh $rel_l1_thresh \
--dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
--vae-sp \
--enable_teacache
--enable_teacache
# Inference with STA only
num_gpus=1
mask_strategy_file_path=assets/mask_strategy.json
export MODEL_BASE=data/hunyuan
CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
fastvideo/sample/sample_t2v_hunyuan_STA.py \
--height 768 \
--width 1280 \
--num_frames 117 \
--num_inference_steps 50 \
--guidance_scale 1 \
--embedded_cfg_scale 6 \
--flow_shift 7 \
--flow-reverse \
--prompt ./assets/prompt.txt \
--seed 12345 \
--output_path outputs_video/hunyuan_STA/ \
--model_path $MODEL_BASE \
--mask_strategy_file_path $mask_strategy_file_path \
--dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
--vae-sp \
--enable_torch_compile