diff --git a/fastvideo/models/hunyuan/inference.py b/fastvideo/models/hunyuan/inference.py index 98206d7ed..9ce0a42eb 100644 --- a/fastvideo/models/hunyuan/inference.py +++ b/fastvideo/models/hunyuan/inference.py @@ -18,23 +18,12 @@ from fastvideo.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_emb from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline -try: - import xfuser - from xfuser.core.distributed import ( - get_sequence_parallel_world_size, - get_sequence_parallel_rank, - get_sp_group, - initialize_model_parallel, - init_distributed_environment - ) -except: - xfuser = None - get_sequence_parallel_world_size = None - get_sequence_parallel_rank = None - get_sp_group = None - initialize_model_parallel = None - init_distributed_environment = None +from fastvideo.utils.parallel_states import ( + initialize_sequence_parallel_state, + nccl_info, +) +from fastvideo.utils.communications import all_gather, all_to_all_4D def parallelize_transformer(pipe): transformer = pipe.transformer @@ -53,34 +42,30 @@ def parallelize_transformer(pipe): guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. return_dict: bool = True, ): - if x.shape[-2] // 2 % get_sequence_parallel_world_size() == 0: + if x.shape[-2] // 2 % nccl_info.sp_size == 0: # try to split x by height split_dim = -2 - elif x.shape[-1] // 2 % get_sequence_parallel_world_size() == 0: + elif x.shape[-1] // 2 % nccl_info.sp_size == 0: # try to split x by width split_dim = -1 else: - raise ValueError(f"Cannot split video sequence into ulysses_degree x ring_degree ({get_sequence_parallel_world_size()}) parts evenly") + raise ValueError(f"Cannot split video sequence into ulysses_degree x ring_degree ({nccl_info.sp_size}) parts evenly") # patch sizes for the temporal, height, and width dimensions are 1, 2, and 2. temporal_size, h, w = x.shape[2], x.shape[3] // 2, x.shape[4] // 2 - x = torch.chunk(x, get_sequence_parallel_world_size(),dim=split_dim)[get_sequence_parallel_rank()] + x = torch.chunk(x, nccl_info.sp_size,dim=split_dim)[nccl_info.rank_within_group] + print("### x.shape", x.shape) dim_thw = freqs_cos.shape[-1] freqs_cos = freqs_cos.reshape(temporal_size, h, w, dim_thw) - freqs_cos = torch.chunk(freqs_cos, get_sequence_parallel_world_size(),dim=split_dim - 1)[get_sequence_parallel_rank()] + freqs_cos = torch.chunk(freqs_cos, nccl_info.sp_size,dim=split_dim - 1)[nccl_info.rank_within_group] freqs_cos = freqs_cos.reshape(-1, dim_thw) dim_thw = freqs_sin.shape[-1] freqs_sin = freqs_sin.reshape(temporal_size, h, w, dim_thw) - freqs_sin = torch.chunk(freqs_sin, get_sequence_parallel_world_size(),dim=split_dim - 1)[get_sequence_parallel_rank()] + freqs_sin = torch.chunk(freqs_sin, nccl_info.sp_size,dim=split_dim - 1)[nccl_info.rank_within_group] freqs_sin = freqs_sin.reshape(-1, dim_thw) - from xfuser.core.long_ctx_attention import xFuserLongContextAttention - - for block in transformer.double_blocks + transformer.single_blocks: - block.hybrid_seq_parallel_attn = xFuserLongContextAttention() - output = original_forward( x, t, @@ -95,7 +80,7 @@ def parallelize_transformer(pipe): return_dict = not isinstance(output, tuple) sample = output["x"] - sample = get_sp_group().all_gather(sample, dim=split_dim) + sample = all_gather(sample, dim=split_dim) output["x"] = sample return output @@ -153,28 +138,8 @@ class Inference(object): logger.info(f"Got text-to-video model root path: {pretrained_model_path}") # ==================== Initialize Distributed Environment ================ - if False: - if args.ulysses_degree > 1 or args.ring_degree > 1: - assert xfuser is not None, \ - "Ulysses Attention and Ring Attention requires xfuser package." - - assert args.use_cpu_offload is False, \ - "Cannot enable use_cpu_offload in the distributed environment." - - dist.init_process_group("nccl") - - assert dist.get_world_size() == args.ring_degree * args.ulysses_degree, \ - "number of GPUs should be equal to ring_degree * ulysses_degree." - - init_distributed_environment(rank=dist.get_rank(), world_size=dist.get_world_size()) - - initialize_model_parallel( - sequence_parallel_degree=dist.get_world_size(), - ring_degree=args.ring_degree, - ulysses_degree=args.ulysses_degree, - ) - device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") - #else: + if nccl_info.sp_size > 1: + device = torch.device(f"cuda:{os.environ['LOCAL_RANK']}") if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" @@ -522,12 +487,8 @@ class HunyuanVideoSampler(Inference): num_images_per_prompt (int): The number of images per prompt. Default is 1. infer_steps (int): The number of inference steps. Default is 100. """ - if False: - if self.parallel_args['ulysses_degree'] > 1 or self.parallel_args['ring_degree'] > 1: - assert seed is not None, \ - "You have to set a seed in the distributed environment, please rerun with --seed ." - - parallelize_transformer(self.pipeline) + if nccl_info.sp_size > 1: + parallelize_transformer(self.pipeline) out_dict = dict() diff --git a/fastvideo/models/hunyuan/modules/attenion.py b/fastvideo/models/hunyuan/modules/attenion.py index 445487937..00009c4c5 100644 --- a/fastvideo/models/hunyuan/modules/attenion.py +++ b/fastvideo/models/hunyuan/modules/attenion.py @@ -13,6 +13,10 @@ except ImportError: flash_attn = None flash_attn_varlen_func = None _flash_attn_forward = None + +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info +from fastvideo.utils.communications import all_gather, all_to_all_4D +from fastvideo.utils.logging import ForkedPdb MEMORY_LAYOUT = { @@ -157,7 +161,6 @@ def attention( def parallel_attention( - hybrid_seq_parallel_attn, q, k, v, @@ -166,18 +169,61 @@ def parallel_attention( cu_seqlens_q, cu_seqlens_kv ): - attn1 = hybrid_seq_parallel_attn( - None, - q[:, :img_q_len, :, :], - k[:, :img_kv_len, :, :], - v[:, :img_kv_len, :, :], - dropout_p=0.0, - causal=False, - joint_tensor_query=q[:,img_q_len:cu_seqlens_q[1]], - joint_tensor_key=k[:,img_kv_len:cu_seqlens_kv[1]], - joint_tensor_value=v[:,img_kv_len:cu_seqlens_kv[1]], - joint_strategy="rear", - ) + # 1GPU torch.Size([1, 11264, 24, 128]) tensor([ 0, 11275, 11520], device='cuda:0', dtype=torch.int32) + # 2GPU torch.Size([1, 5632, 24, 128]) tensor([ 0, 5643, 5888], device='cuda:0', dtype=torch.int32) + query, encoder_query = q[:, :img_q_len, :, :], q[:,img_q_len:cu_seqlens_q[1]] + key, encoder_key = k[:, :img_kv_len, :, :], k[:,img_kv_len:cu_seqlens_kv[1]] + value, encoder_value = v[:, :img_kv_len, :, :], v[:,img_kv_len:cu_seqlens_kv[1]] + print(query.shape, encoder_query.shape) + + if True: + # batch_size, seq_len, attn_heads, head_dim + query = all_to_all_4D(query, scatter_dim=2, gather_dim=1) + key = all_to_all_4D(key, scatter_dim=2, gather_dim=1) + value = all_to_all_4D(value, scatter_dim=2, gather_dim=1) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads) + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + + query, key, value = ( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2) + ) + encoder_query, encoder_key, encoder_value = ( + encoder_query.transpose(1, 2), + encoder_key.transpose(1, 2), + encoder_value.transpose(1, 2), + ) + # [b, h, s, d] + sequence_length = query.size(2) + encoder_sequence_length = encoder_query.size(2) + + # Hint: please check encoder_query.shape + query = torch.cat([query, encoder_query], dim=2) + key = torch.cat([key, encoder_key], dim=2) + value = torch.cat([value, encoder_value], dim=2) + + hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + (sequence_length, encoder_sequence_length), dim=2 + ) + # B, H, S, D + hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous() + hidden_states = hidden_states.transpose(1,2)#.flatten(2, 3) + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.transpose(1, 2)#.flatten(2, 3) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + + attn1 = torch.cat([hidden_states, encoder_hidden_states], dim=1) + #ForkedPdb().set_trace() + if flash_attn.__version__ >= '2.7.0': attn2, *_ = _flash_attn_forward( q[:,cu_seqlens_q[1]:], diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py index 646a42d03..7a9c77a4e 100644 --- a/fastvideo/models/hunyuan/modules/models.py +++ b/fastvideo/models/hunyuan/modules/models.py @@ -17,6 +17,9 @@ from .mlp_layers import MLP, MLPEmbedder, FinalLayer from .modulate_layers import ModulateDiT, modulate, apply_gate from .token_refiner import SingleTokenRefiner +from fastvideo.utils.parallel_states import ( + nccl_info, +) class MMDoubleStreamBlock(nn.Module): """ @@ -200,7 +203,17 @@ class MMDoubleStreamBlock(nn.Module): ), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, img.shape[0]:{img.shape[0]}" # attention computation start - if not self.hybrid_seq_parallel_attn: + if nccl_info.sp_size > 1: + attn = parallel_attention( + q, + k, + v, + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv + ) + else: attn = attention( q, k, @@ -211,17 +224,6 @@ class MMDoubleStreamBlock(nn.Module): max_seqlen_kv=max_seqlen_kv, batch_size=img_k.shape[0], ) - else: - attn = parallel_attention( - self.hybrid_seq_parallel_attn, - q, - k, - v, - img_q_len=img_q.shape[1], - img_kv_len=img_k.shape[1], - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv - ) # attention computation end @@ -364,7 +366,17 @@ class MMSingleStreamBlock(nn.Module): ), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}" # attention computation start - if not self.hybrid_seq_parallel_attn: + if nccl_info.sp_size > 1: + attn = parallel_attention( + q, + k, + v, + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv + ) + else: attn = attention( q, k, @@ -375,17 +387,7 @@ class MMSingleStreamBlock(nn.Module): max_seqlen_kv=max_seqlen_kv, batch_size=x.shape[0], ) - else: - attn = parallel_attention( - self.hybrid_seq_parallel_attn, - q, - k, - v, - img_q_len=img_q.shape[1], - img_kv_len=img_k.shape[1], - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv - ) + # attention computation end # Compute activation in mlp stream, cat again and run second linear layer. diff --git a/fastvideo/sample/sample_t2v_hunyuan.py b/fastvideo/sample/sample_t2v_hunyuan.py index 8700e98e4..54513a201 100644 --- a/fastvideo/sample/sample_t2v_hunyuan.py +++ b/fastvideo/sample/sample_t2v_hunyuan.py @@ -15,7 +15,29 @@ from diffusers.utils import export_to_video from fastvideo.models.hunyuan.utils.file_utils import save_videos_grid from fastvideo.models.hunyuan.inference import HunyuanVideoSampler +import torch.distributed as dist + +from fastvideo.utils.parallel_states import ( + initialize_sequence_parallel_state, + nccl_info, +) + +def initialize_distributed(): + local_rank = int(os.getenv("RANK", 0)) + world_size = int(os.getenv("WORLD_SIZE", 1)) + print("world_size", world_size) + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="nccl", init_method="env://", world_size=world_size, rank=local_rank + ) + initialize_sequence_parallel_state(world_size) + + def main(args): + initialize_distributed() + print(nccl_info.sp_size) + device = torch.cuda.current_device() + print(args) models_root_path = Path(args.model_path) if not models_root_path.exists(): @@ -32,21 +54,23 @@ def main(args): args = hunyuan_video_sampler.args # Start sampling - outputs = hunyuan_video_sampler.predict( - prompt=args.prompts[0], - height=args.height, - width=args.width, - video_length=args.num_frames, - seed=args.seed, - negative_prompt=args.neg_prompt, - infer_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - num_videos_per_prompt=args.num_videos, - flow_shift=args.flow_shift, - batch_size=args.batch_size, - embedded_guidance_scale=args.embedded_cfg_scale - ) - samples = outputs['samples'] + samples = [] + for prompt in args.prompts: + outputs = hunyuan_video_sampler.predict( + prompt=prompt, + height=args.height, + width=args.width, + video_length=args.num_frames, + seed=args.seed, + negative_prompt=args.neg_prompt, + infer_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + num_videos_per_prompt=args.num_videos, + flow_shift=args.flow_shift, + batch_size=args.batch_size, + embedded_guidance_scale=args.embedded_cfg_scale + ) + samples.append(outputs['samples'][0]) for prompt, video in zip(args.prompts, samples): videos = rearrange(video.unsqueeze(0), "b c t h w -> t b c h w") diff --git a/fastvideo/sample/sample_t2v_hunyuan_no_sp.py b/fastvideo/sample/sample_t2v_hunyuan_no_sp.py new file mode 100644 index 000000000..ced952082 --- /dev/null +++ b/fastvideo/sample/sample_t2v_hunyuan_no_sp.py @@ -0,0 +1,126 @@ +import os +import imageio +import time +from einops import rearrange + +import torch +import torchvision +import numpy as np +from pathlib import Path +from loguru import logger +from datetime import datetime +import argparse +from diffusers.utils import export_to_video + +from fastvideo.models.hunyuan.utils.file_utils import save_videos_grid +from fastvideo.models.hunyuan.inference import HunyuanVideoSampler + +def main(args): + print(args) + models_root_path = Path(args.model_path) + if not models_root_path.exists(): + raise ValueError(f"`models_root` not exists: {models_root_path}") + + # Create save folder to save the samples + save_path = args.output_path + os.makedirs(os.path.dirname(save_path), exist_ok=True) + + # Load models + hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args) + + # Get the updated args + args = hunyuan_video_sampler.args + + # Start sampling + samples = [] + for prompt in args.prompts: + outputs = hunyuan_video_sampler.predict( + prompt=prompt, + height=args.height, + width=args.width, + video_length=args.num_frames, + seed=args.seed, + negative_prompt=args.neg_prompt, + infer_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + num_videos_per_prompt=args.num_videos, + flow_shift=args.flow_shift, + batch_size=args.batch_size, + embedded_guidance_scale=args.embedded_cfg_scale + ) + samples.append(outputs['samples'][0]) + + for prompt, video in zip(args.prompts, samples): + videos = rearrange(video.unsqueeze(0), "b c t h w -> t b c h w") + outputs = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=6) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + outputs.append((x * 255).numpy().astype(np.uint8)) + os.makedirs(os.path.dirname(args.output_path), exist_ok=True) + imageio.mimsave(args.output_path + f"{prompt[:100]}.mp4", outputs, fps=args.fps) + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # Basic parameters + parser.add_argument("--prompts", nargs="+", default=[]) + parser.add_argument("--num_frames", type=int, default=16) + parser.add_argument("--height", type=int, default=256) + parser.add_argument("--width", type=int, default=256) + parser.add_argument("--num_inference_steps", type=int, default=50) + parser.add_argument("--model_path", type=str, default="data/hunyuan") + parser.add_argument("--output_path", type=str, default="./outputs/video") + parser.add_argument("--fps", type=int, default=24) + + # Additional parameters + parser.add_argument("--denoise-type", type=str, default="flow", help="Denoise type for noised inputs.") + parser.add_argument("--seed", type=int, default=None, help="Seed for evaluation.") + parser.add_argument("--neg_prompt", type=str, default=None, help="Negative prompt for sampling.") + parser.add_argument("--guidance_scale", type=float, default=1.0, help="Classifier free guidance scale.") + parser.add_argument("--embedded_cfg_scale", type=float, default=6.0, help="Embedded classifier free guidance scale.") + parser.add_argument("--flow_shift", type=int, default=7, help="Flow shift parameter.") + parser.add_argument("--batch_size", type=int, default=1, help="Batch size for inference.") + parser.add_argument("--num_videos", type=int, default=1, help="Number of videos to generate per prompt.") + parser.add_argument("--load-key", type=str, default="module", help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.") + parser.add_argument("--use-cpu-offload", action="store_true", help="Use CPU offload for the model load.") + parser.add_argument("--dit-weight", type=str, default="data/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt") + parser.add_argument("--reproduce", action="store_true", help="Enable reproducibility by setting random seeds and deterministic algorithms.") + parser.add_argument("--disable-autocast", action="store_true", help="Disable autocast for denoising loop and vae decoding in pipeline sampling.") + + # Flow Matching + parser.add_argument("--flow-reverse", action="store_true", help="If reverse, learning/sampling from t=1 -> t=0.") + parser.add_argument("--flow-solver", type=str, default="euler", help="Solver for flow matching.") + parser.add_argument("--use-linear-quadratic-schedule", action="store_true", + help="Use linear quadratic schedule for flow matching. Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)") + parser.add_argument("--linear-schedule-end", type=int, default=25, + help="End step for linear quadratic schedule for flow matching.") + + # Model parameters + parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill") + parser.add_argument("--latent-channels", type=int, default=16) + parser.add_argument("--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.") + + parser.add_argument("--vae", type=str, default="884-16c-hy") + parser.add_argument("--vae-precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--vae-tiling", action="store_true", default=True) + + parser.add_argument("--text-encoder", type=str, default="llm") + parser.add_argument("--text-encoder-precision", type=str, default="fp16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--text-states-dim", type=int, default=4096) + parser.add_argument("--text-len", type=int, default=256) + parser.add_argument("--tokenizer", type=str, default="llm") + parser.add_argument("--prompt-template", type=str, default="dit-llm-encode") + parser.add_argument("--prompt-template-video", type=str, default="dit-llm-encode-video") + parser.add_argument("--hidden-state-skip-layer", type=int, default=2) + parser.add_argument("--apply-final-norm", action="store_true") + + parser.add_argument("--text-encoder-2", type=str, default="clipL") + parser.add_argument("--text-encoder-precision-2", type=str, default="fp16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--text-states-dim-2", type=int, default=768) + parser.add_argument("--tokenizer-2", type=str, default="clipL") + parser.add_argument("--text-len-2", type=int, default=77) + + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/inference_hunyuan_no_sp.sh b/inference_hunyuan_no_sp.sh index 50b2ddb4a..af9aa71c0 100644 --- a/inference_hunyuan_no_sp.sh +++ b/inference_hunyuan_no_sp.sh @@ -1,4 +1,4 @@ -python3 fastvideo/sample/sample_t2v_hunyuan.py \ +python3 fastvideo/sample/sample_t2v_hunyuan_no_sp.py \ --height 500 \ --width 700 \ --num_frames 29 \ @@ -7,4 +7,5 @@ python3 fastvideo/sample/sample_t2v_hunyuan.py \ --embedded_cfg_scale 6 \ --flow-reverse \ --prompts "A cat walks on the grass, realistic style." \ + --prompts "A dog runs in the park, realistic style." \ --output_path outputs_video/hunyuan/ diff --git a/inference_hunyuan_sp.sh b/inference_hunyuan_sp.sh new file mode 100644 index 000000000..5f4c44465 --- /dev/null +++ b/inference_hunyuan_sp.sh @@ -0,0 +1,13 @@ +num_gpus=2 + +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/sample/sample_t2v_hunyuan.py \ + --height 500 \ + --width 700 \ + --num_frames 29 \ + --num_inference_steps 50 \ + --guidance_scale 1 \ + --embedded_cfg_scale 6 \ + --flow-reverse \ + --prompts "A cat walks on the grass, realistic style." \ + --output_path outputs_video/hunyuan/