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

Merge branch 'hl/diffusers' of https://github.com/jzhang38/FastVideo-OSP into peiyuan

Browse Source
This commit is contained in:
rlsu9
2024-11-01 04:44:23 +04:00
parent 5975deed62 db2c68e879
commit 97f48ef433
13 changed files with 367 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/sample/generate_synthetic.py
+1 -1
View File
@@ -19,7 +19,7 @@ def generate_video_and_latent(pipe, prompt, height, width, num_frames, num_infer
generator=generator,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
output_type="latent_and_video"
return_all_states=True,
)
return video[0], latent[0], prompt_embed[0], prompt_attention_mask[0]
fastvideo/sample/monkey_patch/sp_attention.py
+146
View File
@@ -0,0 +1,146 @@
import inspect
import math
from typing import Callable, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
from torch import nn
from einops import rearrange, repeat
from diffusers.image_processor import IPAdapterMaskProcessor
from diffusers.utils import deprecate, logging
from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available
from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph
from diffusers.models.embeddings import apply_rotary_emb
import diffusers
from types import MethodType
from diffusers.models.attention_processor import Attention
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
from fastvideo.utils.communications import all_to_all_BHSD, all_gather_BHSD
class NewMochiAttnProcessor2_0:
"""Attention processor used in Mochi."""
def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
image_rotary_emb: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# [b, s, h * d]
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
# [b, s, h=24, d=128]
query = query.unflatten(2, (attn.heads, -1))
key = key.unflatten(2, (attn.heads, -1))
value = value.unflatten(2, (attn.heads, -1))
if attn.norm_q is not None:
query = attn.norm_q(query)
if attn.norm_k is not None:
key = attn.norm_k(key)
# [b, 256, h * d]
encoder_query = attn.add_q_proj(encoder_hidden_states)
encoder_key = attn.add_k_proj(encoder_hidden_states)
encoder_value = attn.add_v_proj(encoder_hidden_states)
# [b, 256, h=24, d=128]
encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
if attn.norm_added_q is not None:
encoder_query = attn.norm_added_q(encoder_query)
if attn.norm_added_k is not None:
encoder_key = attn.norm_added_k(encoder_key)
if image_rotary_emb is not None:
def apply_rotary_emb(x, freqs_cos, freqs_sin):
x_even = x[..., 0::2].float()
x_odd = x[..., 1::2].float()
cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
return torch.stack([cos, sin], dim=-1).flatten(-2)
query = apply_rotary_emb(query, *image_rotary_emb)
key = apply_rotary_emb(key, *image_rotary_emb)
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)
total_sequence_length = sequence_length * nccl_info.world_size
encoder_sequence_length = encoder_query.size(2)
print("sequence_length: ", sequence_length)
print("encoder_sequence_length: ", encoder_sequence_length)
# if sp query slice, encoder_query full
# first all-to-all query -> full query, then concat encoder_query
if get_sequence_parallel_state():
query = all_to_all_BHSD(query, scatter_dim=1, gather_dim=2).contiguous()
key = all_to_all_BHSD(key, scatter_dim=1, gather_dim=2).contiguous()
value = all_to_all_BHSD(value, scatter_dim=1, gather_dim=2).contiguous()
if get_sequence_parallel_state():
def shrink_head(encoder_state):
world_size, rank = nccl_info.world_size, nccl_info.rank
return rearrange(encoder_state, 'b (n h) s d -> b n h s d', n=world_size,
h=encoder_state.shape[1] // world_size).contiguous()[:, rank, ...]
encoder_query = shrink_head(encoder_query)
encoder_key = shrink_head(encoder_key)
encoder_value = shrink_head(encoder_value)
# Hint: please check encoder_query.shape
print("query.shape: ", query.shape)
print("encoder_query.shape: ", 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)
print("hidden_states.shape: ", hidden_states.shape)
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
print("after attn before sp.shape", hidden_states.shape)
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
(total_sequence_length, encoder_sequence_length), dim=2
)
print("after attn before sp.shape", hidden_states.shape)
if get_sequence_parallel_state():
hidden_states = all_to_all_BHSD(hidden_states, scatter_dim=2, gather_dim=1).contiguous()
encoder_hidden_states = all_gather_BHSD(encoder_hidden_states, dim=1).contiguous()
print("after attn after sp.shape", hidden_states.shape)
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)
print("after trans.shape", hidden_states.shape, encoder_hidden_states.shape)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if hasattr(attn, "to_add_out"):
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
return hidden_states, encoder_hidden_states
fastvideo/sample/monkey_patch/sp_rope.py
+32
View File
@@ -0,0 +1,32 @@
import inspect
import math
from typing import Callable, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
from torch import nn
from einops import rearrange, repeat
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
class NewMochiRoPE(nn.Module):
def _get_positions(
self,
num_frames: int,
height: int,
width: int,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
scale = (self.target_area / (height * width)) ** 0.5
t = torch.arange(num_frames, device=device, dtype=dtype)
if get_sequence_parallel_state():
rank = nccl_info.rank
t += rank * num_frames
h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
return positions
fastvideo/sample/mochi_pipe.py → fastvideo/sample/pipeline_mochi.py
+23 -9
View File
@@ -32,7 +32,8 @@ from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
from einops import rearrange
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
if is_torch_xla_available():
import torch_xla.core.xla_model as xm
@@ -492,6 +493,7 @@ class MochiPipeline(DiffusionPipeline):
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 256,
return_all_states = False,
):
r"""
Function invoked when calling the pipeline for generation.
@@ -629,7 +631,11 @@ class MochiPipeline(DiffusionPipeline):
generator,
latents,
)
world_size, rank = nccl_info.world_size, nccl_info.rank
if get_sequence_parallel_state():
latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
latents = latents[:, :, rank, :, :, :]
# 5. Prepare timestep
# from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
threshold_noise = 0.025
@@ -693,6 +699,14 @@ class MochiPipeline(DiffusionPipeline):
if XLA_AVAILABLE:
xm.mark_step()
if get_sequence_parallel_state():
latents_shape = list(latents.shape)
full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
torch.distributed.all_gather_into_tensor(all_latents, latents)
latents_list = list(all_latents.chunk(world_size, dim=0))
latents = torch.cat(latents_list, dim=2)
if output_type == "latent":
video = latents
else:
@@ -712,14 +726,14 @@ class MochiPipeline(DiffusionPipeline):
latents = latents / self.vae.config.scaling_factor
video = self.vae.decode(latents, return_dict=False)[0]
# TODO: Change pil hardcode
video = self.video_processor.postprocess_video(video, output_type='pil')
video = self.video_processor.postprocess_video(video, output_type=output_type)
if return_all_states:
return video, latents, prompt_embeds, prompt_attention_mask
# Offload all models
self.maybe_free_model_hooks()
if output_type == "latent_and_video":
return video, latents, prompt_embeds, prompt_attention_mask
if not return_dict:
return (video,)
return MochiPipelineOutput(frames=video)
return MochiPipelineOutput(frames=video)
fastvideo/sample/sample_t2v_mochi_sp.py
+67
View File
@@ -0,0 +1,67 @@
import torch
from fastvideo.sample.pipeline_mochi import MochiPipeline
import torch.distributed as dist
from diffusers.utils import export_to_video, load_image, load_video
from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
import argparse
import os
import diffusers
from fastvideo.sample.monkey_patch.sp_attention import NewMochiAttnProcessor2_0
from fastvideo.sample.monkey_patch.sp_rope import NewMochiRoPE
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 replace_hf_ckpt_with_new_ckpt():
initialize_distributed()
diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__
diffusers.models.transformers.transformer_mochi.MochiRoPE._get_positions = NewMochiRoPE._get_positions
def main(args):
replace_hf_ckpt_with_new_ckpt()
print(nccl_info.world_size)
torch.manual_seed(args.seed)
weight_dtype = torch.bfloat16
device = torch.cuda.current_device()
pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
pipe.enable_vae_tiling()
pipe.to(device)
#pipe.enable_model_cpu_offload()
# Generate videos from the input prompt
video = pipe(
prompt=args.prompt,
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
).frames[0]
dist.barrier()
if nccl_info.rank <= 0:
export_to_video(video, args.output_path, fps=30)
if __name__ == "__main__":
# arg parse
parser = argparse.ArgumentParser()
parser.add_argument("--prompt", type=str, default="A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.")
parser.add_argument("--num_frames", type=int, default=163)
parser.add_argument("--height", type=int, default=480)
parser.add_argument("--width", type=int, default=848)
parser.add_argument("--num_inference_steps", type=int, default=64)
parser.add_argument("--guidance_scale", type=float, default=4.5)
parser.add_argument("--model_path", type=str, default="data/mochi")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--output_path", type=str, default="./outputs.mp4")
args = parser.parse_args()
main(args)
fastvideo/train.py
+25 -14
View File
@@ -28,7 +28,7 @@ import accelerate
import torch
from torch.nn import functional as F
import transformers
from accelerate import Accelerator
from accelerate import Accelerator, init_empty_weights, DistributedType
from diffusers.utils.torch_utils import is_compiled_module
from accelerate.logging import get_logger
from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
@@ -50,7 +50,7 @@ from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_func
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
check_min_version("0.24.0")
check_min_version("0.31.0")
logger = get_logger(__name__)
@@ -73,6 +73,7 @@ class ForkedPdb(pdb.Pdb):
@torch.inference_mode()
def log_validation(args, transformer,vae, accelerator, weight_dtype, global_step, ema=False):
#TODO
logger.info(f"Running validation....\n")
transformer = accelerator.unwrap_model(transformer)
mochi_pipeline = MochiPipeline.from_pretrained(args.pretrained_model_name_or_path,vae=vae, transformer=transformer, torch_dtype=weight_dtype).to(device=accelerator.device)
@@ -229,30 +230,41 @@ def main(args):
def load_model_hook(models, input_dir):
# TODO
if args.use_ema:
load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), MochiTransformer3DModel)
ema_transformer.load_state_dict(load_model.state_dict())
ema_transformer.to(accelerator.device)
del load_model
for i in range(len(models)):
# pop models so that they are not loaded again
model = models.pop()
if not accelerator.distributed_type == DistributedType.DEEPSPEED:
while len(models) > 0:
model = models.pop()
# load diffusers style into model
load_model = MochiTransformer3DModel.from_pretrained(input_dir, subfolder="transformer")
model.register_to_config(**load_model.config)
model.load_state_dict(load_model.state_dict())
del load_model
if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
transformer_ = unwrap_model(model)
else:
raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
else:
with init_empty_weights():
transformer_ = MochiTransformer3DModel.from_config(
args.pretrained_model_name_or_path, subfolder="transformer"
)
init_under_meta = True
load_model = MochiTransformer3DModel.from_pretrained(os.path.join(input_dir, "transformer"))
transformer_.register_to_config(**load_model.config)
transformer_.load_state_dict(load_model.state_dict(), assign=init_under_meta)
del load_model
if args.mixed_precision == "fp16":
cast_training_params([transformer_])
accelerator.register_save_state_pre_hook(save_model_hook)
accelerator.register_load_state_pre_hook(load_model_hook)
# Enable TF32 for faster training on Ampere GPUs,
# cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
if args.allow_tf32:
if args.allow_tf32 and torch.cuda.is_available():
torch.backends.cuda.matmul.allow_tf32 = True
if args.scale_lr:
@@ -267,7 +279,6 @@ def main(args):
params_to_optimize = transformer.parameters()
# TODO: Other optmizer in cogvideoX-factory
optimizer = torch.optim.AdamW(
params_to_optimize,
lr=args.learning_rate,
fastvideo/utils/communications.py
+51 -2
View File
@@ -85,10 +85,10 @@ class _AllToAll(torch.autograd.Function):
None,
)
def all_to_all_SBH(
def all_to_all_BHSD(
input_: torch.Tensor,
scatter_dim: int = 1,
gather_dim: int = 0,
gather_dim: int = 2,
):
return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all)
@@ -99,6 +99,55 @@ def all_to_all_BSND(
):
return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all)
class _AllGather(torch.autograd.Function):
"""All-gather communication with autograd support.
Args:
input_: input tensor
dim: dimension along which to concatenate
"""
@staticmethod
def forward(ctx, input_, dim):
ctx.dim = dim
world_size = nccl_info.world_size
group = nccl_info.group
input_size = list(input_.size())
ctx.input_size = input_size[dim]
tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
input_ = input_.contiguous()
dist.all_gather(tensor_list, input_, group=group)
output = torch.cat(tensor_list, dim=dim)
return output
@staticmethod
def backward(ctx, grad_output):
world_size = nccl_info.world_size
rank = nccl_info.rank
dim = ctx.dim
input_size = ctx.input_size
sizes = [input_size] * world_size
grad_input_list = torch.split(grad_output, sizes, dim=dim)
grad_input = grad_input_list[rank]
return grad_input, None
def all_gather_BHSD(input_: torch.Tensor, dim: int = 1):
"""Performs an all-gather operation on the input tensor along the specified dimension.
Args:
input_ (torch.Tensor): Input tensor of shape [B, H, S, D].
dim (int, optional): Dimension along which to concatenate. Defaults to 1.
Returns:
torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'.
"""
return _AllGather.apply(input_, dim)
def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
fastvideo/utils/parallel_states.py
-1
View File
@@ -28,7 +28,6 @@ def initialize_sequence_parallel_group(sequence_parallel_size):
rank = int(os.getenv('RANK', '0'))
world_size = int(os.getenv("WORLD_SIZE", '1'))
assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(world_size, sequence_parallel_size)
# hccl
nccl_info.world_size = sequence_parallel_size
nccl_info.rank = rank
num_sequence_parallel_groups: int = world_size // sequence_parallel_size
infenrece_mochi.sh
+1 -1
View File
@@ -8,4 +8,4 @@ python fastvideo/sample/sample_t2v_mochi.py \
--guidance_scale 4.5 \
--seed 12346 \
--transformer_path data/outputs/debug_1e-5/checkpoint-150/transformer \
--output_path step150__79.mp4
--output_path step150__79.mp4
infenrece_mochi_sp.sh
+14
View File
@@ -0,0 +1,14 @@
#!/bin/bash
num_gpus=2
torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
fastvideo/sample/sample_t2v_mochi_sp.py \
--model_path data/mochi \
--prompt "A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere." \
--num_frames 163 \
--height 480 \
--width 848 \
--num_inference_steps 10 \
--guidance_scale 4.5 \
--seed 42
pyproject.toml
+1 -1
View File
@@ -21,7 +21,7 @@ dependencies = [
"timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0",
"scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0",
"moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3",
"watch", "gpustat"
"watch", "gpustat", "bitsandbytes", "prodigyopt"
]
[project.optional-dependencies]
t2v_debug_multi.sh
+4 -3
View File
@@ -9,13 +9,14 @@ accelerate launch \
--train_batch_size=1 \
--num_latent_t 14 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=4 \
--gradient_accumulation_steps=1 \
--max_train_steps=200 \
--learning_rate=1e-6 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--optimizer="AdamW" \
--mixed_precision="bf16" \
--checkpointing_steps=50 \
--checkpointing_steps=1 \
--allow_tf32 \
--tile_sample_stride 192 \
--ema_start_step 0 \
@@ -24,4 +25,4 @@ accelerate launch \
--enable_tiling \
--sp_size 1 \
--train_sp_batch_size 1 \
--output_dir="data/outputs/debug_1e-6"
--output_dir="data/outputs/debug"
t2v_debug_single.sh
+2
View File
@@ -8,12 +8,14 @@ accelerate launch \
--data_merge_path "data/Mochi-Synthetic-Data/merge.txt" \
--gradient_checkpointing \
--train_batch_size=1 \
--num_latent_t 14 \
--dataloader_num_workers 1 \
--gradient_accumulation_steps=1 \
--max_train_steps=1000000 \
--learning_rate=1e-4 \
--lr_scheduler="constant" \
--lr_warmup_steps=0 \
--optimizer="AdamW" \
--mixed_precision="bf16" \
--checkpointing_steps=250 \
--allow_tf32 \