diff --git a/fastvideo/models/hunyuan/__pycache__/__init__.cpython-310.pyc b/fastvideo/models/hunyuan/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ed77d81592dc9454f8d489c4e6e313df7890457 Binary files /dev/null and b/fastvideo/models/hunyuan/__pycache__/__init__.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/__pycache__/__init__.cpython-312.pyc b/fastvideo/models/hunyuan/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7112ab999d006eb8ad2eef1472cbd88ac7cdabf6 Binary files /dev/null and b/fastvideo/models/hunyuan/__pycache__/__init__.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/__pycache__/constants.cpython-310.pyc b/fastvideo/models/hunyuan/__pycache__/constants.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55ed82665a0c348a8092a91825ac43f78ef876ce Binary files /dev/null and b/fastvideo/models/hunyuan/__pycache__/constants.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/__pycache__/constants.cpython-312.pyc b/fastvideo/models/hunyuan/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97fb45f98fa55cc0b8c8202a4f0296cd5dcc4520 Binary files /dev/null and b/fastvideo/models/hunyuan/__pycache__/constants.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/diffusion/__init__.py b/fastvideo/models/hunyuan/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..754bc48ec2d61ebd4bc794e9bc0e87977b011b49 --- /dev/null +++ b/fastvideo/models/hunyuan/diffusion/__init__.py @@ -0,0 +1,3 @@ +# ruff: noqa: F401 +from .pipelines import HunyuanVideoPipeline +from .schedulers import FlowMatchDiscreteScheduler diff --git a/fastvideo/models/hunyuan/diffusion/pipelines/__init__.py b/fastvideo/models/hunyuan/diffusion/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fa109e14d0f65e3197e3048c2fe590edb3fe2e67 --- /dev/null +++ b/fastvideo/models/hunyuan/diffusion/pipelines/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa: F401 +from .pipeline_hunyuan_video import HunyuanVideoPipeline diff --git a/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py b/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py new file mode 100644 index 0000000000000000000000000000000000000000..d7e90f678a9430d8cd40b2cadfe65e2690e27705 --- /dev/null +++ b/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py @@ -0,0 +1,1010 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +import inspect +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.configuration_utils import FrozenDict +from diffusers.image_processor import VaeImageProcessor +from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin +from diffusers.models import AutoencoderKL +from diffusers.models.lora import adjust_lora_scale_text_encoder +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, logging, + replace_example_docstring, scale_lora_layers) +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange + +from fastvideo.utils.communications import all_gather +from fastvideo.utils.parallel_states import (get_sequence_parallel_state, + nccl_info) + +from ...constants import PRECISION_TO_TYPE +from ...modules import HYVideoDiffusionTransformer +from ...text_encoder import TextEncoder +from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """""" + + +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), + keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = (guidance_rescale * noise_pred_rescaled + + (1 - guidance_rescale) * noise_cfg) + return noise_cfg + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError( + "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values" + ) + if timesteps is not None: + accepts_timesteps = "timesteps" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set( + inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +@dataclass +class HunyuanVideoPipelineOutput(BaseOutput): + videos: Union[torch.Tensor, np.ndarray] + + +class HunyuanVideoPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-video generation using HunyuanVideo. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + vae ([`AutoencoderKL`]): + Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. + text_encoder ([`TextEncoder`]): + Frozen text-encoder. + text_encoder_2 ([`TextEncoder`]): + Frozen text-encoder_2. + transformer ([`HYVideoDiffusionTransformer`]): + A `HYVideoDiffusionTransformer` to denoise the encoded video latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `unet` to denoise the encoded image latents. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae" + _optional_components = ["text_encoder_2"] + _exclude_from_cpu_offload = ["transformer"] + _callback_tensor_inputs = [ + "latents", "prompt_embeds", "negative_prompt_embeds" + ] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: TextEncoder, + transformer: HYVideoDiffusionTransformer, + scheduler: KarrasDiffusionSchedulers, + text_encoder_2: Optional[TextEncoder] = None, + progress_bar_config: Dict[str, Any] = None, + args=None, + ): + super().__init__() + + # ========================================================================================== + if progress_bar_config is None: + progress_bar_config = {} + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + self._progress_bar_config.update(progress_bar_config) + + self.args = args + # ========================================================================================== + + if (hasattr(scheduler.config, "steps_offset") + and scheduler.config.steps_offset != 1): + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`" + f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " + "to update the config accordingly as leaving `steps_offset` might led to incorrect results" + " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub," + " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`" + " file") + deprecate("steps_offset!=1", + "1.0.0", + deprecation_message, + standard_warn=False) + new_config = dict(scheduler.config) + new_config["steps_offset"] = 1 + scheduler._internal_dict = FrozenDict(new_config) + + if (hasattr(scheduler.config, "clip_sample") + and scheduler.config.clip_sample is True): + deprecation_message = ( + f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`." + " `clip_sample` should be set to False in the configuration file. Please make sure to update the" + " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in" + " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very" + " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file" + ) + deprecate("clip_sample not set", + "1.0.0", + deprecation_message, + standard_warn=False) + new_config = dict(scheduler.config) + new_config["clip_sample"] = False + scheduler._internal_dict = FrozenDict(new_config) + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + transformer=transformer, + scheduler=scheduler, + text_encoder_2=text_encoder_2, + ) + self.vae_scale_factor = 2**(len(self.vae.config.block_out_channels) - + 1) + self.image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor) + + def encode_prompt( + self, + prompt, + device, + num_videos_per_prompt, + do_classifier_free_guidance, + negative_prompt=None, + prompt_embeds: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_attention_mask: Optional[torch.Tensor] = None, + lora_scale: Optional[float] = None, + clip_skip: Optional[int] = None, + text_encoder: Optional[TextEncoder] = None, + data_type: Optional[str] = "image", + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + device: (`torch.device`): + torch device + num_videos_per_prompt (`int`): + number of videos that should be generated per prompt + do_classifier_free_guidance (`bool`): + whether to use classifier free guidance or not + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the video generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + attention_mask (`torch.Tensor`, *optional*): + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + negative_attention_mask (`torch.Tensor`, *optional*): + lora_scale (`float`, *optional*): + A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + text_encoder (TextEncoder, *optional*): + data_type (`str`, *optional*): + """ + if text_encoder is None: + text_encoder = self.text_encoder + + # set lora scale so that monkey patched LoRA + # function of text encoder can correctly access it + if lora_scale is not None and isinstance(self, LoraLoaderMixin): + self._lora_scale = lora_scale + + # dynamically adjust the LoRA scale + if not USE_PEFT_BACKEND: + adjust_lora_scale_text_encoder(text_encoder.model, lora_scale) + else: + scale_lora_layers(text_encoder.model, lora_scale) + + if prompt_embeds is None: + # textual inversion: process multi-vector tokens if necessary + if isinstance(self, TextualInversionLoaderMixin): + prompt = self.maybe_convert_prompt(prompt, + text_encoder.tokenizer) + + text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) + if clip_skip is None: + prompt_outputs = text_encoder.encode(text_inputs, + data_type=data_type, + device=device) + prompt_embeds = prompt_outputs.hidden_state + else: + prompt_outputs = text_encoder.encode( + text_inputs, + output_hidden_states=True, + data_type=data_type, + device=device, + ) + # Access the `hidden_states` first, that contains a tuple of + # all the hidden states from the encoder layers. Then index into + # the tuple to access the hidden states from the desired layer. + prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + + 1)] + # We also need to apply the final LayerNorm here to not mess with the + # representations. The `last_hidden_states` that we typically use for + # obtaining the final prompt representations passes through the LayerNorm + # layer. + prompt_embeds = text_encoder.model.text_model.final_layer_norm( + prompt_embeds) + + attention_mask = prompt_outputs.attention_mask + if attention_mask is not None: + attention_mask = attention_mask.to(device) + bs_embed, seq_len = attention_mask.shape + attention_mask = attention_mask.repeat(1, + num_videos_per_prompt) + attention_mask = attention_mask.view( + bs_embed * num_videos_per_prompt, seq_len) + + if text_encoder is not None: + prompt_embeds_dtype = text_encoder.dtype + elif self.transformer is not None: + prompt_embeds_dtype = self.transformer.dtype + else: + prompt_embeds_dtype = prompt_embeds.dtype + + prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, + device=device) + + if prompt_embeds.ndim == 2: + bs_embed, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt) + prompt_embeds = prompt_embeds.view( + bs_embed * num_videos_per_prompt, -1) + else: + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view( + bs_embed * num_videos_per_prompt, seq_len, -1) + + return ( + prompt_embeds, + negative_prompt_embeds, + attention_mask, + negative_attention_mask, + ) + + def decode_latents(self, latents, enable_tiling=True): + deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead" + deprecate("decode_latents", + "1.0.0", + deprecation_message, + standard_warn=False) + + latents = 1 / self.vae.config.scaling_factor * latents + if enable_tiling: + self.vae.enable_tiling() + image = self.vae.decode(latents, return_dict=False)[0] + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16 + if image.ndim == 4: + image = image.cpu().permute(0, 2, 3, 1).float() + else: + image = image.cpu().float() + return image + + def prepare_extra_func_kwargs(self, func, kwargs): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + extra_step_kwargs = {} + + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_step_kwargs[k] = v + return extra_step_kwargs + + def check_inputs( + self, + prompt, + height, + width, + video_length, + callback_steps, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + vae_ver="88-4c-sd", + ): + if height % 8 != 0 or width % 8 != 0: + raise ValueError( + f"`height` and `width` have to be divisible by 8 but are {height} and {width}." + ) + + if video_length is not None: + if "884" in vae_ver: + if video_length != 1 and (video_length - 1) % 4 != 0: + raise ValueError( + f"`video_length` has to be 1 or a multiple of 4 but is {video_length}." + ) + elif "888" in vae_ver: + if video_length != 1 and (video_length - 1) % 8 != 0: + raise ValueError( + f"`video_length` has to be 1 or a multiple of 8 but is {video_length}." + ) + + if callback_steps is not None and (not isinstance(callback_steps, int) + or callback_steps <= 0): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}.") + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs + for k in callback_on_step_end_tensor_inputs): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two.") + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) + and not isinstance(prompt, list)): + raise ValueError( + f"`prompt` has to be of type `str` or `list` but is {type(prompt)}" + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}.") + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + video_length, + dtype, + device, + generator, + latents=None, + ): + shape = ( + batch_size, + num_channels_latents, + video_length, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + latents = randn_tensor(shape, + generator=generator, + device=device, + dtype=dtype) + else: + latents = latents.to(device) + + # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler + if hasattr(self.scheduler, "init_noise_sigma"): + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + return latents + + # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding + def get_guidance_scale_embedding( + self, + w: torch.Tensor, + embedding_dim: int = 512, + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """ + See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298 + + Args: + w (`torch.Tensor`): + Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings. + embedding_dim (`int`, *optional*, defaults to 512): + Dimension of the embeddings to generate. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + Data type of the generated embeddings. + + Returns: + `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`. + """ + assert len(w.shape) == 1 + w = w * 1000.0 + + half_dim = embedding_dim // 2 + emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb) + emb = w.to(dtype)[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: # zero pad + emb = torch.nn.functional.pad(emb, (0, 1)) + assert emb.shape == (w.shape[0], embedding_dim) + return emb + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def clip_skip(self): + return self._clip_skip + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + # return self._guidance_scale > 1 and self.transformer.config.time_cond_proj_dim is None + return self._guidance_scale > 1 + + @property + def cross_attention_kwargs(self): + return self._cross_attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]], + height: int, + width: int, + video_length: int, + data_type: str = "video", + num_inference_steps: int = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, + List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_attention_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Union[Callable[[int, int, Dict], + None], PipelineCallback, + MultiPipelineCallbacks, ]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + vae_ver: str = "88-4c-sd", + enable_tiling: bool = False, + enable_vae_sp: bool = False, + n_tokens: Optional[int] = None, + embedded_guidance_scale: Optional[float] = None, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`): + The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + video_length (`int`): + The number of frames in the generated video. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide what to not include in image generation. If not defined, you need to + pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies + to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor is generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If + not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. + + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a + plain tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in + [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + clip_skip (`int`, *optional*): + Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that + the output of the pre-final layer will be used for computing the prompt embeddings. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~HunyuanVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated images and the + second element is a list of `bool`s indicating whether the corresponding generated image contains + "not-safe-for-work" (nsfw) content. + """ + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", + ) + + if isinstance(callback_on_step_end, + (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 0. Default height and width to unet + # height = height or self.transformer.config.sample_size * self.vae_scale_factor + # width = width or self.transformer.config.sample_size * self.vae_scale_factor + # to deal with lora scaling and other possible forward hooks + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + video_length, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + callback_on_step_end_tensor_inputs, + vae_ver=vae_ver, + ) + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = (torch.device(f"cuda:{dist.get_rank()}") + if dist.is_initialized() else self._execution_device) + + # 3. Encode input prompt + lora_scale = (self.cross_attention_kwargs.get("scale", None) + if self.cross_attention_kwargs is not None else None) + + ( + prompt_embeds, + negative_prompt_embeds, + prompt_mask, + negative_prompt_mask, + ) = self.encode_prompt( + prompt, + device, + num_videos_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + attention_mask=attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_attention_mask=negative_attention_mask, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + data_type=data_type, + ) + if self.text_encoder_2 is not None: + ( + prompt_embeds_2, + negative_prompt_embeds_2, + prompt_mask_2, + negative_prompt_mask_2, + ) = self.encode_prompt( + prompt, + device, + num_videos_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=None, + attention_mask=None, + negative_prompt_embeds=None, + negative_attention_mask=None, + lora_scale=lora_scale, + clip_skip=self.clip_skip, + text_encoder=self.text_encoder_2, + data_type=data_type, + ) + else: + prompt_embeds_2 = None + negative_prompt_embeds_2 = None + prompt_mask_2 = None + negative_prompt_mask_2 = None + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + if prompt_mask is not None: + prompt_mask = torch.cat([negative_prompt_mask, prompt_mask]) + if prompt_embeds_2 is not None: + prompt_embeds_2 = torch.cat( + [negative_prompt_embeds_2, prompt_embeds_2]) + if prompt_mask_2 is not None: + prompt_mask_2 = torch.cat( + [negative_prompt_mask_2, prompt_mask_2]) + + # 4. Prepare timesteps + extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.set_timesteps, {"n_tokens": n_tokens}) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + **extra_set_timesteps_kwargs, + ) + if "884" in vae_ver: + video_length = (video_length - 1) // 4 + 1 + elif "888" in vae_ver: + video_length = (video_length - 1) // 8 + 1 + else: + video_length = video_length + + # 5. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + video_length, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + if get_sequence_parallel_state(): + latents = rearrange(latents, + "b t (n s) h w -> b t n s h w", + n=world_size).contiguous() + latents = latents[:, :, rank, :, :, :] + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + { + "generator": generator, + "eta": eta + }, + ) + + target_dtype = PRECISION_TO_TYPE[self.args.precision] + autocast_enabled = (target_dtype != + torch.float32) and not self.args.disable_autocast + vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not self.args.disable_autocast + + # 7. Denoising loop + num_warmup_steps = len( + timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + # if is_progress_bar: + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = (torch.cat( + [latents] * + 2) if self.do_classifier_free_guidance else latents) + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t) + + t_expand = t.repeat(latent_model_input.shape[0]) + guidance_expand = (torch.tensor( + [embedded_guidance_scale] * latent_model_input.shape[0], + dtype=torch.float32, + device=device, + ).to(target_dtype) * 1000.0 if embedded_guidance_scale + is not None else None) + # predict the noise residual + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled): + # concat prompt_embeds_2 and prompt_embeds. Mismatch fill with zeros + if prompt_embeds_2.shape[-1] != prompt_embeds.shape[-1]: + prompt_embeds_2 = F.pad( + prompt_embeds_2, + (0, prompt_embeds.shape[2] - + prompt_embeds_2.shape[1]), + value=0, + ).unsqueeze(1) + encoder_hidden_states = torch.cat( + [prompt_embeds_2, prompt_embeds], dim=1) + noise_pred = self.transformer( # For an input image (129, 192, 336) (1, 256, 256) + latent_model_input, # [2, 16, 33, 24, 42] + encoder_hidden_states, + t_expand, # [2] + prompt_mask, # [2, 256]fpdb + guidance=guidance_expand, + return_dict=False, + )[0] + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * ( + noise_pred_text - noise_pred_uncond) + + if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg( + noise_pred, + noise_pred_text, + guidance_rescale=self.guidance_rescale, + ) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, + t, + latents, + **extra_step_kwargs, + return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end( + self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop( + "prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and + (i + 1) % self.scheduler.order == 0): + if progress_bar is not None: + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + if get_sequence_parallel_state(): + latents = all_gather(latents, dim=2) + + if not output_type == "latent": + expand_temporal_dim = False + if len(latents.shape) == 4: + if isinstance(self.vae, AutoencoderKLCausal3D): + latents = latents.unsqueeze(2) + expand_temporal_dim = True + elif len(latents.shape) == 5: + pass + else: + raise ValueError( + f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}." + ) + + if (hasattr(self.vae.config, "shift_factor") + and self.vae.config.shift_factor): + latents = (latents / self.vae.config.scaling_factor + + self.vae.config.shift_factor) + else: + latents = latents / self.vae.config.scaling_factor + + with torch.autocast(device_type="cuda", + dtype=vae_dtype, + enabled=vae_autocast_enabled): + if enable_tiling: + self.vae.enable_tiling() + if enable_vae_sp: + self.vae.enable_parallel() + image = self.vae.decode(latents, + return_dict=False, + generator=generator)[0] + + if expand_temporal_dim or image.shape[2] == 1: + image = image.squeeze(2) + + else: + image = latents + + image = (image / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 + image = image.cpu().float() + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return image + + return HunyuanVideoPipelineOutput(videos=image) diff --git a/fastvideo/models/hunyuan/diffusion/schedulers/__init__.py b/fastvideo/models/hunyuan/diffusion/schedulers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2238803fe36dae8bc019e5f50e4cb889f03b72eb --- /dev/null +++ b/fastvideo/models/hunyuan/diffusion/schedulers/__init__.py @@ -0,0 +1,2 @@ +# ruff: noqa: F401 +from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler diff --git a/fastvideo/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py b/fastvideo/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..69ed5317e6e76ee4c31aa75a58d6c6b8a7a36cc5 --- /dev/null +++ b/fastvideo/models/hunyuan/diffusion/schedulers/scheduling_flow_match_discrete.py @@ -0,0 +1,248 @@ +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, logging + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + reverse (`bool`, defaults to `True`): + Whether to reverse the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + reverse: bool = True, + solver: str = "euler", + n_tokens: Optional[int] = None, + ): + sigmas = torch.linspace(1, 0, num_train_timesteps + 1) + + if not reverse: + sigmas = sigmas.flip(0) + + self.sigmas = sigmas + # the value fed to model + self.timesteps = (sigmas[:-1] * + num_train_timesteps).to(dtype=torch.float32) + + self._step_index = None + self._begin_index = None + + self.supported_solver = ["euler"] + if solver not in self.supported_solver: + raise ValueError( + f"Solver {solver} not supported. Supported solvers: {self.supported_solver}" + ) + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device] = None, + n_tokens: int = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + """ + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1) + sigmas = self.sd3_time_shift(sigmas) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to( + dtype=torch.float32, device=device) + + # Reset step index + self._step_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def scale_model_input(self, + sample: torch.Tensor, + timestep: Optional[int] = None) -> torch.Tensor: + return sample + + def sd3_time_shift(self, t: torch.Tensor): + return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if (isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor)): + raise ValueError(( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep."), ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.to(torch.float32) * dt + else: + raise ValueError( + f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}" + ) + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return (prev_sample, ) + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/fastvideo/models/hunyuan/modules/__init__.py b/fastvideo/models/hunyuan/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c85b51a4cdd41ffa8f936dab1c7a401ad7c5c0b6 --- /dev/null +++ b/fastvideo/models/hunyuan/modules/__init__.py @@ -0,0 +1,25 @@ +from .models import HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer + + +def load_model(args, in_channels, out_channels, factor_kwargs): + """load hunyuan video model + + Args: + args (dict): model args + in_channels (int): input channels number + out_channels (int): output channels number + factor_kwargs (dict): factor kwargs + + Returns: + model (nn.Module): The hunyuan video model + """ + if args.model in HUNYUAN_VIDEO_CONFIG.keys(): + model = HYVideoDiffusionTransformer( + in_channels=in_channels, + out_channels=out_channels, + **HUNYUAN_VIDEO_CONFIG[args.model], + **factor_kwargs, + ) + return model + else: + raise NotImplementedError() diff --git a/fastvideo/models/hunyuan/modules/__pycache__/activation_layers.cpython-310.pyc b/fastvideo/models/hunyuan/modules/__pycache__/activation_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46fdd8f17fc7910e8fb0d48881b741fec0114297 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/activation_layers.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/activation_layers.cpython-312.pyc b/fastvideo/models/hunyuan/modules/__pycache__/activation_layers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4186ea5b9d56f58fae9c0a62d5726b2b78d0c7f7 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/activation_layers.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/attenion.cpython-312.pyc b/fastvideo/models/hunyuan/modules/__pycache__/attenion.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e1eccb6de03bd7f457f8a8517d4e9d83ebdf579 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/attenion.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/embed_layers.cpython-310.pyc b/fastvideo/models/hunyuan/modules/__pycache__/embed_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03e5a22ce6ff74d240e0f51853f5ac658bceac73 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/embed_layers.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/mlp_layers.cpython-310.pyc b/fastvideo/models/hunyuan/modules/__pycache__/mlp_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ec9b1cad2773b27b8fb8afd0f4ff7b1d5f35b73 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/mlp_layers.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/mlp_layers.cpython-312.pyc b/fastvideo/models/hunyuan/modules/__pycache__/mlp_layers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1d8b8b9778f7df51e5c9352f5b37ca1e67af525 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/mlp_layers.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/modulate_layers.cpython-310.pyc b/fastvideo/models/hunyuan/modules/__pycache__/modulate_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9b9d106296cd324d62446fc9aff120d79199c94 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/modulate_layers.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/modulate_layers.cpython-312.pyc b/fastvideo/models/hunyuan/modules/__pycache__/modulate_layers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f436fb924245f76c26760bf96466e3ec15d3ae5 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/modulate_layers.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/norm_layers.cpython-310.pyc b/fastvideo/models/hunyuan/modules/__pycache__/norm_layers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7818003b0bef2c5334e650287576e26f6376203 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/norm_layers.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/posemb_layers.cpython-312.pyc b/fastvideo/models/hunyuan/modules/__pycache__/posemb_layers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ded442594cad702a778eac76edf0874cbd6a1e98 Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/posemb_layers.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/modules/__pycache__/token_refiner.cpython-310.pyc b/fastvideo/models/hunyuan/modules/__pycache__/token_refiner.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a20a4d5fd7d749a9c52a591cf2fd8e2ff373ed9f Binary files /dev/null and b/fastvideo/models/hunyuan/modules/__pycache__/token_refiner.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/modules/activation_layers.py b/fastvideo/models/hunyuan/modules/activation_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..f8774c26ceef6081482ca0dbbf930b207d4ac03b --- /dev/null +++ b/fastvideo/models/hunyuan/modules/activation_layers.py @@ -0,0 +1,23 @@ +import torch.nn as nn + + +def get_activation_layer(act_type): + """get activation layer + + Args: + act_type (str): the activation type + + Returns: + torch.nn.functional: the activation layer + """ + if act_type == "gelu": + return lambda: nn.GELU() + elif act_type == "gelu_tanh": + # Approximate `tanh` requires torch >= 1.13 + return lambda: nn.GELU(approximate="tanh") + elif act_type == "relu": + return nn.ReLU + elif act_type == "silu": + return nn.SiLU + else: + raise ValueError(f"Unknown activation type: {act_type}") diff --git a/fastvideo/models/hunyuan/modules/attenion.py b/fastvideo/models/hunyuan/modules/attenion.py new file mode 100644 index 0000000000000000000000000000000000000000..a010388357a4a17a9f69205614fca7178918e82f --- /dev/null +++ b/fastvideo/models/hunyuan/modules/attenion.py @@ -0,0 +1,90 @@ +import torch +import torch.nn.functional as F + +from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad +from fastvideo.utils.communications import all_gather, all_to_all_4D +from fastvideo.utils.parallel_states import (get_sequence_parallel_state, + nccl_info) + + +def attention( + q, + k, + v, + drop_rate=0, + attn_mask=None, + causal=False, +): + + qkv = torch.stack([q, k, v], dim=2) + + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.bool() + + x = flash_attn_no_pad(qkv, + attn_mask, + causal=causal, + dropout_p=drop_rate, + softmax_scale=None) + + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out + + +def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask): + # 1GPU torch.Size([1, 11264, 24, 128]) tensor([ 0, 11275, 11520], device='cuda:0', dtype=torch.int32) + # 2GPU torch.Size([1, 5632, 24, 128]) tensor([ 0, 5643, 5888], device='cuda:0', dtype=torch.int32) + query, encoder_query = q + key, encoder_key = k + value, encoder_value = v + if get_sequence_parallel_state(): + # batch_size, seq_len, attn_heads, head_dim + query = all_to_all_4D(query, scatter_dim=2, gather_dim=1) + key = all_to_all_4D(key, scatter_dim=2, gather_dim=1) + value = all_to_all_4D(value, scatter_dim=2, gather_dim=1) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow( + dim, nccl_info.rank_within_group * local_heads, local_heads) + + encoder_query = shrink_head(encoder_query, dim=2) + encoder_key = shrink_head(encoder_key, dim=2) + encoder_value = shrink_head(encoder_value, dim=2) + # [b, s, h, d] + + sequence_length = query.size(1) + encoder_sequence_length = encoder_query.size(1) + + # Hint: please check encoder_query.shape + 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) + + 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, encoder_hidden_states = hidden_states.split_with_sizes( + (sequence_length, encoder_sequence_length), dim=1) + if get_sequence_parallel_state(): + hidden_states = all_to_all_4D(hidden_states, + scatter_dim=1, + gather_dim=2) + encoder_hidden_states = all_gather(encoder_hidden_states, + dim=2).contiguous() + hidden_states = hidden_states.to(query.dtype) + encoder_hidden_states = encoder_hidden_states.to(query.dtype) + + attn = torch.cat([hidden_states, encoder_hidden_states], dim=1) + + b, s, a, d = attn.shape + attn = attn.reshape(b, s, -1) + + return attn diff --git a/fastvideo/models/hunyuan/modules/embed_layers.py b/fastvideo/models/hunyuan/modules/embed_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..d2cb9bb5a6544b05a8a8a418ff6bdbaaadeabf79 --- /dev/null +++ b/fastvideo/models/hunyuan/modules/embed_layers.py @@ -0,0 +1,163 @@ +import math + +import torch +import torch.nn as nn + +from ..utils.helpers import to_2tuple + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding + + Image to Patch Embedding using Conv2d + + A convolution based approach to patchifying a 2D image w/ embedding projection. + + Based on the impl in https://github.com/google-research/vision_transformer + + Hacked together by / Copyright 2020 Ross Wightman + + Remove the _assert function in forward function to be compatible with multi-resolution images. + """ + + def __init__( + self, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True, + bias=True, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + self.flatten = flatten + + self.proj = nn.Conv3d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + **factory_kwargs, + ) + nn.init.xavier_uniform_( + self.proj.weight.view(self.proj.weight.size(0), -1)) + if bias: + nn.init.zeros_(self.proj.bias) + + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class TextProjection(nn.Module): + """ + Projects text embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, + in_channels, + hidden_size, + act_layer, + dtype=None, + device=None): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.linear_1 = nn.Linear( + in_features=in_channels, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + self.act_1 = act_layer() + self.linear_2 = nn.Linear( + in_features=hidden_size, + out_features=hidden_size, + bias=True, + **factory_kwargs, + ) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * + torch.arange(start=0, end=half, dtype=torch.float32) / + half).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, + hidden_size, + bias=True, + **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding(t, self.frequency_embedding_size, + self.max_period).type( + self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py new file mode 100644 index 0000000000000000000000000000000000000000..759897e2792f73b5c302f56b18526c88c8e7b99a --- /dev/null +++ b/fastvideo/models/hunyuan/modules/models.py @@ -0,0 +1,750 @@ +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin +from einops import rearrange + +from fastvideo.models.hunyuan.modules.posemb_layers import \ + get_nd_rotary_pos_embed +from fastvideo.utils.parallel_states import nccl_info + +from .activation_layers import get_activation_layer +from .attenion import parallel_attention +from .embed_layers import PatchEmbed, TextProjection, TimestepEmbedder +from .mlp_layers import MLP, FinalLayer, MLPEmbedder +from .modulate_layers import ModulateDiT, apply_gate, modulate +from .norm_layers import get_norm_layer +from .posemb_layers import apply_rotary_emb +from .token_refiner import SingleTokenRefiner + + +class MMDoubleStreamBlock(nn.Module): + """ + A multimodal dit block with separate modulation for + text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206 + (Flux.1): https://github.com/black-forest-labs/flux + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float, + mlp_act_type: str = "gelu_tanh", + qk_norm: bool = True, + qk_norm_type: str = "rms", + qkv_bias: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.img_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.img_norm1 = nn.LayerNorm(hidden_size, + elementwise_affine=False, + eps=1e-6, + **factory_kwargs) + + self.img_attn_qkv = nn.Linear(hidden_size, + hidden_size * 3, + bias=qkv_bias, + **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.img_attn_q_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.img_attn_k_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.img_attn_proj = nn.Linear(hidden_size, + hidden_size, + bias=qkv_bias, + **factory_kwargs) + + self.img_norm2 = nn.LayerNorm(hidden_size, + elementwise_affine=False, + eps=1e-6, + **factory_kwargs) + self.img_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + + self.txt_mod = ModulateDiT( + hidden_size, + factor=6, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.txt_norm1 = nn.LayerNorm(hidden_size, + elementwise_affine=False, + eps=1e-6, + **factory_kwargs) + + self.txt_attn_qkv = nn.Linear(hidden_size, + hidden_size * 3, + bias=qkv_bias, + **factory_kwargs) + self.txt_attn_q_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.txt_attn_k_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.txt_attn_proj = nn.Linear(hidden_size, + hidden_size, + bias=qkv_bias, + **factory_kwargs) + + self.txt_norm2 = nn.LayerNorm(hidden_size, + elementwise_affine=False, + eps=1e-6, + **factory_kwargs) + self.txt_mlp = MLP( + hidden_size, + mlp_hidden_dim, + act_layer=get_activation_layer(mlp_act_type), + bias=True, + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple = None, + text_mask: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(vec).chunk(6, dim=-1) + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(vec).chunk(6, dim=-1) + + # Prepare image for attention. + img_modulated = self.img_norm1(img) + img_modulated = modulate(img_modulated, + shift=img_mod1_shift, + scale=img_mod1_scale) + img_qkv = self.img_attn_qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, + "B L (K H D) -> K B L H D", + K=3, + H=self.heads_num) + # Apply QK-Norm if needed + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + + # Apply RoPE if needed. + if freqs_cis is not None: + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow( + dim, nccl_info.rank_within_group * local_heads, + local_heads) + + freqs_cis = ( + shrink_head(freqs_cis[0], dim=0), + shrink_head(freqs_cis[1], dim=0), + ) + + img_qq, img_kk = apply_rotary_emb(img_q, + img_k, + freqs_cis, + head_first=False) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + # Prepare txt for attention. + txt_modulated = self.txt_norm1(txt) + txt_modulated = modulate(txt_modulated, + shift=txt_mod1_shift, + scale=txt_mod1_scale) + txt_qkv = self.txt_attn_qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, + "B L (K H D) -> K B L H D", + K=3, + H=self.heads_num) + # Apply QK-Norm if needed. + txt_q = self.txt_attn_q_norm(txt_q).to(txt_v) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_v) + + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + ) + + # attention computation end + + img_attn, txt_attn = attn[:, :img.shape[1]], attn[:, img.shape[1]:] + + # Calculate the img blocks. + img = img + apply_gate(self.img_attn_proj(img_attn), + gate=img_mod1_gate) + img = img + apply_gate( + self.img_mlp( + modulate(self.img_norm2(img), + shift=img_mod2_shift, + scale=img_mod2_scale)), + gate=img_mod2_gate, + ) + + # Calculate the txt blocks. + txt = txt + apply_gate(self.txt_attn_proj(txt_attn), + gate=txt_mod1_gate) + txt = txt + apply_gate( + self.txt_mlp( + modulate(self.txt_norm2(txt), + shift=txt_mod2_shift, + scale=txt_mod2_scale)), + gate=txt_mod2_gate, + ) + + return img, txt + + +class MMSingleStreamBlock(nn.Module): + """ + A DiT block with parallel linear layers as described in + https://arxiv.org/abs/2302.05442 and adapted modulation interface. + Also refer to (SD3): https://arxiv.org/abs/2403.03206 + (Flux.1): https://github.com/black-forest-labs/flux + """ + + def __init__( + self, + hidden_size: int, + heads_num: int, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + qk_norm: bool = True, + qk_norm_type: str = "rms", + qk_scale: float = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.deterministic = False + self.hidden_size = hidden_size + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + self.mlp_hidden_dim = mlp_hidden_dim + self.scale = qk_scale or head_dim**-0.5 + + # qkv and mlp_in + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim, + **factory_kwargs) + # proj and mlp_out + self.linear2 = nn.Linear(hidden_size + mlp_hidden_dim, hidden_size, + **factory_kwargs) + + qk_norm_layer = get_norm_layer(qk_norm_type) + self.q_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.k_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + + self.pre_norm = nn.LayerNorm(hidden_size, + elementwise_affine=False, + eps=1e-6, + **factory_kwargs) + + self.mlp_act = get_activation_layer(mlp_act_type)() + self.modulation = ModulateDiT( + hidden_size, + factor=3, + act_layer=get_activation_layer("silu"), + **factory_kwargs, + ) + self.hybrid_seq_parallel_attn = None + + def enable_deterministic(self): + self.deterministic = True + + def disable_deterministic(self): + self.deterministic = False + + def forward( + self, + x: torch.Tensor, + vec: torch.Tensor, + txt_len: int, + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + text_mask: torch.Tensor = None, + ) -> torch.Tensor: + mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1) + x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale) + qkv, mlp = torch.split(self.linear1(x_mod), + [3 * self.hidden_size, self.mlp_hidden_dim], + dim=-1) + + q, k, v = rearrange(qkv, + "B L (K H D) -> K B L H D", + K=3, + H=self.heads_num) + + # Apply QK-Norm if needed. + q = self.q_norm(q).to(v) + k = self.k_norm(k).to(v) + + def shrink_head(encoder_state, dim): + local_heads = encoder_state.shape[dim] // nccl_info.sp_size + return encoder_state.narrow( + dim, nccl_info.rank_within_group * local_heads, local_heads) + + freqs_cis = (shrink_head(freqs_cis[0], + dim=0), shrink_head(freqs_cis[1], dim=0)) + + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :] + img_v, txt_v = v[:, :-txt_len, :, :], v[:, -txt_len:, :, :] + img_qq, img_kk = apply_rotary_emb(img_q, + img_k, + freqs_cis, + head_first=False) + assert ( + img_qq.shape == img_q.shape and img_kk.shape == img_k.shape + ), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}" + img_q, img_k = img_qq, img_kk + + attn = parallel_attention( + (img_q, txt_q), + (img_k, txt_k), + (img_v, txt_v), + img_q_len=img_q.shape[1], + img_kv_len=img_k.shape[1], + text_mask=text_mask, + ) + + # attention computation end + + # Compute activation in mlp stream, cat again and run second linear layer. + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + apply_gate(output, gate=mod_gate) + + +class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin): + """ + HunyuanVideo Transformer backbone + + Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline. + + Reference: + [1] Flux.1: https://github.com/black-forest-labs/flux + [2] MMDiT: http://arxiv.org/abs/2403.03206 + + Parameters + ---------- + args: argparse.Namespace + The arguments parsed by argparse. + patch_size: list + The size of the patch. + in_channels: int + The number of input channels. + out_channels: int + The number of output channels. + hidden_size: int + The hidden size of the transformer backbone. + heads_num: int + The number of attention heads. + mlp_width_ratio: float + The ratio of the hidden size of the MLP in the transformer block. + mlp_act_type: str + The activation function of the MLP in the transformer block. + depth_double_blocks: int + The number of transformer blocks in the double blocks. + depth_single_blocks: int + The number of transformer blocks in the single blocks. + rope_dim_list: list + The dimension of the rotary embedding for t, h, w. + qkv_bias: bool + Whether to use bias in the qkv linear layer. + qk_norm: bool + Whether to use qk norm. + qk_norm_type: str + The type of qk norm. + guidance_embed: bool + Whether to use guidance embedding for distillation. + text_projection: str + The type of the text projection, default is single_refiner. + use_attention_mask: bool + Whether to use attention mask for text encoder. + dtype: torch.dtype + The dtype of the model. + device: torch.device + The device of the model. + """ + + @register_to_config + def __init__( + self, + patch_size: list = [1, 2, 2], + in_channels: int = 4, # Should be VAE.config.latent_channels. + out_channels: int = None, + hidden_size: int = 3072, + heads_num: int = 24, + mlp_width_ratio: float = 4.0, + mlp_act_type: str = "gelu_tanh", + mm_double_blocks_depth: int = 20, + mm_single_blocks_depth: int = 40, + rope_dim_list: List[int] = [16, 56, 56], + qkv_bias: bool = True, + qk_norm: bool = True, + qk_norm_type: str = "rms", + guidance_embed: bool = False, # For modulation. + text_projection: str = "single_refiner", + use_attention_mask: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, + rope_theta: int = 256, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.unpatchify_channels = self.out_channels + self.guidance_embed = guidance_embed + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + # Text projection. Default to linear projection. + # Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831 + self.use_attention_mask = use_attention_mask + self.text_projection = text_projection + + if hidden_size % heads_num != 0: + raise ValueError( + f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}" + ) + pe_dim = hidden_size // heads_num + if sum(rope_dim_list) != pe_dim: + raise ValueError( + f"Got {rope_dim_list} but expected positional dim {pe_dim}") + self.hidden_size = hidden_size + self.heads_num = heads_num + + # image projection + self.img_in = PatchEmbed(self.patch_size, self.in_channels, + self.hidden_size, **factory_kwargs) + + # text projection + if self.text_projection == "linear": + self.txt_in = TextProjection( + self.config.text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs, + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + self.config.text_states_dim, + hidden_size, + heads_num, + depth=2, + **factory_kwargs, + ) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}") + + # time modulation + self.time_in = TimestepEmbedder(self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs) + + # text modulation + self.vector_in = MLPEmbedder(self.config.text_states_dim_2, + self.hidden_size, **factory_kwargs) + + # guidance modulation + self.guidance_in = (TimestepEmbedder( + self.hidden_size, get_activation_layer("silu"), **factory_kwargs) + if guidance_embed else None) + + # double blocks + self.double_blocks = nn.ModuleList([ + MMDoubleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(mm_double_blocks_depth) + ]) + + # single blocks + self.single_blocks = nn.ModuleList([ + MMSingleStreamBlock( + self.hidden_size, + self.heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_act_type=mlp_act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + **factory_kwargs, + ) for _ in range(mm_single_blocks_depth) + ]) + + self.final_layer = FinalLayer( + self.hidden_size, + self.patch_size, + self.out_channels, + get_activation_layer("silu"), + **factory_kwargs, + ) + + def enable_deterministic(self): + for block in self.double_blocks: + block.enable_deterministic() + for block in self.single_blocks: + block.enable_deterministic() + + def disable_deterministic(self): + for block in self.double_blocks: + block.disable_deterministic() + for block in self.single_blocks: + block.disable_deterministic() + + def get_rotary_pos_embed(self, rope_sizes): + target_ndim = 3 + + head_dim = self.hidden_size // self.heads_num + rope_dim_list = self.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [ + head_dim // target_ndim for _ in range(target_ndim) + ] + assert ( + sum(rope_dim_list) == head_dim + ), "sum(rope_dim_list) should equal to head_dim of attention layer" + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=self.rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + # x: torch.Tensor, + # t: torch.Tensor, # Should be in range(0, 1000). + # text_states: torch.Tensor = None, + # text_mask: torch.Tensor = None, # Now we don't use it. + # text_states_2: Optional[torch.Tensor] = None, # Text embedding for modulation. + # guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. + # return_dict: bool = True, + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_attention_mask: torch.Tensor, + output_features=False, + output_features_stride=8, + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + guidance=None, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + if guidance is None: + guidance = torch.tensor([6016.0], + device=hidden_states.device, + dtype=torch.bfloat16) + img = x = hidden_states + text_mask = encoder_attention_mask + t = timestep + txt = encoder_hidden_states[:, 1:] + text_states_2 = encoder_hidden_states[:, 0, :self.config. + text_states_dim_2] + _, _, ot, oh, ow = x.shape # codespell:ignore + tt, th, tw = ( + ot // self.patch_size[0], # codespell:ignore + oh // self.patch_size[1], # codespell:ignore + ow // self.patch_size[2], # codespell:ignore + ) + original_tt = nccl_info.sp_size * tt + freqs_cos, freqs_sin = self.get_rotary_pos_embed((original_tt, th, tw)) + # Prepare modulation vectors. + vec = self.time_in(t) + + # text modulation + vec = vec + self.vector_in(text_states_2) + + # guidance modulation + if self.guidance_embed: + if guidance is None: + raise ValueError( + "Didn't get guidance strength for guidance distilled model." + ) + + # our timestep_embedding is merged into guidance_in(TimestepEmbedder) + vec = vec + self.guidance_in(guidance) + + # Embed image and text. + img = self.img_in(img) + if self.text_projection == "linear": + txt = self.txt_in(txt) + elif self.text_projection == "single_refiner": + txt = self.txt_in(txt, t, + text_mask if self.use_attention_mask else None) + else: + raise NotImplementedError( + f"Unsupported text_projection: {self.text_projection}") + + txt_seq_len = txt.shape[1] + img_seq_len = img.shape[1] + + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + # --------------------- Pass through DiT blocks ------------------------ + for _, block in enumerate(self.double_blocks): + double_block_args = [img, txt, vec, freqs_cis, text_mask] + + img, txt = block(*double_block_args) + + # Merge txt and img to pass through single stream blocks. + x = torch.cat((img, txt), 1) + if output_features: + features_list = [] + if len(self.single_blocks) > 0: + for _, block in enumerate(self.single_blocks): + single_block_args = [ + x, + vec, + txt_seq_len, + (freqs_cos, freqs_sin), + text_mask, + ] + + x = block(*single_block_args) + if output_features and _ % output_features_stride == 0: + features_list.append(x[:, :img_seq_len, ...]) + + img = x[:, :img_seq_len, ...] + + # ---------------------------- Final layer ------------------------------ + img = self.final_layer(img, + vec) # (N, T, patch_size ** 2 * out_channels) + + img = self.unpatchify(img, tt, th, tw) + assert not return_dict, "return_dict is not supported." + if output_features: + features_list = torch.stack(features_list, dim=0) + else: + features_list = None + return (img, features_list) + + def unpatchify(self, x, t, h, w): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.unpatchify_channels + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs + + def params_count(self): + counts = { + "double": + sum([ + sum(p.numel() for p in block.img_attn_qkv.parameters()) + + sum(p.numel() for p in block.img_attn_proj.parameters()) + + sum(p.numel() for p in block.img_mlp.parameters()) + + sum(p.numel() for p in block.txt_attn_qkv.parameters()) + + sum(p.numel() for p in block.txt_attn_proj.parameters()) + + sum(p.numel() for p in block.txt_mlp.parameters()) + for block in self.double_blocks + ]), + "single": + sum([ + sum(p.numel() for p in block.linear1.parameters()) + + sum(p.numel() for p in block.linear2.parameters()) + for block in self.single_blocks + ]), + "total": + sum(p.numel() for p in self.parameters()), + } + counts["attn+mlp"] = counts["double"] + counts["single"] + return counts + + +################################################################################# +# HunyuanVideo Configs # +################################################################################# + +HUNYUAN_VIDEO_CONFIG = { + "HYVideo-T/2": { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + }, + "HYVideo-T/2-cfgdistill": { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, + }, +} diff --git a/fastvideo/models/hunyuan/modules/modulate_layers.py b/fastvideo/models/hunyuan/modules/modulate_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..c7a36369296ff50cb2b99ed0a0d5932bb64f781d --- /dev/null +++ b/fastvideo/models/hunyuan/modules/modulate_layers.py @@ -0,0 +1,156 @@ +from typing import Callable + +import torch +import torch.nn as nn + + +class ModulateDiT(nn.Module): + """Modulation layer for DiT.""" + + def __init__( + self, + hidden_size: int, + factor: int, + act_layer: Callable, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.act = act_layer() + self.linear = nn.Linear(hidden_size, + factor * hidden_size, + bias=True, + **factory_kwargs) + # Zero-initialize the modulation + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.act(x)) + + +def modulate(x, shift=None, scale=None): + """modulate by shift and scale + + Args: + x (torch.Tensor): input tensor. + shift (torch.Tensor, optional): shift tensor. Defaults to None. + scale (torch.Tensor, optional): scale tensor. Defaults to None. + + Returns: + torch.Tensor: the output tensor after modulate. + """ + if scale is None and shift is None: + return x + elif shift is None: + return x * (1 + scale.unsqueeze(1)) + elif scale is None: + return x + shift.unsqueeze(1) + else: + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def apply_gate(x, gate=None, tanh=False): + """AI is creating summary for apply_gate + + Args: + x (torch.Tensor): input tensor. + gate (torch.Tensor, optional): gate tensor. Defaults to None. + tanh (bool, optional): whether to use tanh function. Defaults to False. + + Returns: + torch.Tensor: the output tensor after apply gate. + """ + if gate is None: + return x + if tanh: + return x * gate.unsqueeze(1).tanh() + else: + return x * gate.unsqueeze(1) + + +def ckpt_wrapper(module): + + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + + return ckpt_forward + + +class RMSNorm(nn.Module): + + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError( + f"Norm layer {norm_layer} is not implemented") diff --git a/fastvideo/models/hunyuan/modules/norm_layers.py b/fastvideo/models/hunyuan/modules/norm_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..0e1e590e7f17de1947bda3910e35a26a48a44d56 --- /dev/null +++ b/fastvideo/models/hunyuan/modules/norm_layers.py @@ -0,0 +1,79 @@ +import torch +import torch.nn as nn + + +class RMSNorm(nn.Module): + + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +def get_norm_layer(norm_layer): + """ + Get the normalization layer. + + Args: + norm_layer (str): The type of normalization layer. + + Returns: + norm_layer (nn.Module): The normalization layer. + """ + if norm_layer == "layer": + return nn.LayerNorm + elif norm_layer == "rms": + return RMSNorm + else: + raise NotImplementedError( + f"Norm layer {norm_layer} is not implemented") diff --git a/fastvideo/models/hunyuan/modules/posemb_layers.py b/fastvideo/models/hunyuan/modules/posemb_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..2c92471e450319c9916d9ab364d07eccc83ab237 --- /dev/null +++ b/fastvideo/models/hunyuan/modules/posemb_layers.py @@ -0,0 +1,314 @@ +from typing import List, Tuple, Union + +import torch + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x, ) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): + """ + Get n-D meshgrid with start, stop and num. + + Args: + start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop, + step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num + should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in + n-tuples. + *args: See above. + dim (int): Dimension of the meshgrid. Defaults to 2. + + Returns: + grid (np.ndarray): [dim, ...] + """ + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0, ) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D] + grid = torch.stack(grid, dim=0) # [dim, W, H, D] + + return grid + + +################################################################################# +# Rotary Positional Embedding Functions # +################################################################################# +# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 + + +def reshape_for_broadcast( + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], + x: torch.Tensor, + head_first=False, +): + """ + Reshape frequency tensor for broadcasting it with another tensor. + + This function reshapes the frequency tensor to have the same shape as the target tensor 'x' + for the purpose of broadcasting the frequency tensor during element-wise operations. + + Notes: + When using FlashMHAModified, head_first should be False. + When using Attention, head_first should be True. + + Args: + freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. + x (torch.Tensor): Target tensor for broadcasting compatibility. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + torch.Tensor: Reshaped frequency tensor. + + Raises: + AssertionError: If the frequency tensor doesn't match the expected shape. + AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. + """ + ndim = x.ndim + assert 0 <= 1 < ndim + + if isinstance(freqs_cis, tuple): + # freqs_cis: (cos, sin) in real space + if head_first: + assert freqs_cis[0].shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis[0].shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}" + shape = [ + d if i == 1 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) + else: + # freqs_cis: values in complex space + if head_first: + assert freqs_cis.shape == ( + x.shape[-2], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [ + d if i == ndim - 2 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + else: + assert freqs_cis.shape == ( + x.shape[1], + x.shape[-1], + ), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}" + shape = [ + d if i == 1 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + return freqs_cis.view(*shape) + + +def rotate_half(x): + x_real, x_imag = (x.float().reshape(*x.shape[:-1], -1, + 2).unbind(-1)) # [B, S, H, D//2] + return torch.stack([-x_imag, x_real], dim=-1).flatten(3) + + +def apply_rotary_emb( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + head_first: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Apply rotary embeddings to input tensors using the given frequency tensor. + + This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided + frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor + is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are + returned as real tensors. + + Args: + xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] + xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] + freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. + head_first (bool): head dimension first (except batch dim) or not. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. + + """ + xk_out = None + if isinstance(freqs_cis, tuple): + cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] + cos, sin = cos.to(xq.device), sin.to(xq.device) + # real * cos - imag * sin + # imag * cos + real * sin + xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) + xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) + else: + # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) + xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, + 2)) # [B, S, H, D//2] + freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to( + xq.device) # [S, D//2] --> [1, S, 1, D//2] + # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) + # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) + xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, + 2)) # [B, S, H, D//2] + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) + + return xq_out, xk_out + + +def get_nd_rotary_pos_embed( + rope_dim_list, + start, + *args, + theta=10000.0, + use_real=False, + theta_rescale_factor: Union[float, List[float]] = 1.0, + interpolation_factor: Union[float, List[float]] = 1.0, +): + """ + This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure. + + Args: + rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n. + sum(rope_dim_list) should equal to head_dim of attention layer. + start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start, + args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. + *args: See above. + theta (float): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers. + Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real + part and an imaginary part separately. + theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0. + + Returns: + pos_embed (torch.Tensor): [HW, D/2] + """ + + grid = get_meshgrid_nd(start, *args, + dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H] + + if isinstance(theta_rescale_factor, int) or isinstance( + theta_rescale_factor, float): + theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list) + elif isinstance(theta_rescale_factor, + list) and len(theta_rescale_factor) == 1: + theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list) + assert len(theta_rescale_factor) == len( + rope_dim_list + ), "len(theta_rescale_factor) should equal to len(rope_dim_list)" + + if isinstance(interpolation_factor, int) or isinstance( + interpolation_factor, float): + interpolation_factor = [interpolation_factor] * len(rope_dim_list) + elif isinstance(interpolation_factor, + list) and len(interpolation_factor) == 1: + interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list) + assert len(interpolation_factor) == len( + rope_dim_list + ), "len(interpolation_factor) should equal to len(rope_dim_list)" + + # use 1/ndim of dimensions to encode grid_axis + embs = [] + for i in range(len(rope_dim_list)): + emb = get_1d_rotary_pos_embed( + rope_dim_list[i], + grid[i].reshape(-1), + theta, + use_real=use_real, + theta_rescale_factor=theta_rescale_factor[i], + interpolation_factor=interpolation_factor[i], + ) # 2 x [WHD, rope_dim_list[i]] + embs.append(emb) + + if use_real: + cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2) + sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2) + return cos, sin + else: + emb = torch.cat(embs, dim=1) # (WHD, D/2) + return emb + + +def get_1d_rotary_pos_embed( + dim: int, + pos: Union[torch.FloatTensor, int], + theta: float = 10000.0, + use_real: bool = False, + theta_rescale_factor: float = 1.0, + interpolation_factor: float = 1.0, +) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """ + Precompute the frequency tensor for complex exponential (cis) with given dimensions. + (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.) + + This function calculates a frequency tensor with complex exponential using the given dimension 'dim' + and the end index 'end'. The 'theta' parameter scales the frequencies. + The returned tensor contains complex values in complex64 data type. + + Args: + dim (int): Dimension of the frequency tensor. + pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar + theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0. + use_real (bool, optional): If True, return real part and imaginary part separately. + Otherwise, return complex numbers. + theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0. + + Returns: + freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2] + freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D] + """ + if isinstance(pos, int): + pos = torch.arange(pos).float() + + # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning + # has some connection to NTK literature + if theta_rescale_factor != 1.0: + theta *= theta_rescale_factor**(dim / (dim - 2)) + + freqs = 1.0 / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim) + ) # [D/2] + # assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}" + freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2] + if use_real: + freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D] + freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D] + return freqs_cos, freqs_sin + else: + freqs_cis = torch.polar(torch.ones_like(freqs), + freqs) # complex64 # [S, D/2] + return freqs_cis diff --git a/fastvideo/models/hunyuan/modules/token_refiner.py b/fastvideo/models/hunyuan/modules/token_refiner.py new file mode 100644 index 0000000000000000000000000000000000000000..ea6cb9c4d2cea6d41a786b1ae50146c44fecfd0b --- /dev/null +++ b/fastvideo/models/hunyuan/modules/token_refiner.py @@ -0,0 +1,230 @@ +from typing import Optional + +import torch +import torch.nn as nn +from einops import rearrange + +from .activation_layers import get_activation_layer +from .attenion import attention +from .embed_layers import TextProjection, TimestepEmbedder +from .mlp_layers import MLP +from .modulate_layers import apply_gate +from .norm_layers import get_norm_layer + + +class IndividualTokenRefinerBlock(nn.Module): + + def __init__( + self, + hidden_size, + heads_num, + mlp_width_ratio: str = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.heads_num = heads_num + head_dim = hidden_size // heads_num + mlp_hidden_dim = int(hidden_size * mlp_width_ratio) + + self.norm1 = nn.LayerNorm(hidden_size, + elementwise_affine=True, + eps=1e-6, + **factory_kwargs) + self.self_attn_qkv = nn.Linear(hidden_size, + hidden_size * 3, + bias=qkv_bias, + **factory_kwargs) + qk_norm_layer = get_norm_layer(qk_norm_type) + self.self_attn_q_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.self_attn_k_norm = (qk_norm_layer( + head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) + if qk_norm else nn.Identity()) + self.self_attn_proj = nn.Linear(hidden_size, + hidden_size, + bias=qkv_bias, + **factory_kwargs) + + self.norm2 = nn.LayerNorm(hidden_size, + elementwise_affine=True, + eps=1e-6, + **factory_kwargs) + act_layer = get_activation_layer(act_type) + self.mlp = MLP( + in_channels=hidden_size, + hidden_channels=mlp_hidden_dim, + act_layer=act_layer, + drop=mlp_drop_rate, + **factory_kwargs, + ) + + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, + 2 * hidden_size, + bias=True, + **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward( + self, + x: torch.Tensor, + c: torch. + Tensor, # timestep_aware_representations + context_aware_representations + attn_mask: torch.Tensor = None, + ): + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1) + + norm_x = self.norm1(x) + qkv = self.self_attn_qkv(norm_x) + q, k, v = rearrange(qkv, + "B L (K H D) -> K B L H D", + K=3, + H=self.heads_num) + # Apply QK-Norm if needed + q = self.self_attn_q_norm(q).to(v) + k = self.self_attn_k_norm(k).to(v) + + # Self-Attention + attn = attention(q, k, v, attn_mask=attn_mask) + + x = x + apply_gate(self.self_attn_proj(attn), gate_msa) + + # FFN Layer + x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp) + + return x + + +class IndividualTokenRefiner(nn.Module): + + def __init__( + self, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.blocks = nn.ModuleList([ + IndividualTokenRefinerBlock( + hidden_size=hidden_size, + heads_num=heads_num, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) for _ in range(depth) + ]) + + def forward( + self, + x: torch.Tensor, + c: torch.LongTensor, + mask: Optional[torch.Tensor] = None, + ): + mask = mask.clone().bool() + # avoid attention weight become NaN + mask[:, 0] = True + for block in self.blocks: + x = block(x, c, mask) + return x + + +class SingleTokenRefiner(nn.Module): + """ + A single token refiner block for llm text embedding refine. + """ + + def __init__( + self, + in_channels, + hidden_size, + heads_num, + depth, + mlp_width_ratio: float = 4.0, + mlp_drop_rate: float = 0.0, + act_type: str = "silu", + qk_norm: bool = False, + qk_norm_type: str = "layer", + qkv_bias: bool = True, + attn_mode: str = "torch", + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.attn_mode = attn_mode + assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner." + + self.input_embedder = nn.Linear(in_channels, + hidden_size, + bias=True, + **factory_kwargs) + + act_layer = get_activation_layer(act_type) + # Build timestep embedding layer + self.t_embedder = TimestepEmbedder(hidden_size, act_layer, + **factory_kwargs) + # Build context embedding layer + self.c_embedder = TextProjection(in_channels, hidden_size, act_layer, + **factory_kwargs) + + self.individual_token_refiner = IndividualTokenRefiner( + hidden_size=hidden_size, + heads_num=heads_num, + depth=depth, + mlp_width_ratio=mlp_width_ratio, + mlp_drop_rate=mlp_drop_rate, + act_type=act_type, + qk_norm=qk_norm, + qk_norm_type=qk_norm_type, + qkv_bias=qkv_bias, + **factory_kwargs, + ) + + def forward( + self, + x: torch.Tensor, + t: torch.LongTensor, + mask: Optional[torch.LongTensor] = None, + ): + timestep_aware_representations = self.t_embedder(t) + + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [b, s1, 1] + context_aware_representations = (x * mask_float).sum( + dim=1) / mask_float.sum(dim=1) + context_aware_representations = self.c_embedder( + context_aware_representations) + c = timestep_aware_representations + context_aware_representations + + x = self.input_embedder(x) + + x = self.individual_token_refiner(x, c, mask) + + return x diff --git a/fastvideo/models/hunyuan/text_encoder/__init__.py b/fastvideo/models/hunyuan/text_encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..18cd88e183ffd0fdce11d9876711c0c86dd3d9e7 --- /dev/null +++ b/fastvideo/models/hunyuan/text_encoder/__init__.py @@ -0,0 +1,353 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn +from transformers import AutoModel, AutoTokenizer, CLIPTextModel, CLIPTokenizer +from transformers.utils import ModelOutput + +from ..constants import PRECISION_TO_TYPE, TEXT_ENCODER_PATH, TOKENIZER_PATH + + +def use_default(value, default): + return value if value is not None else default + + +def load_text_encoder( + text_encoder_type, + text_encoder_precision=None, + text_encoder_path=None, + logger=None, + device=None, +): + if text_encoder_path is None: + text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] + if logger is not None: + logger.info( + f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}" + ) + + if text_encoder_type == "clipL": + text_encoder = CLIPTextModel.from_pretrained(text_encoder_path) + text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm + elif text_encoder_type == "llm": + text_encoder = AutoModel.from_pretrained(text_encoder_path, + low_cpu_mem_usage=True) + text_encoder.final_layer_norm = text_encoder.norm + else: + raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") + # from_pretrained will ensure that the model is in eval mode. + + if text_encoder_precision is not None: + text_encoder = text_encoder.to( + dtype=PRECISION_TO_TYPE[text_encoder_precision]) + + text_encoder.requires_grad_(False) + + if logger is not None: + logger.info(f"Text encoder to dtype: {text_encoder.dtype}") + + if device is not None: + text_encoder = text_encoder.to(device) + + return text_encoder, text_encoder_path + + +def load_tokenizer(tokenizer_type, + tokenizer_path=None, + padding_side="right", + logger=None): + if tokenizer_path is None: + tokenizer_path = TOKENIZER_PATH[tokenizer_type] + if logger is not None: + logger.info( + f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}") + + if tokenizer_type == "clipL": + tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, + max_length=77) + elif tokenizer_type == "llm": + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, + padding_side=padding_side) + else: + raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}") + + return tokenizer, tokenizer_path + + +@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_type: str, + max_length: int, + text_encoder_precision: Optional[str] = None, + text_encoder_path: Optional[str] = None, + tokenizer_type: Optional[str] = None, + tokenizer_path: Optional[str] = None, + output_key: Optional[str] = None, + use_attention_mask: bool = True, + input_max_length: Optional[int] = None, + 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, + logger=None, + device=None, + ): + super().__init__() + self.text_encoder_type = text_encoder_type + self.max_length = max_length + self.precision = text_encoder_precision + self.model_path = text_encoder_path + self.tokenizer_type = (tokenizer_type if tokenizer_type is not None + else text_encoder_type) + self.tokenizer_path = (tokenizer_path if tokenizer_path is not None + else 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.input_max_length = (input_max_length if input_max_length + is not None else max_length) + 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.logger = logger + + 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 text_encoder_type: + self.output_key = output_key or "last_hidden_state" + elif "clip" in text_encoder_type: + self.output_key = output_key or "pooler_output" + elif "llm" in text_encoder_type or "glm" in text_encoder_type: + self.output_key = output_key or "last_hidden_state" + else: + raise ValueError( + f"Unsupported text encoder type: {text_encoder_type}") + + self.model, self.model_path = load_text_encoder( + text_encoder_type=self.text_encoder_type, + text_encoder_precision=self.precision, + text_encoder_path=self.model_path, + logger=self.logger, + device=device, + ) + self.dtype = self.model.dtype + self.device = self.model.device + + self.tokenizer, self.tokenizer_path = load_tokenizer( + tokenizer_type=self.tokenizer_type, + tokenizer_path=self.tokenizer_path, + padding_side="right", + logger=self.logger, + ) + + 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) + outputs = self.model( + input_ids=batch_encoding["input_ids"].to(device), + attention_mask=attention_mask, + output_hidden_states=output_hidden_states + or hidden_state_skip_layer is not None, + ) + if hidden_state_skip_layer is not None: + last_hidden_state = outputs.hidden_states[-( + hidden_state_skip_layer + 1)] + # Real last hidden state already has layer norm applied. So here we only apply it + # for intermediate layers. + if hidden_state_skip_layer > 0 and self.apply_final_norm: + last_hidden_state = self.model.final_layer_norm( + last_hidden_state) + else: + last_hidden_state = outputs[self.output_key] + + # Remove hidden states of instruction tokens, only keep prompt tokens. + if self.use_template: + if data_type == "image": + crop_start = self.prompt_template.get("crop_start", -1) + elif data_type == "video": + crop_start = self.prompt_template_video.get("crop_start", -1) + else: + raise ValueError(f"Unsupported data type: {data_type}") + if crop_start > 0: + last_hidden_state = last_hidden_state[:, crop_start:] + attention_mask = (attention_mask[:, crop_start:] + if use_attention_mask else None) + + if output_hidden_states: + return TextEncoderModelOutput(last_hidden_state, attention_mask, + outputs.hidden_states) + return TextEncoderModelOutput(last_hidden_state, attention_mask) + + def forward( + self, + text, + use_attention_mask=None, + output_hidden_states=False, + do_sample=False, + hidden_state_skip_layer=None, + return_texts=False, + ): + batch_encoding = self.text2tokens(text) + return self.encode( + batch_encoding, + use_attention_mask=use_attention_mask, + output_hidden_states=output_hidden_states, + do_sample=do_sample, + hidden_state_skip_layer=hidden_state_skip_layer, + return_texts=return_texts, + ) diff --git a/fastvideo/models/hunyuan/text_encoder/__pycache__/__init__.cpython-310.pyc b/fastvideo/models/hunyuan/text_encoder/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5b0370b6970697cc5ac3b239306cf46f49ebb74 Binary files /dev/null and b/fastvideo/models/hunyuan/text_encoder/__pycache__/__init__.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/text_encoder/__pycache__/__init__.cpython-312.pyc b/fastvideo/models/hunyuan/text_encoder/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07843f1fc9a830136e40223ff81fd96c3f85c9f1 Binary files /dev/null and b/fastvideo/models/hunyuan/text_encoder/__pycache__/__init__.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/utils/__init__.py b/fastvideo/models/hunyuan/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/fastvideo/models/hunyuan/utils/__pycache__/__init__.cpython-310.pyc b/fastvideo/models/hunyuan/utils/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dfa52b2b500f5c9d17adf03861000fcf210dcd35 Binary files /dev/null and b/fastvideo/models/hunyuan/utils/__pycache__/__init__.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/utils/__pycache__/__init__.cpython-312.pyc b/fastvideo/models/hunyuan/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16826a8a298869d2e375f9b5d1ba1881f92b725d Binary files /dev/null and b/fastvideo/models/hunyuan/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/utils/__pycache__/helpers.cpython-310.pyc b/fastvideo/models/hunyuan/utils/__pycache__/helpers.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca631b163511c8a7848f6028d1c2b451570f6d17 Binary files /dev/null and b/fastvideo/models/hunyuan/utils/__pycache__/helpers.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/utils/__pycache__/helpers.cpython-312.pyc b/fastvideo/models/hunyuan/utils/__pycache__/helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd27bb1bba0799be8dadcd6033159ca57b627ae2 Binary files /dev/null and b/fastvideo/models/hunyuan/utils/__pycache__/helpers.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/utils/data_utils.py b/fastvideo/models/hunyuan/utils/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..524118179cc8fae51ddee464c901b685ce56caac --- /dev/null +++ b/fastvideo/models/hunyuan/utils/data_utils.py @@ -0,0 +1,14 @@ +import math + + +def align_to(value, alignment): + """align height, width according to alignment + + Args: + value (int): height or width + alignment (int): target alignment factor + + Returns: + int: the aligned value + """ + return int(math.ceil(value / alignment) * alignment) diff --git a/fastvideo/models/hunyuan/utils/file_utils.py b/fastvideo/models/hunyuan/utils/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c87a95eff2b27f36f91a853236def21a32e1eae5 --- /dev/null +++ b/fastvideo/models/hunyuan/utils/file_utils.py @@ -0,0 +1,75 @@ +import os +from pathlib import Path + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +CODE_SUFFIXES = { + ".py", # Python codes + ".sh", # Shell scripts + ".yaml", + ".yml", # Configuration files +} + + +def safe_dir(path): + """ + Create a directory (or the parent directory of a file) if it does not exist. + + Args: + path (str or Path): Path to the directory. + + Returns: + path (Path): Path object of the directory. + """ + path = Path(path) + path.mkdir(exist_ok=True, parents=True) + return path + + +def safe_file(path): + """ + Create the parent directory of a file if it does not exist. + + Args: + path (str or Path): Path to the file. + + Returns: + path (Path): Path object of the file. + """ + path = Path(path) + path.parent.mkdir(exist_ok=True, parents=True) + return path + + +def save_videos_grid(videos: torch.Tensor, + path: str, + rescale=False, + n_rows=1, + fps=24): + """save videos by video tensor + copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61 + + Args: + videos (torch.Tensor): video tensor predicted by the model + path (str): path to save video + rescale (bool, optional): rescale the video tensor from [-1, 1] to . Defaults to False. + n_rows (int, optional): Defaults to 1. + fps (int, optional): video save fps. Defaults to 8. + """ + videos = rearrange(videos, "b c t h w -> t b c h w") + outputs = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = torch.clamp(x, 0, 1) + x = (x * 255).numpy().astype(np.uint8) + outputs.append(x) + + os.makedirs(os.path.dirname(path), exist_ok=True) + imageio.mimsave(path, outputs, fps=fps) diff --git a/fastvideo/models/hunyuan/utils/helpers.py b/fastvideo/models/hunyuan/utils/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..7b5b8731f5d5ece663868b51a9a55ac6a3c0ecf1 --- /dev/null +++ b/fastvideo/models/hunyuan/utils/helpers.py @@ -0,0 +1,41 @@ +import collections.abc +from itertools import repeat + + +def _ntuple(n): + + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + x = tuple(x) + if len(x) == 1: + x = tuple(repeat(x[0], n)) + return x + return tuple(repeat(x, n)) + + return parse + + +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) + + +def as_tuple(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + if x is None or isinstance(x, (int, float, str)): + return (x, ) + else: + raise ValueError(f"Unknown type {type(x)}") + + +def as_list_of_2tuple(x): + x = as_tuple(x) + if len(x) == 1: + x = (x[0], x[0]) + assert len(x) % 2 == 0, f"Expect even length, got {len(x)}." + lst = [] + for i in range(0, len(x), 2): + lst.append((x[i], x[i + 1])) + return lst diff --git a/fastvideo/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py b/fastvideo/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1a6f46c15a98065a987a42563bf19994bd9673fb --- /dev/null +++ b/fastvideo/models/hunyuan/utils/preprocess_text_encoder_tokenizer_utils.py @@ -0,0 +1,41 @@ +import argparse + +import torch +from transformers import AutoProcessor, LlavaForConditionalGeneration + + +def preprocess_text_encoder_tokenizer(args): + + processor = AutoProcessor.from_pretrained(args.input_dir) + model = LlavaForConditionalGeneration.from_pretrained( + args.input_dir, + torch_dtype=torch.float16, + low_cpu_mem_usage=True, + ).to(0) + + model.language_model.save_pretrained(f"{args.output_dir}") + processor.tokenizer.save_pretrained(f"{args.output_dir}") + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + type=str, + required=True, + help="The path to the llava-llama-3-8b-v1_1-transformers.", + ) + parser.add_argument( + "--output_dir", + type=str, + default="", + help="The output path of the llava-llama-3-8b-text-encoder-tokenizer." + "if '', the parent dir of output will be the same as input dir.", + ) + args = parser.parse_args() + + if len(args.output_dir) == 0: + args.output_dir = "/".join(args.input_dir.split("/")[:-1]) + + preprocess_text_encoder_tokenizer(args) diff --git a/fastvideo/models/hunyuan/vae/__init__.py b/fastvideo/models/hunyuan/vae/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f570423e62924ab6f4d691d7648b7e63f95aceae --- /dev/null +++ b/fastvideo/models/hunyuan/vae/__init__.py @@ -0,0 +1,68 @@ +from pathlib import Path + +import torch + +from ..constants import PRECISION_TO_TYPE, VAE_PATH +from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D + + +def load_vae( + vae_type: str = "884-16c-hy", + vae_precision: str = None, + sample_size: tuple = None, + vae_path: str = None, + logger=None, + device=None, +): + """the function to load the 3D VAE model + + Args: + vae_type (str): the type of the 3D VAE model. Defaults to "884-16c-hy". + vae_precision (str, optional): the precision to load vae. Defaults to None. + sample_size (tuple, optional): the tiling size. Defaults to None. + vae_path (str, optional): the path to vae. Defaults to None. + logger (_type_, optional): logger. Defaults to None. + device (_type_, optional): device to load vae. Defaults to None. + """ + if vae_path is None: + vae_path = VAE_PATH[vae_type] + + if logger is not None: + logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}") + config = AutoencoderKLCausal3D.load_config(vae_path) + if sample_size: + vae = AutoencoderKLCausal3D.from_config(config, + sample_size=sample_size) + else: + vae = AutoencoderKLCausal3D.from_config(config) + + vae_ckpt = Path(vae_path) / "pytorch_model.pt" + assert vae_ckpt.exists(), f"VAE checkpoint not found: {vae_ckpt}" + + ckpt = torch.load(vae_ckpt, map_location=vae.device) + if "state_dict" in ckpt: + ckpt = ckpt["state_dict"] + if any(k.startswith("vae.") for k in ckpt.keys()): + ckpt = { + k.replace("vae.", ""): v + for k, v in ckpt.items() if k.startswith("vae.") + } + vae.load_state_dict(ckpt) + + spatial_compression_ratio = vae.config.spatial_compression_ratio + time_compression_ratio = vae.config.time_compression_ratio + + if vae_precision is not None: + vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision]) + + vae.requires_grad_(False) + + if logger is not None: + logger.info(f"VAE to dtype: {vae.dtype}") + + if device is not None: + vae = vae.to(device) + + vae.eval() + + return vae, vae_path, spatial_compression_ratio, time_compression_ratio diff --git a/fastvideo/models/hunyuan/vae/__pycache__/__init__.cpython-310.pyc b/fastvideo/models/hunyuan/vae/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd60a9bbbf7b2347cd884ec9288c50c4caa536c1 Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/__init__.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/__init__.cpython-312.pyc b/fastvideo/models/hunyuan/vae/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ce47e8ce1802b762eb647907b48fe6ccc550fb9 Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/__init__.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/autoencoder_kl_causal_3d.cpython-310.pyc b/fastvideo/models/hunyuan/vae/__pycache__/autoencoder_kl_causal_3d.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..96b14227317caf67581a502966a9e6a8a67db898 Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/autoencoder_kl_causal_3d.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/autoencoder_kl_causal_3d.cpython-312.pyc b/fastvideo/models/hunyuan/vae/__pycache__/autoencoder_kl_causal_3d.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15643dd93e19aacab536efe950a0847d1d55e74a Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/autoencoder_kl_causal_3d.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/unet_causal_3d_blocks.cpython-310.pyc b/fastvideo/models/hunyuan/vae/__pycache__/unet_causal_3d_blocks.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa1e523f161c744a2dbb545117e6afc97697f386 Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/unet_causal_3d_blocks.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/unet_causal_3d_blocks.cpython-312.pyc b/fastvideo/models/hunyuan/vae/__pycache__/unet_causal_3d_blocks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc5e4c8dfd2c145f1019876d409e451095a80433 Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/unet_causal_3d_blocks.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/vae.cpython-310.pyc b/fastvideo/models/hunyuan/vae/__pycache__/vae.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db8f38f6cae9a82f2826a84f1ca42af6af300bdc Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/vae.cpython-310.pyc differ diff --git a/fastvideo/models/hunyuan/vae/__pycache__/vae.cpython-312.pyc b/fastvideo/models/hunyuan/vae/__pycache__/vae.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..459ebb770e72615cb0c45b0b62112892a5952192 Binary files /dev/null and b/fastvideo/models/hunyuan/vae/__pycache__/vae.cpython-312.pyc differ diff --git a/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py b/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py new file mode 100644 index 0000000000000000000000000000000000000000..1461fb1645634fb81c36d122890f4c260c14342f --- /dev/null +++ b/fastvideo/models/hunyuan/vae/autoencoder_kl_causal_3d.py @@ -0,0 +1,831 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== +from dataclasses import dataclass +from math import prod +from typing import Dict, Optional, Tuple, Union + +import torch +import torch.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 + +try: + # This diffusers is modified and packed in the mirror. + from diffusers.loaders import FromOriginalVAEMixin +except ImportError: + # Use this to be compatible with the original diffusers. + from diffusers.loaders.single_file_model import ( + FromOriginalModelMixin as FromOriginalVAEMixin, ) + +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention, + AttentionProcessor, AttnAddedKVProcessor, AttnProcessor) +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils.accelerate_utils import apply_forward_hook + +from .vae import (BaseOutput, DecoderCausal3D, DecoderOutput, + DiagonalGaussianDistribution, EncoderCausal3D) + + +@dataclass +class DecoderOutput2(BaseOutput): + sample: torch.FloatTensor + posterior: Optional[DiagonalGaussianDistribution] = None + + +class AutoencoderKLCausal3D(ModelMixin, ConfigMixin, FromOriginalVAEMixin): + r""" + A VAE model with KL loss for encoding images/videos into latents and decoding latent representations into images/videos. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str] = ("DownEncoderBlockCausal3D", ), + up_block_types: Tuple[str] = ("UpDecoderBlockCausal3D", ), + block_out_channels: Tuple[int] = (64, ), + layers_per_block: int = 1, + act_fn: str = "silu", + latent_channels: int = 4, + norm_num_groups: int = 32, + sample_size: int = 32, + sample_tsize: int = 64, + scaling_factor: float = 0.18215, + force_upcast: float = True, + spatial_compression_ratio: int = 8, + time_compression_ratio: int = 4, + mid_block_add_attention: bool = True, + ): + super().__init__() + + self.time_compression_ratio = time_compression_ratio + + self.encoder = EncoderCausal3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + double_z=True, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.decoder = DecoderCausal3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + time_compression_ratio=time_compression_ratio, + spatial_compression_ratio=spatial_compression_ratio, + mid_block_add_attention=mid_block_add_attention, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, + 2 * latent_channels, + kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, + latent_channels, + kernel_size=1) + + self.use_slicing = False + self.use_spatial_tiling = False + self.use_temporal_tiling = False + self.use_parallel = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // time_compression_ratio + + self.tile_sample_min_size = self.config.sample_size + sample_size = (self.config.sample_size[0] if isinstance( + self.config.sample_size, + (list, tuple)) else self.config.sample_size) + self.tile_latent_min_size = int( + sample_size / (2**(len(self.config.block_out_channels) - 1))) + self.tile_overlap_factor = 0.25 + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (EncoderCausal3D, DecoderCausal3D)): + module.gradient_checkpointing = value + + def enable_temporal_tiling(self, use_tiling: bool = True): + self.use_temporal_tiling = use_tiling + + def disable_temporal_tiling(self): + self.enable_temporal_tiling(False) + + def enable_spatial_tiling(self, use_tiling: bool = True): + self.use_spatial_tiling = use_tiling + + def disable_spatial_tiling(self): + self.enable_spatial_tiling(False) + + def enable_tiling(self, use_tiling: bool = True): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger videos. + """ + self.enable_spatial_tiling(use_tiling) + self.enable_temporal_tiling(use_tiling) + + def disable_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.disable_spatial_tiling() + self.disable_temporal_tiling() + + def enable_parallel(self): + r""" + Enable sequence parallelism for the model. This will allow the vae to decode (with tiling) in parallel. + """ + self.use_parallel = True + + def enable_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.use_slicing = True + + def disable_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing + decoding in one step. + """ + self.use_slicing = False + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor( + return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, + processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor( + self, + processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], + _remove_lora=False, + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, + processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor(processor.pop(f"{name}.processor"), + _remove_lora=_remove_lora) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, + processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS + for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS + for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + @apply_forward_hook + def encode( + self, + x: torch.FloatTensor, + return_dict: bool = True + ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + """ + Encode a batch of images/videos into latents. + + Args: + x (`torch.FloatTensor`): Input batch of images/videos. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded images/videos. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + assert len(x.shape) == 5, "The input tensor should have 5 dimensions." + + if self.use_temporal_tiling and x.shape[2] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x, return_dict=return_dict) + + if self.use_spatial_tiling and ( + x.shape[-1] > self.tile_sample_min_size + or x.shape[-2] > self.tile_sample_min_size): + return self.spatial_tiled_encode(x, return_dict=return_dict) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self.encoder(x) + + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode( + self, + z: torch.FloatTensor, + return_dict: bool = True + ) -> Union[DecoderOutput, torch.FloatTensor]: + assert len(z.shape) == 5, "The input tensor should have 5 dimensions." + + if self.use_parallel: + return self.parallel_tiled_decode(z, return_dict=return_dict) + + if self.use_temporal_tiling and z.shape[2] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z, return_dict=return_dict) + + if self.use_spatial_tiling and ( + z.shape[-1] > self.tile_latent_min_size + or z.shape[-2] > self.tile_latent_min_size): + return self.spatial_tiled_decode(z, return_dict=return_dict) + + z = self.post_quant_conv(z) + dec = self.decoder(z) + + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode(self, + z: torch.FloatTensor, + return_dict: bool = True, + generator=None) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Decode a batch of images/videos. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [ + self._decode(z_slice).sample for z_slice in z.split(1) + ] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded, ) + + return DecoderOutput(sample=decoded) + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, + blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( + 1 - y / blend_extent) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, + blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( + 1 - x / blend_extent) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, + blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * ( + 1 - x / blend_extent) + b[:, :, x, :, :] * (x / blend_extent) + return b + + def spatial_tiled_encode( + self, + x: torch.FloatTensor, + return_dict: bool = True, + return_moments: bool = False, + ) -> AutoencoderKLOutput: + r"""Encode a batch of images/videos using a tiled encoder. + + When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several + steps. This is useful to keep memory use constant regardless of image/videos size. The end result of tiled encoding is + different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the + tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the + output, but they should be much less noticeable. + + Args: + x (`torch.FloatTensor`): Input batch of images/videos. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + overlap_size = int(self.tile_sample_min_size * + (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * + self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split video into tiles and encode them separately. + rows = [] + for i in range(0, x.shape[-2], overlap_size): + row = [] + for j in range(0, x.shape[-1], overlap_size): + tile = x[:, :, :, i:i + self.tile_sample_min_size, + j:j + self.tile_sample_min_size, ] + tile = self.encoder(tile) + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + moments = torch.cat(result_rows, dim=-2) + if return_moments: + return moments + + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def spatial_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True + ) -> Union[DecoderOutput, torch.FloatTensor]: + r""" + Decode a batch of images/videos using a tiled decoder. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + """ + overlap_size = int(self.tile_latent_min_size * + (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * + self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[-2], overlap_size): + row = [] + for j in range(0, z.shape[-1], overlap_size): + tile = z[:, :, :, i:i + self.tile_latent_min_size, + j:j + self.tile_latent_min_size, ] + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=-2) + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + def temporal_tiled_encode(self, + x: torch.FloatTensor, + return_dict: bool = True) -> AutoencoderKLOutput: + + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * + (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_tsize * + self.tile_overlap_factor) + t_limit = self.tile_latent_min_tsize - blend_extent + + # Split the video into tiles and encode them separately. + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i:i + self.tile_sample_min_tsize + 1, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_sample_min_size + or tile.shape[-2] > self.tile_sample_min_size): + tile = self.spatial_tiled_encode(tile, return_moments=True) + else: + tile = self.encoder(tile) + tile = self.quant_conv(tile) + if i > 0: + tile = tile[:, :, 1:, :, :] + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + + moments = torch.cat(result_row, dim=2) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior, ) + + return AutoencoderKLOutput(latent_dist=posterior) + + def temporal_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True + ) -> Union[DecoderOutput, torch.FloatTensor]: + # Split z into overlapping tiles and decode them separately. + + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * + (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_tsize * + self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - blend_extent + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i:i + self.tile_latent_min_tsize + 1, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_latent_min_size + or tile.shape[-2] > self.tile_latent_min_size): + decoded = self.spatial_tiled_decode(tile, + return_dict=True).sample + else: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + if i > 0: + decoded = decoded[:, :, 1:, :, :] + row.append(decoded) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + else: + result_row.append(tile[:, :, :t_limit + 1, :, :]) + + dec = torch.cat(result_row, dim=2) + if not return_dict: + return (dec, ) + + return DecoderOutput(sample=dec) + + def _parallel_data_generator(self, gathered_results, + gathered_dim_metadata): + global_idx = 0 + for i, per_rank_metadata in enumerate(gathered_dim_metadata): + _start_shape = 0 + for shape in per_rank_metadata: + mul_shape = prod(shape) + yield (gathered_results[i, _start_shape:_start_shape + + mul_shape].reshape(shape), global_idx) + _start_shape += mul_shape + global_idx += 1 + + def parallel_tiled_decode(self, + z: torch.FloatTensor, + return_dict: bool = True + ) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs + """ + world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group + B, C, T, H, W = z.shape + + # Calculate parameters + t_overlap_size = int(self.tile_latent_min_tsize * + (1 - self.tile_overlap_factor)) + t_blend_extent = int(self.tile_sample_min_tsize * + self.tile_overlap_factor) + t_limit = self.tile_sample_min_tsize - t_blend_extent + + s_overlap_size = int(self.tile_latent_min_size * + (1 - self.tile_overlap_factor)) + s_blend_extent = int(self.tile_sample_min_size * + self.tile_overlap_factor) + s_row_limit = self.tile_sample_min_size - s_blend_extent + + # Calculate tile dimensions + num_t_tiles = (T + t_overlap_size - 1) // t_overlap_size + num_h_tiles = (H + s_overlap_size - 1) // s_overlap_size + num_w_tiles = (W + s_overlap_size - 1) // s_overlap_size + total_spatial_tiles = num_h_tiles * num_w_tiles + total_tiles = num_t_tiles * total_spatial_tiles + + # Calculate tiles per rank and padding + tiles_per_rank = (total_tiles + world_size - 1) // world_size + start_tile_idx = rank * tiles_per_rank + end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles) + + local_results = [] + local_dim_metadata = [] + # Process assigned tiles + for local_idx, global_idx in enumerate( + range(start_tile_idx, end_tile_idx)): + # Convert flat index to 3D indices + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + + # Calculate positions + t_start = t_idx * t_overlap_size + h_start = h_idx * s_overlap_size + w_start = w_idx * s_overlap_size + + # Extract and process tile + tile = z[:, :, t_start:t_start + self.tile_latent_min_tsize + 1, + h_start:h_start + self.tile_latent_min_size, + w_start:w_start + self.tile_latent_min_size] + + # Process tile + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + + if t_start > 0: + decoded = decoded[:, :, 1:, :, :] + + # Store metadata + shape = decoded.shape + # Store decoded data (flattened) + decoded_flat = decoded.reshape(-1) + local_results.append(decoded_flat) + local_dim_metadata.append(shape) + + results = torch.cat(local_results, dim=0).contiguous() + del local_results + torch.cuda.empty_cache() + # first gather size to pad the results + local_size = torch.tensor([results.size(0)], + device=results.device, + dtype=torch.int64) + all_sizes = [ + torch.zeros(1, device=results.device, dtype=torch.int64) + for _ in range(world_size) + ] + dist.all_gather(all_sizes, local_size) + max_size = max(size.item() for size in all_sizes) + padded_results = torch.zeros(max_size, device=results.device) + padded_results[:results.size(0)] = results + del results + torch.cuda.empty_cache() + # Gather all results + gathered_dim_metadata = [None] * world_size + gathered_results = torch.zeros_like(padded_results).repeat( + world_size, *[1] * len(padded_results.shape) + ).contiguous( + ) # use contiguous to make sure it won't copy data in the following operations + dist.all_gather_into_tensor(gathered_results, padded_results) + dist.all_gather_object(gathered_dim_metadata, local_dim_metadata) + # Process gathered results + data = [[[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)] + for _ in range(num_t_tiles)] + for current_data, global_idx in self._parallel_data_generator( + gathered_results, gathered_dim_metadata): + t_idx = global_idx // total_spatial_tiles + spatial_idx = global_idx % total_spatial_tiles + h_idx = spatial_idx // num_w_tiles + w_idx = spatial_idx % num_w_tiles + data[t_idx][h_idx][w_idx] = current_data + # Merge results + result_slices = [] + last_slice_data = None + for i, tem_data in enumerate(data): + slice_data = self._merge_spatial_tiles(tem_data, s_blend_extent, + s_row_limit) + if i > 0: + slice_data = self.blend_t(last_slice_data, slice_data, + t_blend_extent) + result_slices.append(slice_data[:, :, :t_limit, :, :]) + else: + result_slices.append(slice_data[:, :, :t_limit + 1, :, :]) + last_slice_data = slice_data + dec = torch.cat(result_slices, dim=2) + + if not return_dict: + return (dec, ) + return DecoderOutput(sample=dec) + + def _merge_spatial_tiles(self, spatial_rows, blend_extent, row_limit): + """Helper function to merge spatial tiles with blending""" + result_rows = [] + for i, row in enumerate(spatial_rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(spatial_rows[i - 1][j], tile, + blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + def forward( + self, + sample: torch.FloatTensor, + sample_posterior: bool = False, + return_dict: bool = True, + return_posterior: bool = False, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput2, torch.FloatTensor]: + r""" + Args: + sample (`torch.FloatTensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z).sample + + if not return_dict: + if return_posterior: + return (dec, posterior) + else: + return (dec, ) + if return_posterior: + return DecoderOutput2(sample=dec, posterior=posterior) + else: + return DecoderOutput2(sample=dec) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections + def fuse_qkv_projections(self): + """ + Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, + key, value) are fused. For cross-attention modules, key and value projection matrices are fused. + + + + This API is 🧪 experimental. + + + """ + self.original_attn_processors = None + + for _, attn_processor in self.attn_processors.items(): + if "Added" in str(attn_processor.__class__.__name__): + raise ValueError( + "`fuse_qkv_projections()` is not supported for models having added KV projections." + ) + + self.original_attn_processors = self.attn_processors + + for module in self.modules(): + if isinstance(module, Attention): + module.fuse_projections(fuse=True) + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections + def unfuse_qkv_projections(self): + """Disables the fused QKV projection if enabled. + + + + This API is 🧪 experimental. + + + + """ + if self.original_attn_processors is not None: + self.set_attn_processor(self.original_attn_processors) diff --git a/fastvideo/models/hunyuan/vae/unet_causal_3d_blocks.py b/fastvideo/models/hunyuan/vae/unet_causal_3d_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..37dce39932f3c492e0adba605c7af5a4c81f403a --- /dev/null +++ b/fastvideo/models/hunyuan/vae/unet_causal_3d_blocks.py @@ -0,0 +1,829 @@ +# Copyright 2024 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Modified from diffusers==0.29.2 +# +# ============================================================================== + +from typing import Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import Attention, SpatialNorm +from diffusers.models.normalization import AdaGroupNorm, RMSNorm +from diffusers.utils import logging +from einops import rearrange +from torch import nn + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def prepare_causal_attention_mask(n_frame: int, + n_hw: int, + dtype, + device, + batch_size: int = None): + seq_len = n_frame * n_hw + mask = torch.full((seq_len, seq_len), + float("-inf"), + dtype=dtype, + device=device) + for i in range(seq_len): + i_frame = i // n_hw + mask[i, :(i_frame + 1) * n_hw] = 0 + if batch_size is not None: + mask = mask.unsqueeze(0).expand(batch_size, -1, -1) + return mask + + +class CausalConv3d(nn.Module): + """ + Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations. + This maintains temporal causality in video generation tasks. + """ + + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]] = 1, + dilation: Union[int, Tuple[int, int, int]] = 1, + pad_mode="replicate", + **kwargs, + ): + super().__init__() + + self.pad_mode = pad_mode + padding = ( + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size // 2, + kernel_size - 1, + 0, + ) # W, H, T + self.time_causal_padding = padding + + self.conv = nn.Conv3d(chan_in, + chan_out, + kernel_size, + stride=stride, + dilation=dilation, + **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) + return self.conv(x) + + +class UpsampleCausal3D(nn.Module): + """ + A 3D upsampling layer with an optional convolution. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + kernel_size: Optional[int] = None, + padding=1, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + interpolate=True, + upsample_factor=(2, 2, 2), + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + self.interpolate = interpolate + self.upsample_factor = upsample_factor + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + conv = None + if use_conv_transpose: + raise NotImplementedError + elif use_conv: + if kernel_size is None: + kernel_size = 3 + conv = CausalConv3d(self.channels, + self.out_channels, + kernel_size=kernel_size, + bias=bias) + + if name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + output_size: Optional[int] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + raise NotImplementedError + + if self.use_conv_transpose: + return self.conv(hidden_states) + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if self.interpolate: + B, C, T, H, W = hidden_states.shape + first_h, other_h = hidden_states.split((1, T - 1), dim=2) + if output_size is None: + if T > 1: + other_h = F.interpolate(other_h, + scale_factor=self.upsample_factor, + mode="nearest") + + first_h = first_h.squeeze(2) + first_h = F.interpolate(first_h, + scale_factor=self.upsample_factor[1:], + mode="nearest") + first_h = first_h.unsqueeze(2) + else: + raise NotImplementedError + + if T > 1: + hidden_states = torch.cat((first_h, other_h), dim=2) + else: + hidden_states = first_h + + # If the input is bfloat16, we cast back to bfloat16 + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + if self.use_conv: + if self.name == "conv": + hidden_states = self.conv(hidden_states) + else: + hidden_states = self.Conv2d_0(hidden_states) + + return hidden_states + + +class DownsampleCausal3D(nn.Module): + """ + A 3D downsampling layer with an optional convolution. + """ + + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + padding: int = 1, + name: str = "conv", + kernel_size=3, + norm_type=None, + eps=None, + elementwise_affine=None, + bias=True, + stride=2, + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = stride + self.name = name + + if norm_type == "ln_norm": + self.norm = nn.LayerNorm(channels, eps, elementwise_affine) + elif norm_type == "rms_norm": + self.norm = RMSNorm(channels, eps, elementwise_affine) + elif norm_type is None: + self.norm = None + else: + raise ValueError(f"unknown norm_type: {norm_type}") + + if use_conv: + conv = CausalConv3d( + self.channels, + self.out_channels, + kernel_size=kernel_size, + stride=stride, + bias=bias, + ) + else: + raise NotImplementedError + + if name == "conv": + self.Conv2d_0 = conv + self.conv = conv + elif name == "Conv2d_0": + self.conv = conv + else: + self.conv = conv + + def forward(self, + hidden_states: torch.FloatTensor, + scale: float = 1.0) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.norm is not None: + hidden_states = self.norm(hidden_states.permute(0, 2, 3, + 1)).permute( + 0, 3, 1, 2) + + assert hidden_states.shape[1] == self.channels + + hidden_states = self.conv(hidden_states) + + return hidden_states + + +class ResnetBlockCausal3D(nn.Module): + r""" + A Resnet block. + """ + + def __init__( + self, + *, + in_channels: int, + out_channels: Optional[int] = None, + conv_shortcut: bool = False, + dropout: float = 0.0, + temb_channels: int = 512, + groups: int = 32, + groups_out: Optional[int] = None, + pre_norm: bool = True, + eps: float = 1e-6, + non_linearity: str = "swish", + skip_time_act: bool = False, + # default, scale_shift, ada_group, spatial + time_embedding_norm: str = "default", + kernel: Optional[torch.FloatTensor] = None, + output_scale_factor: float = 1.0, + use_in_shortcut: Optional[bool] = None, + up: bool = False, + down: bool = False, + conv_shortcut_bias: bool = True, + conv_3d_out_channels: Optional[int] = None, + ): + super().__init__() + self.pre_norm = pre_norm + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.up = up + self.down = down + self.output_scale_factor = output_scale_factor + self.time_embedding_norm = time_embedding_norm + self.skip_time_act = skip_time_act + + linear_cls = nn.Linear + + if groups_out is None: + groups_out = groups + + if self.time_embedding_norm == "ada_group": + self.norm1 = AdaGroupNorm(temb_channels, + in_channels, + groups, + eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm1 = SpatialNorm(in_channels, temb_channels) + else: + self.norm1 = torch.nn.GroupNorm(num_groups=groups, + num_channels=in_channels, + eps=eps, + affine=True) + + self.conv1 = CausalConv3d(in_channels, + out_channels, + kernel_size=3, + stride=1) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + self.time_emb_proj = linear_cls(temb_channels, out_channels) + elif self.time_embedding_norm == "scale_shift": + self.time_emb_proj = linear_cls(temb_channels, + 2 * out_channels) + elif (self.time_embedding_norm == "ada_group" + or self.time_embedding_norm == "spatial"): + self.time_emb_proj = None + else: + raise ValueError( + f"Unknown time_embedding_norm : {self.time_embedding_norm} " + ) + else: + self.time_emb_proj = None + + if self.time_embedding_norm == "ada_group": + self.norm2 = AdaGroupNorm(temb_channels, + out_channels, + groups_out, + eps=eps) + elif self.time_embedding_norm == "spatial": + self.norm2 = SpatialNorm(out_channels, temb_channels) + else: + self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, + num_channels=out_channels, + eps=eps, + affine=True) + + self.dropout = torch.nn.Dropout(dropout) + conv_3d_out_channels = conv_3d_out_channels or out_channels + self.conv2 = CausalConv3d(out_channels, + conv_3d_out_channels, + kernel_size=3, + stride=1) + + self.nonlinearity = get_activation(non_linearity) + + self.upsample = self.downsample = None + if self.up: + self.upsample = UpsampleCausal3D(in_channels, use_conv=False) + elif self.down: + self.downsample = DownsampleCausal3D(in_channels, + use_conv=False, + name="op") + + self.use_in_shortcut = (self.in_channels != conv_3d_out_channels if + use_in_shortcut is None else use_in_shortcut) + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = CausalConv3d( + in_channels, + conv_3d_out_channels, + kernel_size=1, + stride=1, + bias=conv_shortcut_bias, + ) + + def forward( + self, + input_tensor: torch.FloatTensor, + temb: torch.FloatTensor, + scale: float = 1.0, + ) -> torch.FloatTensor: + hidden_states = input_tensor + + if (self.time_embedding_norm == "ada_group" + or self.time_embedding_norm == "spatial"): + hidden_states = self.norm1(hidden_states, temb) + else: + hidden_states = self.norm1(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = self.upsample(input_tensor, scale=scale) + hidden_states = self.upsample(hidden_states, scale=scale) + elif self.downsample is not None: + input_tensor = self.downsample(input_tensor, scale=scale) + hidden_states = self.downsample(hidden_states, scale=scale) + + hidden_states = self.conv1(hidden_states) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb, scale)[:, :, None, None] + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + if (self.time_embedding_norm == "ada_group" + or self.time_embedding_norm == "spatial"): + hidden_states = self.norm2(hidden_states, temb) + else: + hidden_states = self.norm2(hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + + hidden_states) / self.output_scale_factor + + return output_tensor + + +def get_down_block3d( + down_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + temb_channels: int, + add_downsample: bool, + downsample_stride: int, + resnet_eps: float, + resnet_act_fn: str, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + downsample_padding: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + downsample_type: Optional[str] = None, + dropout: float = 0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = (down_block_type[7:] + if down_block_type.startswith("UNetRes") else + down_block_type) + if down_block_type == "DownEncoderBlockCausal3D": + return DownEncoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + add_downsample=add_downsample, + downsample_stride=downsample_stride, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block3d( + up_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + add_upsample: bool, + upsample_scale_factor: Tuple, + resnet_eps: float, + resnet_act_fn: str, + resolution_idx: Optional[int] = None, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + upsample_type: Optional[str] = None, + dropout: float = 0.0, +) -> nn.Module: + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = (up_block_type[7:] + if up_block_type.startswith("UNetRes") else up_block_type) + if up_block_type == "UpDecoderBlockCausal3D": + return UpDecoderBlockCausal3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + upsample_scale_factor=upsample_scale_factor, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + temb_channels=temb_channels, + ) + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlockCausal3D(nn.Module): + """ + A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks. + """ + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + attn_groups: Optional[int] = None, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + ): + super().__init__() + resnet_groups = (resnet_groups if resnet_groups is not None else min( + in_channels // 4, 32)) + self.add_attention = add_attention + + if attn_groups is None: + attn_groups = (resnet_groups + if resnet_time_scale_shift == "default" else None) + + # there is always at least one resnet + resnets = [ + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=attn_groups, + spatial_norm_dim=(temb_channels + if resnet_time_scale_shift + == "spatial" else None), + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + )) + else: + attentions.append(None) + + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None) -> torch.FloatTensor: + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + B, C, T, H, W = hidden_states.shape + hidden_states = rearrange(hidden_states, + "b c f h w -> b (f h w) c") + attention_mask = prepare_causal_attention_mask( + T, + H * W, + hidden_states.dtype, + hidden_states.device, + batch_size=B) + hidden_states = attn(hidden_states, + temb=temb, + attention_mask=attention_mask) + hidden_states = rearrange(hidden_states, + "b (f h w) c -> b c f h w", + f=T, + h=H, + w=W) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class DownEncoderBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_stride: int = 2, + downsample_padding: int = 1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlockCausal3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList([ + DownsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + stride=downsample_stride, + ) + ]) + else: + self.downsamplers = None + + def forward(self, + hidden_states: torch.FloatTensor, + scale: float = 1.0) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=None, scale=scale) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale) + + return hidden_states + + +class UpDecoderBlockCausal3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + upsample_scale_factor=(2, 2, 2), + temb_channels: Optional[int] = None, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlockCausal3D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + )) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList([ + UpsampleCausal3D( + out_channels, + use_conv=True, + out_channels=out_channels, + upsample_factor=upsample_scale_factor, + ) + ]) + else: + self.upsamplers = None + + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb=temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states diff --git a/fastvideo/models/hunyuan/vae/vae.py b/fastvideo/models/hunyuan/vae/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..117da9afe7dd336a6f68242842a8d0c3e7133e31 --- /dev/null +++ b/fastvideo/models/hunyuan/vae/vae.py @@ -0,0 +1,385 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn +from diffusers.models.attention_processor import SpatialNorm +from diffusers.utils import BaseOutput, is_torch_version +from diffusers.utils.torch_utils import randn_tensor + +from .unet_causal_3d_blocks import (CausalConv3d, UNetMidBlockCausal3D, + get_down_block3d, get_up_block3d) + + +@dataclass +class DecoderOutput(BaseOutput): + r""" + Output of decoding method. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The decoded output sample from the last layer of the model. + """ + + sample: torch.FloatTensor + + +class EncoderCausal3D(nn.Module): + r""" + The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, + block_out_channels[0], + kernel_size=3, + stride=1) + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_downsample_layers = int( + np.log2(spatial_compression_ratio)) + num_time_downsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_downsample = bool( + i < num_spatial_downsample_layers) + add_time_downsample = bool( + i >= + (len(block_out_channels) - 1 - num_time_downsample_layers) + and not is_final_block) + else: + raise ValueError( + f"Unsupported time_compression_ratio: {time_compression_ratio}." + ) + + downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1) + downsample_stride_T = (2, ) if add_time_downsample else (1, ) + downsample_stride = tuple(downsample_stride_T + + downsample_stride_HW) + down_block = get_down_block3d( + down_block_type, + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=bool(add_spatial_downsample + or add_time_downsample), + downsample_stride=downsample_stride, + resnet_eps=1e-6, + downsample_padding=0, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=None, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + ) + + # out + self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], + num_groups=norm_num_groups, + eps=1e-6) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = CausalConv3d(block_out_channels[-1], + conv_out_channels, + kernel_size=3) + + def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor: + r"""The forward method of the `EncoderCausal3D` class.""" + assert len( + sample.shape) == 5, "The input tensor should have 5 dimensions" + + sample = self.conv_in(sample) + + # down + for down_block in self.down_blocks: + sample = down_block(sample) + + # middle + sample = self.mid_block(sample) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DecoderCausal3D(nn.Module): + r""" + The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D", ), + block_out_channels: Tuple[int, ...] = (64, ), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", # group, spatial + mid_block_add_attention=True, + time_compression_ratio: int = 4, + spatial_compression_ratio: int = 8, + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.conv_in = CausalConv3d(in_channels, + block_out_channels[-1], + kernel_size=3, + stride=1) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + temb_channels = in_channels if norm_type == "spatial" else None + + # mid + self.mid_block = UNetMidBlockCausal3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default" + if norm_type == "group" else norm_type, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + num_spatial_upsample_layers = int( + np.log2(spatial_compression_ratio)) + num_time_upsample_layers = int(np.log2(time_compression_ratio)) + + if time_compression_ratio == 4: + add_spatial_upsample = bool(i < num_spatial_upsample_layers) + add_time_upsample = bool( + i >= len(block_out_channels) - 1 - num_time_upsample_layers + and not is_final_block) + else: + raise ValueError( + f"Unsupported time_compression_ratio: {time_compression_ratio}." + ) + + upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, + 1) + upsample_scale_factor_T = (2, ) if add_time_upsample else (1, ) + upsample_scale_factor = tuple(upsample_scale_factor_T + + upsample_scale_factor_HW) + up_block = get_up_block3d( + up_block_type, + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + prev_output_channel=None, + add_upsample=bool(add_spatial_upsample or add_time_upsample), + upsample_scale_factor=upsample_scale_factor, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + attention_head_dim=output_channel, + temb_channels=temb_channels, + resnet_time_scale_shift=norm_type, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_type == "spatial": + self.conv_norm_out = SpatialNorm(block_out_channels[0], + temb_channels) + else: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=1e-6) + self.conv_act = nn.SiLU() + self.conv_out = CausalConv3d(block_out_channels[0], + out_channels, + kernel_size=3) + + self.gradient_checkpointing = False + + def forward( + self, + sample: torch.FloatTensor, + latent_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + r"""The forward method of the `DecoderCausal3D` class.""" + assert len( + sample.shape) == 5, "The input tensor should have 5 dimensions." + + sample = self.conv_in(sample) + + upscale_dtype = next(iter(self.up_blocks.parameters())).dtype + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), + sample, + latent_embeds, + use_reentrant=False, + ) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), + sample, + latent_embeds, + use_reentrant=False, + ) + else: + # middle + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(self.mid_block), sample, + latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), sample, latent_embeds) + else: + # middle + sample = self.mid_block(sample, latent_embeds) + sample = sample.to(upscale_dtype) + + # up + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds) + + # post-process + if latent_embeds is None: + sample = self.conv_norm_out(sample) + else: + sample = self.conv_norm_out(sample, latent_embeds) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + +class DiagonalGaussianDistribution(object): + + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + + def sample( + self, + generator: Optional[torch.Generator] = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + else: + reduce_dim = list(range(1, self.mean.ndim)) + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=reduce_dim, + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + dim=reduce_dim, + ) + + def nll(self, + sample: torch.Tensor, + dims: Tuple[int, ...] = [1, 2, 3]) -> torch.Tensor: + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + + def mode(self) -> torch.Tensor: + return self.mean diff --git a/scripts/finetune/finetune_g2rpo_hps.sh b/scripts/finetune/finetune_g2rpo_hps.sh new file mode 100644 index 0000000000000000000000000000000000000000..e579760543ad1189897a38b9b6fd576801d3e1c0 --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_hps.sh @@ -0,0 +1,35 @@ +torchrun --nnodes=1 --nproc_per_node=4 --node_rank 0 \ + fastvideo/train_g2rpo_hps.py \ + --seed 42 \ + --pretrained_model_name_or_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux \ + --hps_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --data_json_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/rl_embeddings/videos2caption.json \ + --gradient_checkpointing \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 50 \ + --cfg 0.0 \ + --output_dir /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp/hps_multistep \ + --h 1024 \ + --w 1024 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 2 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 3 4 5 6 7 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_hps_clip.sh b/scripts/finetune/finetune_g2rpo_hps_clip.sh new file mode 100644 index 0000000000000000000000000000000000000000..2e15aad1d7354a075e6cf8a3de7dc5a2e9a786b4 --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_hps_clip.sh @@ -0,0 +1,42 @@ +export NNODES=${NODE_COUNT:-2} +export PROC_PER_NODE=${PROC_PER_NODE:-8} +export MASTER_ADDR=${MASTER_ADDR} +export NODE_RANK=${NODE_RANK} +export MASTER_PORT=29533 + +torchrun --nnodes=2 --nproc_per_node=$PROC_PER_NODE --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT \ + fastvideo/train_g2rpo_hps_clip_merge.py \ + --seed 42 \ + --pretrained_model_name_or_path ckpt/flux \ + --hps_path ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --clip_score_path ckpt/clip_score \ + --data_json_path data/rl_embeddings/videos2caption.json \ + --gradient_checkpointing \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 50 \ + --cfg 0.0 \ + --output_dir save_exp/hps_clip \ + --h 720 \ + --w 720 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 12 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_hps_clip_merge.sh b/scripts/finetune/finetune_g2rpo_hps_clip_merge.sh new file mode 100644 index 0000000000000000000000000000000000000000..67369ee3dfd2dd9037d96b87beb3f6f649fe072b --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_hps_clip_merge.sh @@ -0,0 +1,38 @@ +CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nnodes=1 --nproc_per_node=4 --node_rank 0 \ + fastvideo/train_g2rpo_hps_clip_merge.py \ + --seed 42 \ + --pretrained_model_name_or_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux \ + --resume_ckpt /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp/hps_clip_merge_step/ckpt/checkpoint-200-0 \ + --hps_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --clip_score_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/clip_score \ + --data_json_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/rl_embeddings/videos2caption.json \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --init_steps 200 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 10 \ + --cfg 0.0 \ + --output_dir /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp/hps_clip_merge_resume_200 \ + --h 1024 \ + --w 1024 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 2 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 \ + --eta_step_merge_list 1 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_hps_merge.sh b/scripts/finetune/finetune_g2rpo_hps_merge.sh new file mode 100644 index 0000000000000000000000000000000000000000..d8197156f689ebf47d0bbd90b01841cf5a374c26 --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_hps_merge.sh @@ -0,0 +1,36 @@ +CUDA_VISIBLE_DEVICES=1,2,3 torchrun --nnodes=1 --nproc_per_node=3 --node_rank 0 \ + fastvideo/train_g2rpo_hps_merge.py \ + --seed 42 \ + --pretrained_model_name_or_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux \ + --hps_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --data_json_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/rl_embeddings/videos2caption.json \ + --gradient_checkpointing \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 50 \ + --cfg 0.0 \ + --output_dir /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp/hps_merge_step_test \ + --h 1024 \ + --w 1024 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 2 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 4 \ + --eta_step_merge_list 1 1 2 3 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_rfpt.sh b/scripts/finetune/finetune_g2rpo_rfpt.sh new file mode 100644 index 0000000000000000000000000000000000000000..2d58452cd3ed6f8d5adf05325965aab8ce3234fe --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_rfpt.sh @@ -0,0 +1,41 @@ +export NNODES=${NODE_COUNT:-2} +export PROC_PER_NODE=${PROC_PER_NODE:-8} +export MASTER_ADDR=${MASTER_ADDR} +export NODE_RANK=${NODE_RANK} +export MASTER_PORT=29513 + +torchrun --nnodes=1 --nproc_per_node=4 --node_rank 0 \ + fastvideo/train_g2rpo_rfpt.py \ + --seed 42 \ + --pretrained_model_name_or_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux \ + --hps_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --data_json_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/datasets/flux_rl_embeddings/videos2caption.json \ + --gradient_checkpointing \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 50 \ + --cfg 0.0 \ + --output_dir /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp_rlpt/hps_gt \ + --h 1024 \ + --w 1024 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 12 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 3 4 5 6 7 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_rlpt.sh b/scripts/finetune/finetune_g2rpo_rlpt.sh new file mode 100644 index 0000000000000000000000000000000000000000..4a51ebf5b6714c7f40f598e344ee1cce2858e3fe --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_rlpt.sh @@ -0,0 +1,51 @@ +#!/bin/bash + +# LAION-220k Image-Text RL Training Configuration +echo "[INFO] Starting LAION-220k Image-Text RL Training..." +echo "[INFO] Using GT images from /data2/dataset/laion-220k/images" +echo "[INFO] Using embeddings from data/laion_rl_embeddings/videos2caption.json" + +torchrun --nproc_per_node=8 --master_port 11451 \ + fastvideo/train_grpo_rlpt.py \ + --seed 42 \ + --pretrained_model_name_or_path ckpt/flux \ + --hps_path ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --clip_score_path ckpt/clip_score \ + --data_json_path data/laion_rl_embeddings/videos2caption.json \ + --image_data_dir /data2/dataset/laion-220k/images \ + --log_file save_exp/laion_hps_clip_mse/training_logs.csv \ + --gradient_checkpointing \ + --train_batch_size 2 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 2 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 10 \ + --cfg 0.0 \ + --output_dir save_exp/laion_hps_clip_mse \ + --h 512 \ + --w 512 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 8 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 3 \ + --granular_list 1 \ + --use_hps_reward \ + --use_clip_reward \ + --use_mse_reward \ + --hps_reward_weight 1.0 \ + --clip_reward_weight 1.0 \ + --mse_reward_weight 1.0 \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_rlpt_dino.sh b/scripts/finetune/finetune_g2rpo_rlpt_dino.sh new file mode 100644 index 0000000000000000000000000000000000000000..05496a56869dd5eee00f4ca2f04c248f8c9264be --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_rlpt_dino.sh @@ -0,0 +1,42 @@ +export NNODES=${NODE_COUNT:-2} +export PROC_PER_NODE=${PROC_PER_NODE:-8} +export MASTER_ADDR=${MASTER_ADDR} +export NODE_RANK=${NODE_RANK} +export MASTER_PORT=29513 + +torchrun --nnodes=1 --nproc_per_node=4 --node_rank 0 \ + fastvideo/train_g2rpo_rlpt_dino.py \ + --seed 42 \ + --pretrained_model_name_or_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux \ + --hps_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --dino_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/dinov2 \ + --data_json_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/datasets/flux_rl_embeddings/videos2caption.json \ + --gradient_checkpointing \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 301 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 50 \ + --cfg 0.0 \ + --output_dir /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp_rlpt/hps_gt \ + --h 1024 \ + --w 1024 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 12 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 3 4 5 6 7 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/finetune/finetune_g2rpo_rlpt_from_noise.sh b/scripts/finetune/finetune_g2rpo_rlpt_from_noise.sh new file mode 100644 index 0000000000000000000000000000000000000000..ca59f17606cb2abf1483ea402b371b6dd9d15c4b --- /dev/null +++ b/scripts/finetune/finetune_g2rpo_rlpt_from_noise.sh @@ -0,0 +1,42 @@ +export NNODES=${NODE_COUNT:-2} +export PROC_PER_NODE=${PROC_PER_NODE:-8} +export MASTER_ADDR=${MASTER_ADDR} +export NODE_RANK=${NODE_RANK} +export MASTER_PORT=29513 + +torchrun --nnodes=1 --nproc_per_node=4 --node_rank 0 \ + fastvideo/train_g2rpo_rlpt_from_noise.py \ + --seed 42 \ + --pretrained_model_name_or_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux \ + --hps_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/hps/HPS_v2.1_compressed.pt \ + --hps_clip_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/CLIP-ViT-H-14-laion2B-s32B-b79K/open_clip_pytorch_model.bin \ + --dino_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/dinov2 \ + --data_json_path /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/datasets/flux_rl_embeddings/videos2caption.json \ + --gradient_checkpointing \ + --train_batch_size 1 \ + --num_latent_t 1 \ + --sp_size 1 \ + --train_sp_batch_size 1 \ + --dataloader_num_workers 4 \ + --max_train_steps 151 \ + --learning_rate 2e-6 \ + --mixed_precision bf16 \ + --checkpointing_steps 30 \ + --cfg 0.0 \ + --output_dir /mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/save_exp_rlpt/from_noise \ + --h 1024 \ + --w 1024 \ + --t 1 \ + --sampling_steps 16 \ + --eta 0.7 \ + --lr_warmup_steps 0 \ + --sampler_seed 1223627 \ + --max_grad_norm 1.0 \ + --weight_decay 0.0001 \ + --num_generations 12 \ + --shift 3 \ + --init_same_noise \ + --clip_range 1e-4 \ + --adv_clip_max 5.0 \ + --eta_step_list 0 1 2 3 4 5 6 7 \ + --granular_list 1 \ \ No newline at end of file diff --git a/scripts/huggingface/download_hf.py b/scripts/huggingface/download_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..0af23c69d8d13a090927ed9180c0b31fd01885eb --- /dev/null +++ b/scripts/huggingface/download_hf.py @@ -0,0 +1,43 @@ +import argparse + +from huggingface_hub import hf_hub_download, snapshot_download + + +# set args for repo_id, local_dir, repo_type, + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Download a dataset or model from the Hugging Face Hub") + parser.add_argument("--repo_id", + type=str, + help="The ID of the repository to download") + parser.add_argument( + "--local_dir", + type=str, + help="The local directory to download the repository to", + ) + parser.add_argument( + "--repo_type", + type=str, + default="model", + help="The type of repository to download (dataset or model)", + ) + parser.add_argument("--file_name", + type=str, + help="The file name to download") + args = parser.parse_args() + if args.file_name: + hf_hub_download( + repo_id=args.repo_id, + filename=args.file_name, + repo_type=args.repo_type, + local_dir=args.local_dir, + ) + else: + snapshot_download( + repo_id=args.repo_id, + local_dir=args.local_dir, + repo_type=args.repo_type, + local_dir_use_symlinks=False, + resume_download=True, + ) diff --git a/scripts/inference/infer.py b/scripts/inference/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..1bc0d220bcc375949e8d523aa777995854a2b2bf --- /dev/null +++ b/scripts/inference/infer.py @@ -0,0 +1,28 @@ +import torch +from diffusers import FluxPipeline +from diffusers import FluxTransformer2DModel +from safetensors.torch import load_file + +device = "cuda:0" + +model_path = "ckpt/g2rpo/diffusion_pytorch_model.safetensors" +flux_path = "ckpt/flux" + +pipe = FluxPipeline.from_pretrained(flux_path, use_safetensors=True, torch_dtype=torch.float16) +model_state_dict = load_file(model_path) +pipe.transformer.load_state_dict(model_state_dict, strict=True) +pipe = pipe.to(device) + +prompt = "A golden Labrador retriever is leaping excitedly on the green grass, chasing a soap bubble that glows with a rainbow in the sun, National Geographic photography style" + +image = pipe( + prompt, + guidance_scale=3.5, + height=1024, + width=1024, + num_inference_steps=50, + max_sequence_length=512, +).images[0] + +save_path = "g2rpo.png" +image.save(save_path) \ No newline at end of file diff --git a/scripts/preprocess/preprocess_flux_rl_embeddings.sh b/scripts/preprocess/preprocess_flux_rl_embeddings.sh new file mode 100644 index 0000000000000000000000000000000000000000..0c3cb3370ca7335402b228de3a94c1b6f06ea8cb --- /dev/null +++ b/scripts/preprocess/preprocess_flux_rl_embeddings.sh @@ -0,0 +1,9 @@ +GPU_NUM=4 # 2,4,8 +MODEL_PATH="/mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/ckpt/flux" +OUTPUT_DIR="/mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/flux_rl_embeddings" + +torchrun --nproc_per_node=$GPU_NUM --master_port 19002 \ + fastvideo/data_preprocess/preprocess_flux_rfpt_embedding.py \ + --model_path $MODEL_PATH \ + --output_dir $OUTPUT_DIR \ + --prompt_dir "/mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/datasets/Aesthetics-Part01" \ No newline at end of file diff --git a/scripts/preprocess/preprocess_flux_rlpt_embeddings.sh b/scripts/preprocess/preprocess_flux_rlpt_embeddings.sh new file mode 100644 index 0000000000000000000000000000000000000000..fa7dfd228e51f6a24173a90419810794029d2f8d --- /dev/null +++ b/scripts/preprocess/preprocess_flux_rlpt_embeddings.sh @@ -0,0 +1,15 @@ +GPU_NUM=8 # 2,4,8 +MODEL_PATH="./ckpt/flux" +OUTPUT_DIR="data/laion_rl_embeddings" # Updated for LAION dataset +PROMPT_DIR="/data2/dataset/laion-220k/short_captions.txt" # Path to LAION captions + +echo "[INFO] Processing LAION-220k dataset captions..." +echo "[INFO] Model path: $MODEL_PATH" +echo "[INFO] Output directory: $OUTPUT_DIR" +echo "[INFO] Prompt file: $PROMPT_DIR" + +torchrun --nproc_per_node=$GPU_NUM --master_port 19002 \ + fastvideo/data_preprocess/preprocess_flux_embedding_rlpt.py \ + --model_path $MODEL_PATH \ + --output_dir $OUTPUT_DIR \ + --prompt_dir $PROMPT_DIR \ No newline at end of file diff --git a/scripts/preprocess/preprocess_flux_rlpt_embeddings_old.sh b/scripts/preprocess/preprocess_flux_rlpt_embeddings_old.sh new file mode 100644 index 0000000000000000000000000000000000000000..f496fa193873ad6925b48e1ae33954aeb98d8d56 --- /dev/null +++ b/scripts/preprocess/preprocess_flux_rlpt_embeddings_old.sh @@ -0,0 +1,9 @@ +GPU_NUM=4 # 2,4,8 +MODEL_PATH="./ckpt/flux" +OUTPUT_DIR="/mnt/dolphinfs/ssd_pool/docker/user/hadoop-videogen-hl/hadoop-camera3d/zhangshengjun/checkpoints/G2RPO/rl_embeddings" + +torchrun --nproc_per_node=$GPU_NUM --master_port 19002 \ + fastvideo/data_preprocess/preprocess_flux_embedding.py \ + --model_path $MODEL_PATH \ + --output_dir $OUTPUT_DIR \ + --prompt_dir "./prompts.txt" \ No newline at end of file diff --git a/scripts/preprocess/preprocess_qwen_image_rl_embeddings.sh b/scripts/preprocess/preprocess_qwen_image_rl_embeddings.sh new file mode 100644 index 0000000000000000000000000000000000000000..831dce273f901aad7c1083e57360c538c20806b7 --- /dev/null +++ b/scripts/preprocess/preprocess_qwen_image_rl_embeddings.sh @@ -0,0 +1,19 @@ + + +# pip install diffusers==0.35.0 peft==0.17.0 transformers==4.56.0 + +# GPU 6 is faulty +export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5" +GPU_NUM=6 # 2,4,8 +MODEL_PATH="./data/QwenImage" +OUTPUT_DIR="./data/qwenimage_rl_embeddings" + +# Change to source_code directory if not already there +cd "$(dirname "$0")/../.." + +torchrun --nproc_per_node=$GPU_NUM --master_port 19002 \ + fastvideo/data_preprocess/preprocess_qwenimage_embedding.py \ + --model_path $MODEL_PATH \ + --output_dir $OUTPUT_DIR \ + --prompt_dir "./assets/prompts.txt" + diff --git a/wandb/run-20260124_110332-48ji4pg6/files/config.yaml b/wandb/run-20260124_110332-48ji4pg6/files/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..66ae1c68fc61373f35586333ec00698fba85e1bd --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/files/config.yaml @@ -0,0 +1,87 @@ +_wandb: + value: + cli_version: 0.18.5 + m: [] + python_version: 3.10.19 + t: + "1": + - 1 + - 11 + - 41 + - 49 + - 55 + - 71 + - 83 + - 98 + "2": + - 1 + - 11 + - 41 + - 49 + - 55 + - 63 + - 71 + - 83 + - 98 + "3": + - 13 + - 23 + - 55 + "4": 3.10.19 + "5": 0.18.5 + "6": 4.46.1 + "8": + - 5 + "12": 0.18.5 + "13": linux-x86_64 +allow_tf32: + value: true +logdir: + value: logs +mixed_precision: + value: bf16 +num_checkpoint_limit: + value: 5 +num_epochs: + value: 300 +pretrained: + value: + model: ./data/StableDiffusion + revision: main +prompt_fn: + value: imagenet_animals +resume_from: + value: "" +reward_fn: + value: hpsv2 +run_name: + value: 2026.01.24_11.03.30 +sample: + value: + batch_size: 1 + eta: 1 + guidance_scale: 5 + num_batches_per_epoch: 2 + num_steps: 50 +save_freq: + value: 20 +seed: + value: 42 +train: + value: + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1e-08 + adam_weight_decay: 0.0001 + adv_clip_max: 5 + batch_size: 1 + cfg: true + clip_range: 0.0001 + gradient_accumulation_steps: 1 + learning_rate: 1e-05 + max_grad_norm: 1 + num_inner_epochs: 1 + timestep_fraction: 1 + use_8bit_adam: false +use_lora: + value: false diff --git a/wandb/run-20260124_110332-48ji4pg6/files/output.log b/wandb/run-20260124_110332-48ji4pg6/files/output.log new file mode 100644 index 0000000000000000000000000000000000000000..11befa0202d7fb344172a5394f911cff8336c027 --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/files/output.log @@ -0,0 +1,82 @@ +I0124 11:03:33.182662 129821359810368 train_g2rpo_sd_merge.py:478] +allow_tf32: true +logdir: logs +mixed_precision: bf16 +num_checkpoint_limit: 5 +num_epochs: 300 +pretrained: + model: ./data/StableDiffusion + revision: main +prompt_fn: imagenet_animals +prompt_fn_kwargs: {} +resume_from: '' +reward_fn: hpsv2 +run_name: 2026.01.24_11.03.30 +sample: + batch_size: 1 + eta: 1.0 + guidance_scale: 5.0 + num_batches_per_epoch: 2 + num_steps: 50 +save_freq: 20 +seed: 42 +train: + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_epsilon: 1.0e-08 + adam_weight_decay: 0.0001 + adv_clip_max: 5 + batch_size: 1 + cfg: true + clip_range: 0.0001 + gradient_accumulation_steps: 1 + learning_rate: 1.0e-05 + max_grad_norm: 1.0 + num_inner_epochs: 1 + timestep_fraction: 1.0 + use_8bit_adam: false +use_lora: false + +Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:02<00:00, 2.60it/s] +/home/zsj/anaconda3/envs/g2rpo/lib/python3.10/site-packages/timm/models/layers/__init__.py:48: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers + warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.layers", FutureWarning) +I0124 11:03:36.525150 129821359810368 factory.py:159] Loaded ViT-H-14 model config. +I0124 11:03:41.666221 129821359810368 factory.py:207] Loading pretrained ViT-H-14 weights (./data/hps/open_clip_pytorch_model.bin). +I0124 11:03:46.713310 129821359810368 train_g2rpo_sd_merge.py:670] ***** Running E-GRPO (G2RPO) Training for Stable Diffusion ***** +I0124 11:03:46.714378 129821359810368 train_g2rpo_sd_merge.py:671] Num Epochs = 300 +I0124 11:03:46.714534 129821359810368 train_g2rpo_sd_merge.py:672] Num generations per prompt = 4 +I0124 11:03:46.714650 129821359810368 train_g2rpo_sd_merge.py:673] Eta step list = [0, 1, 2, 3, 4, 5, 6, 7] +I0124 11:03:46.714740 129821359810368 train_g2rpo_sd_merge.py:674] Eta step merge list = [1, 1, 1, 2, 2, 2, 3, 3] +I0124 11:03:46.714825 129821359810368 train_g2rpo_sd_merge.py:675] Granular list = [1] +Traceback (most recent call last): + File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 1001, in + app.run(main) + File "/home/zsj/anaconda3/envs/g2rpo/lib/python3.10/site-packages/absl/app.py", line 316, in run + _run_main(main, args) + File "/home/zsj/anaconda3/envs/g2rpo/lib/python3.10/site-packages/absl/app.py", line 261, in _run_main + sys.exit(main(argv)) + File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 772, in main + eval_latents, anchor_latents = run_anchor_sample_step( + File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 229, in run_anchor_sample_step + prev_sample, _, _ = ddim_step_with_logprob_merge( + File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 165, in ddim_step_with_logprob_merge + std_dev_t = _left_broadcast(std_dev_t.view(-1), sample.shape).to(sample.device) + File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 98, in _left_broadcast + return t.reshape(t.shape + (1,) * (len(shape) - t.ndim)).broadcast_to(shape) +RuntimeError: The expanded size of the tensor (1) must match the existing size (16384) at non-singleton dimension 0. Target sizes: [1, 4, 64, 64]. Tensor sizes: [16384, 1, 1, 1] +[rank0]: Traceback (most recent call last): +[rank0]: File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 1001, in +[rank0]: app.run(main) +[rank0]: File "/home/zsj/anaconda3/envs/g2rpo/lib/python3.10/site-packages/absl/app.py", line 316, in run +[rank0]: _run_main(main, args) +[rank0]: File "/home/zsj/anaconda3/envs/g2rpo/lib/python3.10/site-packages/absl/app.py", line 261, in _run_main +[rank0]: sys.exit(main(argv)) +[rank0]: File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 772, in main +[rank0]: eval_latents, anchor_latents = run_anchor_sample_step( +[rank0]: File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 229, in run_anchor_sample_step +[rank0]: prev_sample, _, _ = ddim_step_with_logprob_merge( +[rank0]: File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 165, in ddim_step_with_logprob_merge +[rank0]: std_dev_t = _left_broadcast(std_dev_t.view(-1), sample.shape).to(sample.device) +[rank0]: File "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", line 98, in _left_broadcast +[rank0]: return t.reshape(t.shape + (1,) * (len(shape) - t.ndim)).broadcast_to(shape) +[rank0]: RuntimeError: The expanded size of the tensor (1) must match the existing size (16384) at non-singleton dimension 0. Target sizes: [1, 4, 64, 64]. Tensor sizes: [16384, 1, 1, 1] diff --git a/wandb/run-20260124_110332-48ji4pg6/files/requirements.txt b/wandb/run-20260124_110332-48ji4pg6/files/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..ee5c7ffa6079b296e15f3c9ff9edceed1bfe0802 --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/files/requirements.txt @@ -0,0 +1,189 @@ +scipy==1.13.0 +regex==2024.9.11 +sentencepiece==0.2.0 +six==1.16.0 +anyio==4.11.0 +nvidia-cuda-nvrtc-cu12==12.6.77 +scikit-video==1.1.11 +platformdirs==4.5.0 +mypy==1.11.1 +ruff==0.6.5 +charset-normalizer==3.4.4 +torch==2.9.0+cu126 +av==13.1.0 +pillow==10.2.0 +gpustat==1.1.1 +torchvision==0.24.0+cu126 +multidict==6.7.0 +torchmetrics==1.5.1 +aiohttp==3.13.1 +transformers==4.46.1 +decord==0.6.0 +wcwidth==0.2.14 +sphinx-lint==1.0.0 +nvidia-cuda-runtime-cu12==12.6.77 +pytz==2025.2 +codespell==2.3.0 +hpsv2==1.2.0 +mypy_extensions==1.1.0 +numpy==1.26.3 +omegaconf==2.3.0 +Markdown==3.9 +tzdata==2025.2 +pandas==2.2.3 +pytorch-lightning==2.4.0 +aiosignal==1.4.0 +aiohappyeyeballs==2.6.1 +python-dateutil==2.9.0.post0 +seaborn==0.13.2 +beautifulsoup4==4.12.3 +isort==5.13.2 +httpx==0.28.1 +certifi==2025.10.5 +ml_collections==1.1.0 +nvidia-cudnn-cu12==9.10.2.21 +hf-xet==1.2.0 +requests==2.31.0 +inflect==6.0.4 +iniconfig==2.1.0 +braceexpand==0.1.7 +h5py==3.12.1 +wandb==0.18.5 +protobuf==3.20.3 +ninja==1.13.0 +kiwisolver==1.4.9 +networkx==3.3 +packaging==25.0 +fvcore==0.1.5.post20221221 +pyparsing==3.2.5 +starlette==0.41.3 +frozenlist==1.8.0 +docker-pycreds==0.4.0 +Werkzeug==3.1.3 +MarkupSafe==2.1.5 +einops==0.8.0 +sentry-sdk==2.42.0 +PyYAML==6.0.1 +nvidia-nccl-cu12==2.27.5 +datasets==4.3.0 +polib==1.2.0 +safetensors==0.6.2 +async-timeout==5.0.1 +setproctitle==1.3.7 +clint==0.5.1 +matplotlib==3.9.2 +propcache==0.4.1 +termcolor==3.1.0 +antlr4-python3-runtime==4.9.3 +cycler==0.12.1 +fastvideo==1.2.0 +toml==0.10.2 +xxhash==3.6.0 +wheel==0.44.0 +albumentations==1.4.20 +fastapi==0.115.3 +nvidia-cufft-cu12==11.3.0.4 +yarl==1.22.0 +psutil==7.1.0 +tensorboard-data-server==0.7.2 +pydantic==2.9.2 +nvidia-nvtx-cu12==12.6.77 +portalocker==3.2.0 +triton==3.5.0 +annotated-types==0.7.0 +proglog==0.1.12 +nvidia-cusparselt-cu12==0.7.1 +yapf==0.32.0 +Jinja2==3.1.6 +types-requests==2.32.4.20250913 +lightning-utilities==0.15.2 +grpcio==1.75.1 +uvicorn==0.32.0 +typing_extensions==4.15.0 +nvidia-nvjitlink-cu12==12.6.85 +watch==0.2.7 +moviepy==1.0.3 +timm==1.0.11 +pytest-split==0.8.0 +gdown==5.2.0 +types-setuptools==80.9.0.20250822 +nvidia-cusolver-cu12==11.7.1.2 +types-PyYAML==6.0.12.20250915 +pip==25.2 +qwen-vl-utils==0.0.14 +soupsieve==2.8 +zipp==3.23.0 +flash_attn==2.8.3 +yacs==0.1.8 +diffusers==0.32.0 +pluggy==1.6.0 +opencv-python-headless==4.11.0.86 +mpmath==1.3.0 +test_tube==0.7.5 +stringzilla==4.2.1 +fonttools==4.60.1 +nvidia-ml-py==13.580.82 +parameterized==0.9.0 +loguru==0.7.3 +tabulate==0.9.0 +idna==3.6 +iopath==0.1.10 +decorator==4.4.2 +nvidia-cufile-cu12==1.11.1.6 +threadpoolctl==3.6.0 +pyarrow==21.0.0 +httpcore==1.0.9 +hydra-core==1.3.2 +multiprocess==0.70.16 +contourpy==1.3.2 +clip==1.0 +tqdm==4.66.5 +open_clip_torch==3.2.0 +accelerate==1.0.1 +gitdb==4.0.12 +importlib_metadata==8.7.0 +nvidia-cublas-cu12==12.6.4.1 +h11==0.16.0 +filelock==3.19.1 +liger_kernel==0.4.1 +click==8.3.0 +urllib3==2.2.0 +imageio-ffmpeg==0.5.1 +setuptools==80.9.0 +joblib==1.5.2 +tensorboard==2.20.0 +attrs==25.4.0 +future==1.0.0 +albucore==0.0.19 +fsspec==2025.9.0 +sympy==1.14.0 +eval_type_backport==0.2.2 +pydantic_core==2.23.4 +sniffio==1.3.1 +nvidia-nvshmem-cu12==3.3.20 +exceptiongroup==1.3.0 +smmap==5.0.2 +tomli==2.0.2 +ftfy==6.3.0 +dill==0.4.0 +pytest==7.2.0 +PySocks==1.7.1 +nvidia-curand-cu12==10.3.7.77 +tokenizers==0.20.1 +args==0.1.0 +fairscale==0.4.13 +peft==0.13.2 +webdataset==1.0.2 +huggingface-hub==0.26.1 +GitPython==3.1.45 +pytorchvideo==0.1.5 +scikit-learn==1.5.2 +bitsandbytes==0.48.1 +nvidia-cusparse-cu12==12.5.4.2 +nvidia-cuda-cupti-cu12==12.6.80 +imageio==2.36.0 +pydub==0.25.1 +image-reward==1.5 +absl-py==2.3.1 +blessed==1.22.0 +torchdiffeq==0.2.4 diff --git a/wandb/run-20260124_110332-48ji4pg6/files/wandb-metadata.json b/wandb/run-20260124_110332-48ji4pg6/files/wandb-metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..c098f2ddbd8e55daa21441e7297befc9fb7f653f --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/files/wandb-metadata.json @@ -0,0 +1,96 @@ +{ + "os": "Linux-6.8.0-85-generic-x86_64-with-glibc2.35", + "python": "3.10.19", + "startedAt": "2026-01-24T03:03:32.116013Z", + "args": [ + "--config", + "fastvideo/config_sd/base.py", + "--eta_step_list", + "0,1,2,3,4,5,6,7", + "--eta_step_merge_list", + "1,1,1,2,2,2,3,3", + "--granular_list", + "1", + "--num_generations", + "4", + "--eta", + "1.0", + "--init_same_noise" + ], + "program": "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py", + "codePath": "fastvideo/train_g2rpo_sd_merge.py", + "email": "zhangemail1428@163.com", + "root": "/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code", + "host": "abc", + "username": "zsj", + "executable": "/home/zsj/anaconda3/envs/g2rpo/bin/python", + "codePathLocal": "fastvideo/train_g2rpo_sd_merge.py", + "cpu_count": 48, + "cpu_count_logical": 96, + "gpu": "NVIDIA RTX 5880 Ada Generation", + "gpu_count": 8, + "disk": { + "/": { + "total": "1006773899264", + "used": "812118953984" + } + }, + "memory": { + "total": "540697260032" + }, + "cpu": { + "count": 48, + "countLogical": 96 + }, + "gpu_nvidia": [ + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + }, + { + "name": "NVIDIA RTX 5880 Ada Generation", + "memoryTotal": "51527024640", + "cudaCores": 14080, + "architecture": "Ada" + } + ], + "cudaVersion": "12.9" +} \ No newline at end of file diff --git a/wandb/run-20260124_110332-48ji4pg6/files/wandb-summary.json b/wandb/run-20260124_110332-48ji4pg6/files/wandb-summary.json new file mode 100644 index 0000000000000000000000000000000000000000..1476d70fd36aa7b2a81bf4776ad653af3dd34436 --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/files/wandb-summary.json @@ -0,0 +1 @@ +{"_wandb":{"runtime":14}} \ No newline at end of file diff --git a/wandb/run-20260124_110332-48ji4pg6/logs/debug-core.log b/wandb/run-20260124_110332-48ji4pg6/logs/debug-core.log new file mode 100644 index 0000000000000000000000000000000000000000..b99b795dfe8da318cdebc496b350fc20d254f08f --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/logs/debug-core.log @@ -0,0 +1,12 @@ +{"time":"2026-01-24T11:03:30.985219213+08:00","level":"INFO","msg":"started logging, with flags","port-filename":"/tmp/tmp7mt3x3cn/port-696886.txt","pid":696886,"debug":false,"disable-analytics":false} +{"time":"2026-01-24T11:03:30.985248026+08:00","level":"INFO","msg":"FeatureState","shutdownOnParentExitEnabled":false} +{"time":"2026-01-24T11:03:30.985986152+08:00","level":"INFO","msg":"server is running","addr":{"IP":"127.0.0.1","Port":33449,"Zone":""}} +{"time":"2026-01-24T11:03:30.986068877+08:00","level":"INFO","msg":"Will exit if parent process dies.","ppid":696886} +{"time":"2026-01-24T11:03:31.176029507+08:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"127.0.0.1:57966"} +{"time":"2026-01-24T11:03:32.120879812+08:00","level":"INFO","msg":"handleInformInit: received","streamId":"48ji4pg6","id":"127.0.0.1:57966"} +{"time":"2026-01-24T11:03:32.240634919+08:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"48ji4pg6","id":"127.0.0.1:57966"} +{"time":"2026-01-24T11:03:47.064995589+08:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"127.0.0.1:57966"} +{"time":"2026-01-24T11:03:47.065112628+08:00","level":"INFO","msg":"connection: Close: initiating connection closure","id":"127.0.0.1:57966"} +{"time":"2026-01-24T11:03:47.065143097+08:00","level":"INFO","msg":"server is shutting down"} +{"time":"2026-01-24T11:03:47.065180994+08:00","level":"INFO","msg":"connection: Close: connection successfully closed","id":"127.0.0.1:57966"} +{"time":"2026-01-24T11:03:49.250437548+08:00","level":"INFO","msg":"Parent process exited, terminating service process."} diff --git a/wandb/run-20260124_110332-48ji4pg6/logs/debug-internal.log b/wandb/run-20260124_110332-48ji4pg6/logs/debug-internal.log new file mode 100644 index 0000000000000000000000000000000000000000..b0dd7ffceaa4858fe4d4a2a8321e740c292f8e0e --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/logs/debug-internal.log @@ -0,0 +1,15 @@ +{"time":"2026-01-24T11:03:32.121096288+08:00","level":"INFO","msg":"using version","core version":"0.18.5"} +{"time":"2026-01-24T11:03:32.121116393+08:00","level":"INFO","msg":"created symlink","path":"/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/wandb/run-20260124_110332-48ji4pg6/logs/debug-core.log"} +{"time":"2026-01-24T11:03:32.240548096+08:00","level":"INFO","msg":"created new stream","id":"48ji4pg6"} +{"time":"2026-01-24T11:03:32.240626186+08:00","level":"INFO","msg":"stream: started","id":"48ji4pg6"} +{"time":"2026-01-24T11:03:32.240790819+08:00","level":"INFO","msg":"sender: started","stream_id":"48ji4pg6"} +{"time":"2026-01-24T11:03:32.24072294+08:00","level":"INFO","msg":"writer: Do: started","stream_id":{"value":"48ji4pg6"}} +{"time":"2026-01-24T11:03:32.24073378+08:00","level":"INFO","msg":"handler: started","stream_id":{"value":"48ji4pg6"}} +{"time":"2026-01-24T11:03:33.015727057+08:00","level":"INFO","msg":"Starting system monitor"} +{"time":"2026-01-24T11:03:47.065133577+08:00","level":"INFO","msg":"stream: closing","id":"48ji4pg6"} +{"time":"2026-01-24T11:03:47.065209431+08:00","level":"INFO","msg":"Stopping system monitor"} +{"time":"2026-01-24T11:03:47.066463879+08:00","level":"INFO","msg":"Stopped system monitor"} +{"time":"2026-01-24T11:03:47.394224304+08:00","level":"WARN","msg":"No job ingredients found, not creating job artifact"} +{"time":"2026-01-24T11:03:47.394263573+08:00","level":"WARN","msg":"No source type found, not creating job artifact"} +{"time":"2026-01-24T11:03:47.394275391+08:00","level":"INFO","msg":"sender: sendDefer: no job artifact to save"} +{"time":"2026-01-24T11:03:48.337170595+08:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"} diff --git a/wandb/run-20260124_110332-48ji4pg6/logs/debug.log b/wandb/run-20260124_110332-48ji4pg6/logs/debug.log new file mode 100644 index 0000000000000000000000000000000000000000..0d402a38a90d957c648bdc0bd01dd5d261a889a3 --- /dev/null +++ b/wandb/run-20260124_110332-48ji4pg6/logs/debug.log @@ -0,0 +1,27 @@ +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Current SDK version is 0.18.5 +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Configure stats pid to 696886 +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Loading settings from /home/zsj/.config/wandb/settings +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Loading settings from /data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/wandb/settings +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Loading settings from environment variables: {} +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Applying setup settings: {'mode': None, '_disable_service': None} +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Inferring run settings from compute environment: {'program_relpath': 'fastvideo/train_g2rpo_sd_merge.py', 'program_abspath': '/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py', 'program': '/data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/fastvideo/train_g2rpo_sd_merge.py'} +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_setup.py:_flush():79] Applying login settings: {} +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_init.py:_log_setup():534] Logging user logs to /data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/wandb/run-20260124_110332-48ji4pg6/logs/debug.log +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_init.py:_log_setup():535] Logging internal logs to /data1/zsj/SceneDPO/Rebuttal/E-GRPO/scoure_code/wandb/run-20260124_110332-48ji4pg6/logs/debug-internal.log +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_init.py:init():621] calling init triggers +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_init.py:init():628] wandb.init called with sweep_config: {} +config: {} +2026-01-24 11:03:32,112 INFO MainThread:696886 [wandb_init.py:init():671] starting backend +2026-01-24 11:03:32,113 INFO MainThread:696886 [wandb_init.py:init():675] sending inform_init request +2026-01-24 11:03:32,115 INFO MainThread:696886 [backend.py:_multiprocessing_setup():104] multiprocessing start_methods=fork,spawn,forkserver, using: spawn +2026-01-24 11:03:32,115 INFO MainThread:696886 [wandb_init.py:init():688] backend started and connected +2026-01-24 11:03:32,118 INFO MainThread:696886 [wandb_init.py:init():783] updated telemetry +2026-01-24 11:03:32,119 INFO MainThread:696886 [wandb_init.py:init():816] communicating run to backend with 90.0 second timeout +2026-01-24 11:03:33,006 INFO MainThread:696886 [wandb_init.py:init():867] starting run threads in backend +2026-01-24 11:03:33,178 INFO MainThread:696886 [wandb_run.py:_console_start():2463] atexit reg +2026-01-24 11:03:33,178 INFO MainThread:696886 [wandb_run.py:_redirect():2311] redirect: wrap_raw +2026-01-24 11:03:33,179 INFO MainThread:696886 [wandb_run.py:_redirect():2376] Wrapping output streams. +2026-01-24 11:03:33,179 INFO MainThread:696886 [wandb_run.py:_redirect():2401] Redirects installed. +2026-01-24 11:03:33,180 INFO MainThread:696886 [wandb_init.py:init():911] run started, returning control to user process +2026-01-24 11:03:33,180 INFO MainThread:696886 [wandb_run.py:_config_callback():1390] config_cb None None {'allow_tf32': True, 'logdir': 'logs', 'mixed_precision': 'bf16', 'num_checkpoint_limit': 5, 'num_epochs': 300, 'pretrained': {'model': './data/StableDiffusion', 'revision': 'main'}, 'prompt_fn': 'imagenet_animals', 'prompt_fn_kwargs': {}, 'resume_from': '', 'reward_fn': 'hpsv2', 'run_name': '2026.01.24_11.03.30', 'sample': {'batch_size': 1, 'eta': 1.0, 'guidance_scale': 5.0, 'num_batches_per_epoch': 2, 'num_steps': 50}, 'save_freq': 20, 'seed': 42, 'train': {'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_epsilon': 1e-08, 'adam_weight_decay': 0.0001, 'adv_clip_max': 5, 'batch_size': 1, 'cfg': True, 'clip_range': 0.0001, 'gradient_accumulation_steps': 1, 'learning_rate': 1e-05, 'max_grad_norm': 1.0, 'num_inner_epochs': 1, 'timestep_fraction': 1.0, 'use_8bit_adam': False}, 'use_lora': False} +2026-01-24 11:03:47,065 WARNING MsgRouterThr:696886 [router.py:message_loop():77] message_loop has been closed