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

Merge branch 'rebased-refactor' of https://github.com/SolitaryThinker/FastVideo into rebased-refactor

Browse Source
This commit is contained in:
Peiyuan Zhang
2025-03-16 03:05:04 +00:00
parent 83348c6ed0 691f9d1064
commit 4e056b92cd
155 changed files with 14197 additions and 2724 deletions
Download Patch File Download Diff File Expand all files Collapse all files
assets/prompt.txt
+1 -1
View File
@@ -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
fastvideo/attention/__init__.py
-1
View File
@@ -1 +0,0 @@
from .distributed_attn import DistributedAttention
fastvideo/inference.py
-250
View File
@@ -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
fastvideo/layers/pooler.py
-322
View File
@@ -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)
fastvideo/layers/resampler.py
-269
View File
@@ -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
fastvideo/loader/loader.py
-452
View File
@@ -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)
fastvideo/model_loader.py
-273
View File
@@ -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]
fastvideo/models/__init__.py
-11
View File
@@ -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",
]
fastvideo/models/dits/__init__.py
-4
View File
@@ -1,4 +0,0 @@
from fastvideo.models.dits.hunyuanvideo import HunyuanVideoDiT
__all__ = ["HunyuanVideoDiT"]
fastvideo/models/hunyuan/inference.py
-22
View File
@@ -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":
fastvideo/models/hunyuan/modules/attenion.py
+5 -9
View File
@@ -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)
fastvideo/models/hunyuan/modules/models.py
+12 -2
View File
@@ -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:
fastvideo/models/hunyuan/modules/posemb_layers.py
-1
View File
@@ -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:
fastvideo/models/hunyuan/modules/token_refiner.py
+4 -1
View File
@@ -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
fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py
+2 -14
View File
@@ -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
fastvideo/pipelines/__init__.py
-70
View File
@@ -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",
]
fastvideo/pipelines/composed/text_to_video.py
-113
View File
@@ -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
fastvideo/pipelines/implementations/__init__.py
-8
View File
@@ -1,8 +0,0 @@
"""
Concrete implementations of diffusion pipelines.
This package contains concrete implementations of diffusion pipelines for
specific models or architectures.
"""
__all__ = []
fastvideo/pipelines/implementations/hunyuan/hunyuan_pipeline.py
-125
View File
@@ -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)
fastvideo/pipelines/pipeline_registry.py
-69
View File
@@ -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)
fastvideo/pipelines/stages/__init__.py
-28
View File
@@ -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",
]
fastvideo/sample/new_sample_hunyuan.py
-103
View File
@@ -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)
fastvideo/sample/sample_t2v_hunyuan.py
+7 -16
View File
@@ -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")
fastvideo/sample/v1_fastvideo_inference.py
-101
View File
@@ -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)
fastvideo/v1/attention/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .flash_attn import (DistributedAttention, LocalAttention)
__all__ = ["DistributedAttention", "LocalAttention"]
fastvideo/attention/distributed_attn.py → fastvideo/v1/attention/flash_attn.py
+7 -4
View File
@@ -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,
fastvideo/distributed/__init__.py → fastvideo/v1/distributed/__init__.py
View File
fastvideo/distributed/communication_op.py → fastvideo/v1/distributed/communication_op.py
View File
fastvideo/distributed/device_communicators/__init__.py → fastvideo/v1/distributed/device_communicators/__init__.py
View File
fastvideo/distributed/device_communicators/base_device_communicator.py → fastvideo/v1/distributed/device_communicators/base_device_communicator.py
View File
fastvideo/distributed/device_communicators/cuda_communicator.py → fastvideo/v1/distributed/device_communicators/cuda_communicator.py
+1 -1
View File
@@ -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
fastvideo/distributed/device_communicators/pynccl.py → fastvideo/v1/distributed/device_communicators/pynccl.py
+4 -4
View File
@@ -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__)
fastvideo/distributed/device_communicators/pynccl_wrapper.py → fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py
+2 -2
View File
@@ -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__)
fastvideo/distributed/parallel_state.py → fastvideo/v1/distributed/parallel_state.py
+281 -33
View File
@@ -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.
"""
fastvideo/distributed/utils.py → fastvideo/v1/distributed/utils.py
+2 -2
View File
@@ -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__)
fastvideo/envs.py → fastvideo/v1/envs.py
View File
fastvideo/inference_args.py → fastvideo/v1/inference_args.py
+40 -2
View File
@@ -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(
fastvideo/inference_engine.py → fastvideo/v1/inference_engine.py
+16 -17
View File
@@ -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,
fastvideo/layers/__init__.py → fastvideo/v1/layers/__init__.py
View File
fastvideo/layers/activation.py → fastvideo/v1/layers/activation.py
+1 -2
View File
@@ -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")
fastvideo/loader/__init__.py → fastvideo/v1/layers/conv.py
View File
fastvideo/layers/layernorm.py → fastvideo/v1/layers/layernorm.py
View File
fastvideo/layers/linear.py → fastvideo/v1/layers/linear.py
+4 -4
View File
@@ -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__)
fastvideo/layers/mlp.py → fastvideo/v1/layers/mlp.py
+2 -2
View File
@@ -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
fastvideo/layers/rotary_embedding.py → fastvideo/v1/layers/rotary_embedding.py
+40 -2
View File
@@ -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
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
fastvideo/layers/utils.py → fastvideo/v1/layers/utils.py
View File
fastvideo/layers/visual_embedding.py → fastvideo/v1/layers/visual_embedding.py
+3 -3
View File
@@ -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
fastvideo/layers/vocab_parallel_embedding.py → fastvideo/v1/layers/vocab_parallel_embedding.py
+4 -4
View File
@@ -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
fastvideo/v1/loader/__init__.py
View File
fastvideo/loader/fsdp_load.py → fastvideo/v1/loader/fsdp_load.py
+4 -5
View File
@@ -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(
fastvideo/v1/loader/loader.py
+342
View File
@@ -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()
fastvideo/v1/loader/utils.py
+184
View File
@@ -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": "<your-repo-name>--<config-name>",
# "AutoModel": "<your-repo-name>--<config-name>",
# "AutoModelFor<Task>": "<your-repo-name>--<config-name>",
# },
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
fastvideo/loader/weight_utils.py → fastvideo/v1/loader/weight_utils.py
+103 -103
View File
@@ -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(
fastvideo/logger.py → fastvideo/v1/logger.py
+5 -5
View File
@@ -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)
fastvideo/logging_utils/__init__.py → fastvideo/v1/logging_utils/__init__.py
+1 -1
View File
@@ -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",
fastvideo/logging_utils/formatter.py → fastvideo/v1/logging_utils/formatter.py
View File
fastvideo/v1/models/__init__.py
+57
View File
@@ -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",
]
fastvideo/models/dits/base.py → fastvideo/v1/models/dits/base.py
View File
fastvideo/models/dits/hunyuanvideo.py → fastvideo/v1/models/dits/hunyuanvideo.py
+12 -13
View File
@@ -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
fastvideo/models/encoders/clip.py → fastvideo/v1/models/encoders/clip.py
+252 -9
View File
@@ -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):
fastvideo/v1/models/encoders/llama.py
+428
View File
@@ -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
fastvideo/v1/models/encoders/utils.py
+438
View File
@@ -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}"
fastvideo/models/encoders/vision.py → fastvideo/v1/models/encoders/vision.py
+3 -3
View File
@@ -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__)
fastvideo/v1/models/hf_transformer_utils.py
+259
View File
@@ -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"
fastvideo/v1/models/loader.py
+325
View File
@@ -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)
fastvideo/models/parameter.py → fastvideo/v1/models/parameter.py
+4 -4
View File
@@ -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
fastvideo/models/dits/registry.py → fastvideo/v1/models/registry.py
+30 -12
View File
@@ -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()
})
fastvideo/v1/models/schedulers/base.py
View File
fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py
+239
View File
@@ -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
fastvideo/v1/models/text_encoder.py
+255
View File
@@ -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,
)
fastvideo/models/utils.py → fastvideo/v1/models/utils.py
+45 -2
View File
@@ -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)
fastvideo/v1/models/vaes/common.py
+469
View File
@@ -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
fastvideo/v1/models/vaes/hunyuanvae.py
+799
View File
@@ -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
fastvideo/v1/pipelines/__init__.py
+244
View File
@@ -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",
]
fastvideo/pipelines/composed/__init__.py → fastvideo/v1/pipelines/composed/__init__.py
+1 -1
View File
@@ -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,
)
fastvideo/pipelines/composed/composed_pipeline_base.py → fastvideo/v1/pipelines/composed/composed_pipeline_base.py
+12 -4
View File
@@ -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.
fastvideo/pipelines/implementations/hunyuan/__init__.py → fastvideo/v1/pipelines/implementations/hunyuan/__init__.py
+3 -3
View File
@@ -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",
]
fastvideo/pipelines/hunyuan/constants.py → fastvideo/v1/pipelines/implementations/hunyuan/constants.py
View File
fastvideo/pipelines/implementations/hunyuan/hunyuan_loader.py → fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_loader.py
+60 -16
View File
@@ -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
fastvideo/v1/pipelines/implementations/hunyuan/hunyuan_pipeline.py
+245
View File
@@ -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
fastvideo/pipelines/pipeline_batch_info.py → fastvideo/v1/pipelines/pipeline_batch_info.py
View File
fastvideo/v1/pipelines/pipeline_registry.py
+91
View File
@@ -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())
fastvideo/v1/pipelines/stages/__init__.py
+28
View File
@@ -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",
]
fastvideo/pipelines/stages/base.py → fastvideo/v1/pipelines/stages/base.py
+4 -4
View File
@@ -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:
fastvideo/pipelines/stages/conditioning.py → fastvideo/v1/pipelines/stages/conditioning.py
+4 -4
View File
@@ -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__)
fastvideo/pipelines/stages/decoding.py → fastvideo/v1/pipelines/stages/decoding.py
+9 -27
View File
@@ -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)
fastvideo/pipelines/stages/denoising.py → fastvideo/v1/pipelines/stages/denoising.py
+6 -5
View File
@@ -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__)
fastvideo/pipelines/stages/input_validation.py → fastvideo/v1/pipelines/stages/input_validation.py
+4 -4
View File
@@ -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__)
fastvideo/pipelines/stages/latent_preparation.py → fastvideo/v1/pipelines/stages/latent_preparation.py
+4 -4
View File
@@ -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__)
fastvideo/pipelines/stages/post_processing.py → fastvideo/v1/pipelines/stages/post_processing.py
+4 -4
View File
@@ -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__)
fastvideo/pipelines/stages/prompt_encoding.py → fastvideo/v1/pipelines/stages/prompt_encoding.py
+12 -6
View File
@@ -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
fastvideo/pipelines/stages/timestep_preparation.py → fastvideo/v1/pipelines/stages/timestep_preparation.py
+4 -4
View File
@@ -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__)
fastvideo/platforms/__init__.py → fastvideo/v1/platforms/__init__.py
View File
fastvideo/platforms/cuda.py → fastvideo/v1/platforms/cuda.py
+4 -4
View File
@@ -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
fastvideo/platforms/interface.py → fastvideo/v1/platforms/interface.py
+2 -2
View File
@@ -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):
fastvideo/sample/v2_fastvideo_inference.py → fastvideo/v1/sample/v1_fastvideo_inference.py
+8 -17
View File
@@ -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()
fastvideo/v1/distributed_example.py → fastvideo/v1/tests/distributed_example.py
+1 -1
View File
@@ -1,4 +1,4 @@
from fastvideo.logger import init_logger
from fastvideo.v1.logger import init_logger
logger = init_logger(__name__)
import os
fastvideo/v1/distributed_run.sh → fastvideo/v1/tests/distributed_run.sh
View File
fastvideo/v1/tests/test_clip_encoder.py
+207
View File
@@ -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()
fastvideo/v1/tests/test_hunyuan_vae.py
+161
View File
@@ -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()

Some files were not shown because too many files have changed in this diff Show More