diff --git a/assets/prompt.txt b/assets/prompt.txt index 65ce027f4..d10865e32 100644 --- a/assets/prompt.txt +++ b/assets/prompt.txt @@ -1,5 +1,5 @@ -A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature. Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting. +A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature. 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. A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones. A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1 diff --git a/fastvideo/attention/__init__.py b/fastvideo/attention/__init__.py deleted file mode 100644 index c0000061c..000000000 --- a/fastvideo/attention/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .distributed_attn import DistributedAttention \ No newline at end of file diff --git a/fastvideo/inference.py b/fastvideo/inference.py deleted file mode 100644 index 45ee70c06..000000000 --- a/fastvideo/inference.py +++ /dev/null @@ -1,250 +0,0 @@ -import os -import random -import time -from pathlib import Path -from typing import Dict, List, Optional, Union, Any - -import torch -from safetensors.torch import load_file as safetensors_load_file -import json - -from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE -# from fastvideo.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline -from fastvideo.pipelines.pipeline_base import DiffusionPipelineBase -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler -# from fastvideo.models.hunyuan.modules import load_model -# from fastvideo.models.hunyuan.text_encoder import TextEncoder -from fastvideo.models.hunyuan.utils.data_utils import align_to -# from fastvideo.models.hunyuan.vae import load_vae -# from fastvideo.utils.parallel_states import nccl_info -from fastvideo.inference_args import InferenceArgs -from fastvideo.platforms import current_platform -from fastvideo.logger import init_logger -from fastvideo.pipelines.loader import PipelineLoader, get_pipeline_loader -from fastvideo.distributed.parallel_state import get_sp_group - -logger = init_logger(__name__) - - -class DiffusionInference: - """ - Unified inference class that works with any diffusion pipeline. - This combines model loading, pipeline creation, and inference in a flexible way. - """ - - def __init__(self, - args: InferenceArgs, - pipeline: DiffusionPipelineBase, - default_negative_prompt: str = NEGATIVE_PROMPT): - """ - Initialize the inference class with a pipeline and args. - - Args: - args: The inference arguments - pipeline: The diffusion pipeline to use - default_negative_prompt: The default negative prompt to use - """ - self.args = args - self.pipeline = pipeline - self.default_negative_prompt = default_negative_prompt - - - - @classmethod - def load_pipeline(cls, inference_args: InferenceArgs): - """ - Create an inference instance from pretrained model components. - - Args: - inference_args: The inference arguments containing model path and other settings - - Returns: - A DiffusionInference instance ready for inference - """ - logger.debug("start loading pipeline") - models_root_path = Path(inference_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 = inference_args.output_path - os.makedirs(os.path.dirname(save_path), exist_ok=True) - - # Log which platform we're using - print(f"Using platform: {current_platform.__class__.__name__}") - assert current_platform.is_cuda(), "CUDA is not available" - - # init distributed - # _initialize_distributed() - - - # get loader for the pipeline - - # Create pipeline using the components - pipeline_loader = get_pipeline_loader(inference_args) - pipeline = pipeline_loader.load_pipeline( - inference_args=inference_args, - ) - - return cls(inference_args, pipeline, default_negative_prompt=NEGATIVE_PROMPT) - - def predict(self, prompt: str, inference_args: InferenceArgs) -> Dict[str, Any]: - """ - Run inference with the pipeline. - - Returns: - Dictionary with generated samples and metadata - """ - - out_dict = dict() - - batch_size = inference_args.batch_size - num_videos_per_prompt = inference_args.num_videos - seed = inference_args.seed - height = inference_args.height - width = inference_args.width - video_length = inference_args.num_frames - # prompt = inference_args.prompt - negative_prompt = inference_args.neg_prompt - infer_steps = inference_args.num_inference_steps - guidance_scale = inference_args.guidance_scale - flow_shift = inference_args.flow_shift - # mask_strategy = inference_args.mask_strategy - embedded_guidance_scale = inference_args.embedded_cfg_scale - - - # ======================================================================== - # Arguments: seed - # ======================================================================== - if isinstance(seed, torch.Tensor): - seed = seed.tolist() - if seed is None: - seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)] - elif isinstance(seed, int): - seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)] - elif isinstance(seed, (list, tuple)): - if len(seed) == batch_size: - seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)] - elif len(seed) == batch_size * num_videos_per_prompt: - seeds = [int(s) for s in seed] - else: - raise ValueError( - f"Length of seed must be equal to number of prompt(batch_size) or " - f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.") - else: - raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.") - # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... - generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds] - out_dict["seeds"] = seeds - - # ======================================================================== - # Arguments: target_width, target_height, target_video_length - # ======================================================================== - if width <= 0 or height <= 0 or video_length <= 0: - raise ValueError( - f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}" - ) - if (video_length - 1) % 4 != 0: - raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}") - - logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})") - - target_height = align_to(height, 16) - target_width = align_to(width, 16) - target_video_length = video_length - - out_dict["size"] = (target_height, target_width, target_video_length) - - # ======================================================================== - # Arguments: prompt, new_prompt, negative_prompt - # ======================================================================== - if not isinstance(prompt, str): - raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") - prompt = [prompt.strip()] - - # negative prompt - if negative_prompt is None or negative_prompt == "": - negative_prompt = self.default_negative_prompt - if not isinstance(negative_prompt, str): - raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}") - negative_prompt = [negative_prompt.strip()] - - # ======================================================================== - # Scheduler - # ======================================================================== - scheduler = FlowMatchDiscreteScheduler( - shift=flow_shift, - reverse=self.args.flow_reverse, - solver=self.args.flow_solver, - ) - self.pipeline.scheduler = scheduler - - if "884" in self.args.vae: - latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8] - elif "888" in self.args.vae: - latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8] - n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - - # ======================================================================== - # Print infer args - # ======================================================================== - debug_str = f""" - height: {target_height} - width: {target_width} - video_length: {target_video_length} - prompt: {prompt} - neg_prompt: {negative_prompt} - seed: {seed} - infer_steps: {infer_steps} - num_videos_per_prompt: {num_videos_per_prompt} - guidance_scale: {guidance_scale} - n_tokens: {n_tokens} - flow_shift: {flow_shift} - embedded_guidance_scale: {embedded_guidance_scale}""" - logger.info(debug_str) - # return - # sp_group = get_sp_group() - # local_rank = sp_group.rank - import os - local_rank = os.environ.get("LOCAL_RANK", 0) - device = torch.device(f"cuda:{local_rank}") - batch = ForwardBatch( - prompt=prompt, - negative_prompt=negative_prompt, - num_videos_per_prompt=num_videos_per_prompt, - height=inference_args.height, - width=inference_args.width, - num_frames=inference_args.num_frames, - num_inference_steps=inference_args.num_inference_steps, - guidance_scale=inference_args.guidance_scale, - generator=generator, - eta=0.0, - n_tokens=n_tokens, - data_type="video" if inference_args.num_frames > 1 else "image", - device=device, - extra={}, # Any additional parameters - ) - - print('===============================================') - print(batch) - print('===============================================') - print('===============================================') - print('===============================================') - print(inference_args) - - # ======================================================================== - # Pipeline inference - # ======================================================================== - start_time = time.time() - samples = self.pipeline.forward( - batch=batch, - inference_args=inference_args, - )[0] - out_dict["samples"] = samples - out_dict["prompts"] = prompt - - gen_time = time.time() - start_time - logger.info(f"Success, time: {gen_time}") - - return out_dict \ No newline at end of file diff --git a/fastvideo/layers/pooler.py b/fastvideo/layers/pooler.py deleted file mode 100644 index 0012636ef..000000000 --- a/fastvideo/layers/pooler.py +++ /dev/null @@ -1,322 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from enum import IntEnum -from typing import List, Optional, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -from transformers import PretrainedConfig -from typing_extensions import assert_never - -from vllm.config import PoolerConfig -from vllm.model_executor.pooling_metadata import (PoolingMetadata, - PoolingTensors) -from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput -from vllm.transformers_utils.config import ( - get_cross_encoder_activation_function) - - -class PoolingType(IntEnum): - """Enumeration for different types of pooling methods.""" - LAST = 0 - ALL = 1 - CLS = 2 - STEP = 3 - MEAN = 4 - - -class SimplePooler(nn.Module): - """A layer that pools specific information from hidden states. - - This layer does the following: - 1. Extracts specific tokens or aggregates data based on pooling method. - 2. Normalizes output if specified. - 3. Returns structured results as `PoolerOutput`. - - Attributes: - pooling_type: The type of pooling to use. - normalize: Whether to normalize the pooled data. - """ - - @staticmethod - def from_pooling_type( - pooling_type: PoolingType, - *, - normalize: bool, - softmax: bool, - step_tag_id: Optional[int] = None, - returned_token_ids: Optional[List[int]] = None, - ) -> "SimplePooler": - if pooling_type == PoolingType.LAST: - assert step_tag_id is None and returned_token_ids is None - return LastPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.ALL: - assert step_tag_id is None and returned_token_ids is None - return AllPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.CLS: - assert step_tag_id is None and returned_token_ids is None - return CLSPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.MEAN: - assert step_tag_id is None and returned_token_ids is None - return MeanPool(normalize=normalize, softmax=softmax) - if pooling_type == PoolingType.STEP: - return StepPool(normalize=normalize, - softmax=softmax, - step_tag_id=step_tag_id, - returned_token_ids=returned_token_ids) - - assert_never(pooling_type) - - def __init__(self, *, normalize: bool, softmax: bool) -> None: - super().__init__() - - self.head = PoolerHead(normalize=normalize, softmax=softmax) - - def get_prompt_lens( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> torch.Tensor: - return PoolingTensors.from_pooling_metadata( - pooling_metadata, hidden_states.device).prompt_lens - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - raise NotImplementedError - - def build_output(self, data: torch.Tensor) -> PoolingSequenceGroupOutput: - return PoolingSequenceGroupOutput(data) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - pooled_data = self.extract_states(hidden_states, pooling_metadata) - pooled_data = self.head(pooled_data) - pooled_outputs = [self.build_output(data) for data in pooled_data] - return PoolerOutput(outputs=pooled_outputs) - - -class CLSPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - first_token_flat_indices = torch.zeros_like(prompt_lens) - first_token_flat_indices[1:] += torch.cumsum(prompt_lens, dim=0)[:-1] - return hidden_states[first_token_flat_indices] - - -class LastPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - last_token_flat_indices = torch.cumsum(prompt_lens, dim=0) - 1 - return hidden_states[last_token_flat_indices] - - -class AllPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - offset = 0 - pooled_data = list[torch.Tensor]() - for prompt_len in prompt_lens: - pooled_data.append(hidden_states[offset:offset + prompt_len]) - offset += prompt_len - - return pooled_data - - -class MeanPool(SimplePooler): - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - cumsum = torch.cumsum(hidden_states, dim=0) - start_indices = torch.cat([ - torch.tensor([0], device=hidden_states.device), - torch.cumsum(prompt_lens[:-1], dim=0) - ]) - end_indices = torch.cumsum(prompt_lens, dim=0) - return (cumsum[end_indices - 1] - cumsum[start_indices] + - hidden_states[start_indices]) / prompt_lens.unsqueeze(1) - - -class StepPool(SimplePooler): - - def __init__( - self, - *, - normalize: bool, - softmax: bool, - step_tag_id: Optional[int] = None, - returned_token_ids: Optional[List[int]] = None, - ): - super().__init__(normalize=normalize, softmax=softmax) - - self.step_tag_id = step_tag_id - self.returned_token_ids = returned_token_ids - - def extract_states( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> Union[list[torch.Tensor], torch.Tensor]: - prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata) - - returned_token_ids = self.returned_token_ids - if returned_token_ids is not None and len(returned_token_ids) > 0: - hidden_states = hidden_states[:, returned_token_ids] - - step_tag_id = self.step_tag_id - - offset = 0 - pooled_data = list[torch.Tensor]() - for prompt_len, seq_data_i in zip(prompt_lens, - pooling_metadata.seq_data.values()): - pooled_data_i = hidden_states[offset:offset + prompt_len] - if step_tag_id is not None: - token_ids = torch.tensor(seq_data_i.prompt_token_ids) - pooled_data_i = pooled_data_i[token_ids == step_tag_id] - - offset += prompt_len - pooled_data.append(pooled_data_i) - - return pooled_data - - -class PoolerHead(nn.Module): - - def __init__(self, *, normalize: bool, softmax: bool) -> None: - super().__init__() - - self.normalize = normalize - self.softmax = softmax - - def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor]): - if self.normalize: - if isinstance(pooled_data, list): - pooled_data = [ - F.normalize(data, p=2, dim=1) for data in pooled_data - ] - else: - pooled_data = F.normalize(pooled_data, p=2, dim=1) - - if self.softmax: - if isinstance(pooled_data, list): - pooled_data = [F.softmax(data, dim=-1) for data in pooled_data] - else: - pooled_data = F.softmax(pooled_data, dim=-1) - - return pooled_data - - -class Pooler(nn.Module): - - @classmethod - def from_config_with_defaults( - cls, - pooler_config: PoolerConfig, - pooling_type: PoolingType, - normalize: bool, - softmax: bool, - step_tag_id: Optional[int] = None, - returned_token_ids: Optional[List[int]] = None, - ) -> SimplePooler: - return SimplePooler.from_pooling_type( - pooling_type=PoolingType[pooler_config.pooling_type] - if pooler_config.pooling_type is not None else pooling_type, - normalize=pooler_config.normalize - if pooler_config.normalize is not None else normalize, - softmax=pooler_config.softmax - if pooler_config.softmax is not None else softmax, - step_tag_id=pooler_config.step_tag_id - if pooler_config.step_tag_id is not None else step_tag_id, - returned_token_ids=pooler_config.returned_token_ids - if pooler_config.returned_token_ids is not None else - returned_token_ids, - ) - - -class CrossEncodingPooler(nn.Module): - """A layer that pools specific information from hidden states. - - This layer does the following: - 1. Extracts specific tokens or aggregates data based on pooling method. - 2. Normalizes output if specified. - 3. Returns structured results as `PoolerOutput`. - - Attributes: - pooling_type: The type of pooling to use. - normalize: Whether to normalize the pooled data. - """ - - def __init__( - self, - config: PretrainedConfig, - classifier: nn.Module, - pooler: Optional[nn.Module] = None, - ): - super().__init__() - self.classifier = classifier - self.pooler = pooler - self.default_activation_function = \ - get_cross_encoder_activation_function(config) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - """Pools sentence pair scores from the hidden_states.""" - - prompt_lens = PoolingTensors.from_pooling_metadata( - pooling_metadata, hidden_states.device).prompt_lens - - offset = 0 - pooled_data_lst = [] - for prompt_len in prompt_lens: - pooled_data_i = hidden_states[offset:offset + prompt_len] - - if self.pooler is not None: - final_shape_tensor = self.pooler(pooled_data_i) - else: - final_shape_tensor = self.classifier(pooled_data_i) - - pooled_data_lst.append(final_shape_tensor) - offset += prompt_len - - pooled_output = torch.stack(pooled_data_lst) - - if self.pooler is not None: - # apply classifier once on the full batch if possible - pooled_output = self.classifier(pooled_output) - - scores = self.default_activation_function(pooled_output).squeeze(-1) - - pooled_outputs = [PoolingSequenceGroupOutput(data) for data in scores] - return PoolerOutput(outputs=pooled_outputs) diff --git a/fastvideo/layers/resampler.py b/fastvideo/layers/resampler.py deleted file mode 100644 index 4c9860006..000000000 --- a/fastvideo/layers/resampler.py +++ /dev/null @@ -1,269 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py -# https://huggingface.co/Qwen/Qwen-7B/blob/main/modeling_qwen.py -# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 -# -# Copyright 2023 The Qwen team. -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Shared resampler perceiver network used in multimodal models and -related helpers for sincos positional embeddings. - -Example models: Qwen (Qwen-VL), MiniCPM-V 2.0 -""" -import math -from functools import partial -from typing import Callable, Optional, Tuple, Union - -import numpy as np -import torch -import torch.nn.functional as F -from torch import nn - -from vllm.model_executor.layers.linear import ReplicatedLinear -from vllm.model_executor.layers.quantization import QuantizationConfig - -DEFAULT_LN = partial(nn.LayerNorm, eps=1e-6) - - -def get_abs_pos(abs_pos: torch.Tensor, tgt_size: Union[torch.Tensor, - int]) -> torch.Tensor: - # abs_pos: L, C - # tgt_size: (H, W) - # return: M, C - src_size = int(math.sqrt(abs_pos.size(0))) - dtype = abs_pos.dtype - if isinstance(tgt_size, int): - tgt_size = (tgt_size, tgt_size) - if (src_size == tgt_size[0] and src_size == tgt_size[1]): - return abs_pos - return (F.interpolate( - abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2), - size=(tgt_size[0], tgt_size[1]), - mode="bicubic", - align_corners=False, - ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)) - - -# sin/cos positional embedding helpers are adapted from: -# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 -def get_1d_sincos_pos_embed_from_grid( - embed_dim: int, pos: np.ndarray, - version: Tuple[int, int] = (2, 0)) -> torch.Tensor: - """ - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size (M,) / (H, W) - out: (M, D) / (H, W, D) - """ - assert embed_dim % 2 == 0 - omega = np.arange(embed_dim // 2, dtype=np.float32) - omega /= embed_dim / 2.0 - omega = 1.0 / 10000**omega # (D/2,) - - if version == (2, 0): - pos = pos.reshape(-1) # (M,) - out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product - emb_sin = np.sin(out) # (M, D/2) - emb_cos = np.cos(out) # (M, D/2) - emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) - else: - out = np.einsum("hw,d->hwd", pos, omega) # (H, W, D/2), outer product - emb_sin = np.sin(out) # (H, W, D/2) - emb_cos = np.cos(out) # (H, W, D/2) - emb = np.concatenate([emb_sin, emb_cos], axis=-1) # (H, W, D) - return emb - - -def get_2d_sincos_pos_embed_from_grid( - embed_dim: int, grid: np.ndarray, - version: Tuple[int, int] = (2, 0)) -> torch.Tensor: - assert embed_dim % 2 == 0 - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid( - embed_dim // 2, grid[0], version) # (H*W, D/2) or (H, W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid( - embed_dim // 2, grid[1], version) # (H*W, D/2) or (H, W, D/2) - - if version == (2, 0): - emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) - else: - emb = np.concatenate([emb_h, emb_w], axis=-1) # (H, W, D) - return emb - - -def get_2d_sincos_pos_embed( - embed_dim: int, - grid_size: Union[int, Tuple[int, int]], - cls_token: bool = False, - version: Tuple[int, int] = (2, 0), -) -> torch.Tensor: - """ - grid_size: int of the grid height and width - return: - pos_embed: [grid_size*grid_size, embed_dim] or - [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) - """ - if isinstance(grid_size, int): - grid_h_size, grid_w_size = grid_size, grid_size - else: - grid_h_size, grid_w_size = grid_size[0], grid_size[1] - - grid_h = np.arange(grid_h_size, dtype=np.float32) - grid_w = np.arange(grid_w_size, dtype=np.float32) - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) - assert isinstance(grid, np.ndarray) and \ - grid.shape == (2, grid_h_size, grid_w_size) - - if version == (2, 0): - grid = grid.reshape([2, 1, grid_h_size, grid_w_size]) - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version) - if cls_token: - pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], - axis=0) - else: - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, version) - return pos_embed - - -class BaseResampler(nn.Module): - """ - A 2D perceiver-resampler network with one cross attention layers by - (grid_size**2) learnable queries and 2d sincos pos_emb. - Outputs: - A tensor with the shape of (grid_size**2, embed_dim) - """ - - def __init__(self, - num_queries: int, - embed_dim: int, - num_heads: int, - kv_dim: Optional[int] = None, - norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN, - do_post_projection: bool = True, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "") -> None: - super().__init__() - - self.num_queries = num_queries - self.embed_dim = embed_dim - self.num_heads = num_heads - - self.query = nn.Parameter(torch.empty(self.num_queries, embed_dim)) - - if kv_dim is not None and kv_dim != embed_dim: - self.kv_proj = ReplicatedLinear(kv_dim, - embed_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.kv_proj") - else: - # Maintain the same return value with ReplicatedLinear.forward - self.kv_proj = lambda *args, **kwargs: ( # type: ignore # noqa - nn.Identity()(*args, **kwargs), - None, - ) - self.attn = nn.MultiheadAttention(embed_dim, num_heads) - self.ln_q = norm_layer(embed_dim) - self.ln_kv = norm_layer(embed_dim) - self.do_post_projection = do_post_projection - self.ln_post = norm_layer(embed_dim) if do_post_projection else None - self.proj = nn.Parameter( - (embed_dim**-0.5) * - torch.empty(embed_dim, embed_dim)) if do_post_projection else None - - def _repeat(self, query, N: int): - return query.unsqueeze(1).repeat(1, N, 1) - - -class Resampler2(BaseResampler): - """Resampler-perceiver network to be used for a variety of model types, - e.g., Qwen-vl / Minicpmv 2.0. The main difference is the addition of the - do_post_projection arg, which indicates whether or not there should be - a post layer normalization and projector after the attention. This is - present in minicpmv2.0, but not qwen-vl. - """ - - def __init__(self, - grid_size: int, - embed_dim: int, - num_heads: int, - kv_dim: Optional[int] = None, - norm_layer: Callable[[int], nn.LayerNorm] = DEFAULT_LN, - adaptive: bool = False, - do_post_projection: bool = True, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "") -> None: - super().__init__(grid_size**2, - embed_dim, - num_heads, - kv_dim, - norm_layer, - do_post_projection=do_post_projection, - quant_config=quant_config, - prefix=prefix) - - self.adaptive = adaptive - pos_embed_arr = get_2d_sincos_pos_embed(embed_dim, - grid_size, - version=(2, 0)) - - self.pos_embed = nn.Parameter( - torch.from_numpy(pos_embed_arr).requires_grad_(False)) - - def forward( - self, - x: torch.Tensor, - tgt_sizes: Optional[torch.Tensor] = None, - attn_mask: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if tgt_sizes is None: - tgt_sizes = int(math.sqrt(x.size(1))) - if self.adaptive: - pos_embed_arr = get_2d_sincos_pos_embed(self.embed_dim, - tgt_sizes, - version=(2, 0)) - pos_embed = torch.from_numpy(pos_embed_arr).to(device=x.device, - dtype=x.dtype) - else: - pos_embed = get_abs_pos(self.pos_embed, - tgt_sizes).to(device=x.device, - dtype=x.dtype) - - x, _ = self.kv_proj(x) - x = self.ln_kv(x).permute(1, 0, 2) - - N = x.shape[1] - q = self.ln_q(self.query) - out = self.attn( - self._repeat(q, N) + self.pos_embed.unsqueeze(1), - x + pos_embed.unsqueeze(1), - x, - attn_mask=attn_mask, - )[0] - x = out.permute(1, 0, 2) - if self.do_post_projection: - x = self.ln_post(x) - x = x @ self.proj - return x diff --git a/fastvideo/loader/loader.py b/fastvideo/loader/loader.py deleted file mode 100644 index 1db66e20e..000000000 --- a/fastvideo/loader/loader.py +++ /dev/null @@ -1,452 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# Adapted from vllm -# Copyright 2023 The vLLM Authors. -# Copyright 2025 The FastVideo Authors. - -# ruff: noqa: SIM117 -import collections -import copy -import dataclasses -import fnmatch -import glob -import inspect -import itertools -import math -import os -import time -import warnings -from abc import ABC, abstractmethod -from contextlib import contextmanager -from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional, - Tuple, cast) - -import gguf -import huggingface_hub -import numpy as np -import torch -from huggingface_hub import HfApi -from torch import nn -from transformers import AutoModelForCausalLM -from transformers.utils import SAFE_WEIGHTS_INDEX_NAME - -from vllm.attention import Attention -from vllm.config import (LoadConfig, LoadFormat, ModelConfig, ParallelConfig, - VllmConfig, set_current_vllm_config) -from vllm.distributed import (get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size) -from vllm.envs import VLLM_USE_MODELSCOPE -from vllm.logger import init_logger -from vllm.model_executor.layers.linear import (LinearBase, - MergedColumnParallelLinear, - QKVParallelLinear, - ReplicatedLinear, - RowParallelLinear) -from vllm.model_executor.layers.quantization.base_config import ( - QuantizeMethodBase) -from vllm.model_executor.model_loader.tensorizer import ( - TensorizerConfig, is_vllm_tensorized, load_with_tensorizer, - serialize_vllm_model, tensorizer_weights_iterator) -from vllm.model_executor.model_loader.utils import (ParamMapping, - configure_quant_config, - get_model_architecture, - set_default_torch_dtype) -from vllm.model_executor.model_loader.weight_utils import ( - download_safetensors_index_file_from_hf, download_weights_from_hf, - filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, - get_gguf_extra_tensor_names, get_lock, gguf_quant_weights_iterator, - initialize_dummy_weights, np_cache_weights_iterator, pt_weights_iterator, - runai_safetensors_weights_iterator, safetensors_weights_iterator) -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.transformers_utils.s3_utils import glob as s3_glob -from vllm.transformers_utils.utils import is_s3 -from vllm.utils import is_pin_memory_available - - -@contextmanager -def device_loading_context(module: torch.nn.Module, - target_device: torch.device): - if target_device.type == "cpu": - # If target is CPU, no need to move anything - yield module - return - - original_device_states: Dict[str, torch.device] = {} - - # Store original device states and move parameters to GPU if they're on CPU - for name, p in module.named_parameters(): - if p.device.type == "cpu": - original_device_states[name] = p.device - p.data = p.data.to(target_device) - # Parameters already on target device are not touched - - try: - yield module - - finally: - # Restore parameters to their original devices, ignoring new parameters - pin_memory = is_pin_memory_available() - for name, p in module.named_parameters(): - if name in original_device_states: - original_device: torch.device = original_device_states[name] - if original_device.type == "cpu": - # `torch.empty_like` does not support `pin_memory` argument - cpu_data = torch.empty_strided( - size=p.data.size(), - stride=p.data.stride(), - dtype=p.data.dtype, - layout=p.data.layout, - device="cpu", - pin_memory=pin_memory, - ) - cpu_data.copy_(p.data) - p.data = cpu_data - else: - p.data = p.data.to(original_device) - # New parameters or parameters already on target device are untouched - - -logger = init_logger(__name__) - - -def _initialize_model( - vllm_config: VllmConfig, - *, - prefix: str = "", -) -> nn.Module: - """Initialize a model with the given configurations.""" - model_config = vllm_config.model_config - model_class, _ = get_model_architecture(model_config) - - if vllm_config.quant_config is not None: - configure_quant_config(vllm_config.quant_config, model_class) - - signatures = inspect.signature(model_class.__init__) - all_params = [param.name for param in signatures.parameters.values()] - if "vllm_config" in all_params and "prefix" in all_params: - # new-style model class - with set_current_vllm_config(vllm_config, check_compile=True): - return model_class(vllm_config=vllm_config, prefix=prefix) - - msg = ("vLLM model class should accept `vllm_config` and `prefix` as " - "input arguments. Possibly you have an old-style model class" - " registered from out of tree and it is used for new vLLM version. " - "Check https://docs.vllm.ai/en/latest/design/arch_overview.html " - "for the design and update the model class accordingly.") - warnings.warn(msg, DeprecationWarning, stacklevel=2) - - logger.warning( - "Trying to guess the arguments for old-style model class %s", - model_class, - ) - # try to be compatible with old-style model class - kwargs = {} - if "prefix" in all_params: - kwargs["prefix"] = prefix - if "config" in all_params: - kwargs["config"] = model_config.hf_config - if "cache_config" in all_params: - kwargs["cache_config"] = vllm_config.cache_config - if "quant_config" in all_params: - kwargs["quant_config"] = vllm_config.quant_config - if "lora_config" in all_params: - kwargs["lora_config"] = vllm_config.lora_config - if "scheduler_config" in all_params: - kwargs["scheduler_config"] = vllm_config.scheduler_config - with set_current_vllm_config(vllm_config, check_compile=True): - return model_class(**kwargs) - - -def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig, - target_device: torch.device) -> None: - for _, module in model.named_modules(): - quant_method = getattr(module, "quant_method", None) - if isinstance(quant_method, QuantizeMethodBase): - # When quant methods need to process weights after loading - # (for repacking, quantizing, etc), they expect parameters - # to be on the global target device. This scope is for the - # case where cpu offloading is used, where we will move the - # parameters onto device for processing and back off after. - with device_loading_context(module, target_device): - quant_method.process_weights_after_loading(module) - - # Currently only used by MLA. - # NOTE: This intentionally happens after other modules so we can easily - # decompress the weights for MLA. - for _, module in model.named_modules(): - if isinstance(module, Attention) and \ - hasattr(module, "process_weights_after_loading"): - # TODO(lucas): see if there is a way to unify the signatures - # of process_weights_after_loading - module.process_weights_after_loading(model_config.dtype) - - -class BaseModelLoader(ABC): - """Base class for model loaders.""" - - def __init__(self, load_config: LoadConfig): - self.load_config = load_config - - @abstractmethod - def download_model(self, model_config: ModelConfig) -> None: - """Download a model so that it can be immediately loaded.""" - raise NotImplementedError - - @abstractmethod - def load_model(self, *, vllm_config: VllmConfig) -> nn.Module: - """Load a model with the given configurations.""" - raise NotImplementedError - - -class DefaultModelLoader(BaseModelLoader): - """Model loader that can load different file types from disk.""" - - @dataclasses.dataclass - class Source: - """A source for weights.""" - - model_or_path: str - """The model ID or path.""" - - revision: Optional[str] - """The optional model revision.""" - - prefix: str = "" - """A prefix to prepend to all weights.""" - - fall_back_to_pt: bool = True - """Whether .pt weights can be used.""" - - allow_patterns_overrides: Optional[list[str]] = None - """If defined, weights will load exclusively using these patterns.""" - - counter_before_loading_weights: float = 0.0 - counter_after_loading_weights: float = 0.0 - - def __init__(self, load_config: LoadConfig): - super().__init__(load_config) - if load_config.model_loader_extra_config: - raise ValueError(f"Model loader extra config is not supported for " - f"load format {load_config.load_format}") - - def _maybe_download_from_modelscope( - self, model: str, revision: Optional[str]) -> Optional[str]: - """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True. - - Returns the path to the downloaded model, or None if the model is not - downloaded from ModelScope.""" - if VLLM_USE_MODELSCOPE: - # download model from ModelScope hub, - # lazy import so that modelscope is not required for normal use. - # pylint: disable=C. - from modelscope.hub.snapshot_download import snapshot_download - - if not os.path.exists(model): - # Use file lock to prevent multiple processes from - # downloading the same model weights at the same time. - with get_lock(model, self.load_config.download_dir): - model_path = snapshot_download( - model_id=model, - cache_dir=self.load_config.download_dir, - local_files_only=huggingface_hub.constants. - HF_HUB_OFFLINE, - revision=revision, - ignore_file_pattern=self.load_config.ignore_patterns, - ) - else: - model_path = model - return model_path - return None - - def _prepare_weights( - self, - model_name_or_path: str, - revision: Optional[str], - fall_back_to_pt: bool, - allow_patterns_overrides: Optional[list[str]], - ) -> Tuple[str, List[str], bool]: - """Prepare weights for the model. - - If the model is not local, it will be downloaded.""" - model_name_or_path = (self._maybe_download_from_modelscope( - model_name_or_path, revision) or model_name_or_path) - - is_local = os.path.isdir(model_name_or_path) - load_format = self.load_config.load_format - use_safetensors = False - index_file = SAFE_WEIGHTS_INDEX_NAME - # Some quantized models use .pt files for storing the weights. - if load_format == LoadFormat.AUTO: - allow_patterns = ["*.safetensors", "*.bin"] - elif load_format == LoadFormat.SAFETENSORS: - use_safetensors = True - allow_patterns = ["*.safetensors"] - elif load_format == LoadFormat.MISTRAL: - use_safetensors = True - allow_patterns = ["consolidated*.safetensors"] - index_file = "consolidated.safetensors.index.json" - elif load_format == LoadFormat.PT: - allow_patterns = ["*.pt"] - elif load_format == LoadFormat.NPCACHE: - allow_patterns = ["*.bin"] - else: - raise ValueError(f"Unknown load_format: {load_format}") - - if fall_back_to_pt: - allow_patterns += ["*.pt"] - - if allow_patterns_overrides is not None: - allow_patterns = allow_patterns_overrides - - if not is_local: - hf_folder = download_weights_from_hf( - model_name_or_path, - self.load_config.download_dir, - allow_patterns, - revision, - ignore_patterns=self.load_config.ignore_patterns, - ) - else: - hf_folder = model_name_or_path - - hf_weights_files: List[str] = [] - for pattern in allow_patterns: - hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) - if len(hf_weights_files) > 0: - if pattern == "*.safetensors": - use_safetensors = True - break - - if use_safetensors: - # For models like Mistral-7B-Instruct-v0.3 - # there are both sharded safetensors files and a consolidated - # safetensors file. Using both breaks. - # Here, we download the `model.safetensors.index.json` and filter - # any files not found in the index. - if not is_local: - download_safetensors_index_file_from_hf( - model_name_or_path, - index_file, - self.load_config.download_dir, - revision, - ) - hf_weights_files = filter_duplicate_safetensors_files( - hf_weights_files, hf_folder, index_file) - else: - hf_weights_files = filter_files_not_needed_for_inference( - hf_weights_files) - - if len(hf_weights_files) == 0: - raise RuntimeError( - f"Cannot find any model weights with `{model_name_or_path}`") - - return hf_folder, hf_weights_files, use_safetensors - - def _get_weights_iterator( - self, source: "Source" - ) -> Generator[Tuple[str, torch.Tensor], None, None]: - """Get an iterator for the model weights based on the load format.""" - hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( - source.model_or_path, source.revision, source.fall_back_to_pt, - source.allow_patterns_overrides) - if self.load_config.load_format == LoadFormat.NPCACHE: - # Currently np_cache only support *.bin checkpoints - assert use_safetensors is False - weights_iterator = np_cache_weights_iterator( - source.model_or_path, - self.load_config.download_dir, - hf_folder, - hf_weights_files, - ) - elif use_safetensors: - weights_iterator = safetensors_weights_iterator(hf_weights_files) - else: - weights_iterator = pt_weights_iterator(hf_weights_files) - - if current_platform.is_tpu(): - # In PyTorch XLA, we should call `xm.mark_step` frequently so that - # not too many ops are accumulated in the XLA program. - import torch_xla.core.xla_model as xm - - def _xla_weights_iterator(iterator: Generator): - for weights in iterator: - yield weights - xm.mark_step() - - weights_iterator = _xla_weights_iterator(weights_iterator) - - if self.counter_before_loading_weights == 0.0: - self.counter_before_loading_weights = time.perf_counter() - # Apply the prefix. - return ((source.prefix + name, tensor) - for (name, tensor) in weights_iterator) - - def _get_all_weights( - self, - model_config: ModelConfig, - model: nn.Module, - ) -> Generator[Tuple[str, torch.Tensor], None, None]: - primary_weights = DefaultModelLoader.Source( - model_config.model, - model_config.revision, - prefix="", - fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", - True), - allow_patterns_overrides=getattr(model, "allow_patterns_overrides", - None), - ) - yield from self._get_weights_iterator(primary_weights) - - secondary_weights = cast( - Iterable[DefaultModelLoader.Source], - getattr(model, "secondary_weights", ()), - ) - for source in secondary_weights: - yield from self._get_weights_iterator(source) - - def download_model(self, model_config: ModelConfig) -> None: - self._prepare_weights(model_config.model, - model_config.revision, - fall_back_to_pt=True, - allow_patterns_overrides=None) - - def load_model(self, vllm_config: VllmConfig) -> nn.Module: - device_config = vllm_config.device_config - model_config = vllm_config.model_config - target_device = torch.device(device_config.device) - with set_default_torch_dtype(model_config.dtype): - with target_device: - model = _initialize_model(vllm_config=vllm_config) - - weights_to_load = {name for name, _ in model.named_parameters()} - loaded_weights = model.load_weights( - self._get_all_weights(model_config, model)) - self.counter_after_loading_weights = time.perf_counter() - logger.info( - "Loading weights took %.2f seconds", - self.counter_after_loading_weights - - self.counter_before_loading_weights) - # We only enable strict check for non-quantized models - # that have loaded weights tracking currently. - if model_config.quantization is None and loaded_weights is not None: - weights_not_loaded = weights_to_load - loaded_weights - if weights_not_loaded: - raise ValueError( - "Following weights were not initialized from " - f"checkpoint: {weights_not_loaded}") - - _process_weights_after_loading(model, model_config, target_device) - - return model.eval() - - - - - -def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: - """Get a model loader based on the load format.""" - if isinstance(load_config.load_format, type): - return load_config.load_format(load_config) - - return DefaultModelLoader(load_config) \ No newline at end of file diff --git a/fastvideo/model_loader.py b/fastvideo/model_loader.py deleted file mode 100644 index 0b90d183e..000000000 --- a/fastvideo/model_loader.py +++ /dev/null @@ -1,273 +0,0 @@ -""" -Model loader for diffusion models. - -This module provides a base class and implementations for loading model components -needed by diffusion pipelines. -""" - -from abc import ABC, abstractmethod -from typing import Any, Dict, Optional, Type, Union -import os -import torch -import traceback - -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger - -logger = init_logger(__name__) - - -class ModelLoader(ABC): - """ - Abstract base class for model loaders. - - A model loader is responsible for loading the components needed by a diffusion - pipeline, such as the UNet, VAE, text encoder, etc. - """ - - @abstractmethod - def load_components( - self, - inference_args: InferenceArgs, - ) -> Dict[str, Any]: - """ - Load the components needed by a diffusion pipeline. - - Args: - inference_args: The inference arguments. - - Returns: - A dictionary of component names to component objects. - """ - pass - - -class HunYuanModelLoader(ModelLoader): - """ - Model loader for HunYuan models. - - This loader loads the components needed by HunYuan diffusion pipelines. - """ - - def load_components( - self, - inference_args: InferenceArgs, - ) -> Dict[str, Any]: - """ - Load the components needed by a HunYuan diffusion pipeline. - - Args: - inference_args: The inference arguments. - - Returns: - A dictionary of component names to component objects. - - Raises: - FileNotFoundError: If a required model file is not found. - RuntimeError: If there is an error loading a model component. - """ - try: - from fastvideo.models.unet import UNet3DConditionModel - from fastvideo.models.vae import AutoencoderKLCausal3D - from fastvideo.models.text_encoder import TextEncoder - from fastvideo.models.tokenizer import Tokenizer - from fastvideo.models.scheduler import DDIMScheduler - - logger.info(f"Loading model components from {inference_args.model_path}") - - # Check if model path exists - if not os.path.exists(inference_args.model_path): - raise FileNotFoundError(f"Model path not found: {inference_args.model_path}") - - # Check if DiT weight exists - if inference_args.dit_weight and not os.path.exists(inference_args.dit_weight): - raise FileNotFoundError(f"DiT weight not found: {inference_args.dit_weight}") - - # Load UNet - try: - logger.info("Loading UNet...") - unet = UNet3DConditionModel.from_pretrained( - inference_args.dit_weight, - subfolder="", - load_key=inference_args.load_key, - ) - logger.info("UNet loaded successfully") - except Exception as e: - logger.error(f"Error loading UNet: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load UNet: {e}") - - # Load VAE - try: - logger.info("Loading VAE...") - vae = AutoencoderKLCausal3D.from_pretrained( - inference_args.model_path, - subfolder="vae", - use_safetensors=True, - ) - logger.info("VAE loaded successfully") - except Exception as e: - logger.error(f"Error loading VAE: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load VAE: {e}") - - # Configure VAE for sequence parallelism if needed - if inference_args.vae_sp: - try: - logger.info("Configuring VAE for sequence parallelism") - vae.enable_sequence_parallel() - logger.info("VAE configured for sequence parallelism") - except Exception as e: - logger.error(f"Error configuring VAE for sequence parallelism: {e}") - logger.error(traceback.format_exc()) - logger.warning("Continuing without sequence parallelism for VAE") - - # Load text encoders - try: - logger.info("Loading primary text encoder...") - text_encoder = TextEncoder.from_pretrained( - inference_args.model_path, - subfolder="text_encoder", - use_safetensors=True, - ) - logger.info("Primary text encoder loaded successfully") - except Exception as e: - logger.error(f"Error loading primary text encoder: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load primary text encoder: {e}") - - try: - logger.info("Loading secondary text encoder...") - text_encoder_2 = TextEncoder.from_pretrained( - inference_args.model_path, - subfolder="text_encoder_2", - use_safetensors=True, - ) - logger.info("Secondary text encoder loaded successfully") - except Exception as e: - logger.error(f"Error loading secondary text encoder: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load secondary text encoder: {e}") - - # Load tokenizers - try: - logger.info("Loading primary tokenizer...") - tokenizer = Tokenizer.from_pretrained( - inference_args.model_path, - subfolder="tokenizer", - ) - logger.info("Primary tokenizer loaded successfully") - except Exception as e: - logger.error(f"Error loading primary tokenizer: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load primary tokenizer: {e}") - - try: - logger.info("Loading secondary tokenizer...") - tokenizer_2 = Tokenizer.from_pretrained( - inference_args.model_path, - subfolder="tokenizer_2", - ) - logger.info("Secondary tokenizer loaded successfully") - except Exception as e: - logger.error(f"Error loading secondary tokenizer: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load secondary tokenizer: {e}") - - # Load scheduler - try: - logger.info("Loading scheduler...") - scheduler = DDIMScheduler.from_pretrained( - inference_args.model_path, - subfolder="scheduler", - ) - logger.info("Scheduler loaded successfully") - except Exception as e: - logger.error(f"Error loading scheduler: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to load scheduler: {e}") - - # Move models to the appropriate device - device = torch.device(f"cuda:{os.environ.get('LOCAL_RANK', '0')}") - logger.info(f"Moving models to device: {device}") - - try: - unet.to(device) - vae.to(device) - text_encoder.to(device) - text_encoder_2.to(device) - logger.info("Models moved to device successfully") - except Exception as e: - logger.error(f"Error moving models to device: {e}") - logger.error(traceback.format_exc()) - raise RuntimeError(f"Failed to move models to device: {e}") - - # Set precision - try: - logger.info(f"Setting UNet precision to {inference_args.precision}") - if inference_args.precision == "bf16": - unet = unet.to(torch.bfloat16) - elif inference_args.precision == "fp16": - unet = unet.to(torch.float16) - - logger.info(f"Setting VAE precision to {inference_args.vae_precision}") - if inference_args.vae_precision == "fp16": - vae = vae.to(torch.float16) - - logger.info(f"Setting text encoder precision to {inference_args.text_encoder_precision}") - if inference_args.text_encoder_precision == "fp16": - text_encoder = text_encoder.to(torch.float16) - - logger.info(f"Setting secondary text encoder precision to {inference_args.text_encoder_precision_2}") - if inference_args.text_encoder_precision_2 == "fp16": - text_encoder_2 = text_encoder_2.to(torch.float16) - - logger.info("Precision set successfully") - except Exception as e: - logger.error(f"Error setting precision: {e}") - logger.error(traceback.format_exc()) - logger.warning("Continuing with default precision") - - # Return the components - return { - "unet": unet, - "vae": vae, - "text_encoder": text_encoder, - "text_encoder_2": text_encoder_2, - "tokenizer": tokenizer, - "tokenizer_2": tokenizer_2, - "scheduler": scheduler, - } - except Exception as e: - logger.error(f"Unhandled error in load_components: {e}") - logger.error(traceback.format_exc()) - raise - - -# Factory function to get the appropriate model loader -def get_model_loader(model_type: str) -> Type[ModelLoader]: - """ - Get the appropriate model loader for the given model type. - - Args: - model_type: The type of model to load. - - Returns: - The model loader class. - - Raises: - ValueError: If the model type is not recognized. - """ - model_loaders = { - "hunyuan": HunYuanModelLoader, - "HYVideo-T/2-cfgdistill": HunYuanModelLoader, - } - - if model_type not in model_loaders: - available_models = list(model_loaders.keys()) - raise ValueError( - f"Model type '{model_type}' not recognized. " - f"Available model types: {available_models}" - ) - - return model_loaders[model_type] \ No newline at end of file diff --git a/fastvideo/models/__init__.py b/fastvideo/models/__init__.py deleted file mode 100644 index 2c63b0b0d..000000000 --- a/fastvideo/models/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from fastvideo.models.parameter import (BasevLLMParameter, - PackedvLLMParameter) -from fastvideo.models.utils import set_random_seed - -__all__ = [ - "set_random_seed", - "BasevLLMParameter", - "PackedvLLMParameter", -] diff --git a/fastvideo/models/dits/__init__.py b/fastvideo/models/dits/__init__.py deleted file mode 100644 index 768f632bd..000000000 --- a/fastvideo/models/dits/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ - -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT - -__all__ = ["HunyuanVideoDiT"] diff --git a/fastvideo/models/hunyuan/inference.py b/fastvideo/models/hunyuan/inference.py index c034f330c..62694e6d9 100644 --- a/fastvideo/models/hunyuan/inference.py +++ b/fastvideo/models/hunyuan/inference.py @@ -58,22 +58,6 @@ class Inference(object): args (argparse.Namespace): The arguments for the pipeline. device (int): The device for inference. Default is 0. """ - from fastvideo.pipelines.loader import get_pipeline_loader - loader = get_pipeline_loader(args) - vae, vae_kwargs, text_encoder, text_encoder_2, transformer, scheduler = \ - loader.load_components(args) - return cls( - args=args, - vae=vae, - vae_kwargs=vae_kwargs, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - model=transformer, - use_cpu_offload=args.use_cpu_offload, - device=device, - logger=logger, - parallel_args=None, - ) # ======================================================================== logger.info(f"Got text-to-video model root path: {pretrained_model_path}") @@ -223,7 +207,6 @@ class Inference(object): else: raise ValueError(f"Invalid model path: {dit_weight}") - print('model_path', model_path) if not model_path.exists(): raise ValueError(f"model_path not exists: {model_path}") logger.info(f"Loading torch model {model_path}...") @@ -313,11 +296,6 @@ class HunyuanVideoSampler(Inference): progress_bar_config=None, data_type="video", ): - print("Loading diffusion pipeline...") - # print('vae', vae) - print('text_encoder', text_encoder) - print('text_encoder_2', text_encoder_2) - print('model', model) """Load the denoising scheduler for inference.""" if scheduler is None: if args.denoise_type == "flow": diff --git a/fastvideo/models/hunyuan/modules/attenion.py b/fastvideo/models/hunyuan/modules/attenion.py index a5760c1a0..699b3ae1d 100644 --- a/fastvideo/models/hunyuan/modules/attenion.py +++ b/fastvideo/models/hunyuan/modules/attenion.py @@ -9,7 +9,6 @@ except ImportError: sliding_tile_attention = None from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad -from flash_attn import flash_attn_func from fastvideo.utils.communications import all_gather, all_to_all_4D from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info @@ -95,17 +94,14 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy= hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2) else: - txt_length = text_mask.sum().int().item() - query = torch.cat([query, encoder_query[:, :txt_length]], dim=1) - key = torch.cat([key, encoder_key[:, :txt_length]], dim=1) - value = torch.cat([value, encoder_value[:, :txt_length]], dim=1) + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) # B, S, 3, H, D - # qkv = torch.stack([query, key, value], dim=2) + qkv = torch.stack([query, key, value], dim=2) attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) - # hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) - hidden_states = flash_attn_func(query, key, value, causal=False, dropout_p=0.0, softmax_scale=None) - hidden_states = torch.cat([hidden_states, encoder_query[:, txt_length:, :]], dim=1) + hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), dim=1) diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py index 0d1c9e084..6b1706ad0 100644 --- a/fastvideo/models/hunyuan/modules/models.py +++ b/fastvideo/models/hunyuan/modules/models.py @@ -131,12 +131,12 @@ class MMDoubleStreamBlock(nn.Module): # Prepare image for attention. img_modulated = self.img_norm1(img) img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) - img_qkv = self.img_attn_qkv(img_modulated) img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) # Apply QK-Norm if needed img_q = self.img_attn_q_norm(img_q).to(img_v) img_k = self.img_attn_k_norm(img_k).to(img_v) + # Apply RoPE if needed. if freqs_cis is not None: @@ -148,12 +148,12 @@ class MMDoubleStreamBlock(nn.Module): shrink_head(freqs_cis[0], dim=0), shrink_head(freqs_cis[1], dim=0), ) + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" img_q, img_k = img_qq, img_kk - # Prepare txt for attention. txt_modulated = self.txt_norm1(txt) txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) @@ -162,6 +162,7 @@ class MMDoubleStreamBlock(nn.Module): # Apply QK-Norm if needed. txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + attn = parallel_attention( (img_q, txt_q), (img_k, txt_k), @@ -173,7 +174,9 @@ class MMDoubleStreamBlock(nn.Module): ) # attention computation end + img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:] + # Calculate the img blocks. img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate) img = img + apply_gate( @@ -540,6 +543,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): # text modulation vec = vec + self.vector_in(text_states_2) + # guidance modulation if self.guidance_embed: if guidance is None: @@ -547,6 +551,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): # our timestep_embedding is merged into guidance_in(TimestepEmbedder) vec = vec + self.guidance_in(guidance) + # Embed image and text. img = self.img_in(img) if self.text_projection == "linear": @@ -555,11 +560,13 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None) else: raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + txt_seq_len = txt.shape[1] img_seq_len = img.shape[1] freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None # --------------------- Pass through DiT blocks ------------------------ + for index, block in enumerate(self.double_blocks): double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]] img, txt = block(*double_block_args) @@ -580,9 +587,12 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): x = block(*single_block_args) if output_features and _ % output_features_stride == 0: features_list.append(x[:, :img_seq_len, ...]) + img = x[:, :img_seq_len, ...] + # ---------------------------- Final layer ------------------------------ img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + img = self.unpatchify(img, tt, th, tw) assert not return_dict, "return_dict is not supported." if output_features: diff --git a/fastvideo/models/hunyuan/modules/posemb_layers.py b/fastvideo/models/hunyuan/modules/posemb_layers.py index 314aa126d..34dcdb5d5 100644 --- a/fastvideo/models/hunyuan/modules/posemb_layers.py +++ b/fastvideo/models/hunyuan/modules/posemb_layers.py @@ -160,7 +160,6 @@ def apply_rotary_emb( cos, sin = cos.to(xq.device), sin.to(xq.device) # real * cos - imag * sin # imag * cos + real * sin - xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) else: diff --git a/fastvideo/models/hunyuan/modules/token_refiner.py b/fastvideo/models/hunyuan/modules/token_refiner.py index 1bc127fea..6aafc1dba 100644 --- a/fastvideo/models/hunyuan/modules/token_refiner.py +++ b/fastvideo/models/hunyuan/modules/token_refiner.py @@ -77,6 +77,7 @@ class IndividualTokenRefinerBlock(nn.Module): # Self-Attention attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) # FFN Layer @@ -191,9 +192,11 @@ class SingleTokenRefiner(nn.Module): else: mask_float = mask.float().unsqueeze(-1) # [b, s1, 1] context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) - context_aware_representations = self.c_embedder(context_aware_representations.to(x.dtype)) + context_aware_representations = self.c_embedder(context_aware_representations) c = timestep_aware_representations + context_aware_representations x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py index 0a14520a1..8c2bf87fb 100644 --- a/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py +++ b/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py @@ -21,15 +21,11 @@ from math import prod from typing import Dict, Optional, Tuple, Union import torch -import torch.nn as nn import torch.distributed as dist +import torch.nn as nn from diffusers.configuration_utils import ConfigMixin, register_to_config from fastvideo.utils.parallel_states import nccl_info -from fastvideo.distributed.parallel_state import get_sp_group -from fastvideo.logger import init_logger - -logger = init_logger(__name__) try: # This diffusers is modified and packed in the mirror. @@ -311,7 +307,6 @@ class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: assert len(z.shape) == 5, "The input tensor should have 5 dimensions." - logger.info(f"use_parallel: {self.use_parallel}") if self.use_parallel: return self.parallel_tiled_decode(z, return_dict=return_dict) @@ -581,12 +576,7 @@ class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): """ Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs """ - # world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group - - sp_group = get_sp_group() - world_size = sp_group.world_size - rank = sp_group.rank_in_group - + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group B, C, T, H, W = z.shape # Calculate parameters @@ -660,8 +650,6 @@ class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): gathered_results = torch.zeros_like(padded_results).repeat( world_size, *[1] * len(padded_results.shape)).contiguous( ) # use contiguous to make sure it won't copy data in the following operations - # gathered_results = sp_group.all_gather(padded_results) - # gathered_dim_metadata = sp_group.broadcast_object_list(gathered_dim_metadata, local_dim_metadata) dist.all_gather_into_tensor(gathered_results, padded_results) dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) # Process gathered results diff --git a/fastvideo/pipelines/__init__.py b/fastvideo/pipelines/__init__.py deleted file mode 100644 index 18b53bbc4..000000000 --- a/fastvideo/pipelines/__init__.py +++ /dev/null @@ -1,70 +0,0 @@ -""" -Diffusion pipelines for FastVideo. - -This package contains diffusion pipelines for generating videos and images. -""" - -from typing import Dict, Optional, Type, Any - -# First, import the registry -from fastvideo.pipelines.pipeline_registry import PipelineRegistry, register_pipeline -from fastvideo.inference_args import InferenceArgs - -# Then import the base classes -from fastvideo.pipelines.composed.composed_pipeline_base import ( - ComposedPipelineBase, - DiffusionPipelineOutput -) - -def get_pipeline_type(inference_args: InferenceArgs) -> str: - # hardcode for now - return "hunyuan_video" - -def build_pipeline(inference_args: InferenceArgs) -> ComposedPipelineBase: - """ - Only works with valid hf diffusers configs. (model_index.json) - We want to build a pipeline based on the inference args. - 1. load the correct hf config from disk or download from hub - 2. based on the config, determine the pipeline class and pipeline loader class - 3. parse the config to get the model component (vae, text_encoders, etc...) - names and paths - 4. the pipeline loader class will use the model component names and paths to load - the models. - 5. the pipeline class will be composed of the models returned by the pipeline loader - """ - # Get pipeline type - pipeline_type = get_pipeline_type(inference_args) - - if not hasattr(inference_args, 'model_type') or not inference_args.model_type: - inference_args.model_type = pipeline_type - - # TODO(will): Get pipeline loader (hardcoded for now) - from fastvideo.pipelines.loader import get_pipeline_loader - pipeline_loader = get_pipeline_loader(inference_args) - - # Load and return pipeline - return pipeline_loader.load_pipeline_v2(inference_args) - - -def list_available_pipelines() -> Dict[str, Type[Any]]: - """ - List all available pipeline types. - - Returns: - A dictionary of pipeline names to pipeline classes. - """ - return PipelineRegistry.list() - - -# Import all pipeline implementations to register them -# These imports should be at the end to avoid circular imports -from fastvideo.pipelines.implementations.hunyuan import HunyuanVideoPipeline - -__all__ = [ - "create_pipeline", - "list_available_pipelines", - "ComposedPipelineBase", - "DiffusionPipelineOutput", - "register_pipeline", - "HunyuanVideoPipeline", -] \ No newline at end of file diff --git a/fastvideo/pipelines/composed/text_to_video.py b/fastvideo/pipelines/composed/text_to_video.py deleted file mode 100644 index fbc18bb4f..000000000 --- a/fastvideo/pipelines/composed/text_to_video.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -Text-to-video pipeline implementation. - -This module contains a standard text-to-video pipeline implementation that can be -used as a base for more specific implementations. -""" - -from typing import Any, Dict, List, Optional, Union, Tuple -import torch - -from fastvideo.pipelines.composed.composed_pipeline_base import ComposedPipelineBase -from fastvideo.pipelines.stages import ( - InputValidationStage, - PromptEncodingStage, - TimestepPreparationStage, - LatentPreparationStage, - ConditioningStage, - DenoisingStage, - DecodingStage, - PostProcessingStage, -) -# Import from the main pipelines package to avoid circular imports -from fastvideo.pipelines import register_pipeline - - -@register_pipeline("text-to-video") -class TextToVideoPipeline(ComposedPipelineBase): - """ - Standard text-to-video pipeline. - - This pipeline implements a standard text-to-video diffusion process with the - following stages: - 1. Input validation - 2. Prompt encoding - 3. Timestep preparation - 4. Latent preparation - 5. Conditioning (classifier-free guidance) - 6. Denoising - 7. Decoding - 8. Post-processing - - This pipeline supports models with either a single text encoder or dual text encoders. - Different model architectures may use secondary encoders for various purposes such as - additional conditioning information, temporal control, stylistic aspects, or other - model-specific features. - """ - - is_video_pipeline = True - - def __init__( - self, - input_validation_stage: InputValidationStage, - prompt_encoding_stage: PromptEncodingStage, - timestep_preparation_stage: TimestepPreparationStage, - latent_preparation_stage: LatentPreparationStage, - conditioning_stage: ConditioningStage, - denoising_stage: DenoisingStage, - decoding_stage: DecodingStage, - post_processing_stage: Optional[PostProcessingStage] = None, - secondary_prompt_encoding_stage: Optional[PromptEncodingStage] = None, - ): - """ - Initialize the text-to-video pipeline. - - Args: - input_validation_stage: The input validation stage. - prompt_encoding_stage: The primary prompt encoding stage. - timestep_preparation_stage: The timestep preparation stage. - latent_preparation_stage: The latent preparation stage. - conditioning_stage: The conditioning stage. - denoising_stage: The denoising stage. - decoding_stage: The decoding stage. - post_processing_stage: The post-processing stage (optional). - secondary_prompt_encoding_stage: Optional secondary prompt encoding stage for models - with dual text encoders. The purpose of this encoder varies by model architecture - and may handle additional conditioning aspects such as temporal control, style, - or other model-specific features. - """ - super().__init__() - - # Add the stages in order - self.add_stage(input_validation_stage) - self.add_stage(prompt_encoding_stage) - - # Add optional secondary prompt encoding stage if provided - if secondary_prompt_encoding_stage is not None: - self.add_stage(secondary_prompt_encoding_stage) - - self.add_stage(timestep_preparation_stage) - self.add_stage(latent_preparation_stage) - self.add_stage(conditioning_stage) - self.add_stage(denoising_stage) - self.add_stage(decoding_stage) - - if post_processing_stage is not None: - self.add_stage(post_processing_stage) - - @property - def components(self) -> Dict[str, Any]: - """Get all components used by this pipeline.""" - components = {} - - # Add components from each stage - for stage in self._stages: - for attr_name in dir(stage): - if attr_name.startswith("__") or attr_name.startswith("_"): - continue - - attr = getattr(stage, attr_name) - if isinstance(attr, torch.nn.Module): - components[attr_name] = attr - - return components \ No newline at end of file diff --git a/fastvideo/pipelines/implementations/__init__.py b/fastvideo/pipelines/implementations/__init__.py deleted file mode 100644 index 616c6d670..000000000 --- a/fastvideo/pipelines/implementations/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -""" -Concrete implementations of diffusion pipelines. - -This package contains concrete implementations of diffusion pipelines for -specific models or architectures. -""" - -__all__ = [] \ No newline at end of file diff --git a/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py b/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py deleted file mode 100644 index d634bea20..000000000 --- a/fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py +++ /dev/null @@ -1,125 +0,0 @@ -""" -HunYuan video diffusion pipeline implementation. - -This module contains an implementation of the HunYuan video diffusion pipeline -using the modular pipeline architecture. -""" - -from typing import Any, Dict, List, Optional, Union, Tuple -import torch -from diffusers.image_processor import VaeImageProcessor - -from fastvideo.pipelines.composed import ComposedPipelineBase -from fastvideo.pipelines.stages import ( - InputValidationStage, - PromptEncodingStage, - TimestepPreparationStage, - LatentPreparationStage, - ConditioningStage, - DenoisingStage, - DecodingStage, - PostProcessingStage, -) -from fastvideo.pipelines import register_pipeline -from fastvideo.pipelines.stages.prompt_encoding import PromptEncodingStage -# from fastvideo.pipelines.stages.timestep_preparation import FlowMatchingTimestepPreparationStage -from fastvideo.inference_args import InferenceArgs -# from fastvideo.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from diffusers.utils import BaseOutput -from typing import Union, Optional, Tuple, List -import numpy as np -from dataclasses import dataclass - -from fastvideo.logger import init_logger -logger = init_logger(__name__) - -# class HunyuanLatentPreparationStage(LatentPreparationStage): -# def _call_implementation(self, batch: ForwardBatch, inference_args: InferenceArgs) -> ForwardBatch: -# "custom logic for HunYuan latent preparation" -# pass - - -# class hunyuanloader(PipelineLoader): -# def load_components(self, inference_args: InferenceArgs): -# pass - -@dataclass -class DiffusionPipelineOutput(BaseOutput): - videos: Union[torch.Tensor, np.ndarray] - - -@register_pipeline("hunyuan-video") -class HunyuanVideoPipeline(ComposedPipelineBase): - - def setup_pipeline(self, inference_args: InferenceArgs): - self._stages = [] - self.add_stage("input_validation_stage", - InputValidationStage()) - self.add_stage("prompt_encoding_stage_primary", - PromptEncodingStage(is_secondary=False)) - self.add_stage("prompt_encoding_stage_secondary", - PromptEncodingStage(is_secondary=True)) - self.add_stage("conditioning_stage", - ConditioningStage()) - self.add_stage("timestep_preparation_stage", - TimestepPreparationStage()) - self.add_stage("latent_preparation_stage", - LatentPreparationStage()) - self.add_stage("denoising_stage", - DenoisingStage()) - self.add_stage("decoding_stage", - DecodingStage()) - - def initialize_pipeline(self, inference_args: InferenceArgs): - assert len(self._stages) > 0, "Pipeline stages are not set" - assert len(self._modules) > 0, "Pipeline modules are not set" - - - vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1) - inference_args.vae_scale_factor = vae_scale_factor - - self.image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) - self.register_modules({"image_processor": self.image_processor}) - - - num_channels_latents = self.transformer.in_channels - inference_args.num_channels_latents = num_channels_latents - - - - def adjust_video_length(self, batch: ForwardBatch, inference_args: InferenceArgs): - """Adjust video length based on VAE version""" - video_length = batch.num_frames - vae_ver = inference_args.vae - logger.info(f"Adjusting video length for VAE version: {vae_ver}") - if "884" in vae_ver: - batch.num_frames = (video_length - 1) // 4 + 1 - elif "888" in vae_ver: - batch.num_frames = (video_length - 1) // 8 + 1 - return batch - - - @torch.no_grad() - def forward(self, batch: ForwardBatch, inference_args: InferenceArgs) -> ForwardBatch: - logger.info(f"Running pipeline stages: {self._stage_name_mapping.keys()}") - logger.info(f"Batch: {batch}") - # for stage in self._stages: - # batch = stage(batch, inference_args) - - # or - - batch = self.input_validation_stage(batch, inference_args) - batch = self.prompt_encoding_stage_primary(batch, inference_args) - batch = self.prompt_encoding_stage_secondary(batch, inference_args) - batch = self.conditioning_stage(batch, inference_args) - batch = self.timestep_preparation_stage(batch, inference_args) - - # custom logic - batch = self.adjust_video_length(batch, inference_args) - - batch = self.latent_preparation_stage(batch, inference_args) - batch = self.denoising_stage(batch, inference_args) - batch = self.decoding_stage(batch, inference_args) - - return DiffusionPipelineOutput(videos=batch.videos) diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py deleted file mode 100644 index 218d77a4b..000000000 --- a/fastvideo/pipelines/pipeline_registry.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -Registry for pipeline implementations. - -This module provides a registry for pipeline implementations, allowing them to be -looked up by name. -""" - -from typing import Dict, Type, Any, Optional, Callable, TypeVar, cast - -# Define a type variable for the pipeline class -T = TypeVar('T') - -class PipelineRegistry: - """Registry for pipeline implementations.""" - - _registry: Dict[str, Type[Any]] = {} - - @classmethod - def register(cls, name: str) -> Callable[[Type[T]], Type[T]]: - """ - Register a pipeline implementation. - - Args: - name: The name to register the pipeline under. - - Returns: - A decorator that registers the pipeline. - """ - def decorator(pipeline_cls: Type[T]) -> Type[T]: - cls._registry[name] = pipeline_cls - return pipeline_cls - - return decorator - - @classmethod - def get(cls, name: str) -> Optional[Type[Any]]: - """ - Get a pipeline implementation by name. - - Args: - name: The name of the pipeline. - - Returns: - The pipeline class, or None if not found. - """ - return cls._registry.get(name) - - @classmethod - def list(cls) -> Dict[str, Type[Any]]: - """ - List all registered pipelines. - - Returns: - A dictionary of pipeline names to pipeline classes. - """ - return cls._registry.copy() - - -def register_pipeline(name: str) -> Callable[[Type[T]], Type[T]]: - """ - Register a pipeline implementation. - - Args: - name: The name to register the pipeline under. - - Returns: - A decorator that registers the pipeline. - """ - return PipelineRegistry.register(name) \ No newline at end of file diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py deleted file mode 100644 index db8612a98..000000000 --- a/fastvideo/pipelines/stages/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -""" -Pipeline stages for diffusion models. - -This package contains the various stages that can be composed to create -complete diffusion pipelines. -""" - -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.input_validation import InputValidationStage -from fastvideo.pipelines.stages.prompt_encoding import PromptEncodingStage -from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage -from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage -from fastvideo.pipelines.stages.conditioning import ConditioningStage -from fastvideo.pipelines.stages.denoising import DenoisingStage -from fastvideo.pipelines.stages.decoding import DecodingStage -from fastvideo.pipelines.stages.post_processing import PostProcessingStage - -__all__ = [ - "PipelineStage", - "InputValidationStage", - "PromptEncodingStage", - "TimestepPreparationStage", - "LatentPreparationStage", - "ConditioningStage", - "DenoisingStage", - "DecodingStage", - "PostProcessingStage", -] \ No newline at end of file diff --git a/fastvideo/sample/new_sample_hunyuan.py b/fastvideo/sample/new_sample_hunyuan.py deleted file mode 100644 index 0ac839bea..000000000 --- a/fastvideo/sample/new_sample_hunyuan.py +++ /dev/null @@ -1,103 +0,0 @@ -import argparse -import os -from pathlib import Path - -import imageio -import numpy as np -import torch -import torchvision -from einops import rearrange - -# Fix the import path -from fastvideo.inference import DiffusionInference -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser -from fastvideo.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.logger import init_logger -import torch.distributed as dist -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state -logger = init_logger(__name__) - -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(inference_args: InferenceArgs): - # initialize_distributed() - # print(nccl_info.sp_size) - # local_rank = int(os.environ.get("LOCAL_RANK", 0)) - # rank = int(os.environ.get("RANK", 0)) - # world_size = int(os.environ.get("WORLD_SIZE", 1)) - - # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}") - # init_distributed_environment( - # world_size=world_size, - # rank=rank, - # local_rank=local_rank - # ) - # print(inference_args.sp_size) - - # Initialize tensor model parallel groups - # initialize_model_parallel( - # sequence_model_parallel_size=inference_args.sp_size - # ) - initialize_distributed() - - - print('Loading pipeline') - # Create inference object using the updated API - diffusion = DiffusionInference.load_pipeline( - # model_path=models_root_path, - inference_args, - # model_loader_cls=ModelLoader, - # pipeline_cls=HunyuanVideoPipeline - ) - print('Pipeline loaded') - # return - - # Load prompts - if inference_args.prompt.endswith('.txt'): - with open(inference_args.prompt) as f: - prompts = [line.strip() for line in f.readlines()] - else: - prompts = [inference_args.prompt] - - # Process each prompt - for prompt in prompts: - outputs = diffusion.predict( - prompt=prompt, - inference_args=inference_args, - ) - - # Process outputs - videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - - # Save video - os.makedirs(os.path.dirname(inference_args.output_path), exist_ok=True) - imageio.mimsave( - os.path.join(inference_args.output_path, f"{prompt[:100]}.mp4"), - frames, - fps=inference_args.fps - ) - - -if __name__ == "__main__": - parser = FlexibleArgumentParser() - InferenceArgs.add_cli_args(parser) - args = parser.parse_args() - inference_args = InferenceArgs.from_cli_args(args) - - # Validate arguments - if inference_args.vae_sp and not inference_args.vae_tiling: - raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - - main(inference_args) diff --git a/fastvideo/sample/sample_t2v_hunyuan.py b/fastvideo/sample/sample_t2v_hunyuan.py index 40cce827d..0e595dc3c 100644 --- a/fastvideo/sample/sample_t2v_hunyuan.py +++ b/fastvideo/sample/sample_t2v_hunyuan.py @@ -11,8 +11,7 @@ from einops import rearrange from fastvideo.models.hunyuan.inference import HunyuanVideoSampler from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser + def initialize_distributed(): local_rank = int(os.getenv("RANK", 0)) @@ -23,21 +22,21 @@ def initialize_distributed(): initialize_sequence_parallel_state(world_size) -def main(inference_args: InferenceArgs): +def main(args): initialize_distributed() print(nccl_info.sp_size) - print(inference_args) - models_root_path = Path(inference_args.model_path) + 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 = inference_args.output_path + 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=inference_args) + hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args) # Get the updated args args = hunyuan_video_sampler.args @@ -74,15 +73,7 @@ def main(inference_args: InferenceArgs): if __name__ == "__main__": - parser = FlexibleArgumentParser() - InferenceArgs.add_cli_args(parser) - args = parser.parse_args() - inference_args = InferenceArgs.from_cli_args(args) - # process for vae sequence parallel - if inference_args.vae_sp and not inference_args.vae_tiling: - raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - main(inference_args) - exit() + parser = argparse.ArgumentParser() # Basic parameters parser.add_argument("--prompt", type=str, help="prompt file for inference") diff --git a/fastvideo/sample/v1_fastvideo_inference.py b/fastvideo/sample/v1_fastvideo_inference.py deleted file mode 100644 index 7eec79cc2..000000000 --- a/fastvideo/sample/v1_fastvideo_inference.py +++ /dev/null @@ -1,101 +0,0 @@ -import argparse -import os -from pathlib import Path - -import imageio -import numpy as np -import torch -import torchvision -from einops import rearrange - -# Fix the import path -from fastvideo.inference_engine import InferenceEngine -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser -from fastvideo.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.logger import init_logger -import torch.distributed as dist -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state -logger = init_logger(__name__) - -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(inference_args: InferenceArgs): - # initialize_distributed() - # print(nccl_info.sp_size) - # local_rank = int(os.environ.get("LOCAL_RANK", 0)) - # rank = int(os.environ.get("RANK", 0)) - # world_size = int(os.environ.get("WORLD_SIZE", 1)) - - # logger.info(f"Initializing process: rank={rank}, local_rank={local_rank}, world_size={world_size}") - # init_distributed_environment( - # world_size=world_size, - # rank=rank, - # local_rank=local_rank - # ) - # print(inference_args.sp_size) - - # Initialize tensor model parallel groups - # initialize_model_parallel( - # sequence_model_parallel_size=inference_args.sp_size - # ) - initialize_distributed() - - - print('Creating engine') - # Create inference object using the updated API - engine = InferenceEngine.create_engine( - inference_args, - ) - print('Engine created') - # return - - # Load prompts - if inference_args.prompt.endswith('.txt'): - with open(inference_args.prompt) as f: - prompts = [line.strip() for line in f.readlines()] - else: - prompts = [inference_args.prompt] - - # Process each prompt - for prompt in prompts: - outputs = engine.run( - prompt=prompt, - inference_args=inference_args, - ) - - # Process outputs - videos = rearrange(outputs["samples"], "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - - # Save video - os.makedirs(os.path.dirname(inference_args.output_path), exist_ok=True) - imageio.mimsave( - os.path.join(inference_args.output_path, f"{prompt[:100]}.mp4"), - frames, - fps=inference_args.fps - ) - - -if __name__ == "__main__": - parser = FlexibleArgumentParser() - InferenceArgs.add_cli_args(parser) - args = parser.parse_args() - inference_args = InferenceArgs.from_cli_args(args) - inference_args.check_inference_args() - - # Validate arguments - if inference_args.vae_sp and not inference_args.vae_tiling: - raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.") - - main(inference_args) \ No newline at end of file diff --git a/fastvideo/v1/attention/__init__.py b/fastvideo/v1/attention/__init__.py new file mode 100644 index 000000000..7dc81b89b --- /dev/null +++ b/fastvideo/v1/attention/__init__.py @@ -0,0 +1,3 @@ +from .flash_attn import (DistributedAttention, LocalAttention) + +__all__ = ["DistributedAttention", "LocalAttention"] \ No newline at end of file diff --git a/fastvideo/attention/distributed_attn.py b/fastvideo/v1/attention/flash_attn.py similarity index 93% rename from fastvideo/attention/distributed_attn.py rename to fastvideo/v1/attention/flash_attn.py index 169578e1a..5ba432389 100644 --- a/fastvideo/attention/distributed_attn.py +++ b/fastvideo/v1/attention/flash_attn.py @@ -1,9 +1,11 @@ +from itertools import accumulate +from typing import List, Optional + import torch import torch.nn as nn -from typing import Optional -from fastvideo.distributed.communication_op import sequence_model_parallel_all_to_all_4D, sequence_model_parallel_all_gather -from fastvideo.distributed.parallel_state import get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size -from flash_attn import flash_attn_func +from fastvideo.v1.distributed.communication_op import sequence_model_parallel_all_to_all_4D, sequence_model_parallel_all_gather +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size +from flash_attn import flash_attn_func, flash_attn_varlen_func class DistributedAttention(nn.Module): @@ -33,6 +35,7 @@ class DistributedAttention(nn.Module): self.dropout_rate = dropout_rate self.causal = causal self.softmax_scale = softmax_scale + def forward( self, q: torch.Tensor, diff --git a/fastvideo/distributed/__init__.py b/fastvideo/v1/distributed/__init__.py similarity index 100% rename from fastvideo/distributed/__init__.py rename to fastvideo/v1/distributed/__init__.py diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/v1/distributed/communication_op.py similarity index 100% rename from fastvideo/distributed/communication_op.py rename to fastvideo/v1/distributed/communication_op.py diff --git a/fastvideo/distributed/device_communicators/__init__.py b/fastvideo/v1/distributed/device_communicators/__init__.py similarity index 100% rename from fastvideo/distributed/device_communicators/__init__.py rename to fastvideo/v1/distributed/device_communicators/__init__.py diff --git a/fastvideo/distributed/device_communicators/base_device_communicator.py b/fastvideo/v1/distributed/device_communicators/base_device_communicator.py similarity index 100% rename from fastvideo/distributed/device_communicators/base_device_communicator.py rename to fastvideo/v1/distributed/device_communicators/base_device_communicator.py diff --git a/fastvideo/distributed/device_communicators/cuda_communicator.py b/fastvideo/v1/distributed/device_communicators/cuda_communicator.py similarity index 98% rename from fastvideo/distributed/device_communicators/cuda_communicator.py rename to fastvideo/v1/distributed/device_communicators/cuda_communicator.py index 264c0e3bb..50494566c 100644 --- a/fastvideo/distributed/device_communicators/cuda_communicator.py +++ b/fastvideo/v1/distributed/device_communicators/cuda_communicator.py @@ -33,7 +33,7 @@ class CudaCommunicator(DeviceCommunicatorBase): # lazy import to avoid documentation build error # from vllm.distributed.device_communicators.custom_all_reduce import ( # CustomAllreduce) - from fastvideo.distributed.device_communicators.pynccl import ( + from fastvideo.v1.distributed.device_communicators.pynccl import ( PyNcclCommunicator) self.pynccl_comm: Optional[PyNcclCommunicator] = None diff --git a/fastvideo/distributed/device_communicators/pynccl.py b/fastvideo/v1/distributed/device_communicators/pynccl.py similarity index 97% rename from fastvideo/distributed/device_communicators/pynccl.py rename to fastvideo/v1/distributed/device_communicators/pynccl.py index 2ab53e02e..3d85a831c 100644 --- a/fastvideo/distributed/device_communicators/pynccl.py +++ b/fastvideo/v1/distributed/device_communicators/pynccl.py @@ -8,12 +8,12 @@ import torch import torch.distributed as dist from torch.distributed import ProcessGroup, ReduceOp -from fastvideo.distributed.device_communicators.pynccl_wrapper import ( +from fastvideo.v1.distributed.device_communicators.pynccl_wrapper import ( NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum, ncclRedOpTypeEnum, ncclUniqueId) -from fastvideo.distributed.utils import StatelessProcessGroup -from fastvideo.logger import init_logger -from fastvideo.utils.utils import current_stream +from fastvideo.v1.distributed.utils import StatelessProcessGroup +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import current_stream logger = init_logger(__name__) diff --git a/fastvideo/distributed/device_communicators/pynccl_wrapper.py b/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py similarity index 99% rename from fastvideo/distributed/device_communicators/pynccl_wrapper.py rename to fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py index 97c3a7a34..cc3770a9e 100644 --- a/fastvideo/distributed/device_communicators/pynccl_wrapper.py +++ b/fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py @@ -30,8 +30,8 @@ from typing import Any, Dict, List, Optional import torch from torch.distributed import ReduceOp -from fastvideo.logger import init_logger -from fastvideo.utils.utils import find_nccl_library +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import find_nccl_library logger = init_logger(__name__) diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/v1/distributed/parallel_state.py similarity index 80% rename from fastvideo/distributed/parallel_state.py rename to fastvideo/v1/distributed/parallel_state.py index 467d8917d..0d00bbe33 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/v1/distributed/parallel_state.py @@ -40,15 +40,15 @@ import torch import torch.distributed from torch.distributed import Backend, ProcessGroup -import fastvideo.envs as envs -from fastvideo.distributed.device_communicators.base_device_communicator import ( +import fastvideo.v1.envs as envs +from fastvideo.v1.distributed.device_communicators.base_device_communicator import ( DeviceCommunicatorBase) -from fastvideo.distributed.utils import StatelessProcessGroup -from fastvideo.logger import init_logger -# from fastvideo.utils import (direct_register_custom_op, resolve_obj_by_qualname, +from fastvideo.v1.distributed.utils import StatelessProcessGroup +from fastvideo.v1.logger import init_logger +# from fastvideo.v1.utils import (direct_register_custom_op, resolve_obj_by_qualname, # supports_custom_op) -from fastvideo.distributed.device_communicators.cuda_communicator import ( +from fastvideo.v1.distributed.device_communicators.cuda_communicator import ( CudaCommunicator) @@ -192,7 +192,7 @@ class GroupCoordinator: assert self.cpu_group is not None assert self.device_group is not None - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform # TODO: fix it for other platforms if current_platform.is_cuda_alike(): @@ -220,7 +220,7 @@ class GroupCoordinator: # self.mq_broadcaster = MessageQueue.create_from_process_group( # self.cpu_group, 1 << 22, 6) - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform # self.use_custom_op_call = current_platform.is_cuda_alike() self.use_custom_op_call = False @@ -270,7 +270,7 @@ class GroupCoordinator: # only cuda uses this function, # so we don't abstract it into the base class maybe_ca_context = nullcontext() - from fastvideo.distributed.device_communicators.cuda_communicator import ( + from fastvideo.v1.distributed.device_communicators.cuda_communicator import ( CudaCommunicator) if self.device_communicator is not None: assert isinstance(self.device_communicator, CudaCommunicator) @@ -850,7 +850,6 @@ def get_sp_group() -> GroupCoordinator: def initialize_model_parallel( tensor_model_parallel_size: int = 1, - pipeline_model_parallel_size: int = 1, sequence_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: @@ -885,15 +884,15 @@ def initialize_model_parallel( get_world_group().device_group) # Ensure the world size is compatible with the parallelism configuration - assert world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * sequence_model_parallel_size) == 0, \ - f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) * " \ - f"pipeline_model_parallel_size ({pipeline_model_parallel_size}) * " \ - f"sequence_model_parallel_size ({sequence_model_parallel_size})" + # assert world_size % (tensor_model_parallel_size * pipeline_model_parallel_size * sequence_model_parallel_size) == 0, \ + # f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size}) * " \ + # f"pipeline_model_parallel_size ({pipeline_model_parallel_size}) * " \ + # f"sequence_model_parallel_size ({sequence_model_parallel_size})" # Check for incompatible parallelism configurations - if sequence_model_parallel_size > 1: - assert tensor_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with tensor parallelism (TP). Please use SP=1 or TP=1." - assert pipeline_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with pipeline parallelism (PP). Please use SP=1 or PP=1." + # if sequence_model_parallel_size > 1: + # assert tensor_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with tensor parallelism (TP). Please use SP=1 or TP=1." + # assert pipeline_model_parallel_size == 1, "Sequence parallelism (SP) is incompatible with pipeline parallelism (PP). Please use SP=1 or PP=1." # Build the tensor model-parallel groups. num_tensor_model_parallel_groups: int = (world_size // @@ -933,21 +932,6 @@ def initialize_model_parallel( backend, group_name="sp") - # Build the pipeline model-parallel groups. - num_pipeline_model_parallel_groups: int = (world_size // - pipeline_model_parallel_size) - global _PP - assert _PP is None, ( - "pipeline model parallel group is already initialized") - group_ranks = [] - for i in range(num_pipeline_model_parallel_groups): - ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) - group_ranks.append(ranks) - _PP = init_model_parallel_group(group_ranks, - get_world_group().local_rank, - backend, - group_name="pp") - def get_sequence_model_parallel_world_size(): """Return world size for the sequence model parallel group.""" @@ -1079,7 +1063,7 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False): import ray # Lazy import Ray ray.shutdown() gc.collect() - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform if not current_platform.is_cpu(): torch.cuda.empty_cache() try: @@ -1172,3 +1156,267 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup], aggregated_data += rank_data return [x == 1 for x in aggregated_data.tolist()] + + +def initialize_tensor_parallel_group( + tensor_model_parallel_size: int = 1, + backend: Optional[str] = None, + group_name_suffix: str = "" +) -> GroupCoordinator: + """Initialize a tensor parallel group for a specific model. + + This function creates a tensor parallel group that can be used with the + patch_tensor_parallel_group context manager. It allows different models + to use different tensor parallelism configurations. + + Arguments: + tensor_model_parallel_size: number of GPUs used for tensor model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for tensor parallelism that can be used with + the patch_tensor_parallel_group context manager. + + Example usage: + ```python + # Initialize tensor parallel group for model1 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Use tensor parallelism for model1 + with patch_tensor_parallel_group(tp_group_model1): + # Run model1 with tensor parallelism + output1 = model1(input1) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert world_size % tensor_model_parallel_size == 0, \ + f"World size ({world_size}) must be divisible by tensor_model_parallel_size ({tensor_model_parallel_size})" + + # Build the tensor model-parallel groups. + num_tensor_model_parallel_groups: int = (world_size // tensor_model_parallel_size) + tp_group_ranks = [] + for i in range(num_tensor_model_parallel_groups): + ranks = list( + range(i * tensor_model_parallel_size, + (i + 1) * tensor_model_parallel_size)) + tp_group_ranks.append(ranks) + + # Create TP group coordinator with a unique name + group_name = f"tp_{group_name_suffix}" if group_name_suffix else "tp" + tp_group = init_model_parallel_group(tp_group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=True, + group_name=group_name) + + return tp_group + + +def initialize_sequence_parallel_group( + sequence_model_parallel_size: int = 1, + backend: Optional[str] = None, + group_name_suffix: str = "" +) -> GroupCoordinator: + """Initialize a sequence parallel group for a specific model. + + This function creates a sequence parallel group that can be used with the + patch_sequence_parallel_group context manager. It allows different models + to use different sequence parallelism configurations. + + Arguments: + sequence_model_parallel_size: number of GPUs used for sequence model parallelism. + backend: communication backend to use. + group_name_suffix: optional suffix to make the group name unique. + + Returns: + A GroupCoordinator for sequence parallelism that can be used with + the patch_sequence_parallel_group context manager. + + Example usage: + ```python + # Initialize sequence parallel group for model2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Use sequence parallelism for model2 + with patch_sequence_parallel_group(sp_group_model2): + # Run model2 with sequence parallelism + output2 = model2(input2) + ``` + """ + # Get world size and rank. Ensure some consistencies. + assert torch.distributed.is_initialized() + world_size: int = torch.distributed.get_world_size() + backend = backend or torch.distributed.get_backend( + get_world_group().device_group) + + # Ensure the world size is compatible with the parallelism configuration + assert world_size % sequence_model_parallel_size == 0, \ + f"World size ({world_size}) must be divisible by sequence_model_parallel_size ({sequence_model_parallel_size})" + + # Build the sequence model-parallel groups. + num_sequence_model_parallel_groups: int = (world_size // sequence_model_parallel_size) + sp_group_ranks = [] + + for i in range(num_sequence_model_parallel_groups): + # Create groups of consecutive ranks + ranks = list(range(i * sequence_model_parallel_size, + (i + 1) * sequence_model_parallel_size)) + sp_group_ranks.append(ranks) + + # Create SP group coordinator with a unique name + group_name = f"sp_{group_name_suffix}" if group_name_suffix else "sp" + sp_group = init_model_parallel_group(sp_group_ranks, + get_world_group().local_rank, + backend, + group_name=group_name) + + return sp_group + + +_SP_STATE_PATCHED = False + +@contextmanager +def patch_sequence_parallel_group(sp_group: GroupCoordinator): + """Patch the sp group temporarily until this function ends. + + This method allows running a model with SP while another model uses TP. + + Args: + sp_group (GroupCoordinator): the sp group coordinator + """ + global _SP_STATE_PATCHED + assert not _SP_STATE_PATCHED, "Should not call when it's already patched" + + _SP_STATE_PATCHED = True + old_sp_group = get_sp_group() + global _SP + _SP = sp_group + try: + yield + finally: + # restore the original state + _SP_STATE_PATCHED = False + _SP = old_sp_group + + +_PP_STATE_PATCHED = False + +@contextmanager +def patch_pipeline_parallel_group(pp_group: GroupCoordinator): + """Patch the pp group temporarily until this function ends. + + This method allows running a model with a different PP configuration. + + Args: + pp_group (GroupCoordinator): the pp group coordinator + """ + global _PP_STATE_PATCHED + assert not _PP_STATE_PATCHED, "Should not call when it's already patched" + + _PP_STATE_PATCHED = True + old_pp_group = get_pp_group() + global _PP + _PP = pp_group + try: + yield + finally: + # restore the original state + _PP_STATE_PATCHED = False + _PP = old_pp_group + + +# Example of how to use the independent parallelism functions +""" +Here's a complete example of how to use the independent parallelism functions +for different models: + +```python +import torch +from fastvideo.v1.distributed.parallel_state import ( + init_distributed_environment, + initialize_tensor_parallel_group, + initialize_sequence_parallel_group, + patch_tensor_parallel_group, + patch_sequence_parallel_group, + destroy_model_parallel, + destroy_distributed_environment +) + +# Initialize the distributed environment +init_distributed_environment( + world_size=8, + rank=torch.distributed.get_rank(), + distributed_init_method="tcp://localhost:12345", + local_rank=torch.distributed.get_rank() % torch.cuda.device_count() +) + +try: + # Create a tensor parallel group for model1 with TP=4 + tp_group_model1 = initialize_tensor_parallel_group( + tensor_model_parallel_size=4, + group_name_suffix="model1" + ) + + # Create a sequence parallel group for model2 with SP=2 + sp_group_model2 = initialize_sequence_parallel_group( + sequence_model_parallel_size=2, + group_name_suffix="model2" + ) + + # Create another tensor parallel group for model3 with TP=2 + tp_group_model3 = initialize_tensor_parallel_group( + tensor_model_parallel_size=2, + group_name_suffix="model3" + ) + + # Use model1 with tensor parallelism + with patch_tensor_parallel_group(tp_group_model1): + # Inside this context, get_tp_group() returns tp_group_model1 + # Run model1 with tensor parallelism + output1 = model1(input1) + + # Use model2 with sequence parallelism + with patch_sequence_parallel_group(sp_group_model2): + # Inside this context, get_sp_group() returns sp_group_model2 + # Run model2 with sequence parallelism + output2 = model2(input2) + + # Use model3 with a different tensor parallelism configuration + with patch_tensor_parallel_group(tp_group_model3): + # Inside this context, get_tp_group() returns tp_group_model3 + # Run model3 with tensor parallelism + output3 = model3(input3) + + # You can switch between models as needed + with patch_tensor_parallel_group(tp_group_model1): + # Back to using model1 + more_output1 = model1(more_input1) + +finally: + # Clean up + destroy_model_parallel() + destroy_distributed_environment() +``` + +This approach allows you to: +1. Create separate parallel groups for each model +2. Use different parallelism strategies for different models +3. Switch between models as needed +4. Use unique group names to avoid conflicts + +Note that each model can use its own optimal parallelism strategy without +interfering with other models. +""" diff --git a/fastvideo/distributed/utils.py b/fastvideo/v1/distributed/utils.py similarity index 99% rename from fastvideo/distributed/utils.py rename to fastvideo/v1/distributed/utils.py index 3b311a7e6..4286ab819 100644 --- a/fastvideo/distributed/utils.py +++ b/fastvideo/v1/distributed/utils.py @@ -18,8 +18,8 @@ from torch.distributed.distributed_c10d import (Backend, PrefixStore, is_nccl_available) from torch.distributed.rendezvous import rendezvous -import fastvideo.envs as envs -from fastvideo.logger import init_logger +import fastvideo.v1.envs as envs +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/envs.py b/fastvideo/v1/envs.py similarity index 100% rename from fastvideo/envs.py rename to fastvideo/v1/envs.py diff --git a/fastvideo/inference_args.py b/fastvideo/v1/inference_args.py similarity index 93% rename from fastvideo/inference_args.py rename to fastvideo/v1/inference_args.py index 6f0c51022..5ba176b30 100644 --- a/fastvideo/inference_args.py +++ b/fastvideo/v1/inference_args.py @@ -17,8 +17,8 @@ import argparse import dataclasses -from fastvideo.logger import init_logger -from fastvideo.utils.utils import FlexibleArgumentParser +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import FlexibleArgumentParser from typing import List, Optional @@ -29,9 +29,16 @@ class InferenceArgs: # Model and path configuration model_path: str model: str = "HYVideo-T/2-cfgdistill" + use_v1_text_encoder: bool = False + use_v1_vae: bool = False + use_v1_transformer: bool = False dit_weight: Optional[str] = None model_dir: Optional[str] = None + # HuggingFace specific parameters + trust_remote_code: bool = False + revision: Optional[str] = None + # Parallelism tp_size: int = 1 sp_size: int = 1 @@ -126,12 +133,29 @@ class InferenceArgs: output_path: str = "outputs/" seed: int = 1024 seeds: Optional[List[int]] = None + device_str: Optional[str] = None + device = None def __post_init__(self): pass @staticmethod def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument( + "--use-v1-text-encoder", + action="store_true", + help="Use the v1 text encoder", + ) + parser.add_argument( + "--use-v1-vae", + action="store_true", + help="Use the v1 vae", + ) + parser.add_argument( + "--use-v1-transformer", + action="store_true", + help="Use the v1 transformer", + ) # Model and path configuration parser.add_argument( "--model-path", @@ -155,6 +179,20 @@ class InferenceArgs: type=str, help="Directory containing StepVideo model", ) + + # HuggingFace specific parameters + parser.add_argument( + "--trust-remote-code", + action="store_true", + default=InferenceArgs.trust_remote_code, + help="Trust remote code when loading HuggingFace models", + ) + parser.add_argument( + "--revision", + type=str, + default=InferenceArgs.revision, + help="The specific model version to use (can be a branch name, tag name, or commit id)", + ) # Parallelism parser.add_argument( diff --git a/fastvideo/inference_engine.py b/fastvideo/v1/inference_engine.py similarity index 89% rename from fastvideo/inference_engine.py rename to fastvideo/v1/inference_engine.py index 6d0315273..285bb6219 100644 --- a/fastvideo/inference_engine.py +++ b/fastvideo/v1/inference_engine.py @@ -7,20 +7,16 @@ This module provides classes and functions for running inference with diffusion import os import time import torch -import numpy as np -from typing import Any, Dict, List, Optional, Union, Type -from PIL import Image -import uuid -import random +from typing import Any, Dict -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines import ComposedPipelineBase, build_pipeline -from fastvideo.pipelines.pipeline_registry import PipelineRegistry -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.logger import init_logger -# TODO(will): remove -from fastvideo.models.hunyuan.utils.data_utils import align_to -from fastvideo.models.hunyuan.constants import NEGATIVE_PROMPT +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines import ComposedPipelineBase, build_pipeline +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.logger import init_logger +# TODO(will): remove, check if this is hunyuan specific +from fastvideo.v1.utils import align_to +# TODO(will): remove, move this to hunyuan stage +from fastvideo.v1.pipelines.implementations.hunyuan.constants import NEGATIVE_PROMPT logger = init_logger(__name__) @@ -72,6 +68,11 @@ class InferenceEngine: """ try: logger.info(f"Building pipeline...") + # TODO(will): probably a better place to set device_str? + local_rank = os.environ.get("LOCAL_RANK", 0) + device_str = f"cuda:{local_rank}" + inference_args.device_str = device_str + inference_args.device = torch.device(device_str) # TODO(will): I don't really like this api. # it should be something closer to pipeline_cls.from_pretrained(...) # this way for training we can just do pipeline_cls.from_pretrained( @@ -153,7 +154,7 @@ class InferenceEngine: negative_prompt = [negative_prompt.strip()] - # from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler + # from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler # scheduler = FlowMatchDiscreteScheduler( # shift=flow_shift, # reverse=self.inference_args.flow_reverse, @@ -190,9 +191,7 @@ class InferenceEngine: # return # sp_group = get_sp_group() # local_rank = sp_group.rank - import os - local_rank = os.environ.get("LOCAL_RANK", 0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(inference_args.device_str) batch = ForwardBatch( prompt=prompt, negative_prompt=negative_prompt, diff --git a/fastvideo/layers/__init__.py b/fastvideo/v1/layers/__init__.py similarity index 100% rename from fastvideo/layers/__init__.py rename to fastvideo/v1/layers/__init__.py diff --git a/fastvideo/layers/activation.py b/fastvideo/v1/layers/activation.py similarity index 99% rename from fastvideo/layers/activation.py rename to fastvideo/v1/layers/activation.py index dc7db7d22..6217ff1a5 100644 --- a/fastvideo/layers/activation.py +++ b/fastvideo/v1/layers/activation.py @@ -11,8 +11,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size) from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform -from vllm.utils import LazyDict +from fastvideo.v1.platforms import current_platform @CustomOp.register("fatrelu_and_mul") diff --git a/fastvideo/loader/__init__.py b/fastvideo/v1/layers/conv.py similarity index 100% rename from fastvideo/loader/__init__.py rename to fastvideo/v1/layers/conv.py diff --git a/fastvideo/layers/layernorm.py b/fastvideo/v1/layers/layernorm.py similarity index 100% rename from fastvideo/layers/layernorm.py rename to fastvideo/v1/layers/layernorm.py diff --git a/fastvideo/layers/linear.py b/fastvideo/v1/layers/linear.py similarity index 99% rename from fastvideo/layers/linear.py rename to fastvideo/v1/layers/linear.py index c8550d09e..365261101 100644 --- a/fastvideo/layers/linear.py +++ b/fastvideo/v1/layers/linear.py @@ -8,25 +8,25 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter -from fastvideo.distributed import (divide, get_tensor_model_parallel_rank, +from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, split_tensor_along_last_dim, tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce) -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger # TODO(will): remove this import by copying the definition from vLLM then # manually import each quantization method we want to use. Refer to SGLang from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase) # yapf: disable -from fastvideo.models.parameter import (BasevLLMParameter, +from fastvideo.v1.models.parameter import (BasevLLMParameter, BlockQuantScaleParameter, PackedColumnParameter, PackedvLLMParameter, PerTensorScaleParameter, RowvLLMParameter) # yapf: enable -from fastvideo.models.utils import set_weight_attrs +from fastvideo.v1.models.utils import set_weight_attrs logger = init_logger(__name__) diff --git a/fastvideo/layers/mlp.py b/fastvideo/v1/layers/mlp.py similarity index 90% rename from fastvideo/layers/mlp.py rename to fastvideo/v1/layers/mlp.py index db2d2ce01..68529c33a 100644 --- a/fastvideo/layers/mlp.py +++ b/fastvideo/v1/layers/mlp.py @@ -1,5 +1,5 @@ -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.activation import get_act_fn +from fastvideo.v1.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.activation import get_act_fn import torch import torch.nn as nn from typing import Optional diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/v1/layers/rotary_embedding.py similarity index 93% rename from fastvideo/layers/rotary_embedding.py rename to fastvideo/v1/layers/rotary_embedding.py index fe9aca306..f557ba1ae 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/v1/layers/rotary_embedding.py @@ -30,7 +30,7 @@ import torch.nn as nn from transformers import PretrainedConfig from vllm.model_executor.custom_op import CustomOp -from fastvideo.distributed.parallel_state import get_sp_group +from fastvideo.v1.distributed.parallel_state import get_sp_group def _rotate_neox(x: torch.Tensor) -> torch.Tensor: x1 = x[..., :x.shape[-1] // 2] @@ -500,4 +500,42 @@ def get_rotary_pos_embed( sp_rank=sp_rank, sp_world_size=sp_world_size ) - return freqs_cos, freqs_sin \ No newline at end of file + return freqs_cos, freqs_sin + +_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {} + +def get_rope( + head_size: int, + rotary_dim: int, + max_position: int, + base: int, + is_neox_style: bool = True, + rope_scaling: Optional[Dict[str, Any]] = None, + dtype: Optional[torch.dtype] = None, + partial_rotary_factor: float = 1.0, +) -> RotaryEmbedding: + if dtype is None: + dtype = torch.get_default_dtype() + if rope_scaling is not None: + # Transforms every value that is a list into a tuple for caching calls + rope_scaling_tuple = { + k: tuple(v) if isinstance(v, list) else v + for k, v in rope_scaling.items() + } + rope_scaling_args = tuple(rope_scaling_tuple.items()) + else: + rope_scaling_args = None + if partial_rotary_factor < 1.0: + rotary_dim = int(rotary_dim * partial_rotary_factor) + key = (head_size, rotary_dim, max_position, base, is_neox_style, + rope_scaling_args, dtype) + if key in _ROPE_DICT: + return _ROPE_DICT[key] + + if rope_scaling is None: + rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base, + is_neox_style, dtype) + else: + raise ValueError(f"Unknown RoPE scaling {rope_scaling}") + _ROPE_DICT[key] = rotary_emb + return rotary_emb diff --git a/fastvideo/layers/utils.py b/fastvideo/v1/layers/utils.py similarity index 100% rename from fastvideo/layers/utils.py rename to fastvideo/v1/layers/utils.py diff --git a/fastvideo/layers/visual_embedding.py b/fastvideo/v1/layers/visual_embedding.py similarity index 96% rename from fastvideo/layers/visual_embedding.py rename to fastvideo/v1/layers/visual_embedding.py index 1f5f8a19f..13a61aef1 100644 --- a/fastvideo/layers/visual_embedding.py +++ b/fastvideo/v1/layers/visual_embedding.py @@ -1,10 +1,10 @@ import torch import torch.nn as nn import math -from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.activation import get_act_fn +from fastvideo.v1.layers.linear import ReplicatedLinear from typing import Optional -from fastvideo.layers.mlp import MLP +from fastvideo.v1.layers.mlp import MLP class PatchEmbed(nn.Module): """2D Image to Patch Embedding diff --git a/fastvideo/layers/vocab_parallel_embedding.py b/fastvideo/v1/layers/vocab_parallel_embedding.py similarity index 98% rename from fastvideo/layers/vocab_parallel_embedding.py rename to fastvideo/v1/layers/vocab_parallel_embedding.py index f65dfc3cb..338fb2ac9 100644 --- a/fastvideo/layers/vocab_parallel_embedding.py +++ b/fastvideo/v1/layers/vocab_parallel_embedding.py @@ -7,14 +7,14 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter -from vllm.distributed import (divide, get_tensor_model_parallel_rank, +from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding) -from vllm.model_executor.parameter import BasevLLMParameter -from vllm.model_executor.utils import set_weight_attrs -from vllm.platforms import current_platform +from fastvideo.v1.models.parameter import BasevLLMParameter +from fastvideo.v1.models.utils import set_weight_attrs +from fastvideo.v1.platforms import current_platform DEFAULT_VOCAB_PADDING_SIZE = 64 diff --git a/fastvideo/v1/loader/__init__.py b/fastvideo/v1/loader/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/loader/fsdp_load.py b/fastvideo/v1/loader/fsdp_load.py similarity index 97% rename from fastvideo/loader/fsdp_load.py rename to fastvideo/v1/loader/fsdp_load.py index 10e7753d2..e2096b0ac 100644 --- a/fastvideo/loader/fsdp_load.py +++ b/fastvideo/v1/loader/fsdp_load.py @@ -3,17 +3,17 @@ # Copyright 2025 The FastVideo Authors. from collections import defaultdict -from typing import Any, Callable, Dict, Generator, List, Optional, Tuple +from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Type from itertools import chain import torch from torch import nn from torch.distributed import DeviceMesh, init_device_mesh -from fastvideo.distributed.parallel_state import get_sequence_model_parallel_world_size +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed._tensor import distribute_tensor from torch.nn.modules.module import _IncompatibleKeys from vllm.model_executor.model_loader.weight_utils import safetensors_weights_iterator -from fastvideo.models.dits.registry import ModelRegistry + import contextlib import re @@ -79,14 +79,13 @@ def get_param_names_mapping(mapping_dict: Dict[str, str]) -> Callable[[str], str # TODO(PY): add compile option def load_fsdp_model( - model_name: str, + model_cls: Type[nn.Module], init_params: Dict[str, Any], weight_dir_list: List[str], device: torch.device, cpu_offload: bool = False, default_dtype: Optional[torch.dtype] = torch.bfloat16, ) -> torch.nn.Module: - model_cls, _ = ModelRegistry.resolve_model_cls(model_name) with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) device_mesh = init_device_mesh( diff --git a/fastvideo/v1/loader/loader.py b/fastvideo/v1/loader/loader.py new file mode 100644 index 000000000..1980aefaf --- /dev/null +++ b/fastvideo/v1/loader/loader.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from vllm +# Copyright 2023 The vLLM Authors. +# Copyright 2025 The FastVideo Authors. + +# ruff: noqa: SIM117 +import collections +import copy +import dataclasses +import fnmatch +import glob +import inspect +import itertools +import math +import os +import time +import warnings +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional, + Tuple, cast) + +import huggingface_hub +import torch +from torch import nn +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.loader.utils import set_default_torch_dtype +from fastvideo.v1.loader.weight_utils import ( + download_safetensors_index_file_from_hf, download_weights_from_hf, + filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, + pt_weights_iterator, + safetensors_weights_iterator) +from fastvideo.v1.platforms import current_platform +from fastvideo.v1.models.registry import ModelRegistry + +from fastvideo.v1.inference_args import InferenceArgs + +@contextmanager +def device_loading_context(module: torch.nn.Module, + target_device: torch.device): + if target_device.type == "cpu": + # If target is CPU, no need to move anything + yield module + return + + original_device_states: Dict[str, torch.device] = {} + + # Store original device states and move parameters to GPU if they're on CPU + for name, p in module.named_parameters(): + if p.device.type == "cpu": + original_device_states[name] = p.device + p.data = p.data.to(target_device) + # Parameters already on target device are not touched + + try: + yield module + + finally: + # Restore parameters to their original devices, ignoring new parameters + pin_memory = True + for name, p in module.named_parameters(): + if name in original_device_states: + original_device: torch.device = original_device_states[name] + if original_device.type == "cpu": + # `torch.empty_like` does not support `pin_memory` argument + cpu_data = torch.empty_strided( + size=p.data.size(), + stride=p.data.stride(), + dtype=p.data.dtype, + layout=p.data.layout, + device="cpu", + pin_memory=pin_memory, + ) + cpu_data.copy_(p.data) + p.data = cpu_data + else: + p.data = p.data.to(original_device) + # New parameters or parameters already on target device are untouched + + +logger = init_logger(__name__) + + +def _initialize_model( + model_path: str, + model_config, +) -> nn.Module: + """Initialize a model with the given configurations. + + Args: + model_name: The name of the model architecture + model_path: Path to the model directory + component_name: Optional component name (e.g., "text_encoder", "vae") + If provided, will load the config from the component subdirectory + + Returns: + The initialized model + """ + architectures = getattr(model_config, "architectures", []) + if len(architectures) == 0: + raise ValueError("Model config does not contain a valid model architecture") + logger.info("Trying to load model from diffusers format") + class_name = model_config.pop("_class_name") + architectures = [class_name] + if architectures is None: + raise ValueError("Model config does not contain a valid model architecture") + model_cls, _ = ModelRegistry.resolve_model_cls(architectures) + logger.info(f"Loading model_cls: {model_cls}") + + return model_cls(model_config) + + +class BaseModelLoader(ABC): + """Base class for model loaders.""" + + @abstractmethod + def download_model(self, inference_args: InferenceArgs) -> None: + """Download a model so that it can be immediately loaded.""" + raise NotImplementedError + + @abstractmethod + def load_model(self, *, inference_args: InferenceArgs) -> nn.Module: + """Load a model with the given configurations.""" + raise NotImplementedError + + +class DefaultModelLoader(BaseModelLoader): + """Model loader that can load different file types from disk.""" + + @dataclasses.dataclass + class Source: + """A source for weights.""" + + model_or_path: str + """The model ID or path.""" + + revision: Optional[str] + """The optional model revision.""" + + prefix: str = "" + """A prefix to prepend to all weights.""" + + fall_back_to_pt: bool = True + """Whether .pt weights can be used.""" + + allow_patterns_overrides: Optional[list[str]] = None + """If defined, weights will load exclusively using these patterns.""" + + counter_before_loading_weights: float = 0.0 + counter_after_loading_weights: float = 0.0 + + def _prepare_weights( + self, + model_name_or_path: str, + revision: Optional[str], + fall_back_to_pt: bool, + allow_patterns_overrides: Optional[list[str]], + ) -> Tuple[str, List[str], bool]: + """Prepare weights for the model. + + If the model is not local, it will be downloaded.""" + # model_name_or_path = (self._maybe_download_from_modelscope( + # model_name_or_path, revision) or model_name_or_path) + + is_local = os.path.isdir(model_name_or_path) + assert is_local, "Model path must be a local directory" + # load_format = self.load_config.load_format + # load_format = LoadFormat.AUTO + use_safetensors = False + index_file = SAFE_WEIGHTS_INDEX_NAME + allow_patterns = ["*.safetensors", "*.bin"] + # Some quantized models use .pt files for storing the weights. + # if load_format == LoadFormat.AUTO: + # allow_patterns = ["*.safetensors", "*.bin"] + # elif load_format == LoadFormat.SAFETENSORS: + # use_safetensors = True + # allow_patterns = ["*.safetensors"] + # elif load_format == LoadFormat.MISTRAL: + # use_safetensors = True + # allow_patterns = ["consolidated*.safetensors"] + # index_file = "consolidated.safetensors.index.json" + # elif load_format == LoadFormat.PT: + # allow_patterns = ["*.pt"] + # elif load_format == LoadFormat.NPCACHE: + # allow_patterns = ["*.bin"] + # else: + # raise ValueError(f"Unknown load_format: {load_format}") + + if fall_back_to_pt: + allow_patterns += ["*.pt"] + + if allow_patterns_overrides is not None: + allow_patterns = allow_patterns_overrides + + if not is_local: + hf_folder = download_weights_from_hf( + model_name_or_path, + self.load_config.download_dir, + allow_patterns, + revision, + ignore_patterns=self.load_config.ignore_patterns, + ) + else: + hf_folder = model_name_or_path + + hf_weights_files: List[str] = [] + for pattern in allow_patterns: + hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + if len(hf_weights_files) > 0: + if pattern == "*.safetensors": + use_safetensors = True + break + + if use_safetensors: + # For models like Mistral-7B-Instruct-v0.3 + # there are both sharded safetensors files and a consolidated + # safetensors file. Using both breaks. + # Here, we download the `model.safetensors.index.json` and filter + # any files not found in the index. + if not is_local: + download_safetensors_index_file_from_hf( + model_name_or_path, + index_file, + self.load_config.download_dir, + revision, + ) + hf_weights_files = filter_duplicate_safetensors_files( + hf_weights_files, hf_folder, index_file) + else: + hf_weights_files = filter_files_not_needed_for_inference( + hf_weights_files) + + if len(hf_weights_files) == 0: + raise RuntimeError( + f"Cannot find any model weights with `{model_name_or_path}`") + + return hf_folder, hf_weights_files, use_safetensors + + def _get_weights_iterator( + self, source: "Source" + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + """Get an iterator for the model weights based on the load format.""" + hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( + source.model_or_path, source.revision, source.fall_back_to_pt, + source.allow_patterns_overrides) + if use_safetensors: + weights_iterator = safetensors_weights_iterator(hf_weights_files) + else: + weights_iterator = pt_weights_iterator(hf_weights_files) + + # if current_platform.is_tpu(): + # # In PyTorch XLA, we should call `xm.mark_step` frequently so that + # # not too many ops are accumulated in the XLA program. + # import torch_xla.core.xla_model as xm + + # def _xla_weights_iterator(iterator: Generator): + # for weights in iterator: + # yield weights + # xm.mark_step() + + # weights_iterator = _xla_weights_iterator(weights_iterator) + + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + # Apply the prefix. + return ((source.prefix + name, tensor) + for (name, tensor) in weights_iterator) + + def _get_all_weights( + self, + model_config: Dict[str, Any], + model: nn.Module, + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + primary_weights = DefaultModelLoader.Source( + model_config.model, + "", + prefix="", + fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", + True), + allow_patterns_overrides=getattr(model, "allow_patterns_overrides", + None), + ) + yield from self._get_weights_iterator(primary_weights) + + secondary_weights = cast( + Iterable[DefaultModelLoader.Source], + getattr(model, "secondary_weights", ()), + ) + for source in secondary_weights: + yield from self._get_weights_iterator(source) + + def download_model(self, inference_args: InferenceArgs) -> None: + raise NotImplementedError + # self._prepare_weights(inference_args.model, + # inference_args.revision, + # fall_back_to_pt=True, + # allow_patterns_overrides=None) + + def load_model(self, model_path: str, + model_config: Dict[str, Any], + inference_args: InferenceArgs) -> nn.Module: + logger.info(f"Loading model on device: {inference_args.device_str}") + target_device = torch.device(inference_args.device_str) + # TODO(will): add support for other dtypes + with set_default_torch_dtype(torch.float16): + with target_device: + model = _initialize_model(model_path, model_config) + + weights_to_load = {name for name, _ in model.named_parameters()} + model_config.model = model_path + loaded_weights = model.load_weights( + self._get_all_weights(model_config, model)) + self.counter_after_loading_weights = time.perf_counter() + logger.info( + "Loading weights took %.2f seconds", + self.counter_after_loading_weights - + self.counter_before_loading_weights) + # We only enable strict check for non-quantized models + # that have loaded weights tracking currently. + # if loaded_weights is not None: + weights_not_loaded = weights_to_load - loaded_weights + if weights_not_loaded: + raise ValueError( + "Following weights were not initialized from " + f"checkpoint: {weights_not_loaded}") + + # TODO(will): add support for training/finetune + return model.eval() + + + + + +def get_model_loader() -> BaseModelLoader: + """Get a model loader based on the load format.""" + # if isinstance(load_config.load_format, type): + # return load_config.load_format(load_config) + + return DefaultModelLoader() \ No newline at end of file diff --git a/fastvideo/v1/loader/utils.py b/fastvideo/v1/loader/utils.py new file mode 100644 index 000000000..a105ceb06 --- /dev/null +++ b/fastvideo/v1/loader/utils.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utilities for selecting and loading models.""" +import contextlib +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Type + +import torch +import transformers +from torch import nn +from transformers.dynamic_module_utils import get_class_from_dynamic_module + +from vllm.config import ModelConfig, ModelImpl +from fastvideo.v1.logger import init_logger +from vllm.model_executor.models import ModelRegistry +from vllm.model_executor.models.adapters import (as_classification_model, + as_embedding_model, + as_reward_model) + +logger = init_logger(__name__) + + +@contextlib.contextmanager +def set_default_torch_dtype(dtype: torch.dtype): + """Sets the default torch dtype to the given dtype.""" + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + yield + torch.set_default_dtype(old_dtype) + + +def load_hf_config_from_subdir(model_path: str, component_name: str, trust_remote_code: bool = False, revision: Optional[str] = None): + """ + Load a HuggingFace config from a component subdirectory. + + Args: + model_path: Path to the model directory + component_name: Name of the component subdirectory (e.g., "text_encoder", "vae") + trust_remote_code: Whether to trust remote code when loading the config + revision: Optional revision to use when loading from HuggingFace Hub + + Returns: + The loaded HuggingFace config + """ + import os + from transformers import AutoConfig + + component_path = os.path.join(model_path, component_name) + + if not os.path.exists(component_path): + raise ValueError(f"Component directory {component_path} does not exist") + + config_path = os.path.join(component_path, "config.json") + if not os.path.exists(config_path): + raise ValueError(f"Config file {config_path} does not exist") + + logger.info(f"Loading config from {config_path}") + config = AutoConfig.from_pretrained( + component_path, + trust_remote_code=trust_remote_code, + revision=revision + ) + + return config + + +def is_transformers_impl_compatible( + arch: str, + module: Optional[transformers.PreTrainedModel] = None) -> bool: + mod = module or getattr(transformers, arch, None) + if mod is None: + return False + if hasattr(mod, "supports_backend"): + return mod.is_backend_compatible() + else: + return mod._supports_flex_attn + + +def resolve_transformers_fallback(model_config: ModelConfig, + architectures: list[str]): + for i, arch in enumerate(architectures): + if arch == "TransformersModel": + continue + auto_map: dict[str, str] = getattr(model_config.hf_config, "auto_map", + None) or dict() + # Make sure that config class is always initialized before model class, + # otherwise the model class won't be able to access the config class, + # the expected auto_map should have correct order like: + # "auto_map": { + # "AutoConfig": "--", + # "AutoModel": "--", + # "AutoModelFor": "--", + # }, + auto_modules = { + name: get_class_from_dynamic_module(module, model_config.model) + for name, module in sorted(auto_map.items(), key=lambda x: x[0]) + } + custom_model_module = auto_modules.get("AutoModel") + # TODO(Isotr0py): Further clean up these raises. + # perhaps handled them in _ModelRegistry._raise_for_unsupported? + if model_config.model_impl == ModelImpl.TRANSFORMERS: + if not is_transformers_impl_compatible(arch, custom_model_module): + raise ValueError( + f"The Transformers implementation of {arch} is not " + "compatible with vLLM.") + architectures[i] = "TransformersModel" + if model_config.model_impl == ModelImpl.AUTO: + if not is_transformers_impl_compatible(arch, custom_model_module): + raise ValueError( + f"{arch} has no vLLM implementation and the Transformers " + "implementation is not compatible with vLLM.") + logger.warning( + "%s has no vLLM implementation, falling back to Transformers " + "implementation. Some features may not be supported and " + "performance may not be optimal.", arch) + architectures[i] = "TransformersModel" + return architectures + + +def get_model_architecture( + model_config: ModelConfig) -> Tuple[Type[nn.Module], str]: + architectures = getattr(model_config.hf_config, "architectures", []) + + # Special handling for quantized Mixtral. + # FIXME(woosuk): This is a temporary hack. + mixtral_supported = [ + "fp8", "compressed-tensors", "gptq_marlin", "awq_marlin" + ] + + if (model_config.quantization is not None + and model_config.quantization not in mixtral_supported + and "MixtralForCausalLM" in architectures): + architectures = ["QuantMixtralForCausalLM"] + + vllm_supported_archs = ModelRegistry.get_supported_archs() + is_vllm_supported = any(arch in vllm_supported_archs + for arch in architectures) + if (not is_vllm_supported + or model_config.model_impl == ModelImpl.TRANSFORMERS): + architectures = resolve_transformers_fallback(model_config, + architectures) + + model_cls, arch = ModelRegistry.resolve_model_cls(architectures) + if model_config.task == "embed": + model_cls = as_embedding_model(model_cls) + elif model_config.task == "classify": + model_cls = as_classification_model(model_cls) + elif model_config.task == "reward": + model_cls = as_reward_model(model_cls) + + return model_cls, arch + + +def get_architecture_class_name(model_config: ModelConfig) -> str: + return get_model_architecture(model_config)[1] + + +@dataclass +class ParamMapping: + """ + A class to handle parameter mapping for model weight loading. + It creates a bidirectional mapping between packed parameters and their + constituent parts. + """ + packed_mapping: Dict[str, List[str]] + inverse_packed_mapping: Dict[str, Tuple[str, + int]] = field(default_factory=dict) + + def __post_init__(self): + for packed_name, sub_params in self.packed_mapping.items(): + # Skip self-contained cases (e.g., {"W_pack": ["W_pack"]}) + if len(sub_params) == 1 and sub_params[0] == packed_name: + continue + for index, param_name in enumerate(sub_params): + self.inverse_packed_mapping[param_name] = ( + packed_name, + index, + ) + + def get_sub_modules(self, + module_name: str) -> Optional[Tuple[str, List[str]]]: + for key, value in self.packed_mapping.items(): + if module_name.endswith(key): + return key, value + return None diff --git a/fastvideo/loader/weight_utils.py b/fastvideo/v1/loader/weight_utils.py similarity index 80% rename from fastvideo/loader/weight_utils.py rename to fastvideo/v1/loader/weight_utils.py index 3e9b88a68..e03c82fa7 100644 --- a/fastvideo/loader/weight_utils.py +++ b/fastvideo/v1/loader/weight_utils.py @@ -25,23 +25,23 @@ from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download from safetensors.torch import load_file, safe_open, save_file from tqdm.auto import tqdm -from vllm.config import LoadConfig, ModelConfig -from vllm.distributed import get_tensor_model_parallel_rank -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization import (QuantizationConfig, - get_quantization_config) -from vllm.platforms import current_platform -from vllm.utils import PlaceholderModule +# from vllm.config import LoadConfig, ModelConfig +from fastvideo.v1.distributed import get_tensor_model_parallel_rank +from fastvideo.v1.logger import init_logger +# from vllm.model_executor.layers.quantization import (QuantizationConfig, +# get_quantization_config) +from fastvideo.v1.platforms import current_platform +# from vllm.utils import PlaceholderModule -try: - from runai_model_streamer import SafetensorsStreamer -except (ImportError, OSError): - # see https://github.com/run-ai/runai-model-streamer/issues/26 - # OSError will be raised on arm64 platform - runai_model_streamer = PlaceholderModule( - "runai_model_streamer") # type: ignore[assignment] - SafetensorsStreamer = runai_model_streamer.placeholder_attr( - "SafetensorsStreamer") +# try: +# from runai_model_streamer import SafetensorsStreamer +# except (ImportError, OSError): +# # see https://github.com/run-ai/runai-model-streamer/issues/26 +# # OSError will be raised on arm64 platform +# runai_model_streamer = PlaceholderModule( +# "runai_model_streamer") # type: ignore[assignment] +# SafetensorsStreamer = runai_model_streamer.placeholder_attr( +# "SafetensorsStreamer") logger = init_logger(__name__) @@ -137,89 +137,89 @@ def convert_bin_to_safetensor_file( # TODO(woosuk): Move this to other place. -def get_quant_config(model_config: ModelConfig, - load_config: LoadConfig) -> QuantizationConfig: +# def get_quant_config(model_config: ModelConfig, +# load_config: LoadConfig) -> QuantizationConfig: - quant_cls = get_quantization_config(model_config.quantization) +# quant_cls = get_quantization_config(model_config.quantization) - # GGUF doesn't have config file - if model_config.quantization == "gguf": - return quant_cls.from_config({}) +# # GGUF doesn't have config file +# if model_config.quantization == "gguf": +# return quant_cls.from_config({}) - # Read the quantization config from the HF model config, if available. - hf_quant_config = getattr(model_config.hf_config, "quantization_config", - None) - # some vision model may keep quantization_config in their text_config - hf_text_config = getattr(model_config.hf_config, "text_config", None) - if hf_quant_config is None and hf_text_config is not None: - hf_quant_config = getattr(hf_text_config, "quantization_config", None) - if hf_quant_config is None: - # compressed-tensors uses a compressions_config - hf_quant_config = getattr(model_config.hf_config, "compression_config", - None) - if hf_quant_config is not None: - return quant_cls.from_config(hf_quant_config) - # In case of bitsandbytes/QLoRA, get quant config from the adapter model. - if model_config.quantization == "bitsandbytes": - if (not load_config.model_loader_extra_config - or "qlora_adapter_name_or_path" - not in load_config.model_loader_extra_config): - return quant_cls.from_config({"adapter_name_or_path": ""}) - model_name_or_path = load_config.model_loader_extra_config[ - "qlora_adapter_name_or_path"] +# # Read the quantization config from the HF model config, if available. +# hf_quant_config = getattr(model_config.hf_config, "quantization_config", +# None) +# # some vision model may keep quantization_config in their text_config +# hf_text_config = getattr(model_config.hf_config, "text_config", None) +# if hf_quant_config is None and hf_text_config is not None: +# hf_quant_config = getattr(hf_text_config, "quantization_config", None) +# if hf_quant_config is None: +# # compressed-tensors uses a compressions_config +# hf_quant_config = getattr(model_config.hf_config, "compression_config", +# None) +# if hf_quant_config is not None: +# return quant_cls.from_config(hf_quant_config) +# # In case of bitsandbytes/QLoRA, get quant config from the adapter model. +# if model_config.quantization == "bitsandbytes": +# if (not load_config.model_loader_extra_config +# or "qlora_adapter_name_or_path" +# not in load_config.model_loader_extra_config): +# return quant_cls.from_config({"adapter_name_or_path": ""}) +# model_name_or_path = load_config.model_loader_extra_config[ +# "qlora_adapter_name_or_path"] - else: - model_name_or_path = model_config.model - is_local = os.path.isdir(model_name_or_path) - if not is_local: - # Download the config files. - with get_lock(model_name_or_path, load_config.download_dir): - hf_folder = snapshot_download( - model_name_or_path, - revision=model_config.revision, - allow_patterns="*.json", - cache_dir=load_config.download_dir, - local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, - tqdm_class=DisabledTqdm, - ) - else: - hf_folder = model_name_or_path +# else: +# model_name_or_path = model_config.model +# is_local = os.path.isdir(model_name_or_path) +# if not is_local: +# # Download the config files. +# with get_lock(model_name_or_path, load_config.download_dir): +# hf_folder = snapshot_download( +# model_name_or_path, +# revision=model_config.revision, +# allow_patterns="*.json", +# cache_dir=load_config.download_dir, +# local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, +# tqdm_class=DisabledTqdm, +# ) +# else: +# hf_folder = model_name_or_path - possible_config_filenames = quant_cls.get_config_filenames() +# possible_config_filenames = quant_cls.get_config_filenames() - # If the quantization config is not found, use the default config. - if not possible_config_filenames: - return quant_cls() +# # If the quantization config is not found, use the default config. +# if not possible_config_filenames: +# return quant_cls() - config_files = glob.glob(os.path.join(hf_folder, "*.json")) +# config_files = glob.glob(os.path.join(hf_folder, "*.json")) - quant_config_files = [ - f for f in config_files if any( - f.endswith(x) for x in possible_config_filenames) - ] - if len(quant_config_files) == 0: - raise ValueError( - f"Cannot find the config file for {model_config.quantization}") - if len(quant_config_files) > 1: - raise ValueError( - f"Found multiple config files for {model_config.quantization}: " - f"{quant_config_files}") +# quant_config_files = [ +# f for f in config_files if any( +# f.endswith(x) for x in possible_config_filenames) +# ] +# if len(quant_config_files) == 0: +# raise ValueError( +# f"Cannot find the config file for {model_config.quantization}") +# if len(quant_config_files) > 1: +# raise ValueError( +# f"Found multiple config files for {model_config.quantization}: " +# f"{quant_config_files}") - quant_config_file = quant_config_files[0] - with open(quant_config_file) as f: - config = json.load(f) +# quant_config_file = quant_config_files[0] +# with open(quant_config_file) as f: +# config = json.load(f) - if model_config.quantization == "bitsandbytes": - config["adapter_name_or_path"] = model_name_or_path - elif model_config.quantization == "modelopt": - if config["producer"]["name"] == "modelopt": - return quant_cls.from_config(config) - else: - raise ValueError( - f"Unsupported quantization config" - f" found for {model_config.quantization} in {f}.") +# if model_config.quantization == "bitsandbytes": +# config["adapter_name_or_path"] = model_name_or_path +# elif model_config.quantization == "modelopt": +# if config["producer"]["name"] == "modelopt": +# return quant_cls.from_config(config) +# else: +# raise ValueError( +# f"Unsupported quantization config" +# f" found for {model_config.quantization} in {f}.") - return quant_cls.from_config(config) +# return quant_cls.from_config(config) def download_weights_from_hf( @@ -436,21 +436,21 @@ def safetensors_weights_iterator( yield name, param -def runai_safetensors_weights_iterator( - hf_weights_files: List[str] -) -> Generator[Tuple[str, torch.Tensor], None, None]: - """Iterate over the weights in the model safetensor files.""" - enable_tqdm = not torch.distributed.is_initialized( - ) or torch.distributed.get_rank() == 0 - with SafetensorsStreamer() as streamer: - for st_file in tqdm( - hf_weights_files, - desc="Loading safetensors using Runai Model Streamer", - disable=not enable_tqdm, - bar_format=_BAR_FORMAT, - ): - streamer.stream_file(st_file) - yield from streamer.get_tensors() +# def runai_safetensors_weights_iterator( +# hf_weights_files: List[str] +# ) -> Generator[Tuple[str, torch.Tensor], None, None]: +# """Iterate over the weights in the model safetensor files.""" +# enable_tqdm = not torch.distributed.is_initialized( +# ) or torch.distributed.get_rank() == 0 +# with SafetensorsStreamer() as streamer: +# for st_file in tqdm( +# hf_weights_files, +# desc="Loading safetensors using Runai Model Streamer", +# disable=not enable_tqdm, +# bar_format=_BAR_FORMAT, +# ): +# streamer.stream_file(st_file) +# yield from streamer.get_tensors() def pt_weights_iterator( diff --git a/fastvideo/logger.py b/fastvideo/v1/logger.py similarity index 96% rename from fastvideo/logger.py rename to fastvideo/v1/logger.py index 8855baa85..180c19bfb 100644 --- a/fastvideo/logger.py +++ b/fastvideo/v1/logger.py @@ -4,7 +4,7 @@ # Copyright 2023 The vLLM Authors. # Copyright 2023 The FastVideo Authors. -"""Logging configuration for fastvideo.""" +"""Logging configuration for fastvideo.v1.""" import datetime import json import logging @@ -17,7 +17,7 @@ from os import path from types import MethodType from typing import Any, Optional, cast -import fastvideo.envs as envs +import fastvideo.v1.envs as envs FASTVIDEO_CONFIGURE_LOGGING = envs.FASTVIDEO_CONFIGURE_LOGGING FASTVIDEO_LOGGING_CONFIG_PATH = envs.FASTVIDEO_LOGGING_CONFIG_PATH @@ -31,7 +31,7 @@ _DATE_FORMAT = "%m-%d %H:%M:%S" DEFAULT_LOGGING_CONFIG = { "formatters": { "fastvideo": { - "class": "fastvideo.logging_utils.NewLineFormatter", + "class": "fastvideo.v1.logging_utils.NewLineFormatter", "datefmt": _DATE_FORMAT, "format": _FORMAT, }, @@ -124,8 +124,8 @@ def _configure_fastvideo_root_logger() -> None: for formatter in logging_config.get("formatters", {}).values(): # This provides backwards compatibility after #10134. - if formatter.get("class") == "fastvideo.logging.NewLineFormatter": - formatter["class"] = "fastvideo.logging_utils.NewLineFormatter" + if formatter.get("class") == "fastvideo.v1.logging.NewLineFormatter": + formatter["class"] = "fastvideo.v1.logging_utils.NewLineFormatter" if logging_config: dictConfig(logging_config) diff --git a/fastvideo/logging_utils/__init__.py b/fastvideo/v1/logging_utils/__init__.py similarity index 54% rename from fastvideo/logging_utils/__init__.py rename to fastvideo/v1/logging_utils/__init__.py index 784da68c4..4a2511681 100644 --- a/fastvideo/logging_utils/__init__.py +++ b/fastvideo/v1/logging_utils/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -from fastvideo.logging_utils.formatter import NewLineFormatter +from fastvideo.v1.logging_utils.formatter import NewLineFormatter __all__ = [ "NewLineFormatter", diff --git a/fastvideo/logging_utils/formatter.py b/fastvideo/v1/logging_utils/formatter.py similarity index 100% rename from fastvideo/logging_utils/formatter.py rename to fastvideo/v1/logging_utils/formatter.py diff --git a/fastvideo/v1/models/__init__.py b/fastvideo/v1/models/__init__.py new file mode 100644 index 000000000..76d2ac2dc --- /dev/null +++ b/fastvideo/v1/models/__init__.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch.nn as nn +from typing import Dict, Any +from fastvideo.v1.loader.loader import get_model_loader +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +import os +from transformers import PretrainedConfig + +logger = init_logger(__name__) + +def get_transformers_model(model_config: PretrainedConfig, model_path: str, inference_args: InferenceArgs) -> nn.Module: + """ + Load a transformers model. + + Args: + architecture: The model architecture name + model_path: Path to the model directory + inference_args: Inference arguments + component_name: Optional component name (e.g., "text_encoder", "vae") + If provided, will load the config from the component subdirectory + + Returns: + The loaded model + """ + model_loader = get_model_loader() + return model_loader.load_model(model_config, model_path, inference_args) + +def get_diffusers_model(model_config: Dict[str, Any], model_path: str, inference_args: InferenceArgs) -> nn.Module: + """ + Load a diffusers model. + """ + raise NotImplementedError("Diffusers models are not supported yet.") + +def get_scheduler(module_path: str, architecture: str, inference_args: InferenceArgs) -> Dict: + """Create a scheduler based on the inference args. Can be overridden by subclasses.""" + if hasattr(inference_args, 'denoise_type') and inference_args.denoise_type == "flow": + # TODO(will): add schedulers to register or create a new scheduler registry + # TODO(will): default to config file but allow override through + # inference args. Currently only uses inference args. + from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchDiscreteScheduler + return FlowMatchDiscreteScheduler( + shift=inference_args.flow_shift, + reverse=inference_args.flow_reverse, + solver=inference_args.flow_solver, + ) + else: + raise ValueError(f"Invalid denoise type: {inference_args.denoise_type}") + +__all__ = [ + "set_random_seed", + "BasevLLMParameter", + "PackedvLLMParameter", + "get_model", + "get_scheduler", +] diff --git a/fastvideo/models/dits/base.py b/fastvideo/v1/models/dits/base.py similarity index 100% rename from fastvideo/models/dits/base.py rename to fastvideo/v1/models/dits/base.py diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py similarity index 97% rename from fastvideo/models/dits/hunyuanvideo.py rename to fastvideo/v1/models/dits/hunyuanvideo.py index f407e645d..96431b546 100644 --- a/fastvideo/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -1,16 +1,16 @@ import torch import torch.nn as nn from typing import Optional, Tuple, List -from fastvideo.attention.distributed_attn import DistributedAttention, LocalAttention -from fastvideo.layers.linear import ReplicatedLinear -from fastvideo.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift -from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify -from fastvideo.layers.rotary_embedding import _apply_rotary_emb, get_rotary_pos_embed -from fastvideo.distributed.parallel_state import get_sequence_model_parallel_world_size, get_sequence_model_parallel_rank -# TODO: RMSNorm .... -from fastvideo.models.hunyuan.modules.norm_layers import RMSNorm -from fastvideo.layers.mlp import MLP -from fastvideo.models.dits.base import BaseDiT +from fastvideo.v1.attention.flash_attn import DistributedAttention, LocalAttention +from fastvideo.v1.layers.linear import ReplicatedLinear +from fastvideo.v1.layers.layernorm import LayerNormScaleShift, ScaleResidual, ScaleResidualLayerNormScaleShift +from fastvideo.v1.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection, unpatchify +from fastvideo.v1.layers.rotary_embedding import _apply_rotary_emb, get_rotary_pos_embed +from fastvideo.v1.distributed.parallel_state import get_sequence_model_parallel_world_size, get_sequence_model_parallel_rank +# TODO(will-PY-refactor): RMSNorm .... +from fastvideo.v1.v0_reference_src.models.hunyuan.modules.norm_layers import RMSNorm +from fastvideo.v1.layers.mlp import MLP +from fastvideo.v1.models.dits.base import BaseDiT @@ -329,7 +329,7 @@ class MMSingleStreamBlock(nn.Module): -class HunyuanVideoDiT(BaseDiT): +class HunyuanVideoTransformer3DModel(BaseDiT): """ HunyuanVideo Transformer backbone adapted for distributed training. @@ -799,9 +799,8 @@ class FinalLayer(nn.Module): ) def forward(self, x, c): - # What the fuck HF? Why you change the scale and shift order here??? + # What the heck HF? Why you change the scale and shift order here??? scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1) x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) x, _ = self.linear(x) return x - diff --git a/fastvideo/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py similarity index 56% rename from fastvideo/models/encoders/clip.py rename to fastvideo/v1/models/encoders/clip.py index 907b30f11..b141db13f 100644 --- a/fastvideo/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -5,20 +5,30 @@ from typing import Iterable, Optional, Set, Tuple, Union import torch import torch.nn as nn -from transformers import CLIPVisionConfig +from transformers import CLIPVisionConfig, CLIPTextConfig +from transformers.modeling_outputs import BaseModelOutputWithPooling +# from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask from vllm.attention.layer import MultiHeadAttention -from fastvideo.distributed import divide, get_tensor_model_parallel_world_size -from fastvideo.layers.activation import get_act_fn -from fastvideo.layers.linear import (ColumnParallelLinear, +# from fastvideo.v1.attention.flash_attn import LocalAttention +from fastvideo.v1.distributed import divide, get_tensor_model_parallel_world_size +from fastvideo.v1.layers.activation import get_act_fn +from fastvideo.v1.layers.linear import (ColumnParallelLinear, QKVParallelLinear, RowParallelLinear) -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.model_loader.weight_utils import default_weight_loader +# TODO: support quantization +# from vllm.model_executor.layers.quantization import QuantizationConfig +from fastvideo.v1.loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsQuant from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) + +class QuantizationConfig: + pass class CLIPEncoderInfo(VisionEncoderInfo[CLIPVisionConfig]): @@ -88,6 +98,48 @@ class CLIPVisionEmbeddings(nn.Module): return embeddings +class CLIPTextEmbeddings(nn.Module): + + def __init__(self, config: CLIPTextConfig): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + max_position_embedding = self.position_embedding.weight.shape[0] + + if seq_length > max_position_embedding: + raise ValueError( + f"Sequence length must be less than max_position_embeddings (got `sequence length`: " + f"{seq_length} and max_position_embeddings: {max_position_embedding}" + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + class CLIPAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" @@ -143,8 +195,14 @@ class CLIPAttention(nn.Module): qkv_states, _ = self.qkv_proj(hidden_states) query_states, key_states, value_states = qkv_states.chunk(3, dim=-1) - out = self.attn(query_states, key_states, value_states) - attn_output, _ = self.out_proj(out) + # use flash_attn_func + from flash_attn import flash_attn_func + query_states = query_states.reshape(query_states.shape[0], query_states.shape[1], self.num_heads_per_partition, self.head_dim) + key_states = key_states.reshape(key_states.shape[0], key_states.shape[1], self.num_heads_per_partition, self.head_dim) + value_states = value_states.reshape(value_states.shape[0], value_states.shape[1], self.num_heads_per_partition, self.head_dim) + attn_output = flash_attn_func(query_states, key_states, value_states,softmax_scale=self.scale, causal=True) + attn_output = attn_output.reshape(attn_output.shape[0], attn_output.shape[1], self.num_heads_per_partition * self.head_dim) + attn_output, _ = self.out_proj(attn_output) return attn_output, None @@ -183,7 +241,7 @@ class CLIPEncoderLayer(nn.Module): def __init__( self, - config: CLIPVisionConfig, + config: CLIPTextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: @@ -263,6 +321,191 @@ class CLIPEncoder(nn.Module): if return_all_hidden_states: return hidden_states_pool return hidden_states + + + +class CLIPTextTransformer(nn.Module): + + def __init__(self, + config: CLIPTextConfig, + quant_config: Optional[QuantizationConfig] = None, + *, + num_hidden_layers_override: Optional[int] = None, + prefix: str = ""): + super().__init__() + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPTextEmbeddings(config) + + self.encoder = CLIPEncoder(config, + quant_config=quant_config, + num_hidden_layers_override=num_hidden_layers_override, + prefix=prefix) + + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + # For `pooled_output` computation + self.eos_token_id = config.eos_token_id + + # For attention mask, it differs between `flash_attention_2` and other attention implementations + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + Returns: + + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if input_ids is None: + raise ValueError("You have to specify input_ids") + + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + + hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids) + + # CLIP's text model uses causal mask, prepare it here. + # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324 + # causal_attention_mask = _create_4d_causal_attention_mask( + # input_shape, hidden_states.dtype, device=hidden_states.device + # ) + + # # expand attention_mask + # if attention_mask is not None and not self._use_flash_attention_2: + # raise NotImplementedError("attention_mask is not supported for CLIPTextTransformer") + # # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + # attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype) + + encoder_outputs = self.encoder( + inputs_embeds=hidden_states, + # attention_mask=attention_mask, + # causal_attention_mask=causal_attention_mask, + # output_attentions=output_attentions, + return_all_hidden_states=output_hidden_states, + # return_dict=return_dict, + ) + + last_hidden_state = encoder_outputs[-1] + last_hidden_state = self.final_layer_norm(last_hidden_state) + + if self.eos_token_id == 2: + # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here. + # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added + # ------------------------------------------------------------ + # text_embeds.shape = [batch_size, sequence_length, transformer.width] + # take features from the eot embedding (eot_token is the highest number in each sequence) + # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14 + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1), + ] + else: + # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) + pooled_output = last_hidden_state[ + torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), + # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) + # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) + (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id) + .int() + .argmax(dim=-1), + ] + + if not return_dict: + return (last_hidden_state, pooled_output) + encoder_outputs[1:] + + # return last_hidden_state + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs, + # attentions=encoder_outputs.attentions, + ) + + +class CLIPTextModel(nn.Module): + + def __init__( + self, + config: CLIPTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.text_model = CLIPTextTransformer( + config=config, + quant_config=quant_config, + prefix=prefix) + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + return self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=None, + ) + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + + # Define mapping for stacked parameters + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + # Handle q_proj, k_proj, v_proj -> qkv_proj mapping + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name in name: + # Replace the weight name with the parameter name + model_param_name = name.replace(weight_name, param_name) + + if model_param_name in params_dict: + param = params_dict[model_param_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(model_param_name) + break + else: + # Use default weight loader for all other parameters + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params class CLIPVisionTransformer(nn.Module): diff --git a/fastvideo/v1/models/encoders/llama.py b/fastvideo/v1/models/encoders/llama.py new file mode 100644 index 000000000..d5b8169c1 --- /dev/null +++ b/fastvideo/v1/models/encoders/llama.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only LLaMA model compatible with HuggingFace weights.""" +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union + +import torch +from torch import nn +from transformers import LlamaConfig +from transformers.modeling_outputs import BaseModelOutputWithPast + +from vllm.attention.layer import MultiHeadAttention + +from fastvideo.v1.distributed import get_tensor_model_parallel_world_size +from fastvideo.v1.layers.activation import SiluAndMul +from fastvideo.v1.layers.layernorm import RMSNorm +from fastvideo.v1.layers.linear import (MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear) +# from vllm.model_executor.layers.quantization import QuantizationConfig + +from fastvideo.v1.layers.rotary_embedding import get_rope +from fastvideo.v1.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.v1.loader.weight_utils import ( + default_weight_loader, maybe_remap_kv_scale_name) + +from .utils import (extract_layer_index) + +class QuantizationConfig: + pass + + +class LlamaMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + # output_size=intermediate_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + x, _ = self.gate_up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj(x) + return x + + +class LlamaAttention(nn.Module): + + def __init__(self, + config: LlamaConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + bias_o_proj: bool = False, + prefix: str = "") -> None: + super().__init__() + layer_idx = extract_layer_index(prefix) + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr(config, "head_dim", + self.hidden_size // self.total_num_heads) + # Phi models introduced a partial_rotary_factor parameter in the config + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias_o_proj, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + is_neox_style = True + is_gguf = quant_config and quant_config.get_name() == "gguf" + if is_gguf and config.model_type == "llama": + is_neox_style = False + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + ) + + self.attn = MultiHeadAttention(self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # print(f"query_states sum before apply_rotary_pos_emb: {q.float().sum()}") + q, k = self.rotary_emb(positions, q, k) + # attn_output = self.attn(q, k, v) + # use flash_attn_func + from flash_attn import flash_attn_func + # reshape q, k, v to (batch_size, seq_len, num_heads, head_dim) + batch_size = q.shape[0] + seq_len = q.shape[1] + q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + # import pdb; pdb.set_trace() + attn_output = flash_attn_func(q, k, v, softmax_scale=self.scaling, causal=True) + print(f"shape of attn_output: {attn_output.shape}") + print(f"num_kv_heads: {self.num_kv_heads}") + print(f"num_heads: {self.num_heads}") + print(f"head_dim: {self.head_dim}") + attn_output = attn_output.reshape(batch_size, seq_len, self.num_heads * self.head_dim) + + output, _ = self.o_proj(attn_output) + return output + + +class LlamaDecoderLayer(nn.Module): + + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings) + max_position_embeddings = getattr(config, "max_position_embeddings", + 8192) + # Support abacusai/Smaug-72B-v0.1 with attention_bias + # Support internlm/internlm-7b with bias + attention_bias = getattr(config, "attention_bias", False) or getattr( + config, "bias", False) + bias_o_proj = attention_bias + # support internlm/internlm3-8b with qkv_bias + if hasattr(config, 'qkv_bias'): + attention_bias = config.qkv_bias + + self.self_attn = LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", + config.num_attention_heads), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + bias_o_proj=bias_o_proj, + prefix=f"{prefix}.self_attn", + ) + self.mlp = LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, + eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, + eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual) + + hidden_states = self.self_attn(positions=positions, + hidden_states=hidden_states) + + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class LlamaModel(nn.Module): + + def __init__(self, + config: LlamaConfig, + prefix: str = "", + layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer): + super().__init__() + + quant_config = None + lora_config = None + + self.config = config + self.quant_config = quant_config + lora_vocab = (lora_config.lora_extra_vocab_size * + (lora_config.max_loras or 1)) if lora_config else 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=quant_config, + ) + + self.layers = nn.ModuleList([ + layer_type( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}" + ) for i in range(config.num_hidden_layers) + ]) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: Optional[torch.Tensor], + positions: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + ) -> torch.Tensor: + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + + if positions is None: + positions = torch.arange( + 0, hidden_states.shape[1], device=hidden_states.device + ).unsqueeze(0) + + all_hidden_states = () if output_hidden_states else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + hidden_states, residual = layer(positions, hidden_states, residual) + + hidden_states, _ = self.norm(hidden_states, residual) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + # TODO(will): maybe unify the output format with other models and use + # our own class + output = BaseModelOutputWithPast( + last_hidden_state=hidden_states, + # past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + # attentions=all_self_attns, + ) + + return output + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if ("rotary_emb.cos_cached" in name + or "rotary_emb.sin_cached" in name): + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if (self.quant_config is not None and + (scale_name := self.quant_config.get_cache_scale(name))): + # Loading kv cache quantization scales + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else + loaded_weight[0]) + weight_loader(param, loaded_weight) + loaded_params.add(scale_name) + continue + if "scale" in name: + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/fastvideo/v1/models/encoders/utils.py b/fastvideo/v1/models/encoders/utils.py new file mode 100644 index 000000000..db1d24fe5 --- /dev/null +++ b/fastvideo/v1/models/encoders/utils.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 + +import itertools +from dataclasses import dataclass, field +from typing import (Callable, Dict, Iterable, List, Literal, Mapping, Optional, + Protocol, Set, Tuple, Union, overload) + +import torch +import torch.nn as nn +from torch.func import functional_call +from transformers import PretrainedConfig + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.loader.weight_utils import default_weight_loader + +logger = init_logger(__name__) + +WeightsMapping = Mapping[str, Optional[str]] +"""If a key maps to a value of `None`, the corresponding weight is ignored.""" + + +@dataclass +class WeightsMapper: + """Maps the name of each weight if they match the following patterns.""" + + orig_to_new_substr: WeightsMapping = field(default_factory=dict) + orig_to_new_prefix: WeightsMapping = field(default_factory=dict) + orig_to_new_suffix: WeightsMapping = field(default_factory=dict) + + def _map_name(self, key: str) -> Optional[str]: + for substr, new_key in self.orig_to_new_substr.items(): + if substr in key: + if new_key is None: + return None + + key = key.replace(substr, new_key, 1) + + for prefix, new_key in self.orig_to_new_prefix.items(): + if key.startswith(prefix): + if new_key is None: + return None + + key = key.replace(prefix, new_key, 1) + + for suffix, new_key in self.orig_to_new_suffix.items(): + if key.endswith(suffix): + if new_key is None: + return None + + key = new_key.join(key.rsplit(suffix, 1)) + + return key + + def apply( + self, weights: Iterable[Tuple[str, torch.Tensor]] + ) -> Iterable[Tuple[str, torch.Tensor]]: + return ((out_name, data) for name, data in weights + if (out_name := self._map_name(name)) is not None) + + +class AutoWeightsLoader: + """ + Helper class to load weights into a :class:`torch.nn.Module`. It is able + to automatically detect child modules and parameters while iterating over + the weights only once. + + The weight loading logic for individual modules can be overridden + by defining a ``load_weights`` method. + + Similarly, the weight loading logic for individual parameters can be + overridden by defining a ``weight_loader`` method. + + Detailed weight loading information can be viewed by setting the + environment variable ``VLLM_LOGGING_LEVEL=DEBUG``. + """ + + def __init__( + self, + module: nn.Module, + *, + skip_prefixes: Optional[List[str]] = None, + ignore_unexpected_prefixes: Optional[List[str]] = None, + ) -> None: + super().__init__() + + self.module = module + self.skip_prefixes = skip_prefixes or [] + self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or [] + + def _groupby_prefix( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[Tuple[str, Iterable[Tuple[str, torch.Tensor]]]]: + weights_by_parts = ((weight_name.split(".", 1), weight_data) + for weight_name, weight_data in weights) + + for prefix, group in itertools.groupby(weights_by_parts, + key=lambda x: x[0][0]): + yield ( + prefix, + # Because maxsplit=1 in weight_name.split(...), + # the length of `parts` must either be 1 or 2 + (("" if len(parts) == 1 else parts[1], weights_data) + for parts, weights_data in group), + ) + + def _get_qualname(self, prefix: str, rest: str) -> str: + if prefix == "": + return rest + if rest == "": + return prefix + + return ".".join((prefix, rest)) + + def _can_skip(self, qualname: str) -> bool: + return any(qualname.startswith(p) for p in self.skip_prefixes) + + def _can_ignore_unexpected(self, qualname: str) -> bool: + return any( + qualname.startswith(p) for p in self.ignore_unexpected_prefixes) + + def _load_param( + self, + base_prefix: str, + param: nn.Parameter, + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[str]: + for weight_name, weight_data in weights: + weight_qualname = self._get_qualname(base_prefix, weight_name) + + if self._can_skip(weight_qualname): + logger.debug("Skipping weight %s", weight_qualname) + + continue + + if weight_name != "": + if self._can_ignore_unexpected(weight_qualname): + logger.debug("Ignoring weight %s", weight_qualname) + + continue + + raise ValueError( + f"Attempted to load nested weight '{weight_qualname}' " + f"into a single parameter '{base_prefix}'") + + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, weight_data) + + logger.debug("Loaded weight %s with shape %s", weight_qualname, + param.shape) + + yield weight_qualname + + def _load_module( + self, + base_prefix: str, + module: nn.Module, + weights: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[str]: + if isinstance(module, PPMissingLayer): + return + + # Avoid infinite recursion since this function is typically + # called inside load_weights of the module itself + if module != self.module: + module_load_weights = getattr(module, "load_weights", None) + if callable(module_load_weights): + loaded_params = module_load_weights(weights) + if loaded_params is None: + logger.warning( + "Unable to collect loaded parameters " + "for module %s", module) + else: + yield from map( + lambda x: self._get_qualname(base_prefix, x), + loaded_params, + ) + + child_modules = dict(module.named_children()) + child_params = dict(module.named_parameters(recurse=False)) + + for child_prefix, child_weights in self._groupby_prefix(weights): + prefix = self._get_qualname(base_prefix, child_prefix) + + if child_prefix in child_modules: + if self._can_skip(prefix + "."): + logger.debug("Skipping module %s", prefix) + + continue + + yield from self._load_module(prefix, + child_modules[child_prefix], + child_weights) + elif child_prefix in child_params: + if self._can_skip(prefix): + logger.debug("Skipping param %s", prefix) + + continue + + yield from self._load_param(prefix, child_params[child_prefix], + child_weights) + else: + can_skip_module = self._can_skip(prefix + ".") + can_skip_param = self._can_skip(prefix) + if can_skip_module or can_skip_param: + logger.debug("Skipping missing %s", prefix) + + continue + + can_ignore_module = self._can_ignore_unexpected(prefix + ".") + can_ignore_param = self._can_ignore_unexpected(prefix) + if can_ignore_module or can_ignore_param: + logger.debug("Ignoring missing %s", prefix) + + continue + + msg = (f"There is no module or parameter named '{prefix}' " + f"in {type(self.module).__name__}") + raise ValueError(msg) + + def load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + *, + mapper: Optional[WeightsMapper] = None, + ) -> Set[str]: + if mapper is not None: + weights = mapper.apply(weights) + + autoloaded_weights = set(self._load_module("", self.module, weights)) + return autoloaded_weights + + +@overload +def flatten_bn(x: torch.Tensor) -> torch.Tensor: + ... + + +@overload +def flatten_bn(x: List[torch.Tensor]) -> List[torch.Tensor]: + ... + + +@overload +def flatten_bn( + x: Union[List[torch.Tensor], torch.Tensor], + *, + concat: Literal[True], +) -> torch.Tensor: + ... + + +@overload +def flatten_bn( + x: Union[List[torch.Tensor], torch.Tensor], + *, + concat: bool = False, +) -> Union[List[torch.Tensor], torch.Tensor]: + ... + + +def flatten_bn( + x: Union[List[torch.Tensor], torch.Tensor], + *, + concat: bool = False, +) -> Union[List[torch.Tensor], torch.Tensor]: + """ + Flatten the ``B`` and ``N`` dimensions of batched multimodal inputs. + + The input tensor should have shape ``(B, N, ...)```. + """ + if isinstance(x, torch.Tensor): + return x.flatten(0, 1) + + if concat: + return torch.cat(x) + + return [x_n for x_b in x for x_n in x_b] + + +class LayerFn(Protocol): + + def __call__(self, prefix: str) -> torch.nn.Module: + ... + + +class PPMissingLayer(torch.nn.Identity): + """ + A placeholder layer for missing layers in a pipeline parallel model. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + + +_CPU_OFFLOAD_BYTES = 0 +_CPU_OFFLOAD_MAX_BYTES = 0 + + +def set_cpu_offload_max_bytes(max_bytes: int) -> None: + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + _CPU_OFFLOAD_BYTES = 0 + _CPU_OFFLOAD_MAX_BYTES = max_bytes + + +def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module: + device = next(module.parameters()).device + + if device == torch.device("cpu"): + return module + + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + return module + + pin_memory = True + + # offload parameters to CPU + # use pin_memory if possible, which helps cudagraph capture speed + offloaded_parameters = False + for p in module.parameters(): + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + # we use per-parameter offloading + # one module might have some parameters offloaded and some not + break + + # `torch.empty_like` does not support `pin_memory` argument + cpu_data = torch.empty_strided(size=p.data.size(), + stride=p.data.stride(), + dtype=p.data.dtype, + layout=p.data.layout, + device='cpu', + pin_memory=pin_memory) + cpu_data.copy_(p.data) + p.data = cpu_data + _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() + offloaded_parameters = True + + if offloaded_parameters: + original_forward = module.forward + + def forward(*args, **kwargs): + module.forward = original_forward + device_state = { + # here we blindly call `to(device)` + # if the parameter is already on the device, it will be a no-op + k: v.to(device, non_blocking=True) + for k, v in module.state_dict().items() + } + output = functional_call(module, + device_state, + args=args, + kwargs=kwargs) + module.forward = forward + return output + + module.forward = forward + + return module + + +def maybe_prefix(prefix: str, name: str) -> str: + """Add a prefix to a name if the prefix is non-empty. + + Args: + prefix: The prefix to add. If empty, no prefix will be added. + name: The name to potentially prefix. + + Returns: + The string "prefix.name" if prefix was non-empty, otherwise just "name". + """ + return name if not prefix else f"{prefix}.{name}" + + +def extract_layer_index(layer_name: str) -> int: + """ + Extract the layer index from the module name. + Examples: + - "encoder.layers.0" -> 0 + - "encoder.layers.1.self_attn" -> 1 + - "2.self_attn" -> 2 + - "model.encoder.layers.0.sub.1" -> ValueError + """ + subnames = layer_name.split(".") + int_vals: List[int] = [] + for subname in subnames: + try: + int_vals.append(int(subname)) + except ValueError: + continue + assert len(int_vals) == 1, (f"layer name {layer_name} should" + " only contain one integer") + return int_vals[0] + + +def cast_overflow_tensors( + tensors: torch.Tensor, + offset: float = 1000, +) -> torch.Tensor: + if tensors.isinf().any() or tensors.isnan().any(): + clamp_value = torch.finfo(tensors.dtype).max - offset + tensors = torch.clamp(tensors, min=-clamp_value, max=clamp_value) + return tensors + +def make_layers( + num_hidden_layers: int, + layer_fn: LayerFn, + prefix: str, +) -> torch.nn.ModuleList: + """Make a list of layers with the given layer function. + + Args: + num_hidden_layers: Number of layers to create. + layer_fn: Function that creates a layer given a prefix. + prefix: Prefix for layer names. + + Returns: + A ModuleList containing all layers. + """ + modules = torch.nn.ModuleList([ + maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}")) + for idx in range(num_hidden_layers) + ]) + return modules + + +def add_prefix(name: str, prefix: str) -> str: + """Add a weight path prefix to a module name. + + Args: + name: base module name. + prefix: weight prefix str to added to the front of `name` concatenated with `.`. + + Returns: + The string `prefix.name` if prefix is non-empty, otherwise just `name`. + """ + return name if not prefix else f"{prefix}.{name}" diff --git a/fastvideo/models/encoders/vision.py b/fastvideo/v1/models/encoders/vision.py similarity index 97% rename from fastvideo/models/encoders/vision.py rename to fastvideo/v1/models/encoders/vision.py index 778e5640b..b0921f2a4 100644 --- a/fastvideo/models/encoders/vision.py +++ b/fastvideo/v1/models/encoders/vision.py @@ -6,11 +6,11 @@ from typing import Final, Generic, Optional, Protocol, TypeVar, Union import torch from transformers import PretrainedConfig -import vllm.envs as envs +import fastvideo.v1.envs as envs from vllm.attention.selector import (backend_name_to_enum, get_global_forced_attn_backend) -from vllm.logger import init_logger -from vllm.platforms import _Backend, current_platform +from fastvideo.v1.logger import init_logger +from fastvideo.v1.platforms import _Backend, current_platform logger = init_logger(__name__) diff --git a/fastvideo/v1/models/hf_transformer_utils.py b/fastvideo/v1/models/hf_transformer_utils.py new file mode 100644 index 000000000..7c22b971c --- /dev/null +++ b/fastvideo/v1/models/hf_transformer_utils.py @@ -0,0 +1,259 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utilities for Huggingface Transformers.""" + +import contextlib +import os +import warnings +from pathlib import Path +from typing import Dict, Optional, Type, Union, Any +import json + +from huggingface_hub import snapshot_download +from transformers import ( + AutoConfig, + AutoProcessor, + AutoTokenizer, + PretrainedConfig, + PreTrainedTokenizer, + PreTrainedTokenizerFast, +) +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES + +# from fastvideo.v1.models.configs import ChatGLMConfig, DbrxConfig, ExaoneConfig, Qwen2_5_VLConfig + +_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = { + # ChatGLMConfig.model_type: ChatGLMConfig, + # DbrxConfig.model_type: DbrxConfig, + # ExaoneConfig.model_type: ExaoneConfig, + # Qwen2_5_VLConfig.model_type: Qwen2_5_VLConfig, +} + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def download_from_hf(model_path: str): + if os.path.exists(model_path): + return model_path + + return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"]) + + +def get_hf_config( + model: str, + trust_remote_code: bool, + revision: Optional[str] = None, + model_override_args: Optional[dict] = None, + inference_args: Optional[dict] = None, + **kwargs, +): + is_gguf = check_gguf_file(model) + if is_gguf: + kwargs["gguf_file"] = model + model = Path(model).parent + + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + if config.model_type in _CONFIG_REGISTRY: + config_class = _CONFIG_REGISTRY[config.model_type] + config = config_class.from_pretrained(model, revision=revision) + # NOTE(HandH1998): Qwen2VL requires `_name_or_path` attribute in `config`. + setattr(config, "_name_or_path", model) + if model_override_args: + config.update(model_override_args) + + # Special architecture mapping check for GGUF models + if is_gguf: + if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + raise RuntimeError(f"Can't get gguf config for {config.model_type}.") + model_type = MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type] + config.update({"architectures": [model_type]}) + + return config + + +def get_diffusers_config( + model: str, + inference_args: Optional[dict] = None, +) -> Dict[str, Any]: + """Gets a configuration for the given diffusers model. + + Args: + model: The model name or path. + inference_args: Optional inference arguments to override in the config. + + Returns: + The loaded configuration. + """ + # Check if the model path exists + if os.path.exists(model): + config_file = os.path.join(model, "config.json") + if os.path.exists(config_file): + try: + # Load the config directly from the file + with open(config_file, "r") as f: + config_dict = json.load(f) + + # TODO(will): apply any overrides from inference args + return config_dict + except Exception as e: + raise RuntimeError(f"Failed to load diffusers config from {config_file}: {e}") + else: + raise RuntimeError(f"Diffusers config file not found at {model}") + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# NOTE: The ordering here is important. Some models have two of these and we +# have a preference for which value gets used. +CONTEXT_LENGTH_KEYS = [ + "max_sequence_length", + "seq_length", + "max_seq_len", + "model_max_length", + "max_position_embeddings", +] + + +def get_context_length(config): + """Get the context length of a model from a huggingface model configs.""" + text_config = config + rope_scaling = getattr(text_config, "rope_scaling", None) + if rope_scaling: + rope_scaling_factor = rope_scaling.get("factor", 1) + if "original_max_position_embeddings" in rope_scaling: + rope_scaling_factor = 1 + if rope_scaling.get("rope_type", None) == "llama3": + rope_scaling_factor = 1 + else: + rope_scaling_factor = 1 + + for key in CONTEXT_LENGTH_KEYS: + val = getattr(text_config, key, None) + if val is not None: + return int(rope_scaling_factor * val) + return 2048 + + +# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file. +_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer" + + +def get_tokenizer( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + """Gets a tokenizer for the given model name via Huggingface.""" + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + + is_gguf = check_gguf_file(tokenizer_name) + if is_gguf: + kwargs["gguf_file"] = tokenizer_name + tokenizer_name = Path(tokenizer_name).parent + + try: + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + tokenizer_revision=tokenizer_revision, + clean_up_tokenization_spaces=False, + **kwargs, + ) + except TypeError as e: + # The LLaMA tokenizer causes a protobuf error in some environments. + err_msg = ( + "Failed to load the tokenizer. If you are using a LLaMA V1 model " + f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the " + "original tokenizer." + ) + raise RuntimeError(err_msg) from e + except ValueError as e: + # If the error pertains to the tokenizer class not existing or not + # currently being imported, suggest using the --trust-remote-code flag. + if not trust_remote_code and ( + "does not exist or is not currently imported." in str(e) + or "requires you to execute the tokenizer file" in str(e) + ): + err_msg = ( + "Failed to load the tokenizer. If the tokenizer is a custom " + "tokenizer not yet available in the HuggingFace transformers " + "library, consider setting `trust_remote_code=True` in LLM " + "or using the `--trust-remote-code` flag in the CLI." + ) + raise RuntimeError(err_msg) from e + else: + raise e + + if not isinstance(tokenizer, PreTrainedTokenizerFast): + warnings.warn( + "Using a slow tokenizer. This might cause a significant " + "slowdown. Consider using a fast tokenizer instead." + ) + + attach_additional_stop_token_ids(tokenizer) + return tokenizer + + +def get_processor( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: Optional[str] = None, + **kwargs, +): + processor = AutoProcessor.from_pretrained( + tokenizer_name, + *args, + trust_remote_code=trust_remote_code, + tokenizer_revision=tokenizer_revision, + **kwargs, + ) + + attach_additional_stop_token_ids(processor.tokenizer) + return processor + + +def attach_additional_stop_token_ids(tokenizer): + # Special handling for stop token <|eom_id|> generated by llama 3 tool use. + if "<|eom_id|>" in tokenizer.get_added_vocab(): + tokenizer.additional_stop_token_ids = set( + [tokenizer.get_added_vocab()["<|eom_id|>"]] + ) + else: + tokenizer.additional_stop_token_ids = None + + +def check_gguf_file(model: Union[str, os.PathLike]) -> bool: + """Check if the file is a GGUF model.""" + model = Path(model) + if not model.is_file(): + return False + elif model.suffix == ".gguf": + return True + + with open(model, "rb") as f: + header = f.read(4) + return header == b"GGUF" diff --git a/fastvideo/v1/models/loader.py b/fastvideo/v1/models/loader.py new file mode 100644 index 000000000..2d2257a98 --- /dev/null +++ b/fastvideo/v1/models/loader.py @@ -0,0 +1,325 @@ +from abc import ABC, abstractmethod + +import torch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +import os +from fastvideo.v1.distributed.parallel_state import initialize_sequence_parallel_group +import glob +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.loader.loader import get_model_loader +from transformers import PretrainedConfig, AutoTokenizer +from fastvideo.v1.models.hf_transformer_utils import get_hf_config, get_diffusers_config +from fastvideo.v1.models import get_scheduler +from fastvideo.v1.models.registry import ModelRegistry +from safetensors.torch import load_file as safetensors_load_file + +import json + + +logger = init_logger(__name__) + + +class ComponentLoader(ABC): + """Base class for loading a specific type of model component.""" + + def __init__(self, device=None): + self.device = device + + @abstractmethod + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """ + Load the component based on the model path, architecture, and inference args. + + Args: + model_path: Path to the component model + architecture: Architecture of the component model + inference_args: Inference arguments + + Returns: + The loaded component + """ + raise NotImplementedError + + @classmethod + def for_module_type(cls, module_type: str, transformers_or_diffusers: str) -> 'ComponentLoader': + """ + Factory method to create a component loader for a specific module type. + + Args: + module_type: Type of module (e.g., "vae", "text_encoder", "transformer", "scheduler") + transformers_or_diffusers: Whether the module is from transformers or diffusers + + Returns: + A component loader for the specified module type + """ + # Map of module types to their loader classes and expected library + module_loaders = { + "scheduler": (SchedulerLoader, "diffusers"), + "transformer": (TransformerLoader, "diffusers"), + "vae": (VAELoader, "diffusers"), + "text_encoder": (TextEncoderLoader, "transformers"), + "text_encoder_2": (TextEncoderLoader, "transformers"), + "tokenizer": (TokenizerLoader, "transformers"), + "tokenizer_2": (TokenizerLoader, "transformers"), + } + + if module_type in module_loaders: + loader_cls, expected_library = module_loaders[module_type] + # Assert that the library matches what's expected for this module type + assert transformers_or_diffusers == expected_library, f"{module_type} must be loaded from {expected_library}, got {transformers_or_diffusers}" + return loader_cls() + + # For unknown module types, use a generic loader + logger.warning(f"No specific loader found for module type: {module_type}. Using generic loader.") + return GenericComponentLoader(transformers_or_diffusers) + +class TextEncoderLoader(ComponentLoader): + """Loader for text encoders.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the text encoders based on the model path, architecture, and inference args.""" + # should always use v1 here. If inference_args.use_v1_text_encoder is False, + # the pipeline will overwrite this text encoder with a v0 text encoder + # during initialize_encoders() + return self.load_v1(model_path, architecture, inference_args) + + def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the text encoders based on the model path, architecture, and inference args.""" + model_config: PretrainedConfig = get_hf_config( + model=model_path, + trust_remote_code=inference_args.trust_remote_code, + revision=inference_args.revision, + model_override_args=None, + inference_args=inference_args, + ) + logger.info(f"HF Model config: {model_config}") + model_loader = get_model_loader() + model = model_loader.load_model(model_path, model_config, inference_args) + return model + + +class TokenizerLoader(ComponentLoader): + """Loader for tokenizers.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the tokenizer based on the model path, architecture, and inference args.""" + logger.info(f"Loading tokenizer from {model_path}") + + tokenizer = AutoTokenizer.from_pretrained( + model_path, + # TODO(will): pass these tokenizer kwargs from inference args? Maybe + # other method of config? + padding_size='right', + ) + logger.info(f"Loaded tokenizer: {tokenizer.__class__.__name__}") + return tokenizer + + +class VAELoader(ComponentLoader): + """Loader for VAE.""" + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the VAE based on the model path, architecture, and inference args.""" + use_v1 = inference_args.use_v1_vae + if not use_v1: + return self.load_v0(model_path, architecture, inference_args) + else: + return self.load_v1(model_path, architecture, inference_args) + + def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Custom VAE loading for Hunyuan""" + # TODO(will): replace this with abstracted model + from fastvideo.v1.v0_reference_src.models.hunyuan.vae import load_vae + vae, _, s_ratio, t_ratio = load_vae( + inference_args.vae, + inference_args.vae_precision, + logger=logger, + device=self.device + ) + vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} + vae.kwargs = vae_kwargs + return vae + + def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the VAE based on the model path, architecture, and inference args.""" + # TODO(will): move this to a constants file + from fastvideo.v1.utils import PRECISION_TO_TYPE + + config = get_diffusers_config(model=model_path) + + class_name = config.pop("_class_name") + assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported." + config.pop("_diffusers_version") + + vae_cls, _ = ModelRegistry.resolve_model_cls(class_name) + + vae = vae_cls(**config).to(inference_args.device) + + # Find all safetensors files + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) + # TODO(PY) + assert len(safetensors_list) == 1, f"Found {len(safetensors_list)} safetensors files in {path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict(loaded) + dtype = PRECISION_TO_TYPE[inference_args.vae_precision] + vae = vae.eval().to(dtype) + + # TODO(will): should we define hunyuan vae config class? + vae_kwargs = { + "s_ratio": config["spatial_compression_ratio"], + "t_ratio": config["temporal_compression_ratio"], + } + + vae.kwargs = vae_kwargs + + return vae + +class TransformerLoader(ComponentLoader): + """Loader for transformer.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the transformer based on the model path, architecture, and inference args.""" + use_v1 = inference_args.use_v1_transformer + if not use_v1: + return self.load_v0(model_path, architecture, inference_args) + else: + return self.load_v1(model_path, architecture, inference_args) + + def load_v0(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Custom transformer loading for Hunyuan""" + # TODO(will): replace this with abstracted model + from fastvideo.v1.v0_reference_src.models.hunyuan.modules import load_model + from fastvideo.v1.utils import PRECISION_TO_TYPE + # Disable gradient + torch.set_grad_enabled(False) + + # =========================== Build main model =========================== + logger.info("Building model...") + factor_kwargs = {"device": self.device, "dtype": PRECISION_TO_TYPE[inference_args.precision]} + in_channels = inference_args.latent_channels + out_channels = inference_args.latent_channels + + model = load_model( + inference_args, + in_channels=in_channels, + out_channels=out_channels, + factor_kwargs=factor_kwargs, + ) + model = model.to(self.device) + model = self._load_transformer_state_dict(inference_args, model, inference_args.model_path) + if inference_args.enable_torch_compile: + model = torch.compile(model) + model.eval() + return model + + def load_v1(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the transformer based on the model path, architecture, and inference args.""" + model_config = get_diffusers_config(model=model_path) + cls_name = model_config.pop("_class_name") + if cls_name is None: + raise ValueError(f"Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") + model_config.pop("_diffusers_version") + + model_cls, _ = ModelRegistry.resolve_model_cls(cls_name) + + # Find all safetensors files + safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors")) + if not safetensors_list: + raise ValueError(f"No safetensors files found in {model_path}") + + logger.info(f"Loading model from {len(safetensors_list)} safetensors files in {model_path}") + + # initialize_sequence_parallel_group(inference_args.sp_size) + + # Load the model using FSDP loader + logger.info(f"Loading model from {cls_name}") + model = load_fsdp_model( + model_cls=model_cls, + init_params=model_config, + weight_dir_list=safetensors_list, + device=inference_args.device, + cpu_offload=inference_args.use_cpu_offload + ) + + total_params = sum(p.numel() for p in model.parameters()) + logger.info(f"Loaded model with {total_params / 1e9:.2f}B parameters") + + model.eval() + return model + +class SchedulerLoader(ComponentLoader): + """Loader for scheduler.""" + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load the scheduler based on the model path, architecture, and inference args.""" + + scheduler = get_scheduler( + module_path=model_path, + architecture=architecture, + inference_args=inference_args, + ) + logger.info(f"Scheduler loaded: {scheduler}") + return scheduler + +class GenericComponentLoader(ComponentLoader): + """Generic loader for components that don't have a specific loader.""" + + def __init__(self, library="transformers"): + super().__init__() + self.library = library + + def load(self, model_path: str, architecture: str, inference_args: InferenceArgs): + """Load a generic component based on the model path, architecture, and inference args.""" + logger.warning(f"Using generic loader for {model_path} with library {self.library}") + + if self.library == "transformers": + from transformers import AutoModel + + model = AutoModel.from_pretrained( + model_path, + trust_remote_code=inference_args.trust_remote_code, + revision=inference_args.revision, + ) + logger.info(f"Loaded generic transformers model: {model.__class__.__name__}") + return model + elif self.library == "diffusers": + logger.warning(f"Generic loading for diffusers components is not fully implemented") + from fastvideo.v1.models.hf_transformer_utils import get_diffusers_config + + model_config = get_diffusers_config(model=model_path) + logger.info(f"Diffusers Model config: {model_config}") + # This is a placeholder - in a real implementation, you'd need to handle this properly + return None + else: + raise ValueError(f"Unsupported library: {self.library}") + +class PipelineComponentLoader: + """ + Utility class for loading pipeline components. + This replaces the chain of if-else statements in load_pipeline_module. + """ + + @staticmethod + def load_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, + architecture: str, inference_args: InferenceArgs): + """ + Load a pipeline module. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + architecture: Architecture of the component model + inference_args: Inference arguments + + Returns: + The loaded module + """ + logger.info(f"Loading {module_name} using {transformers_or_diffusers} from {component_model_path}") + + # Get the appropriate loader for this module type + loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) + + # Load the module + return loader.load(component_model_path, architecture, inference_args) diff --git a/fastvideo/models/parameter.py b/fastvideo/v1/models/parameter.py similarity index 98% rename from fastvideo/models/parameter.py rename to fastvideo/v1/models/parameter.py index 30e901e98..11fc6f990 100644 --- a/fastvideo/models/parameter.py +++ b/fastvideo/v1/models/parameter.py @@ -7,9 +7,9 @@ from typing import Callable, Optional, Union import torch from torch.nn import Parameter -from fastvideo.distributed import get_tensor_model_parallel_rank -from fastvideo.logger import init_logger -from fastvideo.models.utils import _make_synced_weight_loader +from fastvideo.v1.distributed import get_tensor_model_parallel_rank +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.utils import _make_synced_weight_loader __all__ = [ "BasevLLMParameter", "PackedvLLMParameter", "PerTensorScaleParameter", @@ -431,4 +431,4 @@ def _adjust_shard_indexes_for_packing(shard_size, shard_offset, packed_factor, shard_size=shard_size, shard_offset=shard_offset, marlin_tile_size=marlin_tile_size) - return shard_size, shard_offset + return shard_size, shard_offset \ No newline at end of file diff --git a/fastvideo/models/dits/registry.py b/fastvideo/v1/models/registry.py similarity index 86% rename from fastvideo/models/dits/registry.py rename to fastvideo/v1/models/registry.py index 8a1a2b248..b665a6ccc 100644 --- a/fastvideo/models/dits/registry.py +++ b/fastvideo/v1/models/registry.py @@ -11,24 +11,40 @@ import importlib from functools import lru_cache import cloudpickle from torch import nn -from fastvideo.logger import logger +from fastvideo.v1.logger import logger -# huggingface class name: (fastvideo module name, fastvideo class name) -_TEXT_TO_VIDEO_MODELS = { - "HunyuanVideoTransformer3DModel": ("hunyuanvideo", "HunyuanVideoDiT"), +# huggingface class name: (component_name, fastvideo module name, fastvideo class name) +_TEXT_TO_VIDEO_DIT_MODELS = { + "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), } -_IMAGE_TO_VIDEO_MODELS = { - +_IMAGE_TO_VIDEO_DIT_MODELS = { + # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), +} + +_TEXT_ENCODER_MODELS = { + "CLIPTextModel": ("encoders", "clip", "CLIPTextModel"), + "LlamaModel": ("encoders", "llama", "LlamaModel"), +} + +_IMAGE_ENCODER_MODELS = { + # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), +} + +_VAE_MODELS = { + "AutoencoderKLHunyuanVideo": ("vaes", "hunyuanvae", "AutoencoderKLHunyuanVideo"), } _FAST_VIDEO_MODELS = { - **_TEXT_TO_VIDEO_MODELS, - **_IMAGE_TO_VIDEO_MODELS, + **_TEXT_TO_VIDEO_DIT_MODELS, + **_IMAGE_TO_VIDEO_DIT_MODELS, + **_TEXT_ENCODER_MODELS, + **_IMAGE_ENCODER_MODELS, + **_VAE_MODELS, } _SUBPROCESS_COMMAND = [ - sys.executable, "-m", "fastvideo.models.dits.registry" + sys.executable, "-m", "fastvideo.v1.models.dits.registry" ] @@ -111,6 +127,7 @@ class _LazyRegisteredModel(_BaseRegisteredModel): Represents a model that has not been imported in the main process. """ module_name: str + component_name: str class_name: str # Performed in another process to avoid initializing CUDA @@ -128,7 +145,7 @@ def _try_load_model_cls( model_arch: str, model: _BaseRegisteredModel, ) -> Optional[Type[nn.Module]]: - from vllm.platforms import current_platform + from fastvideo.v1.platforms import current_platform current_platform.verify_model_arch(model_arch) try: return model.load_model_cls() @@ -266,8 +283,9 @@ class _ModelRegistry: ModelRegistry = _ModelRegistry({ model_arch: _LazyRegisteredModel( - module_name=f"fastvideo.models.dits.{mod_relname}", + module_name=f"fastvideo.v1.models.{component_name}.{mod_relname}", + component_name=component_name, class_name=cls_name, ) - for model_arch, (mod_relname, cls_name) in _FAST_VIDEO_MODELS.items() + for model_arch, (component_name, mod_relname, cls_name) in _FAST_VIDEO_MODELS.items() }) \ No newline at end of file diff --git a/fastvideo/v1/models/schedulers/base.py b/fastvideo/v1/models/schedulers/base.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py new file mode 100644 index 000000000..e9e96edb1 --- /dev/null +++ b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -0,0 +1,239 @@ +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, logging + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + reverse: bool = True, + solver: str = "euler", + n_tokens: Optional[int] = None, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32) + + self._step_index = None + self._begin_index = None + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}") + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device] = None, + n_tokens: int = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + sigmas = self.sd3_time_shift(sigmas) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device) + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + return sample + + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor)): + raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.to(torch.float32) * dt + else: + raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}") + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample, ) + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps \ No newline at end of file diff --git a/fastvideo/v1/models/text_encoder.py b/fastvideo/v1/models/text_encoder.py new file mode 100644 index 000000000..8268d224b --- /dev/null +++ b/fastvideo/v1/models/text_encoder.py @@ -0,0 +1,255 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn + +from transformers.utils import ModelOutput +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) + +def use_default(value, default): + return value if value is not None else default + +@dataclass +class TextEncoderModelOutput(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + text_outputs (`list`, *optional*, returned when `return_texts=True` is passed): + List of decoded texts. + """ + + hidden_state: torch.FloatTensor = None + attention_mask: Optional[torch.LongTensor] = None + hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None + text_outputs: Optional[list] = None + + +class TextEncoder(nn.Module): + + def __init__( + self, + text_encoder, + tokenizer, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + prompt_template: Optional[dict] = None, + prompt_template_video: Optional[dict] = None, + hidden_state_skip_layer: Optional[int] = None, + apply_final_norm: bool = False, + reproduce: bool = False, + device=None, + ): + super().__init__() + # TODO(will): check if there's a cleaner way to do this + self.text_encoder_type = text_encoder.config.architectures[0] + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.use_attention_mask = use_attention_mask + if prompt_template_video is not None: + assert (use_attention_mask is True), "Attention mask is True required when training videos." + self.prompt_template = prompt_template + self.prompt_template_video = prompt_template_video + self.hidden_state_skip_layer = hidden_state_skip_layer + self.apply_final_norm = apply_final_norm + self.reproduce = reproduce + + self.use_template = self.prompt_template is not None + if self.use_template: + assert (isinstance(self.prompt_template, dict) and "template" in self.prompt_template + ), f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}" + assert "{}" in str(self.prompt_template["template"]), ( + "`prompt_template['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template['template']}") + + self.use_video_template = self.prompt_template_video is not None + if self.use_video_template: + if self.prompt_template_video is not None: + assert ( + isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video + ), f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}" + assert "{}" in str(self.prompt_template_video["template"]), ( + "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, " + f"got {self.prompt_template_video['template']}") + + if "T5" in self.text_encoder_type: + self.output_key = output_key or "last_hidden_state" + elif "CLIPTextModel" in self.text_encoder_type: + self.output_key = output_key or "pooler_output" + elif "LlamaModel" in self.text_encoder_type or "glm" in self.text_encoder_type: + self.output_key = output_key or "last_hidden_state" + else: + raise ValueError(f"Unsupported text encoder type: {self.text_encoder_type}") + + self.model = text_encoder + # self.dtype = self.model.dtype + self.device = device + + self.tokenizer = tokenizer + + def __repr__(self): + return f"{self.text_encoder_type} ({self.precision} - {self.model_path})" + + @staticmethod + def apply_text_to_template(text, template, prevent_empty_text=True): + """ + Apply text to template. + + Args: + text (str): Input text. + template (str or list): Template string or list of chat conversation. + prevent_empty_text (bool): If True, we will prevent the user text from being empty + by adding a space. Defaults to True. + """ + if isinstance(template, str): + # Will send string to tokenizer. Used for llm + return template.format(text) + else: + raise TypeError(f"Unsupported template type: {type(template)}") + + def text2tokens(self, text, data_type="image"): + """ + Tokenize the input text. + + Args: + text (str or list): Input text. + """ + tokenize_input_type = "str" + if self.use_template: + if data_type == "image": + prompt_template = self.prompt_template["template"] + elif data_type == "video": + prompt_template = self.prompt_template_video["template"] + else: + raise ValueError(f"Unsupported data type: {data_type}") + if isinstance(text, (list, tuple)): + text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text] + if isinstance(text[0], list): + tokenize_input_type = "list" + elif isinstance(text, str): + text = self.apply_text_to_template(text, prompt_template) + if isinstance(text, list): + tokenize_input_type = "list" + else: + raise TypeError(f"Unsupported text type: {type(text)}") + + kwargs = dict( + truncation=True, + max_length=self.max_length, + padding="max_length", + return_tensors="pt", + ) + if tokenize_input_type == "str": + return self.tokenizer( + text, + return_length=False, + return_overflowing_tokens=False, + return_attention_mask=True, + **kwargs, + ) + elif tokenize_input_type == "list": + return self.tokenizer.apply_chat_template( + text, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + **kwargs, + ) + else: + raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}") + + def encode( + self, + batch_encoding, + use_attention_mask=None, + output_hidden_states=False, + do_sample=None, + hidden_state_skip_layer=None, + return_texts=False, + data_type="image", + device=None, + ): + """ + Args: + batch_encoding (dict): Batch encoding from tokenizer. + use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask. + Defaults to None. + output_hidden_states (bool): Whether to output hidden states. If False, return the value of + self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer, + output_hidden_states will be set True. Defaults to False. + do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None. + When self.produce is False, do_sample is set to True by default. + hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer. + If None, self.output_key will be used. Defaults to None. + return_texts (bool): Whether to return the decoded texts. Defaults to False. + """ + device = self.model.device if device is None else device + use_attention_mask = use_default(use_attention_mask, self.use_attention_mask) + hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer) + do_sample = use_default(do_sample, not self.reproduce) + attention_mask = (batch_encoding["attention_mask"].to(device) if use_attention_mask else None) + # print(f"batch_encoding: {batch_encoding}") + # logger.info(f"attention_mask: {attention_mask}") + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(device), + # attention_mask=attention_mask, + output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None, + ) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm(last_hidden_state) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_template: + if data_type == "image": + crop_start = self.prompt_template.get("crop_start", -1) + elif data_type == "video": + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = (attention_mask[:, crop_start:] if use_attention_mask else None) + + if output_hidden_states: + return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + def forward( + self, + text, + use_attention_mask=None, + output_hidden_states=False, + do_sample=False, + hidden_state_skip_layer=None, + return_texts=False, + ): + batch_encoding = self.text2tokens(text) + return self.encode( + batch_encoding, + use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, + do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, + return_texts=return_texts, + ) diff --git a/fastvideo/models/utils.py b/fastvideo/v1/models/utils.py similarity index 55% rename from fastvideo/models/utils.py rename to fastvideo/v1/models/utils.py index 31d230b2e..75ea7818b 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/v1/models/utils.py @@ -5,9 +5,52 @@ from typing import Any, Dict, Optional import torch +# TODO(PY): move it elsewhere +def auto_attributes(init_func): + """ + Decorator that automatically adds all initialization arguments as object attributes. + + Example: + @auto_attributes + def __init__(self, a=1, b=2): + pass + + # This will automatically set: + # - self.a = 1 and self.b = 2 + # - self.config.a = 1 and self.config.b = 2 + """ + def wrapper(self, *args, **kwargs): + # Get the function signature + import inspect + signature = inspect.signature(init_func) + parameters = signature.parameters + + # Get parameter names (excluding 'self') + param_names = list(parameters.keys())[1:] + + # Bind arguments to parameters + bound_args = signature.bind(self, *args, **kwargs) + bound_args.apply_defaults() + + # Create config object if it doesn't exist + if not hasattr(self, 'config'): + self.config = type('Config', (), {})() + + # Set attributes on self and self.config + for name in param_names: + if name in bound_args.arguments: + value = bound_args.arguments[name] + setattr(self, name, value) + setattr(self.config, name, value) + + # Call the original __init__ function + return init_func(self, *args, **kwargs) + + return wrapper + def set_random_seed(seed: int) -> None: - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform current_platform.seed_everything(seed) @@ -39,7 +82,7 @@ def set_weight_attrs( # This sometimes causes OOM errors during model loading. To avoid this, # we sync the param tensor after its weight loader is called. # TODO(woosuk): Remove this hack once we have a better solution. - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform if current_platform.is_tpu() and key == "weight_loader": value = _make_synced_weight_loader(value) setattr(weight, key, value) diff --git a/fastvideo/v1/models/vaes/common.py b/fastvideo/v1/models/vaes/common.py new file mode 100644 index 000000000..f3ecec89f --- /dev/null +++ b/fastvideo/v1/models/vaes/common.py @@ -0,0 +1,469 @@ +import torch +from typing import Optional, Tuple +import numpy as np +import torch.nn as nn +from abc import abstractmethod, ABC +import torch.distributed as dist +from fastvideo.v1.distributed import tensor_model_parallel_all_gather, get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size +from math import prod +class ParallelTiledVAE(ABC): + def __init__(self, *args, **kwargs): + # Check if subclass has defined all required properties + required_attributes = [ + 'tile_sample_min_height', + 'tile_sample_min_width', + 'tile_sample_min_num_frames', + 'tile_sample_stride_height', + 'tile_sample_stride_width', + 'tile_sample_stride_num_frames', + 'use_tiling' + ] + + for attr in required_attributes: + if not hasattr(self, attr): + raise AttributeError(f"Subclasses of ParallelVAE must define '{attr}' property") + + @abstractmethod + def _encode(self, *args, **kwargs): + pass + + @abstractmethod + def _decode(self, *args, **kwargs): + pass + + + def encode(self, x: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + + if self.use_tiling and num_frames > self.tile_sample_min_num_frames: + latents = self.tiled_encode(x) + elif self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height): + latents = self.spatial_tiled_encode(x) + else: + latents = self._encode(x) + return DiagonalGaussianDistribution(latents) + + + + def decode(self, z: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + + if self.use_tiling and get_sequence_model_parallel_world_size() > 1: + return self.parallel_tiled_decode(z) + if self.use_tiling and num_frames > tile_latent_min_num_frames: + return self.tiled_decode(z) + + if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height): + return self.spatial_tiled_decode(z) + + return self._decode(z) + + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def spatial_tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + r"""Encode a batch of images using a tiled encoder. + + Args: + x (`torch.Tensor`): Input batch of videos. + + Returns: + `torch.Tensor`: + The latent representation of the encoded videos. + """ + batch_size, num_channels, num_frames, height, width = x.shape + latent_height = height // self.spatial_compression_ratio + latent_width = width // self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = tile_latent_min_height - tile_latent_stride_height + blend_width = tile_latent_min_width - tile_latent_stride_width + + # Split x into overlapping tiles and encode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, self.tile_sample_stride_height): + row = [] + for j in range(0, width, self.tile_sample_stride_width): + tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width] + tile = self._encode(tile) + row.append(tile) + rows.append(row) + + return self._merge_spatial_tiles(rows, blend_height, blend_width,tile_latent_stride_height, tile_latent_stride_width) + + def _parallel_data_generator(self, gathered_results, + gathered_dim_metadata): + global_idx = 0 + for i, per_rank_metadata in enumerate(gathered_dim_metadata): + _start_shape = 0 + for shape in per_rank_metadata: + mul_shape = prod(shape) + yield (gathered_results[i, _start_shape:_start_shape + + mul_shape].reshape(shape), global_idx) + _start_shape += mul_shape + global_idx += 1 + + def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor: + """ + Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs + """ + world_size, rank = get_sequence_model_parallel_world_size(), get_sequence_model_parallel_rank() + B, C, T, H, W = z.shape + + # Calculate parameters + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + # Calculate tile dimensions + num_t_tiles = (T + tile_latent_stride_num_frames - 1) // tile_latent_stride_num_frames + num_h_tiles = (H + tile_latent_stride_height - 1) // tile_latent_stride_height + num_w_tiles = (W + tile_latent_stride_width - 1) // tile_latent_stride_width + total_spatial_tiles = num_h_tiles * num_w_tiles + total_tiles = num_t_tiles * total_spatial_tiles + + # Calculate tiles per rank and padding + tiles_per_rank = (total_tiles + world_size - 1) // world_size + start_tile_idx = rank * tiles_per_rank + end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) + + local_results = [] + local_dim_metadata = [] + # Process assigned tiles + for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + + # Calculate positions + t_start = t_idx * tile_latent_stride_num_frames + h_start = h_idx * tile_latent_stride_height + w_start = w_idx * tile_latent_stride_width + + # Extract and process tile + tile = z[:, :, t_start:t_start + tile_latent_min_num_frames + 1, + h_start:h_start + tile_latent_min_height, + w_start:w_start + tile_latent_min_width] + + # Process tile + tile = self._decode(tile) + + if t_start > 0: + tile = tile[:, :, 1:, :, :] + + # Store metadata + shape = tile.shape + # Store decoded data (flattened) + decoded_flat = tile.reshape(-1) + local_results.append(decoded_flat) + local_dim_metadata.append(shape) + + results = torch.cat(local_results, dim=0).contiguous() + del local_results + torch.cuda.empty_cache() + # first gather size to pad the results + local_size = torch.tensor([results.size(0)], + device=results.device, + dtype=torch.int64) + all_sizes = [ + torch.zeros(1, device=results.device, dtype=torch.int64) + for _ in range(world_size) + ] + dist.all_gather(all_sizes, local_size) + max_size = max(size.item() for size in all_sizes) + padded_results = torch.zeros(max_size, device=results.device) + padded_results[:results.size(0)] = results + del results + torch.cuda.empty_cache() + # Gather all results + gathered_dim_metadata = [None] * world_size + gathered_results = torch.zeros_like(padded_results).repeat( + world_size, *[1] * len(padded_results.shape) + ).contiguous() # use contiguous to make sure it won't copy data in the following operations + # TODO (PY): use fastvideo distributed methods + dist.all_gather_into_tensor(gathered_results, padded_results) + dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) + # Process gathered results + data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] + for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator( + gathered_results, gathered_dim_metadata): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + data[t_idx][h_idx][w_idx] = current_data + # Merge results + result_slices = [] + last_slice_data = None + for i, tem_data in enumerate(data): + slice_data = self._merge_spatial_tiles(tem_data, blend_height, blend_width, + self.tile_sample_stride_height, self.tile_sample_stride_width) + if i > 0: + slice_data = self.blend_t(last_slice_data, slice_data, blend_num_frames) + result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames, :, :]) + else: + result_slices.append(slice_data[:, :, :self.tile_sample_stride_num_frames + 1, :, :]) + last_slice_data = slice_data + dec = torch.cat(result_slices, dim=2) + + return dec + + def _merge_spatial_tiles(self, tiles, blend_height, blend_width, stride_height, stride_width): + """Helper function to merge spatial tiles with blending""" + result_rows = [] + for i, row in enumerate(tiles): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(tiles[i - 1][j], tile, + blend_height) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_width) + result_row.append(tile[:, :, :, :stride_height, :stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def spatial_tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + r""" + Decode a batch of images using a tiled decoder. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + + Returns: + `torch.Tensor`: + The decoded images. + """ + + batch_size, num_channels, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, height, tile_latent_stride_height): + row = [] + for j in range(0, width, tile_latent_stride_width): + tile = z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + decoded = self._decode(tile) + row.append(decoded) + rows.append(row) + return self._merge_spatial_tiles(rows, blend_height, blend_width,self.tile_sample_stride_height, self.tile_sample_stride_width) + + def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = x.shape + latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1 + + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames + + row = [] + for i in range(0, num_frames, self.tile_sample_stride_num_frames): + tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :] + if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width): + tile = self.spatial_tiled_encode(tile) + else: + tile = self._encode(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :]) + else: + result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :]) + + enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames] + return enc + + def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = z.shape + num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1 + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio + tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio + blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames + + row = [] + for i in range(0, num_frames, tile_latent_stride_num_frames): + tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :] + if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self._decode(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + result_row = [] + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_num_frames) + result_row.append(tile[:, :, : self.tile_sample_stride_num_frames, :, :]) + else: + result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]) + + dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames] + return dec + + + + def enable_tiling( + self, + tile_sample_min_height: Optional[int] = None, + tile_sample_min_width: Optional[int] = None, + tile_sample_min_num_frames: Optional[int] = None, + tile_sample_stride_height: Optional[float] = None, + tile_sample_stride_width: Optional[float] = None, + tile_sample_stride_num_frames: Optional[float] = None, + ) -> None: + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + + Args: + tile_sample_min_height (`int`, *optional*): + The minimum height required for a sample to be separated into tiles across the height dimension. + tile_sample_min_width (`int`, *optional*): + The minimum width required for a sample to be separated into tiles across the width dimension. + tile_sample_min_num_frames (`int`, *optional*): + The minimum number of frames required for a sample to be separated into tiles across the frame + dimension. + tile_sample_stride_height (`int`, *optional*): + The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are + no tiling artifacts produced across the height dimension. + tile_sample_stride_width (`int`, *optional*): + The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling + artifacts produced across the width dimension. + tile_sample_stride_num_frames (`int`, *optional*): + The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts + produced across the frame dimension. + """ + self.use_tiling = True + self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height + self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width + self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames + self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height + self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width + self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames + + + + def disable_tiling(self) -> None: + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_tiling = False + + + +# adapted from https://github.com/huggingface/diffusers/blob/e7ffeae0a191f710881d1fbde00cd6ff025e81f2/src/diffusers/models/autoencoders/vae.py#L691 +class DiagonalGaussianDistribution(object): + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + + def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = torch.randn( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=[1, 2, 3], + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=[1, 2, 3], + ) + + def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean \ No newline at end of file diff --git a/fastvideo/v1/models/vaes/hunyuanvae.py b/fastvideo/v1/models/vaes/hunyuanvae.py new file mode 100644 index 000000000..32df14429 --- /dev/null +++ b/fastvideo/v1/models/vaes/hunyuanvae.py @@ -0,0 +1,799 @@ +# Copyright 2024 The Hunyuan Team, The HuggingFace Team and The FastVideo Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint + +from fastvideo.v1.models.vaes.common import DiagonalGaussianDistribution, ParallelTiledVAE +from fastvideo.v1.layers.activation import get_act_fn + +from fastvideo.v1.models.utils import auto_attributes + +def prepare_causal_attention_mask( + num_frames: int, height_width: int, dtype: torch.dtype, device: torch.device, batch_size: int = None +) -> torch.Tensor: + indices = torch.arange(1, num_frames + 1, dtype=torch.int32, device=device) + indices_blocks = indices.repeat_interleave(height_width) + x, y = torch.meshgrid(indices_blocks, indices_blocks, indexing="xy") + mask = torch.where(x <= y, 0, -float("inf")).to(dtype=dtype) + + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + +class HunyuanVAEAttention(nn.Module): + def __init__(self, in_channels, heads, dim_head, eps, norm_num_groups, bias): + super().__init__() + self.in_channels = in_channels + self.heads = heads + self.dim_head = dim_head + self.eps = eps + self.norm_num_groups = norm_num_groups + self.bias = bias + + inner_dim = heads * dim_head + + # Define the projection layers + self.to_q = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_k = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_v = nn.Linear(in_channels, inner_dim, bias=bias) + self.to_out = nn.Sequential( + nn.Linear(inner_dim, in_channels, bias=bias) + ) + + # Optional normalization layers + self.group_norm = nn.GroupNorm(norm_num_groups, in_channels, eps=eps, affine=True) + + def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + residual = hidden_states + + batch_size, sequence_length, _ = hidden_states.shape + + hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + # Project to query, key, value + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + # Reshape for multi-head attention + head_dim = self.dim_head + + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + # Perform scaled dot-product attention + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + # Reshape back + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # Linear projection + hidden_states = self.to_out(hidden_states) + + # Residual connection and rescale + hidden_states = hidden_states + residual + + return hidden_states + +class HunyuanVideoCausalConv3d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int, int]] = 3, + stride: Union[int, Tuple[int, int, int]] = 1, + padding: Union[int, Tuple[int, int, int]] = 0, + dilation: Union[int, Tuple[int, int, int]] = 1, + bias: bool = True, + pad_mode: str = "replicate", + ) -> None: + super().__init__() + + kernel_size = (kernel_size, kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size + + self.pad_mode = pad_mode + self.time_causal_padding = ( + kernel_size[0] // 2, + kernel_size[0] // 2, + kernel_size[1] // 2, + kernel_size[1] // 2, + kernel_size[2] - 1, + 0, + ) + + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = F.pad(hidden_states, self.time_causal_padding, mode=self.pad_mode) + return self.conv(hidden_states) + + +class HunyuanVideoUpsampleCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + kernel_size: int = 3, + stride: int = 1, + bias: bool = True, + upsample_factor: Tuple[float, float, float] = (2, 2, 2), + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + self.upsample_factor = upsample_factor + + self.conv = HunyuanVideoCausalConv3d(in_channels, out_channels, kernel_size, stride, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_frames = hidden_states.size(2) + + first_frame, other_frames = hidden_states.split((1, num_frames - 1), dim=2) + first_frame = F.interpolate( + first_frame.squeeze(2), scale_factor=self.upsample_factor[1:], mode="nearest" + ).unsqueeze(2) + + if num_frames > 1: + # See: https://github.com/pytorch/pytorch/issues/81665 + # Unless you have a version of pytorch where non-contiguous implementation of F.interpolate + # is fixed, this will raise either a runtime error, or fail silently with bad outputs. + # If you are encountering an error here, make sure to try running encoding/decoding with + # `vae.enable_tiling()` first. If that doesn't work, open an issue at: + # https://github.com/huggingface/diffusers/issues + other_frames = other_frames.contiguous() + other_frames = F.interpolate(other_frames, scale_factor=self.upsample_factor, mode="nearest") + hidden_states = torch.cat((first_frame, other_frames), dim=2) + else: + hidden_states = first_frame + + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoDownsampleCausal3D(nn.Module): + def __init__( + self, + channels: int, + out_channels: Optional[int] = None, + padding: int = 1, + kernel_size: int = 3, + bias: bool = True, + stride=2, + ) -> None: + super().__init__() + out_channels = out_channels or channels + + self.conv = HunyuanVideoCausalConv3d(channels, out_channels, kernel_size, stride, padding, bias=bias) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv(hidden_states) + return hidden_states + + +class HunyuanVideoResnetBlockCausal3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: Optional[int] = None, + dropout: float = 0.0, + groups: int = 32, + eps: float = 1e-6, + non_linearity: str = "silu", + ) -> None: + super().__init__() + out_channels = out_channels or in_channels + + self.nonlinearity = get_act_fn(non_linearity) + + self.norm1 = nn.GroupNorm(groups, in_channels, eps=eps, affine=True) + self.conv1 = HunyuanVideoCausalConv3d(in_channels, out_channels, 3, 1, 0) + + self.norm2 = nn.GroupNorm(groups, out_channels, eps=eps, affine=True) + self.dropout = nn.Dropout(dropout) + self.conv2 = HunyuanVideoCausalConv3d(out_channels, out_channels, 3, 1, 0) + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = HunyuanVideoCausalConv3d(in_channels, out_channels, 1, 1, 0) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.contiguous() + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + hidden_states = hidden_states + residual + return hidden_states + + +class HunyuanVideoMidBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_attention: bool = True, + attention_head_dim: int = 1, + ) -> None: + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + # There is always at least one resnet + resnets = [ + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ] + attentions = [] + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + HunyuanVAEAttention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + eps=resnet_eps, + norm_num_groups=resnet_groups, + bias=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(self.resnets[0], hidden_states) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask( + num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size + ) + hidden_states = attn(hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) + + else: + hidden_states = self.resnets[0](hidden_states) + + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.permute(0, 2, 3, 4, 1).flatten(1, 3) + attention_mask = prepare_causal_attention_mask( + num_frames, height * width, hidden_states.dtype, hidden_states.device, batch_size=batch_size + ) + hidden_states = attn(hidden_states, attention_mask=attention_mask) + hidden_states = hidden_states.unflatten(1, (num_frames, height, width)).permute(0, 4, 1, 2, 3) + + hidden_states = resnet(hidden_states) + + return hidden_states + + +class HunyuanVideoDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + HunyuanVideoDownsampleCausal3D( + out_channels, + out_channels=out_channels, + padding=downsample_padding, + stride=downsample_stride, + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + for resnet in self.resnets: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) + else: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "silu", + resnet_groups: int = 32, + add_upsample: bool = True, + upsample_scale_factor: Tuple[int, int, int] = (2, 2, 2), + ) -> None: + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + HunyuanVideoResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + non_linearity=resnet_act_fn, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + HunyuanVideoUpsampleCausal3D( + out_channels, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and self.gradient_checkpointing: + for resnet in self.resnets: + hidden_states = self._gradient_checkpointing_func(resnet, hidden_states) + + else: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states + + +class HunyuanVideoEncoder3D(nn.Module): + r""" + Causal encoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + temporal_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ) -> None: + super().__init__() + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + if down_block_type != "HunyuanVideoDownBlock3D": + raise ValueError(f"Unsupported down_block_type: {down_block_type}") + + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(temporal_compression_ratio)) + + if temporal_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool( + i >= (len(block_out_channels) - 1 - num_time_downsample_layers) and not is_final_block + ) + elif temporal_compression_ratio == 8: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i < num_time_downsample_layers) + else: + raise ValueError(f"Unsupported time_compression_ratio: {temporal_compression_ratio}") + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2,) if add_time_downsample else (1,) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) + + down_block = HunyuanVideoDownBlock3D( + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample or add_time_downsample), + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + downsample_stride=downsample_stride, + downsample_padding=0, + ) + + self.down_blocks.append(down_block) + + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + for down_block in self.down_blocks: + hidden_states = self._gradient_checkpointing_func(down_block, hidden_states) + + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) + else: + for down_block in self.down_blocks: + hidden_states = down_block(hidden_states) + + hidden_states = self.mid_block(hidden_states) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class HunyuanVideoDecoder3D(nn.Module): + r""" + Causal decoder for 3D video-like data introduced in [Hunyuan Video](https://huggingface.co/papers/2412.03603). + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: Tuple[int, ...] = (128, 256, 512, 512), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = HunyuanVideoCausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = HunyuanVideoMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + if up_block_type != "HunyuanVideoUpBlock3D": + raise ValueError(f"Unsupported up_block_type: {up_block_type}") + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool( + i >= len(block_out_channels) - 1 - num_time_upsample_layers and not is_final_block + ) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}") + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) + upsample_scale_factor_T = (2,) if add_time_upsample else (1,) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) + + up_block = HunyuanVideoUpBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + ) + + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = HunyuanVideoCausalConv3d(block_out_channels[0], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states) + + for up_block in self.up_blocks: + hidden_states = self._gradient_checkpointing_func(up_block, hidden_states) + else: + hidden_states = self.mid_block(hidden_states) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states) + + # post-process + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return hidden_states + + +class AutoencoderKLHunyuanVideo(nn.Module, ParallelTiledVAE): + r""" + A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. + Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = True + @auto_attributes + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 16, + down_block_types: Tuple[str, ...] = ( + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ), + up_block_types: Tuple[str, ...] = ( + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ), + block_out_channels: Tuple[int] = (128, 256, 512, 512), + layers_per_block: int = 2, + act_fn: str = "silu", + norm_num_groups: int = 32, + scaling_factor: float = 0.476986, + spatial_compression_ratio: int = 8, + temporal_compression_ratio: int = 4, + mid_block_add_attention: bool = True, + load_encoder: bool = True, + load_decoder: bool = True, + ) -> None: + super().__init__() + + # TODO(will): only pass in config. We do this by manually defining a + # config for hunyuan vae + self.block_out_channels = block_out_channels + + if load_encoder: + self.encoder = HunyuanVideoEncoder3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + double_z=True, + mid_block_add_attention=mid_block_add_attention, + temporal_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + ) + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + + if load_decoder: + self.decoder = HunyuanVideoDecoder3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=temporal_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + + + # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent + # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the + # intermediate tiles together, the memory requirement can be lowered. + self.use_tiling = True + + # The minimal tile height and width for spatial tiling to be used + self.tile_sample_min_height = 256 + self.tile_sample_min_width = 256 + self.tile_sample_min_num_frames = 16 + + # The minimal distance between two spatial tiles + self.tile_sample_stride_height = 192 + self.tile_sample_stride_width = 192 + self.tile_sample_stride_num_frames = 12 + + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + x = self.encoder(x) + enc = self.quant_conv(x) + return enc + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z) + return dec + \ No newline at end of file diff --git a/fastvideo/v1/pipelines/__init__.py b/fastvideo/v1/pipelines/__init__.py new file mode 100644 index 000000000..e100e85d1 --- /dev/null +++ b/fastvideo/v1/pipelines/__init__.py @@ -0,0 +1,244 @@ +""" +Diffusion pipelines for fastvideo.v1. + +This package contains diffusion pipelines for generating videos and images. +""" +import os +import json +from copy import deepcopy + +from typing import Dict, Optional, Type, Any + +from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from huggingface_hub import snapshot_download +from transformers import PretrainedConfig +from fastvideo.v1.models.hf_transformer_utils import get_hf_config, get_diffusers_config +from fastvideo.v1.models import get_scheduler +import glob +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.loader import PipelineComponentLoader + +logger = init_logger(__name__) + +# Then import the base classes +from fastvideo.v1.pipelines.composed.composed_pipeline_base import ( + ComposedPipelineBase, + DiffusionPipelineOutput +) + +def get_pipeline_type(inference_args: InferenceArgs) -> str: + # hardcode for now + return "hunyuan_video" + +def maybe_download_model(model_path: str) -> str: + """ + Check if the model path is a Hugging Face Hub model ID and download it if needed. + + Args: + model_path: Local path or Hugging Face Hub model ID + + Returns: + Local path to the model + """ + + # If the path exists locally, return it + if os.path.exists(model_path): + logger.info(f"Model already exists locally at {model_path}") + return model_path + + # Otherwise, assume it's a HF Hub model ID and try to download it + try: + logger.info(f"Downloading model snapshot from HF Hub for {model_path}...") + local_path = snapshot_download( + repo_id=model_path, + # allow_patterns=["*.json", "*.bin", "*.safetensors", "*.pt", "*.pth", "*.ckpt"], + ignore_patterns=["*.onnx", "*.msgpack"], + ) + logger.info(f"Downloaded model to {local_path}") + return local_path + except Exception as e: + raise ValueError(f"Could not find model at {model_path} and failed to download from HF Hub: {e}") + +def verify_model_config_and_directory(model_path: str) -> dict: + """ + Verify that the model directory contains a valid diffusers configuration. + + Args: + model_path: Path to the model directory + + Returns: + The loaded model configuration as a dictionary + """ + + # Check for model_index.json which is required for diffusers models + config_path = os.path.join(model_path, "model_index.json") + if not os.path.exists(config_path): + raise ValueError( + f"Model directory {model_path} does not contain model_index.json. " + "Only Hugging Face diffusers format is supported." + ) + + # Check for transformer and vae directories + transformer_dir = os.path.join(model_path, "transformer") + vae_dir = os.path.join(model_path, "vae") + + if not os.path.exists(transformer_dir): + raise ValueError(f"Model directory {model_path} does not contain a transformer/ directory.") + + if not os.path.exists(vae_dir): + raise ValueError(f"Model directory {model_path} does not contain a vae/ directory.") + + # Load the config + try: + with open(config_path, "r") as f: + config = json.load(f) + except Exception as e: + raise ValueError(f"Failed to load model configuration from {config_path}: {e}") + + # Verify diffusers version exists + if "_diffusers_version" not in config: + raise ValueError(f"model_index.json does not contain _diffusers_version") + + logger.info(f"Diffusers version: {config['_diffusers_version']}") + return config + +def load_pipeline_module(module_name: str, component_model_path: str, transformers_or_diffusers: str, architecture: str, inference_args: InferenceArgs) -> Any: + """ + Load a pipeline module using the appropriate loader. + + Args: + module_name: Name of the module (e.g., "vae", "text_encoder", "transformer", "scheduler") + component_model_path: Path to the component model + transformers_or_diffusers: Whether the module is from transformers or diffusers + architecture: Architecture of the component model + inference_args: Inference arguments + + Returns: + The loaded module + """ + return PipelineComponentLoader.load_module( + module_name=module_name, + component_model_path=component_model_path, + transformers_or_diffusers=transformers_or_diffusers, + architecture=architecture, + inference_args=inference_args + ) + +def load_pipeline_modules(model_path: str, config: Dict, inference_args: InferenceArgs) -> dict[str, Any]: + """ + Load the pipeline modules from the config. + + Args: + config: The model_index.json config + inference_args: Inference arguments + + Returns: + Dictionary mapping module names to loaded modules + """ + logger.info(f"Loading pipeline modules from config: {config}") + modules_config = deepcopy(config) + + # remove keys that are not pipeline modules + modules_config.pop("_class_name") + modules_config.pop("_diffusers_version") + + # some sanity checks + assert len(modules_config) > 1, "model_index.json must contain at least one pipeline module" + + required_modules = ["vae", "text_encoder", "transformer", "scheduler", "tokenizer"] + for module_name in required_modules: + if module_name not in modules_config: + raise ValueError(f"model_index.json must contain a {module_name} module") + logger.info(f"Diffusers config passed sanity checks") + + + # all the component models used by the pipeline + pipeline_modules = {} + for module_name, (transformers_or_diffusers, architecture) in modules_config.items(): + component_model_path = os.path.join(model_path, module_name) + module = load_pipeline_module( + module_name, + component_model_path, + transformers_or_diffusers, + architecture, + inference_args, + ) + + pipeline_modules[module_name] = module + + + # Check if all required modules were loaded + for module_name in required_modules: + if module_name not in pipeline_modules or pipeline_modules[module_name] is None: + logger.warning(f"Required module {module_name} was not loaded properly") + + return pipeline_modules + +def build_pipeline(inference_args: InferenceArgs) -> ComposedPipelineBase: + """ + Only works with valid hf diffusers configs. (model_index.json) + We want to build a pipeline based on the inference args mode_path: + 1. download the model from the hub if it's not already downloaded + 2. verify the model config and directory + 3. based on the config, determine the pipeline class + 4. parse the config to get the model components (vae, text_encoders, etc...) + 5. the pipeline loader class will use the model component names and paths to load + 6. the pipeline class will be composed of the models returned by the pipeline loader + """ + # Get pipeline type + model_path = inference_args.model_path + model_path = maybe_download_model(model_path) + # inference_args.downloaded_model_path = model_path + logger.info(f"Model path: {model_path}") + config = verify_model_config_and_directory(model_path) + + pipeline_architecture = config.get("_class_name") + if pipeline_architecture is None: + raise ValueError("Model config does not contain a _class_name attribute. " + "Only diffusers format is supported.") + + pipeline_cls, pipeline_architecture = PipelineRegistry.resolve_pipeline_cls(pipeline_architecture) + + # instantiate the pipeline + pipeline = pipeline_cls() + + try: + pipeline_modules = load_pipeline_modules(model_path, config, inference_args) + except Exception as e: + logger.error(f"Failed to load pipeline modules: {e}") + raise ValueError(f"Failed to load pipeline modules: {e}") + + logger.info(f"Initializing encoders") + pipeline.initialize_encoders(pipeline_modules, inference_args) + + logger.info(f"Registering modules") + pipeline.register_modules(pipeline_modules) + + logger.info(f"Setting up pipeline") + pipeline.setup_pipeline(inference_args) + + logger.info(f"Initializing pipeline") + pipeline.initialize_pipeline(inference_args) + + # pipeline is now initialized and ready to use + return pipeline + + + +def list_available_pipelines() -> Dict[str, Type[Any]]: + """ + List all available pipeline types. + + Returns: + A dictionary of pipeline names to pipeline classes. + """ + return PipelineRegistry.list() + +__all__ = [ + "build_pipeline", + "list_available_pipelines", + "ComposedPipelineBase", + "DiffusionPipelineOutput", +] \ No newline at end of file diff --git a/fastvideo/pipelines/composed/__init__.py b/fastvideo/v1/pipelines/composed/__init__.py similarity index 84% rename from fastvideo/pipelines/composed/__init__.py rename to fastvideo/v1/pipelines/composed/__init__.py index d451c39df..1ba564a34 100644 --- a/fastvideo/pipelines/composed/__init__.py +++ b/fastvideo/v1/pipelines/composed/__init__.py @@ -4,7 +4,7 @@ Composed pipelines for diffusion models. This package contains pipelines that are composed of multiple stages. """ -from fastvideo.pipelines.composed.composed_pipeline_base import ( +from fastvideo.v1.pipelines.composed.composed_pipeline_base import ( ComposedPipelineBase, DiffusionPipelineOutput, ) diff --git a/fastvideo/pipelines/composed/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed/composed_pipeline_base.py similarity index 89% rename from fastvideo/pipelines/composed/composed_pipeline_base.py rename to fastvideo/v1/pipelines/composed/composed_pipeline_base.py index fd04d81e0..669b4ed4f 100644 --- a/fastvideo/pipelines/composed/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed/composed_pipeline_base.py @@ -11,10 +11,10 @@ import torch import numpy as np from tqdm.auto import tqdm -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.stages import PipelineStage -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages import PipelineStage +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -63,6 +63,14 @@ class ComposedPipelineBase(ABC): """ ... + @abstractmethod + def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs): + """ + Initialize the encoders. Will remove the encoders/tokenizers modules from the + modules. Will add the TextEncoder or ImageEncoder to the modules. + """ + ... + def register_modules(self, modules: Dict[str, Any]): """ Register modules with the pipeline and its stages. diff --git a/fastvideo/pipelines/implementations/hunyuan/__init__.py b/fastvideo/v1/pipelines/implementations/hunyuan/__init__.py similarity index 52% rename from fastvideo/pipelines/implementations/hunyuan/__init__.py rename to fastvideo/v1/pipelines/implementations/hunyuan/__init__.py index 976e6fe6c..70c8c01d6 100644 --- a/fastvideo/pipelines/implementations/hunyuan/__init__.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/__init__.py @@ -4,10 +4,10 @@ HunYuan pipeline implementations. This package contains implementations of diffusion pipelines for HunYuan models. """ -from fastvideo.pipelines.implementations.hunyuan.hunyuan_pipeline import ( - HunyuanVideoPipeline, +from fastvideo.v1.pipelines.implementations.hunyuan.hunyuan_pipeline import ( + EntryClass, ) __all__ = [ - "HunyuanVideoPipeline", + "EntryClass", ] \ No newline at end of file diff --git a/fastvideo/pipelines/hunyuan/constants.py b/fastvideo/v1/pipelines/implementations/hunyuan/constants.py similarity index 100% rename from fastvideo/pipelines/hunyuan/constants.py rename to fastvideo/v1/pipelines/implementations/hunyuan/constants.py diff --git a/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py similarity index 79% rename from fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py rename to fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py index a2b71fec1..1e8bd13b0 100644 --- a/fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py +++ b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py @@ -3,20 +3,20 @@ import torch.nn as nn from pathlib import Path from safetensors.torch import load_file as safetensors_load_file -from fastvideo.pipelines.loader import PipelineLoader -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.implementations.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.loader import PipelineLoader +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.implementations.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE +from fastvideo.v1.logger import init_logger from typing import Tuple, Optional - +from fastvideo.v1.models.dits import HunyuanVideoTransformer3DModel import os import glob import json -from fastvideo.loader.fsdp_load import load_fsdp_model - -from fastvideo.models.hunyuan.text_encoder import TextEncoder - -from fastvideo.platforms import current_platform +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.vaes import AutoencoderKLHunyuanVideo +from safetensors.torch import load_file as safetensors_load_file +from fastvideo.v1.platforms import current_platform logger = init_logger(__name__) @@ -113,7 +113,7 @@ class HunyuanPipelineLoader(PipelineLoader): """Custom transformer loading for Hunyuan""" # Path to model files - # TODO(PY): remove this hardcode + # TODO(PY/Will): remove this hardcode path = "data/hunyuanvideo_community/transformer" path = Path(path) @@ -130,7 +130,7 @@ class HunyuanPipelineLoader(PipelineLoader): config.pop("_class_name") if "_diffusers_version" in config: config.pop("_diffusers_version") - + # Find all safetensors files safetensors_list = glob.glob(os.path.join(str(path), "*.safetensors")) if not safetensors_list: @@ -139,8 +139,9 @@ class HunyuanPipelineLoader(PipelineLoader): logger.info(f"Loading model from {len(safetensors_list)} safetensors files in {path}") # Load the model using FSDP loader + # TOOD(Will) new registry for DiT model = load_fsdp_model( - model_name="HunyuanVideoTransformer3DModel", + model_cls=HunyuanVideoTransformer3DModel, init_params=config, weight_dir_list=safetensors_list, device=self.device, @@ -157,7 +158,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_transformer_v0(self, inference_args: InferenceArgs): """Custom transformer loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.modules import load_model + from fastvideo.v1.models.hunyuan.modules import load_model # Disable gradient torch.set_grad_enabled(False) @@ -183,12 +184,55 @@ class HunyuanPipelineLoader(PipelineLoader): def load_vae(self, inference_args: InferenceArgs): """Custom VAE loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.vae import load_vae + use_v1_loader = True + if use_v1_loader: + return self.load_vae_v1(inference_args) + else: + return self.load_vae_v0(inference_args) + + def load_vae_v1(self, inference_args: InferenceArgs): + """Custom VAE loading for Hunyuan""" + path = "data/hunyuanvideo_community/vae" + path = Path(path) + # TODO(PY/Will): remove this hardcode + # Load config file + config_path = path / "config.json" + if not config_path.exists(): + raise ValueError(f"Config file not found at {config_path}") + + with open(config_path, "r") as f: + config = json.load(f) + + # Clean up config + + class_name = config.pop("_class_name") + if "_diffusers_version" in config: + config.pop("_diffusers_version") + vae = AutoencoderKLHunyuanVideo(**config).to(self.device) + + # Find all safetensors files + safetensors_list = glob.glob(os.path.join(str(path), "*.safetensors")) + # TODO(PY) + assert len(safetensors_list) == 1, f"Found {len(safetensors_list)} safetensors files in {path}" + loaded = safetensors_load_file(safetensors_list[0]) + vae.load_state_dict(loaded) + dtype = PRECISION_TO_TYPE[inference_args.vae_precision] + vae = vae.eval().to(dtype) + + return vae, {"s_ratio": config["spatial_compression_ratio"], "t_ratio": config["temporal_compression_ratio"]} + + + + + def load_vae_v0(self, inference_args: InferenceArgs): + """Custom VAE loading for Hunyuan""" + # TODO(will): replace this with abstracted model + from fastvideo.v1.models.hunyuan.vae import load_vae vae, _, s_ratio, t_ratio = load_vae( inference_args.vae, inference_args.vae_precision, logger=logger, - device=self.device if not inference_args.use_cpu_offload else "cpu", + device=self.device ) vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} return vae, vae_kwargs diff --git a/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py new file mode 100644 index 000000000..b78b5691e --- /dev/null +++ b/fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py @@ -0,0 +1,245 @@ +""" +HunYuan video diffusion pipeline implementation. + +This module contains an implementation of the HunYuan video diffusion pipeline +using the modular pipeline architecture. +""" + +from typing import Union, Any, Dict +import torch +from diffusers.image_processor import VaeImageProcessor + +from fastvideo.v1.pipelines.composed import ComposedPipelineBase +from fastvideo.v1.pipelines.stages import ( + InputValidationStage, + PromptEncodingStage, + TimestepPreparationStage, + LatentPreparationStage, + ConditioningStage, + DenoisingStage, + DecodingStage, + PostProcessingStage, +) +from fastvideo.v1.pipelines.stages.prompt_encoding import PromptEncodingStage +# from fastvideo.v1.pipelines.stages.timestep_preparation import FlowMatchingTimestepPreparationStage +from fastvideo.v1.inference_args import InferenceArgs +# from fastvideo.v1.pipelines.composed.composed_pipeline_base import DiffusionPipelineOutput +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +# TODO(will): move PRECISION_TO_TYPE to better place +from .constants import PROMPT_TEMPLATE, PRECISION_TO_TYPE + +from diffusers.utils import BaseOutput +import numpy as np +from dataclasses import dataclass + +from fastvideo.v1.logger import init_logger +logger = init_logger(__name__) + +# class HunyuanLatentPreparationStage(LatentPreparationStage): +# def _call_implementation(self, batch: ForwardBatch, inference_args: InferenceArgs) -> ForwardBatch: +# "custom logic for HunYuan latent preparation" +# pass + + +# class hunyuanloader(PipelineLoader): +# def load_components(self, inference_args: InferenceArgs): +# pass + +@dataclass +class DiffusionPipelineOutput(BaseOutput): + videos: Union[torch.Tensor, np.ndarray] + + +class HunyuanVideoPipeline(ComposedPipelineBase): + + def initialize_encoders(self, modules: Dict[str, Any], inference_args: InferenceArgs): + use_v1 = inference_args.use_v1_text_encoder + if use_v1: + self.initialize_encoders_v1(modules, inference_args) + else: + self.initialize_encoders_v0(modules, inference_args) + + def initialize_encoders_v0(self, modules: Dict[str, Any], inference_args: InferenceArgs): + from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder + # Text encoder + if inference_args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) + elif inference_args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + max_length = inference_args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video] + if inference_args.prompt_template_video is not None else None) + + text_encoder = TextEncoder( + text_encoder_type=inference_args.text_encoder, + max_length=max_length, + text_encoder_precision=inference_args.text_encoder_precision, + tokenizer_type=inference_args.tokenizer, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=inference_args.hidden_state_skip_layer, + apply_final_norm=inference_args.apply_final_norm, + reproduce=inference_args.reproduce, + logger=logger, + device=self.device if not inference_args.use_cpu_offload else "cpu", + ) + text_encoder_2 = None + if inference_args.text_encoder_2 is not None: + text_encoder_2 = TextEncoder( + text_encoder_type=inference_args.text_encoder_2, + max_length=inference_args.text_len_2, + text_encoder_precision=inference_args.text_encoder_precision_2, + tokenizer_type=inference_args.tokenizer_2, + reproduce=inference_args.reproduce, + logger=logger, + device=self.device if not inference_args.use_cpu_offload else "cpu", + ) + modules["text_encoder"] = text_encoder + modules["text_encoder_2"] = text_encoder_2 + + def initialize_encoders_v1(self, modules: Dict[str, Any], inference_args: InferenceArgs): + """ + Initialize the encoders. Will remove the encoders/tokenizers modules from the + modules. Will add the TextEncoder or ImageEncoder to the modules. + """ + from fastvideo.v1.models.text_encoder import TextEncoder + if inference_args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template_video].get("crop_start", 0) + elif inference_args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[inference_args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + + max_length = inference_args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[inference_args.prompt_template] if inference_args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[inference_args.prompt_template_video] + if inference_args.prompt_template_video is not None else None) + + encoder_1 = modules.pop("text_encoder") + assert encoder_1 is not None, "Text encoder is not found" + encoder_1.to(inference_args.device) + encoder_1.to(dtype=PRECISION_TO_TYPE[inference_args.text_encoder_precision]) + encoder_1.requires_grad_(False) + + tokenizer_1 = modules.pop("tokenizer") + assert tokenizer_1 is not None, "Tokenizer is not found" + + text_encoder = TextEncoder( + text_encoder=encoder_1, + tokenizer=tokenizer_1, + # text_encoder_type="text_encoder" + max_length=max_length, + # text_encoder_precision=inference_args.text_encoder_precision, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=inference_args.hidden_state_skip_layer, + apply_final_norm=inference_args.apply_final_norm, + reproduce=inference_args.reproduce, + device=inference_args.device if not inference_args.use_cpu_offload else "cpu", + ) + + encoder_2 = modules.pop("text_encoder_2") + assert encoder_2 is not None, "Text encoder 2 is not found" + encoder_2.to(inference_args.device) + encoder_2.to(dtype=PRECISION_TO_TYPE[inference_args.text_encoder_precision]) + encoder_2.requires_grad_(False) + + tokenizer_2 = modules.pop("tokenizer_2") + assert tokenizer_2 is not None, "Tokenizer 2 is not found" + + text_encoder_2 = TextEncoder( + text_encoder=encoder_2, + tokenizer=tokenizer_2, + # text_encoder_type="text_encoder_2", + max_length=inference_args.text_len_2, + # text_encoder_precision=inference_args.text_encoder_precision, + reproduce=inference_args.reproduce, + device=inference_args.device if not inference_args.use_cpu_offload else "cpu", + ) + modules["text_encoder"] = text_encoder + modules["text_encoder_2"] = text_encoder_2 + + + def setup_pipeline(self, inference_args: InferenceArgs): + self.add_stage("input_validation_stage", + InputValidationStage()) + self.add_stage("prompt_encoding_stage_primary", + PromptEncodingStage(is_secondary=False)) + self.add_stage("prompt_encoding_stage_secondary", + PromptEncodingStage(is_secondary=True)) + self.add_stage("conditioning_stage", + ConditioningStage()) + self.add_stage("timestep_preparation_stage", + TimestepPreparationStage()) + self.add_stage("latent_preparation_stage", + LatentPreparationStage()) + self.add_stage("denoising_stage", + DenoisingStage()) + self.add_stage("decoding_stage", + DecodingStage()) + + def initialize_pipeline(self, inference_args: InferenceArgs): + assert len(self._stages) > 0, "Pipeline stages are not set" + assert len(self._modules) > 0, "Pipeline modules are not set" + + + vae_scale_factor = 2**(len(self.vae.block_out_channels) - 1) + inference_args.vae_scale_factor = vae_scale_factor + + self.image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + self.register_modules({"image_processor": self.image_processor}) + + + num_channels_latents = self.transformer.in_channels + inference_args.num_channels_latents = num_channels_latents + + + + def adjust_video_length(self, batch: ForwardBatch, inference_args: InferenceArgs): + """Adjust video length based on VAE version""" + video_length = batch.num_frames + vae_ver = inference_args.vae + logger.info(f"Adjusting video length for VAE version: {vae_ver}") + if "884" in vae_ver: + batch.num_frames = (video_length - 1) // 4 + 1 + elif "888" in vae_ver: + batch.num_frames = (video_length - 1) // 8 + 1 + return batch + + + @torch.no_grad() + def forward(self, batch: ForwardBatch, inference_args: InferenceArgs) -> ForwardBatch: + logger.info(f"Running pipeline stages: {self._stage_name_mapping.keys()}") + logger.info(f"Batch: {batch}") + # for stage in self._stages: + # batch = stage(batch, inference_args) + + # or + + batch = self.input_validation_stage(batch, inference_args) + batch = self.prompt_encoding_stage_primary(batch, inference_args) + batch = self.prompt_encoding_stage_secondary(batch, inference_args) + batch = self.conditioning_stage(batch, inference_args) + batch = self.timestep_preparation_stage(batch, inference_args) + + # custom logic + batch = self.adjust_video_length(batch, inference_args) + + batch = self.latent_preparation_stage(batch, inference_args) + batch = self.denoising_stage(batch, inference_args) + batch = self.decoding_stage(batch, inference_args) + + return DiffusionPipelineOutput(videos=batch.videos) + +EntryClass = HunyuanVideoPipeline \ No newline at end of file diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py similarity index 100% rename from fastvideo/pipelines/pipeline_batch_info.py rename to fastvideo/v1/pipelines/pipeline_batch_info.py diff --git a/fastvideo/v1/pipelines/pipeline_registry.py b/fastvideo/v1/pipelines/pipeline_registry.py new file mode 100644 index 000000000..10d601093 --- /dev/null +++ b/fastvideo/v1/pipelines/pipeline_registry.py @@ -0,0 +1,91 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/models/registry.py +# and https://github.com/sgl-project/sglang/blob/v0.4.3/python/sglang/srt/models/registry.py + +import importlib +import pkgutil +from dataclasses import dataclass, field +from functools import lru_cache +from typing import AbstractSet, Dict, List, Optional, Tuple, Type, Union + +from fastvideo.v1.logger import init_logger + +import torch.nn as nn + +logger = init_logger(__name__) + + +@dataclass +class _PipelineRegistry: + # Keyed by pipeline_arch + pipelines: Dict[str, Union[Type[nn.Module], str]] = field(default_factory=dict) + + def get_supported_archs(self) -> AbstractSet[str]: + return self.pipelines.keys() + + def _raise_for_unsupported(self, architectures: List[str]): + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Pipeline architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + raise ValueError( + f"Pipeline architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_pipeline_cls(self, pipeline_arch: str) -> Optional[Type[nn.Module]]: + if pipeline_arch not in self.pipelines: + return None + + return self.pipelines[pipeline_arch] + + def resolve_pipeline_cls( + self, + architecture: str, + ) -> Tuple[Type[nn.Module], str]: + if not architecture: + logger.warning("No pipeline architecture is specified") + + pipeline_cls = self._try_load_pipeline_cls(architecture) + if pipeline_cls is not None: + return (pipeline_cls, architecture) + + return self._raise_for_unsupported(architecture) + + +@lru_cache() +def import_pipeline_classes(): + pipeline_arch_name_to_cls = {} + package_name = "fastvideo.v1.pipelines.implementations" + package = importlib.import_module(package_name) + for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."): + if ispkg: + try: + module = importlib.import_module(name) + except Exception as e: + logger.warning(f"Ignore import error when loading {name}. " f"{e}") + continue + if hasattr(module, "EntryClass"): + entry = module.EntryClass + print(entry) + print(entry.__name__) + if isinstance( + entry, list + ): # To support multiple pipeline classes in one module + for tmp in entry: + assert ( + tmp.__name__ not in pipeline_arch_name_to_cls + ), f"Duplicated pipeline implementation for {tmp.__name__}" + pipeline_arch_name_to_cls[tmp.__name__] = tmp + else: + assert ( + entry.__name__ not in pipeline_arch_name_to_cls + ), f"Duplicated pipeline implementation for {entry.__name__}" + pipeline_arch_name_to_cls[entry.__name__] = entry + return pipeline_arch_name_to_cls + + +PipelineRegistry = _PipelineRegistry(import_pipeline_classes()) \ No newline at end of file diff --git a/fastvideo/v1/pipelines/stages/__init__.py b/fastvideo/v1/pipelines/stages/__init__.py new file mode 100644 index 000000000..696f1b337 --- /dev/null +++ b/fastvideo/v1/pipelines/stages/__init__.py @@ -0,0 +1,28 @@ +""" +Pipeline stages for diffusion models. + +This package contains the various stages that can be composed to create +complete diffusion pipelines. +""" + +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.stages.input_validation import InputValidationStage +from fastvideo.v1.pipelines.stages.prompt_encoding import PromptEncodingStage +from fastvideo.v1.pipelines.stages.timestep_preparation import TimestepPreparationStage +from fastvideo.v1.pipelines.stages.latent_preparation import LatentPreparationStage +from fastvideo.v1.pipelines.stages.conditioning import ConditioningStage +from fastvideo.v1.pipelines.stages.denoising import DenoisingStage +from fastvideo.v1.pipelines.stages.decoding import DecodingStage +from fastvideo.v1.pipelines.stages.post_processing import PostProcessingStage + +__all__ = [ + "PipelineStage", + "InputValidationStage", + "PromptEncodingStage", + "TimestepPreparationStage", + "LatentPreparationStage", + "ConditioningStage", + "DenoisingStage", + "DecodingStage", + "PostProcessingStage", +] \ No newline at end of file diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/v1/pipelines/stages/base.py similarity index 94% rename from fastvideo/pipelines/stages/base.py rename to fastvideo/v1/pipelines/stages/base.py index 83ee6da17..f858725f6 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/v1/pipelines/stages/base.py @@ -11,9 +11,9 @@ import time import traceback from typing import Dict, Any -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -36,7 +36,7 @@ class PipelineStage(ABC): """ self._enable_logging = enable_logging self._stage_name = self.__class__.__name__ - self._logger = init_logger(f"fastvideo.pipelines.stages.{self._stage_name}") + self._logger = init_logger(f"fastvideo.v1.pipelines.stages.{self._stage_name}") @property def device(self) -> torch.device: diff --git a/fastvideo/pipelines/stages/conditioning.py b/fastvideo/v1/pipelines/stages/conditioning.py similarity index 90% rename from fastvideo/pipelines/stages/conditioning.py rename to fastvideo/v1/pipelines/stages/conditioning.py index c907e92d8..9765a2f13 100644 --- a/fastvideo/pipelines/stages/conditioning.py +++ b/fastvideo/v1/pipelines/stages/conditioning.py @@ -5,10 +5,10 @@ Conditioning stage for diffusion pipelines. import torch from typing import Optional -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/v1/pipelines/stages/decoding.py similarity index 63% rename from fastvideo/pipelines/stages/decoding.py rename to fastvideo/v1/pipelines/stages/decoding.py index 61775ef11..bd8d6bd5c 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/v1/pipelines/stages/decoding.py @@ -3,17 +3,11 @@ Decoding stage for diffusion pipelines. """ import torch -from typing import Optional, Union, List, Tuple -import numpy as np - -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D -from fastvideo.logger import init_logger -from diffusers.utils import BaseOutput - +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.utils import PRECISION_TO_TYPE +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -50,18 +44,6 @@ class DecodingStage(PipelineStage): vae_dtype = PRECISION_TO_TYPE[inference_args.vae_precision] vae_autocast_enabled = (vae_dtype != torch.float32) and not inference_args.disable_autocast - # Handle different latent shapes - expand_temporal_dim = False - if len(latents.shape) == 4: - if isinstance(self.vae, AutoencoderKLCausal3D): - latents = latents.unsqueeze(2) - expand_temporal_dim = True - elif len(latents.shape) == 5: - pass - else: - raise ValueError( - f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}." - ) # Apply scaling/shifting if needed if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor): @@ -73,11 +55,11 @@ class DecodingStage(PipelineStage): with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): if inference_args.vae_tiling: self.vae.enable_tiling() - image = self.vae.decode(latents, return_dict=False, generator=batch.generator)[0] + # if inference_args.vae_sp: + # self.vae.enable_parallel() + image = self.vae.decode(latents) + - # Handle temporal dimension if needed - if expand_temporal_dim or image.shape[2] == 1: - image = image.squeeze(2) # Normalize image to [0, 1] range image = (image / 2 + 0.5).clamp(0, 1) diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py similarity index 96% rename from fastvideo/pipelines/stages/denoising.py rename to fastvideo/v1/pipelines/stages/denoising.py index a7101b32b..6f2997c60 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -8,13 +8,14 @@ from typing import Optional, Dict, Any, List from tqdm.auto import tqdm from einops import rearrange -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.utils import PRECISION_TO_TYPE +# TODO(will-refactor): change this to fastvideo.distributed from fastvideo.utils.parallel_states import nccl_info, get_sequence_parallel_state from fastvideo.utils.communications import all_gather -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/v1/pipelines/stages/input_validation.py similarity index 94% rename from fastvideo/pipelines/stages/input_validation.py rename to fastvideo/v1/pipelines/stages/input_validation.py index 0cf5e8929..2108c7279 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/v1/pipelines/stages/input_validation.py @@ -5,10 +5,10 @@ Input validation stage for diffusion pipelines. from typing import Optional, Union, List import torch -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger import random logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/v1/pipelines/stages/latent_preparation.py similarity index 94% rename from fastvideo/pipelines/stages/latent_preparation.py rename to fastvideo/v1/pipelines/stages/latent_preparation.py index ea83f1116..f54817227 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/v1/pipelines/stages/latent_preparation.py @@ -5,11 +5,11 @@ Latent preparation stage for diffusion pipelines. import torch from typing import Optional, Tuple -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs from diffusers.utils.torch_utils import randn_tensor -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/post_processing.py b/fastvideo/v1/pipelines/stages/post_processing.py similarity index 85% rename from fastvideo/pipelines/stages/post_processing.py rename to fastvideo/v1/pipelines/stages/post_processing.py index 4de05674f..8bce92fba 100644 --- a/fastvideo/pipelines/stages/post_processing.py +++ b/fastvideo/v1/pipelines/stages/post_processing.py @@ -6,11 +6,11 @@ import torch import numpy as np from typing import Optional, Union, List, Tuple -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs from diffusers.utils import BaseOutput -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/pipelines/stages/prompt_encoding.py b/fastvideo/v1/pipelines/stages/prompt_encoding.py similarity index 87% rename from fastvideo/pipelines/stages/prompt_encoding.py rename to fastvideo/v1/pipelines/stages/prompt_encoding.py index 4785ffe8f..9117e8ae0 100644 --- a/fastvideo/pipelines/stages/prompt_encoding.py +++ b/fastvideo/v1/pipelines/stages/prompt_encoding.py @@ -7,11 +7,11 @@ This module contains implementations of prompt encoding stages for diffusion pip from typing import Any, Dict, List, Optional, Union, Tuple import torch -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.stages import PipelineStage -from fastvideo.models.hunyuan.text_encoder import TextEncoder -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages import PipelineStage +from fastvideo.v1.models.text_encoder import TextEncoder +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -87,13 +87,19 @@ class PromptEncodingStage(PipelineStage): attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len) if text_encoder is not None: - prompt_embeds_dtype = text_encoder.dtype + # TODO(will-refactor): use text_encoder.dtype + prompt_embeds_dtype = torch.float16 elif self.transformer is not None: prompt_embeds_dtype = self.transformer.dtype else: prompt_embeds_dtype = prompt_embeds.dtype prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + # print("prompt_embeds", type(prompt_embeds)) + # logger.info(f"prompt_embeds shape: {prompt_embeds.shape}") + if prompt_embeds.ndim == 1: + prompt_embeds = prompt_embeds.unsqueeze(0) + # attention_mask = attention_mask.unsqueeze(0) if prompt_embeds.ndim == 2: bs_embed, _ = prompt_embeds.shape diff --git a/fastvideo/pipelines/stages/timestep_preparation.py b/fastvideo/v1/pipelines/stages/timestep_preparation.py similarity index 93% rename from fastvideo/pipelines/stages/timestep_preparation.py rename to fastvideo/v1/pipelines/stages/timestep_preparation.py index f16ae3224..e5b87bbd5 100644 --- a/fastvideo/pipelines/stages/timestep_preparation.py +++ b/fastvideo/v1/pipelines/stages/timestep_preparation.py @@ -8,10 +8,10 @@ from typing import Any, Dict, List, Optional, Union, Tuple import torch import inspect -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.stages.base import PipelineStage +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/platforms/__init__.py b/fastvideo/v1/platforms/__init__.py similarity index 100% rename from fastvideo/platforms/__init__.py rename to fastvideo/v1/platforms/__init__.py diff --git a/fastvideo/platforms/cuda.py b/fastvideo/v1/platforms/cuda.py similarity index 97% rename from fastvideo/platforms/cuda.py rename to fastvideo/v1/platforms/cuda.py index f02a4b415..f9bb18a5b 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/v1/platforms/cuda.py @@ -12,9 +12,9 @@ import torch from typing_extensions import ParamSpec # import custom ops, trigger op registration -import fastvideo.envs as envs -from fastvideo.logger import init_logger -from fastvideo.utils import import_pynvml +import fastvideo.v1.envs as envs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import import_pynvml from .interface import DeviceCapability, Platform, PlatformEnum, _Backend @@ -115,7 +115,7 @@ class CudaPlatformBase(Platform): @classmethod def get_device_communicator_cls(cls) -> str: - return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa + return "fastvideo.v1.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa # NVML utils diff --git a/fastvideo/platforms/interface.py b/fastvideo/v1/platforms/interface.py similarity index 97% rename from fastvideo/platforms/interface.py rename to fastvideo/v1/platforms/interface.py index 8b061535f..d92e563f4 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/v1/platforms/interface.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union import numpy as np import torch -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) @@ -186,7 +186,7 @@ class Platform: """ Get device specific communicator class for distributed communication. """ - return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa + return "fastvideo.v1.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa class UnspecifiedPlatform(Platform): diff --git a/fastvideo/sample/v2_fastvideo_inference.py b/fastvideo/v1/sample/v1_fastvideo_inference.py similarity index 74% rename from fastvideo/sample/v2_fastvideo_inference.py rename to fastvideo/v1/sample/v1_fastvideo_inference.py index fe47e09e9..bfd9560fb 100644 --- a/fastvideo/sample/v2_fastvideo_inference.py +++ b/fastvideo/v1/sample/v1_fastvideo_inference.py @@ -9,23 +9,13 @@ import torchvision from einops import rearrange # Fix the import path -from fastvideo.inference_engine import InferenceEngine -from fastvideo.inference_args import InferenceArgs -from fastvideo.utils.utils import FlexibleArgumentParser -from fastvideo.distributed import init_distributed_environment, initialize_model_parallel -from fastvideo.logger import init_logger -import torch.distributed as dist -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state +from fastvideo.v1.inference_engine import InferenceEngine +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.utils import FlexibleArgumentParser +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) -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(inference_args: InferenceArgs): # initialize_distributed() # print(nccl_info.sp_size) @@ -43,9 +33,10 @@ def main(inference_args: InferenceArgs): # Initialize tensor model parallel groups initialize_model_parallel( - sequence_model_parallel_size=inference_args.sp_size + sequence_model_parallel_size=inference_args.sp_size, + tensor_model_parallel_size=inference_args.tp_size, ) - initialize_sequence_parallel_state(world_size) + # initialize_sequence_parallel_state(world_size) # initialize_distributed() diff --git a/fastvideo/v1/distributed_example.py b/fastvideo/v1/tests/distributed_example.py similarity index 90% rename from fastvideo/v1/distributed_example.py rename to fastvideo/v1/tests/distributed_example.py index 2e15baa6c..c37d8a0f0 100644 --- a/fastvideo/v1/distributed_example.py +++ b/fastvideo/v1/tests/distributed_example.py @@ -1,4 +1,4 @@ -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) import os diff --git a/fastvideo/v1/distributed_run.sh b/fastvideo/v1/tests/distributed_run.sh similarity index 100% rename from fastvideo/v1/distributed_run.sh rename to fastvideo/v1/tests/distributed_run.sh diff --git a/fastvideo/v1/tests/test_clip_encoder.py b/fastvideo/v1/tests/test_clip_encoder.py new file mode 100644 index 000000000..c7a8297af --- /dev/null +++ b/fastvideo/v1/tests/test_clip_encoder.py @@ -0,0 +1,207 @@ +from fastvideo.v1.models.encoders.clip import CLIPTextModel +from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder, load_text_encoder, load_tokenizer +import os +import torch +import torch.nn as nn +import argparse +import numpy as np +from fastvideo.v1.logger import init_logger +from transformers import AutoConfig +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel +# from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer + +logger = init_logger(__name__) + +def initialize_identical_weights(model1, model2, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + for name2, param2 in params2.items(): + if 'weight' in name2: + # Reset seed to get same initialization + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + for name2, param2 in params2.items(): + if 'bias' in name2: + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + param2.data = param2.data.to(torch.bfloat16) + + logger.info("Both models initialized with identical weights in bfloat16") + return model1, model2 + +def setup_args(): + parser = argparse.ArgumentParser(description='CLIP Text Encoder Test') + parser.add_argument('--model-path', type=str, default="openai/clip-vit-large-patch14", + help='Path to the CLIP model') + parser.add_argument('--precision', type=str, default="float16", + help='Precision to use for the model (float32, float16, bfloat16)') + return parser.parse_args() + +def test_clip_encoder(): + init_distributed_environment(world_size=1, rank=0, distributed_init_method="env://", local_rank=0, backend="nccl") + initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1, backend="nccl") + args = setup_args() + + # Set fixed random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + # Initialize the two model implementations + logger.info(f"Loading models from {args.model_path}") + model_path = "data/hunyuanvideo-community/HunyuanVideo/text_encoder_2" + + # config = json.load(open(os.path.join(model_path, "config.json"))) + + hf_config = AutoConfig.from_pretrained(model_path) + print(hf_config) + print(hf_config.use_return_dict) + + # Load our implementation using the loader from text_encoder/__init__.py + model1, _ = load_text_encoder( + text_encoder_type="clipL", + text_encoder_precision='fp16', + text_encoder_path=model_path, + logger=logger, + device=device + ) + + from fastvideo.v1.models.loader import TextEncoderLoader + from fastvideo.v1.loader.loader import DefaultModelLoader + loader = DefaultModelLoader() + args.device_str = "cuda:0" + model2 = loader.load_model(model_path, hf_config, args) + + # Load the HuggingFace implementation directly + # model2 = CLIPTextModel(hf_config) + model2 = model2.to(torch.float16) + model2 = model2.to(device) + model2.eval() + + # Sanity check weights between the two models + logger.info("Comparing model weights for sanity check...") + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Check number of parameters + logger.info(f"Model1 has {len(params1)} parameters") + logger.info(f"Model2 has {len(params2)} parameters") + + # Compare a few key parameters + + # weight_diffs = [] + # for (name1, param1), (name2, param2) in zip( + # sorted(params1.items()), sorted(params2.items()) + # ): + # # if len(weight_diffs) < 5: # Just check a few parameters + # max_diff = torch.max(torch.abs(param1 - param2)).item() + # mean_diff = torch.mean(torch.abs(param1 - param2)).item() + # weight_diffs.append((name1, name2, max_diff, mean_diff)) + # logger.info(f"Parameter: {name1} vs {name2}") + # logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") + + # Load tokenizer + tokenizer, _ = load_tokenizer( + tokenizer_type="clipL", + tokenizer_path=args.model_path, + logger=logger + ) + + # Test with some sample prompts + prompts = [ + "a photo of a cat", + "a beautiful landscape with mountains", + "an astronaut riding a horse on the moon" + ] + + logger.info("Testing CLIP text encoder with sample prompts") + + with torch.no_grad(): + for prompt in prompts: + logger.info(f"Testing prompt: '{prompt}'") + + # Tokenize the prompt + tokens = tokenizer( + prompt, + padding="max_length", + max_length=77, + truncation=True, + return_tensors="pt" + ).to(device) + + # Get embeddings from our implementation + outputs1 = model1( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + logger.info(f"Testing model2") + + # Get embeddings from HuggingFace implementation + outputs2 = model2( + input_ids=tokens.input_ids, + # attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + # Compare last hidden states + last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask==1] + last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask==1] + # print("last_hidden_state1", last_hidden_state1) + # print("last_hidden_state2", last_hidden_state2) + + assert last_hidden_state1.shape == last_hidden_state2.shape, \ + f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" + + max_diff_hidden = torch.max(torch.abs(last_hidden_state1 - last_hidden_state2)) + mean_diff_hidden = torch.mean(torch.abs(last_hidden_state1 - last_hidden_state2)) + + logger.info(f"Maximum difference in last hidden states: {max_diff_hidden.item()}") + logger.info(f"Mean difference in last hidden states: {mean_diff_hidden.item()}") + + # Compare pooler outputs + pooler_output1 = outputs1.pooler_output + pooler_output2 = outputs2.pooler_output + + assert pooler_output1.shape == pooler_output2.shape, \ + f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}" + + max_diff_pooler = torch.max(torch.abs(pooler_output1 - pooler_output2)) + mean_diff_pooler = torch.mean(torch.abs(pooler_output1 - pooler_output2)) + + logger.info(f"Maximum difference in pooler outputs: {max_diff_pooler.item()}") + logger.info(f"Mean difference in pooler outputs: {mean_diff_pooler.item()}") + + # Check if outputs are similar (allowing for small numerical differences) + assert max_diff_hidden < 1e-4, \ + f"Hidden states differ significantly: max diff = {max_diff_hidden.item()}" + assert max_diff_pooler < 1e-4, \ + f"Pooler outputs differ significantly: max diff = {max_diff_pooler.item()}" + + logger.info("Test passed! Both CLIP text encoder implementations produce similar outputs.") + logger.info("Test completed successfully") + +if __name__ == "__main__": + test_clip_encoder() \ No newline at end of file diff --git a/fastvideo/v1/tests/test_hunyuan_vae.py b/fastvideo/v1/tests/test_hunyuan_vae.py new file mode 100644 index 000000000..e73624379 --- /dev/null +++ b/fastvideo/v1/tests/test_hunyuan_vae.py @@ -0,0 +1,161 @@ +from fastvideo.v1.models.vaes.hunyuanvae import AutoencoderKLHunyuanVideo as MyHunyuanVAE +from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE + +import os +import torch +import torch.nn as nn +import argparse +import numpy as np +import json +from fastvideo.v1.logger import init_logger +from safetensors.torch import load_file +logger = init_logger(__name__) + +def initialize_identical_weights(model1, model2, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + for name2, param2 in params2.items(): + if 'weight' in name2: + # Reset seed to get same initialization + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + for name2, param2 in params2.items(): + if 'bias' in name2: + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + param2.data = param2.data.to(torch.bfloat16) + + logger.info("Both models initialized with identical weights in bfloat16") + return model1, model2 + +def setup_args(): + parser = argparse.ArgumentParser(description='HunyuanVAE Test') + parser.add_argument('--in-channels', type=int, default=4, + help='Number of input channels') + parser.add_argument('--out-channels', type=int, default=4, + help='Number of output channels') + parser.add_argument('--latent-channels', type=int, default=4, + help='Number of latent channels') + return parser.parse_args() + +def test_hunyuan_vae(): + args = setup_args() + + # Set fixed random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Model parameters + in_channels = args.in_channels + out_channels = args.out_channels + latent_channels = args.latent_channels + + device = torch.device("cuda:0") + print + # Initialize the two model implementations + path = "data/hunyuanvideo-community/HunyuanVideo/vae" + config_path = os.path.join(path, "config.json") + config = json.load(open(config_path)) + config.pop("_class_name") + config.pop("_diffusers_version") + model1 = MyHunyuanVAE( + **config + ).to(torch.bfloat16) + + model2 = DiffusersHunyuanVAE( + **config + ).to(torch.bfloat16) + + loaded = load_file(os.path.join(path, "diffusion_pytorch_model.safetensors")) + model1.load_state_dict(loaded) + model2.load_state_dict(loaded) + + + # Set both models to eval mode + model1.eval() + model2.eval() + + # Move to GPU + model1 = model1.to(device) + model2 = model2.to(device) + + + model1.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_min_num_frames=8, + tile_sample_stride_height=16, + tile_sample_stride_width=16, + tile_sample_stride_num_frames=4 + ) + model2.enable_tiling( + tile_sample_min_height=32, + tile_sample_min_width=32, + tile_sample_min_num_frames=8, + tile_sample_stride_height=16, + tile_sample_stride_width=16, + tile_sample_stride_num_frames=4 + ) + + # Create identical inputs for both models + batch_size = 1 + + # Video input [B, C, T, H, W] + input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16) + + # Disable gradients for inference + with torch.no_grad(): + # Test encoding + logger.info("Testing encoding...") + latent1 = model1.encode(input_tensor).mean + print("--------------------------------") + latent2 = model2.encode(input_tensor).latent_dist.mean + # Check if latents have the same shape + assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" + assert latent1.shape == latent2.shape, f"Latent shapes don't match: {latent1.shape} vs {latent2.shape}" + # Check if latents are similar + max_diff_encode = torch.max(torch.abs(latent1 - latent2)) + mean_diff_encode = torch.mean(torch.abs(latent1 - latent2)) + logger.info(f"Maximum difference between encoded latents: {max_diff_encode.item()}") + logger.info(f"Mean difference between encoded latents: {mean_diff_encode.item()}") + assert max_diff_encode < 1e-4, f"Encoded latents differ significantly: max diff = {max_diff_encode.item()}" + # Test decoding + logger.info("Testing decoding...") + output1 = model1.decode(latent1) + output2 = model2.decode(latent2).sample + # Check if outputs have the same shape + assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" + + # Check if outputs are similar + max_diff_decode = torch.max(torch.abs(output1 - output2)) + mean_diff_decode = torch.mean(torch.abs(output1 - output2)) + logger.info(f"Maximum difference between decoded outputs: {max_diff_decode.item()}") + logger.info(f"Mean difference between decoded outputs: {mean_diff_decode.item()}") + assert max_diff_decode < 1e-4, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}" + + + logger.info("Test passed! Both VAE implementations produce similar outputs.") + logger.info("Test completed successfully") + +if __name__ == "__main__": + test_hunyuan_vae() \ No newline at end of file diff --git a/fastvideo/v1/test_hunyuanvideo.py b/fastvideo/v1/tests/test_hunyuanvideo.py similarity index 94% rename from fastvideo/v1/test_hunyuanvideo.py rename to fastvideo/v1/tests/test_hunyuanvideo.py index ed68b8d89..1194510ab 100644 --- a/fastvideo/v1/test_hunyuanvideo.py +++ b/fastvideo/v1/tests/test_hunyuanvideo.py @@ -5,8 +5,8 @@ import torch.nn as nn import argparse import numpy as np -from fastvideo.logger import init_logger -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, get_sequence_model_parallel_rank, @@ -16,13 +16,13 @@ from fastvideo.distributed.parallel_state import ( cleanup_dist_env_and_memory ) -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state +from fastvideo.v1.utils.parallel_states import initialize_sequence_parallel_state from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.device_mesh import init_device_mesh -from fastvideo.loader.fsdp_load import shard_model -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT -from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer -from fastvideo.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel +from fastvideo.v1.loader.fsdp_load import shard_model +from fastvideo.v1.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit +from fastvideo.v1.models.hunyuan.modules.models import HYVideoDiffusionTransformer +from fastvideo.v1.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel logger = init_logger(__name__) def initialize_identical_weights(model1, model2, seed=42): @@ -116,7 +116,7 @@ def test_hunyuanvideo_distributed(): patch_size = [1, 2, 2] torch.cuda.set_device(f"cuda:{local_rank}") # Initialize the two model implementations - model1 = HunyuanVideoDiT( + model1 = HunyuanVideoDit( patch_size=2, patch_size_t=1, in_channels=4, diff --git a/fastvideo/v1/test_hunyuanvideo_load.py b/fastvideo/v1/tests/test_hunyuanvideo_load.py similarity index 94% rename from fastvideo/v1/test_hunyuanvideo_load.py rename to fastvideo/v1/tests/test_hunyuanvideo_load.py index e04c1915b..03789aa1b 100644 --- a/fastvideo/v1/test_hunyuanvideo_load.py +++ b/fastvideo/v1/tests/test_hunyuanvideo_load.py @@ -4,8 +4,8 @@ import torch.nn as nn import argparse import numpy as np -from fastvideo.logger import init_logger -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, get_sequence_model_parallel_rank, @@ -15,11 +15,11 @@ from fastvideo.distributed.parallel_state import ( cleanup_dist_env_and_memory ) import json -from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT -from fastvideo.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG -from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer -from fastvideo.loader.fsdp_load import load_fsdp_model -from fastvideo.models.hunyuan.inference import Inference +from fastvideo.v1.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit +from fastvideo.v1.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG +from fastvideo.v1.models.hunyuan.modules.models import HYVideoDiffusionTransformer +from fastvideo.v1.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.models.hunyuan.inference import Inference import glob logger = init_logger(__name__) diff --git a/fastvideo/v1/tests/test_llama_encoder.py b/fastvideo/v1/tests/test_llama_encoder.py new file mode 100644 index 000000000..11153d195 --- /dev/null +++ b/fastvideo/v1/tests/test_llama_encoder.py @@ -0,0 +1,201 @@ +from fastvideo.v1.models.encoders.llama import LlamaModel +from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder, load_text_encoder, load_tokenizer +import os +import torch +import torch.nn as nn +import argparse +import numpy as np +from fastvideo.v1.logger import init_logger +from transformers import AutoConfig +from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel + +logger = init_logger(__name__) + +def initialize_identical_weights(model1, model2, seed=42): + """Initialize both models with identical weights using a fixed seed for reproducibility.""" + # Get all parameters from both models + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Initialize each layer with identical values + with torch.no_grad(): + # Initialize weights + for name1, param1 in params1.items(): + if 'weight' in name1: + # Set seed before each weight initialization + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + + for name2, param2 in params2.items(): + if 'weight' in name2: + # Reset seed to get same initialization + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + + # Initialize biases + for name1, param1 in params1.items(): + if 'bias' in name1: + torch.manual_seed(seed) + nn.init.normal_(param1, mean=0.0, std=0.05) + param1.data = param1.data.to(torch.bfloat16) + + for name2, param2 in params2.items(): + if 'bias' in name2: + torch.manual_seed(seed) + nn.init.normal_(param2, mean=0.0, std=0.05) + param2.data = param2.data.to(torch.bfloat16) + + logger.info("Both models initialized with identical weights in bfloat16") + return model1, model2 + +def setup_args(): + parser = argparse.ArgumentParser(description='LLaMA Encoder Test') + parser.add_argument('--model-path', type=str, default="meta-llama/Llama-2-7b-hf", + help='Path to the LLaMA model') + parser.add_argument('--precision', type=str, default="float16", + help='Precision to use for the model (float32, float16, bfloat16)') + return parser.parse_args() + +def test_llama_encoder(): + init_distributed_environment(world_size=1, rank=0, distributed_init_method="env://", local_rank=0, backend="nccl") + initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1, backend="nccl") + args = setup_args() + + # Set fixed random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + # Initialize the two model implementations + logger.info(f"Loading models from {args.model_path}") + model_path = "data/hunyuanvideo-community/HunyuanVideo/text_encoder" + + hf_config = AutoConfig.from_pretrained(model_path) + print(hf_config) + + # Load our implementation using the loader from text_encoder/__init__.py + model1, _ = load_text_encoder( + text_encoder_type="llm", + text_encoder_precision='fp16', + text_encoder_path=model_path, + logger=logger, + device=device + ) + + from fastvideo.v1.models.loader import TextEncoderLoader + from fastvideo.v1.loader.loader import DefaultModelLoader + loader = DefaultModelLoader() + args.device_str = "cuda:0" + model2 = loader.load_model(model_path, hf_config, args) + + # Convert to float16 and move to device + model2 = model2.to(torch.float16) + model2 = model2.to(device) + model2.eval() + + # Sanity check weights between the two models + logger.info("Comparing model weights for sanity check...") + params1 = dict(model1.named_parameters()) + params2 = dict(model2.named_parameters()) + + # Check number of parameters + logger.info(f"Model1 has {len(params1)} parameters") + logger.info(f"Model2 has {len(params2)} parameters") + + # Compare a few key parameters + weight_diffs = [] + # check if embed_tokens are the same + print(model1.embed_tokens.weight.shape, model2.embed_tokens.weight.shape) + assert torch.allclose(model1.embed_tokens.weight, model2.embed_tokens.weight) + weights = ["layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight"] + # for (name1, param1), (name2, param2) in zip( + # sorted(params1.items()), sorted(params2.items()) + # ): + for l in range(hf_config.num_hidden_layers): + for w in weights: + name1 = w.format(l) + name2 = w.format(l) + p1 = params1[name1] + p2 = params2[name2] + print(type(p2)) + if "gate_up" in name2: + print("skipping gate_up") + continue + try: + logger.info(f"Parameter: {name1} vs {name2}") + max_diff = torch.max(torch.abs(p1 - p2)).item() + mean_diff = torch.mean(torch.abs(p1 - p2)).item() + weight_diffs.append((name1, name2, max_diff, mean_diff)) + logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") + except Exception as e: + logger.info(f"Error comparing {name1} and {name2}: {e}") + + tokenizer_path = "data/hunyuanvideo-community/HunyuanVideo/tokenizer" + # Load tokenizer + tokenizer, _ = load_tokenizer( + tokenizer_type="llm", + tokenizer_path=tokenizer_path, + logger=logger + ) + + # Test with some sample prompts + prompts = [ + "Once upon a time", + # "The quick brown fox jumps over", + # "In a galaxy far, far away" + ] + + logger.info("Testing LLaMA encoder with sample prompts") + + with torch.no_grad(): + for prompt in prompts: + logger.info(f"Testing prompt: '{prompt}'") + + # Tokenize the prompt + tokens = tokenizer( + prompt, + padding="max_length", + max_length=128, + truncation=True, + return_tensors="pt" + ).to(device) + + # Get outputs from our implementation + # filter out padding input_ids + # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] + # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] + outputs1 = model1( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + print("--------------------------------") + logger.info(f"Testing model2") + + # Get outputs from HuggingFace implementation + outputs2 = model2( + input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True + ) + + # Compare last hidden states + last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask==1] + last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask==1] + + assert last_hidden_state1.shape == last_hidden_state2.shape, \ + f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" + + max_diff_hidden = torch.max(torch.abs(last_hidden_state1 - last_hidden_state2)) + mean_diff_hidden = torch.mean(torch.abs(last_hidden_state1 - last_hidden_state2)) + + logger.info(f"Maximum difference in last hidden states: {max_diff_hidden.item()}") + logger.info(f"Mean difference in last hidden states: {mean_diff_hidden.item()}") + + + logger.info("Test passed! Both LLaMA encoder implementations produce similar outputs.") + logger.info("Test completed successfully") + +if __name__ == "__main__": + test_llama_encoder() diff --git a/fastvideo/v1/tp_example.py b/fastvideo/v1/tests/tp_example.py similarity index 96% rename from fastvideo/v1/tp_example.py rename to fastvideo/v1/tests/tp_example.py index 88377a188..c6b27fe4e 100644 --- a/fastvideo/v1/tp_example.py +++ b/fastvideo/v1/tests/tp_example.py @@ -4,8 +4,8 @@ import torch.nn as nn import torch.distributed as dist import argparse -from fastvideo.logger import init_logger -from fastvideo.distributed.parallel_state import ( +from fastvideo.v1.logger import init_logger +from fastvideo.v1.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, get_tensor_model_parallel_rank, @@ -14,11 +14,11 @@ from fastvideo.distributed.parallel_state import ( destroy_distributed_environment, cleanup_dist_env_and_memory ) -from fastvideo.layers.linear import ( +from fastvideo.v1.layers.linear import ( ColumnParallelLinear, RowParallelLinear ) -from fastvideo.distributed.communication_op import ( +from fastvideo.v1.distributed.communication_op import ( tensor_model_parallel_all_reduce, tensor_model_parallel_all_gather ) diff --git a/fastvideo/utils/utils.py b/fastvideo/v1/utils.py similarity index 93% rename from fastvideo/utils/utils.py rename to fastvideo/v1/utils.py index 7c9896432..8f0741ed7 100644 --- a/fastvideo/utils/utils.py +++ b/fastvideo/v1/utils.py @@ -1,8 +1,9 @@ import torch -import fastvideo.envs as envs +import fastvideo.v1.envs as envs import inspect -from fastvideo.logger import init_logger +from fastvideo.v1.logger import init_logger import argparse +import math import sys from typing import List, Dict, Union, Type, Any, TypeVar import yaml @@ -13,6 +14,14 @@ logger = init_logger(__name__) T = TypeVar("T") +# TODO(will): used to convert inference_args.precision to torch.dtype. Find a +# cleaner way to do this. +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + def find_nccl_library() -> str: """ @@ -63,7 +72,7 @@ def current_stream() -> torch.cuda.Stream: the underlying hypothesis is that we do not call `torch._C._cuda_setStream` from C/C++ code. """ - from fastvideo.platforms import current_platform + from fastvideo.v1.platforms import current_platform global _current_stream if _current_stream is None: # when this function is called before any stream is set, @@ -288,3 +297,16 @@ def warn_for_unimplemented_methods(cls: Type[T]) -> Type[T]: type.__setattr__(cls, '__init__', wrapped_init) return cls + + +def align_to(value, alignment): + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) diff --git a/fastvideo/pipelines/loader.py b/fastvideo/v1/v0_reference_src/loader.py similarity index 90% rename from fastvideo/pipelines/loader.py rename to fastvideo/v1/v0_reference_src/loader.py index 146b3232e..e1f8fa323 100644 --- a/fastvideo/pipelines/loader.py +++ b/fastvideo/v1/v0_reference_src/loader.py @@ -1,12 +1,12 @@ from abc import ABC, abstractmethod import torch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler -from fastvideo.logger import init_logger -from fastvideo.platforms import current_platform -from fastvideo.distributed.parallel_state import get_sp_group -from fastvideo.pipelines.composed import ComposedPipelineBase +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler +from fastvideo.v1.logger import init_logger +from fastvideo.v1.platforms import current_platform +from fastvideo.v1.distributed.parallel_state import get_sp_group +from fastvideo.v1.pipelines.composed import ComposedPipelineBase import os import json @@ -26,14 +26,14 @@ def resolve_pipeline_cls(inference_args: InferenceArgs): """Resolve the pipeline class based on the inference args.""" # TODO(will): resolve the pipeline class based on the model path # read from hf config file _class_name - from fastvideo.pipelines.hunyuan.hunyuan_video import HunyuanVideoPipeline + from fastvideo.v1.pipelines.hunyuan.hunyuan_video import HunyuanVideoPipeline return HunyuanVideoPipeline def resolve_pipeline_cls_v2(inference_args: InferenceArgs): """Resolve the pipeline class based on the inference args.""" # TODO(will): resolve the pipeline class based on the model path # read from hf config file _class_name - from fastvideo.pipelines.implementations.hunyuan import HunyuanVideoPipeline + from fastvideo.v1.pipelines.implementations.hunyuan import HunyuanVideoPipeline return HunyuanVideoPipeline @@ -246,11 +246,19 @@ class PipelineLoader(ABC): solver=inference_args.flow_solver, ) return None - + + @staticmethod + def load_pipeline_from_config(config: dict): + """Load a pipeline from a config.""" + pipeline_cls = resolve_pipeline_cls(inference_args) + assert issubclass(pipeline_cls, ComposedPipelineBase) + logger.info(f"Pipeline class: {pipeline_cls}") + + def get_pipeline_loader(inference_args: InferenceArgs) -> PipelineLoader: """Get a pipeline loader based on the inference args.""" - from fastvideo.pipelines.implementations.hunyuan.hunyuan_loader import HunyuanPipelineLoader + from fastvideo.v1.pipelines.implementations.hunyuan.hunyuan_loader import HunyuanPipelineLoader # Map model types to their specific loaders loader_map = { "hunyuan_video": HunyuanPipelineLoader, diff --git a/fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py b/fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py new file mode 100644 index 000000000..c0334a9e6 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/flash_attn_no_pad.py @@ -0,0 +1,28 @@ +from einops import rearrange +from flash_attn import flash_attn_varlen_qkvpacked_func +from flash_attn.bert_padding import pad_input, unpad_input + + +def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None): + # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27 + batch_size = qkv.shape[0] + seqlen = qkv.shape[1] + nheads = qkv.shape[-2] + x = rearrange(qkv, "b s three h d -> b s (three h d)") + x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask) + + x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads) + output_unpad = flash_attn_varlen_qkvpacked_func( + x_unpad, + cu_seqlens, + max_s, + dropout_p, + softmax_scale=softmax_scale, + causal=causal, + ) + output = rearrange( + pad_input(rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen), + "b s (h d) -> b s h d", + h=nheads, + ) + return output \ No newline at end of file diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py new file mode 100644 index 000000000..53c9b0628 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/constants.py @@ -0,0 +1,94 @@ +import os + +import torch + +__all__ = [ + "C_SCALE", + "PROMPT_TEMPLATE", + "MODEL_BASE", + "PRECISIONS", + "NORMALIZATION_TYPE", + "ACTIVATION_TYPE", + "VAE_PATH", + "TEXT_ENCODER_PATH", + "TOKENIZER_PATH", + "TEXT_PROJECTION", + "DATA_TYPE", + "NEGATIVE_PROMPT", +] + +PRECISION_TO_TYPE = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +# =================== Constant Values ===================== +# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid +# overflow error when tensorboard logging values. +C_SCALE = 1_000_000_000_000_000 + +# When using decoder-only models, we must provide a prompt template to instruct the text encoder +# on how to generate the text. +# -------------------------------------------------------------------- +PROMPT_TEMPLATE_ENCODE = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " + "1. The main content and theme of the video." + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." + "4. background environment, light, style and atmosphere." + "5. camera angles, movements, and transitions used in the video:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") + +NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion" + +PROMPT_TEMPLATE = { + "dit-llm-encode": { + "template": PROMPT_TEMPLATE_ENCODE, + "crop_start": 36, + }, + "dit-llm-encode-video": { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, + }, +} + +# ======================= Model ====================== +PRECISIONS = {"fp32", "fp16", "bf16"} +NORMALIZATION_TYPE = {"layer", "rms"} +ACTIVATION_TYPE = {"relu", "silu", "gelu", "gelu_tanh"} + +# =================== Model Path ===================== +MODEL_BASE = os.getenv("MODEL_BASE", "./data/hunyuan") + +# =================== Data ======================= +DATA_TYPE = {"image", "video", "image_video"} + +# 3D VAE +# VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"} +VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/vae"} + +# Text Encoder +TEXT_ENCODER_PATH = { + "clipL": f"{MODEL_BASE}/text_encoder_2", + "llm": f"{MODEL_BASE}/text_encoder", +} + +# Tokenizer +# TOKENIZER_PATH = { +# "clipL": f"{MODEL_BASE}/text_encoder_2", +# "llm": f"{MODEL_BASE}/text_encoder", +# } +TOKENIZER_PATH = { + "clipL": f"{MODEL_BASE}/tokenizer_2", + "llm": f"{MODEL_BASE}/tokenizer", +} + +TEXT_PROJECTION = { + "linear", # Default, an nn.Linear() layer + "single_refiner", # Single TokenRefiner. Refer to LI-DiT +} diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/__init__.py new file mode 100644 index 000000000..754bc48ec --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/__init__.py @@ -0,0 +1,3 @@ +# ruff: noqa: F401 +from .pipelines import HunyuanVideoPipeline +from .schedulers import FlowMatchDiscreteScheduler diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/__init__.py new file mode 100644 index 000000000..fa109e14d --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa: F401 +from .pipeline_hunyuan_video import HunyuanVideoPipeline diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py new file mode 100644 index 000000000..2efc9fecf --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py @@ -0,0 +1,931 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +import inspect +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.configuration_utils import FrozenDict +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.models import AutoencoderKL +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, logging, replace_example_docstring, + scale_lora_layers) +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange + +from fastvideo.v1.utils.communications import all_gather +from fastvideo.v1.utils.parallel_states import get_sequence_parallel_state, nccl_info + +from ...constants import PRECISION_TO_TYPE +from ...modules import HYVideoDiffusionTransformer +from ...text_encoder import TextEncoder +from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """""" + + +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = (guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg) + return noise_cfg + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler.") + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +@dataclass +class HunyuanVideoPipelineOutput(BaseOutput): + videos: Union[torch.Tensor, np.ndarray] + + +class HunyuanVideoPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-video generation using HunyuanVideo. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. + text_encoder ([`TextEncoder`]): + Frozen text-encoder. + text_encoder_2 ([`TextEncoder`]): + Frozen text-encoder_2. + transformer ([`HYVideoDiffusionTransformer`]): + A `HYVideoDiffusionTransformer` to denoise the encoded video latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" + _optional_components = ["text_encoder_2"] + _exclude_from_cpu_offload = ["transformer"] + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: TextEncoder, + transformer: HYVideoDiffusionTransformer, + scheduler: KarrasDiffusionSchedulers, + text_encoder_2: Optional[TextEncoder] = None, + progress_bar_config: Dict[str, Any] = None, + args=None, + ): + super().__init__() + + # ========================================================================================== + if progress_bar_config is None: + progress_bar_config = {} + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + self._progress_bar_config.update(progress_bar_config) + + self.args = args + # ========================================================================================== + + if (hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1): + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" + f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " + "to update the config accordingly as leaving `steps_offset` might led to incorrect results" + " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," + " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" + " file") + deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["steps_offset"] = 1 + scheduler._internal_dict = FrozenDict(new_config) + + if (hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True): + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." + " `clip_sample` should be set to False in the configuration file. Please make sure to update the" + " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" + " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" + " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file") + deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False) + new_config = dict(scheduler.config) + new_config["clip_sample"] = False + scheduler._internal_dict = FrozenDict(new_config) + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + transformer=transformer, + scheduler=scheduler, + text_encoder_2=text_encoder_2, + ) + self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + def encode_prompt( + self, + prompt, + device, + num_videos_per_prompt, + do_classifier_free_guidance, + negative_prompt=None, + prompt_embeds: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_attention_mask: Optional[torch.Tensor] = None, + lora_scale: Optional[float] = None, + clip_skip: Optional[int] = None, + text_encoder: Optional[TextEncoder] = None, + data_type: Optional[str] = "image", + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_videos_per_prompt (`int`): + number of videos that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the video generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + attention_mask (`torch.Tensor`, *optional*): + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + negative_attention_mask (`torch.Tensor`, *optional*): + lora_scale (`float`, *optional*): + A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + text_encoder (TextEncoder, *optional*): + data_type (`str`, *optional*): + """ + if text_encoder is None: + text_encoder = self.text_encoder + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if not USE_PEFT_BACKEND: + adjust_lora_scale_text_encoder(text_encoder.model, lora_scale) + else: + scale_lora_layers(text_encoder.model, lora_scale) + + if prompt_embeds is None: + # textual inversion: process multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer) + + text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) + if clip_skip is None: + prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type, device=device) + prompt_embeds = prompt_outputs.hidden_state + else: + prompt_outputs = text_encoder.encode( + text_inputs, + output_hidden_states=True, + data_type=data_type, + device=device, + ) + # Access the `hidden_states` first, that contains a tuple of + # all the hidden states from the encoder layers. Then index into + # the tuple to access the hidden states from the desired layer. + prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)] + # We also need to apply the final LayerNorm here to not mess with the + # representations. The `last_hidden_states` that we typically use for + # obtaining the final prompt representations passes through the LayerNorm + # layer. + prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds) + + attention_mask = prompt_outputs.attention_mask + if attention_mask is not None: + attention_mask = attention_mask.to(device) + bs_embed, seq_len = attention_mask.shape + attention_mask = attention_mask.repeat(1, num_videos_per_prompt) + attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len) + + if text_encoder is not None: + prompt_embeds_dtype = text_encoder.dtype + elif self.transformer is not None: + prompt_embeds_dtype = self.transformer.dtype + else: + prompt_embeds_dtype = prompt_embeds.dtype + + prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + + if prompt_embeds.ndim == 2: + bs_embed, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1) + else: + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + + return ( + prompt_embeds, + negative_prompt_embeds, + attention_mask, + negative_attention_mask, + ) + + def decode_latents(self, latents, enable_tiling=True): + deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead" + deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False) + + latents = 1 / self.vae.config.scaling_factor * latents + if enable_tiling: + self.vae.enable_tiling() + image = self.vae.decode(latents, return_dict=False)[0] + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + if image.ndim == 4: + image = image.cpu().permute(0, 2, 3, 1).float() + else: + image = image.cpu().float() + return image + + def prepare_extra_func_kwargs(self, func, kwargs): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + extra_step_kwargs = {} + + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_step_kwargs[k] = v + return extra_step_kwargs + + def check_inputs( + self, + prompt, + height, + width, + video_length, + callback_steps, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + vae_ver="88-4c-sd", + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if video_length is not None: + if "884" in vae_ver: + if video_length != 1 and (video_length - 1) % 4 != 0: + raise ValueError(f"`video_length` has to be 1 or a multiple of 4 but is {video_length}.") + elif "888" in vae_ver: + if video_length != 1 and (video_length - 1) % 8 != 0: + raise ValueError(f"`video_length` has to be 1 or a multiple of 8 but is {video_length}.") + + if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0): + raise ValueError(f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}.") + if callback_on_step_end_tensor_inputs is not None and not all(k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two.") + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined.") + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError(f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}.") + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + video_length, + dtype, + device, + generator, + latents=None, + ): + shape = ( + batch_size, + num_channels_latents, + video_length, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators.") + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler + if hasattr(self.scheduler, "init_noise_sigma"): + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding + def get_guidance_scale_embedding( + self, + w: torch.Tensor, + embedding_dim: int = 512, + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """ + See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 + + Args: + w (`torch.Tensor`): + Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. + embedding_dim (`int`, *optional*, defaults to 512): + Dimension of the embeddings to generate. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + Data type of the generated embeddings. + + Returns: + `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. + """ + assert len(w.shape) == 1 + w = w * 1000.0 + + half_dim = embedding_dim // 2 + emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) + emb = w.to(dtype)[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: # zero pad + emb = torch.nn.functional.pad(emb, (0, 1)) + assert emb.shape == (w.shape[0], embedding_dim) + return emb + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def clip_skip(self): + return self._clip_skip + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None + return self._guidance_scale > 1 + + @property + def cross_attention_kwargs(self): + return self._cross_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]], + height: int, + width: int, + video_length: int, + data_type: str = "video", + num_inference_steps: int = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, + MultiPipelineCallbacks, ]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + vae_ver: str = "88-4c-sd", + enable_tiling: bool = False, + enable_vae_sp: bool = False, + n_tokens: Optional[int] = None, + embedded_guidance_scale: Optional[float] = None, + mask_strategy: Optional[Dict[str, list]] = None, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`): + The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + video_length (`int`): + The number of frames in the generated video. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide what to not include in image generation. If not defined, you need to + pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies + to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If + not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. + + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a + plain tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in + [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~HunyuanVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated images and the + second element is a list of `bool`s indicating whether the corresponding generated image contains + "not-safe-for-work" (nsfw) content. + """ + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 0. Default height and width to unet + # height = height or self.transformer.config.sample_size * self.vae_scale_factor + # width = width or self.transformer.config.sample_size * self.vae_scale_factor + # to deal with lora scaling and other possible forward hooks + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + video_length, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + vae_ver=vae_ver, + ) + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = (torch.device(f"cuda:{dist.get_rank()}") if dist.is_initialized() else self._execution_device) + + # 3. Encode input prompt + lora_scale = (self.cross_attention_kwargs.get("scale", None) + if self.cross_attention_kwargs is not None else None) + + ( + prompt_embeds, + negative_prompt_embeds, + prompt_mask, + negative_prompt_mask, + ) = self.encode_prompt( + prompt, + device, + num_videos_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + attention_mask=attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_attention_mask=negative_attention_mask, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + data_type=data_type, + ) + if self.text_encoder_2 is not None: + ( + prompt_embeds_2, + negative_prompt_embeds_2, + prompt_mask_2, + negative_prompt_mask_2, + ) = self.encode_prompt( + prompt, + device, + num_videos_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=None, + attention_mask=None, + negative_prompt_embeds=None, + negative_attention_mask=None, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + text_encoder=self.text_encoder_2, + data_type=data_type, + ) + else: + prompt_embeds_2 = None + negative_prompt_embeds_2 = None + prompt_mask_2 = None + negative_prompt_mask_2 = None + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + if prompt_mask is not None: + prompt_mask = torch.cat([negative_prompt_mask, prompt_mask]) + if prompt_embeds_2 is not None: + prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2]) + if prompt_mask_2 is not None: + prompt_mask_2 = torch.cat([negative_prompt_mask_2, prompt_mask_2]) + + # 4. Prepare timesteps + extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs(self.scheduler.set_timesteps, + {"n_tokens": n_tokens}) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + **extra_set_timesteps_kwargs, + ) + if "884" in vae_ver: + video_length = (video_length - 1) // 4 + 1 + elif "888" in vae_ver: + video_length = (video_length - 1) // 8 + 1 + else: + video_length = video_length + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + video_length, + prompt_embeds.dtype, + device, + generator, + latents, + ) + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + 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, :, :, :] + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + { + "generator": generator, + "eta": eta + }, + ) + + target_dtype = PRECISION_TO_TYPE[self.args.precision] + autocast_enabled = (target_dtype != torch.float32) and not self.args.disable_autocast + vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision] + vae_autocast_enabled = (vae_dtype != torch.float32) and not self.args.disable_autocast + + # 7. Denoising loop + 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): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = (torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents) + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + t_expand = t.repeat(latent_model_input.shape[0]) + guidance_expand = (torch.tensor( + [embedded_guidance_scale] * latent_model_input.shape[0], + dtype=torch.float32, + device=device, + ).to(target_dtype) * 1000.0 if embedded_guidance_scale is not None else None) + # predict the noise residual + with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled): + # concat prompt_embeds_2 and prompt_embeds. Mismatch fill with zeros + if prompt_embeds_2.shape[-1] != prompt_embeds.shape[-1]: + prompt_embeds_2 = F.pad( + prompt_embeds_2, + (0, prompt_embeds.shape[2] - prompt_embeds_2.shape[1]), + value=0, + ).unsqueeze(1) + encoder_hidden_states = torch.cat([prompt_embeds_2, prompt_embeds], dim=1) + noise_pred = self.transformer( # For an input image (129, 192, 336) (1, 256, 256) + latent_model_input, + encoder_hidden_states, + t_expand, + prompt_mask, + mask_strategy=mask_strategy[i], + guidance=guidance_expand, + return_dict=False, + )[0] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + + if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg( + noise_pred, + noise_pred_text, + guidance_rescale=self.guidance_rescale, + ) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if progress_bar is not None: + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + + if not output_type == "latent": + expand_temporal_dim = False + if len(latents.shape) == 4: + if isinstance(self.vae, AutoencoderKLCausal3D): + latents = latents.unsqueeze(2) + expand_temporal_dim = True + elif len(latents.shape) == 5: + pass + else: + raise ValueError( + f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}.") + + if (hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor): + latents = (latents / self.vae.config.scaling_factor + self.vae.config.shift_factor) + else: + latents = latents / self.vae.config.scaling_factor + + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): + if enable_tiling: + self.vae.enable_tiling() + if enable_vae_sp: + self.vae.enable_parallel() + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + + if expand_temporal_dim or image.shape[2] == 1: + image = image.squeeze(2) + + else: + image = latents + + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 + image = image.cpu().float() + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return image + + return HunyuanVideoPipelineOutput(videos=image) diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/__init__.py new file mode 100644 index 000000000..2238803fe --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa: F401 +from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py new file mode 100644 index 000000000..bf032ab0e --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py @@ -0,0 +1,239 @@ +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, logging + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + reverse: bool = True, + solver: str = "euler", + n_tokens: Optional[int] = None, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32) + + self._step_index = None + self._begin_index = None + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError(f"Solver {solver} not supported. Supported solvers: {self.supported_solver}") + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device] = None, + n_tokens: int = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + sigmas = self.sd3_time_shift(sigmas) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(dtype=torch.float32, device=device) + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: + return sample + + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor)): + raise ValueError(("Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.to(torch.float32) * dt + else: + raise ValueError(f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}") + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample, ) + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/idle_config.py b/fastvideo/v1/v0_reference_src/models/hunyuan/idle_config.py new file mode 100644 index 000000000..7d13c9d0a --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/idle_config.py @@ -0,0 +1,380 @@ +# ruff: noqa: F405, F403 +import argparse +import re + +from .constants import * +from .modules.models import HUNYUAN_VIDEO_CONFIG + + +def parse_args(namespace=None): + parser = argparse.ArgumentParser(description="HunyuanVideo inference script") + + parser = add_network_args(parser) + parser = add_extra_models_args(parser) + parser = add_denoise_schedule_args(parser) + parser = add_inference_args(parser) + parser = add_parallel_args(parser) + + args = parser.parse_args(namespace=namespace) + args = sanity_check_args(args) + + return args + + +def add_network_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="HunyuanVideo network args") + + # Main model + group.add_argument( + "--model", + type=str, + choices=list(HUNYUAN_VIDEO_CONFIG.keys()), + default="HYVideo-T/2-cfgdistill", + ) + group.add_argument( + "--latent-channels", + type=str, + default=16, + help="Number of latent channels of DiT. If None, it will be determined by `vae`. If provided, " + "it still needs to match the latent channels of the VAE model.", + ) + group.add_argument( + "--precision", + type=str, + default="bf16", + choices=PRECISIONS, + help="Precision mode. Options: fp32, fp16, bf16. Applied to the backbone model and optimizer.", + ) + + # RoPE + group.add_argument("--rope-theta", type=int, default=256, help="Theta used in RoPE.") + return parser + + +def add_extra_models_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Extra models args, including vae, text encoders and tokenizers)") + + # - VAE + group.add_argument( + "--vae", + type=str, + default="884-16c-hy", + choices=list(VAE_PATH), + help="Name of the VAE model.", + ) + group.add_argument( + "--vae-precision", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the VAE model.", + ) + group.add_argument( + "--vae-tiling", + action="store_true", + help="Enable tiling for the VAE model to save GPU memory.", + ) + group.set_defaults(vae_tiling=True) + + group.add_argument( + "--text-encoder", + type=str, + default="llm", + choices=list(TEXT_ENCODER_PATH), + help="Name of the text encoder model.", + ) + group.add_argument( + "--text-encoder-precision", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the text encoder model.", + ) + group.add_argument( + "--text-states-dim", + type=int, + default=4096, + help="Dimension of the text encoder hidden states.", + ) + group.add_argument("--text-len", type=int, default=256, help="Maximum length of the text input.") + group.add_argument( + "--tokenizer", + type=str, + default="llm", + choices=list(TOKENIZER_PATH), + help="Name of the tokenizer model.", + ) + group.add_argument( + "--prompt-template", + type=str, + default="dit-llm-encode", + choices=PROMPT_TEMPLATE, + help="Image prompt template for the decoder-only text encoder model.", + ) + group.add_argument( + "--prompt-template-video", + type=str, + default="dit-llm-encode-video", + choices=PROMPT_TEMPLATE, + help="Video prompt template for the decoder-only text encoder model.", + ) + group.add_argument( + "--hidden-state-skip-layer", + type=int, + default=2, + help="Skip layer for hidden states.", + ) + group.add_argument( + "--apply-final-norm", + action="store_true", + help="Apply final normalization to the used text encoder hidden states.", + ) + + # - CLIP + group.add_argument( + "--text-encoder-2", + type=str, + default="clipL", + choices=list(TEXT_ENCODER_PATH), + help="Name of the second text encoder model.", + ) + group.add_argument( + "--text-encoder-precision-2", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the second text encoder model.", + ) + group.add_argument( + "--text-states-dim-2", + type=int, + default=768, + help="Dimension of the second text encoder hidden states.", + ) + group.add_argument( + "--tokenizer-2", + type=str, + default="clipL", + choices=list(TOKENIZER_PATH), + help="Name of the second tokenizer model.", + ) + group.add_argument( + "--text-len-2", + type=int, + default=77, + help="Maximum length of the second text input.", + ) + + return parser + + +def add_denoise_schedule_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Denoise schedule args") + + group.add_argument( + "--denoise-type", + type=str, + default="flow", + help="Denoise type for noised inputs.", + ) + + # Flow Matching + group.add_argument( + "--flow-shift", + type=float, + default=7.0, + help="Shift factor for flow matching schedulers.", + ) + group.add_argument( + "--flow-reverse", + action="store_true", + help="If reverse, learning/sampling from t=1 -> t=0.", + ) + group.add_argument( + "--flow-solver", + type=str, + default="euler", + help="Solver for flow matching.", + ) + group.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)", + ) + group.add_argument( + "--linear-schedule-end", + type=int, + default=25, + help="End step for linear quadratic schedule for flow matching.", + ) + + return parser + + +def add_inference_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Inference args") + + # ======================== Model loads ======================== + group.add_argument( + "--model-base", + type=str, + default="ckpts", + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--dit-weight", + type=str, + default="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt", + help="Path to the HunyuanVideo model. If None, search the model in the args.model_root." + "1. If it is a file, load the model directly." + "2. If it is a directory, search the model in the directory. Support two types of models: " + "1) named `pytorch_model_*.pt`" + "2) named `*_model_states.pt`, where * can be `mp_rank_00`.", + ) + group.add_argument( + "--model-resolution", + type=str, + default="540p", + choices=["540p", "720p"], + help="Root path of all the models, including t2v models and extra models.", + ) + group.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.", + ) + group.add_argument( + "--use-cpu-offload", + action="store_true", + help="Use CPU offload for the model load.", + ) + + # ======================== Inference general setting ======================== + group.add_argument( + "--batch-size", + type=int, + default=1, + help="Batch size for inference and evaluation.", + ) + group.add_argument( + "--infer-steps", + type=int, + default=50, + help="Number of denoising steps for inference.", + ) + group.add_argument( + "--disable-autocast", + action="store_true", + help="Disable autocast for denoising loop and vae decoding in pipeline sampling.", + ) + group.add_argument( + "--save-path", + type=str, + default="./results", + help="Path to save the generated samples.", + ) + group.add_argument( + "--save-path-suffix", + type=str, + default="", + help="Suffix for the directory of saved samples.", + ) + group.add_argument( + "--name-suffix", + type=str, + default="", + help="Suffix for the names of saved samples.", + ) + group.add_argument( + "--num-videos", + type=int, + default=1, + help="Number of videos to generate for each prompt.", + ) + # ---sample size--- + group.add_argument( + "--video-size", + type=int, + nargs="+", + default=(720, 1280), + help="Video size for training. If a single value is provided, it will be used for both height " + "and width. If two values are provided, they will be used for height and width " + "respectively.", + ) + group.add_argument( + "--video-length", + type=int, + default=129, + help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1", + ) + # --- prompt --- + group.add_argument( + "--prompt", + type=str, + default=None, + help="Prompt for sampling during evaluation.", + ) + group.add_argument( + "--seed-type", + type=str, + default="auto", + choices=["file", "random", "fixed", "auto"], + help="Seed type for evaluation. If file, use the seed from the CSV file. If random, generate a " + "random seed. If fixed, use the fixed seed given by `--seed`. If auto, `csv` will use the " + "seed column if available, otherwise use the fixed `seed` value. `prompt` will use the " + "fixed `seed` value.", + ) + group.add_argument("--seed", type=int, default=None, help="Seed for evaluation.") + + # Classifier-Free Guidance + group.add_argument("--neg-prompt", type=str, default=None, help="Negative prompt for sampling.") + group.add_argument("--cfg-scale", type=float, default=1.0, help="Classifier free guidance scale.") + group.add_argument( + "--embedded-cfg-scale", + type=float, + default=6.0, + help="Embedded classifier free guidance scale.", + ) + + group.add_argument( + "--reproduce", + action="store_true", + help="Enable reproducibility by setting random seeds and deterministic algorithms.", + ) + + return parser + + +def add_parallel_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Parallel args") + + # ======================== Model loads ======================== + group.add_argument( + "--ulysses-degree", + type=int, + default=1, + help="Ulysses degree.", + ) + group.add_argument( + "--ring-degree", + type=int, + default=1, + help="Ulysses degree.", + ) + + return parser + + +def sanity_check_args(args): + # VAE channels + vae_pattern = r"\d{2,3}-\d{1,2}c-\w+" + if not re.match(vae_pattern, args.vae): + raise ValueError(f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'.") + vae_channels = int(args.vae.split("-")[1][:-1]) + if args.latent_channels is None: + args.latent_channels = vae_channels + if vae_channels != args.latent_channels: + raise ValueError(f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels}).") + return args diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/inference.py b/fastvideo/v1/v0_reference_src/models/hunyuan/inference.py new file mode 100644 index 000000000..6dcadc690 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/inference.py @@ -0,0 +1,504 @@ +import os +import random +import time +from pathlib import Path + +import torch +from loguru import logger +from safetensors.torch import load_file as safetensors_load_file + +from fastvideo.v1.models.hunyuan.constants import NEGATIVE_PROMPT, PRECISION_TO_TYPE, PROMPT_TEMPLATE +from fastvideo.v1.models.hunyuan.diffusion.pipelines import HunyuanVideoPipeline +from fastvideo.v1.models.hunyuan.diffusion.schedulers import FlowMatchDiscreteScheduler +from fastvideo.v1.models.hunyuan.modules import load_model +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.utils.data_utils import align_to +from fastvideo.v1.models.hunyuan.vae import load_vae +from fastvideo.v1.utils.parallel_states import nccl_info + + +class Inference(object): + + def __init__( + self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + use_cpu_offload=False, + device=None, + logger=None, + parallel_args=None, + ): + self.vae = vae + self.vae_kwargs = vae_kwargs + + self.text_encoder = text_encoder + self.text_encoder_2 = text_encoder_2 + + self.model = model + self.pipeline = pipeline + self.use_cpu_offload = use_cpu_offload + + self.args = args + self.device = (device if device is not None else "cuda" if torch.cuda.is_available() else "cpu") + self.logger = logger + self.parallel_args = parallel_args + + @classmethod + def from_pretrained(cls, pretrained_model_path, args, device=None, **kwargs): + """ + Initialize the Inference pipeline. + + Args: + pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints. + args (argparse.Namespace): The arguments for the pipeline. + device (int): The device for inference. Default is 0. + """ + from fastvideo.v1.pipelines.loader import get_pipeline_loader + loader = get_pipeline_loader(args) + vae, vae_kwargs, text_encoder, text_encoder_2, transformer, scheduler = \ + loader.load_components(args) + return cls( + args=args, + vae=vae, + vae_kwargs=vae_kwargs, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + model=transformer, + use_cpu_offload=args.use_cpu_offload, + device=device, + logger=logger, + parallel_args=None, + ) + # ======================================================================== + logger.info(f"Got text-to-video model root path: {pretrained_model_path}") + + # ==================== Initialize Distributed Environment ================ + 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" + + parallel_args = None # {"ulysses_degree": args.ulysses_degree, "ring_degree": args.ring_degree} + + # ======================== Get the args path ============================= + + # Disable gradient + torch.set_grad_enabled(False) + + # =========================== Build main model =========================== + logger.info("Building model...") + factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]} + in_channels = args.latent_channels + out_channels = args.latent_channels + + model = load_model( + args, + in_channels=in_channels, + out_channels=out_channels, + factor_kwargs=factor_kwargs, + ) + 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 ======================== + # VAE + vae, _, s_ratio, t_ratio = load_vae( + args.vae, + args.vae_precision, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} + + # Text encoder + if args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template_video].get("crop_start", 0) + elif args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + max_length = args.text_len + crop_start + + # prompt_template + prompt_template = (PROMPT_TEMPLATE[args.prompt_template] if args.prompt_template is not None else None) + + # prompt_template_video + prompt_template_video = (PROMPT_TEMPLATE[args.prompt_template_video] + if args.prompt_template_video is not None else None) + + text_encoder = TextEncoder( + text_encoder_type=args.text_encoder, + max_length=max_length, + text_encoder_precision=args.text_encoder_precision, + tokenizer_type=args.tokenizer, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=args.hidden_state_skip_layer, + apply_final_norm=args.apply_final_norm, + reproduce=args.reproduce, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + text_encoder_2 = None + if args.text_encoder_2 is not None: + text_encoder_2 = TextEncoder( + text_encoder_type=args.text_encoder_2, + max_length=args.text_len_2, + text_encoder_precision=args.text_encoder_precision_2, + tokenizer_type=args.tokenizer_2, + reproduce=args.reproduce, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + + return cls( + args=args, + vae=vae, + vae_kwargs=vae_kwargs, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + model=model, + use_cpu_offload=args.use_cpu_offload, + device=device, + logger=logger, + parallel_args=parallel_args, + ) + + @staticmethod + def load_state_dict(args, model, pretrained_model_path): + load_key = args.load_key + dit_weight = Path(args.dit_weight) + + if dit_weight is None: + model_dir = pretrained_model_path / f"t2v_{args.model_resolution}" + files = list(model_dir.glob("*.pt")) + if len(files) == 0: + raise ValueError(f"No model weights found in {model_dir}") + if str(files[0]).startswith("pytorch_model_"): + model_path = dit_weight / f"pytorch_model_{load_key}.pt" + bare_model = True + elif any(str(f).endswith("_model_states.pt") for f in files): + files = [f for f in files if str(f).endswith("_model_states.pt")] + model_path = files[0] + if len(files) > 1: + logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}") + bare_model = False + else: + raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: " + f"{list(map(str, files))}. When given a directory as --dit-weight, only " + f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and " + f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a " + f"specific weight file, please provide the full path to the file.") + else: + if dit_weight.is_dir(): + files = list(dit_weight.glob("*.pt")) + if len(files) == 0: + raise ValueError(f"No model weights found in {dit_weight}") + if str(files[0]).startswith("pytorch_model_"): + model_path = dit_weight / f"pytorch_model_{load_key}.pt" + bare_model = True + elif any(str(f).endswith("_model_states.pt") for f in files): + files = [f for f in files if str(f).endswith("_model_states.pt")] + model_path = files[0] + if len(files) > 1: + logger.warning(f"Multiple model weights found in {dit_weight}, using {model_path}") + bare_model = False + else: + raise ValueError(f"Invalid model path: {dit_weight} with unrecognized weight format: " + f"{list(map(str, files))}. When given a directory as --dit-weight, only " + f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and " + f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a " + f"specific weight file, please provide the full path to the file.") + elif dit_weight.is_file(): + model_path = dit_weight + bare_model = "unknown" + else: + raise ValueError(f"Invalid model path: {dit_weight}") + + print('model_path', model_path) + if not model_path.exists(): + raise ValueError(f"model_path not exists: {model_path}") + logger.info(f"Loading torch model {model_path}...") + if model_path.suffix == ".safetensors": + # Use safetensors library for .safetensors files + state_dict = safetensors_load_file(model_path) + elif model_path.suffix == ".pt": + # Use torch for .pt files + state_dict = torch.load(model_path, map_location=lambda storage, loc: storage) + else: + raise ValueError(f"Unsupported file format: {model_path}") + + if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}.") + model.load_state_dict(state_dict, strict=True) + return model + + @staticmethod + def parse_size(size): + if isinstance(size, int): + size = [size] + if not isinstance(size, (list, tuple)): + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + if len(size) == 1: + size = [size[0], size[0]] + if len(size) != 2: + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + return size + + +class HunyuanVideoSampler(Inference): + + def __init__( + self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + use_cpu_offload=False, + device=0, + logger=None, + parallel_args=None, + ): + super().__init__( + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=text_encoder_2, + pipeline=pipeline, + use_cpu_offload=use_cpu_offload, + device=device, + logger=logger, + parallel_args=parallel_args, + ) + + self.pipeline = self.load_diffusion_pipeline( + args=args, + vae=self.vae, + text_encoder=self.text_encoder, + text_encoder_2=self.text_encoder_2, + model=self.model, + device=self.device, + ) + + self.default_negative_prompt = NEGATIVE_PROMPT + + def load_diffusion_pipeline( + self, + args, + vae, + text_encoder, + text_encoder_2, + model, + scheduler=None, + device=None, + progress_bar_config=None, + data_type="video", + ): + print("Loading diffusion pipeline...") + # print('vae', vae) + print('text_encoder', text_encoder) + print('text_encoder_2', text_encoder_2) + print('model', model) + """Load the denoising scheduler for inference.""" + if scheduler is None: + if args.denoise_type == "flow": + scheduler = FlowMatchDiscreteScheduler( + shift=args.flow_shift, + reverse=args.flow_reverse, + solver=args.flow_solver, + ) + else: + raise ValueError(f"Invalid denoise type {args.denoise_type}") + + pipeline = HunyuanVideoPipeline( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + transformer=model, + scheduler=scheduler, + progress_bar_config=progress_bar_config, + args=args, + ) + if self.use_cpu_offload: + pipeline.enable_sequential_cpu_offload() + else: + pipeline = pipeline.to(device) + + return pipeline + + @torch.no_grad() + def predict( + self, + prompt, + height=192, + width=336, + video_length=129, + seed=None, + negative_prompt=None, + infer_steps=50, + guidance_scale=6, + flow_shift=5.0, + embedded_guidance_scale=None, + batch_size=1, + num_videos_per_prompt=1, + mask_strategy=None, + **kwargs, + ): + """ + Predict the image/video from the given text. + + Args: + prompt (str or List[str]): The input text. + kwargs: + height (int): The height of the output video. Default is 192. + width (int): The width of the output video. Default is 336. + video_length (int): The frame number of the output video. Default is 129. + seed (int or List[str]): The random seed for the generation. Default is a random integer. + negative_prompt (str or List[str]): The negative text prompt. Default is an empty string. + guidance_scale (float): The guidance scale for the generation. Default is 6.0. + 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. + """ + + out_dict = dict() + + # ======================================================================== + # Arguments: seed + # ======================================================================== + if isinstance(seed, torch.Tensor): + seed = seed.tolist() + if seed is None: + seeds = [random.randint(0, 1_000_000) for _ in range(batch_size * num_videos_per_prompt)] + elif isinstance(seed, int): + seeds = [seed + i for _ in range(batch_size) for i in range(num_videos_per_prompt)] + elif isinstance(seed, (list, tuple)): + if len(seed) == batch_size: + seeds = [int(seed[i]) + j for i in range(batch_size) for j in range(num_videos_per_prompt)] + elif len(seed) == batch_size * num_videos_per_prompt: + seeds = [int(s) for s in seed] + else: + raise ValueError( + f"Length of seed must be equal to number of prompt(batch_size) or " + f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}.") + else: + raise ValueError(f"Seed must be an integer, a list of integers, or None, got {seed}.") + # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... + generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds] + out_dict["seeds"] = seeds + + # ======================================================================== + # Arguments: target_width, target_height, target_video_length + # ======================================================================== + if width <= 0 or height <= 0 or video_length <= 0: + raise ValueError( + f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}" + ) + if (video_length - 1) % 4 != 0: + raise ValueError(f"`video_length-1` must be a multiple of 4, got {video_length}") + + logger.info(f"Input (height, width, video_length) = ({height}, {width}, {video_length})") + + target_height = align_to(height, 16) + target_width = align_to(width, 16) + target_video_length = video_length + + out_dict["size"] = (target_height, target_width, target_video_length) + + # ======================================================================== + # Arguments: prompt, new_prompt, negative_prompt + # ======================================================================== + if not isinstance(prompt, str): + raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") + prompt = [prompt.strip()] + + # negative prompt + if negative_prompt is None or negative_prompt == "": + negative_prompt = self.default_negative_prompt + if not isinstance(negative_prompt, str): + raise TypeError(f"`negative_prompt` must be a string, but got {type(negative_prompt)}") + negative_prompt = [negative_prompt.strip()] + + # ======================================================================== + # Scheduler + # ======================================================================== + scheduler = FlowMatchDiscreteScheduler( + shift=flow_shift, + reverse=self.args.flow_reverse, + solver=self.args.flow_solver, + ) + self.pipeline.scheduler = scheduler + + if "884" in self.args.vae: + latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8] + elif "888" in self.args.vae: + latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + # ======================================================================== + # Print infer args + # ======================================================================== + debug_str = f""" + height: {target_height} + width: {target_width} + video_length: {target_video_length} + prompt: {prompt} + neg_prompt: {negative_prompt} + seed: {seed} + infer_steps: {infer_steps} + num_videos_per_prompt: {num_videos_per_prompt} + guidance_scale: {guidance_scale} + n_tokens: {n_tokens} + flow_shift: {flow_shift} + embedded_guidance_scale: {embedded_guidance_scale}""" + logger.debug(debug_str) + + # ======================================================================== + # Pipeline inference + # ======================================================================== + start_time = time.time() + samples = self.pipeline( + prompt=prompt, + height=target_height, + width=target_width, + video_length=target_video_length, + num_inference_steps=infer_steps, + guidance_scale=guidance_scale, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + output_type="pil", + n_tokens=n_tokens, + embedded_guidance_scale=embedded_guidance_scale, + data_type="video" if target_video_length > 1 else "image", + is_progress_bar=True, + vae_ver=self.args.vae, + enable_tiling=self.args.vae_tiling, + enable_vae_sp=self.args.vae_sp, + mask_strategy=mask_strategy, + )[0] + out_dict["samples"] = samples + out_dict["prompts"] = prompt + + gen_time = time.time() - start_time + logger.info(f"Success, time: {gen_time}") + + return out_dict diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py new file mode 100644 index 000000000..c85b51a4c --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/__init__.py @@ -0,0 +1,25 @@ +from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer + + +def load_model(args, in_channels, out_channels, factor_kwargs): + """load hunyuan video model + + Args: + args (dict): model args + in_channels (int): input channels number + out_channels (int): output channels number + factor_kwargs (dict): factor kwargs + + Returns: + model (nn.Module): The hunyuan video model + """ + if args.model in HUNYUAN_VIDEO_CONFIG.keys(): + model = HYVideoDiffusionTransformer( + in_channels=in_channels, + out_channels=out_channels, + **HUNYUAN_VIDEO_CONFIG[args.model], + **factor_kwargs, + ) + return model + else: + raise NotImplementedError() diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/activation_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/activation_layers.py new file mode 100644 index 000000000..f8774c26c --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/activation_layers.py @@ -0,0 +1,23 @@ +import torch.nn as nn + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py new file mode 100644 index 000000000..0567b08aa --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/attenion.py @@ -0,0 +1,128 @@ +import torch +import torch.nn.functional as F +from einops import rearrange + +try: + from st_attn import sliding_tile_attention +except ImportError: + print("Could not load Sliding Tile Attention.") + sliding_tile_attention = None + +from fastvideo.v1.v0_reference_src.models.flash_attn_no_pad import flash_attn_no_pad +from flash_attn import flash_attn_func +from fastvideo.v1.v0_reference_src.utils.communications import all_gather, all_to_all_4D +from fastvideo.v1.v0_reference_src.utils.parallel_states import get_sequence_parallel_state, nccl_info + + +def attention( + q, + k, + v, + drop_rate=0, + attn_mask=None, + causal=False, +): + + qkv = torch.stack([q, k, v], dim=2) + + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.bool() + + x = flash_attn_no_pad(qkv, attn_mask, causal=causal, dropout_p=drop_rate, softmax_scale=None) + + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +def tile(x, sp_size): + x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80) + return rearrange(x, + "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", + n_t=5, + n_h=6, + n_w=10, + ts_t=6, + ts_h=8, + ts_w=8) + + +def untile(x, sp_size): + x = rearrange(x, + "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", + n_t=5, + n_h=6, + n_w=10, + ts_t=6, + ts_h=8, + ts_w=8) + return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80) + + +def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None): + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + text_length = text_mask.sum() + + if get_sequence_parallel_state(): + # 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) + # [b, s, h, d] + + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + if mask_strategy[0] 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], dim=1).transpose(1, 2) + value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2) + + head_num = query.size(1) + current_rank = nccl_info.rank_within_group + start_head = current_rank * head_num + windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)] + + hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2) + else: + txt_length = text_mask.sum().int().item() + query = torch.cat([query, encoder_query[:, :txt_length]], dim=1) + key = torch.cat([key, encoder_key[:, :txt_length]], dim=1) + value = torch.cat([value, encoder_value[:, :txt_length]], dim=1) + # B, S, 3, H, D + # qkv = torch.stack([query, key, value], dim=2) + + attn_mask = F.pad(text_mask, (sequence_length, 0), value=True) + # hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) + hidden_states = flash_attn_func(query, key, value, causal=False, dropout_p=0.0, softmax_scale=None) + hidden_states = torch.cat([hidden_states, encoder_query[:, txt_length:, :]], dim=1) + + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length), + dim=1) + + if mask_strategy[0] is not None: + hidden_states = untile(hidden_states, nccl_info.sp_size) + + if get_sequence_parallel_state(): + hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous() + + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + + attn = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, a, d = attn.shape + attn = attn.reshape(b, s, -1) + + return attn diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/embed_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/embed_layers.py new file mode 100644 index 000000000..38aad89cb --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/embed_layers.py @@ -0,0 +1,150 @@ +import math + +import torch +import torch.nn as nn + +from ..utils.helpers import to_2tuple + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + **factory_kwargs, + ) + nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1)) + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / + half).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/mlp_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/mlp_layers.py new file mode 100644 index 000000000..921b85252 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/mlp_layers.py @@ -0,0 +1,107 @@ +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from ..utils.helpers import to_2tuple +from .modulate_layers import modulate + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_channels, hidden_channels, bias=bias[0], **factory_kwargs) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = (norm_layer(hidden_channels, **factory_kwargs) if norm_layer is not None else nn.Identity()) + self.fc2 = linear_layer(hidden_channels, out_features, bias=bias[1], **factory_kwargs) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +# +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + if isinstance(patch_size, int): + self.linear = nn.Linear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + **factory_kwargs, + ) + else: + self.linear = nn.Linear( + hidden_size, + patch_size[0] * patch_size[1] * patch_size[2] * out_channels, + bias=True, + ) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py new file mode 100644 index 000000000..7f4e0d2ad --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/models.py @@ -0,0 +1,657 @@ +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin +from einops import rearrange + +from fastvideo.v1.v0_reference_src.models.hunyuan.modules.posemb_layers import get_nd_rotary_pos_embed +from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info + +from .activation_layers import get_activation_layer +from .attenion import parallel_attention +from .embed_layers import PatchEmbed, TextProjection, TimestepEmbedder +from .mlp_layers import MLP, FinalLayer, MLPEmbedder +from .modulate_layers import ModulateDiT, apply_gate, modulate +from .norm_layers import get_norm_layer +from .posemb_layers import apply_rotary_emb +from .token_refiner import SingleTokenRefiner + + +class MMDoubleStreamBlock(nn.Module): + """ + A multimodal dit block with separate modulation for + text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206 + (Flux.1): https://github.com/black-forest-labs/flux + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_tanh", + qk_norm: bool = True, + qk_norm_type: str = "rms", + qkv_bias: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.img_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.img_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.img_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.img_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + + self.txt_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.txt_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + self.txt_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.txt_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.txt_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + self.txt_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple = None, + text_mask: torch.Tensor = None, + mask_strategy=None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec).chunk(6, dim=-1) + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec).chunk(6, dim=-1) + + # Prepare image for attention. + img_modulated = self.img_norm1(img) + img_modulated = modulate(img_modulated, shift=img_mod1_shift, scale=img_mod1_scale) + + img_qkv = self.img_attn_qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # Apply QK-Norm if needed + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + # Apply RoPE if needed. + if freqs_cis is not None: + + 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) + + freqs_cis = ( + shrink_head(freqs_cis[0], dim=0), + shrink_head(freqs_cis[1], dim=0), + ) + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + + # Prepare txt for attention. + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate(txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale) + txt_qkv = self.txt_attn_qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # Apply QK-Norm if needed. + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + mask_strategy=mask_strategy, + ) + + # attention computation end + img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:] + # Calculate the img blocks. + img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate) + img = img + apply_gate( + self.img_mlp(modulate(self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale)), + gate=img_mod2_gate, + ) + + # Calculate the txt blocks. + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate) + txt = txt + apply_gate( + self.txt_mlp(modulate(self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale)), + gate=txt_mod2_gate, + ) + return img, txt + + +class MMSingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + Also refer to (SD3): https://arxiv.org/abs/2403.03206 + (Flux.1): https://github.com/black-forest-labs/flux + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim**-0.5 + + # qkv and mlp_in + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs) + # proj and mlp_out + self.linear2 = nn.Linear(hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs) + + self.mlp_act = get_activation_layer(mlp_act_type)() + self.modulation = ModulateDiT( + hidden_size, + factor=3, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + text_mask: torch.Tensor = 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) + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + 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) + + freqs_cis = ( + shrink_head(freqs_cis[0], dim=0), + shrink_head(freqs_cis[1], dim=0), + ) + + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False) + assert (img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + mask_strategy=mask_strategy, + ) + + # attention computation end + + # Compute activation in mlp stream, cat again and run second linear layer. + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + apply_gate(output, gate=mod_gate) + + +class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): + """ + HunyuanVideo Transformer backbone + + Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline. + + Reference: + [1] Flux.1: https://github.com/black-forest-labs/flux + [2] MMDiT: http://arxiv.org/abs/2403.03206 + + Parameters + ---------- + args: argparse.Namespace + The arguments parsed by argparse. + patch_size: list + The size of the patch. + in_channels: int + The number of input channels. + out_channels: int + The number of output channels. + hidden_size: int + The hidden size of the transformer backbone. + heads_num: int + The number of attention heads. + mlp_width_ratio: float + The ratio of the hidden size of the MLP in the transformer block. + mlp_act_type: str + The activation function of the MLP in the transformer block. + depth_double_blocks: int + The number of transformer blocks in the double blocks. + depth_single_blocks: int + The number of transformer blocks in the single blocks. + rope_dim_list: list + The dimension of the rotary embedding for t, h, w. + qkv_bias: bool + Whether to use bias in the qkv linear layer. + qk_norm: bool + Whether to use qk norm. + qk_norm_type: str + The type of qk norm. + guidance_embed: bool + Whether to use guidance embedding for distillation. + text_projection: str + The type of the text projection, default is single_refiner. + use_attention_mask: bool + Whether to use attention mask for text encoder. + dtype: torch.dtype + The dtype of the model. + device: torch.device + The device of the model. + """ + + @register_to_config + def __init__( + self, + patch_size: list = [1, 2, 2], + in_channels: int = 4, # Should be VAE.config.latent_channels. + out_channels: int = None, + hidden_size: int = 3072, + heads_num: int = 24, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + mm_double_blocks_depth: int = 20, + mm_single_blocks_depth: int = 40, + rope_dim_list: List[int] = [16, 56, 56], + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = "rms", + guidance_embed: bool = False, # For modulation. + text_projection: str = "single_refiner", + use_attention_mask: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, + rope_theta: int = 256, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.unpatchify_channels = self.out_channels + self.guidance_embed = guidance_embed + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + # Text projection. Default to linear projection. + # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831 + self.use_attention_mask = use_attention_mask + self.text_projection = text_projection + + if hidden_size % heads_num != 0: + raise ValueError(f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}") + pe_dim = hidden_size // heads_num + if sum(rope_dim_list) != pe_dim: + raise ValueError(f"Got {rope_dim_list} but expected positional dim {pe_dim}") + self.hidden_size = hidden_size + self.heads_num = heads_num + + # image projection + self.img_in = PatchEmbed(self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs) + + # text projection + if self.text_projection == "linear": + self.txt_in = TextProjection( + self.config.text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs, + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + self.config.text_states_dim, + hidden_size, + heads_num, + depth=2, + **factory_kwargs, + ) + else: + raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + + # time modulation + self.time_in = TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs) + + # text modulation + self.vector_in = MLPEmbedder(self.config.text_states_dim_2, self.hidden_size, **factory_kwargs) + + # guidance modulation + self.guidance_in = (TimestepEmbedder(self.hidden_size, get_activation_layer("silu"), **factory_kwargs) + if guidance_embed else None) + + # double blocks + self.double_blocks = nn.ModuleList([ + MMDoubleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(mm_double_blocks_depth) + ]) + + # single blocks + self.single_blocks = nn.ModuleList([ + MMSingleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs, + ) for _ in range(mm_single_blocks_depth) + ]) + + self.final_layer = FinalLayer( + self.hidden_size, + self.patch_size, + self.out_channels, + get_activation_layer("silu"), + **factory_kwargs, + ) + + def enable_deterministic(self): + for block in self.double_blocks: + block.enable_deterministic() + for block in self.single_blocks: + block.enable_deterministic() + + def disable_deterministic(self): + for block in self.double_blocks: + block.disable_deterministic() + for block in self.single_blocks: + block.disable_deterministic() + + def get_rotary_pos_embed(self, rope_sizes): + target_ndim = 3 + + head_dim = self.hidden_size // self.heads_num + rope_dim_list = self.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)] + assert (sum(rope_dim_list) == head_dim), "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=self.rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + # x: torch.Tensor, + # t: torch.Tensor, # Should be in range(0, 1000). + # text_states: torch.Tensor = None, + # text_mask: torch.Tensor = None, # Now we don't use it. + # text_states_2: Optional[torch.Tensor] = None, # Text embedding for modulation. + # guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. + # return_dict: bool = True, + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_attention_mask: torch.Tensor, + mask_strategy=None, + output_features=False, + output_features_stride=8, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + guidance=None, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if guidance is None: + guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16) + if mask_strategy is None: + mask_strategy = [[None] * self.heads_num + for _ in range(len(self.double_blocks) + len(self.single_blocks))] + img = x = hidden_states + text_mask = encoder_attention_mask + t = timestep + txt = encoder_hidden_states[:, 1:] + text_states_2 = encoder_hidden_states[:, 0, :self.config.text_states_dim_2] + _, _, ot, oh, ow = x.shape # codespell:ignore + tt, th, tw = ( + ot // self.patch_size[0], # codespell:ignore + oh // self.patch_size[1], # codespell:ignore + ow // self.patch_size[2], # codespell:ignore + ) + original_tt = nccl_info.sp_size * tt + freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw)) + # Prepare modulation vectors. + vec = self.time_in(t) + + # text modulation + vec = vec + self.vector_in(text_states_2) + # guidance modulation + if self.guidance_embed: + if guidance is None: + raise ValueError("Didn't get guidance strength for guidance distilled model.") + + # our timestep_embedding is merged into guidance_in(TimestepEmbedder) + vec = vec + self.guidance_in(guidance) + # Embed image and text. + img = self.img_in(img) + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None) + else: + raise NotImplementedError(f"Unsupported text_projection: {self.text_projection}") + txt_seq_len = txt.shape[1] + img_seq_len = img.shape[1] + + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + # --------------------- Pass through DiT blocks ------------------------ + for index, block in enumerate(self.double_blocks): + double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]] + img, txt = block(*double_block_args) + # 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): + single_block_args = [ + x, + vec, + txt_seq_len, + (freqs_cos, freqs_sin), + text_mask, + mask_strategy[index + len(self.double_blocks)], + ] + x = block(*single_block_args) + if output_features and _ % output_features_stride == 0: + features_list.append(x[:, :img_seq_len, ...]) + img = x[:, :img_seq_len, ...] + # ---------------------------- Final layer ------------------------------ + img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels) + img = self.unpatchify(img, tt, th, tw) + assert not return_dict, "return_dict is not supported." + if output_features: + features_list = torch.stack(features_list, dim=0) + else: + features_list = None + return (img, features_list) + + def unpatchify(self, x, t, h, w): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.unpatchify_channels + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs + + def params_count(self): + counts = { + "double": + sum([ + sum(p.numel() + for p in block.img_attn_qkv.parameters()) + sum(p.numel() + for p in block.img_attn_proj.parameters()) + + sum(p.numel() for p in block.img_mlp.parameters()) + sum(p.numel() + for p in block.txt_attn_qkv.parameters()) + + sum(p.numel() for p in block.txt_attn_proj.parameters()) + sum(p.numel() + for p in block.txt_mlp.parameters()) + for block in self.double_blocks + ]), + "single": + sum([ + sum(p.numel() for p in block.linear1.parameters()) + sum(p.numel() for p in block.linear2.parameters()) + for block in self.single_blocks + ]), + "total": + sum(p.numel() for p in self.parameters()), + } + counts["attn+mlp"] = counts["double"] + counts["single"] + return counts + + +################################################################################# +# HunyuanVideo Configs # +################################################################################# + +HUNYUAN_VIDEO_CONFIG = { + "HYVideo-T/2": { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + }, + "HYVideo-T/2-cfgdistill": { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, + }, +} diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/modulate_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/modulate_layers.py new file mode 100644 index 000000000..c02ef3187 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/modulate_layers.py @@ -0,0 +1,152 @@ +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) + elif scale is None: + return x + shift.unsqueeze(1) + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + if tanh: + return x * gate.unsqueeze(1).tanh() + else: + return x * gate.unsqueeze(1) + + +def ckpt_wrapper(module): + + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + + +class RMSNorm(nn.Module): + + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/norm_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/norm_layers.py new file mode 100644 index 000000000..5bc997059 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/norm_layers.py @@ -0,0 +1,78 @@ +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/posemb_layers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/posemb_layers.py new file mode 100644 index 000000000..314aa126d --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/posemb_layers.py @@ -0,0 +1,290 @@ +from typing import List, Tuple, Union + +import torch + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x, ) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0, ) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + + +def reshape_for_broadcast( + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + x: torch.Tensor, + head_first=False, +): + """ + Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + else: + assert freqs_cis[0].shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [d if i == ndim - 2 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + else: + assert freqs_cis.shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = (x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + use_real=False, + theta_rescale_factor: Union[float, List[float]] = 1.0, + interpolation_factor: Union[float, List[float]] = 1.0, +): + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor**(dim / (dim - 2)) + + freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) # [D/2] + # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}" + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2] + return freqs_cis diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/modules/token_refiner.py b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/token_refiner.py new file mode 100644 index 000000000..1bc127fea --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/modules/token_refiner.py @@ -0,0 +1,199 @@ +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from .activation_layers import get_activation_layer +from .attenion import attention +from .embed_layers import TextProjection, TimestepEmbedder +from .mlp_layers import MLP +from .modulate_layers import apply_gate +from .norm_layers import get_norm_layer + + +class IndividualTokenRefinerBlock(nn.Module): + + def __init__( + self, + hidden_size, + heads_num, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + self.self_attn_qkv = nn.Linear(hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.self_attn_k_norm = (qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.self_attn_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs) + + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: torch.Tensor = None, + ): + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) + # Apply QK-Norm if needed + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + + # Self-Attention + attn = attention(q, k, v, attn_mask=attn_mask) + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + + # FFN Layer + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + + return x + + +class IndividualTokenRefiner(nn.Module): + + def __init__( + self, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.blocks = nn.ModuleList([ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(depth) + ]) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ): + mask = mask.clone().bool() + # avoid attention weight become NaN + mask[:, 0] = True + for block in self.blocks: + x = block(x, c, mask) + return x + + +class SingleTokenRefiner(nn.Module): + """ + A single token refiner block for llm text embedding refine. + """ + + def __init__( + self, + in_channels, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + attn_mode: str = "torch", + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.attn_mode = attn_mode + assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner." + + self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True, **factory_kwargs) + + act_layer = get_activation_layer(act_type) + # Build timestep embedding layer + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs) + # Build context embedding layer + self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, **factory_kwargs) + + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ): + timestep_aware_representations = self.t_embedder(t) + + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [b, s1, 1] + context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder(context_aware_representations.to(x.dtype)) + c = timestep_aware_representations + context_aware_representations + + x = self.input_embedder(x) + x = self.individual_token_refiner(x, c, mask) + return x diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/prompt_rewrite.py b/fastvideo/v1/v0_reference_src/models/hunyuan/prompt_rewrite.py new file mode 100644 index 000000000..556b8aaf2 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/prompt_rewrite.py @@ -0,0 +1,52 @@ +normal_mode_prompt = """Normal mode - Video Recaption Task: + +You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description. + +0. Preserve ALL information, including style words and technical terms. + +1. If the input is in Chinese, translate the entire description to English. + +2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences. + +3. If the input does not include style, lighting, atmosphere, you can make reasonable associations. + +4. Output ALL must be in English. + +Given Input: +input: "{input}" +""" + +master_mode_prompt = """Master mode - Video Recaption Task: + +You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description. + +0. Preserve ALL information, including style words and technical terms. + +1. If the input is in Chinese, translate the entire description to English. + +2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences. + +3. If the input does not include style, lighting, atmosphere, you can make reasonable associations. + +4. Output ALL must be in English. + +Given Input: +input: "{input}" +""" + + +def get_rewrite_prompt(ori_prompt, mode="Normal"): + if mode == "Normal": + prompt = normal_mode_prompt.format(input=ori_prompt) + elif mode == "Master": + prompt = master_mode_prompt.format(input=ori_prompt) + else: + raise Exception("Only supports Normal and Normal", mode) + return prompt + + +ori_prompt = "一只小狗在草地上奔跑。" +normal_prompt = get_rewrite_prompt(ori_prompt, mode="Normal") +master_prompt = get_rewrite_prompt(ori_prompt, mode="Master") + +# Then you can use the normal_prompt or master_prompt to access the hunyuan-large rewrite model to get the final prompt. diff --git a/fastvideo/models/encoders/encoder.py b/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py similarity index 99% rename from fastvideo/models/encoders/encoder.py rename to fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py index aadfaa87f..bf6f8023c 100644 --- a/fastvideo/models/encoders/encoder.py +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/text_encoder/__init__.py @@ -50,8 +50,8 @@ def load_text_encoder( def load_tokenizer(tokenizer_type, tokenizer_path=None, padding_side="right", logger=None): - if tokenizer_path is None: - tokenizer_path = TOKENIZER_PATH[tokenizer_type] + # if tokenizer_path is None: + tokenizer_path = TOKENIZER_PATH[tokenizer_type] if logger is not None: logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}") diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/utils/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/utils/data_utils.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/data_utils.py new file mode 100644 index 000000000..524118179 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/data_utils.py @@ -0,0 +1,14 @@ +import math + + +def align_to(value, alignment): + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/utils/file_utils.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/file_utils.py new file mode 100644 index 000000000..82e244f03 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/file_utils.py @@ -0,0 +1,71 @@ +import os +from pathlib import Path + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +CODE_SUFFIXES = { + ".py", # Python codes + ".sh", # Shell scripts + ".yaml", + ".yml", # Configuration files +} + + +def safe_dir(path): + """ + Create a directory (or the parent directory of a file) if it does not exist. + + Args: + path (str or Path): Path to the directory. + + Returns: + path (Path): Path object of the directory. + """ + path = Path(path) + path.mkdir(exist_ok=True, parents=True) + return path + + +def safe_file(path): + """ + Create the parent directory of a file if it does not exist. + + Args: + path (str or Path): Path to the file. + + Returns: + path (Path): Path object of the file. + """ + path = Path(path) + path.parent.mkdir(exist_ok=True, parents=True) + return path + + +def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, fps=24): + """save videos by video tensor + copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61 + + Args: + videos (torch.Tensor): video tensor predicted by the model + path (str): path to save video + rescale (bool, optional): rescale the video tensor from [-1, 1] to . Defaults to False. + n_rows (int, optional): Defaults to 1. + fps (int, optional): video save fps. Defaults to 8. + """ + videos = rearrange(videos, "b c t h w -> t b c h w") + outputs = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = torch.clamp(x, 0, 1) + x = (x * 255).numpy().astype(np.uint8) + outputs.append(x) + + os.makedirs(os.path.dirname(path), exist_ok=True) + imageio.mimsave(path, outputs, fps=fps) diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/utils/helpers.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/helpers.py new file mode 100644 index 000000000..7b5b8731f --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/helpers.py @@ -0,0 +1,41 @@ +import collections.abc +from itertools import repeat + + +def _ntuple(n): + + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + + return parse + + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + + +def as_tuple(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + if x is None or isinstance(x, (int, float, str)): + return (x, ) + else: + raise ValueError(f"Unknown type {type(x)}") + + +def as_list_of_2tuple(x): + x = as_tuple(x) + if len(x) == 1: + x = (x[0], x[0]) + assert len(x) % 2 == 0, f"Expect even length, got {len(x)}." + lst = [] + for i in range(0, len(x), 2): + lst.append((x[i], x[i + 1])) + return lst diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py new file mode 100644 index 000000000..1a6f46c15 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py @@ -0,0 +1,41 @@ +import argparse + +import torch +from transformers import AutoProcessor, LlavaForConditionalGeneration + + +def preprocess_text_encoder_tokenizer(args): + + processor = AutoProcessor.from_pretrained(args.input_dir) + model = LlavaForConditionalGeneration.from_pretrained( + args.input_dir, + torch_dtype=torch.float16, + low_cpu_mem_usage=True, + ).to(0) + + model.language_model.save_pretrained(f"{args.output_dir}") + processor.tokenizer.save_pretrained(f"{args.output_dir}") + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="The path to the llava-llama-3-8b-v1_1-transformers.", + ) + parser.add_argument( + "--output_dir", + type=str, + default="", + help="The output path of the llava-llama-3-8b-text-encoder-tokenizer." + "if '', the parent dir of output will be the same as input dir.", + ) + args = parser.parse_args() + + if len(args.output_dir) == 0: + args.output_dir = "/".join(args.input_dir.split("/")[:-1]) + + preprocess_text_encoder_tokenizer(args) diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py new file mode 100644 index 000000000..aa3de3c84 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/__init__.py @@ -0,0 +1,64 @@ +from pathlib import Path + +import torch + +from ..constants import PRECISION_TO_TYPE, VAE_PATH +from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D + + +def load_vae( + vae_type: str = "884-16c-hy", + vae_precision: str = None, + sample_size: tuple = None, + vae_path: str = None, + logger=None, + device=None, +): + """the function to load the 3D VAE model + + Args: + vae_type (str): the type of the 3D VAE model. Defaults to "884-16c-hy". + vae_precision (str, optional): the precision to load vae. Defaults to None. + sample_size (tuple, optional): the tiling size. Defaults to None. + vae_path (str, optional): the path to vae. Defaults to None. + logger (_type_, optional): logger. Defaults to None. + device (_type_, optional): device to load vae. Defaults to None. + """ + if vae_path is None: + vae_path = VAE_PATH[vae_type] + + if logger is not None: + logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}") + config = AutoencoderKLCausal3D.load_config(vae_path) + if sample_size: + vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size) + else: + vae = AutoencoderKLCausal3D.from_config(config) + + vae_ckpt = Path("/home/ubuntu/src/FastVideo/data/hunyuan/hunyuan-video-t2v-720p/vae/pytorch_model.pt") + assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}" + + ckpt = torch.load(vae_ckpt, map_location=vae.device) + if "state_dict" in ckpt: + ckpt = ckpt["state_dict"] + if any(k.startswith("vae.") for k in ckpt.keys()): + ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} + vae.load_state_dict(ckpt) + + spatial_compression_ratio = vae.config.spatial_compression_ratio + time_compression_ratio = vae.config.time_compression_ratio + + if vae_precision is not None: + vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision]) + + vae.requires_grad_(False) + + if logger is not None: + logger.info(f"VAE to dtype: {vae.dtype}") + + if device is not None: + vae = vae.to(device) + + vae.eval() + + return vae, vae_path, spatial_compression_ratio, time_compression_ratio diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py new file mode 100644 index 000000000..b0973853b --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/autoencoder_kl_causal_3d.py @@ -0,0 +1,776 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +from dataclasses import dataclass +from math import prod +from typing import Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.distributed as dist +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info +from fastvideo.v1.distributed.parallel_state import get_sp_group +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) + +try: + # This diffusers is modified and packed in the mirror. + from diffusers.loaders import FromOriginalVAEMixin +except ImportError: + # Use this to be compatible with the original diffusers. + from diffusers.loaders.single_file_model import ( + FromOriginalModelMixin as FromOriginalVAEMixin, ) + +from diffusers.models.attention_processor import (ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention, + AttentionProcessor, AttnAddedKVProcessor, AttnProcessor) +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils.accelerate_utils import apply_forward_hook + +from .vae import BaseOutput, DecoderCausal3D, DecoderOutput, DiagonalGaussianDistribution, EncoderCausal3D + + +@dataclass +class DecoderOutput2(BaseOutput): + sample: torch.FloatTensor + posterior: Optional[DiagonalGaussianDistribution] = None + + +class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): + r""" + A VAE model with KL loss for encoding images/videos into latents and decoding latent representations into images/videos. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D", ), + up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D", ), + block_out_channels: Tuple[int] = (64, ), + layers_per_block: int = 1, + act_fn: str = "silu", + latent_channels: int = 4, + norm_num_groups: int = 32, + sample_size: int = 32, + sample_tsize: int = 64, + scaling_factor: float = 0.18215, + force_upcast: float = True, + spatial_compression_ratio: int = 8, + time_compression_ratio: int = 4, + mid_block_add_attention: bool = True, + ): + super().__init__() + + self.time_compression_ratio = time_compression_ratio + + self.encoder = EncoderCausal3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + double_z=True, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.decoder = DecoderCausal3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + self.use_slicing = False + self.use_spatial_tiling = False + self.use_temporal_tiling = False + self.use_parallel = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // time_compression_ratio + + self.tile_sample_min_size = self.config.sample_size + sample_size = (self.config.sample_size[0] if isinstance(self.config.sample_size, + (list, tuple)) else self.config.sample_size) + self.tile_latent_min_size = int(sample_size / (2**(len(self.config.block_out_channels) - 1))) + self.tile_overlap_factor = 0.25 + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (EncoderCausal3D, DecoderCausal3D)): + module.gradient_checkpointing = value + + def enable_temporal_tiling(self, use_tiling: bool = True): + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger videos. + """ + self.enable_spatial_tiling(use_tiling) + self.enable_temporal_tiling(use_tiling) + + def disable_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.disable_spatial_tiling() + self.disable_temporal_tiling() + + def enable_parallel(self): + r""" + Enable sequence parallelism for the model. This will allow the vae to decode (with tiling) in parallel. + """ + self.use_parallel = True + + def enable_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.use_slicing = True + + def disable_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor( + self, + processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], + _remove_lora=False, + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes.") + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor(processor.pop(f"{name}.processor"), _remove_lora=_remove_lora) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + @apply_forward_hook + def encode(self, + x: torch.FloatTensor, + return_dict: bool = True) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + """ + Encode a batch of images/videos into latents. + + Args: + x (`torch.FloatTensor`): Input batch of images/videos. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded images/videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + assert len(x.shape) == 5, "The input tensor should have 5 dimensions." + + if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x, return_dict=return_dict) + + if self.use_spatial_tiling and (x.shape[-1] > self.tile_sample_min_size + or x.shape[-2] > self.tile_sample_min_size): + return self.spatial_tiled_encode(x, return_dict=return_dict) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self.encoder(x) + + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + assert len(z.shape) == 5, "The input tensor should have 5 dimensions." + + logger.info(f"use_parallel: {self.use_parallel}") + if self.use_parallel: + return self.parallel_tiled_decode(z, return_dict=return_dict) + + if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z, return_dict=return_dict) + + if self.use_spatial_tiling and (z.shape[-1] > self.tile_latent_min_size + or z.shape[-2] > self.tile_latent_min_size): + return self.spatial_tiled_decode(z, return_dict=return_dict) + + z = self.post_quant_conv(z) + dec = self.decoder(z) + + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode(self, + z: torch.FloatTensor, + return_dict: bool = True, + generator=None) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Decode a batch of images/videos. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded, ) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, + y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, + x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, + x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, + x, :, :] * (x / blend_extent) + return b + + def spatial_tiled_encode( + self, + x: torch.FloatTensor, + return_dict: bool = True, + return_moments: bool = False, + ) -> AutoencoderKLOutput: + r"""Encode a batch of images/videos using a tiled encoder. + + When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several + steps. This is useful to keep memory use constant regardless of image/videos size. The end result of tiled encoding is + different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the + tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the + output, but they should be much less noticeable. + + Args: + x (`torch.FloatTensor`): Input batch of images/videos. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split video into tiles and encode them separately. + rows = [] + for i in range(0, x.shape[-2], overlap_size): + row = [] + for j in range(0, x.shape[-1], overlap_size): + tile = x[:, :, :, i:i + self.tile_sample_min_size, j:j + self.tile_sample_min_size, ] + tile = self.encoder(tile) + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + moments = torch.cat(result_rows, dim=-2) + if return_moments: + return moments + + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def spatial_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + r""" + Decode a batch of images/videos using a tiled decoder. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[-2], overlap_size): + row = [] + for j in range(0, z.shape[-1], overlap_size): + tile = z[:, :, :, i:i + self.tile_latent_min_size, j:j + self.tile_latent_min_size, ] + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + def temporal_tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput: + + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_latent_min_tsize - blend_extent + + # Split the video into tiles and encode them separately. + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i:i + self.tile_sample_min_tsize + 1, :, :] + if self.use_spatial_tiling and (tile.shape[-1] > self.tile_sample_min_size + or tile.shape[-2] > self.tile_sample_min_size): + tile = self.spatial_tiled_encode(tile, return_moments=True) + else: + tile = self.encoder(tile) + tile = self.quant_conv(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + + moments = torch.cat(result_row, dim=2) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def temporal_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + # Split z into overlapping tiles and decode them separately. + + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - blend_extent + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :] + if self.use_spatial_tiling and (tile.shape[-1] > self.tile_latent_min_size + or tile.shape[-2] > self.tile_latent_min_size): + decoded = self.spatial_tiled_decode(tile, return_dict=True).sample + else: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + + dec = torch.cat(result_row, dim=2) + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + def _parallel_data_generator(self, gathered_results, gathered_dim_metadata): + global_idx = 0 + for i, per_rank_metadata in enumerate(gathered_dim_metadata): + _start_shape = 0 + for shape in per_rank_metadata: + mul_shape = prod(shape) + yield (gathered_results[i, _start_shape:_start_shape + mul_shape].reshape(shape), global_idx) + _start_shape += mul_shape + global_idx += 1 + + def parallel_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs + """ + # world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + + sp_group = get_sp_group() + world_size = sp_group.world_size + rank = sp_group.rank_in_group + + B, C, T, H, W = z.shape + + # Calculate parameters + t_overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) + t_blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - t_blend_extent + + s_overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + s_blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + s_row_limit = self.tile_sample_min_size - s_blend_extent + + # Calculate tile dimensions + num_t_tiles = (T + t_overlap_size - 1) // t_overlap_size + num_h_tiles = (H + s_overlap_size - 1) // s_overlap_size + num_w_tiles = (W + s_overlap_size - 1) // s_overlap_size + total_spatial_tiles = num_h_tiles * num_w_tiles + total_tiles = num_t_tiles * total_spatial_tiles + + # Calculate tiles per rank and padding + tiles_per_rank = (total_tiles + world_size - 1) // world_size + start_tile_idx = rank * tiles_per_rank + end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) + + local_results = [] + local_dim_metadata = [] + # Process assigned tiles + for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)): + # Convert flat index to 3D indices + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + + # Calculate positions + t_start = t_idx * t_overlap_size + h_start = h_idx * s_overlap_size + w_start = w_idx * s_overlap_size + + # Extract and process tile + tile = z[:, :, t_start:t_start + self.tile_latent_min_tsize + 1, + h_start:h_start + self.tile_latent_min_size, w_start:w_start + self.tile_latent_min_size] + + # Process tile + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + + if t_start > 0: + decoded = decoded[:, :, 1:, :, :] + + # Store metadata + shape = decoded.shape + # Store decoded data (flattened) + decoded_flat = decoded.reshape(-1) + local_results.append(decoded_flat) + local_dim_metadata.append(shape) + + results = torch.cat(local_results, dim=0).contiguous() + del local_results + torch.cuda.empty_cache() + # first gather size to pad the results + local_size = torch.tensor([results.size(0)], device=results.device, dtype=torch.int64) + all_sizes = [torch.zeros(1, device=results.device, dtype=torch.int64) for _ in range(world_size)] + dist.all_gather(all_sizes, local_size) + max_size = max(size.item() for size in all_sizes) + padded_results = torch.zeros(max_size, device=results.device) + padded_results[:results.size(0)] = results + del results + torch.cuda.empty_cache() + # Gather all results + gathered_dim_metadata = [None] * world_size + gathered_results = torch.zeros_like(padded_results).repeat( + world_size, *[1] * len(padded_results.shape)).contiguous( + ) # use contiguous to make sure it won't copy data in the following operations + # gathered_results = sp_group.all_gather(padded_results) + # gathered_dim_metadata = sp_group.broadcast_object_list(gathered_dim_metadata, local_dim_metadata) + dist.all_gather_into_tensor(gathered_results, padded_results) + dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) + # Process gathered results + data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator(gathered_results, gathered_dim_metadata): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + data[t_idx][h_idx][w_idx] = current_data + # Merge results + result_slices = [] + last_slice_data = None + for i, tem_data in enumerate(data): + slice_data = self._merge_spatial_tiles(tem_data, s_blend_extent, s_row_limit) + if i > 0: + slice_data = self.blend_t(last_slice_data, slice_data, t_blend_extent) + result_slices.append(slice_data[:, :, :t_limit, :, :]) + else: + result_slices.append(slice_data[:, :, :t_limit + 1, :, :]) + last_slice_data = slice_data + dec = torch.cat(result_slices, dim=2) + + if not return_dict: + return (dec, ) + return DecoderOutput(sample=dec) + + def _merge_spatial_tiles(self, spatial_rows, blend_extent, row_limit): + """Helper function to merge spatial tiles with blending""" + result_rows = [] + for i, row in enumerate(spatial_rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(spatial_rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def forward( + self, + sample: torch.FloatTensor, + sample_posterior: bool = False, + return_dict: bool = True, + return_posterior: bool = False, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput2, torch.FloatTensor]: + r""" + Args: + sample (`torch.FloatTensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z).sample + + if not return_dict: + if return_posterior: + return (dec, posterior) + else: + return (dec, ) + if return_posterior: + return DecoderOutput2(sample=dec, posterior=posterior) + else: + return DecoderOutput2(sample=dec) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections + def fuse_qkv_projections(self): + """ + Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, + key, value) are fused. For cross-attention modules, key and value projection matrices are fused. + + + + This API is 🧪 experimental. + + + """ + self.original_attn_processors = None + + for _, attn_processor in self.attn_processors.items(): + if "Added" in str(attn_processor.__class__.__name__): + raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") + + self.original_attn_processors = self.attn_processors + + for module in self.modules(): + if isinstance(module, Attention): + module.fuse_projections(fuse=True) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections + def unfuse_qkv_projections(self): + """Disables the fused QKV projection if enabled. + + + + This API is 🧪 experimental. + + + + """ + if self.original_attn_processors is not None: + self.set_attn_processor(self.original_attn_processors) diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py new file mode 100644 index 000000000..6b16b26bd --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/unet_causal_3d_blocks.py @@ -0,0 +1,760 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from typing import Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import Attention, SpatialNorm +from diffusers.models.normalization import AdaGroupNorm, RMSNorm +from diffusers.utils import logging +from einops import rearrange +from torch import nn + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def prepare_causal_attention_mask(n_frame: int, n_hw: int, dtype, device, batch_size: int = None): + seq_len = n_frame * n_hw + mask = torch.full((seq_len, seq_len), float("-inf"), dtype=dtype, device=device) + for i in range(seq_len): + i_frame = i // n_hw + mask[i, :(i_frame + 1) * n_hw] = 0 + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + + +class CausalConv3d(nn.Module): + """ + Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations. + This maintains temporal causality in video generation tasks. + """ + + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + dilation: Union[int, Tuple[int, int, int]] = 1, + pad_mode="replicate", + **kwargs, + ): + super().__init__() + + self.pad_mode = pad_mode + padding = ( + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size - 1, + 0, + ) # W, H, T + self.time_causal_padding = padding + + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + + +class UpsampleCausal3D(nn.Module): + """ + A 3D upsampling layer with an optional convolution. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + kernel_size: Optional[int] = None, + padding=1, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + interpolate=True, + upsample_factor=(2, 2, 2), + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + self.interpolate = interpolate + self.upsample_factor = upsample_factor + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + conv = None + if use_conv_transpose: + raise NotImplementedError + elif use_conv: + if kernel_size is None: + kernel_size = 3 + conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias) + + if name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + output_size: Optional[int] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + raise NotImplementedError + + if self.use_conv_transpose: + return self.conv(hidden_states) + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if self.interpolate: + B, C, T, H, W = hidden_states.shape + first_h, other_h = hidden_states.split((1, T - 1), dim=2) + if output_size is None: + if T > 1: + other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest") + + first_h = first_h.squeeze(2) + first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest") + first_h = first_h.unsqueeze(2) + else: + raise NotImplementedError + + if T > 1: + hidden_states = torch.cat((first_h, other_h), dim=2) + else: + hidden_states = first_h + + # If the input is bfloat16, we cast back to bfloat16 + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + if self.use_conv: + if self.name == "conv": + hidden_states = self.conv(hidden_states) + else: + hidden_states = self.Conv2d_0(hidden_states) + + return hidden_states + + +class DownsampleCausal3D(nn.Module): + """ + A 3D downsampling layer with an optional convolution. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + kernel_size=3, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + stride=2, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = stride + self.name = name + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + if use_conv: + conv = CausalConv3d( + self.channels, + self.out_channels, + kernel_size=kernel_size, + stride=stride, + bias=bias, + ) + else: + raise NotImplementedError + + if name == "conv": + self.Conv2d_0 = conv + self.conv = conv + elif name == "Conv2d_0": + self.conv = conv + else: + self.conv = conv + + def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + assert hidden_states.shape[1] == self.channels + + hidden_states = self.conv(hidden_states) + + return hidden_states + + +class ResnetBlockCausal3D(nn.Module): + r""" + A Resnet block. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: Optional[int] = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + pre_norm: bool = True, + eps: float = 1e-6, + non_linearity: str = "swish", + skip_time_act: bool = False, + # default, scale_shift, ada_group, spatial + time_embedding_norm: str = "default", + kernel: Optional[torch.FloatTensor] = None, + output_scale_factor: float = 1.0, + use_in_shortcut: Optional[bool] = None, + up: bool = False, + down: bool = False, + conv_shortcut_bias: bool = True, + conv_3d_out_channels: Optional[int] = None, + ): + super().__init__() + self.pre_norm = pre_norm + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.up = up + self.down = down + self.output_scale_factor = output_scale_factor + self.time_embedding_norm = time_embedding_norm + self.skip_time_act = skip_time_act + + linear_cls = nn.Linear + + if groups_out is None: + groups_out = groups + + if self.time_embedding_norm == "ada_group": + self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm1 = SpatialNorm(in_channels, temb_channels) + else: + self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) + + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + self.time_emb_proj = linear_cls(temb_channels, out_channels) + elif self.time_embedding_norm == "scale_shift": + self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels) + elif (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"): + self.time_emb_proj = None + else: + raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ") + else: + self.time_emb_proj = None + + if self.time_embedding_norm == "ada_group": + self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm2 = SpatialNorm(out_channels, temb_channels) + else: + self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) + + self.dropout = torch.nn.Dropout(dropout) + conv_3d_out_channels = conv_3d_out_channels or out_channels + self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1) + + self.nonlinearity = get_activation(non_linearity) + + self.upsample = self.downsample = None + if self.up: + self.upsample = UpsampleCausal3D(in_channels, use_conv=False) + elif self.down: + self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op") + + self.use_in_shortcut = (self.in_channels != conv_3d_out_channels + if use_in_shortcut is None else use_in_shortcut) + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = CausalConv3d( + in_channels, + conv_3d_out_channels, + kernel_size=1, + stride=1, + bias=conv_shortcut_bias, + ) + + def forward( + self, + input_tensor: torch.FloatTensor, + temb: torch.FloatTensor, + scale: float = 1.0, + ) -> torch.FloatTensor: + hidden_states = input_tensor + + if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"): + hidden_states = self.norm1(hidden_states, temb) + else: + hidden_states = self.norm1(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = self.upsample(input_tensor, scale=scale) + hidden_states = self.upsample(hidden_states, scale=scale) + elif self.downsample is not None: + input_tensor = self.downsample(input_tensor, scale=scale) + hidden_states = self.downsample(hidden_states, scale=scale) + + hidden_states = self.conv1(hidden_states) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb, scale)[:, :, None, None] + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + if (self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial"): + hidden_states = self.norm2(hidden_states, temb) + else: + hidden_states = self.norm2(hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + + +def get_down_block3d( + down_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + temb_channels: int, + add_downsample: bool, + downsample_stride: int, + resnet_eps: float, + resnet_act_fn: str, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + downsample_padding: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + downsample_type: Optional[str] = None, + dropout: float = 0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = (down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type) + # if down_block_type == "DownEncoderBlockCausal3D": + return DownEncoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + downsample_stride=downsample_stride, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block3d( + up_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + add_upsample: bool, + upsample_scale_factor: Tuple, + resnet_eps: float, + resnet_act_fn: str, + resolution_idx: Optional[int] = None, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + upsample_type: Optional[str] = None, + dropout: float = 0.0, +) -> nn.Module: + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = (up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type) + # if up_block_type == "UpDecoderBlockCausal3D": + return UpDecoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + upsample_scale_factor=upsample_scale_factor, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlockCausal3D(nn.Module): + """ + A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks. + """ + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + attn_groups: Optional[int] = None, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + ): + super().__init__() + resnet_groups = (resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)) + self.add_attention = add_attention + + if attn_groups is None: + attn_groups = (resnet_groups if resnet_time_scale_shift == "default" else None) + + # there is always at least one resnet + resnets = [ + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=attn_groups, + spatial_norm_dim=(temb_channels if resnet_time_scale_shift == "spatial" else None), + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + else: + attentions.append(None) + + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor: + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + B, C, T, H, W = hidden_states.shape + hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c") + attention_mask = prepare_causal_attention_mask(T, + H * W, + hidden_states.dtype, + hidden_states.device, + batch_size=B) + hidden_states = attn(hidden_states, temb=temb, attention_mask=attention_mask) + hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=T, h=H, w=W) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class DownEncoderBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + DownsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + stride=downsample_stride, + ) + ]) + else: + self.downsamplers = None + + def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, scale=scale) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class UpDecoderBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + upsample_scale_factor=(2, 2, 2), + temb_channels: Optional[int] = None, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + UpsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ]) + else: + self.upsamplers = None + + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states diff --git a/fastvideo/v1/v0_reference_src/models/hunyuan/vae/vae.py b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/vae.py new file mode 100644 index 000000000..1d0f75e51 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/models/hunyuan/vae/vae.py @@ -0,0 +1,342 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn +from diffusers.models.attention_processor import SpatialNorm +from diffusers.utils import BaseOutput, is_torch_version +from diffusers.utils.torch_utils import randn_tensor + +from .unet_causal_3d_blocks import CausalConv3d, UNetMidBlockCausal3D, get_down_block3d, get_up_block3d + + +@dataclass +class DecoderOutput(BaseOutput): + r""" + Output of decoding method. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The decoded output sample from the last layer of the model. + """ + + sample: torch.FloatTensor + + +class EncoderCausal3D(nn.Module): + r""" + The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_downsample = bool(i < num_spatial_downsample_layers) + add_time_downsample = bool(i >= (len(block_out_channels) - 1 - num_time_downsample_layers) + and not is_final_block) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.") + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2, ) if add_time_downsample else (1, ) + downsample_stride = tuple(downsample_stride_T + downsample_stride_HW) + down_block = get_down_block3d( + down_block_type, + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample or add_time_downsample), + downsample_stride=downsample_stride, + resnet_eps=1e-6, + downsample_padding=0, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=None, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + ) + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) + + def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor: + r"""The forward method of the `EncoderCausal3D` class.""" + assert len(sample.shape) == 5, "The input tensor should have 5 dimensions" + + sample = self.conv_in(sample) + + # down + for down_block in self.down_blocks: + sample = down_block(sample) + + # middle + sample = self.mid_block(sample) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DecoderCausal3D(nn.Module): + r""" + The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", # group, spatial + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + temb_channels = in_channels if norm_type == "spatial" else None + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default" if norm_type == "group" else norm_type, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool(i >= len(block_out_channels) - 1 - num_time_upsample_layers + and not is_final_block) + else: + raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.") + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1) + upsample_scale_factor_T = (2, ) if add_time_upsample else (1, ) + upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW) + up_block = get_up_block3d( + up_block_type, + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + prev_output_channel=None, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=temb_channels, + resnet_time_scale_shift=norm_type, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_type == "spatial": + self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) + else: + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3) + + self.gradient_checkpointing = False + + def forward( + self, + sample: torch.FloatTensor, + latent_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + r"""The forward method of the `DecoderCausal3D` class.""" + assert len(sample.shape) == 5, "The input tensor should have 5 dimensions." + + sample = self.conv_in(sample) + + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), + sample, + latent_embeds, + use_reentrant=False, + ) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), + sample, + latent_embeds, + use_reentrant=False, + ) + else: + # middle + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block), sample, latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds) + else: + # middle + sample = self.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds) + + # post-process + if latent_embeds is None: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample, latent_embeds) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DiagonalGaussianDistribution(object): + + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + + def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + reduce_dim = list(range(1, self.mean.ndim)) + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=reduce_dim, + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + self.var / other.var - 1.0 - self.logvar + + other.logvar, + dim=reduce_dim, + ) + + def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean diff --git a/fastvideo/pipelines/pipeline_base.py b/fastvideo/v1/v0_reference_src/pipeline_base.py similarity index 97% rename from fastvideo/pipelines/pipeline_base.py rename to fastvideo/v1/v0_reference_src/pipeline_base.py index fbe75cfc1..5229ff1a8 100644 --- a/fastvideo/pipelines/pipeline_base.py +++ b/fastvideo/v1/v0_reference_src/pipeline_base.py @@ -11,13 +11,13 @@ from tqdm.auto import tqdm from diffusers.utils.torch_utils import randn_tensor -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.inference_args import InferenceArgs -from fastvideo.models.hunyuan.text_encoder import TextEncoder -from fastvideo.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D -from fastvideo.models.hunyuan.constants import PRECISION_TO_TYPE -from fastvideo.distributed.parallel_state import get_sp_group -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.v0_reference_src.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.v0_reference_src.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from fastvideo.v1.v0_reference_src.models.hunyuan.constants import PRECISION_TO_TYPE +from fastvideo.v1.v0_reference_src.distributed.parallel_state import get_sp_group +from fastvideo.v1.logger import init_logger from einops import rearrange from diffusers.utils import BaseOutput import numpy as np @@ -420,7 +420,7 @@ class DiffusionPipelineBase(ABC): # 5. Prepare latent variables batch = self.prepare_latents(batch) - from fastvideo.utils.parallel_states import nccl_info + from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group # world_size = 1 # rank = 0 @@ -541,8 +541,8 @@ class DiffusionPipelineBase(ABC): # if sp_group: # latents = sp_group.all_gather(latents, dim=2) - from fastvideo.utils.communications import all_gather - from fastvideo.utils.parallel_states import get_sequence_parallel_state + from fastvideo.v1.v0_reference_src.utils.communications import all_gather + from fastvideo.v1.v0_reference_src.utils.parallel_states import get_sequence_parallel_state if get_sequence_parallel_state(): latents = all_gather(latents, dim=2) diff --git a/fastvideo/pipelines/implementations/hunyuan/constants.py b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/constants.py similarity index 100% rename from fastvideo/pipelines/implementations/hunyuan/constants.py rename to fastvideo/v1/v0_reference_src/pipelines/hunyuan/constants.py diff --git a/fastvideo/pipelines/hunyuan/hunyuan_loader.py b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_loader.py similarity index 95% rename from fastvideo/pipelines/hunyuan/hunyuan_loader.py rename to fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_loader.py index 6ee639f3f..fd157b92e 100644 --- a/fastvideo/pipelines/hunyuan/hunyuan_loader.py +++ b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_loader.py @@ -4,19 +4,19 @@ from pathlib import Path from safetensors.torch import load_file as safetensors_load_file from ..loader import PipelineLoader -from fastvideo.inference_args import InferenceArgs -from fastvideo.pipelines.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE -from fastvideo.logger import init_logger +from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.pipelines.hunyuan.constants import PRECISION_TO_TYPE, PROMPT_TEMPLATE +from fastvideo.v1.logger import init_logger from typing import Tuple, Optional import os import glob import json -from fastvideo.loader.fsdp_load import load_fsdp_model +from fastvideo.v1.loader.fsdp_load import load_fsdp_model -from fastvideo.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder -from fastvideo.platforms import current_platform +from fastvideo.v1.platforms import current_platform logger = init_logger(__name__) @@ -157,7 +157,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_transformer_v0(self, inference_args: InferenceArgs): """Custom transformer loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.modules import load_model + from fastvideo.v1.models.hunyuan.modules import load_model # Disable gradient torch.set_grad_enabled(False) @@ -183,7 +183,7 @@ class HunyuanPipelineLoader(PipelineLoader): def load_vae(self, inference_args: InferenceArgs): """Custom VAE loading for Hunyuan""" # TODO(will): replace this with abstracted model - from fastvideo.models.hunyuan.vae import load_vae + from fastvideo.v1.models.hunyuan.vae import load_vae vae, _, s_ratio, t_ratio = load_vae( inference_args.vae, inference_args.vae_precision, @@ -221,7 +221,6 @@ class HunyuanPipelineLoader(PipelineLoader): hidden_state_skip_layer=inference_args.hidden_state_skip_layer, apply_final_norm=inference_args.apply_final_norm, reproduce=inference_args.reproduce, - logger=logger, device=self.device if not inference_args.use_cpu_offload else "cpu", ) text_encoder_2 = None @@ -232,7 +231,6 @@ class HunyuanPipelineLoader(PipelineLoader): text_encoder_precision=inference_args.text_encoder_precision_2, tokenizer_type=inference_args.tokenizer_2, reproduce=inference_args.reproduce, - logger=logger, device=self.device if not inference_args.use_cpu_offload else "cpu", ) diff --git a/fastvideo/pipelines/hunyuan/hunyuan_video.py b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_video.py similarity index 93% rename from fastvideo/pipelines/hunyuan/hunyuan_video.py rename to fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_video.py index 3bcae1e33..5a807b2e0 100644 --- a/fastvideo/pipelines/hunyuan/hunyuan_video.py +++ b/fastvideo/v1/v0_reference_src/pipelines/hunyuan/hunyuan_video.py @@ -6,19 +6,19 @@ from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple, Union from diffusers.image_processor import VaeImageProcessor -from fastvideo.pipelines.pipeline_base import DiffusionPipelineBase +from fastvideo.v1.pipelines.pipeline_base import DiffusionPipelineBase # from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.logger import init_logger +from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.v1.logger import init_logger # TODO(will): temporary import from diffusers.models import AutoencoderKL -# from fastvideo.models.encoders.encoder import TextEncoder -from fastvideo.models.hunyuan.text_encoder import TextEncoder +# from fastvideo.v1.models.encoders.encoder import TextEncoder +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder from diffusers.schedulers import KarrasDiffusionSchedulers -from fastvideo.models.hunyuan.modules import HYVideoDiffusionTransformer -from fastvideo.inference_args import InferenceArgs +from fastvideo.v1.models.hunyuan.modules import HYVideoDiffusionTransformer +from fastvideo.v1.inference_args import InferenceArgs logger = init_logger(__name__) diff --git a/fastvideo/v1/v0_reference_src/utils/checkpoint.py b/fastvideo/v1/v0_reference_src/utils/checkpoint.py new file mode 100644 index 000000000..5a4300c13 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/checkpoint.py @@ -0,0 +1,286 @@ +# import +import json +import os + +import torch +import torch.distributed.checkpoint as dist_cp +from peft import get_peft_model_state_dict +from safetensors.torch import load_file, save_file +from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner +from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict +from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import StateDictType + +from fastvideo.v1.utils.logging_ import main_print + + +def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discriminator=False): + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + cpu_state = model.state_dict() + optim_state = FSDP.optim_state_dict( + model, + optimizer, + ) + + # todo move to get_state_dict + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + # save using safetensors + if rank <= 0 and not discriminator: + weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + config_dict = dict(model.config) + config_dict.pop('dtype') + config_path = os.path.join(save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + optimizer_path = os.path.join(save_dir, "optimizer.pt") + torch.save(optim_state, optimizer_path) + else: + weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt") + torch.save(optim_state, optimizer_path) + main_print(f"--> checkpoint saved at step {step}") + + +def save_checkpoint(transformer, rank, output_dir, step): + main_print(f"--> saving checkpoint at step {step}") + with FSDP.state_dict_type( + transformer, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + cpu_state = transformer.state_dict() + # todo move to get_state_dict + if rank <= 0: + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + # save using safetensors + weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + config_dict = dict(transformer.config) + if "dtype" in config_dict: + del config_dict["dtype"] # TODO + config_path = os.path.join(save_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + main_print(f"--> checkpoint saved at step {step}") + + +def save_checkpoint_generator_discriminator( + model, + optimizer, + discriminator, + discriminator_optimizer, + rank, + output_dir, + step, +): + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + cpu_state = model.state_dict() + + # todo move to get_state_dict + save_dir = os.path.join(output_dir, f"checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + hf_weight_dir = os.path.join(save_dir, "hf_weights") + os.makedirs(hf_weight_dir, exist_ok=True) + # save using safetensors + if rank <= 0: + config_dict = dict(model.config) + config_path = os.path.join(hf_weight_dir, "config.json") + # save dict as json + with open(config_path, "w") as f: + json.dump(config_dict, f, indent=4) + weight_path = os.path.join(hf_weight_dir, "diffusion_pytorch_model.safetensors") + save_file(cpu_state, weight_path) + + main_print(f"--> saved HF weight checkpoint at path {hf_weight_dir}") + model_weight_dir = os.path.join(save_dir, "model_weights_state") + os.makedirs(model_weight_dir, exist_ok=True) + model_optimizer_dir = os.path.join(save_dir, "model_optimizer_state") + os.makedirs(model_optimizer_dir, exist_ok=True) + with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): + optim_state = FSDP.optim_state_dict(model, optimizer) + model_state = model.state_dict() + weight_state_dict = {"model": model_state} + dist_cp.save_state_dict( + state_dict=weight_state_dict, + storage_writer=dist_cp.FileSystemWriter(model_weight_dir), + planner=DefaultSavePlanner(), + ) + optimizer_state_dict = {"optimizer": optim_state} + dist_cp.save_state_dict( + state_dict=optimizer_state_dict, + storage_writer=dist_cp.FileSystemWriter(model_optimizer_dir), + planner=DefaultSavePlanner(), + ) + + discriminator_fsdp_state_dir = os.path.join(save_dir, "discriminator_fsdp_state") + os.makedirs(discriminator_fsdp_state_dir, exist_ok=True) + with FSDP.state_dict_type( + discriminator, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + optim_state = FSDP.optim_state_dict(discriminator, discriminator_optimizer) + model_state = discriminator.state_dict() + state_dict = {"optimizer": optim_state, "model": model_state} + if rank <= 0: + discriminator_fsdp_state_fil = os.path.join(discriminator_fsdp_state_dir, "discriminator_state.pt") + torch.save(state_dict, discriminator_fsdp_state_fil) + + main_print("--> saved FSDP state checkpoint") + + +def load_sharded_model(model, optimizer, model_dir, optimizer_dir): + with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT): + weight_state_dict = {"model": model.state_dict()} + + optim_state = load_sharded_optimizer_state_dict( + model_state_dict=weight_state_dict["model"], + optimizer_key="optimizer", + storage_reader=dist_cp.FileSystemReader(optimizer_dir), + ) + optim_state = optim_state["optimizer"] + flattened_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optim_state) + optimizer.load_state_dict(flattened_osd) + dist_cp.load_state_dict( + state_dict=weight_state_dict, + storage_reader=dist_cp.FileSystemReader(model_dir), + planner=DefaultLoadPlanner(), + ) + model_state = weight_state_dict["model"] + model.load_state_dict(model_state) + main_print(f"--> loaded model and optimizer from path {model_dir}") + return model, optimizer + + +def load_full_state_model(model, optimizer, checkpoint_file, rank): + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + discriminator_state = torch.load(checkpoint_file) + model_state = discriminator_state["model"] + if rank <= 0: + optim_state = discriminator_state["optimizer"] + else: + optim_state = None + model.load_state_dict(model_state) + discriminator_optim_state = FSDP.optim_state_dict_to_load(model=model, + optim=optimizer, + optim_state_dict=optim_state) + optimizer.load_state_dict(discriminator_optim_state) + main_print(f"--> loaded discriminator and discriminator optimizer from path {checkpoint_file}") + return model, optimizer + + +def resume_training_generator_discriminator(model, optimizer, discriminator, discriminator_optimizer, checkpoint_dir, + rank): + step = int(checkpoint_dir.split("-")[-1]) + model_weight_dir = os.path.join(checkpoint_dir, "model_weights_state") + model_optimizer_dir = os.path.join(checkpoint_dir, "model_optimizer_state") + model, optimizer = load_sharded_model(model, optimizer, model_weight_dir, model_optimizer_dir) + discriminator_ckpt_file = os.path.join(checkpoint_dir, "discriminator_fsdp_state", "discriminator_state.pt") + discriminator, discriminator_optimizer = load_full_state_model(discriminator, discriminator_optimizer, + discriminator_ckpt_file, rank) + return model, optimizer, discriminator, discriminator_optimizer, step + + +def resume_training(model, optimizer, checkpoint_dir, discriminator=False): + weight_path = os.path.join(checkpoint_dir, "diffusion_pytorch_model.safetensors") + if discriminator: + weight_path = os.path.join(checkpoint_dir, "discriminator_pytorch_model.safetensors") + model_weights = load_file(weight_path) + + with FSDP.state_dict_type( + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + current_state = model.state_dict() + current_state.update(model_weights) + model.load_state_dict(current_state, strict=False) + if discriminator: + optim_path = os.path.join(checkpoint_dir, "discriminator_optimizer.pt") + else: + optim_path = os.path.join(checkpoint_dir, "optimizer.pt") + optimizer_state_dict = torch.load(optim_path, weights_only=False) + optim_state = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optimizer_state_dict) + optimizer.load_state_dict(optim_state) + step = int(checkpoint_dir.split("-")[-1]) + return model, optimizer, step + + +def save_lora_checkpoint(transformer, optimizer, rank, output_dir, step, pipeline): + with FSDP.state_dict_type( + transformer, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), + ): + full_state_dict = transformer.state_dict() + lora_optim_state = FSDP.optim_state_dict( + transformer, + optimizer, + ) + + if rank <= 0: + save_dir = os.path.join(output_dir, f"lora-checkpoint-{step}") + os.makedirs(save_dir, exist_ok=True) + + # save optimizer + optim_path = os.path.join(save_dir, "lora_optimizer.pt") + torch.save(lora_optim_state, optim_path) + # save lora weight + main_print(f"--> saving LoRA checkpoint at step {step}") + transformer_lora_layers = get_peft_model_state_dict(model=transformer, state_dict=full_state_dict) + pipeline.save_lora_weights( + save_directory=save_dir, + transformer_lora_layers=transformer_lora_layers, + is_main_process=True, + ) + # save config + lora_config = { + "step": step, + "lora_params": { + "lora_rank": transformer.config.lora_rank, + "lora_alpha": transformer.config.lora_alpha, + "target_modules": transformer.config.lora_target_modules, + }, + } + config_path = os.path.join(save_dir, "lora_config.json") + with open(config_path, "w") as f: + json.dump(lora_config, f, indent=4) + main_print(f"--> LoRA checkpoint saved at step {step}") + + +def resume_lora_optimizer(transformer, checkpoint_dir, optimizer): + config_path = os.path.join(checkpoint_dir, "lora_config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + optim_path = os.path.join(checkpoint_dir, "lora_optimizer.pt") + optimizer_state_dict = torch.load(optim_path, weights_only=False) + optim_state = FSDP.optim_state_dict_to_load(model=transformer, + optim=optimizer, + optim_state_dict=optimizer_state_dict) + optimizer.load_state_dict(optim_state) + step = config_dict["step"] + main_print(f"--> Successfully resuming LoRA optimizer from step {step}") + return transformer, optimizer, step diff --git a/fastvideo/v1/v0_reference_src/utils/communications.py b/fastvideo/v1/v0_reference_src/utils/communications.py new file mode 100644 index 000000000..38056738a --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/communications.py @@ -0,0 +1,308 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Any, Tuple + +import torch +import torch.distributed as dist +from torch import Tensor + +from fastvideo.v1.v0_reference_src.utils.parallel_states import nccl_info + + +def broadcast(input_: torch.Tensor): + src = nccl_info.group_id * nccl_info.sp_size + dist.broadcast(input_, src=src, group=nccl_info.group) + + +def _all_to_all_4D(input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None) -> torch.tensor: + """ + all-to-all for QKV + + Args: + input (torch.tensor): a tensor sharded along dim scatter dim + scatter_idx (int): default 1 + gather_idx (int): default 2 + group : torch process group + + Returns: + torch.tensor: resharded tensor (bs, seqlen/P, hc, hs) + """ + assert (input.dim() == 4), f"input must be 4D tensor, got {input.dim()} and shape {input.shape}" + + seq_world_size = dist.get_world_size(group) + + if scatter_idx == 2 and gather_idx == 1: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs) + bs, shard_seqlen, hc, hs = input.shape + seqlen = shard_seqlen * seq_world_size + shard_hc = hc // seq_world_size + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs) + input_t = (input.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs).transpose(0, 2).contiguous()) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + torch.cuda.synchronize() + else: + output = input_t + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(seqlen, bs, shard_hc, hs) + + # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs) + output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs) + + return output + + elif scatter_idx == 1 and gather_idx == 2: + # input (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs) + bs, seqlen, shard_hc, hs = input.shape + hc = shard_hc * seq_world_size + shard_seqlen = seqlen // seq_world_size + seq_world_size = dist.get_world_size(group) + + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs) + input_t = (input.reshape(bs, seq_world_size, shard_seqlen, shard_hc, + hs).transpose(0, + 3).transpose(0, + 1).contiguous().reshape(seq_world_size, shard_hc, + shard_seqlen, bs, hs)) + + output = torch.empty_like(input_t) + # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single + # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head + if seq_world_size > 1: + dist.all_to_all_single(output, input_t, group=group) + torch.cuda.synchronize() + else: + output = input_t + + # if scattering the seq-dim, transpose the heads back to the original dimension + output = output.reshape(hc, shard_seqlen, bs, hs) + + # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs) + output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs) + + return output + else: + raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2") + + +class SeqAllToAll4D(torch.autograd.Function): + + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: Tensor, + scatter_idx: int, + gather_idx: int, + ) -> Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + + return _all_to_all_4D(input, scatter_idx, gather_idx, group=group) + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: + return ( + None, + SeqAllToAll4D.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), + None, + None, + ) + + +def all_to_all_4D( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return SeqAllToAll4D.apply(nccl_info.group, input_, scatter_dim, gather_dim) + + +def _all_to_all( + input_: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, + scatter_dim: int, + gather_dim: int, +): + input_list = [t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class _AllToAll(torch.autograd.Function): + """All-to-all communication. + + Args: + input_: input matrix + process_group: communication group + scatter_dim: scatter dimension + gather_dim: gather dimension + """ + + @staticmethod + def forward(ctx, input_, process_group, scatter_dim, gather_dim): + ctx.process_group = process_group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.world_size = dist.get_world_size(process_group) + output = _all_to_all(input_, ctx.world_size, process_group, scatter_dim, gather_dim) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = _all_to_all( + grad_output, + ctx.world_size, + ctx.process_group, + ctx.gather_dim, + ctx.scatter_dim, + ) + return ( + grad_output, + None, + None, + None, + ) + + +def all_to_all( + input_: torch.Tensor, + scatter_dim: int = 2, + gather_dim: int = 1, +): + return _AllToAll.apply(input_, nccl_info.group, scatter_dim, gather_dim) + + +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.sp_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.sp_size + rank = nccl_info.rank_within_group + 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(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_sequence_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): + if nccl_info.sp_size == 1: + return ( + hidden_states, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) + + def prepare(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): + hidden_states = all_to_all(hidden_states, scatter_dim=2, gather_dim=0) + encoder_hidden_states = all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0) + attention_mask = all_to_all(attention_mask, scatter_dim=1, gather_dim=0) + encoder_attention_mask = all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0) + return ( + hidden_states, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) + + sp_size = nccl_info.sp_size + frame = hidden_states.shape[2] + assert frame % sp_size == 0, "frame should be a multiple of sp_size" + + ( + hidden_states, + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) = prepare( + hidden_states, + encoder_hidden_states.repeat(1, sp_size, 1), + attention_mask.repeat(1, sp_size, 1, 1), + encoder_attention_mask.repeat(1, sp_size), + ) + + return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask + + +def sp_parallel_dataloader_wrapper(dataloader, device, train_batch_size, sp_size, train_sp_batch_size): + while True: + for data_item in dataloader: + latents, cond, attn_mask, cond_mask = data_item + latents = latents.to(device) + cond = cond.to(device) + attn_mask = attn_mask.to(device) + cond_mask = cond_mask.to(device) + frame = latents.shape[2] + if frame == 1: + yield latents, cond, attn_mask, cond_mask + else: + latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data( + latents, cond, attn_mask, cond_mask) + assert (train_batch_size * sp_size >= + train_sp_batch_size), "train_batch_size * sp_size should be greater than train_sp_batch_size" + for iter in range(train_batch_size * sp_size // train_sp_batch_size): + st_idx = iter * train_sp_batch_size + ed_idx = (iter + 1) * train_sp_batch_size + encoder_hidden_states = cond[st_idx:ed_idx] + attention_mask = attn_mask[st_idx:ed_idx] + encoder_attention_mask = cond_mask[st_idx:ed_idx] + yield ( + latents[st_idx:ed_idx], + encoder_hidden_states, + attention_mask, + encoder_attention_mask, + ) diff --git a/fastvideo/v1/v0_reference_src/utils/dataset_utils.py b/fastvideo/v1/v0_reference_src/utils/dataset_utils.py new file mode 100644 index 000000000..3c8c9e5a6 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/dataset_utils.py @@ -0,0 +1,342 @@ +import math +import random +from collections import Counter +from typing import List, Optional + +import decord +import torch +import torch.utils +import torch.utils.data +from torch.nn import functional as F +from torch.utils.data import Sampler + +IMG_EXTENSIONS = [".jpg", ".JPG", ".jpeg", ".JPEG", ".png", ".PNG"] + + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + + +class DecordInit(object): + """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader.""" + + def __init__(self, num_threads=1): + self.num_threads = num_threads + self.ctx = decord.cpu(0) + + def __call__(self, filename): + """Perform the Decord initialization. + Args: + results (dict): The resulting dict to be modified and passed + to the next transform in pipeline. + """ + reader = decord.VideoReader(filename, ctx=self.ctx, num_threads=self.num_threads) + return reader + + def __repr__(self): + repr_str = (f"{self.__class__.__name__}(" + f"sr={self.sr}," + f"num_threads={self.num_threads})") + return repr_str + + +def pad_to_multiple(number, ds_stride): + remainder = number % ds_stride + if remainder == 0: + return number + else: + padding = ds_stride - remainder + return number + padding + + +# TODO +class Collate: + + def __init__(self, args): + self.batch_size = args.train_batch_size + self.group_frame = args.group_frame + self.group_resolution = args.group_resolution + + self.max_height = args.max_height + self.max_width = args.max_width + self.ae_stride = args.ae_stride + + self.ae_stride_t = args.ae_stride_t + self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride) + + self.patch_size = args.patch_size + self.patch_size_t = args.patch_size_t + + self.num_frames = args.num_frames + self.use_image_num = args.use_image_num + self.max_thw = (self.num_frames, self.max_height, self.max_width) + + def package(self, batch): + batch_tubes = [i["pixel_values"] for i in batch] # b [c t h w] + input_ids = [i["input_ids"] for i in batch] # b [1 l] + cond_mask = [i["cond_mask"] for i in batch] # b [1 l] + return batch_tubes, input_ids, cond_mask + + def __call__(self, batch): + batch_tubes, input_ids, cond_mask = self.package(batch) + + ds_stride = self.ae_stride * self.patch_size + t_ds_stride = self.ae_stride_t * self.patch_size_t + + pad_batch_tubes, attention_mask, input_ids, cond_mask = self.process( + batch_tubes, + input_ids, + cond_mask, + t_ds_stride, + ds_stride, + self.max_thw, + self.ae_stride_thw, + ) + assert not torch.any(torch.isnan(pad_batch_tubes)), "after pad_batch_tubes" + return pad_batch_tubes, attention_mask, input_ids, cond_mask + + def process( + self, + batch_tubes, + input_ids, + cond_mask, + t_ds_stride, + ds_stride, + max_thw, + ae_stride_thw, + ): + # pad to max multiple of ds_stride + batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] + assert len(batch_input_size) == self.batch_size + if self.group_frame or self.group_resolution or self.batch_size == 1: # + len_each_batch = batch_input_size + idx_length_dict = dict([*zip(list(range(self.batch_size)), len_each_batch)]) + count_dict = Counter(len_each_batch) + if len(count_dict) != 1: + sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1]) + pick_length = sorted_by_value[-1][0] # the highest frequency + candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length] + random_select_batch = [ + random.choice(candidate_batch) for _ in range(len(len_each_batch) - len(candidate_batch)) + ] + print( + batch_input_size, + idx_length_dict, + count_dict, + sorted_by_value, + pick_length, + candidate_batch, + random_select_batch, + ) + pick_idx = candidate_batch + random_select_batch + + batch_tubes = [batch_tubes[i] for i in pick_idx] + batch_input_size = [i.shape for i in batch_tubes] # [(c t h w), (c t h w)] + input_ids = [input_ids[i] for i in pick_idx] # b [1, l] + cond_mask = [cond_mask[i] for i in pick_idx] # b [1, l] + + for i in range(1, self.batch_size): + assert batch_input_size[0] == batch_input_size[i] + max_t = max([i[1] for i in batch_input_size]) + max_h = max([i[2] for i in batch_input_size]) + max_w = max([i[3] for i in batch_input_size]) + else: + max_t, max_h, max_w = max_thw + pad_max_t, pad_max_h, pad_max_w = ( + pad_to_multiple(max_t - 1 + self.ae_stride_t, t_ds_stride), + pad_to_multiple(max_h, ds_stride), + pad_to_multiple(max_w, ds_stride), + ) + pad_max_t = pad_max_t + 1 - self.ae_stride_t + each_pad_t_h_w = [[pad_max_t - i.shape[1], pad_max_h - i.shape[2], pad_max_w - i.shape[3]] for i in batch_tubes] + pad_batch_tubes = [ + F.pad(im, (0, pad_w, 0, pad_h, 0, pad_t), value=0) + for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes) + ] + pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0) + + max_tube_size = [pad_max_t, pad_max_h, pad_max_w] + max_latent_size = [ + ((max_tube_size[0] - 1) // ae_stride_thw[0] + 1), + max_tube_size[1] // ae_stride_thw[1], + max_tube_size[2] // ae_stride_thw[2], + ] + valid_latent_size = [[ + int(math.ceil((i[1] - 1) / ae_stride_thw[0])) + 1, + int(math.ceil(i[2] / ae_stride_thw[1])), + int(math.ceil(i[3] / ae_stride_thw[2])), + ] for i in batch_input_size] + attention_mask = [ + F.pad( + torch.ones(i, dtype=pad_batch_tubes.dtype), + ( + 0, + max_latent_size[2] - i[2], + 0, + max_latent_size[1] - i[1], + 0, + max_latent_size[0] - i[0], + ), + value=0, + ) for i in valid_latent_size + ] + attention_mask = torch.stack(attention_mask) # b t h w + if self.batch_size == 1 or self.group_frame or self.group_resolution: + assert torch.all(attention_mask.bool()) + + input_ids = torch.stack(input_ids) # b 1 l + cond_mask = torch.stack(cond_mask) # b 1 l + + return pad_batch_tubes, attention_mask, input_ids, cond_mask + + +def split_to_even_chunks(indices, lengths, num_chunks, batch_size): + """ + Split a list of indices into `chunks` chunks of roughly equal lengths. + """ + + if len(indices) % num_chunks != 0: + chunks = [indices[i::num_chunks] for i in range(num_chunks)] + else: + num_indices_per_chunk = len(indices) // num_chunks + + chunks = [[] for _ in range(num_chunks)] + chunks_lengths = [0 for _ in range(num_chunks)] + for index in indices: + shortest_chunk = chunks_lengths.index(min(chunks_lengths)) + chunks[shortest_chunk].append(index) + chunks_lengths[shortest_chunk] += lengths[index] + if len(chunks[shortest_chunk]) == num_indices_per_chunk: + chunks_lengths[shortest_chunk] = float("inf") + # return chunks + + pad_chunks = [] + for idx, chunk in enumerate(chunks): + if batch_size != len(chunk): + assert batch_size > len(chunk) + if len(chunk) != 0: + chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))] + else: + chunk = random.choice(pad_chunks) + print(chunks[idx], "->", chunk) + pad_chunks.append(chunk) + return pad_chunks + + +def group_frame_fun(indices, lengths): + # sort by num_frames + indices.sort(key=lambda i: lengths[i], reverse=True) + return indices + + +def megabatch_frame_alignment(megabatches, lengths): + aligned_magabatches = [] + for _, megabatch in enumerate(megabatches): + assert len(megabatch) != 0 + len_each_megabatch = [lengths[i] for i in megabatch] + idx_length_dict = dict([*zip(megabatch, len_each_megabatch)]) + count_dict = Counter(len_each_megabatch) + + # mixed frame length, align megabatch inside + if len(count_dict) != 1: + sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1]) + pick_length = sorted_by_value[-1][0] # the highest frequency + candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length] + random_select_batch = [ + random.choice(candidate_batch) for i in range(len(idx_length_dict) - len(candidate_batch)) + ] + aligned_magabatch = candidate_batch + random_select_batch + aligned_magabatches.append(aligned_magabatch) + # already aligned megabatches + else: + aligned_magabatches.append(megabatch) + + return aligned_magabatches + + +def get_length_grouped_indices( + lengths, + batch_size, + world_size, + generator=None, + group_frame=False, + group_resolution=False, + seed=42, +): + # We need to use torch for the random part as a distributed sampler will set the random seed for torch. + if generator is None: + generator = torch.Generator().manual_seed(seed) # every rank will generate a fixed order but random index + + indices = torch.randperm(len(lengths), generator=generator).tolist() + + # sort dataset according to frame + indices = group_frame_fun(indices, lengths) + + # chunk dataset to megabatches + megabatch_size = world_size * batch_size + megabatches = [indices[i:i + megabatch_size] for i in range(0, len(lengths), megabatch_size)] + + # make sure the length in each magabatch is align with each other + megabatches = megabatch_frame_alignment(megabatches, lengths) + + # aplit aligned megabatch into batches + megabatches = [split_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches] + + # random megabatches to do video-image mix training + indices = torch.randperm(len(megabatches), generator=generator).tolist() + shuffled_megabatches = [megabatches[i] for i in indices] + + # expand indices and return + return [i for megabatch in shuffled_megabatches for batch in megabatch for i in batch] + + +class LengthGroupedSampler(Sampler): + r""" + Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while + keeping a bit of randomness. + """ + + def __init__( + self, + batch_size: int, + rank: int, + world_size: int, + lengths: Optional[List[int]] = None, + group_frame=False, + group_resolution=False, + generator=None, + ): + if lengths is None: + raise ValueError("Lengths must be provided.") + + self.batch_size = batch_size + self.rank = rank + self.world_size = world_size + self.lengths = lengths + self.group_frame = group_frame + self.group_resolution = group_resolution + self.generator = generator + + def __len__(self): + return len(self.lengths) + + def __iter__(self): + indices = get_length_grouped_indices( + self.lengths, + self.batch_size, + self.world_size, + group_frame=self.group_frame, + group_resolution=self.group_resolution, + generator=self.generator, + ) + + def distributed_sampler(lst, rank, batch_size, world_size): + result = [] + index = rank * batch_size + while index < len(lst): + result.extend(lst[index:index + batch_size]) + index += batch_size * world_size + return result + + indices = distributed_sampler(indices, self.rank, self.batch_size, self.world_size) + return iter(indices) diff --git a/fastvideo/v1/v0_reference_src/utils/env_utils.py b/fastvideo/v1/v0_reference_src/utils/env_utils.py new file mode 100644 index 000000000..13591e9aa --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/env_utils.py @@ -0,0 +1,38 @@ +import platform + +import accelerate +import peft +import torch +import transformers +from transformers.utils import is_torch_cuda_available, is_torch_npu_available + +VERSION = "1.2.0" + +if __name__ == "__main__": + info = { + "FastVideo version": VERSION, + "Platform": platform.platform(), + "Python version": platform.python_version(), + "PyTorch version": torch.__version__, + "Transformers version": transformers.__version__, + "Accelerate version": accelerate.__version__, + "PEFT version": peft.__version__, + } + + if is_torch_cuda_available(): + info["PyTorch version"] += " (GPU)" + info["GPU type"] = torch.cuda.get_device_name() + + if is_torch_npu_available(): + info["PyTorch version"] += " (NPU)" + info["NPU type"] = torch.npu.get_device_name() + info["CANN version"] = torch.version.cann # codespell:ignore + + try: + import bitsandbytes + + info["Bitsandbytes version"] = bitsandbytes.__version__ + except Exception: + pass + + print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n") diff --git a/fastvideo/v1/v0_reference_src/utils/fsdp_util.py b/fastvideo/v1/v0_reference_src/utils/fsdp_util.py new file mode 100644 index 000000000..96a1e11ce --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/fsdp_util.py @@ -0,0 +1,132 @@ +# ruff: noqa: E731 +import functools +from functools import partial + +import torch +from peft.utils.other import fsdp_auto_wrap_policy +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (CheckpointImpl, apply_activation_checkpointing, + checkpoint_wrapper) +from torch.distributed.fsdp import MixedPrecision, ShardingStrategy +from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy + +from fastvideo.v1.models.mochi_hf.modeling_mochi import MochiTransformerBlock +from fastvideo.v1.utils.load import get_no_split_modules + +non_reentrant_wrapper = partial( + checkpoint_wrapper, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, +) + +check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock) + + +def apply_fsdp_checkpointing(model, no_split_modules, p=1): + # https://github.com/foundation-model-stack/fms-fsdp/blob/408c7516d69ea9b6bcd4c0f5efab26c0f64b3c2d/fms_fsdp/policies/ac_handler.py#L16 + """apply activation checkpointing to model + returns None as model is updated directly + """ + print("--> applying fdsp activation checkpointing...") + block_idx = 0 + cut_off = 1 / 2 + # when passing p as a fraction number (e.g. 1/3), it will be interpreted + # as a string in argv, thus we need eval("1/3") here for fractions. + p = eval(p) if isinstance(p, str) else p + + def selective_checkpointing(submodule): + nonlocal block_idx + nonlocal cut_off + + if isinstance(submodule, no_split_modules): + block_idx += 1 + if block_idx * p >= cut_off: + cut_off += 1 + return True + return False + + apply_activation_checkpointing( + model, + checkpoint_wrapper_fn=non_reentrant_wrapper, + check_fn=selective_checkpointing, + ) + + +def get_mixed_precision(master_weight_type="fp32"): + weight_type = torch.float32 if master_weight_type == "fp32" else torch.bfloat16 + mixed_precision = MixedPrecision( + param_dtype=weight_type, + # Gradient communication precision. + reduce_dtype=weight_type, + # Buffer precision. + buffer_dtype=weight_type, + cast_forward_inputs=False, + ) + return mixed_precision + + +def get_dit_fsdp_kwargs( + transformer, + sharding_strategy, + use_lora=False, + cpu_offload=False, + master_weight_type="fp32", +): + no_split_modules = get_no_split_modules(transformer) + if use_lora: + auto_wrap_policy = fsdp_auto_wrap_policy + else: + auto_wrap_policy = functools.partial( + transformer_auto_wrap_policy, + transformer_layer_cls=no_split_modules, + ) + + # we use float32 for fsdp but autocast during training + mixed_precision = get_mixed_precision(master_weight_type) + + if sharding_strategy == "full": + sharding_strategy = ShardingStrategy.FULL_SHARD + elif sharding_strategy == "hybrid_full": + sharding_strategy = ShardingStrategy.HYBRID_SHARD + elif sharding_strategy == "none": + sharding_strategy = ShardingStrategy.NO_SHARD + auto_wrap_policy = None + elif sharding_strategy == "hybrid_zero2": + sharding_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2 + + device_id = torch.cuda.current_device() + cpu_offload = (torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None) + fsdp_kwargs = { + "auto_wrap_policy": auto_wrap_policy, + "mixed_precision": mixed_precision, + "sharding_strategy": sharding_strategy, + "device_id": device_id, + "limit_all_gathers": True, + "cpu_offload": cpu_offload, + } + + # Add LoRA-specific settings when LoRA is enabled + if use_lora: + fsdp_kwargs.update({ + "use_orig_params": False, # Required for LoRA memory savings + "sync_module_states": True, + }) + + return fsdp_kwargs, no_split_modules + + +def get_discriminator_fsdp_kwargs(master_weight_type="fp32"): + auto_wrap_policy = None + + # Use existing mixed precision settings + + mixed_precision = get_mixed_precision(master_weight_type) + sharding_strategy = ShardingStrategy.NO_SHARD + device_id = torch.cuda.current_device() + fsdp_kwargs = { + "auto_wrap_policy": auto_wrap_policy, + "mixed_precision": mixed_precision, + "sharding_strategy": sharding_strategy, + "device_id": device_id, + "limit_all_gathers": True, + } + + return fsdp_kwargs diff --git a/fastvideo/v1/v0_reference_src/utils/load.py b/fastvideo/v1/v0_reference_src/utils/load.py new file mode 100644 index 000000000..d305e0106 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/load.py @@ -0,0 +1,339 @@ +import os +from pathlib import Path + +import torch +import torch.nn.functional as F +from diffusers import AutoencoderKLHunyuanVideo, AutoencoderKLMochi +from torch import nn +from transformers import AutoTokenizer, T5EncoderModel + +from fastvideo.v1.models.hunyuan.modules.models import (HYVideoDiffusionTransformer, MMDoubleStreamBlock, + MMSingleStreamBlock) +from fastvideo.v1.models.hunyuan.text_encoder import TextEncoder +from fastvideo.v1.models.hunyuan.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from fastvideo.v1.models.hunyuan_hf.modeling_hunyuan import (HunyuanVideoSingleTransformerBlock, + HunyuanVideoTransformer3DModel, HunyuanVideoTransformerBlock) +from fastvideo.v1.models.mochi_hf.modeling_mochi import MochiTransformer3DModel, MochiTransformerBlock +from fastvideo.v1.utils.logging_ import main_print + +hunyuan_config = { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, +} + +PROMPT_TEMPLATE_ENCODE = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " + "1. The main content and theme of the video." + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." + "4. background environment, light, style and atmosphere." + "5. camera angles, movements, and transitions used in the video:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>") + +NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion" + +PROMPT_TEMPLATE = { + "dit-llm-encode": { + "template": PROMPT_TEMPLATE_ENCODE, + "crop_start": 36, + }, + "dit-llm-encode-video": { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, + }, +} + + +class HunyuanTextEncoderWrapper(nn.Module): + + def __init__(self, pretrained_model_name_or_path, device): + super().__init__() + + text_len = 256 + crop_start = PROMPT_TEMPLATE["dit-llm-encode-video"].get("crop_start", 0) + + max_length = text_len + crop_start + + # prompt_template + prompt_template = PROMPT_TEMPLATE["dit-llm-encode"] + + # prompt_template_video + prompt_template_video = PROMPT_TEMPLATE["dit-llm-encode-video"] + text_encoder_path = os.path.join(pretrained_model_name_or_path, "text_encoder") + self.text_encoder = TextEncoder( + text_encoder_type="llm", + text_encoder_path=text_encoder_path, + max_length=max_length, + text_encoder_precision="fp16", + tokenizer_type="llm", + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=2, + apply_final_norm=False, + reproduce=False, + logger=None, + device=device, + ) + text_encoder_path_2 = os.path.join(pretrained_model_name_or_path, "text_encoder_2") + self.text_encoder_2 = TextEncoder( + text_encoder_type="clipL", + text_encoder_path=text_encoder_path_2, + max_length=77, + text_encoder_precision="fp16", + tokenizer_type="clipL", + reproduce=False, + logger=None, + device=device, + ) + + def encode_(self, prompt, text_encoder, clip_skip=None): + # TODO + device = self.text_encoder.device + data_type = "video" + num_videos_per_prompt = 1 + + text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) + + if clip_skip is None: + prompt_outputs = text_encoder.encode(text_inputs, data_type="video", device=device) + prompt_embeds = prompt_outputs.hidden_state + else: + prompt_outputs = text_encoder.encode( + text_inputs, + output_hidden_states=True, + data_type=data_type, + device=device, + ) + prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)] + + prompt_embeds = text_encoder.model.text_model.final_layer_norm(prompt_embeds) + + attention_mask = prompt_outputs.attention_mask + if attention_mask is not None: + attention_mask = attention_mask.to(device) + bs_embed, seq_len = attention_mask.shape + attention_mask = attention_mask.repeat(1, num_videos_per_prompt) + attention_mask = attention_mask.view(bs_embed * num_videos_per_prompt, seq_len) + + if text_encoder is not None: + prompt_embeds_dtype = text_encoder.dtype + elif self.transformer is not None: + prompt_embeds_dtype = self.transformer.dtype + else: + prompt_embeds_dtype = prompt_embeds.dtype + + prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device) + + if prompt_embeds.ndim == 2: + bs_embed, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1) + else: + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + return (prompt_embeds, attention_mask) + + def encode_prompt(self, prompt): + prompt_embeds, attention_mask = self.encode_(prompt, self.text_encoder) + prompt_embeds_2, attention_mask_2 = self.encode_(prompt, self.text_encoder_2) + prompt_embeds_2 = F.pad( + prompt_embeds_2, + (0, prompt_embeds.shape[2] - prompt_embeds_2.shape[1]), + value=0, + ).unsqueeze(1) + prompt_embeds = torch.cat([prompt_embeds_2, prompt_embeds], dim=1) + return prompt_embeds, attention_mask + + +class MochiTextEncoderWrapper(nn.Module): + + def __init__(self, pretrained_model_name_or_path, device): + super().__init__() + self.text_encoder = T5EncoderModel.from_pretrained(os.path.join(pretrained_model_name_or_path, + "text_encoder")).to(device) + self.tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_model_name_or_path, "tokenizer")) + self.max_sequence_length = 256 + + def encode_prompt(self, prompt): + device = self.text_encoder.device + dtype = self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.max_sequence_length, + truncation=True, + 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) + + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.max_sequence_length - 1:-1]) + main_print(f"Truncated text input: {prompt} to: {removed_text} for model input.") + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.view(batch_size, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + + return prompt_embeds, prompt_attention_mask + + +def load_hunyuan_state_dict(model, dit_model_name_or_path): + load_key = "module" + model_path = dit_model_name_or_path + bare_model = "unknown" + + state_dict = torch.load(model_path, map_location=lambda storage, loc: storage, weights_only=True) + + if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict): + bare_model = False + if bare_model is False: + if load_key in state_dict: + state_dict = state_dict[load_key] + else: + raise KeyError(f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint " + f"are: {list(state_dict.keys())}.") + model.load_state_dict(state_dict, strict=True) + return model + + +def load_transformer( + model_type, + dit_model_name_or_path, + pretrained_model_name_or_path, + master_weight_type, +): + if model_type == "mochi": + if dit_model_name_or_path: + transformer = MochiTransformer3DModel.from_pretrained( + dit_model_name_or_path, + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + else: + transformer = MochiTransformer3DModel.from_pretrained( + pretrained_model_name_or_path, + subfolder="transformer", + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + elif model_type == "hunyuan_hf": + if dit_model_name_or_path: + transformer = HunyuanVideoTransformer3DModel.from_pretrained( + dit_model_name_or_path, + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + else: + transformer = HunyuanVideoTransformer3DModel.from_pretrained( + pretrained_model_name_or_path, + subfolder="transformer", + torch_dtype=master_weight_type, + # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32, + ) + elif model_type == "hunyuan": + transformer = HYVideoDiffusionTransformer( + in_channels=16, + out_channels=16, + **hunyuan_config, + dtype=master_weight_type, + ) + transformer = load_hunyuan_state_dict(transformer, dit_model_name_or_path) + if master_weight_type == torch.bfloat16: + transformer = transformer.bfloat16() + else: + raise ValueError(f"Unsupported model type: {model_type}") + return transformer + + +def load_vae(model_type, pretrained_model_name_or_path): + weight_dtype = torch.float32 + if model_type == "mochi": + vae = AutoencoderKLMochi.from_pretrained(pretrained_model_name_or_path, + subfolder="vae", + torch_dtype=weight_dtype).to("cuda") + autocast_type = torch.bfloat16 + fps = 30 + elif model_type == "hunyuan_hf": + vae = AutoencoderKLHunyuanVideo.from_pretrained(pretrained_model_name_or_path, + subfolder="vae", + torch_dtype=weight_dtype).to("cuda") + autocast_type = torch.bfloat16 + fps = 24 + elif model_type == "hunyuan": + vae_precision = torch.float32 + vae_path = os.path.join(pretrained_model_name_or_path, "hunyuan-video-t2v-720p/vae") + + config = AutoencoderKLCausal3D.load_config(vae_path) + vae = AutoencoderKLCausal3D.from_config(config) + + vae_ckpt = Path(vae_path) / "pytorch_model.pt" + assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}" + + ckpt = torch.load(vae_ckpt, map_location=vae.device, weights_only=True) + if "state_dict" in ckpt: + ckpt = ckpt["state_dict"] + if any(k.startswith("vae.") for k in ckpt.keys()): + ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} + vae.load_state_dict(ckpt) + vae = vae.to(dtype=vae_precision) + vae.requires_grad_(False) + vae = vae.to("cuda") + vae.eval() + autocast_type = torch.float32 + fps = 24 + return vae, autocast_type, fps + + +def load_text_encoder(model_type, pretrained_model_name_or_path, device): + if model_type == "mochi": + text_encoder = MochiTextEncoderWrapper(pretrained_model_name_or_path, device) + elif model_type == "hunyuan" or "hunyuan_hf": + text_encoder = HunyuanTextEncoderWrapper(pretrained_model_name_or_path, device) + else: + raise ValueError(f"Unsupported model type: {model_type}") + return text_encoder + + +def get_no_split_modules(transformer): + # if of type MochiTransformer3DModel + if isinstance(transformer, MochiTransformer3DModel): + return (MochiTransformerBlock, ) + elif isinstance(transformer, HunyuanVideoTransformer3DModel): + return (HunyuanVideoSingleTransformerBlock, HunyuanVideoTransformerBlock) + elif isinstance(transformer, HYVideoDiffusionTransformer): + return (MMDoubleStreamBlock, MMSingleStreamBlock) + else: + raise ValueError(f"Unsupported transformer type: {type(transformer)}") + + +if __name__ == "__main__": + # test encode prompt + device = torch.cuda.current_device() + pretrained_model_name_or_path = "data/hunyuan" + text_encoder = load_text_encoder("hunyuan", pretrained_model_name_or_path, device) + prompt = "A man on stage claps his hands together while facing the audience. The audience, visible in the foreground, holds up mobile devices to record the event, capturing the moment from various angles. The background features a large banner with text identifying the man on stage. Throughout the sequence, the man's expression remains engaged and directed towards the audience. The camera angle remains constant, focusing on capturing the interaction between the man on stage and the audience." + prompt_embeds, attention_mask = text_encoder.encode_prompt(prompt) diff --git a/fastvideo/v1/v0_reference_src/utils/logging_.py b/fastvideo/v1/v0_reference_src/utils/logging_.py new file mode 100644 index 000000000..ee1523f87 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/logging_.py @@ -0,0 +1,24 @@ +import os +import pdb +import sys + + +def main_print(content): + if int(os.environ["LOCAL_RANK"]) <= 0: + print(content) + + +# ForkedPdb().set_trace() +class ForkedPdb(pdb.Pdb): + """A Pdb subclass that may be used + from a forked multiprocessing child + + """ + + def interaction(self, *args, **kwargs): + _stdin = sys.stdin + try: + sys.stdin = open("/dev/stdin") + pdb.Pdb.interaction(self, *args, **kwargs) + finally: + sys.stdin = _stdin diff --git a/fastvideo/v1/v0_reference_src/utils/optimizer.py b/fastvideo/v1/v0_reference_src/utils/optimizer.py new file mode 100644 index 000000000..661f30379 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/optimizer.py @@ -0,0 +1,68 @@ +import torch +from accelerate.logging import get_logger + +logger = get_logger(__name__) + + +def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False): + # Optimizer creation + supported_optimizers = ["adam", "adamw", "prodigy"] + if args.optimizer not in supported_optimizers: + logger.warning( + f"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW" + ) + args.optimizer = "adamw" + + if args.use_8bit_adam and not (args.optimizer.lower() not in ["adam", "adamw"]): + logger.warning(f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was " + f"set to {args.optimizer.lower()}") + + if args.use_8bit_adam: + try: + import bitsandbytes as bnb + except ImportError: + raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.") + + if args.optimizer.lower() == "adamw": + optimizer_class = (bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW) + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + eps=args.adam_epsilon, + weight_decay=args.adam_weight_decay, + ) + elif args.optimizer.lower() == "adam": + optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam + + optimizer = optimizer_class( + params_to_optimize, + betas=(args.adam_beta1, args.adam_beta2), + eps=args.adam_epsilon, + weight_decay=args.adam_weight_decay, + ) + elif args.optimizer.lower() == "prodigy": + try: + import prodigyopt + except ImportError: + raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`") + + optimizer_class = prodigyopt.Prodigy + + if args.learning_rate <= 0.1: + logger.warning( + "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0") + + optimizer = optimizer_class( + params_to_optimize, + lr=args.learning_rate, + betas=(args.adam_beta1, args.adam_beta2), + beta3=args.prodigy_beta3, + weight_decay=args.adam_weight_decay, + eps=args.adam_epsilon, + decouple=args.prodigy_decouple, + use_bias_correction=args.prodigy_use_bias_correction, + safeguard_warmup=args.prodigy_safeguard_warmup, + ) + + return optimizer diff --git a/fastvideo/v1/v0_reference_src/utils/parallel_states.py b/fastvideo/v1/v0_reference_src/utils/parallel_states.py new file mode 100644 index 000000000..c1f4cd751 --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/parallel_states.py @@ -0,0 +1,63 @@ +import os + +import torch.distributed as dist + + +class COMM_INFO: + + def __init__(self): + self.group = None + self.sp_size = 1 + self.global_rank = 0 + self.rank_within_group = 0 + self.group_id = 0 + + +nccl_info = COMM_INFO() +_SEQUENCE_PARALLEL_STATE = False + + +def initialize_sequence_parallel_state(sequence_parallel_size): + global _SEQUENCE_PARALLEL_STATE + if sequence_parallel_size > 1: + _SEQUENCE_PARALLEL_STATE = True + initialize_sequence_parallel_group(sequence_parallel_size) + else: + nccl_info.sp_size = 1 + nccl_info.global_rank = int(os.getenv("RANK", "0")) + nccl_info.rank_within_group = 0 + nccl_info.group_id = int(os.getenv("RANK", "0")) + + +def set_sequence_parallel_state(state): + global _SEQUENCE_PARALLEL_STATE + _SEQUENCE_PARALLEL_STATE = state + + +def get_sequence_parallel_state(): + return _SEQUENCE_PARALLEL_STATE + + +def initialize_sequence_parallel_group(sequence_parallel_size): + """Initialize the sequence parallel group.""" + 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) + nccl_info.sp_size = sequence_parallel_size + nccl_info.global_rank = rank + num_sequence_parallel_groups: int = world_size // sequence_parallel_size + for i in range(num_sequence_parallel_groups): + ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) + group = dist.new_group(ranks) + if rank in ranks: + nccl_info.group = group + nccl_info.rank_within_group = rank - i * sequence_parallel_size + nccl_info.group_id = i + + +def destroy_sequence_parallel_group(): + """Destroy the sequence parallel group.""" + dist.destroy_process_group() diff --git a/fastvideo/v1/v0_reference_src/utils/validation.py b/fastvideo/v1/v0_reference_src/utils/validation.py new file mode 100644 index 000000000..6d815a6ab --- /dev/null +++ b/fastvideo/v1/v0_reference_src/utils/validation.py @@ -0,0 +1,308 @@ +# isort: skip_file +import gc +import os +from typing import List, Optional, Union + +import numpy as np +import torch +from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import export_to_video +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from einops import rearrange +from tqdm import tqdm + +import wandb +from fastvideo.v1.distill.solver import PCMFMScheduler +from fastvideo.v1.models.mochi_hf.pipeline_mochi import (linear_quadratic_schedule, retrieve_timesteps) +from fastvideo.v1.utils.communications import all_gather +from fastvideo.v1.utils.load import load_vae +from fastvideo.v1.utils.parallel_states import (get_sequence_parallel_state, nccl_info) + + +def prepare_latents( + batch_size, + num_channels_latents, + height, + width, + num_frames, + dtype, + device, + generator, + vae_spatial_scale_factor, + vae_temporal_scale_factor, +): + height = height // vae_spatial_scale_factor + width = width // vae_spatial_scale_factor + num_frames = (num_frames - 1) // vae_temporal_scale_factor + 1 + + shape = (batch_size, num_channels_latents, num_frames, height, width) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + +def sample_validation_video( + model_type, + transformer, + vae, + scheduler, + scheduler_type="euler", + height: Optional[int] = None, + width: Optional[int] = None, + num_frames: int = 16, + num_inference_steps: int = 28, + timesteps: List[int] = None, + guidance_scale: float = 4.5, + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + vae_spatial_scale_factor=8, + vae_temporal_scale_factor=6, + num_channels_latents=12, +): + device = vae.device + + batch_size = prompt_embeds.shape[0] + + do_classifier_free_guidance = guidance_scale > 1.0 + if 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 + # TODO: Remove hardcore + latents = prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + prompt_embeds.dtype, + device, + generator, + vae_spatial_scale_factor, + vae_temporal_scale_factor, + ) + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + 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 + sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise) + sigmas = np.array(sigmas) + if scheduler_type == "euler" and model_type == "mochi": #todo + timesteps, num_inference_steps = retrieve_timesteps( + scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + ) + else: + timesteps, num_inference_steps = retrieve_timesteps( + scheduler, + num_inference_steps, + device, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0) + + # 6. Denoising loop + # with self.progress_bar(total=num_inference_steps) as progress_bar: + # write with tqdm instead + # only enable if nccl_info.global_rank == 0 + + with tqdm( + total=num_inference_steps, + disable=nccl_info.rank_within_group != 0, + desc="Validation sampling...", + ) as progress_bar: + for i, t in enumerate(timesteps): + latent_model_input = (torch.cat([latents] * 2) if do_classifier_free_guidance else latents) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + with torch.autocast("cuda", dtype=torch.bfloat16): + noise_pred = transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + timestep=timestep, + encoder_attention_mask=prompt_attention_mask, + return_dict=False, + )[0] + + # Mochi CFG + Sampling runs in FP32 + noise_pred = noise_pred.to(torch.float32) + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0] + latents = latents.to(latents_dtype) + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0): + progress_bar.update() + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + + if output_type == "latent": + video = latents + else: + # unscale/denormalize the latents + # denormalize with the mean and std if available and not None + has_latents_mean = (hasattr(vae.config, "latents_mean") and vae.config.latents_mean is not None) + has_latents_std = (hasattr(vae.config, "latents_std") and vae.config.latents_std is not None) + if has_latents_mean and has_latents_std: + latents_mean = (torch.tensor(vae.config.latents_mean).view(1, 12, 1, 1, + 1).to(latents.device, latents.dtype)) + latents_std = (torch.tensor(vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)) + latents = latents * latents_std / vae.config.scaling_factor + latents_mean + else: + latents = latents / vae.config.scaling_factor + with torch.autocast("cuda", dtype=vae.dtype): + video = vae.decode(latents, return_dict=False)[0] + video_processor = VideoProcessor(vae_scale_factor=vae_spatial_scale_factor) + video = video_processor.postprocess_video(video, output_type=output_type) + + return (video, ) + + +@torch.no_grad() +@torch.autocast("cuda", dtype=torch.bfloat16) +def log_validation( + args, + transformer, + device, + weight_dtype, # TODO + global_step, + scheduler_type="euler", + shift=1.0, + num_euler_timesteps=100, + linear_quadratic_threshold=0.025, + linear_range=0.5, + ema=False, +): + # TODO + print("Running validation....\n") + if args.model_type == "mochi": + vae_spatial_scale_factor = 8 + vae_temporal_scale_factor = 6 + num_channels_latents = 12 + elif args.model_type == "hunyuan" or "hunyuan_hf": + vae_spatial_scale_factor = 8 + vae_temporal_scale_factor = 4 + num_channels_latents = 16 + else: + raise ValueError(f"Model type {args.model_type} not supported") + vae, autocast_type, fps = load_vae(args.model_type, args.pretrained_model_name_or_path) + vae.enable_tiling() + if scheduler_type == "euler": + scheduler = FlowMatchEulerDiscreteScheduler(shift=shift) + else: + linear_quadraic = True if scheduler_type == "pcm_linear_quadratic" else False + scheduler = PCMFMScheduler( + 1000, + shift, + num_euler_timesteps, + linear_quadraic, + linear_quadratic_threshold, + linear_range, + ) + # args.validation_prompt_dir + + validation_guidance_scale_ls = args.validation_guidance_scale.split(",") + validation_guidance_scale_ls = [float(scale) for scale in validation_guidance_scale_ls] + for validation_sampling_step in args.validation_sampling_steps.split(","): + validation_sampling_step = int(validation_sampling_step) + for validation_guidance_scale in validation_guidance_scale_ls: + videos = [] + # prompt_embed are named embed0 to embedN + # check how many embeds are there + embe_dir = os.path.join(args.validation_prompt_dir, "prompt_embed") + mask_dir = os.path.join(args.validation_prompt_dir, "prompt_attention_mask") + embeds = sorted([f for f in os.listdir(embe_dir)]) + masks = sorted([f for f in os.listdir(mask_dir)]) + num_embeds = len(embeds) + validation_prompt_ids = list(range(num_embeds)) + num_sp_groups = int(os.getenv("WORLD_SIZE", "1")) // nccl_info.sp_size + # pad to multiple of groups + if num_embeds % num_sp_groups != 0: + validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups) + num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups + local_prompt_ids = validation_prompt_ids[nccl_info.group_id * + num_embeds_per_group:(nccl_info.group_id + 1) * + num_embeds_per_group] + + for i in local_prompt_ids: + prompt_embed_path = os.path.join(embe_dir, f"{embeds[i]}") + prompt_mask_path = os.path.join(mask_dir, f"{masks[i]}") + prompt_embeds = (torch.load(prompt_embed_path, map_location="cpu", + weights_only=True).to(device).unsqueeze(0)) + prompt_attention_mask = (torch.load(prompt_mask_path, map_location="cpu", + weights_only=True).to(device).unsqueeze(0)) + negative_prompt_embeds = torch.zeros(256, 4096).to(device).unsqueeze(0) + negative_prompt_attention_mask = (torch.zeros(256).bool().to(device).unsqueeze(0)) + generator = torch.Generator(device="cpu").manual_seed(12345) + video = sample_validation_video( + args.model_type, + transformer, + vae, + scheduler, + scheduler_type=scheduler_type, + num_frames=args.num_frames, + height=args.num_height, + width=args.num_width, + num_inference_steps=validation_sampling_step, + guidance_scale=validation_guidance_scale, + generator=generator, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_attention_mask=negative_prompt_attention_mask, + vae_spatial_scale_factor=vae_spatial_scale_factor, + vae_temporal_scale_factor=vae_temporal_scale_factor, + num_channels_latents=num_channels_latents, + )[0] + if nccl_info.rank_within_group == 0: + videos.append(video[0]) + # collect videos from all process to process zero + + gc.collect() + torch.cuda.empty_cache() + # log if main process + torch.distributed.barrier() + all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", "1")))] # remove padded videos + torch.distributed.all_gather_object(all_videos, videos) + if nccl_info.global_rank == 0: + # remove padding + videos = [video for videos in all_videos for video in videos] + videos = videos[:num_embeds] + # linearize all videos + video_filenames = [] + for i, video in enumerate(videos): + filename = os.path.join( + args.output_dir, + f"validation_step_{global_step}_sample_{validation_sampling_step}_guidance_{validation_guidance_scale}_video_{i}.mp4", + ) + export_to_video(video, filename, fps=fps) + video_filenames.append(filename) + + logs = { + f"{'ema_' if ema else ''}validation_sample_{validation_sampling_step}_guidance_{validation_guidance_scale}": + [wandb.Video(filename) for i, filename in enumerate(video_filenames)] + } + wandb.log(logs, step=global_step) diff --git a/pyproject.toml b/pyproject.toml index ac44c82c0..c59afaa69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,14 +18,14 @@ dependencies = [ "scipy==1.14.1", "six==1.16.0", "h5py==3.12.1", # Machine Learning & Transformers - "transformers", "tokenizers", "sentencepiece==0.2.0", + "transformers==4.46.1", "tokenizers==0.20.1", "sentencepiece==0.2.0", "timm==1.0.11", "peft==0.13.2", "diffusers==0.32.0", "bitsandbytes", # Acceleration & Optimization "accelerate==1.0.1", "torch", # Computer Vision & Image Processing - "opencv-python==4.10.0.84", "pillow", "imageio==2.36.0", + "opencv-python==4.10.0.84", "pillow==10.2.0", "imageio==2.36.0", "imageio-ffmpeg==0.5.1", "decord==0.6.0", "einops", # Experiment Tracking & Logging @@ -33,7 +33,7 @@ dependencies = [ # Miscellaneous Utilities "tqdm==4.66.5", "PyYAML==6.0.1", "idna==3.6", "protobuf==5.28.3", - "gradio", "huggingface_hub==0.26.1", "moviepy==1.0.3", "flask", "flask_restful", "aiohttp", + "gradio==5.3.0", "huggingface_hub==0.26.1", "moviepy==1.0.3", "flask", "flask_restful", "aiohttp", # System & Monitoring Tools "gpustat", "watch", @@ -76,4 +76,4 @@ skip ="./data,./wandb,./csrc/sliding_tile_attention/tk" column_limit = 120 [tool.isort] -line_length = 120 +line_length = 120 \ No newline at end of file diff --git a/scripts/inference/inference_hunyuan.sh b/scripts/inference/inference_hunyuan.sh index d42d008b8..2bf74a5e7 100644 --- a/scripts/inference/inference_hunyuan.sh +++ b/scripts/inference/inference_hunyuan.sh @@ -7,34 +7,34 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ --height 720 \ --width 1280 \ --num_frames 125 \ - --num_inference_steps 1 \ + --num_inference_steps 6 \ --guidance_scale 1 \ --embedded_cfg_scale 6 \ --flow_shift 17 \ --flow-reverse \ --prompt ./assets/prompt.txt \ --seed 1024 \ - --output_path outputs_video/hunyuan/vae_sp_v0_1/ \ + --output_path outputs_video/hunyuan/vae_sp/ \ --model_path $MODEL_BASE \ --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ --vae-sp # Inference original Hunyuan -# num_gpus=4 -# export MODEL_BASE=data/hunyuan -# torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ -# fastvideo/sample/sample_t2v_hunyuan.py \ -# --height 720 \ -# --width 1280 \ -# --num_frames 125 \ -# --num_inference_steps 50 \ -# --guidance_scale 1 \ -# --embedded_cfg_scale 6 \ -# --flow_shift 7 \ -# --flow-reverse \ -# --prompt ./assets/prompt.txt \ -# --seed 1024 \ -# --output_path outputs_video/hunyuan/vae_sp/ \ -# --model_path $MODEL_BASE \ -# --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ -# --vae-sp +num_gpus=4 +export MODEL_BASE=data/hunyuan +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/sample/sample_t2v_hunyuan.py \ + --height 720 \ + --width 1280 \ + --num_frames 125 \ + --num_inference_steps 50 \ + --guidance_scale 1 \ + --embedded_cfg_scale 6 \ + --flow_shift 7 \ + --flow-reverse \ + --prompt ./assets/prompt.txt \ + --seed 1024 \ + --output_path outputs_video/hunyuan/vae_sp/ \ + --model_path $MODEL_BASE \ + --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ + --vae-sp diff --git a/scripts/inference/v1_inference_hunyuan.sh b/scripts/inference/v1_inference_hunyuan.sh index 1b229ed8c..bab145435 100644 --- a/scripts/inference/v1_inference_hunyuan.sh +++ b/scripts/inference/v1_inference_hunyuan.sh @@ -1,10 +1,16 @@ #!/bin/bash num_gpus=4 -export MODEL_BASE=data/FastHunyuan +export MODEL_BASE=data/FastHunyuan-diffusers +# export MODEL_BASE=hunyuanvideo-community/HunyuanVideo +# Note that the tp_size and sp_size should be the same and equal to the number +# of GPUs. They are used for different parallel groups. sp_size is used for +# dit model and tp_size is used for encoder models. torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ - fastvideo/sample/v1_fastvideo_inference.py \ + fastvideo/v1/sample/v1_fastvideo_inference.py \ + --use-v1-transformer \ --sp_size 4 \ + --tp_size 4 \ --height 720 \ --width 1280 \ --num_frames 125 \ @@ -15,7 +21,6 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ --flow-reverse \ --prompt ./assets/prompt.txt \ --seed 1024 \ - --output_path outputs_video/hunyuan/vae_sp_v1/ \ + --output_path outputs_video/hunyuan/vae_sp_v2/ \ --model_path $MODEL_BASE \ - --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ --vae-sp diff --git a/scripts/inference/v2_inference_hunyuan.sh b/scripts/inference/v2_inference_hunyuan.sh deleted file mode 100644 index a0eda35d8..000000000 --- a/scripts/inference/v2_inference_hunyuan.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/bin/bash - -num_gpus=4 -export MODEL_BASE=data/FastHunyuan -torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ - fastvideo/sample/v2_fastvideo_inference.py \ - --sp_size 4 \ - --height 720 \ - --width 1280 \ - --num_frames 125 \ - --num_inference_steps 6 \ - --guidance_scale 1 \ - --embedded_cfg_scale 6 \ - --flow_shift 17 \ - --flow-reverse \ - --prompt ./assets/prompt.txt \ - --seed 1024 \ - --output_path outputs_video/hunyuan/vae_sp_v2/ \ - --model_path $MODEL_BASE \ - --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \ - --vae-sp