diff --git a/fastvideo/fsdp_util.py b/fastvideo/fsdp_util.py index 96a459767..1f0e596ad 100644 --- a/fastvideo/fsdp_util.py +++ b/fastvideo/fsdp_util.py @@ -18,7 +18,7 @@ from torch.distributed.fsdp import ( ) from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformerBlock - +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmetricJointBlock from functools import partial from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy @@ -88,7 +88,7 @@ def get_dit_fsdp_kwargs( auto_wrap_policy = functools.partial( transformer_auto_wrap_policy, transformer_layer_cls={ - MochiTransformerBlock, + MochiTransformerBlock, # AsymmetricJointBlock }, ) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py index e6d69dddd..56583533b 100644 --- a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py @@ -7,7 +7,7 @@ import torch.nn as nn import torch.nn.functional as F from einops import rearrange from torch.nn.attention import sdpa_kernel -import fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.context_parallel as cp + from fastvideo.models.mochi_genmo.lib.attn_imports import flash_varlen_attn, sage_attn, sdpa_attn_ctx from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.layers import ( FeedForward, @@ -110,18 +110,9 @@ class AsymmetricAttention(nn.Module): self.proj_y = LoraLinear(dim_x, dim_y, **proj_lora_kwargs) if update_y else nn.Identity() def run_qkv_y(self, y): - cp_rank, cp_size = cp.get_cp_rank_size() - local_heads = self.num_heads // cp_size + local_heads = self.num_heads - if cp.is_cp_active(): - # Only predict local heads. - assert not self.qkv_bias - W_qkv_y = self.qkv_y.weight.view(3, self.num_heads, self.head_dim, self.dim_y) - W_qkv_y = W_qkv_y.narrow(1, cp_rank * local_heads, local_heads) - W_qkv_y = W_qkv_y.reshape(3 * local_heads * self.head_dim, self.dim_y) - qkv_y = F.linear(y, W_qkv_y, None) # (B, L, 3 * local_h * head_dim) - else: - qkv_y = self.qkv_y(y) # (B, L, 3 * dim) + qkv_y = self.qkv_y(y) # (B, L, 3 * dim) qkv_y = qkv_y.view(qkv_y.size(0), qkv_y.size(1), 3, local_heads, self.head_dim) q_y, k_y, v_y = qkv_y.unbind(2) @@ -143,11 +134,12 @@ class AsymmetricAttention(nn.Module): max_seqlen_in_batch: int, ): # Process visual features - x = modulated_rmsnorm(x, scale_x) # (B, M, dim_x) where M = N / cp_group_size + x = modulated_rmsnorm(x, scale_x) # (B, M, dim_x) where M = N qkv_x = self.qkv_x(x) # (B, M, 3 * dim_x) assert qkv_x.dtype == torch.bfloat16 - - qkv_x = cp.all_to_all_collect_tokens(qkv_x, self.num_heads) # (3, B, N, local_h, head_dim) + B, M, _ = qkv_x.size() + qkv_x = qkv_x.view(B, M, 3, self.num_heads, -1) + qkv_x = qkv_x.permute(2, 0, 1, 3, 4) # Split qkv_x into q, k, v q_x, k_x, v_x = qkv_x.unbind(0) # (B, N, local_h, head_dim) @@ -224,9 +216,7 @@ class AsymmetricAttention(nn.Module): cu_seqlens: Optional[torch.Tensor] = None, max_seqlen_in_batch: Optional[int] = None, ): - _, cp_size = cp.get_cp_rank_size() - assert self.num_heads % cp_size == 0 - local_heads = self.num_heads // cp_size + local_heads = self.num_heads local_dim = local_heads * self.head_dim # Check shapes @@ -275,13 +265,12 @@ class AsymmetricAttention(nn.Module): dtype: Data type of the input and output tensors Returns: - x: (B, N, dim_x) tensor of visual tokens where N = M * cp_size + x: (B, N, dim_x) tensor of visual tokens where N = M y: (B, L, dim_y) tensor of text token features """ - _, cp_size = cp.get_cp_rank_size() - local_heads = self.num_heads // cp_size + local_heads = self.num_heads local_dim = local_heads * self.head_dim - N = M * cp_size + N = M # Split sequence into visual and text tokens, adding back padding. if B == 1: @@ -299,10 +288,7 @@ class AsymmetricAttention(nn.Module): # Communicate across context parallel ranks. x = x.view(B, N, local_heads, self.head_dim) - x = cp.all_to_all_collect_heads(x) # (B, M, dim_x = num_heads * head_dim) - if cp.is_cp_active(): - y = cp.all_gather(y) # (cp_size * B, L, local_heads * head_dim) - y = rearrange(y, "(G B) L D -> B L (G D)", G=cp_size, D=local_dim) # (B, L, dim_x) + x = x.view(x.size(0), x.size(1), x.size(2) * x.size(3)) # (B, M, dim_x = num_heads * head_dim) x = self.proj_x(x) y = self.proj_y(y) @@ -704,19 +690,6 @@ class AsymmDiTJoint(ModelMixin, ConfigMixin, PeftAdapterMixin): hidden_states, c, encoder_hidden_states, rope_cos, rope_sin = self.prepare(hidden_states, sigma, encoder_hidden_states, encoder_attention_mask) # [1, 11130, 3072], [1, 3072], [1, 256, 1536], [11130, 24, 64] del encoder_attention_mask - cp_rank, cp_size = cp.get_cp_rank_size() - N = hidden_states.size(1) - M = N // cp_size - assert N % cp_size == 0, f"Visual sequence length ({hidden_states.shape[1]}) must be divisible by cp_size ({cp_size})." - - if cp_size > 1: - hidden_states = hidden_states.narrow(1, cp_rank * M, M) - - assert self.num_heads % cp_size == 0 - local_heads = self.num_heads // cp_size - rope_cos = rope_cos.narrow(1, cp_rank * local_heads, local_heads) - rope_sin = rope_sin.narrow(1, cp_rank * local_heads, local_heads) - for i, block in enumerate(self.blocks): hidden_states, encoder_hidden_states = block( # [1, 11130, 3072], [1, 256, 1536] hidden_states, @@ -733,10 +706,6 @@ class AsymmDiTJoint(ModelMixin, ConfigMixin, PeftAdapterMixin): hidden_states = self.final_layer(hidden_states, c) # (B, M, patch_size ** 2 * out_channels) [1, 11130, 48] - patch = hidden_states.size(2) - - hidden_states = cp.all_gather(hidden_states) - hidden_states = rearrange(hidden_states, "(G B) M P -> B (G M) P", G=cp_size, P=patch) # [1, 11130, 48] hidden_states = rearrange( # [1, 12, 7, 60, 106] hidden_states, "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)", diff --git a/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py b/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py index 568bcabfa..a6560a990 100644 --- a/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py +++ b/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py @@ -602,8 +602,8 @@ class Decoder(nn.Module): x = F.silu(x, inplace=not self.training) else: assert not self.output_nonlinearity # StyleGAN3 omits the to-RGB nonlinearity. - - return self.output_proj(x).contiguous() + new_conv_cache = None + return self.output_proj(x).contiguous(), new_conv_cache def make_broadcastable( diff --git a/fastvideo/models/mochi_hf/pipeline_mochi.py b/fastvideo/models/mochi_hf/pipeline_mochi.py index 79822635a..38e0c8241 100644 --- a/fastvideo/models/mochi_hf/pipeline_mochi.py +++ b/fastvideo/models/mochi_hf/pipeline_mochi.py @@ -249,9 +249,13 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): add_special_tokens=True, return_tensors="pt", ) + text_input_ids = text_inputs.input_ids prompt_attention_mask = text_inputs.attention_mask prompt_attention_mask = prompt_attention_mask.bool().to(device) + if prompt == "" or prompt[-1] == "": + text_input_ids = torch.zeros_like(text_input_ids, device=device) + prompt_attention_mask = torch.zeros_like(prompt_attention_mask, dtype=torch.bool, device=device) untruncated_ids = self.tokenizer( prompt, padding="longest", return_tensors="pt" @@ -363,10 +367,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): " the batch size of `prompt`." ) - ( - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self._get_t5_prompt_embeds( + negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds( prompt=negative_prompt, num_videos_per_prompt=num_videos_per_prompt, max_sequence_length=max_sequence_length, @@ -374,12 +375,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): dtype=dtype, ) - return ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask def check_inputs( self, @@ -503,7 +499,8 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32) + latents = latents.to(dtype) return latents @property @@ -534,8 +531,8 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): negative_prompt: Optional[Union[str, List[str]]] = None, height: Optional[int] = None, width: Optional[int] = None, - num_frames: int = 16, - num_inference_steps: int = 28, + num_frames: int = 19, + num_inference_steps: int = 64, timesteps: List[int] = None, guidance_scale: float = 4.5, num_videos_per_prompt: Optional[int] = 1, @@ -677,12 +674,6 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): max_sequence_length=max_sequence_length, device=device, ) - if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - prompt_attention_mask = torch.cat( - [negative_prompt_attention_mask, prompt_attention_mask], dim=0 - ) - # 4. Prepare latent variables num_channels_latents = self.transformer.config.in_channels latents = self.prepare_latents( @@ -696,6 +687,11 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): generator, latents, ) + + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group if get_sequence_parallel_state(): latents = rearrange( @@ -751,9 +747,9 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin): attention_kwargs=attention_kwargs, return_dict=False, )[0] - # Mochi CFG + Sampling runs in FP32 noise_pred = noise_pred.to(torch.float32) + if self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + self.guidance_scale * ( diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py index fad2dcb38..d5efaefb1 100644 --- a/fastvideo/sample/sample_t2v_mochi.py +++ b/fastvideo/sample/sample_t2v_mochi.py @@ -55,9 +55,9 @@ def main(args): args.linear_range, ) - mochi_genmo = True + mochi_genmo = False if mochi_genmo: - model_path = "/root/fastmochi_genmo/dit.safetensors" + model_path = "/root/weights/dit.safetensors" state_dcit = load_file(model_path) transformer = AsymmDiTJoint() transformer.load_state_dict(state_dcit) @@ -74,6 +74,7 @@ def main(args): pipe = MochiPipeline.from_pretrained( args.model_path, transformer=transformer, scheduler=scheduler ) + pipe.enable_vae_tiling() diff --git a/fastvideo/train.py b/fastvideo/train.py index 93252a9b6..878e59d72 100644 --- a/fastvideo/train.py +++ b/fastvideo/train.py @@ -214,7 +214,7 @@ def main(args): # as these weights are only used for inference, keeping weights in full precision is not required. # Create model: - +f main_print(f"--> loading model from {args.pretrained_model_name_or_path}") # keep the master weight to float32 weight_type = torch.float32 if args.master_weight_type == 'fp32' else torch.bfloat16 @@ -729,6 +729,12 @@ if __name__ == "__main__": default="fp32", help="Weight type to use - fp32 or bf16.", ) + parser.add_argument( + "--Mochi_type", + type=str, + default="hf", + help="Choose Mochi model between hf and genmo(original mochi).", + ) args = parser.parse_args() main(args) diff --git a/scripts/download_hf.py b/scripts/download_hf.py index cef005c71..cd7b1e29e 100644 --- a/scripts/download_hf.py +++ b/scripts/download_hf.py @@ -15,7 +15,7 @@ if __name__ == "__main__": type=str, help="The local directory to download the repository to", ) - parser.add_argument( + parser.add_argument(di "--repo_type", type=str, help="The type of repository to download (dataset or model)",