diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d75870d --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +output/ +*__pycache__/ +samples*/ +runs/ +checkpoints/ +master_ip +logs/ +*.DS_Store +.idea +*.pt +tools/ \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..f4e09a9 --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/configs/hy_vae_config.json b/configs/hy_vae_config.json new file mode 100644 index 0000000..73d4cc0 --- /dev/null +++ b/configs/hy_vae_config.json @@ -0,0 +1,34 @@ +{ + "_class_name": "AutoencoderKLCausal3D", + "_diffusers_version": "0.4.2", + "act_fn": "silu", + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlockCausal3D", + "DownEncoderBlockCausal3D", + "DownEncoderBlockCausal3D", + "DownEncoderBlockCausal3D" + ], + "in_channels": 3, + "latent_channels": 16, + "layers_per_block": 2, + "norm_num_groups": 32, + "out_channels": 3, + "sample_size": 256, + "sample_tsize": 64, + "up_block_types": [ + "UpDecoderBlockCausal3D", + "UpDecoderBlockCausal3D", + "UpDecoderBlockCausal3D", + "UpDecoderBlockCausal3D" + ], + "scaling_factor": 0.476986, + "time_compression_ratio": 4, + "mid_block_add_attention": true, + "mid_block_causal_attn": true +} diff --git a/examples/hyvideo_t2v_example_01.json b/examples/hyvideo_t2v_example_01.json new file mode 100644 index 0000000..6752fef --- /dev/null +++ b/examples/hyvideo_t2v_example_01.json @@ -0,0 +1,391 @@ +{ + "last_node_id": 16, + "last_link_id": 17, + "nodes": [ + { + "id": 5, + "type": "HyVideoDecode", + "pos": [ + 1225, + 313 + ], + "size": [ + 315, + 78 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "VAE", + "link": 6 + }, + { + "name": "samples", + "type": "LATENT", + "link": 4 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 13 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoDecode" + }, + "widgets_values": [ + true + ] + }, + { + "id": 7, + "type": "HyVideoVAELoader", + "pos": [ + 1206, + 71 + ], + "size": [ + 379.166748046875, + 82 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "compile_args", + "type": "COMPILEARGS", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "vae", + "type": "VAE", + "links": [ + 6 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoVAELoader" + }, + "widgets_values": [ + "hyvid\\hunyuan_video_vae_bf16.safetensors", + "bf16" + ] + }, + { + "id": 4, + "type": "HyVideoTextEncode", + "pos": [ + 616, + 428 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "text_encoders", + "type": "HYVIDTEXTENCODER", + "link": 17 + } + ], + "outputs": [ + { + "name": "hyvid_embeds", + "type": "HYVIDEMBEDS", + "links": [ + 12 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoTextEncode" + }, + "widgets_values": [ + "A cat walks on the grass, realistic style", + "bad quality video, slow motion", + true + ] + }, + { + "id": 3, + "type": "HyVideoSampler", + "pos": [ + 768, + 106 + ], + "size": [ + 315, + 246 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "HYVIDEOMODEL", + "link": 2 + }, + { + "name": "hyvid_embeds", + "type": "HYVIDEMBEDS", + "link": 12 + } + ], + "outputs": [ + { + "name": "samples", + "type": "LATENT", + "links": [ + 4 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoSampler" + }, + "widgets_values": [ + 512, + 320, + 49, + 30, + 4, + 1008437192147815, + "fixed", + true + ] + }, + { + "id": 11, + "type": "VHS_VideoCombine", + "pos": [ + 1634, + 75 + ], + "size": [ + 632.9049682617188, + 707.0656127929688 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 13 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "HunyuanVideo", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "HunyuanVideo_00008.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 16 + }, + "muted": false + } + } + }, + { + "id": 1, + "type": "HyVideoModelLoader", + "pos": [ + 124, + 105 + ], + "size": [ + 509.7506103515625, + 178 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "compile_args", + "type": "COMPILEARGS", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "model", + "type": "HYVIDEOMODEL", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoModelLoader" + }, + "widgets_values": [ + "hyvideo\\hunyuan_video_720_fp8_e4m3fn.safetensors", + "bf16", + "fp8_e4m3fn", + "offload_device", + false, + "sageattn_varlen" + ] + }, + { + "id": 16, + "type": "DownloadAndLoadHyVideoTextEncoder", + "pos": [ + 133, + 407 + ], + "size": [ + 441, + 106 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "hyvid_text_encoder", + "type": "HYVIDTEXTENCODER", + "links": [ + 17 + ] + } + ], + "properties": { + "Node name for S&R": "DownloadAndLoadHyVideoTextEncoder" + }, + "widgets_values": [ + "Kijai/llava-llama-3-8b-text-encoder-tokenizer", + "openai/clip-vit-large-patch14", + "bf16" + ] + } + ], + "links": [ + [ + 2, + 1, + 0, + 3, + 0, + "HYVIDEOMODEL" + ], + [ + 4, + 3, + 0, + 5, + 1, + "LATENT" + ], + [ + 6, + 7, + 0, + 5, + 0, + "VAE" + ], + [ + 12, + 4, + 0, + 3, + 1, + "HYVIDEMBEDS" + ], + [ + 13, + 5, + 0, + 11, + 0, + "IMAGE" + ], + [ + 17, + 16, + 0, + 4, + 0, + "HYVIDTEXTENCODER" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": [ + 105.44467975365966, + 249.98241175433185 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/hyvideo/__init__.py b/hyvideo/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hyvideo/config.py b/hyvideo/config.py new file mode 100644 index 0000000..3f34c33 --- /dev/null +++ b/hyvideo/config.py @@ -0,0 +1,371 @@ +import argparse +from .constants import * +import re +from .modules.models import HUNYUAN_VIDEO_CONFIG + + +def parse_args(namespace=None): + parser = argparse.ArgumentParser(description="HunyuanVideo inference script") + + parser = add_network_args(parser) + parser = add_extra_models_args(parser) + parser = add_denoise_schedule_args(parser) + parser = add_inference_args(parser) + + args = parser.parse_args(namespace=namespace) + args = sanity_check_args(args) + + return args + + +def add_network_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="HunyuanVideo network args") + + # Main model + group.add_argument( + "--model", + type=str, + choices=list(HUNYUAN_VIDEO_CONFIG.keys()), + default="HYVideo-T/2-cfgdistill", + ) + group.add_argument( + "--latent-channels", + type=str, + default=16, + help="Number of latent channels of DiT. If None, it will be determined by `vae`. If provided, " + "it still needs to match the latent channels of the VAE model.", + ) + group.add_argument( + "--precision", + type=str, + default="bf16", + choices=PRECISIONS, + help="Precision mode. Options: fp32, fp16, bf16. Applied to the backbone model and optimizer.", + ) + + # RoPE + group.add_argument( + "--rope-theta", type=int, default=256, help="Theta used in RoPE." + ) + return parser + + +def add_extra_models_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group( + title="Extra models args, including vae, text encoders and tokenizers)" + ) + + # - VAE + group.add_argument( + "--vae", + type=str, + default="884-16c-hy", + choices=list(VAE_PATH), + help="Name of the VAE model.", + ) + group.add_argument( + "--vae-precision", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the VAE model.", + ) + group.add_argument( + "--vae-tiling", + action="store_true", + help="Enable tiling for the VAE model to save GPU memory.", + ) + group.set_defaults(vae_tiling=True) + + group.add_argument( + "--text-encoder", + type=str, + default="llm", + choices=list(TEXT_ENCODER_PATH), + help="Name of the text encoder model.", + ) + group.add_argument( + "--text-encoder-precision", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the text encoder model.", + ) + group.add_argument( + "--text-states-dim", + type=int, + default=4096, + help="Dimension of the text encoder hidden states.", + ) + group.add_argument( + "--text-len", type=int, default=256, help="Maximum length of the text input." + ) + group.add_argument( + "--tokenizer", + type=str, + default="llm", + choices=list(TOKENIZER_PATH), + help="Name of the tokenizer model.", + ) + group.add_argument( + "--prompt-template", + type=str, + default="dit-llm-encode", + choices=PROMPT_TEMPLATE, + help="Image prompt template for the decoder-only text encoder model.", + ) + group.add_argument( + "--prompt-template-video", + type=str, + default="dit-llm-encode-video", + choices=PROMPT_TEMPLATE, + help="Video prompt template for the decoder-only text encoder model.", + ) + group.add_argument( + "--hidden-state-skip-layer", + type=int, + default=2, + help="Skip layer for hidden states.", + ) + group.add_argument( + "--apply-final-norm", + action="store_true", + help="Apply final normalization to the used text encoder hidden states.", + ) + + # - CLIP + group.add_argument( + "--text-encoder-2", + type=str, + default="clipL", + choices=list(TEXT_ENCODER_PATH), + help="Name of the second text encoder model.", + ) + group.add_argument( + "--text-encoder-precision-2", + type=str, + default="fp16", + choices=PRECISIONS, + help="Precision mode for the second text encoder model.", + ) + group.add_argument( + "--text-states-dim-2", + type=int, + default=768, + help="Dimension of the second text encoder hidden states.", + ) + group.add_argument( + "--tokenizer-2", + type=str, + default="clipL", + choices=list(TOKENIZER_PATH), + help="Name of the second tokenizer model.", + ) + group.add_argument( + "--text-len-2", + type=int, + default=77, + help="Maximum length of the second text input.", + ) + + return parser + + +def add_denoise_schedule_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Denoise schedule args") + + group.add_argument( + "--denoise-type", + type=str, + default="flow", + help="Denoise type for noised inputs.", + ) + + # Flow Matching + group.add_argument( + "--flow-shift", + type=float, + default=9.0, + help="Shift factor for flow matching schedulers.", + ) + group.add_argument( + "--flow-reverse", + action="store_true", + help="If reverse, learning/sampling from t=1 -> t=0.", + ) + group.add_argument( + "--flow-solver", + type=str, + default="euler", + help="Solver for flow matching.", + ) + group.add_argument( + "--use-linear-quadratic-schedule", + action="store_true", + help="Use linear quadratic schedule for flow matching." + "Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)", + ) + group.add_argument( + "--linear-schedule-end", + type=int, + default=25, + help="End step for linear quadratic schedule for flow matching.", + ) + + return parser + + +def add_inference_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Inference args") + + # ======================== Model loads ======================== + group.add_argument( + "--model-base", + type=str, + default="ckpts", + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--dit-weight", + type=str, + default="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt", + help="Path to the HunyuanVideo model. If None, search the model in the args.model_root." + "1. If it is a file, load the model directly." + "2. If it is a directory, search the model in the directory. Support two types of models: " + "1) named `pytorch_model_*.pt`" + "2) named `*_model_states.pt`, where * can be `mp_rank_00`.", + ) + group.add_argument( + "--model-resolution", + type=str, + default="540p", + choices=["540p", "720p"], + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--load-key", + type=str, + default="module", + help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.", + ) + group.add_argument( + "--use-cpu-offload", + action="store_true", + help="Use CPU offload for the model load.", + ) + + # ======================== Inference general setting ======================== + group.add_argument( + "--batch-size", + type=int, + default=1, + help="Batch size for inference and evaluation.", + ) + group.add_argument( + "--infer-steps", + type=int, + default=30, + help="Number of denoising steps for inference.", + ) + group.add_argument( + "--disable-autocast", + action="store_true", + help="Disable autocast for denoising loop and vae decoding in pipeline sampling.", + ) + group.add_argument( + "--save-path", + type=str, + default="./results", + help="Path to save the generated samples.", + ) + group.add_argument( + "--save-path-suffix", + type=str, + default="", + help="Suffix for the directory of saved samples.", + ) + group.add_argument( + "--name-suffix", + type=str, + default="", + help="Suffix for the names of saved samples.", + ) + group.add_argument( + "--num-videos", + type=int, + default=1, + help="Number of videos to generate for each prompt.", + ) + # ---sample size--- + group.add_argument( + "--video-size", + type=int, + nargs="+", + default=(720, 1280), + help="Video size for training. If a single value is provided, it will be used for both height " + "and width. If two values are provided, they will be used for height and width " + "respectively.", + ) + group.add_argument( + "--video-length", + type=int, + default=129, + help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1", + ) + # --- prompt --- + group.add_argument( + "--prompt", + type=str, + default=None, + help="Prompt for sampling during evaluation.", + ) + group.add_argument( + "--seed-type", + type=str, + default="auto", + choices=["file", "random", "fixed", "auto"], + help="Seed type for evaluation. If file, use the seed from the CSV file. If random, generate a " + "random seed. If fixed, use the fixed seed given by `--seed`. If auto, `csv` will use the " + "seed column if available, otherwise use the fixed `seed` value. `prompt` will use the " + "fixed `seed` value.", + ) + group.add_argument("--seed", type=int, default=0, help="Seed for evaluation.") + + # Classifier-Free Guidance + group.add_argument( + "--neg-prompt", type=str, default=None, help="Negative prompt for sampling." + ) + group.add_argument( + "--cfg-scale", type=float, default=1.0, help="Classifier free guidance scale." + ) + group.add_argument( + "--embedded-cfg-scale", + type=float, + default=6.0, + help="Embeded classifier free guidance scale.", + ) + + group.add_argument( + "--reproduce", + action="store_true", + help="Enable reproducibility by setting random seeds and deterministic algorithms.", + ) + + return parser + + +def sanity_check_args(args): + # VAE channels + vae_pattern = r"\d{2,3}-\d{1,2}c-\w+" + if not re.match(vae_pattern, args.vae): + raise ValueError( + f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'." + ) + vae_channels = int(args.vae.split("-")[1][:-1]) + if args.latent_channels is None: + args.latent_channels = vae_channels + if vae_channels != args.latent_channels: + raise ValueError( + f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels})." + ) + return args diff --git a/hyvideo/constants.py b/hyvideo/constants.py new file mode 100644 index 0000000..db4db71 --- /dev/null +++ b/hyvideo/constants.py @@ -0,0 +1,91 @@ +import os +import torch + +__all__ = [ + "C_SCALE", + "PROMPT_TEMPLATE", + "MODEL_BASE", + "PRECISIONS", + "NORMALIZATION_TYPE", + "ACTIVATION_TYPE", + "VAE_PATH", + "TEXT_ENCODER_PATH", + "TOKENIZER_PATH", + "TEXT_PROJECTION", + "DATA_TYPE", + "NEGATIVE_PROMPT", +] + +PRECISION_TO_TYPE = { + 'fp32': torch.float32, + 'fp16': torch.float16, + 'bf16': torch.bfloat16, + 'fp8_e4m3fn': torch.float8_e4m3fn, +} + +# =================== Constant Values ===================== +# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid +# overflow error when tensorboard logging values. +C_SCALE = 1_000_000_000_000_000 + +# When using decoder-only models, we must provide a prompt template to instruct the text encoder +# on how to generate the text. +# -------------------------------------------------------------------- +PROMPT_TEMPLATE_ENCODE = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" +) +PROMPT_TEMPLATE_ENCODE_VIDEO = ( + "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: " + "1. The main content and theme of the video." + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects." + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects." + "4. background environment, light, style and atmosphere." + "5. camera angles, movements, and transitions used in the video:<|eot_id|>" + "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>" +) + +NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion" + +PROMPT_TEMPLATE = { + "dit-llm-encode": { + "template": PROMPT_TEMPLATE_ENCODE, + "crop_start": 36, + }, + "dit-llm-encode-video": { + "template": PROMPT_TEMPLATE_ENCODE_VIDEO, + "crop_start": 95, + }, +} + +# ======================= Model ====================== +PRECISIONS = {"fp32", "fp16", "bf16"} +NORMALIZATION_TYPE = {"layer", "rms"} +ACTIVATION_TYPE = {"relu", "silu", "gelu", "gelu_tanh"} + +# =================== Model Path ===================== +MODEL_BASE = os.getenv("MODEL_BASE", "./ckpts") + +# =================== Data ======================= +DATA_TYPE = {"image", "video", "image_video"} + +# 3D VAE +VAE_PATH = {"884-16c-hy": f"{MODEL_BASE}/hunyuan-video-t2v-720p/vae"} + +# Text Encoder +TEXT_ENCODER_PATH = { + "clipL": f"{MODEL_BASE}/text_encoder_2", + "llm": f"{MODEL_BASE}/text_encoder", +} + +# Tokenizer +TOKENIZER_PATH = { + "clipL": f"{MODEL_BASE}/text_encoder_2", + "llm": f"{MODEL_BASE}/text_encoder", +} + +TEXT_PROJECTION = { + "linear", # Default, an nn.Linear() layer + "single_refiner", # Single TokenRefiner. Refer to LI-DiT +} diff --git a/hyvideo/diffusion/__init__.py b/hyvideo/diffusion/__init__.py new file mode 100644 index 0000000..2141aa3 --- /dev/null +++ b/hyvideo/diffusion/__init__.py @@ -0,0 +1,2 @@ +from .pipelines import HunyuanVideoPipeline +from .schedulers import FlowMatchDiscreteScheduler diff --git a/hyvideo/diffusion/pipelines/__init__.py b/hyvideo/diffusion/pipelines/__init__.py new file mode 100644 index 0000000..e44cb61 --- /dev/null +++ b/hyvideo/diffusion/pipelines/__init__.py @@ -0,0 +1 @@ +from .pipeline_hunyuan_video import HunyuanVideoPipeline diff --git a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py new file mode 100644 index 0000000..c783702 --- /dev/null +++ b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py @@ -0,0 +1,566 @@ +# 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 typing import Any, Callable, Dict, List, Optional, Union, Tuple +import torch + +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.configuration_utils import FrozenDict +from diffusers.image_processor import VaeImageProcessor + +from diffusers.schedulers import KarrasDiffusionSchedulers +from diffusers.utils import ( + deprecate, + logging, + replace_example_docstring +) +from diffusers.utils.torch_utils import randn_tensor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline + +from ...modules import HYVideoDiffusionTransformer +from comfy.utils import ProgressBar + +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 + +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: + 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, + transformer: HYVideoDiffusionTransformer, + scheduler: KarrasDiffusionSchedulers, + progress_bar_config: Dict[str, Any] = None, + base_dtype = torch.bfloat16, + ): + 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.base_dtype = base_dtype + # ========================================================================================== + + self.register_modules( + transformer=transformer, + scheduler=scheduler + ) + self.vae_scale_factor = 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + + 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 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, + height: int, + width: int, + video_length: int, + prompt_embed_dict: dict, + num_inference_steps: int = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: float = 7.5, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + denoise_strength: float = 1.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + + 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"], + freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None, + n_tokens: Optional[int] = None, + embedded_guidance_scale: Optional[float] = None, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + 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`. + 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`. + + 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. + 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 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 + + 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 + + batch_size = 1 + device = self._execution_device + + prompt_embeds = prompt_embed_dict["prompt_embeds"] + negative_prompt_embeds = prompt_embed_dict["negative_prompt_embeds"] + prompt_mask = prompt_embed_dict["attention_mask"] + negative_prompt_mask = prompt_embed_dict["negative_attention_mask"] + prompt_embeds_2 = prompt_embed_dict["prompt_embeds_2"] + negative_prompt_embeds_2 = prompt_embed_dict["negative_prompt_embeds_2"] + prompt_mask_2 = prompt_embed_dict["attention_mask_2"] + negative_prompt_mask_2 = prompt_embed_dict["negative_attention_mask_2"] + + # 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 + + + # 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, + ) + + # 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}, + ) + + # 7. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + # if is_progress_bar: + comfy_pbar = ProgressBar(num_inference_steps) + 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(self.base_dtype) + * 1000.0 + if embedded_guidance_scale is not None + else None + ) + + # predict the noise residual + with torch.autocast( + device_type="cuda", dtype=self.base_dtype, enabled=True + ): + noise_pred = self.transformer( # For an input image (129, 192, 336) (1, 256, 256) + latent_model_input, # [2, 16, 33, 24, 42] + t_expand, # [2] + text_states=prompt_embeds, # [2, 256, 4096] + text_mask=prompt_mask, # [2, 256] + text_states_2=prompt_embeds_2, # [2, 768] + freqs_cos=freqs_cis[0], # [seqlen, head_dim] + freqs_sin=freqs_cis[1], # [seqlen, head_dim] + guidance=guidance_expand, + return_dict=True, + )["x"] + + # 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() + comfy_pbar.update(1) + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + #latents = (latents / 2 + 0.5).clamp(0, 1).cpu() + + # Offload all models + self.maybe_free_model_hooks() + + return latents \ No newline at end of file diff --git a/hyvideo/diffusion/schedulers/__init__.py b/hyvideo/diffusion/schedulers/__init__.py new file mode 100644 index 0000000..14f2ba3 --- /dev/null +++ b/hyvideo/diffusion/schedulers/__init__.py @@ -0,0 +1 @@ +from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler diff --git a/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py b/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py new file mode 100644 index 0000000..c507ec4 --- /dev/null +++ b/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py @@ -0,0 +1,257 @@ +# 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 numpy as np +import torch + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput, logging +from diffusers.schedulers.scheduling_utils import SchedulerMixin + + +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/hyvideo/inference.py b/hyvideo/inference.py new file mode 100644 index 0000000..9b7a41b --- /dev/null +++ b/hyvideo/inference.py @@ -0,0 +1,494 @@ +import time +import random + +from pathlib import Path +from loguru import logger + +import torch +from hyvideo.constants import PROMPT_TEMPLATE, NEGATIVE_PROMPT, PRECISION_TO_TYPE +from hyvideo.vae import load_vae +from hyvideo.text_encoder import TextEncoder +from hyvideo.utils.data_utils import align_to +from hyvideo.modules.posemb_layers import get_nd_rotary_pos_embed +from hyvideo.diffusion.schedulers import FlowMatchDiscreteScheduler +from hyvideo.diffusion.pipelines import HunyuanVideoPipeline + +from hyvideo.modules.models import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG +from accelerate import init_empty_weights +from accelerate.utils import set_module_tensor_to_device +import safetensors.torch + +class Inference(object): + def __init__( + self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + use_cpu_offload=False, + device=None, + logger=None, + ): + self.vae = vae + self.vae_kwargs = vae_kwargs + + self.text_encoder = text_encoder + self.text_encoder_2 = text_encoder_2 + + self.model = model + self.pipeline = pipeline + self.use_cpu_offload = use_cpu_offload + + self.args = args + self.device = ( + device + if device is not None + else "cuda" + if torch.cuda.is_available() + else "cpu" + ) + self.logger = logger + + @classmethod + def from_pretrained(cls, pretrained_model_path, args, device=None, **kwargs): + """ + Initialize the Inference pipeline. + + Args: + pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints. + args (argparse.Namespace): The arguments for the pipeline. + device (int): The device for inference. Default is 0. + """ + # ======================================================================== + logger.info(f"Got text-to-video model root path: {pretrained_model_path}") + + # ======================== Get the args path ============================= + + # Set device and disable gradient + #if device is None: + # device = "cuda" if torch.cuda.is_available() else "cpu" + torch.set_grad_enabled(False) + device = "cpu" + + # =========================== Build main model =========================== + logger.info("Building model...") + factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]} + in_channels = args.latent_channels + out_channels = args.latent_channels + + # model = load_model( + # args, + # in_channels=in_channels, + # out_channels=out_channels, + # factor_kwargs=factor_kwargs, + # ) + with init_empty_weights(): + transformer = HYVideoDiffusionTransformer( + args, + in_channels=in_channels, + out_channels=out_channels, + **HUNYUAN_VIDEO_CONFIG[args.model], + **factor_kwargs, + ) + + #model = Inference.load_state_dict(args, model, pretrained_model_path) + model_path = "ckpts/hunyuan-video-t2v-720p/transformers/hunyuan_video_720_fp8_e4m3fn_keep_bias.safetensors" + sd = safetensors.torch.load_file(model_path) + base_dtype = torch.bfloat16 + dtype = torch.float8_e4m3fn + params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"} + for name, param in transformer.named_parameters(): + dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype + set_module_tensor_to_device(transformer, name, device=device, dtype=dtype_to_use, value=sd[name]) + transformer.eval() + + + # ============================= Build extra models ======================== + # VAE + + vae, _, s_ratio, t_ratio = load_vae( + vae_type = "884-16c-hy", + vae_precision = "bf16", + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + vae_kwargs = {"s_ratio": s_ratio, "t_ratio": t_ratio} + + # Text encoder + if args.prompt_template_video is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template_video].get( + "crop_start", 0 + ) + elif args.prompt_template is not None: + crop_start = PROMPT_TEMPLATE[args.prompt_template].get("crop_start", 0) + else: + crop_start = 0 + max_length = args.text_len + crop_start + + # prompt_template + prompt_template = ( + PROMPT_TEMPLATE[args.prompt_template] + if args.prompt_template is not None + else None + ) + + # prompt_template_video + prompt_template_video = ( + PROMPT_TEMPLATE[args.prompt_template_video] + if args.prompt_template_video is not None + else None + ) + + text_encoder = TextEncoder( + text_encoder_type=args.text_encoder, + max_length=max_length, + text_encoder_precision=args.text_encoder_precision, + tokenizer_type=args.tokenizer, + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=args.hidden_state_skip_layer, + apply_final_norm=args.apply_final_norm, + reproduce=args.reproduce, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + text_encoder_2 = None + if args.text_encoder_2 is not None: + text_encoder_2 = TextEncoder( + text_encoder_type=args.text_encoder_2, + max_length=args.text_len_2, + text_encoder_precision=args.text_encoder_precision_2, + tokenizer_type=args.tokenizer_2, + reproduce=args.reproduce, + logger=logger, + device=device if not args.use_cpu_offload else "cpu", + ) + + return cls( + args=args, + vae=vae, + vae_kwargs=vae_kwargs, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + model=transformer, + use_cpu_offload=args.use_cpu_offload, + device=device, + logger=logger, + ) + + + @staticmethod + def parse_size(size): + if isinstance(size, int): + size = [size] + if not isinstance(size, (list, tuple)): + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + if len(size) == 1: + size = [size[0], size[0]] + if len(size) != 2: + raise ValueError(f"Size must be an integer or (height, width), got {size}.") + return size + + +class HunyuanVideoSampler(Inference): + def __init__( + self, + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=None, + pipeline=None, + use_cpu_offload=True, + device=0, + logger=None, + ): + super().__init__( + args, + vae, + vae_kwargs, + text_encoder, + model, + text_encoder_2=text_encoder_2, + pipeline=pipeline, + use_cpu_offload=use_cpu_offload, + device=device, + logger=logger, + ) + + self.pipeline = self.load_diffusion_pipeline( + args=args, + vae=self.vae, + text_encoder=self.text_encoder, + text_encoder_2=self.text_encoder_2, + model=self.model, + device=self.device, + ) + + self.default_negative_prompt = NEGATIVE_PROMPT + + def load_diffusion_pipeline( + self, + args, + vae, + text_encoder, + text_encoder_2, + model, + scheduler=None, + device=None, + progress_bar_config=None, + data_type="video", + ): + """Load the denoising scheduler for inference.""" + if scheduler is None: + if args.denoise_type == "flow": + scheduler = FlowMatchDiscreteScheduler( + shift=args.flow_shift, + reverse=args.flow_reverse, + solver=args.flow_solver, + ) + else: + raise ValueError(f"Invalid denoise type {args.denoise_type}") + + pipeline = HunyuanVideoPipeline( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + transformer=model, + scheduler=scheduler, + progress_bar_config=progress_bar_config, + args=args, + ) + if self.use_cpu_offload: + pipeline.enable_sequential_cpu_offload() + else: + pipeline = pipeline.to(device) + + return pipeline + + def get_rotary_pos_embed(self, video_length, height, width): + target_ndim = 3 + ndim = 5 - 2 + # 884 + if "884" in self.args.vae: + latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8] + elif "888" in self.args.vae: + latents_size = [(video_length - 1) // 8 + 1, height // 8, width // 8] + else: + latents_size = [video_length, height // 8, width // 8] + + if isinstance(self.model.patch_size, int): + assert all(s % self.model.patch_size == 0 for s in latents_size), ( + f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), " + f"but got {latents_size}." + ) + rope_sizes = [s // self.model.patch_size for s in latents_size] + elif isinstance(self.model.patch_size, list): + assert all( + s % self.model.patch_size[idx] == 0 + for idx, s in enumerate(latents_size) + ), ( + f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.model.patch_size}), " + f"but got {latents_size}." + ) + rope_sizes = [ + s // self.model.patch_size[idx] for idx, s in enumerate(latents_size) + ] + + if len(rope_sizes) != target_ndim: + rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis + head_dim = self.model.hidden_size // self.model.heads_num + rope_dim_list = self.model.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.args.rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + + @torch.no_grad() + def predict( + self, + prompt, + height=192, + width=336, + video_length=129, + seed=None, + negative_prompt=None, + infer_steps=50, + guidance_scale=6, + flow_shift=5.0, + embedded_guidance_scale=None, + batch_size=1, + num_videos_per_prompt=1, + **kwargs, + ): + """ + Predict the image/video from the given text. + + Args: + prompt (str or List[str]): The input text. + kwargs: + height (int): The height of the output video. Default is 192. + width (int): The width of the output video. Default is 336. + video_length (int): The frame number of the output video. Default is 129. + seed (int or List[str]): The random seed for the generation. Default is a random integer. + negative_prompt (str or List[str]): The negative text prompt. Default is an empty string. + guidance_scale (float): The guidance scale for the generation. Default is 6.0. + num_images_per_prompt (int): The number of images per prompt. Default is 1. + infer_steps (int): The number of inference steps. Default is 100. + """ + out_dict = dict() + + # ======================================================================== + # Arguments: seed + # ======================================================================== + if isinstance(seed, torch.Tensor): + seed = seed.tolist() + if seed is None: + seeds = [ + random.randint(0, 1_000_000) + for _ in range(batch_size * num_videos_per_prompt) + ] + elif isinstance(seed, int): + seeds = [ + seed + i + for _ in range(batch_size) + for i in range(num_videos_per_prompt) + ] + elif isinstance(seed, (list, tuple)): + if len(seed) == batch_size: + seeds = [ + int(seed[i]) + j + for i in range(batch_size) + for j in range(num_videos_per_prompt) + ] + elif len(seed) == batch_size * num_videos_per_prompt: + seeds = [int(s) for s in seed] + else: + raise ValueError( + f"Length of seed must be equal to number of prompt(batch_size) or " + f"batch_size * num_videos_per_prompt ({batch_size} * {num_videos_per_prompt}), got {seed}." + ) + else: + raise ValueError( + f"Seed must be an integer, a list of integers, or None, got {seed}." + ) + generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds] + out_dict["seeds"] = seeds + + # ======================================================================== + # Arguments: target_width, target_height, target_video_length + # ======================================================================== + if width <= 0 or height <= 0 or video_length <= 0: + raise ValueError( + f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}" + ) + if (video_length - 1) % 4 != 0: + raise ValueError( + f"`video_length-1` must be a multiple of 4, got {video_length}" + ) + + logger.info( + f"Input (height, width, video_length) = ({height}, {width}, {video_length})" + ) + + target_height = align_to(height, 16) + target_width = align_to(width, 16) + target_video_length = video_length + + out_dict["size"] = (target_height, target_width, target_video_length) + + # ======================================================================== + # Arguments: prompt, new_prompt, negative_prompt + # ======================================================================== + if not isinstance(prompt, str): + raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") + prompt = [prompt.strip()] + + # negative prompt + if negative_prompt is None or negative_prompt == "": + negative_prompt = self.default_negative_prompt + if not isinstance(negative_prompt, str): + raise TypeError( + f"`negative_prompt` must be a string, but got {type(negative_prompt)}" + ) + negative_prompt = [negative_prompt.strip()] + + # ======================================================================== + # Scheduler + # ======================================================================== + scheduler = FlowMatchDiscreteScheduler( + shift=flow_shift, + reverse=self.args.flow_reverse, + solver=self.args.flow_solver + ) + self.pipeline.scheduler = scheduler + + # ======================================================================== + # Build Rope freqs + # ======================================================================== + freqs_cos, freqs_sin = self.get_rotary_pos_embed( + target_video_length, target_height, target_width + ) + n_tokens = freqs_cos.shape[0] + + # ======================================================================== + # Print infer args + # ======================================================================== + debug_str = f""" + height: {target_height} + width: {target_width} + video_length: {target_video_length} + prompt: {prompt} + neg_prompt: {negative_prompt} + seed: {seed} + infer_steps: {infer_steps} + num_videos_per_prompt: {num_videos_per_prompt} + guidance_scale: {guidance_scale} + n_tokens: {n_tokens} + flow_shift: {flow_shift} + embedded_guidance_scale: {embedded_guidance_scale}""" + logger.debug(debug_str) + + # ======================================================================== + # Pipeline inference + # ======================================================================== + start_time = time.time() + samples = self.pipeline( + prompt=prompt, + height=target_height, + width=target_width, + video_length=target_video_length, + num_inference_steps=infer_steps, + guidance_scale=guidance_scale, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + output_type="pil", + freqs_cis=(freqs_cos, freqs_sin), + n_tokens=n_tokens, + embedded_guidance_scale=embedded_guidance_scale, + data_type="video" if target_video_length > 1 else "image", + is_progress_bar=True, + vae_ver=self.args.vae, + enable_tiling=self.args.vae_tiling, + )[0] + out_dict["samples"] = samples + out_dict["prompts"] = prompt + + gen_time = time.time() - start_time + logger.info(f"Success, time: {gen_time}") + + return out_dict diff --git a/hyvideo/modules/__init__.py b/hyvideo/modules/__init__.py new file mode 100644 index 0000000..a95842b --- /dev/null +++ b/hyvideo/modules/__init__.py @@ -0,0 +1,3 @@ +from .models import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG + + diff --git a/hyvideo/modules/activation_layers.py b/hyvideo/modules/activation_layers.py new file mode 100644 index 0000000..f8774c2 --- /dev/null +++ b/hyvideo/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/hyvideo/modules/attention.py b/hyvideo/modules/attention.py new file mode 100644 index 0000000..057d44b --- /dev/null +++ b/hyvideo/modules/attention.py @@ -0,0 +1,188 @@ +import importlib.metadata +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from flash_attn.flash_attn_interface import flash_attn_varlen_func +except ImportError: + flash_attn_varlen_func = None +try: + from sageattention import sageattn_varlen + @torch.compiler.disable() + def sageattn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + ): + return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) +except: + pass + +MEMORY_LAYOUT = { + "flash_attn": ( + lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]), + lambda x: x, + ), + "sageattn_varlen": ( + lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]), + lambda x: x, + ), + "sdpa": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), + "sageattn": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), + "vanilla": ( + lambda x: x.transpose(1, 2), + lambda x: x.transpose(1, 2), + ), +} + + +def get_cu_seqlens(text_mask, img_len): + """Calculate cu_seqlens_q, cu_seqlens_kv using text_mask and img_len + + Args: + text_mask (torch.Tensor): the mask of text + img_len (int): the length of image + + Returns: + torch.Tensor: the calculated cu_seqlens for flash attention + """ + batch_size = text_mask.shape[0] + text_len = text_mask.sum(dim=1) + max_len = text_mask.shape[1] + img_len + + cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda") + + for i in range(batch_size): + s = text_len[i] + img_len + s1 = i * max_len + s + s2 = (i + 1) * max_len + cu_seqlens[2 * i + 1] = s1 + cu_seqlens[2 * i + 2] = s2 + + return cu_seqlens + + +def attention( + q, + k, + v, + mode="flash_attn", + drop_rate=0, + attn_mask=None, + causal=False, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + batch_size=1, +): + """ + Perform QKV self attention. + + Args: + q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads. + k (torch.Tensor): Key tensor with shape [b, s1, a, d] + v (torch.Tensor): Value tensor with shape [b, s1, a, d] + mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'. + drop_rate (float): Dropout rate in attention map. (default: 0) + attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla). + (default: None) + causal (bool): Whether to use causal attention. (default: False) + cu_seqlens_q (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch, + used to index into q. + cu_seqlens_kv (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch, + used to index into kv. + max_seqlen_q (int): The maximum sequence length in the batch of q. + max_seqlen_kv (int): The maximum sequence length in the batch of k and v. + + Returns: + torch.Tensor: Output tensor after self attention with shape [b, s, ad] + """ + pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode] + q = pre_attn_layout(q) + k = pre_attn_layout(k) + v = pre_attn_layout(v) + + if mode == "sdpa": + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(q.dtype) + x = F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal + ) + elif mode == "sageattn_varlen": + x = sageattn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + ) + # x with shape [(bxs), a, d] + x = x.view( + batch_size, max_seqlen_q, x.shape[-2], x.shape[-1] + ) # reshape x to [b, s, a, d] + elif mode == "flash_attn": + x = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + ) + # x with shape [(bxs), a, d] + x = x.view( + batch_size, max_seqlen_q, x.shape[-2], x.shape[-1] + ) # reshape x to [b, s, a, d] + elif mode == "vanilla": + scale_factor = 1 / math.sqrt(q.size(-1)) + + b, a, s, _ = q.shape + s1 = k.size(2) + attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device) + if causal: + # Only applied to self attention + assert ( + attn_mask is None + ), "Causal mask and attn_mask cannot be used together" + temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril( + diagonal=0 + ) + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) + attn_bias.to(q.dtype) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) + else: + attn_bias += attn_mask + + # TODO: Maybe force q and k to be float32 to avoid numerical overflow + attn = (q @ k.transpose(-2, -1)) * scale_factor + attn += attn_bias + attn = attn.softmax(dim=-1) + attn = torch.dropout(attn, p=drop_rate, train=True) + x = attn @ v + else: + raise NotImplementedError(f"Unsupported attention mode: {mode}") + + x = post_attn_layout(x) + b, s, a, d = x.shape + out = x.reshape(b, s, -1) + return out diff --git a/hyvideo/modules/embed_layers.py b/hyvideo/modules/embed_layers.py new file mode 100644 index 0000000..3d65ed1 --- /dev/null +++ b/hyvideo/modules/embed_layers.py @@ -0,0 +1,157 @@ +import math +import torch +import torch.nn as nn +from einops import rearrange, repeat + +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/hyvideo/modules/mlp_layers.py b/hyvideo/modules/mlp_layers.py new file mode 100644 index 0000000..24dd2d9 --- /dev/null +++ b/hyvideo/modules/mlp_layers.py @@ -0,0 +1,118 @@ +# Modified from timm library: +# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13 + +from functools import partial + +import torch +import torch.nn as nn + +from .modulate_layers import modulate +from ..utils.helpers import to_2tuple + + +class MLP(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks""" + + def __init__( + self, + in_channels, + hidden_channels=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + device=None, + dtype=None, + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + out_features = out_features or in_channels + hidden_channels = hidden_channels or in_channels + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer( + in_channels, hidden_channels, bias=bias[0], **factory_kwargs + ) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = ( + norm_layer(hidden_channels, **factory_kwargs) + if norm_layer is not None + else nn.Identity() + ) + self.fc2 = linear_layer( + hidden_channels, out_features, bias=bias[1], **factory_kwargs + ) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +# +class MLPEmbedder(nn.Module): + """copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py""" + def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__( + self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + # Just use LayerNorm for the final layer + self.norm_final = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs + ) + if isinstance(patch_size, int): + self.linear = nn.Linear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + **factory_kwargs + ) + else: + self.linear = nn.Linear( + hidden_size, + patch_size[0] * patch_size[1] * patch_size[2] * out_channels, + bias=True, + ) + nn.init.zeros_(self.linear.weight) + nn.init.zeros_(self.linear.bias) + + # Here we don't distinguish between the modulate types. Just use the simple one. + self.adaLN_modulation = nn.Sequential( + act_layer(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs), + ) + # Zero-initialize the modulation + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift=shift, scale=scale) + x = self.linear(x) + return x diff --git a/hyvideo/modules/models.py b/hyvideo/modules/models.py new file mode 100644 index 0000000..5ca086d --- /dev/null +++ b/hyvideo/modules/models.py @@ -0,0 +1,739 @@ +from typing import Any, List, Tuple, Optional, Union, Dict +from einops import rearrange + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from diffusers.models import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from .activation_layers import get_activation_layer +from .norm_layers import get_norm_layer +from .embed_layers import TimestepEmbedder, PatchEmbed, TextProjection +from .attention import attention, get_cu_seqlens +from .posemb_layers import apply_rotary_emb +from .mlp_layers import MLP, MLPEmbedder, FinalLayer +from .modulate_layers import ModulateDiT, modulate, apply_gate +from .token_refiner import SingleTokenRefiner + + +class MMDoubleStreamBlock(nn.Module): + """ + A multimodal dit block with seperate 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, + attention_mode: str = "flash_attn", + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.attention_mode = attention_mode + + 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, + ) + + 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, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + freqs_cis: tuple = 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: + 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) + + # Run actual attention. + q = torch.cat((img_q, txt_q), dim=1) + k = torch.cat((img_k, txt_k), dim=1) + v = torch.cat((img_v, txt_v), dim=1) + assert ( + cu_seqlens_q.shape[0] == 2 * img.shape[0] + 1 + ), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, img.shape[0]:{img.shape[0]}" + attn = attention( + q, + k, + v, + mode=self.attention_mode, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + batch_size=img_k.shape[0], + ) + + img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1] :] + + # Calculate the img bloks. + 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 bloks. + 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, + attention_mode: str = "flash_attn", + ): + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + + self.attention_mode = attention_mode + + 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, + ) + + 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, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + freqs_cis: Tuple[torch.Tensor, 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) + + # Apply RoPE if needed. + if freqs_cis is not None: + img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :] + img_k, txt_k = k[:, :-txt_len, :, :], k[:, -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 + q = torch.cat((img_q, txt_q), dim=1) + k = torch.cat((img_k, txt_k), dim=1) + + # Compute attention. + assert ( + cu_seqlens_q.shape[0] == 2 * x.shape[0] + 1 + ), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}" + attn = attention( + q, + k, + v, + mode=self.attention_mode, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + batch_size=x.shape[0], + ) + + # 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, + text_states_dim: int = 4096, + text_states_dim_2: int = 768, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + attention_mode: str = "flash_attn", + ): + 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 + + # 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 + + self.text_states_dim = text_states_dim + self.text_states_dim_2 = text_states_dim_2 + + 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.text_states_dim, + self.hidden_size, + get_activation_layer("silu"), + **factory_kwargs, + ) + elif self.text_projection == "single_refiner": + self.txt_in = SingleTokenRefiner( + self.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.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, + attention_mode=attention_mode, + **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, + attention_mode=attention_mode, + **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 forward( + self, + 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. + freqs_cos: Optional[torch.Tensor] = None, + freqs_sin: Optional[torch.Tensor] = None, + guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000. + return_dict: bool = True, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + out = {} + img = x + txt = text_states + _, _, ot, oh, ow = x.shape + tt, th, tw = ( + ot // self.patch_size[0], + oh // self.patch_size[1], + ow // self.patch_size[2], + ) + + # 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] + + # Compute cu_squlens and max_seqlen for flash attention + cu_seqlens_q = get_cu_seqlens(text_mask, img_seq_len) + cu_seqlens_kv = cu_seqlens_q + max_seqlen_q = img_seq_len + txt_seq_len + max_seqlen_kv = max_seqlen_q + + 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, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + freqs_cis, + ] + + img, txt = block(*double_block_args) + + # Merge txt and img to pass through single stream blocks. + x = torch.cat((img, txt), 1) + if len(self.single_blocks) > 0: + for _, block in enumerate(self.single_blocks): + single_block_args = [ + x, + vec, + txt_seq_len, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + (freqs_cos, freqs_sin), + ] + + x = block(*single_block_args) + + 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) + if return_dict: + out["x"] = img + return out + return img + + 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/hyvideo/modules/modulate_layers.py b/hyvideo/modules/modulate_layers.py new file mode 100644 index 0000000..93a57c6 --- /dev/null +++ b/hyvideo/modules/modulate_layers.py @@ -0,0 +1,76 @@ +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 diff --git a/hyvideo/modules/norm_layers.py b/hyvideo/modules/norm_layers.py new file mode 100644 index 0000000..d8c73b1 --- /dev/null +++ b/hyvideo/modules/norm_layers.py @@ -0,0 +1,77 @@ +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/hyvideo/modules/posemb_layers.py b/hyvideo/modules/posemb_layers.py new file mode 100644 index 0000000..dfce82c --- /dev/null +++ b/hyvideo/modules/posemb_layers.py @@ -0,0 +1,310 @@ +import torch +from typing import Union, Tuple, List + + +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/hyvideo/modules/token_refiner.py b/hyvideo/modules/token_refiner.py new file mode 100644 index 0000000..4a8d1df --- /dev/null +++ b/hyvideo/modules/token_refiner.py @@ -0,0 +1,236 @@ +from typing import Optional + +from einops import rearrange +import torch +import torch.nn as nn + +from .activation_layers import get_activation_layer +from .attention import attention +from .norm_layers import get_norm_layer +from .embed_layers import TimestepEmbedder, TextProjection +from .attention import attention +from .mlp_layers import MLP +from .modulate_layers import modulate, apply_gate + + +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, mode="sdpa", 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, + ): + self_attn_mask = None + if mask is not None: + batch_size = mask.shape[0] + seq_len = mask.shape[1] + mask = mask.to(x.device) + # batch_size x 1 x seq_len x seq_len + self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat( + 1, 1, seq_len, 1 + ) + # batch_size x 1 x seq_len x seq_len + self_attn_mask_2 = self_attn_mask_1.transpose(2, 3) + # batch_size x 1 x seq_len x seq_len, 1 for broadcasting of heads_num + self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool() + # avoids self-attention weight being NaN for padding tokens + self_attn_mask[:, :, :, 0] = True + + for block in self.blocks: + x = block(x, c, self_attn_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 = "sdpa", + 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 == "sdpa", "Only support 'torch sdpa' 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/hyvideo/prompt_rewrite.py b/hyvideo/prompt_rewrite.py new file mode 100644 index 0000000..974c452 --- /dev/null +++ b/hyvideo/prompt_rewrite.py @@ -0,0 +1,51 @@ +normal_mode_prompt = """Normal mode - Video Recaption Task: + +You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description. + +0. Preserve ALL information, including style words and technical terms. + +1. If the input is in Chinese, translate the entire description to English. + +2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences. + +3. If the input does not include style, lighting, atmosphere, you can make reasonable associations. + +4. Output ALL must be in English. + +Given Input: +input: "{input}" +""" + + +master_mode_prompt = """Master mode - Video Recaption Task: + +You are a large language model specialized in rewriting video descriptions. Your task is to modify the input description. + +0. Preserve ALL information, including style words and technical terms. + +1. If the input is in Chinese, translate the entire description to English. + +2. If the input is just one or two words describing an object or person, provide a brief, simple description focusing on basic visual characteristics. Limit the description to 1-2 short sentences. + +3. If the input does not include style, lighting, atmosphere, you can make reasonable associations. + +4. Output ALL must be in English. + +Given Input: +input: "{input}" +""" + +def get_rewrite_prompt(ori_prompt, mode="Normal"): + if mode == "Normal": + prompt = normal_mode_prompt.format(input=ori_prompt) + elif mode == "Master": + prompt = master_mode_prompt.format(input=ori_prompt) + else: + raise Exception("Only supports Normal and Normal", mode) + return prompt + +ori_prompt = "一只小狗在草地上奔跑。" +normal_prompt = get_rewrite_prompt(ori_prompt, mode="Normal") +master_prompt = get_rewrite_prompt(ori_prompt, mode="Master") + +# Then you can use the normal_prompt or master_prompt to access the hunyuan-large rewrite model to get the final prompt. \ No newline at end of file diff --git a/hyvideo/text_encoder/__init__.py b/hyvideo/text_encoder/__init__.py new file mode 100644 index 0000000..494992b --- /dev/null +++ b/hyvideo/text_encoder/__init__.py @@ -0,0 +1,357 @@ +from dataclasses import dataclass +from typing import Optional, Tuple +from copy import deepcopy + +import torch +import torch.nn as nn +from transformers import CLIPTextModel, CLIPTokenizer, AutoTokenizer, AutoModel +from transformers.utils import ModelOutput + +from ..constants import TEXT_ENCODER_PATH, TOKENIZER_PATH +from ..constants import PRECISION_TO_TYPE + + +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 Ture, 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/hyvideo/utils/__init__.py b/hyvideo/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hyvideo/utils/data_utils.py b/hyvideo/utils/data_utils.py new file mode 100644 index 0000000..583a903 --- /dev/null +++ b/hyvideo/utils/data_utils.py @@ -0,0 +1,15 @@ +import numpy as np +import math + + +def align_to(value, alignment): + """align hight, 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/hyvideo/utils/file_utils.py b/hyvideo/utils/file_utils.py new file mode 100644 index 0000000..2ba3651 --- /dev/null +++ b/hyvideo/utils/file_utils.py @@ -0,0 +1,70 @@ +import os +from pathlib import Path +from einops import rearrange + +import torch +import torchvision +import numpy as np +import imageio + +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/hyvideo/utils/helpers.py b/hyvideo/utils/helpers.py new file mode 100644 index 0000000..72ab8cb --- /dev/null +++ b/hyvideo/utils/helpers.py @@ -0,0 +1,40 @@ +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/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py b/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py new file mode 100644 index 0000000..2908eb2 --- /dev/null +++ b/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py @@ -0,0 +1,46 @@ +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/hyvideo/vae/__init__.py b/hyvideo/vae/__init__.py new file mode 100644 index 0000000..d8b4f98 --- /dev/null +++ b/hyvideo/vae/__init__.py @@ -0,0 +1,61 @@ +from pathlib import Path + +import torch + +from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from ..constants import VAE_PATH, PRECISION_TO_TYPE + +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 fucntion 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"] + vae_ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} + vae.load_state_dict(vae_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/hyvideo/vae/autoencoder_kl_causal_3d.py b/hyvideo/vae/autoencoder_kl_causal_3d.py new file mode 100644 index 0000000..a8a504c --- /dev/null +++ b/hyvideo/vae/autoencoder_kl_causal_3d.py @@ -0,0 +1,626 @@ +# 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 Dict, Optional, Tuple, Union +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from diffusers.configuration_utils import ConfigMixin, register_to_config +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.utils.accelerate_utils import apply_forward_hook +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 .vae import DecoderCausal3D, BaseOutput, 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 + + # only relevant if vae tiling is enabled + sample_tsize = 16 + 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_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_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 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/hyvideo/vae/unet_causal_3d_blocks.py b/hyvideo/vae/unet_causal_3d_blocks.py new file mode 100644 index 0000000..d38ddf4 --- /dev/null +++ b/hyvideo/vae/unet_causal_3d_blocks.py @@ -0,0 +1,797 @@ +# 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 torch import nn +from einops import rearrange + +from diffusers.utils import logging +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import SpatialNorm +from diffusers.models.attention_processor import Attention +from diffusers.models.normalization import AdaGroupNorm +from diffusers.models.normalization import RMSNorm + + +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: + assert False, "Not Implement yet" + if kernel_size is None: + kernel_size = 4 + conv = nn.ConvTranspose2d( + channels, self.out_channels, kernel_size=kernel_size, stride=2, padding=padding, bias=bias + ) + 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: + assert False, "Not Implement yet" + hidden_states = self.norm( + hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) + + 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: + assert False, "Not Implement yet" + other_h = F.interpolate( + other_h, size=output_size, mode="nearest") + + 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: + # assert False, "Not implemented yet" + 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/hyvideo/vae/vae.py b/hyvideo/vae/vae.py new file mode 100644 index 0000000..fa44bc4 --- /dev/null +++ b/hyvideo/vae/vae.py @@ -0,0 +1,373 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import numpy as np +import torch +import torch.nn as nn + +from diffusers.utils import BaseOutput, is_torch_version +from diffusers.utils.torch_utils import randn_tensor +from diffusers.models.attention_processor import SpatialNorm +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) + elif time_compression_ratio == 8: + add_spatial_downsample = bool( + i < num_spatial_downsample_layers) + add_time_downsample = bool(i < num_time_downsample_layers) + else: + raise ValueError( + f"Unsupported time_compression_ratio: {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/nodes.py b/nodes.py new file mode 100644 index 0000000..c9a5ce4 --- /dev/null +++ b/nodes.py @@ -0,0 +1,802 @@ +import os +import torch +import json +from einops import rearrange +from contextlib import nullcontext +from typing import List +from pathlib import Path +from .utils import log, check_diffusers_version, print_memory +from diffusers.video_processor import VideoProcessor + +from .hyvideo.constants import PROMPT_TEMPLATE, NEGATIVE_PROMPT, PRECISION_TO_TYPE +from .hyvideo.vae import load_vae +from .hyvideo.text_encoder import TextEncoder +from .hyvideo.utils.data_utils import align_to +from .hyvideo.modules.posemb_layers import get_nd_rotary_pos_embed +from .hyvideo.diffusion.schedulers import FlowMatchDiscreteScheduler +from .hyvideo.diffusion.pipelines import HunyuanVideoPipeline +from .hyvideo.vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D +from .hyvideo.modules.models import HYVideoDiffusionTransformer +from accelerate import init_empty_weights +from accelerate.utils import set_module_tensor_to_device + +import folder_paths +import comfy.model_management as mm +from comfy.utils import load_torch_file + +script_directory = os.path.dirname(os.path.abspath(__file__)) + +def get_rotary_pos_embed(transformer, video_length, height, width): + target_ndim = 3 + ndim = 5 - 2 + rope_theta = 225 + patch_size = transformer.patch_size + rope_dim_list = transformer.rope_dim_list + hidden_size = transformer.hidden_size + heads_num = transformer.heads_num + head_dim = hidden_size // heads_num + + # 884 + latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8] + + if isinstance(patch_size, int): + assert all(s % patch_size == 0 for s in latents_size), ( + f"Latent size(last {ndim} dimensions) should be divisible by patch size({patch_size}), " + f"but got {latents_size}." + ) + rope_sizes = [s // patch_size for s in latents_size] + elif isinstance(patch_size, list): + assert all( + s % patch_size[idx] == 0 + for idx, s in enumerate(latents_size) + ), ( + f"Latent size(last {ndim} dimensions) should be divisible by patch size({patch_size}), " + f"but got {latents_size}." + ) + rope_sizes = [ + s // patch_size[idx] for idx, s in enumerate(latents_size) + ] + + if len(rope_sizes) != target_ndim: + rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes # time axis + + 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=rope_theta, + use_real=True, + theta_rescale_factor=1, + ) + return freqs_cos, freqs_sin + +#region Model loading +class HyVideoModelLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}), + + "base_precision": (["fp16", "fp32", "bf16"], {"default": "bf16"}), + "quantization": (['disabled', 'fp8_e4m3fn', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6"], {"default": 'disabled', "tooltip": "optional quantization method"}), + "load_device": (["main_device", "offload_device"], {"default": "main_device"}), + "enable_sequential_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "significantly reducing memory usage and slows down the inference"}), + }, + "optional": { + "attention_mode": ([ + "sdpa", + "flash_attn", + "sageattn_varlen", + ], {"default": "flash_attn"}), + "compile_args": ("COMPILEARGS", ), + } + } + + RETURN_TYPES = ("HYVIDEOMODEL",) + RETURN_NAMES = ("model", ) + FUNCTION = "loadmodel" + CATEGORY = "HunyuanVideoWrapper" + + def loadmodel(self, model, base_precision, load_device, quantization, + compile_args=None, attention_mode="sdpa", enable_sequential_cpu_offload=False): + transformer = None + manual_offloading = True + if "sage" in attention_mode: + try: + from sageattention import sageattn_varlen + except Exception as e: + raise ValueError(f"Can't import SageAttention: {str(e)}") + + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + manual_offloading = True + transformer_load_device = device if load_device == "main_device" else offload_device + mm.soft_empty_cache() + + base_dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "fp8_e4m3fn_fast": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[base_precision] + + model_path = folder_paths.get_full_path_or_raise("diffusion_models", model) + sd = load_torch_file(model_path, device=transformer_load_device) + + in_channels = out_channels = 16 + factor_kwargs = {"device": device, "dtype": base_dtype} + HUNYUAN_VIDEO_CONFIG = { + "mm_double_blocks_depth": 20, + "mm_single_blocks_depth": 40, + "rope_dim_list": [16, 56, 56], + "hidden_size": 3072, + "heads_num": 24, + "mlp_width_ratio": 4, + "guidance_embed": True, + } + with init_empty_weights(): + transformer = HYVideoDiffusionTransformer( + in_channels=in_channels, + out_channels=out_channels, + attention_mode=attention_mode, + **HUNYUAN_VIDEO_CONFIG, + **factor_kwargs + ) + + log.info("Using accelerate to load and assign model weights to device...") + if quantization == "fp8_e4m3fn": + dtype = torch.float8_e4m3fn + else: + dtype = base_dtype + params_to_keep = {"norm", "bias", "time_in", "vector_in", "guidance_in", "txt_in", "img_in"} + for name, param in transformer.named_parameters(): + dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype + set_module_tensor_to_device(transformer, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name]) + transformer.eval() + + #compile + if compile_args is not None: + torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"] + for i, block in enumerate(transformer.single_blocks): + transformer.single_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) + for i, block in enumerate(transformer.double_blocks): + transformer.double_blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) + + if "torchao" in quantization: + try: + from torchao.quantization import ( + quantize_, + fpx_weight_only, + float8_dynamic_activation_float8_weight, + int8_dynamic_activation_int8_weight + ) + except: + raise ImportError("torchao is not installed, please install torchao to use fp8dq") + + # def filter_fn(module: nn.Module, fqn: str) -> bool: + # target_submodules = {'attn1', 'ff'} # avoid norm layers, 1.5 at least won't work with quantized norm1 #todo: test other models + # if any(sub in fqn for sub in target_submodules): + # return isinstance(module, nn.Linear) + # return False + + if "fp6" in quantization: #slower for some reason on 4090 + quant_func = fpx_weight_only(3, 2) + elif "fp8dq" in quantization: #very fast on 4090 when compiled + quant_func = float8_dynamic_activation_float8_weight() + elif 'fp8dqrow' in quantization: + from torchao.quantization.quant_api import PerRow + quant_func = float8_dynamic_activation_float8_weight(granularity=PerRow()) + elif 'int8dq' in quantization: + quant_func = int8_dynamic_activation_int8_weight() + + quantize_(transformer, quant_func) + + manual_offloading = False # to disable manual .to(device) calls + log.info(f"Quantized transformer blocks to {quantization}") + + scheduler = FlowMatchDiscreteScheduler( + shift=9.0, #this is not even used? + reverse=True, #has to be true or noise + solver="euler", #has to be euler + ) + + pipe = HunyuanVideoPipeline( + transformer=transformer, + scheduler=scheduler, + progress_bar_config=None + ) + if enable_sequential_cpu_offload: + pipe.enable_sequential_cpu_offload() + manual_offloading = False + + + pipeline = { + "pipe": pipe, + "dtype": base_dtype, + "base_path": model_path, + "cpu_offloading": enable_sequential_cpu_offload, + "model_name": model, + "manual_offloading": manual_offloading, + "quantization": "disabled", + } + return (pipeline,) + +#region load VAE + +class HyVideoVAELoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}), + }, + "optional": { + "precision": (["fp16", "fp32", "bf16"], + {"default": "bf16"} + ), + "compile_args":("COMPILEARGS", ), + } + } + + RETURN_TYPES = ("VAE",) + RETURN_NAMES = ("vae", ) + FUNCTION = "loadmodel" + CATEGORY = "HunyuanVideoWrapper" + DESCRIPTION = "Loads Hunyuan VAE model from 'ComfyUI/models/vae'" + + def loadmodel(self, model_name, precision, compile_args=None): + + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] + with open(os.path.join(script_directory, 'configs', 'hy_vae_config.json')) as f: + vae_config = json.load(f) + model_path = folder_paths.get_full_path("vae", model_name) + vae_sd = load_torch_file(model_path) + + vae = AutoencoderKLCausal3D.from_config(vae_config).to(dtype).to(offload_device) + vae.load_state_dict(vae_sd) + vae.requires_grad_(False) + vae.eval() + #compile + if compile_args is not None: + torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"] + vae = torch.compile(vae, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) + + return (vae,) + + + +class HyVideoTorchCompileSettings: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "backend": (["inductor","cudagraphs"], {"default": "inductor"}), + "fullgraph": ("BOOLEAN", {"default": False, "tooltip": "Enable full graph mode"}), + "mode": (["default", "max-autotune", "max-autotune-no-cudagraphs", "reduce-overhead"], {"default": "default"}), + "dynamic": ("BOOLEAN", {"default": False, "tooltip": "Enable dynamic mode"}), + "dynamo_cache_size_limit": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 1, "tooltip": "torch._dynamo.config.cache_size_limit"}), + }, + } + RETURN_TYPES = ("COMPILEARGS",) + RETURN_NAMES = ("torch_compile_args",) + FUNCTION = "loadmodel" + CATEGORY = "HunyuanVideoWrapper" + DESCRIPTION = "torch.compile settings, when connected to the model loader, torch.compile of the selected layers is attempted. Requires Triton and torch 2.5.0 is recommended" + + def loadmodel(self, backend, fullgraph, mode, dynamic, dynamo_cache_size_limit): + + compile_args = { + "backend": backend, + "fullgraph": fullgraph, + "mode": mode, + "dynamic": dynamic, + "dynamo_cache_size_limit": dynamo_cache_size_limit, + } + + return (compile_args, ) + +#region TextEncode + +class DownloadAndLoadHyVideoTextEncoder: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "llm_model": (["Kijai/llava-llama-3-8b-text-encoder-tokenizer",],), + "clip_model": (["disabled","openai/clip-vit-large-patch14",],), + + "precision": (["fp16", "fp32", "bf16"], + {"default": "bf16"} + ), + }, + } + + RETURN_TYPES = ("HYVIDTEXTENCODER",) + RETURN_NAMES = ("hyvid_text_encoder", ) + FUNCTION = "loadmodel" + CATEGORY = "HunyuanVideoWrapper" + DESCRIPTION = "Loads Hunyuan text_encoder model from 'ComfyUI/models/LLM'" + + def loadmodel(self, llm_model, clip_model, precision): + + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] + if clip_model != "disabled": + clip_model_path = os.path.join(folder_paths.models_dir, "clip", "clip-vit-large-patch14") + if not os.path.exists(clip_model_path): + log.info(f"Downloading clip model to: {clip_model_path}") + from huggingface_hub import snapshot_download + snapshot_download( + repo_id=clip_model, + ignore_patterns=["*.msgpack", "*.bin", "*.h5"], + local_dir=clip_model_path, + local_dir_use_symlinks=False, + ) + + text_encoder_2 = TextEncoder( + text_encoder_path=clip_model_path, + text_encoder_type="clipL", + max_length=77, + text_encoder_precision=precision, + tokenizer_type="clipL", + reproduce=True, + logger=log, + device=device, + ) + else: + text_encoder_2 = None + + download_path = os.path.join(folder_paths.models_dir,"LLM") + base_path = os.path.join(download_path, (llm_model.split("/")[-1])) + if not os.path.exists(base_path): + log.info(f"Downloading model to: {base_path}") + from huggingface_hub import snapshot_download + snapshot_download( + repo_id=llm_model, + local_dir=base_path, + local_dir_use_symlinks=False, + ) + # prompt_template + prompt_template = ( + PROMPT_TEMPLATE["dit-llm-encode"] + ) + # prompt_template_video + prompt_template_video = ( + PROMPT_TEMPLATE["dit-llm-encode-video"] + ) + + text_encoder = TextEncoder( + text_encoder_path=base_path, + text_encoder_type="llm", + max_length=256, + text_encoder_precision=precision, + tokenizer_type="llm", + prompt_template=prompt_template, + prompt_template_video=prompt_template_video, + hidden_state_skip_layer=2, + apply_final_norm=True, + reproduce=True, + logger=log, + device=device, + ) + + + hyvid_text_encoders = { + "text_encoder": text_encoder, + "text_encoder_2": text_encoder_2, + } + + return (hyvid_text_encoders,) + +class HyVideoTextEncode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "text_encoders": ("HYVIDTEXTENCODER",), + "prompt": ("STRING", {"default": "", "multiline": True} ), + "negative_prompt": ("STRING", {"default": "", "multiline": True}), + }, + "optional": { + "force_offload": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("HYVIDEMBEDS", ) + RETURN_NAMES = ("hyvid_embeds",) + FUNCTION = "process" + CATEGORY = "HunyuanVideoWrapper" + + def process(self, text_encoders, prompt, negative_prompt, force_offload=True): + device = mm.text_encoder_device() + offload_device = mm.text_encoder_offload_device() + + text_encoder_1 = text_encoders["text_encoder"] + text_encoder_2 = text_encoders["text_encoder_2"] + + def encode_prompt(self, prompt, negative_prompt, text_encoder): + batch_size = 1 + num_videos_per_prompt = 1 + do_classifier_free_guidance = True + data_type = "video" + + text_inputs = text_encoder.text2tokens(prompt, data_type=data_type) + + prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type, device=device) + prompt_embeds = prompt_outputs.hidden_state + + 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 + ) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + # max_length = prompt_embeds.shape[1] + uncond_input = text_encoder.text2tokens(uncond_tokens, data_type=data_type) + + negative_prompt_outputs = text_encoder.encode( + uncond_input, data_type=data_type, device=device + ) + negative_prompt_embeds = negative_prompt_outputs.hidden_state + + negative_attention_mask = negative_prompt_outputs.attention_mask + if negative_attention_mask is not None: + negative_attention_mask = negative_attention_mask.to(device) + _, seq_len = negative_attention_mask.shape + negative_attention_mask = negative_attention_mask.repeat( + 1, num_videos_per_prompt + ) + negative_attention_mask = negative_attention_mask.view( + batch_size * num_videos_per_prompt, seq_len + ) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to( + dtype=prompt_embeds_dtype, device=device + ) + + if negative_prompt_embeds.ndim == 2: + negative_prompt_embeds = negative_prompt_embeds.repeat( + 1, num_videos_per_prompt + ) + negative_prompt_embeds = negative_prompt_embeds.view( + batch_size * num_videos_per_prompt, -1 + ) + else: + negative_prompt_embeds = negative_prompt_embeds.repeat( + 1, num_videos_per_prompt, 1 + ) + negative_prompt_embeds = negative_prompt_embeds.view( + batch_size * num_videos_per_prompt, seq_len, -1 + ) + + return ( + prompt_embeds, + negative_prompt_embeds, + attention_mask, + negative_attention_mask, + ) + text_encoder_1.to(device) + prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask = encode_prompt(self, prompt, negative_prompt, text_encoder_1) + if force_offload: + text_encoder_1.to(offload_device) + mm.soft_empty_cache() + + if text_encoder_2 is not None: + text_encoder_2.to(device) + prompt_embeds_2, negative_prompt_embeds_2, attention_mask_2, negative_attention_mask_2 = encode_prompt(self, prompt, negative_prompt, text_encoder_2) + if force_offload: + text_encoder_2.to(offload_device) + mm.soft_empty_cache() + else: + prompt_embeds_2 = None + negative_prompt_embeds_2 = None + attention_mask_2 = None + negative_attention_mask_2 = None + + prompt_embeds_dict = { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "attention_mask": attention_mask, + "negative_attention_mask": negative_attention_mask, + "prompt_embeds_2": prompt_embeds_2, + "negative_prompt_embeds_2": negative_prompt_embeds_2, + "attention_mask_2": attention_mask_2, + "negative_attention_mask_2": negative_attention_mask_2, + } + return (prompt_embeds_dict,) + + +#region Sampler +class HyVideoSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("HYVIDEOMODEL",), + "hyvid_embeds": ("HYVIDEMBEDS", ), + "width": ("INT", {"default": 512, "min": 64, "max": 1024, "step": 16}), + "height": ("INT", {"default": 512, "min": 64, "max": 1024, "step": 16}), + "num_frames": ("INT", {"default": 49, "min": 1, "max": 1024, "step": 1}), + "steps": ("INT", {"default": 30, "min": 1}), + "guidance_scale": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "force_offload": ("BOOLEAN", {"default": True}), + + }, + "optional": { + #"samples": ("LATENT", {"tooltip": "init Latents to use for video2video process"} ), + #"denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("samples",) + FUNCTION = "process" + CATEGORY = "HunyuanVideoWrapper" + + def process(self, model, hyvid_embeds, steps, guidance_scale, seed, width, height, num_frames, samples=None, denoise_strength=1.0, force_offload=True): + mm.unload_all_models() + mm.soft_empty_cache() + + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = model["dtype"] + + generator = torch.Generator(device=torch.device("cpu")).manual_seed(seed) + + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + + if width <= 0 or height <= 0 or num_frames <= 0: + raise ValueError( + f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={num_frames}" + ) + if (num_frames - 1) % 4 != 0: + raise ValueError( + f"`video_length-1` must be a multiple of 4, got {num_frames}" + ) + + log.info( + f"Input (height, width, video_length) = ({height}, {width}, {num_frames})" + ) + + target_height = align_to(height, 16) + target_width = align_to(width, 16) + + freqs_cos, freqs_sin = get_rotary_pos_embed( + model["pipe"].transformer, num_frames, target_height, target_width + ) + n_tokens = freqs_cos.shape[0] + + # autocast_context = torch.autocast( + # mm.get_autocast_device(device), dtype=dtype + # ) if any(q in model["quantization"] for q in ("e4m3fn", "GGUF")) else nullcontext() + #with autocast_context: + if not model["cpu_offloading"] and model["manual_offloading"]: + model["pipe"].transformer.to(device) + latents = model["pipe"]( + num_inference_steps=steps, + height = target_height, + width = target_width, + video_length = num_frames, + guidance_scale=guidance_scale, + embedded_guidance_scale=guidance_scale, + latents=latents if samples is not None else None, + denoise_strength=denoise_strength, + prompt_embed_dict=hyvid_embeds, + generator=generator, + freqs_cis=(freqs_cos, freqs_sin), + n_tokens=n_tokens, + ) + + print_memory(device) + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + + if force_offload: + if not model["cpu_offloading"] and model["manual_offloading"]: + model["pipe"].transformer.to(offload_device) + mm.soft_empty_cache() + + return ({ + "samples": latents + },) + + +#region VideoDecode +class HyVideoDecode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "vae": ("VAE",), + "samples": ("LATENT",), + "enable_vae_tiling": ("BOOLEAN", {"default": True, "tooltip": "Drastically reduces memory use but may introduce seams"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "decode" + CATEGORY = "HunyuanVideoWrapper" + + def decode(self, vae, samples, enable_vae_tiling): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + latents = samples["samples"] + generator = torch.Generator(device=torch.device("cpu"))#.manual_seed(seed) + vae.to(device) + + + expand_temporal_dim = False + if len(latents.shape) == 4: + if isinstance(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}." + ) + + latents = latents / vae.config.scaling_factor + latents = latents.to(vae.dtype).to(device) + + if enable_vae_tiling: + vae.enable_tiling() + video = vae.decode( + latents, return_dict=False, generator=generator + )[0] + else: + video = vae.decode( + latents, return_dict=False, generator=generator + )[0] + + if expand_temporal_dim or video.shape[2] == 1: + video = video.squeeze(2) + + vae.to(offload_device) + mm.soft_empty_cache() + + video_processor = VideoProcessor(vae_scale_factor=8) + video_processor.config.do_resize = False + + video = video_processor.postprocess_video(video=video, output_type="pt") + video = video[0].permute(0, 2, 3, 1).cpu().float() + + return (video,) + +class CogVideoLatentPreview: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT",), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "min_val": ("FLOAT", {"default": -0.15, "min": -1.0, "max": 0.0, "step": 0.001}), + "max_val": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001}), + "r_bias": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001}), + "g_bias": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001}), + "b_bias": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.001}), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", ) + RETURN_NAMES = ("images", "latent_rgb_factors",) + FUNCTION = "sample" + CATEGORY = "PyramidFlowWrapper" + + def sample(self, samples, seed, min_val, max_val, r_bias, g_bias, b_bias): + mm.soft_empty_cache() + + latents = samples["samples"].clone() + print("in sample", latents.shape) + latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] + + #[[0.0658900170023352, 0.04687556512203313, -0.056971557475649186], [-0.01265770449940036, -0.02814809569100843, -0.0768912512529372], [0.061456544746314665, 0.0005511617552452358, -0.0652574975291287], [-0.09020669168815276, -0.004755440180558637, -0.023763970904494294], [0.031766964513999865, -0.030959599938418375, 0.08654669098083616], [-0.005981764690055846, -0.08809119252349802, -0.06439852368217663], [-0.0212114426433989, 0.08894281999597677, 0.05155629477559985], [-0.013947446911030725, -0.08987475069900677, -0.08923124751217484], [-0.08235967967978511, 0.07268025379974379, 0.08830486164536037], [-0.08052049179735378, -0.050116143175332195, 0.02023752569687405], [-0.07607527759162447, 0.06827156419895981, 0.08678111754261035], [-0.04689089232553825, 0.017294986041038893, -0.10280492336438908], [-0.06105783150270304, 0.07311850680875913, 0.019995735372550075], [-0.09232589996527711, -0.012869815059053047, -0.04355587834255975], [-0.06679931010802251, 0.018399815879067458, 0.06802404982033876], [-0.013062632927118165, -0.04292991477896661, 0.07476243356192845]] + latent_rgb_factors =[[0.11945946736445662, 0.09919175788574555, -0.004832707433877734], [-0.0011977028264356232, 0.05496505130267682, 0.021321622433638193], [-0.014088548986590666, -0.008701477861945644, -0.020991313281459367], [0.03063921972519621, 0.12186477097625073, 0.0139593690235148], [0.0927403067854673, 0.030293187650929136, 0.05083134241694003], [0.0379112441305742, 0.04935199882777209, 0.058562766246777774], [0.017749911959153715, 0.008839453404921545, 0.036005638019226294], [0.10610119248526109, 0.02339855688237826, 0.057154257614084596], [0.1273639464837117, -0.010959856130713416, 0.043268631260428896], [-0.01873510946881321, 0.08220930648486932, 0.10613256772247093], [0.008429116376722327, 0.07623856561000408, 0.09295712117576727], [0.12938137079617007, 0.12360403483892413, 0.04478930933220116], [0.04565908794779364, 0.041064156741596365, -0.017695041535528512], [0.00019003240570281826, -0.013965147883381978, 0.05329669529635849], [0.08082391586738358, 0.11548306825496074, -0.021464170006615893], [-0.01517932393230994, -0.0057985555313003236, 0.07216646476618871]] + import random + random.seed(seed) + latent_rgb_factors = [[random.uniform(min_val, max_val) for _ in range(3)] for _ in range(16)] + out_factors = latent_rgb_factors + print(latent_rgb_factors) + + latent_rgb_factors_bias = [0.085, 0.137, 0.158] + #latent_rgb_factors_bias = [r_bias, g_bias, b_bias] + + latent_rgb_factors = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1) + latent_rgb_factors_bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype) + + print("latent_rgb_factors", latent_rgb_factors.shape) + + latent_images = [] + for t in range(latents.shape[2]): + latent = latents[:, :, t, :, :] + latent = latent[0].permute(1, 2, 0) + latent_image = torch.nn.functional.linear( + latent, + latent_rgb_factors, + bias=latent_rgb_factors_bias + ) + latent_images.append(latent_image) + latent_images = torch.stack(latent_images, dim=0) + print("latent_images", latent_images.shape) + latent_images_min = latent_images.min() + latent_images_max = latent_images.max() + latent_images = (latent_images - latent_images_min) / (latent_images_max - latent_images_min) + + return (latent_images.float().cpu(), out_factors) + +NODE_CLASS_MAPPINGS = { + "HyVideoSampler": HyVideoSampler, + "HyVideoDecode": HyVideoDecode, + "HyVideoTextEncode": HyVideoTextEncode, + "HyVideoModelLoader": HyVideoModelLoader, + "HyVideoVAELoader": HyVideoVAELoader, + "DownloadAndLoadHyVideoTextEncoder": DownloadAndLoadHyVideoTextEncoder, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "HyVideoSampler": "HunyuanVideo Sampler", + "HyVideoDecode": "HunyuanVideo Decode", + "HyVideoTextEncode": "HunyuanVideo TextEncode", + "HyVideoModelLoader": "HunyuanVideo Model Loader", + "HyVideoVAELoader": "HunyuanVideo VAE Loader", + "DownloadAndLoadHyVideoTextEncoder": "(Down)Load HunyuanVideo TextEncoder", + } diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..503f090 --- /dev/null +++ b/readme.md @@ -0,0 +1,21 @@ +# ComfyUI wrapper nodes for [HunyuanVideo](https://github.com/Tencent/HunyuanVideo) + +## WORK IN PROGRESS + +Transformer and VAE (single files, no autodownload): + +https://huggingface.co/Kijai/HunyuanVideo_comfy/tree/main + +Go to the usual ComfyUI folders (diffusion_models and vae) + +LLM text encoder (has autodownload): + +https://huggingface.co/Kijai/llava-llama-3-8b-text-encoder-tokenizer + +Files go to `ComfyUI/models/LLM/llava-llama-3-8b-text-encoder-tokenizer` + +Clip text encoder (has autodownload) + +For now using the original https://huggingface.co/openai/clip-vit-large-patch14, files (only need the .safetensor from the weights) go to: + +`ComfyUI/models/clip/clip-vit-large-patch14` \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..309cc6e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +accelerate >= 1.1.1 +diffusers >= 0.31.0 +transformers >= 4.39.3 \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..7d3412c --- /dev/null +++ b/utils.py @@ -0,0 +1,22 @@ +import importlib.metadata +import torch +import logging +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +log = logging.getLogger(__name__) + +def check_diffusers_version(): + try: + version = importlib.metadata.version('diffusers') + required_version = '0.31.0' + if version < required_version: + raise AssertionError(f"diffusers version {version} is installed, but version {required_version} or higher is required.") + except importlib.metadata.PackageNotFoundError: + raise AssertionError("diffusers is not installed.") + +def print_memory(device): + memory = torch.cuda.memory_allocated(device) / 1024**3 + max_memory = torch.cuda.max_memory_allocated(device) / 1024**3 + max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 + log.info(f"Allocated memory: {memory=:.3f} GB") + log.info(f"Max allocated memory: {max_memory=:.3f} GB") + log.info(f"Max reserved memory: {max_reserved=:.3f} GB") \ No newline at end of file