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

add orginal mochi but performance bad

Browse Source
This commit is contained in:
“BrianChen1129”
2024-12-09 08:44:24 +00:00
parent 60ce6a62a6
commit 8a41e61cec
7 changed files with 43 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/fsdp_util.py
+2 -2
View File
@@ -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
},
)
fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py
+12 -43
View File
@@ -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)",
fastvideo/models/mochi_genmo/mochi_preview/vae/models.py
+2 -2
View File
@@ -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(
fastvideo/models/mochi_hf/pipeline_mochi.py
+16 -20
View File
@@ -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 * (
fastvideo/sample/sample_t2v_mochi.py
+3 -2
View File
@@ -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()
fastvideo/train.py
+7 -1
View File
@@ -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)
scripts/download_hf.py
+1 -1
View File
@@ -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)",