diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/.gitignore b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/.gitignore
new file mode 100644
index 0000000000000000000000000000000000000000..5df6c9c031d19999b16bf4bbd85d8623e9c6a296
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/.gitignore
@@ -0,0 +1,3 @@
+__pycache__
+/ckpts/**/
+/results
\ No newline at end of file
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/config.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6a153dab1648cedf64dff79d303472511874bd0
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/config.py
@@ -0,0 +1,423 @@
+import argparse
+import re
+
+from .constants import *
+from .modules.models import HUNYUAN_VIDEO_CONFIG
+
+
+def parse_args(namespace=None):
+ parser = argparse.ArgumentParser(description="HunyuanVideo inference script")
+
+ parser = add_network_args(parser)
+ parser = add_extra_models_args(parser)
+ parser = add_denoise_schedule_args(parser)
+ parser = add_inference_args(parser)
+ parser = add_parallel_args(parser)
+ parser = add_algorithm_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-path",
+ type=str,
+ required=True,
+ help="Path of 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-path",
+ type=str,
+ required=True,
+ help="Path of 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-2-path",
+ type=str,
+ default="clip-vit-large-patch14",
+ help="Path of text encoder model",
+ )
+ group.add_argument(
+ "--text-encoder-precision-2",
+ type=str,
+ default="fp16",
+ choices=PRECISIONS,
+ help="Precision mode for the second text encoder model.",
+ )
+ group.add_argument(
+ "--text-states-dim-2",
+ type=int,
+ default=768,
+ help="Dimension of the second text encoder hidden states.",
+ )
+ group.add_argument(
+ "--tokenizer-2",
+ type=str,
+ default="clipL",
+ choices=list(TOKENIZER_PATH),
+ help="Name of the second tokenizer model.",
+ )
+ group.add_argument(
+ "--text-len-2",
+ type=int,
+ default=77,
+ help="Maximum length of the second text input.",
+ )
+
+ return parser
+
+
+def add_denoise_schedule_args(parser: argparse.ArgumentParser):
+ group = parser.add_argument_group(title="Denoise schedule args")
+
+ group.add_argument(
+ "--denoise-type",
+ type=str,
+ default="flow",
+ help="Denoise type for noised inputs.",
+ )
+
+ # Flow Matching
+ group.add_argument(
+ "--flow-shift",
+ type=float,
+ default=7.0,
+ help="Shift factor for flow matching schedulers.",
+ )
+ group.add_argument(
+ "--flow-reverse",
+ action="store_true",
+ help="If reverse, learning/sampling from t=1 -> t=0.",
+ )
+ group.add_argument(
+ "--flow-solver",
+ type=str,
+ default="euler",
+ help="Solver for flow matching.",
+ )
+ group.add_argument(
+ "--use-linear-quadratic-schedule",
+ action="store_true",
+ help="Use linear quadratic schedule for flow matching."
+ "Following MovieGen (https://ai.meta.com/static-resource/movie-gen-research-paper)",
+ )
+ group.add_argument(
+ "--linear-schedule-end",
+ type=int,
+ default=25,
+ help="End step for linear quadratic schedule for flow matching.",
+ )
+
+ return parser
+
+
+def add_inference_args(parser: argparse.ArgumentParser):
+ group = parser.add_argument_group(title="Inference args")
+
+ # ======================== Model loads ========================
+ group.add_argument(
+ "--model-base",
+ type=str,
+ default="ckpts",
+ help="Root path of all the models, including t2v models and extra models.",
+ )
+ group.add_argument(
+ "--dit-weight",
+ type=str,
+ default="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
+ help="Path to the HunyuanVideo model. If None, search the model in the args.model_root."
+ "1. If it is a file, load the model directly."
+ "2. If it is a directory, search the model in the directory. Support two types of models: "
+ "1) named `pytorch_model_*.pt`"
+ "2) named `*_model_states.pt`, where * can be `mp_rank_00`.",
+ )
+ group.add_argument(
+ "--model-resolution",
+ type=str,
+ default="540p",
+ choices=["540p", "720p"],
+ help="Root path of all the models, including t2v models and extra models.",
+ )
+ group.add_argument(
+ "--load-key",
+ type=str,
+ default="module",
+ help="Key to load the model states. 'module' for the main model, 'ema' for the EMA model.",
+ )
+ group.add_argument(
+ "--use-cpu-offload",
+ action="store_true",
+ help="Use CPU offload for the model load.",
+ )
+
+ # ======================== Inference general setting ========================
+ group.add_argument(
+ "--batch-size",
+ type=int,
+ default=1,
+ help="Batch size for inference and evaluation.",
+ )
+ group.add_argument(
+ "--infer-steps",
+ type=int,
+ default=50,
+ help="Number of denoising steps for inference.",
+ )
+ group.add_argument(
+ "--disable-autocast",
+ action="store_true",
+ help="Disable autocast for denoising loop and vae decoding in pipeline sampling.",
+ )
+ group.add_argument(
+ "--save-path",
+ type=str,
+ default="./results",
+ help="Path to save the generated samples.",
+ )
+ group.add_argument(
+ "--save-path-suffix",
+ type=str,
+ default="",
+ help="Suffix for the directory of saved samples.",
+ )
+ group.add_argument(
+ "--name-suffix",
+ type=str,
+ default="",
+ help="Suffix for the names of saved samples.",
+ )
+ group.add_argument(
+ "--num-videos",
+ type=int,
+ default=1,
+ help="Number of videos to generate for each prompt.",
+ )
+ # ---sample size---
+ group.add_argument(
+ "--video-size",
+ type=int,
+ nargs="+",
+ default=(720, 1280),
+ help="Video size for training. If a single value is provided, it will be used for both height "
+ "and width. If two values are provided, they will be used for height and width "
+ "respectively.",
+ )
+ group.add_argument(
+ "--video-length",
+ type=int,
+ default=129,
+ help="How many frames to sample from a video. if using 3d vae, the number should be 4n+1",
+ )
+ # --- prompt ---
+ group.add_argument(
+ "--prompt",
+ type=str,
+ default=None,
+ help="Prompt for sampling during evaluation.",
+ )
+ group.add_argument(
+ "--seed-type",
+ type=str,
+ default="auto",
+ choices=["file", "random", "fixed", "auto"],
+ help="Seed type for evaluation. If file, use the seed from the CSV file. If random, generate a "
+ "random seed. If fixed, use the fixed seed given by `--seed`. If auto, `csv` will use the "
+ "seed column if available, otherwise use the fixed `seed` value. `prompt` will use the "
+ "fixed `seed` value.",
+ )
+ group.add_argument("--seed", type=int, default=None, help="Seed for evaluation.")
+
+ # Classifier-Free Guidance
+ group.add_argument(
+ "--neg-prompt", type=str, default=None, help="Negative prompt for sampling."
+ )
+ group.add_argument(
+ "--cfg-scale", type=float, default=1.0, help="Classifier free guidance scale."
+ )
+ group.add_argument(
+ "--embedded-cfg-scale",
+ type=float,
+ default=6.0,
+ help="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 add_parallel_args(parser: argparse.ArgumentParser):
+ group = parser.add_argument_group(title="Parallel args")
+
+ group.add_argument(
+ "--use_cfg_parallel",
+ action='store_true',
+ help="Use CFG parallel.",
+ )
+ group.add_argument(
+ "--sp-degree",
+ type=int,
+ default=1,
+ help="Seqence parallel degree.",
+ )
+
+ return parser
+
+
+def sanity_check_args(args):
+ # VAE channels
+ vae_pattern = r"\d{2,3}-\d{1,2}c-\w+"
+ if not re.match(vae_pattern, args.vae):
+ raise ValueError(
+ f"Invalid VAE model: {args.vae}. Must be in the format of '{vae_pattern}'."
+ )
+ vae_channels = int(args.vae.split("-")[1][:-1])
+ if args.latent_channels is None:
+ args.latent_channels = vae_channels
+ if vae_channels != args.latent_channels:
+ raise ValueError(
+ f"Latent channels ({args.latent_channels}) must match the VAE channels ({vae_channels})."
+ )
+ return args
+
+
+def add_algorithm_args(parser: argparse.ArgumentParser):
+ group = parser.add_argument_group(title="Algorithm args")
+
+ group.add_argument(
+ "--algorithm",
+ type=str,
+ default=None,
+ choices=[None, 'dit_cache', 'attention_cache'],
+ help="Choose inference algorithm.",
+ )
+
+ return parser
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/constants.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..b8c31b7cfd12f1fd90f4099d6082893452857c65
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/constants.py
@@ -0,0 +1,90 @@
+import os
+import torch
+
+__all__ = [
+ "C_SCALE",
+ "PROMPT_TEMPLATE",
+ "MODEL_BASE",
+ "PRECISIONS",
+ "NORMALIZATION_TYPE",
+ "ACTIVATION_TYPE",
+ "VAE_PATH",
+ "TEXT_ENCODER_PATH",
+ "TOKENIZER_PATH",
+ "TEXT_PROJECTION",
+ "DATA_TYPE",
+ "NEGATIVE_PROMPT",
+]
+
+PRECISION_TO_TYPE = {
+ 'fp32': torch.float32,
+ 'fp16': torch.float16,
+ 'bf16': torch.bfloat16,
+}
+
+# =================== Constant Values =====================
+# Computation scale factor, 1P = 1_000_000_000_000_000. Tensorboard will display the value in PetaFLOPS to avoid
+# overflow error when tensorboard logging values.
+C_SCALE = 1_000_000_000_000_000
+
+# When using decoder-only models, we must provide a prompt template to instruct the text encoder
+# on how to generate the text.
+# --------------------------------------------------------------------
+PROMPT_TEMPLATE_ENCODE = (
+ "<|start_header_id|>system<|end_header_id|>\n\nDescribe the image by detailing the color, shape, size, texture, "
+ "quantity, text, spatial relationships of the objects and background:<|eot_id|>"
+ "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
+)
+PROMPT_TEMPLATE_ENCODE_VIDEO = (
+ "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
+ "1. The main content and theme of the video."
+ "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
+ "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
+ "4. background environment, light, style and atmosphere."
+ "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
+ "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
+)
+
+NEGATIVE_PROMPT = "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion"
+
+PROMPT_TEMPLATE = {
+ "dit-llm-encode": {
+ "template": PROMPT_TEMPLATE_ENCODE,
+ "crop_start": 36,
+ },
+ "dit-llm-encode-video": {
+ "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
+ "crop_start": 95,
+ },
+}
+
+# ======================= Model ======================
+PRECISIONS = {"fp32", "fp16", "bf16"}
+NORMALIZATION_TYPE = {"layer", "rms"}
+ACTIVATION_TYPE = {"relu", "silu", "gelu", "gelu_tanh"}
+
+# =================== Model Path =====================
+MODEL_BASE = os.getenv("MODEL_BASE", "./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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2141aa3dccb5a6b231bf2f3ae6ab864152ffc3ec
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/__init__.py
@@ -0,0 +1,2 @@
+from .pipelines import HunyuanVideoPipeline
+from .schedulers import FlowMatchDiscreteScheduler
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/pipelines/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/pipelines/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e44cb6196fe7f6a7fa821b2bebddf8d1117521fc
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/pipelines/__init__.py
@@ -0,0 +1 @@
+from .pipeline_hunyuan_video import HunyuanVideoPipeline
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
new file mode 100644
index 0000000000000000000000000000000000000000..b704bae753e32aceaab8f28c0fac9f726d9c24b6
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
@@ -0,0 +1,1139 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+
+import numpy as np
+import torch
+import torch_npu
+import torch.distributed as dist
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import BaseOutput
+from diffusers.utils import (
+ USE_PEFT_BACKEND,
+ deprecate,
+ logging,
+ replace_example_docstring,
+ scale_lora_layers,
+ unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+from ...constants import PRECISION_TO_TYPE
+from ...modules import HYVideoDiffusionTransformer
+from ...text_encoder import TextEncoder
+from ...utils.parallel_mgr import (
+ get_cfg_group,
+ get_classifier_free_guidance_world_size,
+ get_classifier_free_guidance_rank
+)
+from ...vae.autoencoder_kl_causal_3d import AutoencoderKLCausal3D
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """"""
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+ """
+ Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+ """
+ std_text = noise_pred_text.std(
+ dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+ )
+ std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+ # rescale the results from guidance (fixes overexposure)
+ noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+ # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+ noise_cfg = (
+ guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+ )
+ return noise_cfg
+
+
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ sigmas: Optional[List[float]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError(
+ "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+ )
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(
+ inspect.signature(scheduler.set_timesteps).parameters.keys()
+ )
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(
+ inspect.signature(scheduler.set_timesteps).parameters.keys()
+ )
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+@dataclass
+class HunyuanVideoPipelineOutput(BaseOutput):
+ videos: Union[torch.Tensor, np.ndarray]
+
+
+class HunyuanVideoPipeline(DiffusionPipeline):
+ r"""
+ Pipeline for text-to-video generation using HunyuanVideo.
+
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+ implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+ Args:
+ vae ([`AutoencoderKL`]):
+ Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+ text_encoder ([`TextEncoder`]):
+ Frozen text-encoder.
+ text_encoder_2 ([`TextEncoder`]):
+ Frozen text-encoder_2.
+ transformer ([`HYVideoDiffusionTransformer`]):
+ A `HYVideoDiffusionTransformer` to denoise the encoded video latents.
+ scheduler ([`SchedulerMixin`]):
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+ """
+
+ model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+ _optional_components = ["text_encoder_2"]
+ _exclude_from_cpu_offload = ["transformer"]
+ _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+ def __init__(
+ self,
+ vae: AutoencoderKL,
+ text_encoder: TextEncoder,
+ transformer: HYVideoDiffusionTransformer,
+ scheduler: KarrasDiffusionSchedulers,
+ text_encoder_2: Optional[TextEncoder] = None,
+ progress_bar_config: Dict[str, Any] = None,
+ args=None,
+ ):
+ super().__init__()
+
+ # ==========================================================================================
+ if progress_bar_config is None:
+ progress_bar_config = {}
+ if not hasattr(self, "_progress_bar_config"):
+ self._progress_bar_config = {}
+ self._progress_bar_config.update(progress_bar_config)
+
+ self.args = args
+ # ==========================================================================================
+
+ if (
+ hasattr(scheduler.config, "steps_offset")
+ and scheduler.config.steps_offset != 1
+ ):
+ deprecation_message = (
+ f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+ f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+ "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+ " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+ " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+ " file"
+ )
+ deprecate(
+ "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
+ )
+ new_config = dict(scheduler.config)
+ new_config["steps_offset"] = 1
+ scheduler._internal_dict = FrozenDict(new_config)
+
+ if (
+ hasattr(scheduler.config, "clip_sample")
+ and scheduler.config.clip_sample is True
+ ):
+ deprecation_message = (
+ f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+ " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+ " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+ " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+ " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+ )
+ deprecate(
+ "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
+ )
+ new_config = dict(scheduler.config)
+ new_config["clip_sample"] = False
+ scheduler._internal_dict = FrozenDict(new_config)
+
+ self.register_modules(
+ vae=vae,
+ text_encoder=text_encoder,
+ transformer=transformer,
+ scheduler=scheduler,
+ text_encoder_2=text_encoder_2,
+ )
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+ def encode_prompt(
+ self,
+ prompt,
+ device,
+ num_videos_per_prompt,
+ do_classifier_free_guidance,
+ negative_prompt=None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_attention_mask: Optional[torch.Tensor] = None,
+ lora_scale: Optional[float] = None,
+ clip_skip: Optional[int] = None,
+ text_encoder: Optional[TextEncoder] = None,
+ data_type: Optional[str] = "image",
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ device: (`torch.device`):
+ torch device
+ num_videos_per_prompt (`int`):
+ number of videos that should be generated per prompt
+ do_classifier_free_guidance (`bool`):
+ whether to use classifier free guidance or not
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the video generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ attention_mask (`torch.Tensor`, *optional*):
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ negative_attention_mask (`torch.Tensor`, *optional*):
+ lora_scale (`float`, *optional*):
+ A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+ clip_skip (`int`, *optional*):
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+ the output of the pre-final layer will be used for computing the prompt embeddings.
+ text_encoder (TextEncoder, *optional*):
+ data_type (`str`, *optional*):
+ """
+ if text_encoder is None:
+ text_encoder = self.text_encoder
+
+ # set lora scale so that monkey patched LoRA
+ # function of text encoder can correctly access it
+ if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+ self._lora_scale = lora_scale
+
+ # dynamically adjust the LoRA scale
+ if not USE_PEFT_BACKEND:
+ adjust_lora_scale_text_encoder(text_encoder.model, lora_scale)
+ else:
+ scale_lora_layers(text_encoder.model, lora_scale)
+
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ # textual inversion: process multi-vector tokens if necessary
+ if isinstance(self, TextualInversionLoaderMixin):
+ prompt = self.maybe_convert_prompt(prompt, text_encoder.tokenizer)
+
+ text_inputs = text_encoder.text2tokens(prompt, data_type=data_type)
+
+ if clip_skip is None:
+ prompt_outputs = text_encoder.encode(
+ text_inputs, data_type=data_type, device=device
+ )
+ prompt_embeds = prompt_outputs.hidden_state
+ else:
+ prompt_outputs = text_encoder.encode(
+ text_inputs,
+ output_hidden_states=True,
+ data_type=data_type,
+ device=device,
+ )
+ # Access the `hidden_states` first, that contains a tuple of
+ # all the hidden states from the encoder layers. Then index into
+ # the tuple to access the hidden states from the desired layer.
+ prompt_embeds = prompt_outputs.hidden_states_list[-(clip_skip + 1)]
+ # We also need to apply the final LayerNorm here to not mess with the
+ # representations. The `last_hidden_states` that we typically use for
+ # obtaining the final prompt representations passes through the LayerNorm
+ # layer.
+ prompt_embeds = text_encoder.model.text_model.final_layer_norm(
+ prompt_embeds
+ )
+
+ attention_mask = prompt_outputs.attention_mask
+ if attention_mask is not None:
+ attention_mask = attention_mask.to(device)
+ bs_embed, seq_len = attention_mask.shape
+ attention_mask = attention_mask.repeat(1, num_videos_per_prompt)
+ attention_mask = attention_mask.view(
+ bs_embed * num_videos_per_prompt, seq_len
+ )
+
+ if text_encoder is not None:
+ prompt_embeds_dtype = text_encoder.dtype
+ elif self.transformer is not None:
+ prompt_embeds_dtype = self.transformer.dtype
+ else:
+ prompt_embeds_dtype = prompt_embeds.dtype
+
+ prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+ if prompt_embeds.ndim == 2:
+ bs_embed, _ = prompt_embeds.shape
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt)
+ prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, -1)
+ else:
+ bs_embed, seq_len, _ = prompt_embeds.shape
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(
+ bs_embed * num_videos_per_prompt, seq_len, -1
+ )
+
+ # get unconditional embeddings for classifier free guidance
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
+ 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
+
+ # textual inversion: process multi-vector tokens if necessary
+ if isinstance(self, TextualInversionLoaderMixin):
+ uncond_tokens = self.maybe_convert_prompt(
+ uncond_tokens, text_encoder.tokenizer
+ )
+
+ 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
+ )
+
+ if text_encoder is not None:
+ if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(text_encoder.model, lora_scale)
+
+ return (
+ prompt_embeds,
+ negative_prompt_embeds,
+ attention_mask,
+ negative_attention_mask,
+ )
+
+ def decode_latents(self, latents, enable_tiling=True):
+ deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+ deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+ latents = 1 / self.vae.config.scaling_factor * latents
+ if enable_tiling:
+ self.vae.enable_tiling()
+ image = self.vae.decode(latents, return_dict=False)[0]
+ else:
+ image = self.vae.decode(latents, return_dict=False)[0]
+ image = (image / 2 + 0.5).clamp(0, 1)
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+ if image.ndim == 4:
+ image = image.cpu().permute(0, 2, 3, 1).float()
+ else:
+ image = image.cpu().float()
+ return image
+
+ def prepare_extra_func_kwargs(self, func, kwargs):
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+ # and should be between [0, 1]
+ extra_step_kwargs = {}
+
+ for k, v in kwargs.items():
+ accepts = k in set(inspect.signature(func).parameters.keys())
+ if accepts:
+ extra_step_kwargs[k] = v
+ return extra_step_kwargs
+
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ video_length,
+ callback_steps,
+ negative_prompt=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ callback_on_step_end_tensor_inputs=None,
+ vae_ver="88-4c-sd",
+ ):
+ if height % 8 != 0 or width % 8 != 0:
+ raise ValueError(
+ f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
+ )
+
+ if video_length is not None:
+ if "884" in vae_ver:
+ if video_length != 1 and (video_length - 1) % 4 != 0:
+ raise ValueError(
+ f"`video_length` has to be 1 or a multiple of 4 but is {video_length}."
+ )
+ elif "888" in vae_ver:
+ if video_length != 1 and (video_length - 1) % 8 != 0:
+ raise ValueError(
+ f"`video_length` has to be 1 or a multiple of 8 but is {video_length}."
+ )
+
+ if callback_steps is not None and (
+ not isinstance(callback_steps, int) or callback_steps <= 0
+ ):
+ raise ValueError(
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+ f" {type(callback_steps)}."
+ )
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs
+ for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (
+ not isinstance(prompt, str) and not isinstance(prompt, list)
+ ):
+ raise ValueError(
+ f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
+ )
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {negative_prompt_embeds.shape}."
+ )
+
+ def prepare_latents(
+ self,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ video_length,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ shape = (
+ batch_size,
+ num_channels_latents,
+ video_length,
+ int(height) // self.vae_scale_factor,
+ int(width) // self.vae_scale_factor,
+ )
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ if latents is None:
+ latents = randn_tensor(
+ shape, generator=generator, device=device, dtype=dtype
+ )
+ else:
+ latents = latents.to(device)
+
+ # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler
+ if hasattr(self.scheduler, "init_noise_sigma"):
+ # scale the initial noise by the standard deviation required by the scheduler
+ latents = latents * self.scheduler.init_noise_sigma
+ return latents
+
+ # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+ def get_guidance_scale_embedding(
+ self,
+ w: torch.Tensor,
+ embedding_dim: int = 512,
+ dtype: torch.dtype = torch.float32,
+ ) -> torch.Tensor:
+ """
+ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+ Args:
+ w (`torch.Tensor`):
+ Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+ embedding_dim (`int`, *optional*, defaults to 512):
+ Dimension of the embeddings to generate.
+ dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+ Data type of the generated embeddings.
+
+ Returns:
+ `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
+ """
+ assert len(w.shape) == 1
+ w = w * 1000.0
+
+ half_dim = embedding_dim // 2
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+ emb = w.to(dtype)[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1))
+ assert emb.shape == (w.shape[0], embedding_dim)
+ return emb
+
+ def process_cfg_split_batch(
+ self,
+ negative_embeds: torch.Tensor,
+ embeds: torch.Tensor,
+ negative_embdes_mask: torch.Tensor = None,
+ embeds_mask: torch.Tensor = None,
+ ):
+ if get_classifier_free_guidance_world_size() == 1:
+ embeds = torch.cat([negative_embeds, embeds], dim=0)
+ elif get_classifier_free_guidance_rank() == 0:
+ embeds = negative_embeds
+ elif get_classifier_free_guidance_rank() == 1:
+ embeds = embeds
+ else:
+ raise ValueError("Invalid classifier free guidance rank")
+
+ if negative_embdes_mask is None:
+ return embeds
+
+ if get_classifier_free_guidance_world_size() == 1:
+ embeds_mask = torch.cat([negative_embdes_mask, embeds_mask], dim=0)
+ elif get_classifier_free_guidance_rank() == 0:
+ embeds_mask = negative_embdes_mask
+ elif get_classifier_free_guidance_rank() == 1:
+ embeds_mask = embeds_mask
+ else:
+ raise ValueError("Invalid classifier free guidance rank")
+ return embeds, embeds_mask
+
+ @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
+
+ @property
+ def cross_attention_kwargs(self):
+ return self._cross_attention_kwargs
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Union[str, List[str]],
+ height: int,
+ width: int,
+ video_length: int,
+ data_type: str = "video",
+ num_inference_steps: int = 50,
+ timesteps: List[int] = None,
+ sigmas: List[float] = None,
+ guidance_scale: float = 7.5,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ num_videos_per_prompt: Optional[int] = 1,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_attention_mask: Optional[torch.Tensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ guidance_rescale: float = 0.0,
+ clip_skip: Optional[int] = None,
+ callback_on_step_end: Optional[
+ Union[
+ Callable[[int, int, Dict], None],
+ PipelineCallback,
+ MultiPipelineCallbacks,
+ ]
+ ] = None,
+ callback_on_step_end_tensor_inputs: List[str] = None,
+ freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
+ vae_ver: str = "88-4c-sd",
+ enable_tiling: bool = False,
+ n_tokens: Optional[int] = None,
+ embedded_guidance_scale: Optional[float] = None,
+ **kwargs,
+ ):
+ r"""
+ The call function to the pipeline for generation.
+
+ Args:
+ prompt (`str` or `List[str]`):
+ The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+ height (`int`):
+ The height in pixels of the generated image.
+ width (`int`):
+ The width in pixels of the generated image.
+ video_length (`int`):
+ The number of frames in the generated video.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+ expense of slower inference.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+ passed will be used. Must be in descending order.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+ their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+ will be used.
+ guidance_scale (`float`, *optional*, defaults to 7.5):
+ A higher guidance scale value encourages the model to generate images closely linked to the text
+ `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+ pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ eta (`float`, *optional*, defaults to 0.0):
+ Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+ to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+ generation deterministic.
+ latents (`torch.Tensor`, *optional*):
+ Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+ tensor is generated by sampling using the supplied random `generator`.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+ provided, text embeddings are generated from the `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+ not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`HunyuanVideoPipelineOutput`] instead of a
+ plain tuple.
+ cross_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+ [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ guidance_rescale (`float`, *optional*, defaults to 0.0):
+ Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
+ Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
+ using zero terminal SNR.
+ clip_skip (`int`, *optional*):
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+ the output of the pre-final layer will be used for computing the prompt embeddings.
+ callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+ A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+ each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
+ DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
+ list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+ `._callback_tensor_inputs` attribute of your pipeline class.
+
+ Examples:
+
+ Returns:
+ [`~HunyuanVideoPipelineOutput`] or `tuple`:
+ If `return_dict` is `True`, [`HunyuanVideoPipelineOutput`] is returned,
+ otherwise a `tuple` is returned where the first element is a list with the generated images and the
+ second element is a list of `bool`s indicating whether the corresponding generated image contains
+ "not-safe-for-work" (nsfw) content.
+ """
+ callback = kwargs.pop("callback", None)
+ callback_steps = kwargs.pop("callback_steps", None)
+ if callback_on_step_end_tensor_inputs is None:
+ callback_on_step_end_tensor_inputs = ["latents"]
+
+ if callback is not None:
+ deprecate(
+ "callback",
+ "1.0.0",
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+ )
+ if callback_steps is not None:
+ deprecate(
+ "callback_steps",
+ "1.0.0",
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+ )
+
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ video_length,
+ callback_steps,
+ negative_prompt,
+ prompt_embeds,
+ negative_prompt_embeds,
+ callback_on_step_end_tensor_inputs,
+ vae_ver=vae_ver,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._guidance_rescale = guidance_rescale
+ self._clip_skip = clip_skip
+ self._cross_attention_kwargs = cross_attention_kwargs
+ self._interrupt = False
+
+ # 2. Define call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = torch.device(f"npu:{dist.get_rank()}") if dist.is_initialized() else self._execution_device
+
+ # 3. Encode input prompt
+ lora_scale = (
+ self.cross_attention_kwargs.get("scale", None)
+ if self.cross_attention_kwargs is not None
+ else None
+ )
+
+ (
+ prompt_embeds,
+ negative_prompt_embeds,
+ prompt_mask,
+ negative_prompt_mask,
+ ) = self.encode_prompt(
+ prompt,
+ device,
+ num_videos_per_prompt,
+ self.do_classifier_free_guidance,
+ negative_prompt,
+ prompt_embeds=prompt_embeds,
+ attention_mask=attention_mask,
+ negative_prompt_embeds=negative_prompt_embeds,
+ negative_attention_mask=negative_attention_mask,
+ lora_scale=lora_scale,
+ clip_skip=self.clip_skip,
+ data_type=data_type,
+ )
+ if self.text_encoder_2 is not None:
+ (
+ prompt_embeds_2,
+ negative_prompt_embeds_2,
+ prompt_mask_2,
+ negative_prompt_mask_2,
+ ) = self.encode_prompt(
+ prompt,
+ device,
+ num_videos_per_prompt,
+ self.do_classifier_free_guidance,
+ negative_prompt,
+ prompt_embeds=None,
+ attention_mask=None,
+ negative_prompt_embeds=None,
+ negative_attention_mask=None,
+ lora_scale=lora_scale,
+ clip_skip=self.clip_skip,
+ text_encoder=self.text_encoder_2,
+ data_type=data_type,
+ )
+ else:
+ prompt_embeds_2 = None
+ negative_prompt_embeds_2 = None
+ prompt_mask_2 = None
+ negative_prompt_mask_2 = None
+
+ # For classifier free guidance, we need to do two forward passes.
+ # Here we concatenate the unconditional and text embeddings into a single batch
+ # to avoid doing two forward passes
+ if self.do_classifier_free_guidance:
+ if prompt_mask is not None:
+ prompt_embeds, prompt_mask = self.process_cfg_split_batch(
+ negative_prompt_embeds,
+ prompt_embeds,
+ negative_prompt_mask,
+ prompt_mask
+ )
+ else:
+ prompt_embeds = self.process_cfg_split_batch(negative_prompt_embeds, prompt_embeds)
+ if prompt_mask_2 is not None:
+ prompt_embeds_2, prompt_mask_2 = self.process_cfg_split_batch(
+ negative_prompt_embeds_2,
+ prompt_embeds_2,
+ negative_prompt_mask_2,
+ prompt_mask_2
+ )
+ else:
+ prompt_embeds_2 = self.process_cfg_split_batch(negative_prompt_embeds_2, prompt_embeds_2)
+
+ # 4. Prepare timesteps
+ extra_set_timesteps_kwargs = self.prepare_extra_func_kwargs(
+ self.scheduler.set_timesteps, {"n_tokens": n_tokens}
+ )
+ timesteps, num_inference_steps = retrieve_timesteps(
+ self.scheduler,
+ num_inference_steps,
+ device,
+ timesteps,
+ sigmas,
+ **extra_set_timesteps_kwargs,
+ )
+
+ if "884" in vae_ver:
+ video_length = (video_length - 1) // 4 + 1
+ elif "888" in vae_ver:
+ video_length = (video_length - 1) // 8 + 1
+ else:
+ video_length = video_length
+
+ target_dtype = PRECISION_TO_TYPE[self.args.precision]
+ prompt_embeds = prompt_embeds.to(target_dtype)
+ prompt_embeds_2 = prompt_embeds_2.to(target_dtype)
+ freqs_cos = freqs_cis[0].to(target_dtype).to(device)
+ freqs_sin = freqs_cis[1].to(target_dtype).to(device)
+
+ # 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,
+ target_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:
+ 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 // get_classifier_free_guidance_world_size()))
+ if self.do_classifier_free_guidance
+ else latents
+ )
+ latent_model_input = self.scheduler.scale_model_input(
+ latent_model_input, t
+ )
+
+ t_expand = t.repeat(latent_model_input.shape[0])
+ guidance_expand = (
+ torch.tensor(
+ [embedded_guidance_scale] * latent_model_input.shape[0],
+ dtype=torch.float32,
+ device=device,
+ ).to(target_dtype)
+ * 1000.0
+ if embedded_guidance_scale is not None
+ else None
+ )
+
+ # predict the noise residual
+ 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_cos, # [seqlen, head_dim]
+ freqs_sin=freqs_sin, # [seqlen, head_dim]
+ guidance=guidance_expand,
+ t_idx=i,
+ return_dict=True,
+ )[
+ "x"
+ ]
+
+ # perform guidance
+ if self.do_classifier_free_guidance:
+ if get_classifier_free_guidance_world_size() == 1:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ elif get_classifier_free_guidance_world_size() == 2:
+ noise_pred_uncond, noise_pred_text = get_cfg_group().all_gather(
+ noise_pred, separate_tensors=True
+ )
+ else:
+ raise ValueError("Invalid classifier free guidance world size")
+ noise_pred = noise_pred_uncond + self.guidance_scale * (
+ noise_pred_text - noise_pred_uncond
+ )
+
+ if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+ # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+ noise_pred = rescale_noise_cfg(
+ noise_pred,
+ noise_pred_text,
+ guidance_rescale=self.guidance_rescale,
+ )
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents = self.scheduler.step(
+ noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+ )[0]
+
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+ negative_prompt_embeds = callback_outputs.pop(
+ "negative_prompt_embeds", negative_prompt_embeds
+ )
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or (
+ (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+ ):
+ if progress_bar is not None:
+ progress_bar.update()
+ if callback is not None and i % callback_steps == 0:
+ step_idx = i // getattr(self.scheduler, "order", 1)
+ callback(step_idx, t, latents)
+
+ if not output_type == "latent":
+ expand_temporal_dim = False
+ if len(latents.shape) == 4:
+ if isinstance(self.vae, AutoencoderKLCausal3D):
+ latents = latents.unsqueeze(2)
+ expand_temporal_dim = True
+ elif len(latents.shape) == 5:
+ pass
+ else:
+ raise ValueError(
+ f"Only support latents with shape (b, c, h, w) or (b, c, f, h, w), but got {latents.shape}."
+ )
+
+ if (
+ hasattr(self.vae.config, "shift_factor")
+ and self.vae.config.shift_factor
+ ):
+ latents = (
+ latents / self.vae.config.scaling_factor
+ + self.vae.config.shift_factor
+ )
+ else:
+ latents = latents / self.vae.config.scaling_factor
+
+ vae_dtype = PRECISION_TO_TYPE[self.args.vae_precision]
+ latents = latents.to(vae_dtype)
+ if enable_tiling:
+ self.vae.enable_tiling()
+ image = self.vae.decode(
+ latents, return_dict=False, generator=generator
+ )[0]
+ else:
+ image = self.vae.decode(
+ latents, return_dict=False, generator=generator
+ )[0]
+
+ if expand_temporal_dim or image.shape[2] == 1:
+ image = image.squeeze(2)
+
+ else:
+ image = latents
+
+ image = (image / 2 + 0.5).clamp(0, 1)
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+ image = image.cpu().float()
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return image
+
+ return HunyuanVideoPipelineOutput(videos=image)
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/schedulers/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/schedulers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..14f2ba33feb0a1a802a9a86818781a2a15140bd6
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/schedulers/__init__.py
@@ -0,0 +1 @@
+from .scheduling_flow_match_discrete import FlowMatchDiscreteScheduler
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ba432028e3e7a52debc4fb5188d74712813d447
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/diffusion/schedulers/scheduling_flow_match_discrete.py
@@ -0,0 +1,256 @@
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, logging
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchDiscreteSchedulerOutput(BaseOutput):
+ """
+ Output class for the scheduler's `step` function output.
+
+ Args:
+ prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+ Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+ denoising loop.
+ """
+
+ prev_sample: torch.FloatTensor
+
+
+class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin):
+ """
+ Euler scheduler.
+
+ This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+ methods the library implements for all schedulers such as loading and saving.
+
+ Args:
+ num_train_timesteps (`int`, defaults to 1000):
+ The number of diffusion steps to train the model.
+ timestep_spacing (`str`, defaults to `"linspace"`):
+ The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+ shift (`float`, defaults to 1.0):
+ The shift value for the timestep schedule.
+ reverse (`bool`, defaults to `True`):
+ Whether to reverse the timestep schedule.
+ """
+
+ _compatibles = []
+ order = 1
+
+ @register_to_config
+ def __init__(
+ self,
+ num_train_timesteps: int = 1000,
+ shift: float = 1.0,
+ reverse: bool = True,
+ solver: str = "euler",
+ n_tokens: Optional[int] = None,
+ ):
+ sigmas = torch.linspace(1, 0, num_train_timesteps + 1)
+
+ if not reverse:
+ sigmas = sigmas.flip(0)
+
+ self.sigmas = sigmas
+ # the value fed to model
+ self.timesteps = (sigmas[:-1] * num_train_timesteps).to(dtype=torch.float32)
+
+ self._step_index = None
+ self._begin_index = None
+
+ self.supported_solver = ["euler"]
+ if solver not in self.supported_solver:
+ raise ValueError(
+ f"Solver {solver} not supported. Supported solvers: {self.supported_solver}"
+ )
+
+ @property
+ def step_index(self):
+ """
+ The index counter for current timestep. It will increase 1 after each scheduler step.
+ """
+ return self._step_index
+
+ @property
+ def begin_index(self):
+ """
+ The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+ """
+ return self._begin_index
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+ def set_begin_index(self, begin_index: int = 0):
+ """
+ Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+ Args:
+ begin_index (`int`):
+ The begin index for the scheduler.
+ """
+ self._begin_index = begin_index
+
+ def _sigma_to_t(self, sigma):
+ return sigma * self.config.num_train_timesteps
+
+ def set_timesteps(
+ self,
+ num_inference_steps: int,
+ device: Union[str, torch.device] = None,
+ n_tokens: int = None,
+ ):
+ """
+ Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+ Args:
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ n_tokens (`int`, *optional*):
+ Number of tokens in the input sequence.
+ """
+ self.num_inference_steps = num_inference_steps
+
+ sigmas = torch.linspace(1, 0, num_inference_steps + 1)
+ sigmas = self.sd3_time_shift(sigmas)
+
+ if not self.config.reverse:
+ sigmas = 1 - sigmas
+
+ self.sigmas = sigmas
+ self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to(
+ dtype=torch.float32, device=device
+ )
+
+ # Reset step index
+ self._step_index = None
+
+ def index_for_timestep(self, timestep, schedule_timesteps=None):
+ if schedule_timesteps is None:
+ schedule_timesteps = self.timesteps
+
+ indices = (schedule_timesteps == timestep).nonzero()
+
+ # The sigma index that is taken for the **very** first `step`
+ # is always the second index (or the last index if there is only 1)
+ # This way we can ensure we don't accidentally skip a sigma in
+ # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+ pos = 1 if len(indices) > 1 else 0
+
+ return indices[pos].item()
+
+ def _init_step_index(self, timestep):
+ if self.begin_index is None:
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.to(self.timesteps.device)
+ self._step_index = self.index_for_timestep(timestep)
+ else:
+ self._step_index = self._begin_index
+
+ def scale_model_input(
+ self, sample: torch.Tensor, timestep: Optional[int] = None
+ ) -> torch.Tensor:
+ return sample
+
+ def sd3_time_shift(self, t: torch.Tensor):
+ return (self.config.shift * t) / (1 + (self.config.shift - 1) * t)
+
+ def step(
+ self,
+ model_output: torch.FloatTensor,
+ timestep: Union[float, torch.FloatTensor],
+ sample: torch.FloatTensor,
+ return_dict: bool = True,
+ ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]:
+ """
+ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+ process from the learned model outputs (most often the predicted noise).
+
+ Args:
+ model_output (`torch.FloatTensor`):
+ The direct output from learned diffusion model.
+ timestep (`float`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.FloatTensor`):
+ A current instance of a sample created by the diffusion process.
+ generator (`torch.Generator`, *optional*):
+ A random number generator.
+ n_tokens (`int`, *optional*):
+ Number of tokens in the input sequence.
+ return_dict (`bool`):
+ Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
+ tuple.
+
+ Returns:
+ [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
+ If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
+ returned, otherwise a tuple is returned where the first element is the sample tensor.
+ """
+
+ if (
+ isinstance(timestep, int)
+ or isinstance(timestep, torch.IntTensor)
+ or isinstance(timestep, torch.LongTensor)
+ ):
+ raise ValueError(
+ (
+ "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+ " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+ " one of the `scheduler.timesteps` as a timestep."
+ ),
+ )
+
+ if self.step_index is None:
+ self._init_step_index(timestep)
+
+ # Upcast to avoid precision issues when computing prev_sample
+ dtype = sample.dtype
+ 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
+ prev_sample = prev_sample.to(dtype)
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/inference.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/inference.py
new file mode 100644
index 0000000000000000000000000000000000000000..e422bcaf27e9c0e4cbf07dc180807f9f0c59aa20
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/inference.py
@@ -0,0 +1,588 @@
+import os
+import random
+import time
+import math
+from pathlib import Path
+from typing import List
+
+import torch
+import torch_npu
+from loguru import logger
+
+from hyvideo.constants import PROMPT_TEMPLATE, NEGATIVE_PROMPT, PRECISION_TO_TYPE
+from hyvideo.diffusion.pipelines import HunyuanVideoPipeline
+from hyvideo.diffusion.schedulers import FlowMatchDiscreteScheduler
+from hyvideo.modules import load_model
+from hyvideo.modules.posemb_layers import get_nd_rotary_pos_embed
+from hyvideo.text_encoder import TextEncoder
+from hyvideo.utils.data_utils import align_to
+from hyvideo.utils.parallel_mgr import (
+ ParallelConfig,
+ init_parallel_env,
+ finalize_parallel_env,
+)
+from hyvideo.vae import load_vae
+
+
+class Inference(object):
+ def __init__(
+ self,
+ args,
+ vae,
+ vae_kwargs,
+ text_encoder,
+ model,
+ text_encoder_2=None,
+ pipeline=None,
+ use_cpu_offload=False,
+ device=None,
+ logger=None,
+ parallel_args=None,
+ ):
+ self.vae = vae
+ self.vae_kwargs = vae_kwargs
+
+ self.text_encoder = text_encoder
+ self.text_encoder_2 = text_encoder_2
+
+ self.model = model
+ self.pipeline = pipeline
+ self.use_cpu_offload = use_cpu_offload
+
+ self.args = args
+ self.device = device if device is not None else "npu"
+
+ self.logger = logger
+ self.parallel_args = parallel_args
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_path, args, device=None, **kwargs):
+ """
+ Initialize the Inference pipeline.
+
+ Args:
+ pretrained_model_path (str or pathlib.Path): The model path, including t2v, text encoder and vae checkpoints.
+ args (argparse.Namespace): The arguments for the pipeline.
+ device (int): The device for inference. Default is 0.
+ """
+ # ========================================================================
+ logger.info(f"Got text-to-video model root path: {pretrained_model_path}")
+
+ # ==================== Initialize Distributed Environment ================
+ world_size = int(os.getenv('WORLD_SIZE', 1))
+ if world_size > 1:
+ sp_degree = args.sp_degree
+ tp_degree = world_size // sp_degree // 2 if args.use_cfg_parallel else world_size // sp_degree
+ parallel_degree = sp_degree * tp_degree * 2 if args.use_cfg_parallel else sp_degree * tp_degree
+ assert parallel_degree == world_size, "Parallel degree should be equal to world size"
+ parallel_config = ParallelConfig(sp_degree=sp_degree,
+ tp_degree=tp_degree,
+ use_cfg_parallel=args.use_cfg_parallel,
+ world_size=world_size)
+ init_parallel_env(parallel_config)
+ else:
+ if device is None:
+ device = "npu" if torch_npu.npu.is_available() else "cpu"
+
+ parallel_args = {"use_cfg_parallel": args.use_cfg_parallel}
+
+ # ======================== Get the args path =============================
+
+ # Disable gradient
+ torch.set_grad_enabled(False)
+
+ # =========================== Build main model ===========================
+ logger.info("Building model...")
+ factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
+ in_channels = args.latent_channels
+ out_channels = args.latent_channels
+
+ model = load_model(
+ args,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ factor_kwargs=factor_kwargs,
+ )
+
+ model = model.to(device)
+ model = Inference.load_state_dict(args, model, pretrained_model_path)
+ model.eval()
+
+ # ============================= Build extra models ========================
+ # VAE
+ vae, _, s_ratio, t_ratio = load_vae(
+ args.vae,
+ args.vae_precision,
+ logger=logger,
+ device=device if not args.use_cpu_offload else "cpu",
+ vae_path=args.vae_path
+ )
+ 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_path=args.text_encoder_path,
+ )
+ if world_size > 1:
+ import deepspeed
+ text_encoder.model = deepspeed.init_inference(
+ text_encoder.model,
+ tensor_parallel={"tp_size": world_size}
+ )
+ 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",
+ text_encoder_path=args.text_encoder_2_path,
+ )
+
+ return cls(
+ args=args,
+ vae=vae,
+ vae_kwargs=vae_kwargs,
+ text_encoder=text_encoder,
+ text_encoder_2=text_encoder_2,
+ model=model,
+ use_cpu_offload=args.use_cpu_offload,
+ device=device,
+ logger=logger,
+ parallel_args=parallel_args
+ )
+
+ @staticmethod
+ def load_state_dict(args, model, pretrained_model_path):
+ load_key = args.load_key
+ dit_weight = Path(args.dit_weight)
+
+ if dit_weight is None:
+ model_dir = pretrained_model_path / f"t2v_{args.model_resolution}"
+ files = list(model_dir.glob("*.pt"))
+ if len(files) == 0:
+ raise ValueError(f"No model weights found in {model_dir}")
+ if str(files[0]).startswith("pytorch_model_"):
+ model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+ bare_model = True
+ elif any(str(f).endswith("_model_states.pt") for f in files):
+ files = [f for f in files if str(f).endswith("_model_states.pt")]
+ model_path = files[0]
+ if len(files) > 1:
+ logger.warning(
+ f"Multiple model weights found in {dit_weight}, using {model_path}"
+ )
+ bare_model = False
+ else:
+ raise ValueError(
+ f"Invalid model path: {dit_weight} with unrecognized weight format: "
+ f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+ f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
+ f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+ f"specific weight file, please provide the full path to the file."
+ )
+ else:
+ if dit_weight.is_dir():
+ files = list(dit_weight.glob("*.pt"))
+ if len(files) == 0:
+ raise ValueError(f"No model weights found in {dit_weight}")
+ if str(files[0]).startswith("pytorch_model_"):
+ model_path = dit_weight / f"pytorch_model_{load_key}.pt"
+ bare_model = True
+ elif any(str(f).endswith("_model_states.pt") for f in files):
+ files = [f for f in files if str(f).endswith("_model_states.pt")]
+ model_path = files[0]
+ if len(files) > 1:
+ logger.warning(
+ f"Multiple model weights found in {dit_weight}, using {model_path}"
+ )
+ bare_model = False
+ else:
+ raise ValueError(
+ f"Invalid model path: {dit_weight} with unrecognized weight format: "
+ f"{list(map(str, files))}. When given a directory as --dit-weight, only "
+ f"`pytorch_model_*.pt`(provided by HunyuanDiT official) and "
+ f"`*_model_states.pt`(saved by deepspeed) can be parsed. If you want to load a "
+ f"specific weight file, please provide the full path to the file."
+ )
+ elif dit_weight.is_file():
+ model_path = dit_weight
+ bare_model = "unknown"
+ else:
+ raise ValueError(f"Invalid model path: {dit_weight}")
+
+ if not model_path.exists():
+ raise ValueError(f"model_path not exists: {model_path}")
+ logger.info(f"Loading torch model {model_path}...")
+ state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
+
+ if bare_model == "unknown" and ("ema" in state_dict or "module" in state_dict):
+ bare_model = False
+ if bare_model is False:
+ if load_key in state_dict:
+ state_dict = state_dict[load_key]
+ else:
+ raise KeyError(
+ f"Missing key: `{load_key}` in the checkpoint: {model_path}. The keys in the checkpoint "
+ f"are: {list(state_dict.keys())}."
+ )
+ model.load_state_dict(state_dict, strict=True)
+ return model
+
+ @staticmethod
+ def parse_size(size):
+ if isinstance(size, int):
+ size = [size]
+ if not isinstance(size, (list, tuple)):
+ raise ValueError(f"Size must be an integer or (height, width), got {size}.")
+ if len(size) == 1:
+ size = [size[0], size[0]]
+ if len(size) != 2:
+ raise ValueError(f"Size must be an integer or (height, width), got {size}.")
+ return size
+
+
+class HunyuanVideoSampler(Inference):
+ def __init__(
+ self,
+ args,
+ vae,
+ vae_kwargs,
+ text_encoder,
+ model,
+ text_encoder_2=None,
+ pipeline=None,
+ use_cpu_offload=False,
+ device=0,
+ logger=None,
+ parallel_args=None
+ ):
+ super().__init__(
+ args,
+ vae,
+ vae_kwargs,
+ text_encoder,
+ model,
+ text_encoder_2=text_encoder_2,
+ pipeline=pipeline,
+ use_cpu_offload=use_cpu_offload,
+ device=device,
+ logger=logger,
+ parallel_args=parallel_args
+ )
+
+ self.pipeline = self.load_diffusion_pipeline(
+ args=args,
+ vae=self.vae,
+ text_encoder=self.text_encoder,
+ text_encoder_2=self.text_encoder_2,
+ model=self.model,
+ device=self.device,
+ )
+
+ self.default_negative_prompt = NEGATIVE_PROMPT
+
+ def load_diffusion_pipeline(
+ self,
+ args,
+ vae,
+ text_encoder,
+ text_encoder_2,
+ model,
+ scheduler=None,
+ device=None,
+ progress_bar_config=None,
+ data_type="video",
+ ):
+ """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(device="npu")
+ 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.0,
+ 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 math.isclose(guidance_scale, 1.0):
+ 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
+
+ def finalize(self):
+ world_size = int(os.getenv('WORLD_SIZE', 1))
+ if world_size > 1:
+ finalize_parallel_env()
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ebe2c3ed5ceb92a380f17d8a91aa2fa08798c50
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/__init__.py
@@ -0,0 +1,26 @@
+from .models import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG
+
+
+def load_model(args, in_channels, out_channels, factor_kwargs):
+ """load hunyuan video model
+
+ Args:
+ args (dict): model args
+ in_channels (int): input channels number
+ out_channels (int): output channels number
+ factor_kwargs (dict): factor kwargs
+
+ Returns:
+ model (nn.Module): The hunyuan video model
+ """
+ if args.model in HUNYUAN_VIDEO_CONFIG.keys():
+ model = HYVideoDiffusionTransformer(
+ args,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ **HUNYUAN_VIDEO_CONFIG[args.model],
+ **factor_kwargs,
+ )
+ return model
+ else:
+ raise NotImplementedError()
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/activation_layers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/activation_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8774c26ceef6081482ca0dbbf930b207d4ac03b
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/attention.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..ddd4c23d075993b1f1dee7a1906019148ea821c9
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/attention.py
@@ -0,0 +1,243 @@
+import math
+
+import torch
+import torch_npu
+from einops import rearrange
+from mindiesd.layers.flash_attn.attention_forward import attention_forward
+
+from hyvideo.utils.comm import all_to_all_4D
+from hyvideo.utils.parallel_mgr import (
+ get_sp_group,
+ get_sequence_parallel_world_size,
+ get_sequence_parallel_rank,
+ get_tp_group,
+ get_tensor_model_parallel_world_size,
+ get_tensor_model_parallel_rank,
+)
+
+USE_LAYSER_ATTENTION = True
+
+
+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="npu")
+
+ 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,
+ 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]
+ """
+ b, q_s, n, d = q.size(0), q.size(1), q.size(2), q.size(3)
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ x = torch_npu.npu_fusion_attention(q, k, v,
+ atten_mask=attn_mask,
+ input_layout="BNSD",
+ scale=1 / math.sqrt(d),
+ head_num=n)[0]
+ x = x.transpose(1, 2).reshape(b, q_s, -1)
+ return x
+
+
+def dit_attention(
+ q,
+ k,
+ v,
+ img_q_len,
+ img_kv_len,
+ cu_seqlens_q=None,
+ cu_seqlens_kv=None
+):
+ """
+ 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]
+ """
+ b, n, d = q.size(0), q.size(2), q.size(3)
+ if get_sequence_parallel_world_size() > 1:
+ x1 = hybrid_seq_parallel_attn(
+ q[:, :img_q_len, :, :],
+ k[:, :img_kv_len, :, :],
+ v[:, :img_kv_len, :, :],
+ joint_tensor_query=q[:, img_q_len:cu_seqlens_q[1]],
+ joint_tensor_key=k[:, img_kv_len:cu_seqlens_kv[1]],
+ joint_tensor_value=v[:, img_kv_len:cu_seqlens_kv[1]],
+ )
+ x2 = torch_npu.npu_fusion_attention(q[:, cu_seqlens_q[1]:],
+ k[:, cu_seqlens_kv[1]:],
+ v[:, cu_seqlens_kv[1]:],
+ input_layout="BSND",
+ scale=1 / math.sqrt(d),
+ head_num=n)[0]
+ x = torch.cat([x1, x2], dim=1)
+ x = rearrange(x, "b s n d -> b s (n d)")
+ else:
+ q_s = q.shape[1]
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+ x1 = torch_npu.npu_fusion_attention(q[:, :, :cu_seqlens_q[1], :],
+ k[:, :, :cu_seqlens_kv[1], :],
+ v[:, :, :cu_seqlens_kv[1], :],
+ input_layout="BNSD",
+ scale=1 / math.sqrt(d),
+ head_num=n)[0]
+ x2 = torch_npu.npu_fusion_attention(q[:, :, cu_seqlens_q[1]:, :],
+ k[:, :, cu_seqlens_kv[1]:, :],
+ v[:, :, cu_seqlens_kv[1]:, :],
+ input_layout="BNSD",
+ scale=1 / math.sqrt(d),
+ head_num=n)[0]
+ x = torch.cat([x1, x2], dim=2)
+ x = x.transpose(1, 2).reshape(b, q_s, -1)
+ return x
+
+
+def hybrid_seq_parallel_attn(
+ query,
+ key,
+ value,
+ joint_tensor_query=None,
+ joint_tensor_key=None,
+ joint_tensor_value=None,
+):
+ sp_size = get_sequence_parallel_world_size()
+ sp_rank = get_sequence_parallel_rank()
+ tp_size = get_tensor_model_parallel_world_size()
+ tp_rank = get_tensor_model_parallel_rank()
+ is_joint = False
+ if (joint_tensor_query is not None and
+ joint_tensor_key is not None and
+ joint_tensor_value is not None):
+ query = torch.cat([query, joint_tensor_query], dim=1)
+ is_joint = True
+ elif (joint_tensor_query is None and
+ joint_tensor_key is None and
+ joint_tensor_value is None):
+ pass
+ else:
+ raise ValueError(
+ f"joint_tensor_query, joint_tensor_key, and joint_tensor_value should be None or not None simultaneously."
+ )
+
+ if is_joint:
+ attn_heads_per_sp_rank = (
+ joint_tensor_key.shape[-2] // sp_size
+ )
+ joint_tensor_key = joint_tensor_key[
+ ...,
+ attn_heads_per_sp_rank
+ * sp_rank: attn_heads_per_sp_rank
+ * (sp_rank + 1),
+ :,
+ ]
+ joint_tensor_value = joint_tensor_value[
+ ...,
+ attn_heads_per_sp_rank
+ * sp_rank: attn_heads_per_sp_rank
+ * (sp_rank + 1),
+ :,
+ ]
+ query_layer = all_to_all_4D(input_=query, scatter_idx=2, gather_idx=1, group=get_sp_group().device_group)
+ key_layer = all_to_all_4D(input_=key, scatter_idx=2, gather_idx=1, group=get_sp_group().device_group)
+ value_layer = all_to_all_4D(input_=value, scatter_idx=2, gather_idx=1, group=get_sp_group().device_group)
+
+ if is_joint:
+ key_layer = torch.cat([key_layer, joint_tensor_key], dim=1)
+ value_layer = torch.cat([value_layer, joint_tensor_value], dim=1)
+
+ if tp_size > 1:
+ query_layer = query_layer.chunk(tp_size, dim=1)[tp_rank]
+
+ if USE_LAYSER_ATTENTION:
+ output = attention_forward(query_layer,
+ key_layer,
+ value_layer,
+ opt_mode="manual", op_type="ascend_laser_attention", layout="BNSD")
+ else:
+ query_layer = query_layer.transpose(1, 2) # BNSD
+ key_layer = key_layer.transpose(1, 2) # BNSD
+ value_layer = value_layer.transpose(1, 2) # BNSD
+ output = torch_npu.npu_fusion_attention(query_layer,
+ key_layer,
+ value_layer,
+ head_num=query_layer.shape[1],
+ input_layout="BNSD",
+ scale=1 / math.sqrt(query_layer.shape[-1]))[0]
+ output = output.transpose(1, 2).contiguous() # BSND
+
+ if tp_size > 1:
+ output = get_tp_group().all_gather(output, dim=1, separate_tensors=False)
+
+ output = all_to_all_4D(input_=output, scatter_idx=1, gather_idx=2, group=get_sp_group().device_group)
+ return output
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/embed_layers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/embed_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ec990cea0765b2bbb08275d3c41cdcb3f239556
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/embed_layers.py
@@ -0,0 +1,156 @@
+import math
+import torch
+import torch.nn as nn
+
+from ..utils.helpers import to_2tuple
+
+
+class PatchEmbed(nn.Module):
+ """2D Image to Patch Embedding
+
+ Image to Patch Embedding using Conv2d
+
+ A convolution based approach to patchifying a 2D image w/ embedding projection.
+
+ Based on the impl in https://github.com/google-research/vision_transformer
+
+ Hacked together by / Copyright 2020 Ross Wightman
+
+ Remove the _assert function in forward function to be compatible with multi-resolution images.
+ """
+
+ def __init__(
+ self,
+ patch_size=16,
+ in_chans=3,
+ embed_dim=768,
+ norm_layer=None,
+ flatten=True,
+ bias=True,
+ dtype=None,
+ device=None,
+ ):
+ factory_kwargs = {"dtype": dtype, "device": device}
+ super().__init__()
+ patch_size = to_2tuple(patch_size)
+ self.patch_size = patch_size
+ self.flatten = flatten
+
+ self.proj = nn.Conv3d(
+ in_chans,
+ embed_dim,
+ kernel_size=patch_size,
+ stride=patch_size,
+ bias=bias,
+ **factory_kwargs
+ )
+ nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1))
+ if bias:
+ nn.init.zeros_(self.proj.bias)
+
+ self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+ def forward(self, x):
+ x = self.proj(x)
+ if self.flatten:
+ x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
+ x = self.norm(x)
+ return x
+
+
+class TextProjection(nn.Module):
+ """
+ Projects text embeddings. Also handles dropout for classifier-free guidance.
+
+ Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+ """
+
+ def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None):
+ factory_kwargs = {"dtype": dtype, "device": device}
+ super().__init__()
+ self.linear_1 = nn.Linear(
+ in_features=in_channels,
+ out_features=hidden_size,
+ bias=True,
+ **factory_kwargs
+ )
+ self.act_1 = act_layer()
+ self.linear_2 = nn.Linear(
+ in_features=hidden_size,
+ out_features=hidden_size,
+ bias=True,
+ **factory_kwargs
+ )
+
+ def forward(self, caption):
+ hidden_states = self.linear_1(caption)
+ hidden_states = self.act_1(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+
+def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+
+ Args:
+ t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
+ dim (int): the dimension of the output.
+ max_period (int): controls the minimum frequency of the embeddings.
+
+ Returns:
+ embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
+
+ .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ """
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period)
+ * torch.arange(start=0, end=half, dtype=torch.float32)
+ / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+
+ def __init__(
+ self,
+ hidden_size,
+ act_layer,
+ frequency_embedding_size=256,
+ max_period=10000,
+ out_size=None,
+ dtype=None,
+ device=None,
+ ):
+ factory_kwargs = {"dtype": dtype, "device": device}
+ super().__init__()
+ self.frequency_embedding_size = frequency_embedding_size
+ self.max_period = max_period
+ if out_size is None:
+ out_size = hidden_size
+
+ self.mlp = nn.Sequential(
+ nn.Linear(
+ frequency_embedding_size, hidden_size, bias=True, **factory_kwargs
+ ),
+ act_layer(),
+ nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs),
+ )
+ nn.init.normal_(self.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.mlp[2].weight, std=0.02)
+
+ def forward(self, t):
+ t_freq = timestep_embedding(
+ t, self.frequency_embedding_size, self.max_period
+ ).type(self.mlp[0].weight.dtype)
+ t_emb = self.mlp(t_freq)
+ return t_emb
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/mlp_layers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/mlp_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..fecf17ad8b87794d4e7834f5e07a5c221c8e183a
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/mlp_layers.py
@@ -0,0 +1,119 @@
+# 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/models.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..b995c289498af393d8e401b91feace223bb94965
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/models.py
@@ -0,0 +1,871 @@
+from typing import Any, List, Tuple, Optional, Union, Dict
+import math
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models import ModelMixin
+from einops import rearrange
+
+from hyvideo.utils.parallel_mgr import (
+ get_sp_group,
+ get_sequence_parallel_world_size,
+ get_sequence_parallel_rank,
+)
+from .activation_layers import get_activation_layer
+from .attention import dit_attention, get_cu_seqlens
+from .embed_layers import TimestepEmbedder, PatchEmbed, TextProjection
+from .mlp_layers import MLP, MLPEmbedder, FinalLayer
+from .modulate_layers import ModulateDiT, apply_gate
+from .norm_layers import get_norm_layer
+from .posemb_layers import apply_rotary_emb
+from .token_refiner import SingleTokenRefiner
+
+
+class MMDoubleStreamBlock(nn.Module):
+ """
+ A multimodal dit block with 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,
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+
+ self.deterministic = False
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+ mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+
+ self.img_mod = ModulateDiT(
+ hidden_size,
+ factor=6,
+ act_layer=get_activation_layer("silu"),
+ **factory_kwargs,
+ )
+
+ self.img_norm1 = get_norm_layer("layer")(hidden_size)
+
+ 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 = get_norm_layer("layer")(hidden_size)
+ 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 = get_norm_layer("layer")(hidden_size)
+
+ 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 = get_norm_layer("layer")(hidden_size)
+ self.txt_mlp = MLP(
+ hidden_size,
+ mlp_hidden_dim,
+ act_layer=get_activation_layer(mlp_act_type),
+ bias=True,
+ **factory_kwargs,
+ )
+ self.attention_cache = None
+
+ def enable_deterministic(self):
+ self.deterministic = True
+
+ def disable_deterministic(self):
+ self.deterministic = False
+
+ def double_stream_attention(
+ self,
+ img, txt,
+ img_mod1_shift,
+ img_mod1_scale,
+ txt_mod1_shift,
+ txt_mod1_scale,
+ freqs_cis,
+ cu_seqlens_q,
+ cu_seqlens_kv,
+ max_seqlen_q,
+ max_seqlen_kv
+ ):
+ # Prepare image for attention.
+ img_modulated = self.img_norm1(img, 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, 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]}"
+
+ # attention computation start
+ attn = dit_attention(
+ q,
+ k,
+ v,
+ img_q_len=img_q.shape[1],
+ img_kv_len=img_k.shape[1],
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv
+ )
+
+ # attention computation end
+ return attn
+
+ 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]:
+ batch_size = img.size(0)
+ (
+ img_mod1_shift,
+ img_mod1_scale,
+ img_mod1_gate,
+ img_mod2_shift,
+ img_mod2_scale,
+ img_mod2_gate,
+ ) = self.img_mod(vec).reshape(batch_size, 6, -1).transpose(0, 1).unbind(dim=0)
+ (
+ txt_mod1_shift,
+ txt_mod1_scale,
+ txt_mod1_gate,
+ txt_mod2_shift,
+ txt_mod2_scale,
+ txt_mod2_gate,
+ ) = self.txt_mod(vec).reshape(batch_size, 6, -1).transpose(0, 1).unbind(dim=0)
+
+ if self.attention_cache:
+ attn = self.attention_cache.apply(
+ self.double_stream_attention,
+ img=img, txt=txt,
+ img_mod1_shift=img_mod1_shift,
+ img_mod1_scale=img_mod1_scale,
+ txt_mod1_shift=txt_mod1_shift,
+ txt_mod1_scale=txt_mod1_scale,
+ freqs_cis=freqs_cis,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv
+ )
+ else:
+ attn = self.double_stream_attention(
+ img=img, txt=txt,
+ img_mod1_shift=img_mod1_shift,
+ img_mod1_scale=img_mod1_scale,
+ txt_mod1_shift=txt_mod1_shift,
+ txt_mod1_scale=txt_mod1_scale,
+ freqs_cis=freqs_cis,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv
+ )
+
+ 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(
+ 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(
+ self.txt_norm2(txt, shift=txt_mod2_shift, scale=txt_mod2_scale)
+ ),
+ gate=txt_mod2_gate,
+ )
+
+ return img, txt
+
+
+class MMSingleStreamBlock(nn.Module):
+ """
+ A DiT block with parallel linear layers as described in
+ https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+ Also refer to (SD3): https://arxiv.org/abs/2403.03206
+ (Flux.1): https://github.com/black-forest-labs/flux
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ heads_num: int,
+ mlp_width_ratio: float = 4.0,
+ mlp_act_type: str = "gelu_tanh",
+ qk_norm: bool = True,
+ qk_norm_type: str = "rms",
+ qk_scale: float = None,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+
+ self.deterministic = False
+ self.hidden_size = hidden_size
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+ mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+ self.mlp_hidden_dim = mlp_hidden_dim
+ self.scale = qk_scale or head_dim ** -0.5
+
+ # qkv and mlp_in
+ self.linear1 = nn.Linear(
+ hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs
+ )
+ # proj and mlp_out
+ self.linear2 = nn.Linear(
+ hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs
+ )
+
+ qk_norm_layer = get_norm_layer(qk_norm_type)
+ self.q_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.k_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ if qk_norm
+ else nn.Identity()
+ )
+
+ self.pre_norm = get_norm_layer("layer")(hidden_size)
+
+ self.mlp_act = get_activation_layer(mlp_act_type)()
+ self.modulation = ModulateDiT(
+ hidden_size,
+ factor=3,
+ act_layer=get_activation_layer("silu"),
+ **factory_kwargs,
+ )
+ self.attention_cache = None
+
+ def enable_deterministic(self):
+ self.deterministic = True
+
+ def disable_deterministic(self):
+ self.deterministic = False
+
+ def single_stream_attention(
+ self,
+ qkv,
+ freqs_cis,
+ txt_len,
+ x,
+ cu_seqlens_q,
+ cu_seqlens_kv,
+ max_seqlen_q,
+ max_seqlen_kv
+ ):
+ 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]}"
+
+ # attention computation start
+ attn = dit_attention(
+ q,
+ k,
+ v,
+ img_q_len=img_q.shape[1],
+ img_kv_len=img_k.shape[1],
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv
+ )
+ # attention computation end
+ return attn
+
+ 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:
+ batch_size = x.size(0)
+ mod_shift, mod_scale, mod_gate = self.modulation(vec).reshape(batch_size, 3, -1).transpose(0, 1).unbind(dim=0)
+ x_mod = 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
+ )
+ if self.attention_cache:
+ attn = self.attention_cache.apply(
+ self.single_stream_attention,
+ qkv=qkv,
+ freqs_cis=freqs_cis,
+ txt_len=txt_len,
+ x=x,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv
+ )
+ else:
+ attn = self.single_stream_attention(
+ qkv=qkv,
+ freqs_cis=freqs_cis,
+ txt_len=txt_len,
+ x=x,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv
+ )
+
+ # 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,
+ args: Any,
+ patch_size: list = None,
+ 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] = None,
+ qkv_bias: bool = True,
+ qk_norm: bool = True,
+ qk_norm_type: str = "rms",
+ guidance_embed: bool = False, # For modulation.
+ text_projection: str = "single_refiner",
+ use_attention_mask: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+
+ if patch_size is None:
+ patch_size = [1, 2, 2]
+ if rope_dim_list is None:
+ rope_dim_list = [16, 56, 56]
+
+ 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 = args.text_states_dim
+ self.text_states_dim_2 = args.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,
+ **factory_kwargs,
+ )
+ for _ in range(mm_double_blocks_depth)
+ ]
+ )
+
+ # single blocks
+ self.single_blocks = nn.ModuleList(
+ [
+ MMSingleStreamBlock(
+ self.hidden_size,
+ self.heads_num,
+ mlp_width_ratio=mlp_width_ratio,
+ mlp_act_type=mlp_act_type,
+ qk_norm=qk_norm,
+ qk_norm_type=qk_norm_type,
+ **factory_kwargs,
+ )
+ for _ in range(mm_single_blocks_depth)
+ ]
+ )
+
+ self.final_layer = FinalLayer(
+ self.hidden_size,
+ self.patch_size,
+ self.out_channels,
+ get_activation_layer("silu"),
+ **factory_kwargs,
+ )
+ self.double_blocks_cache = None
+ self.single_blocks_cache = None
+
+ 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.
+ t_idx: int = 0, # Timestep index
+ return_dict: bool = True,
+ ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+ if get_sequence_parallel_world_size() > 1:
+ if ((x.shape[-1] // 2) * (x.shape[-2] // 2)) % get_sequence_parallel_world_size() != 0:
+ raise ValueError(
+ f"Cannot split video sequence into ulysses_degree ({get_sequence_parallel_world_size()}) parts evenly")
+
+ # patch sizes for the temporal, height, and width dimensions are 1, 2, and 2.
+ temporal_size, h, w = x.shape[2], x.shape[3] // 2, x.shape[4] // 2
+ w_chunk_size = math.gcd(w, get_sequence_parallel_world_size())
+ h_chunk_size = get_sequence_parallel_world_size() // w_chunk_size
+
+ x = torch.chunk(x, w_chunk_size, dim=-1)[get_sequence_parallel_rank() // h_chunk_size]
+ if h_chunk_size > 1:
+ x = torch.chunk(x, h_chunk_size, dim=-2)[get_sequence_parallel_rank() % h_chunk_size]
+
+ dim_thw = freqs_cos.shape[-1]
+ freqs_cos = freqs_cos.reshape(temporal_size, h, w, dim_thw)
+ freqs_cos = torch.chunk(freqs_cos,
+ w_chunk_size, dim=-2)[get_sequence_parallel_rank() // h_chunk_size]
+ if h_chunk_size > 1:
+ freqs_cos = torch.chunk(freqs_cos,
+ h_chunk_size, dim=-3)[get_sequence_parallel_rank() % h_chunk_size]
+ freqs_cos = freqs_cos.reshape(-1, dim_thw)
+ dim_thw = freqs_sin.shape[-1]
+ freqs_sin = freqs_sin.reshape(temporal_size, h, w, dim_thw)
+ freqs_sin = torch.chunk(freqs_sin,
+ w_chunk_size, dim=-2)[get_sequence_parallel_rank() // h_chunk_size]
+ if h_chunk_size > 1:
+ freqs_sin = torch.chunk(freqs_sin,
+ h_chunk_size, dim=-3)[get_sequence_parallel_rank() % h_chunk_size]
+ freqs_sin = freqs_sin.reshape(-1, dim_thw)
+ 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,
+ ]
+
+ if not self.double_blocks_cache:
+ img, txt = block(*double_block_args)
+ else:
+ img, txt = self.double_blocks_cache.apply(block,
+ hidden_states=img,
+ encoder_hidden_states=txt,
+ vec=vec,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv,
+ freqs_cis=freqs_cis,
+ )
+
+ # 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),
+ ]
+
+ if not self.single_blocks_cache:
+ x = block(*single_block_args)
+ else:
+ x = self.single_blocks_cache.apply(block,
+ hidden_states=x,
+ vec=vec,
+ txt_len=txt_seq_len,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv,
+ freqs_cis=(freqs_cos, freqs_sin),
+ )
+
+ 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 get_sequence_parallel_world_size() > 1:
+ if h_chunk_size > 1:
+ h_, w_ = img.shape[-2], img.shape[-1]
+ img = get_sp_group().all_gather(img, dim=-2)
+ b, c, = img.shape[0], img.shape[1]
+ img = img.reshape(
+ b, c, temporal_size, h_chunk_size, w_chunk_size, h_ // h_chunk_size, w_ // w_chunk_size)
+ img = img.transpose(4, 5)
+ img = img.reshape(b, c, temporal_size, h_, w_)
+ else:
+ img = get_sp_group().all_gather(img, dim=-1)
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/modulate_layers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/modulate_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..24f1380dcc5f553bdd89f1f55add7ee520155c06
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/modulate_layers.py
@@ -0,0 +1,77 @@
+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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/norm_layers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/norm_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..69d3f7b445fd48221e0be203ff0746aa7839a15c
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/norm_layers.py
@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+import torch_npu
+
+
+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.
+
+ """
+ return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.eps)[0]
+
+
+class LayerNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.eps = eps
+
+ def forward(self, x, shift, scale):
+ return torch_npu.npu_layer_norm_eval(
+ x, normalized_shape=[self.hidden_size],
+ weight=(1 + scale.unsqueeze(1)), bias=shift.unsqueeze(1), eps=self.eps
+ )
+
+
+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 LayerNorm
+ elif norm_layer == "rms":
+ return RMSNorm
+ else:
+ raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/posemb_layers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/posemb_layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..a06be04b32b99b5a9cf45d4b4821891654c15aec
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/posemb_layers.py
@@ -0,0 +1,305 @@
+from typing import Union, Tuple, List
+import math
+
+import torch
+from mindiesd.layers.embedding import rotary_position_embedding
+
+
+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 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]
+ # real * cos - imag * sin
+ # imag * cos + real * sin
+ xq_out = rotary_position_embedding(xq, cos, sin, "rotated_interleaved")
+ xk_out = rotary_position_embedding(xk, cos, sin, "rotated_interleaved")
+ 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, rope_dim in enumerate(rope_dim_list):
+ emb = get_1d_rotary_pos_embed(
+ rope_dim,
+ 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 not math.isclose(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]
+
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/token_refiner.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/token_refiner.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b9f25c9416e0f74728afb8c6697dccf152c87cf
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/modules/token_refiner.py
@@ -0,0 +1,236 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from .activation_layers import get_activation_layer
+from .attention import attention
+from .embed_layers import TimestepEmbedder, TextProjection
+from .mlp_layers import MLP
+from .modulate_layers import apply_gate
+from .norm_layers import get_norm_layer
+
+
+class IndividualTokenRefinerBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ heads_num,
+ mlp_width_ratio: str = 4.0,
+ mlp_drop_rate: float = 0.0,
+ act_type: str = "silu",
+ qk_norm: bool = False,
+ qk_norm_type: str = "layer",
+ qkv_bias: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+ mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
+
+ self.norm1 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
+ )
+ self.self_attn_qkv = nn.Linear(
+ hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
+ )
+ qk_norm_layer = get_norm_layer(qk_norm_type)
+ self.self_attn_q_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.self_attn_k_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+ if qk_norm
+ else nn.Identity()
+ )
+ self.self_attn_proj = nn.Linear(
+ hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
+ )
+
+ self.norm2 = nn.LayerNorm(
+ hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
+ )
+ act_layer = get_activation_layer(act_type)
+ self.mlp = MLP(
+ in_channels=hidden_size,
+ hidden_channels=mlp_hidden_dim,
+ act_layer=act_layer,
+ drop=mlp_drop_rate,
+ **factory_kwargs,
+ )
+
+ self.adaLN_modulation = nn.Sequential(
+ act_layer(),
+ nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
+ )
+ # Zero-initialize the modulation
+ nn.init.zeros_(self.adaLN_modulation[1].weight)
+ nn.init.zeros_(self.adaLN_modulation[1].bias)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ c: torch.Tensor, # timestep_aware_representations + context_aware_representations
+ attn_mask: torch.Tensor = None,
+ ):
+ gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
+
+ norm_x = self.norm1(x)
+ qkv = self.self_attn_qkv(norm_x)
+ q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
+ # Apply QK-Norm if needed
+ q = self.self_attn_q_norm(q).to(v)
+ k = self.self_attn_k_norm(k).to(v)
+
+ # Self-Attention
+ attn = attention(q, k, v, attn_mask=(~attn_mask))
+
+ x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
+
+ # FFN Layer
+ x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp)
+
+ return x
+
+
+class IndividualTokenRefiner(nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ heads_num,
+ depth,
+ mlp_width_ratio: float = 4.0,
+ mlp_drop_rate: float = 0.0,
+ act_type: str = "silu",
+ qk_norm: bool = False,
+ qk_norm_type: str = "layer",
+ qkv_bias: bool = True,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+ self.blocks = nn.ModuleList(
+ [
+ IndividualTokenRefinerBlock(
+ hidden_size=hidden_size,
+ heads_num=heads_num,
+ mlp_width_ratio=mlp_width_ratio,
+ mlp_drop_rate=mlp_drop_rate,
+ act_type=act_type,
+ qk_norm=qk_norm,
+ qk_norm_type=qk_norm_type,
+ qkv_bias=qkv_bias,
+ **factory_kwargs,
+ )
+ for _ in range(depth)
+ ]
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ c: torch.LongTensor,
+ mask: Optional[torch.Tensor] = None,
+ ):
+ 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 = "torch",
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+ self.attn_mode = attn_mode
+ assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner."
+
+ self.input_embedder = nn.Linear(
+ in_channels, hidden_size, bias=True, **factory_kwargs
+ )
+
+ act_layer = get_activation_layer(act_type)
+ # Build timestep embedding layer
+ self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs)
+ # Build context embedding layer
+ self.c_embedder = TextProjection(
+ in_channels, hidden_size, act_layer, **factory_kwargs
+ )
+
+ self.individual_token_refiner = IndividualTokenRefiner(
+ hidden_size=hidden_size,
+ heads_num=heads_num,
+ depth=depth,
+ mlp_width_ratio=mlp_width_ratio,
+ mlp_drop_rate=mlp_drop_rate,
+ act_type=act_type,
+ qk_norm=qk_norm,
+ qk_norm_type=qk_norm_type,
+ qkv_bias=qkv_bias,
+ **factory_kwargs,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ t: torch.LongTensor,
+ mask: Optional[torch.LongTensor] = None,
+ ):
+ timestep_aware_representations = self.t_embedder(t)
+
+ if mask is None:
+ context_aware_representations = x.mean(dim=1)
+ else:
+ mask_float = mask.float().unsqueeze(-1) # [b, s1, 1]
+ context_aware_representations = (x * mask_float).sum(
+ dim=1
+ ) / mask_float.sum(dim=1)
+ context_aware_representations = self.c_embedder(context_aware_representations)
+ c = timestep_aware_representations + context_aware_representations
+
+ x = self.input_embedder(x)
+
+ x = self.individual_token_refiner(x, c, mask)
+
+ return x
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/text_encoder/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/text_encoder/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..94a5888af4f3291542ae22359055ea1097030071
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/text_encoder/__init__.py
@@ -0,0 +1,356 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from transformers import CLIPTextModel, CLIPTokenizer, AutoTokenizer, AutoModel
+from transformers.utils import ModelOutput
+
+from ..constants import PRECISION_TO_TYPE
+from ..constants import TEXT_ENCODER_PATH, TOKENIZER_PATH
+
+
+def use_default(value, default):
+ return value if value is not None else default
+
+
+def load_text_encoder(
+ text_encoder_type,
+ text_encoder_precision=None,
+ text_encoder_path=None,
+ logger=None,
+ device=None,
+):
+ if text_encoder_path is None:
+ text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type]
+ if logger is not None:
+ logger.info(
+ f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}"
+ )
+
+ if text_encoder_type == "clipL":
+ text_encoder = CLIPTextModel.from_pretrained(text_encoder_path, local_files_only=True)
+ 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, local_files_only=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, local_files_only=True)
+ elif tokenizer_type == "llm":
+ tokenizer = AutoTokenizer.from_pretrained(
+ tokenizer_path, padding_side=padding_side, local_files_only=True
+ )
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/comm.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/comm.py
new file mode 100644
index 0000000000000000000000000000000000000000..287241e5b4de4edea35a214e723248c58233dcdf
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/comm.py
@@ -0,0 +1,95 @@
+import torch
+
+import torch.distributed as dist
+
+
+def all_to_all_4D(
+ input_: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None, use_sync: bool = False
+) -> torch.tensor:
+ """
+ all-to-all for QKV
+
+ Args:
+ input_ (torch.tensor): a tensor sharded along dim scatter dim
+ scatter_idx (int): default 1
+ gather_idx (int): default 2
+ group : torch process group
+ use_sync (bool): whether to synchronize after all-to-all
+
+ Returns:
+ torch.tensor: resharded tensor (bs, seqlen/P, hc, hs)
+ """
+ assert (
+ input_.dim() == 4
+ ), f"input_ must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
+
+ seq_world_size = dist.get_world_size(group)
+
+ if scatter_idx == 2 and gather_idx == 1:
+ # input_ (torch.tensor): a tensor sharded along dim 1 (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs)
+ bs, shard_seqlen, hc, hs = input_.shape
+ seqlen = shard_seqlen * seq_world_size
+ shard_hc = hc // seq_world_size
+
+ # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+ # (bs, seqlen/P, hc, hs) -reshape-> (bs, seq_len/P, P, hc/P, hs) -transpose(0,2)-> (P, seq_len/P, bs, hc/P, hs)
+ input_t = (
+ input_.reshape(bs, shard_seqlen, seq_world_size, shard_hc, hs)
+ .transpose(0, 2)
+ .contiguous()
+ )
+
+ output = torch.empty_like(input_t)
+ # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single
+ # (P, seq_len/P, bs, hc/P, hs) scatter seqlen -all2all-> (P, seq_len/P, bs, hc/P, hs) scatter head
+
+ if seq_world_size > 1:
+ dist.all_to_all_single(output, input_t, group=group)
+ if use_sync:
+ torch.npu.synchronize()
+ else:
+ output = input_t
+ # if scattering the seq-dim, transpose the heads back to the original dimension
+ output = output.reshape(seqlen, bs, shard_hc, hs)
+
+ # (seq_len, bs, hc/P, hs) -reshape-> (bs, seq_len, hc/P, hs)
+ output = output.transpose(0, 1).contiguous().reshape(bs, seqlen, shard_hc, hs)
+
+ return output
+
+ elif scatter_idx == 1 and gather_idx == 2:
+ # input_ (torch.tensor): a tensor sharded along dim 1 (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs)
+ bs, seqlen, shard_hc, hs = input_.shape
+ hc = shard_hc * seq_world_size
+ shard_seqlen = seqlen // seq_world_size
+ seq_world_size = dist.get_world_size(group)
+
+ # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+ # (bs, seqlen, hc/P, hs) -reshape-> (bs, P, seq_len/P, hc/P, hs) -transpose(0, 3)-> (hc/P, P, seqlen/P, bs, hs) -transpose(0, 1) -> (P, hc/P, seqlen/P, bs, hs)
+ input_t = (
+ input_.reshape(bs, seq_world_size, shard_seqlen, shard_hc, hs)
+ .transpose(0, 3)
+ .transpose(0, 1)
+ .contiguous()
+ .reshape(seq_world_size, shard_hc, shard_seqlen, bs, hs)
+ )
+
+ output = torch.empty_like(input_t)
+ # https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all_single
+ # (P, bs x hc/P, seqlen/P, hs) scatter seqlen -all2all-> (P, bs x seq_len/P, hc/P, hs) scatter head
+ if seq_world_size > 1:
+ dist.all_to_all_single(output, input_t, group=group)
+ if use_sync:
+ torch.npu.synchronize()
+ else:
+ output = input_t
+
+ # if scattering the seq-dim, transpose the heads back to the original dimension
+ output = output.reshape(hc, shard_seqlen, bs, hs)
+
+ # (hc, seqlen/N, bs, hs) -tranpose(0,2)-> (bs, seqlen/N, hc, hs)
+ output = output.transpose(0, 2).contiguous().reshape(bs, shard_seqlen, hc, hs)
+
+ return output
+ else:
+ raise RuntimeError("scatter_idx must be 1 or 2 and gather_idx must be 1 or 2")
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/data_utils.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/data_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c049c635fd332a735fcda0024f7cddf5687428e
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/data_utils.py
@@ -0,0 +1,14 @@
+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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/file_utils.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/file_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..82e244f03968eda97428d9129a1952060246243f
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/file_utils.py
@@ -0,0 +1,71 @@
+import os
+from pathlib import Path
+
+import imageio
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+
+CODE_SUFFIXES = {
+ ".py", # Python codes
+ ".sh", # Shell scripts
+ ".yaml",
+ ".yml", # Configuration files
+}
+
+
+def safe_dir(path):
+ """
+ Create a directory (or the parent directory of a file) if it does not exist.
+
+ Args:
+ path (str or Path): Path to the directory.
+
+ Returns:
+ path (Path): Path object of the directory.
+ """
+ path = Path(path)
+ path.mkdir(exist_ok=True, parents=True)
+ return path
+
+
+def safe_file(path):
+ """
+ Create the parent directory of a file if it does not exist.
+
+ Args:
+ path (str or Path): Path to the file.
+
+ Returns:
+ path (Path): Path object of the file.
+ """
+ path = Path(path)
+ path.parent.mkdir(exist_ok=True, parents=True)
+ return path
+
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, fps=24):
+ """save videos by video tensor
+ copy from https://github.com/guoyww/AnimateDiff/blob/e92bd5671ba62c0d774a32951453e328018b7c5b/animatediff/utils/util.py#L61
+
+ Args:
+ videos (torch.Tensor): video tensor predicted by the model
+ path (str): path to save video
+ rescale (bool, optional): rescale the video tensor from [-1, 1] to . Defaults to False.
+ n_rows (int, optional): Defaults to 1.
+ fps (int, optional): video save fps. Defaults to 8.
+ """
+ videos = rearrange(videos, "b c t h w -> t b c h w")
+ outputs = []
+ for x in videos:
+ x = torchvision.utils.make_grid(x, nrow=n_rows)
+ x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+ if rescale:
+ x = (x + 1.0) / 2.0 # -1,1 -> 0,1
+ x = torch.clamp(x, 0, 1)
+ x = (x * 255).numpy().astype(np.uint8)
+ outputs.append(x)
+
+ os.makedirs(os.path.dirname(path), exist_ok=True)
+ imageio.mimsave(path, outputs, fps=fps)
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/group_coordinator.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/group_coordinator.py
new file mode 100644
index 0000000000000000000000000000000000000000..dcc18e78b16b69081fc4f135978736b015488758
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/group_coordinator.py
@@ -0,0 +1,594 @@
+# Copyright 2024 xDiT team.
+# Adapted from
+# https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py
+# Copyright 2023 The vLLM team.
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+import pickle
+from collections import namedtuple
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed
+from diffusers.utils import logging
+from torch.distributed import Backend, ProcessGroup
+
+logger = logging.get_logger(__name__)
+
+TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
+
+
+def _split_tensor_dict(
+ tensor_dict: Dict[str, Union[torch.Tensor, Any]], prefix: str = ""
+) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]:
+ """Split the tensor dictionary into two parts:
+ 1. A list of (key, value) pairs. If the value is a tensor, it is replaced
+ by its metadata.
+ 2. A list of tensors.
+
+ If the Tensor is nested under `tensor_dict["key1"]["key2"]`, the key of its
+ metadata will be "key1%key2".
+ """
+ metadata_list: List[Tuple[str, Any]] = []
+ tensor_list = []
+ for key, value in tensor_dict.items():
+ if "%" in key:
+ logger.error(
+ "Avoid having '%' in key "
+ "as it is used as a separator for nested entries."
+ )
+ if isinstance(value, torch.Tensor):
+ # Note: we cannot use `value.device` here,
+ # because it contains not only the device type but also the device
+ # index (e.g. "npu:0"). We only need the device type.
+ # receiving side will set the device index.
+ device = value.device.type
+ metadata_list.append(
+ (prefix + key, TensorMetadata(device, value.dtype, value.size()))
+ )
+ tensor_list.append(value)
+ elif isinstance(value, dict):
+ if len(value) == 0:
+ metadata_list.append((prefix + key, value))
+ inner_metadata_list, inner_tensor_list = _split_tensor_dict(
+ value, prefix + key + "%"
+ )
+ metadata_list.extend(inner_metadata_list)
+ tensor_list.extend(inner_tensor_list)
+ else:
+ metadata_list.append((prefix + key, value))
+ return metadata_list, tensor_list
+
+
+def _update_nested_dict(nested_dict, flattened_key, value):
+ key_splits = flattened_key.split("%")
+ cur_dict = nested_dict
+ for k in key_splits[:-1]:
+ if k not in cur_dict:
+ cur_dict[k] = {}
+ cur_dict = cur_dict[k]
+ cur_dict[key_splits[-1]] = value
+
+
+class GroupCoordinator:
+ """
+ PyTorch ProcessGroup wrapper for a group of processes.
+ PyTorch ProcessGroup is bound to one specific communication backend,
+ e.g. NCCL, Gloo, MPI, etc.
+ GroupCoordinator takes charge of all the communication operations among
+ the processes in the group. It can route the communication to
+ a specific implementation (e.g. switch allreduce implementation
+ based on the tensor size and npu graph mode).
+ """
+
+ # available attributes:
+ rank: int # global rank
+ ranks: List[int] # global ranks in the group
+ world_size: int # size of the group
+ # difference between `local_rank` and `rank_in_group`:
+ # if we have a group of size 4 across two nodes:
+ # Process | Node | Rank | Local Rank | Rank in Group
+ # 0 | 0 | 0 | 0 | 0
+ # 1 | 0 | 1 | 1 | 1
+ # 2 | 1 | 2 | 0 | 2
+ # 3 | 1 | 3 | 1 | 3
+ local_rank: int # local rank used to assign devices
+ rank_in_group: int # rank inside the group
+ cpu_group: ProcessGroup # group for CPU communication
+ device_group: ProcessGroup # group for device communication
+
+ def __init__(
+ self,
+ group_ranks: List[List[int]],
+ local_rank: int,
+ torch_distributed_backend: Union[str, Backend],
+ ):
+
+ self.rank = torch.distributed.get_rank()
+ self.local_rank = local_rank
+ self.device_group = None
+ self.cpu_group = None
+
+ for ranks in group_ranks:
+ device_group = torch.distributed.new_group(
+ ranks, backend=torch_distributed_backend
+ )
+ # a group with `gloo` backend, to allow direct coordination between
+ # processes through the CPU.
+ cpu_group = torch.distributed.new_group(ranks, backend="gloo")
+ if self.rank in ranks:
+ self.ranks = ranks
+ self.world_size = len(ranks)
+ self.rank_in_group = ranks.index(self.rank)
+ self.device_group = device_group
+ self.cpu_group = cpu_group
+
+ if torch.npu.is_available():
+ self.device = torch.device(f"npu:{local_rank}")
+ else:
+ self.device = torch.device("cpu")
+
+ @property
+ def first_rank(self):
+ """Return the global rank of the first process in the group"""
+ return self.ranks[0]
+
+ @property
+ def last_rank(self):
+ """Return the global rank of the last process in the group"""
+ return self.ranks[-1]
+
+ @property
+ def is_first_rank(self):
+ """Return whether the caller is the first process in the group"""
+ return self.rank == self.first_rank
+
+ @property
+ def is_last_rank(self):
+ """Return whether the caller is the last process in the group"""
+ return self.rank == self.last_rank
+
+ @property
+ def next_rank(self):
+ """Return the global rank of the process that follows the caller"""
+ rank_in_group = self.rank_in_group
+ world_size = self.world_size
+ return self.ranks[(rank_in_group + 1) % world_size]
+
+ @property
+ def prev_rank(self):
+ """Return the global rank of the process that precedes the caller"""
+ rank_in_group = self.rank_in_group
+ world_size = self.world_size
+ return self.ranks[(rank_in_group - 1) % world_size]
+
+ @property
+ def group_next_rank(self):
+ """Return the group rank of the process that follows the caller"""
+ rank_in_group = self.rank_in_group
+ world_size = self.world_size
+ return (rank_in_group + 1) % world_size
+
+ @property
+ def group_prev_rank(self):
+ """Return the group rank of the process that precedes the caller"""
+ rank_in_group = self.rank_in_group
+ world_size = self.world_size
+ return (rank_in_group - 1) % world_size
+
+ @property
+ def skip_rank(self):
+ """Return the global rank of the process that skip connects with the caller"""
+ rank_in_group = self.rank_in_group
+ world_size = self.world_size
+ return self.ranks[(world_size - rank_in_group - 1) % world_size]
+
+ @property
+ def group_skip_rank(self):
+ """Return the group rank of the process that skip connects with the caller"""
+ rank_in_group = self.rank_in_group
+ world_size = self.world_size
+ return (world_size - rank_in_group - 1) % world_size
+
+ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+ """
+ NOTE: This operation will be applied in-place or out-of-place.
+ Always assume this function modifies its input, but use the return
+ value as the output.
+ """
+ # Bypass the function if we are using only 1 GPU.
+ if self.world_size == 1:
+ return input_
+ else:
+ torch.distributed.all_reduce(input_, group=self.device_group)
+ return input_
+
+ def all_gather(
+ self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False
+ ) -> Union[torch.Tensor, List[torch.Tensor]]:
+ world_size = self.world_size
+ # Bypass the function if we are using only 1 GPU.
+ if world_size == 1:
+ return input_
+ if dim < 0:
+ # Convert negative dim to positive.
+ dim += input_.dim()
+ # Allocate output tensor.
+ input_size = list(input_.size())
+ input_size[0] *= world_size
+ output_tensor = torch.empty(
+ input_size, dtype=input_.dtype, device=input_.device
+ )
+ # All-gather.
+ torch.distributed.all_gather_into_tensor(
+ output_tensor, input_, group=self.device_group
+ )
+ if dim != 0:
+ input_size[0] //= world_size
+ output_tensor = output_tensor.reshape([world_size, ] + input_size)
+ output_tensor = output_tensor.movedim(0, dim)
+
+ if separate_tensors:
+ tensor_list = [
+ output_tensor.view(-1)
+ .narrow(0, input_.numel() * i, input_.numel())
+ .view_as(input_)
+ for i in range(world_size)
+ ]
+ return tensor_list
+ else:
+ input_size = list(input_.size())
+ input_size[dim] = input_size[dim] * world_size
+ # Reshape
+ output_tensor = output_tensor.reshape(input_size)
+ return output_tensor
+
+ def gather(self, input_: torch.Tensor, dst: int = 0, dim: int = -1) -> torch.Tensor:
+ """
+ NOTE: We assume that the input tensor is on the same device across
+ all the ranks.
+ NOTE: `dst` is the local rank of the destination rank.
+ """
+ world_size = self.world_size
+ # Bypass the function if we are using only 1 GPU.
+ if world_size == 1:
+ return input_
+ if dim < 0:
+ # Convert negative dim to positive.
+ dim += input_.dim()
+ # Allocate output tensor.
+ if self.rank_in_group == dst:
+ gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+ else:
+ gather_list = None
+ # Gather.
+ torch.distributed.gather(
+ input_, gather_list, dst=self.ranks[dst], group=self.device_group
+ )
+ if self.rank_in_group == dst:
+ output_tensor = torch.cat(gather_list, dim=dim)
+ else:
+ output_tensor = None
+ return output_tensor
+
+ def broadcast(self, input_: torch.Tensor, src: int = 0):
+ """Broadcast the input tensor.
+ NOTE: `src` is the local rank of the source rank.
+ """
+
+ # Bypass the function if we are using only 1 GPU.
+ if self.world_size == 1:
+ return input_
+ # Broadcast.
+ torch.distributed.broadcast(
+ input_, src=self.ranks[src], group=self.device_group
+ )
+ return input_
+
+ def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):
+ """Broadcast the input object.
+ NOTE: `src` is the local rank of the source rank.
+ """
+
+ # Bypass the function if we are using only 1 GPU.
+ if self.world_size == 1:
+ return obj
+ if self.shm_broadcaster is not None:
+ return self.shm_broadcaster.broadcast_object(obj)
+ if self.rank_in_group == src:
+ torch.distributed.broadcast_object_list(
+ [obj], src=self.ranks[src], group=self.cpu_group
+ )
+ return obj
+ else:
+ recv = [None]
+ torch.distributed.broadcast_object_list(
+ recv, src=self.ranks[src], group=self.cpu_group
+ )
+ return recv[0]
+
+ def broadcast_object_list(
+ self, obj_list: List[Any], src: int = 0, group: Optional[ProcessGroup] = None
+ ):
+ """Broadcast the input object list.
+ NOTE: `src` is the local rank of the source rank.
+ """
+
+ # Bypass the function if we are using only 1 GPU.
+ if self.world_size == 1:
+ return obj_list
+ # Broadcast.
+ torch.distributed.broadcast_object_list(
+ obj_list, src=self.ranks[src], group=self.device_group
+ )
+ return obj_list
+
+ def send_object(self, obj: Any, dst: int) -> None:
+ """Send the input object list to the destination rank."""
+ """NOTE: `dst` is the local rank of the destination rank."""
+
+ # Serialize object to tensor and get the size as well
+ object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8)
+
+ size_tensor = torch.tensor(
+ [object_tensor.numel()], dtype=torch.long, device="cpu"
+ )
+
+ # Send object size
+
+ torch.distributed.send(size_tensor, dst=self.ranks[dst], group=self.cpu_group)
+
+ # Send object
+ torch.distributed.send(object_tensor, dst=self.ranks[dst], group=self.cpu_group)
+
+ return None
+
+ def recv_object(self, src: int) -> Any:
+ """Receive the input object list from the source rank."""
+ """NOTE: `src` is the local rank of the source rank."""
+
+ size_tensor = torch.empty(1, dtype=torch.long, device="cpu")
+
+ # Receive object size
+ rank_size = torch.distributed.recv(
+ size_tensor, src=self.ranks[src], group=self.cpu_group
+ )
+
+ # Tensor to receive serialized objects into.
+ object_tensor = torch.empty( # type: ignore[call-overload]
+ size_tensor.item(), # type: ignore[arg-type]
+ dtype=torch.uint8,
+ device="cpu",
+ )
+
+ rank_object = torch.distributed.recv(
+ object_tensor, src=self.ranks[src], group=self.cpu_group
+ )
+
+ obj = pickle.loads(object_tensor.numpy().tobytes())
+
+ return obj
+
+ def broadcast_tensor_dict(
+ self,
+ tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None,
+ src: int = 0,
+ group: Optional[ProcessGroup] = None,
+ metadata_group: Optional[ProcessGroup] = None,
+ ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+ """Broadcast the input tensor dictionary.
+ NOTE: `src` is the local rank of the source rank.
+ """
+ # Bypass the function if we are using only 1 GPU.
+ if not torch.distributed.is_initialized() or self.world_size == 1:
+ return tensor_dict
+
+ group = self.device_group
+ metadata_group = self.cpu_group
+ src = self.ranks[src]
+
+ rank = self.rank
+ if rank == src:
+ metadata_list: List[Tuple[Any, Any]] = []
+ metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
+ # `metadata_list` lives in CPU memory.
+ # `broadcast_object_list` has serialization & deserialization,
+ # all happening on CPU. Therefore, we can use the CPU group.
+ self.broadcast_object(metadata_list, src=src)
+ async_handles = []
+ for tensor in tensor_list:
+ if tensor.numel() == 0:
+ # Skip broadcasting empty tensors.
+ continue
+ if tensor.is_cpu:
+ # use metadata_group for CPU tensors
+ handle = torch.distributed.broadcast(
+ tensor, src=src, group=metadata_group, async_op=True
+ )
+ else:
+ # use group for GPU tensors
+ handle = torch.distributed.broadcast(
+ tensor, src=src, group=group, async_op=True
+ )
+ async_handles.append(handle)
+ for async_handle in async_handles:
+ async_handle.wait()
+
+ else:
+ metadata_list = self.broadcast_object(None, src=src)
+ tensor_dict = {}
+ async_handles = []
+ for key, value in metadata_list:
+ if isinstance(value, TensorMetadata):
+ tensor = torch.empty(
+ value.size, dtype=value.dtype, device=value.device
+ )
+ if tensor.numel() == 0:
+ # Skip broadcasting empty tensors.
+ _update_nested_dict(tensor_dict, key, tensor)
+ continue
+ if tensor.is_cpu:
+ # use metadata_group for CPU tensors
+ handle = torch.distributed.broadcast(
+ tensor, src=src, group=metadata_group, async_op=True
+ )
+ else:
+ # use group for GPU tensors
+ handle = torch.distributed.broadcast(
+ tensor, src=src, group=group, async_op=True
+ )
+ async_handles.append(handle)
+ _update_nested_dict(tensor_dict, key, tensor)
+ else:
+ _update_nested_dict(tensor_dict, key, value)
+ for async_handle in async_handles:
+ async_handle.wait()
+ return tensor_dict
+
+ def send_tensor_dict(
+ self,
+ tensor_dict: Dict[str, Union[torch.Tensor, Any]],
+ dst: Optional[int] = None,
+ ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+ """Send the input tensor dictionary.
+ NOTE: `dst` is the local rank of the source rank.
+ """
+ # Bypass the function if we are using only 1 GPU.
+ if not torch.distributed.is_initialized() or self.world_size == 1:
+ return tensor_dict
+
+ group = self.device_group
+ metadata_group = self.cpu_group
+
+ if dst is None:
+ dst = self.group_next_rank
+
+ metadata_list: List[Tuple[Any, Any]] = []
+ metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
+ # `metadata_list` lives in CPU memory.
+ # `send_object_list` has serialization & deserialization,
+ # all happening on CPU. Therefore, we can use the CPU group.
+ self.send_object(metadata_list, dst=dst)
+ for tensor in tensor_list:
+ if tensor.numel() == 0:
+ # Skip sending empty tensors.
+ continue
+ if tensor.is_cpu:
+ # use metadata_group for CPU tensors
+ torch.distributed.send(
+ tensor, dst=self.ranks[dst], group=metadata_group
+ )
+ else:
+ # use group for GPU tensors
+ torch.distributed.send(tensor, dst=self.ranks[dst], group=group)
+ return None
+
+ def recv_tensor_dict(
+ self, src: Optional[int] = None
+ ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+ """Recv the input tensor dictionary.
+ NOTE: `src` is the local rank of the source rank.
+ """
+ # Bypass the function if we are using only 1 GPU.
+ if not torch.distributed.is_initialized() or self.world_size == 1:
+ return None
+
+ group = self.device_group
+ metadata_group = self.cpu_group
+
+ if src is None:
+ src = self.group_prev_rank
+
+ recv_metadata_list = self.recv_object(src=src)
+ tensor_dict: Dict[str, Any] = {}
+ for key, value in recv_metadata_list:
+ if isinstance(value, TensorMetadata):
+ tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)
+ if tensor.numel() == 0:
+ # Skip broadcasting empty tensors.
+ _update_nested_dict(tensor_dict, key, tensor)
+ continue
+ if tensor.is_cpu:
+ # use metadata_group for CPU tensors
+ torch.distributed.recv(
+ tensor, src=self.ranks[src], group=metadata_group
+ )
+ else:
+ # use group for GPU tensors
+ torch.distributed.recv(tensor, src=self.ranks[src], group=group)
+ _update_nested_dict(tensor_dict, key, tensor)
+ else:
+ _update_nested_dict(tensor_dict, key, value)
+ return tensor_dict
+
+ def barrier(self):
+ """Barrier synchronization among the group.
+ NOTE: don't use `device_group` here! `barrier` in NCCL is
+ terrible because it is internally a broadcast operation with
+ secretly created GPU tensors. It is easy to mess up the current
+ device. Use the CPU group instead.
+ """
+ torch.distributed.barrier(group=self.cpu_group)
+
+ def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+ """Sends a tensor to the destination rank in a non-blocking way"""
+ """NOTE: `dst` is the rank_in_group of the destination rank."""
+ if dst is None:
+ dst = self.group_next_rank
+
+ torch.distributed.send(
+ tensor,
+ self.ranks[dst],
+ group=(
+ self.device_groups[self.rank_in_group % 2]
+ if self.world_size == 2
+ else self.device_group
+ ),
+ )
+
+ def recv(
+ self, size: torch.Size, dtype: torch.dtype, src: Optional[int] = None
+ ) -> torch.Tensor:
+ """Receives a tensor from the src rank."""
+ """NOTE: `src` is the rank_in_group of the source rank."""
+ if src is None:
+ src = self.group_prev_rank
+
+ tensor = torch.empty(size, dtype=dtype, device=self.device)
+ torch.distributed.recv(
+ tensor,
+ self.ranks[src],
+ (
+ self.device_groups[(self.rank_in_group + 1) % 2]
+ if self.world_size == 2
+ else self.device_group
+ ),
+ )
+ return tensor
+
+ def destroy(self):
+ if self.device_group is not None:
+ torch.distributed.destroy_process_group(self.device_group)
+ self.device_group = None
+ if self.cpu_group is not None:
+ torch.distributed.destroy_process_group(self.cpu_group)
+ self.cpu_group = None
+
+
+class SequenceParallelGroupCoordinator(GroupCoordinator):
+ def __init__(
+ self,
+ group_ranks: List[List[int]],
+ local_rank: int,
+ torch_distributed_backend: Union[str, Backend],
+ **kwargs,
+ ):
+ super().__init__(
+ group_ranks=group_ranks,
+ local_rank=local_rank,
+ torch_distributed_backend=torch_distributed_backend,
+ )
+ ulysses_group = kwargs.get("ulysses_group", None)
+ ring_group = kwargs.get("ring_group", None)
+
+ self.ulysses_world_size = self.ring_world_size = 1
+ self.ulysses_rank = self.ring_rank = 0
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/helpers.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/helpers.py
new file mode 100644
index 0000000000000000000000000000000000000000..8768812d1aa27bd38f57d78c69eda9b106dd64eb
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/helpers.py
@@ -0,0 +1,41 @@
+import collections.abc
+
+from itertools import repeat
+
+
+def _ntuple(n):
+ def parse(x):
+ if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+ x = tuple(x)
+ if len(x) == 1:
+ x = tuple(repeat(x[0], n))
+ return x
+ return tuple(repeat(x, n))
+
+ return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+
+
+def as_tuple(x):
+ if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+ return tuple(x)
+ if x is None or isinstance(x, (int, float, str)):
+ return (x,)
+ else:
+ raise ValueError(f"Unknown type {type(x)}")
+
+
+def as_list_of_2tuple(x):
+ x = as_tuple(x)
+ if len(x) == 1:
+ x = (x[0], x[0])
+ assert len(x) % 2 == 0, f"Expect even length, got {len(x)}."
+ lst = []
+ for i in range(0, len(x), 2):
+ lst.append((x[i], x[i + 1]))
+ return lst
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/parallel_mgr.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/parallel_mgr.py
new file mode 100644
index 0000000000000000000000000000000000000000..dee8c0572e126602de064bdc78b3b39ea5c3ac68
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/parallel_mgr.py
@@ -0,0 +1,393 @@
+import os
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch.distributed as dist
+import torch_npu
+from diffusers.utils import logging
+
+from .group_coordinator import GroupCoordinator, SequenceParallelGroupCoordinator
+from .utils import RankGenerator
+
+logger = logging.get_logger(__name__)
+
+_WORLD: Optional[GroupCoordinator] = None
+_TP: Optional[GroupCoordinator] = None
+_SP: Optional[SequenceParallelGroupCoordinator] = None
+_ULYSSES = None
+_RING = None
+_CFG: Optional[GroupCoordinator] = None
+_VAE = None
+
+
+@dataclass
+class ParallelConfig:
+ tp_degree: int = 1
+ sp_degree: int = 1
+ use_cfg_parallel: bool = False
+ tp_degree: int = 1
+ vae_parallel_size: int = 0
+ world_size: int = 1
+
+ def __post_init__(self):
+ if self.use_cfg_parallel:
+ self.cfg_degree = 2
+ else:
+ self.cfg_degree = 1
+ if not self.tp_degree * self.sp_degree * self.cfg_degree <= self.world_size:
+ logger.error(
+ "tp_degree * sp_degree * cfg_degree must be less than or equal to world_size"
+ )
+ if not (self.world_size % (self.tp_degree * self.sp_degree * self.cfg_degree) == 0):
+ logger.error("world_size must be divisible by tp_degree * sp_degree * cfg_degree")
+ if self.vae_parallel_size == 0:
+ self.vae_parallel_size = self.world_size
+
+
+# * QUERY
+def get_world_group() -> GroupCoordinator:
+ if _WORLD is None:
+ logger.error("world group is not initialized")
+ return _WORLD
+
+
+# TP
+def get_tp_group() -> GroupCoordinator:
+ if _TP is None:
+ logger.error("tensor model parallel group is not initialized")
+ return _TP
+
+
+def get_tensor_parallel_state():
+ """Return state for the tensor parallel group."""
+ return _TP is not None
+
+
+def get_tensor_model_parallel_world_size():
+ """Return world size for the tensor model parallel group."""
+ if not get_tensor_parallel_state():
+ return 1
+ return get_tp_group().world_size
+
+
+def get_tensor_model_parallel_rank():
+ """Return my rank for the tensor model parallel group."""
+ if not get_tensor_parallel_state():
+ return 1
+ return get_tp_group().rank_in_group
+
+
+# SP
+def get_sp_group() -> SequenceParallelGroupCoordinator:
+ if _SP is None:
+ logger.error("sequence model parallel group is not initialized")
+ return _SP
+
+
+def get_sequence_parallel_state():
+ """Return state for the sequence parallel group."""
+ return _SP is not None
+
+
+def get_sequence_parallel_world_size():
+ """Return world size for the sequence parallel group."""
+ if not get_sequence_parallel_state():
+ return 1
+ return get_sp_group().world_size
+
+
+def get_sequence_parallel_rank():
+ """Return my rank for the sequence parallel group."""
+ if not get_sequence_parallel_state():
+ return 0
+ return get_sp_group().rank_in_group
+
+
+# CFG
+def get_cfg_group() -> GroupCoordinator:
+ if _CFG is None:
+ logger.error("classifier_free_guidance parallel group is not initialized")
+ return _CFG
+
+
+def get_cfg_state():
+ """Return state for the sequence parallel group."""
+ return _CFG is not None
+
+
+def get_classifier_free_guidance_world_size():
+ """Return world size for the classifier_free_guidance parallel group."""
+ if not get_cfg_state():
+ return 1
+ return get_cfg_group().world_size
+
+
+def get_classifier_free_guidance_rank():
+ """Return my rank for the classifier_free_guidance parallel group."""
+ if not get_cfg_state():
+ return 0
+ return get_cfg_group().rank_in_group
+
+
+# Add VAE getter functions
+def get_vae_parallel_group():
+ if _VAE is None:
+ logger.error("VAE parallel group is not initialized")
+ return _VAE
+
+
+def get_vae_parallel_world_size():
+ """Return world size for the VAE parallel group."""
+ if _VAE is None:
+ return 1
+ return dist.get_world_size(group=_VAE)
+
+
+def get_vae_parallel_rank():
+ """Return my rank for the VAE parallel group."""
+ if _VAE is None:
+ return 0
+ return dist.get_rank(group=_VAE)
+
+
+# * SET
+def init_world_group(
+ ranks: List[int], local_rank: int, backend: str
+) -> GroupCoordinator:
+ return GroupCoordinator(
+ group_ranks=[ranks],
+ local_rank=local_rank,
+ torch_distributed_backend=backend,
+ )
+
+
+def init_distributed_environment(
+ world_size: int = -1,
+ rank: int = -1,
+ distributed_init_method: str = "env://",
+ local_rank: int = -1,
+ backend: str = "hccl",
+):
+ logger.debug(
+ "world_size=%d rank=%d local_rank=%d " "distributed_init_method=%s backend=%s",
+ world_size,
+ rank,
+ local_rank,
+ distributed_init_method,
+ backend,
+ )
+ if not dist.is_initialized():
+ if distributed_init_method is None:
+ logger.error(
+ "distributed_init_method must be provided when initializing "
+ "distributed environment"
+ )
+ # this backend is used for WORLD
+ dist.init_process_group(
+ backend=backend,
+ init_method=distributed_init_method,
+ world_size=world_size,
+ rank=rank,
+ )
+ torch_npu.npu.set_device(dist.get_rank() % torch_npu.npu.device_count())
+ # set the local rank
+ # local_rank is not available in torch ProcessGroup,
+ # see https://github.com/pytorch/pytorch/issues/122816
+ if local_rank == -1:
+ # local rank not set, this usually happens in single-node
+ # setting, where we can use rank as local rank
+ if distributed_init_method == "env://":
+ local_rank = int(os.getenv('LOCAL_RANK', 0))
+ else:
+ local_rank = rank
+ global _WORLD
+ if _WORLD is None:
+ ranks = list(range(dist.get_world_size()))
+ _WORLD = init_world_group(ranks, local_rank, backend)
+ else:
+ if not _WORLD.world_size == dist.get_world_size():
+ logger.error("world group already initialized with a different world size")
+
+
+def model_parallel_is_initialized():
+ """Check if tensor and pipeline parallel groups are initialized."""
+ return (
+ _CFG is not None
+ and _SP is not None
+ and _TP is not None
+ )
+
+
+def init_model_parallel_group(
+ group_ranks: List[List[int]],
+ local_rank: int,
+ backend: str,
+ parallel_mode: str,
+ **kwargs,
+) -> GroupCoordinator:
+ if parallel_mode not in [
+ "tensor",
+ "sequence",
+ "classifier_free_guidance",
+ ]:
+ logger.error(f"parallel_mode {parallel_mode} is not supported")
+ if parallel_mode == "sequence":
+ return SequenceParallelGroupCoordinator(
+ group_ranks=group_ranks,
+ local_rank=local_rank,
+ torch_distributed_backend=backend,
+ **kwargs,
+ )
+ else:
+ return GroupCoordinator(
+ group_ranks=group_ranks,
+ local_rank=local_rank,
+ torch_distributed_backend=backend,
+ )
+
+
+def init_vae_group(
+ vae_parallel_size: int,
+ backend: str,
+):
+ # Initialize VAE group first
+ global _VAE
+ if _VAE is not None:
+ logger.error("VAE parallel group is already initialized")
+ vae_ranks = list(range(vae_parallel_size))
+ _VAE = dist.new_group(
+ ranks=vae_ranks, backend=backend
+ )
+
+
+def initialize_model_parallel(
+ classifier_free_guidance_degree: int = 1,
+ sequence_parallel_degree: int = 1,
+ ulysses_degree: int = 1,
+ ring_degree: int = 1,
+ tensor_parallel_degree: int = 1,
+ vae_parallel_size: int = 0,
+ backend: Optional[str] = None,
+) -> None:
+ """
+ Initialize model parallel groups.
+
+ Arguments:
+ classifier_free_guidance_degree: number of NPUs used for Classifier Free Guidance (CFG)
+ sequence_parallel_degree: number of NPUs used for sequence parallelism.
+ ulysses_degree: number of NPUs used for ulysses sequence parallelism.
+ ring_degree: number of NPUs used for ring sequence parallelism.
+ tensor_parallel_degree: number of NPUs used for tensor parallelism.
+ backend: distributed backend of pytorch collective comm.
+ """
+ # Get world size and rank. Ensure some consistencies.
+ if not dist.is_initialized():
+ logger.error("dist is not initialized")
+ world_size: int = dist.get_world_size()
+ backend = backend or dist.get_backend(get_world_group().device_group)
+
+ if (
+ world_size
+ != classifier_free_guidance_degree
+ * tensor_parallel_degree
+ * sequence_parallel_degree
+ ):
+ raise RuntimeError(
+ f"world_size ({world_size}) is not equal to "
+ f"tensor_parallel_degree ({tensor_parallel_degree}) x "
+ f"sequence_parallel_degree ({sequence_parallel_degree}) x"
+ f"classifier_free_guidance_degree "
+ f"({classifier_free_guidance_degree}) x"
+ )
+
+ rank_generator: RankGenerator = RankGenerator(
+ tensor_parallel_degree,
+ sequence_parallel_degree,
+ classifier_free_guidance_degree,
+ "tp-sp-cfg",
+ )
+
+ global _CFG
+ if _CFG is not None:
+ logger.error("classifier_free_guidance group is already initialized")
+ _CFG = init_model_parallel_group(
+ group_ranks=rank_generator.get_ranks("cfg"),
+ local_rank=get_world_group().local_rank,
+ backend=backend,
+ parallel_mode="classifier_free_guidance",
+ )
+
+ global _SP
+ if _SP is not None:
+ logger.error("sequence parallel group is already initialized")
+
+ _SP = init_model_parallel_group(
+ group_ranks=rank_generator.get_ranks("sp"),
+ local_rank=get_world_group().local_rank,
+ backend=backend,
+ parallel_mode="sequence",
+ )
+
+ global _TP
+ if _TP is not None:
+ logger.error("Tensor parallel group is already initialized")
+ _TP = init_model_parallel_group(
+ group_ranks=rank_generator.get_ranks("tp"),
+ local_rank=get_world_group().local_rank,
+ backend=backend,
+ parallel_mode="tensor",
+ )
+
+ if vae_parallel_size > 0:
+ init_vae_group(vae_parallel_size, backend)
+
+
+def destroy_model_parallel():
+ """Set the groups to none and destroy them."""
+ global _CFG
+ if _CFG:
+ _CFG.destroy()
+ _CFG = None
+
+ global _SP
+ if _SP:
+ _SP.destroy()
+ _SP = None
+
+ global _TP
+ if _TP:
+ _TP.destroy()
+ _TP = None
+
+ global _VAE
+ if _VAE:
+ dist.destroy_process_group(_VAE)
+ _VAE = None
+
+
+def destroy_distributed_environment():
+ global _WORLD
+ if _WORLD:
+ _WORLD.destroy()
+ _WORLD = None
+ if dist.is_initialized():
+ dist.destroy_process_group()
+
+
+def init_parallel_env(parallel_config: ParallelConfig):
+ if not model_parallel_is_initialized():
+ logger.warning("Model parallel is not initialized, initializing...")
+ if not dist.is_initialized():
+ init_distributed_environment()
+ initialize_model_parallel(
+ classifier_free_guidance_degree=parallel_config.cfg_degree,
+ sequence_parallel_degree=parallel_config.sp_degree,
+ tensor_parallel_degree=parallel_config.tp_degree,
+ vae_parallel_size=parallel_config.vae_parallel_size
+ )
+
+
+def finalize_parallel_env():
+ if model_parallel_is_initialized():
+ destroy_model_parallel()
+ destroy_distributed_environment()
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a80abe238b323aa4906152e12c8d201693a8d659
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/preprocess_text_encoder_tokenizer_utils.py
@@ -0,0 +1,48 @@
+import argparse
+
+import torch
+import torch_npu
+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("npu")
+
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/utils.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..26b0478ec47690533d83ac3307c43736090ac368
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/utils/utils.py
@@ -0,0 +1,155 @@
+from typing import List
+
+from diffusers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+def generate_masked_orthogonal_rank_groups(
+ world_size: int, parallel_size: List[int], mask: List[bool]
+) -> List[List[int]]:
+ """Generate orthogonal parallel groups based on the parallel size and mask.
+
+ Arguments:
+ world_size (int): world size
+
+ parallel_size (List[int]):
+ The parallel size of each orthogonal parallel type. For example, if
+ tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4,
+ and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4].
+
+ mask (List[bool]):
+ The mask controls which parallel methods the generated groups represent. If mask[i] is
+ True, it means the generated group contains the i-th parallelism method. For example,
+ if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then
+ the generated group is the `tp-dp` group, if the mask = [False, True, False], then the
+ generated group is the `pp` group.
+ """
+
+ def prefix_product(a: List[int], init=1) -> List[int]:
+ r = [init]
+ for v in a:
+ init = init * v
+ r.append(init)
+ return r
+
+ def inner_product(a: List[int], b: List[int]) -> int:
+ return sum([x * y for x, y in zip(a, b)])
+
+ def decompose(index, shape, stride=None):
+ """
+ This function solve the math problem below:
+ There is an equation:
+ index = sum(idx[i] * stride[i])
+ And given the value of index, stride.
+ Return the idx.
+ This function will used to get the pp/dp/pp_rank
+ from group_index and rank_in_group.
+ """
+ if stride is None:
+ stride = prefix_product(shape)
+ idx = [(index // d) % s for s, d in zip(shape, stride)]
+ # stride is a prefix_product result. And the value of stride[-1]
+ # is not used.
+ if not (
+ sum([x * y for x, y in zip(idx, stride[:-1])]) == index
+ ):
+ logger.error("idx {} with shape {} mismatch the return idx {}".format(index, shape, idx))
+ return idx
+
+ masked_shape = [s for s, m in zip(parallel_size, mask) if m]
+ unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m]
+
+ global_stride = prefix_product(parallel_size)
+ masked_stride = [d for d, m in zip(global_stride, mask) if m]
+ unmasked_stride = [d for d, m in zip(global_stride, mask) if not m]
+
+ group_size = prefix_product(masked_shape)[-1]
+ num_of_group = world_size // group_size
+
+ ranks = []
+ for group_index in range(num_of_group):
+ # get indices from unmaksed for group_index.
+ decomposed_group_idx = decompose(group_index, unmasked_shape)
+ rank = []
+ for rank_in_group in range(group_size):
+ # get indices from masked for rank_in_group.
+ decomposed_rank_idx = decompose(rank_in_group, masked_shape)
+ rank.append(
+ inner_product(decomposed_rank_idx, masked_stride)
+ + inner_product(decomposed_group_idx, unmasked_stride)
+ )
+ ranks.append(rank)
+ return ranks
+
+
+class RankGenerator(object):
+ def __init__(
+ self,
+ tp: int,
+ sp: int,
+ cfg: int,
+ order: str,
+ rank_offset: int = 0,
+ ) -> None:
+ self.tp = tp
+ self.sp = sp
+ self.cfg = cfg
+ self.rank_offset = rank_offset
+ self.world_size = tp * sp * cfg
+
+ self.name_to_size = {
+ "tp": self.tp,
+ "sp": self.sp,
+ "cfg": self.cfg,
+ }
+ order = order.lower()
+
+ for name in self.name_to_size.keys():
+ if name not in order and self.name_to_size[name] != 1:
+ raise RuntimeError(
+ f"The size of ({name}) is ({self.name_to_size[name]}), but you haven't specified the order ({self.order})."
+ )
+ elif name not in order:
+ order = order + "-" + name
+
+ self.order = order
+ self.ordered_size = []
+
+ for token in order.split("-"):
+ self.ordered_size.append(self.name_to_size[token])
+
+ def get_mask(self, order: str, token: str):
+ ordered_token = order.split("-")
+ token = token.split("-")
+ mask = [False] * len(ordered_token)
+ for t in token:
+ mask[ordered_token.index(t)] = True
+ return mask
+
+ def get_ranks(self, token):
+ """Get rank group by input token.
+
+ Arguments:
+ token (str):
+ Specify the ranks type that want to get. If we want
+ to obtain multiple parallel types, we can use a hyphen
+ '-' to separate them. For example, if we want to obtain
+ the TP_DP group, the token should be 'tp-dp'.
+
+ independent_ep (bool: True):
+ This flag controls whether we treat EP and DP independently.
+ EP shares ranks with DP, if we want to get ranks related to
+ EP, we should set the flag. For example, get_ranks('dp', True)
+ will get DP modulo EP group, and get_ranks('dp', False) will
+ get full DP group.
+ """
+ mask = self.get_mask(self.order, token)
+ ranks = generate_masked_orthogonal_rank_groups(
+ self.world_size, self.ordered_size, mask
+ )
+ if self.rank_offset > 0:
+ for rank_group in ranks:
+ for i, _ in enumerate(rank_group):
+ rank_group[i] += self.rank_offset
+ return ranks
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..13f57c005ca43dc2c0a51a9fb2fc1c32ab456fd7
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/__init__.py
@@ -0,0 +1,63 @@
+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"]
+ if any(k.startswith("vae.") for k in ckpt.keys()):
+ ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")}
+ vae.load_state_dict(ckpt)
+
+ spatial_compression_ratio = vae.config.spatial_compression_ratio
+ time_compression_ratio = vae.config.time_compression_ratio
+
+ if vae_precision is not None:
+ vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision])
+
+ vae.requires_grad_(False)
+
+ if logger is not None:
+ logger.info(f"VAE to dtype: {vae.dtype}")
+
+ if device is not None:
+ vae = vae.to(device)
+
+ vae.eval()
+
+ return vae, vae_path, spatial_compression_ratio, time_compression_ratio
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/attention.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..b35e2ab660e8052d8f20a6f98dee60c4723f9615
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/attention.py
@@ -0,0 +1,109 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import torch_npu
+from diffusers.models.attention_processor import Attention
+
+MAX_TOKEN = 2147483647
+
+
+class AttnProcessor:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ scale = 1.0 / math.sqrt(head_dim)
+ hidden_states = torch_npu.npu_fusion_attention(
+ query, key, value,
+ head_num=attn.heads,
+ input_layout="BNSD",
+ scale=scale,
+ pse=None,
+ atten_mask=attention_mask,
+ pre_tockens=MAX_TOKEN,
+ next_tockens=MAX_TOKEN,
+ keep_prob=1.0,
+ sync=False
+ )[0]
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/autoencoder_kl_causal_3d.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/autoencoder_kl_causal_3d.py
new file mode 100644
index 0000000000000000000000000000000000000000..cf65f55d75adc07da6d253352b1de5f96b0f7a7e
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/autoencoder_kl_causal_3d.py
@@ -0,0 +1,669 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, Union
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+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 hyvideo.utils.parallel_mgr import get_vae_parallel_group, get_vae_parallel_world_size, get_vae_parallel_rank
+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
+ 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 = []
+ if get_vae_parallel_world_size() > 1:
+ world_size = get_vae_parallel_world_size()
+ rank = get_vae_parallel_rank()
+ tile_list = []
+ decoded_list = []
+ for i in range(0, z.shape[-2], overlap_size):
+ for j in range(0, z.shape[-1], overlap_size):
+ h_start, h_end = i, i + self.tile_latent_min_size
+ w_start, w_end = j, j + self.tile_latent_min_size
+ tile_list.append((h_start, h_end, w_start, w_end))
+ tile_num = len(tile_list)
+ for idx in range(0, tile_num, world_size):
+ tile_id = idx + rank
+ if tile_id < tile_num:
+ h_start, h_end, w_start, w_end = tile_list[tile_id]
+ tile = z[:, :, :, h_start: h_end, w_start: w_end]
+ else:
+ h_start, h_end, w_start, w_end = tile_list[0]
+ tile = torch.zeros_like(z[:, :, :, h_start: h_end, w_start: w_end])
+ tile = self.post_quant_conv(tile)
+ decoded = self.decoder(tile)
+ padding_right = self.tile_sample_min_size - decoded.shape[-1]
+ padding_bottom = self.tile_sample_min_size - decoded.shape[-2]
+ if padding_right > 0 or padding_bottom > 0:
+ pad_shape = (0, padding_right, 0, padding_bottom)
+ decoded = F.pad(decoded, pad_shape, "constant", 0)
+
+ decoded_size = list(decoded.size())
+ decoded_size[0] *= world_size
+ decoded_tensor = torch.empty(
+ decoded_size, dtype=decoded.dtype, device=decoded.device
+ )
+ # All-gather.
+ torch.distributed.all_gather_into_tensor(
+ decoded_tensor, decoded, group=get_vae_parallel_group()
+ )
+ decoded_tensors = decoded_tensor.chunk(world_size, dim=0)
+ for i, tensor in enumerate(decoded_tensors):
+ tile_id_ = idx + i
+ if tile_id_ < tile_num:
+ h0, h1, w0, w1 = tile_list[tile_id_]
+ h1 = h1 if h1 < z.shape[-2] else z.shape[-2]
+ w1 = w1 if w1 < z.shape[-1] else z.shape[-1]
+ decoded_h = 8 * (h1 - h0)
+ decoded_w = 8 * (w1 - w0)
+ decoded_list.append(tensor[:, :, :, :decoded_h, :decoded_w])
+ row_cnt = int(math.ceil(z.shape[-2] / overlap_size))
+ column_cnt = int(math.ceil(z.shape[-1] / overlap_size))
+ for i in range(row_cnt):
+ row = []
+ for j in range(column_cnt):
+ row.append(decoded_list[i * column_cnt + j])
+ rows.append(row)
+ else:
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/unet_causal_3d_blocks.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/unet_causal_3d_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..a6514f19c6b9d54e36459328bbae741bd96ed5a0
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/unet_causal_3d_blocks.py
@@ -0,0 +1,770 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Modified from diffusers==0.29.2
+#
+# ==============================================================================
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import Attention
+from diffusers.models.attention_processor import SpatialNorm
+from diffusers.models.normalization import AdaGroupNorm
+from diffusers.models.normalization import RMSNorm
+from diffusers.utils import logging
+from einops import rearrange
+from torch import nn
+
+from .attention import AttnProcessor
+
+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
+ indices = torch.arange(seq_len, dtype=torch.int64, device=device)
+ frame_indices = indices // n_hw
+ mask = (frame_indices.unsqueeze(1) >= (indices // n_hw))
+ mask = torch.where(mask, torch.tensor(0., dtype=torch.uint8, device=device),
+ torch.tensor(1, dtype=torch.uint8, device=device))
+
+ if batch_size is not None:
+ mask = mask.unsqueeze(0).expand(batch_size, -1, -1)
+ return mask
+
+
+class CausalConv3d(nn.Module):
+ """
+ Implements a causal 3D convolution layer where each position only depends on previous timesteps and current spatial locations.
+ This maintains temporal causality in video generation tasks.
+ """
+
+ def __init__(
+ self,
+ chan_in,
+ chan_out,
+ kernel_size: Union[int, Tuple[int, int, int]],
+ stride: Union[int, Tuple[int, int, int]] = 1,
+ dilation: Union[int, Tuple[int, int, int]] = 1,
+ pad_mode='replicate',
+ **kwargs
+ ):
+ super().__init__()
+
+ self.pad_mode = pad_mode
+ padding = (
+ kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size - 1, 0) # W, H, T
+ self.time_causal_padding = padding
+
+ self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+ def forward(self, x):
+ x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+ return self.conv(x)
+
+
+class UpsampleCausal3D(nn.Module):
+ """
+ A 3D upsampling layer with an optional convolution.
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ use_conv: bool = False,
+ use_conv_transpose: bool = False,
+ out_channels: Optional[int] = None,
+ name: str = "conv",
+ kernel_size: Optional[int] = None,
+ padding=1,
+ norm_type=None,
+ eps=None,
+ elementwise_affine=None,
+ bias=True,
+ interpolate=True,
+ upsample_factor=(2, 2, 2),
+ ):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.use_conv_transpose = use_conv_transpose
+ self.name = name
+ self.interpolate = interpolate
+ self.upsample_factor = upsample_factor
+
+ if norm_type == "ln_norm":
+ self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+ elif norm_type == "rms_norm":
+ self.norm = RMSNorm(channels, eps, elementwise_affine)
+ elif norm_type is None:
+ self.norm = None
+ else:
+ raise ValueError(f"unknown norm_type: {norm_type}")
+
+ conv = None
+ if use_conv_transpose:
+ raise NotImplementedError
+ elif use_conv:
+ if kernel_size is None:
+ kernel_size = 3
+ conv = CausalConv3d(self.channels, self.out_channels, kernel_size=kernel_size, bias=bias)
+
+ if name == "conv":
+ self.conv = conv
+ else:
+ self.Conv2d_0 = conv
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ output_size: Optional[int] = None,
+ scale: float = 1.0,
+ ) -> torch.FloatTensor:
+ assert hidden_states.shape[1] == self.channels
+
+ if self.norm is not None:
+ raise NotImplementedError
+
+ if self.use_conv_transpose:
+ return self.conv(hidden_states)
+
+ # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+ dtype = hidden_states.dtype
+ if dtype == torch.bfloat16:
+ hidden_states = hidden_states.to(torch.float32)
+
+ # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+ if hidden_states.shape[0] >= 64:
+ hidden_states = hidden_states.contiguous()
+
+ # if `output_size` is passed we force the interpolation output
+ # size and do not make use of `scale_factor=2`
+ if self.interpolate:
+ B, C, T, H, W = hidden_states.shape
+ first_h, other_h = hidden_states.split((1, T - 1), dim=2)
+ if output_size is None:
+ if T > 1:
+ other_h = F.interpolate(other_h, scale_factor=self.upsample_factor, mode="nearest")
+
+ first_h = first_h.squeeze(2)
+ first_h = F.interpolate(first_h, scale_factor=self.upsample_factor[1:], mode="nearest")
+ first_h = first_h.unsqueeze(2)
+ else:
+ raise NotImplementedError
+
+ if T > 1:
+ hidden_states = torch.cat((first_h, other_h), dim=2)
+ else:
+ hidden_states = first_h
+
+ # If the input is bfloat16, we cast back to bfloat16
+ if dtype == torch.bfloat16:
+ hidden_states = hidden_states.to(dtype)
+
+ if self.use_conv:
+ if self.name == "conv":
+ hidden_states = self.conv(hidden_states)
+ else:
+ hidden_states = self.Conv2d_0(hidden_states)
+
+ return hidden_states
+
+
+class DownsampleCausal3D(nn.Module):
+ """
+ A 3D downsampling layer with an optional convolution.
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ use_conv: bool = False,
+ out_channels: Optional[int] = None,
+ padding: int = 1,
+ name: str = "conv",
+ kernel_size=3,
+ norm_type=None,
+ eps=None,
+ elementwise_affine=None,
+ bias=True,
+ stride=2,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.padding = padding
+ stride = stride
+ self.name = name
+
+ if norm_type == "ln_norm":
+ self.norm = nn.LayerNorm(channels, eps, elementwise_affine)
+ elif norm_type == "rms_norm":
+ self.norm = RMSNorm(channels, eps, elementwise_affine)
+ elif norm_type is None:
+ self.norm = None
+ else:
+ raise ValueError(f"unknown norm_type: {norm_type}")
+
+ if use_conv:
+ conv = CausalConv3d(
+ self.channels, self.out_channels, kernel_size=kernel_size, stride=stride, bias=bias
+ )
+ else:
+ raise NotImplementedError
+
+ if name == "conv":
+ self.Conv2d_0 = conv
+ self.conv = conv
+ elif name == "Conv2d_0":
+ self.conv = conv
+ else:
+ self.conv = conv
+
+ def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
+ assert hidden_states.shape[1] == self.channels
+
+ if self.norm is not None:
+ hidden_states = self.norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+ assert hidden_states.shape[1] == self.channels
+
+ hidden_states = self.conv(hidden_states)
+
+ return hidden_states
+
+
+class ResnetBlockCausal3D(nn.Module):
+ r"""
+ A Resnet block.
+ """
+
+ def __init__(
+ self,
+ *,
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ conv_shortcut: bool = False,
+ dropout: float = 0.0,
+ temb_channels: int = 512,
+ groups: int = 32,
+ groups_out: Optional[int] = None,
+ pre_norm: bool = True,
+ eps: float = 1e-6,
+ non_linearity: str = "swish",
+ skip_time_act: bool = False,
+ # default, scale_shift, ada_group, spatial
+ time_embedding_norm: str = "default",
+ kernel: Optional[torch.FloatTensor] = None,
+ output_scale_factor: float = 1.0,
+ use_in_shortcut: Optional[bool] = None,
+ up: bool = False,
+ down: bool = False,
+ conv_shortcut_bias: bool = True,
+ conv_3d_out_channels: Optional[int] = None,
+ ):
+ super().__init__()
+ self.pre_norm = pre_norm
+ self.pre_norm = True
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+ self.up = up
+ self.down = down
+ self.output_scale_factor = output_scale_factor
+ self.time_embedding_norm = time_embedding_norm
+ self.skip_time_act = skip_time_act
+
+ linear_cls = nn.Linear
+
+ if groups_out is None:
+ groups_out = groups
+
+ if self.time_embedding_norm == "ada_group":
+ self.norm1 = AdaGroupNorm(temb_channels, in_channels, groups, eps=eps)
+ elif self.time_embedding_norm == "spatial":
+ self.norm1 = SpatialNorm(in_channels, temb_channels)
+ else:
+ self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+ self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, stride=1)
+
+ if temb_channels is not None:
+ if self.time_embedding_norm == "default":
+ self.time_emb_proj = linear_cls(temb_channels, out_channels)
+ elif self.time_embedding_norm == "scale_shift":
+ self.time_emb_proj = linear_cls(temb_channels, 2 * out_channels)
+ elif self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+ self.time_emb_proj = None
+ else:
+ raise ValueError(f"Unknown time_embedding_norm : {self.time_embedding_norm} ")
+ else:
+ self.time_emb_proj = None
+
+ if self.time_embedding_norm == "ada_group":
+ self.norm2 = AdaGroupNorm(temb_channels, out_channels, groups_out, eps=eps)
+ elif self.time_embedding_norm == "spatial":
+ self.norm2 = SpatialNorm(out_channels, temb_channels)
+ else:
+ self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+
+ self.dropout = torch.nn.Dropout(dropout)
+ conv_3d_out_channels = conv_3d_out_channels or out_channels
+ self.conv2 = CausalConv3d(out_channels, conv_3d_out_channels, kernel_size=3, stride=1)
+
+ self.nonlinearity = get_activation(non_linearity)
+
+ self.upsample = self.downsample = None
+ if self.up:
+ self.upsample = UpsampleCausal3D(in_channels, use_conv=False)
+ elif self.down:
+ self.downsample = DownsampleCausal3D(in_channels, use_conv=False, name="op")
+
+ self.use_in_shortcut = self.in_channels != conv_3d_out_channels if use_in_shortcut is None else use_in_shortcut
+
+ self.conv_shortcut = None
+ if self.use_in_shortcut:
+ self.conv_shortcut = CausalConv3d(
+ in_channels,
+ conv_3d_out_channels,
+ kernel_size=1,
+ stride=1,
+ bias=conv_shortcut_bias,
+ )
+
+ def forward(
+ self,
+ input_tensor: torch.FloatTensor,
+ temb: torch.FloatTensor,
+ scale: float = 1.0,
+ ) -> torch.FloatTensor:
+ hidden_states = input_tensor
+
+ if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+ hidden_states = self.norm1(hidden_states, temb)
+ else:
+ hidden_states = self.norm1(hidden_states)
+
+ hidden_states = self.nonlinearity(hidden_states)
+
+ if self.upsample is not None:
+ # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+ if hidden_states.shape[0] >= 64:
+ input_tensor = input_tensor.contiguous()
+ hidden_states = hidden_states.contiguous()
+ input_tensor = (
+ self.upsample(input_tensor, scale=scale)
+ )
+ hidden_states = (
+ self.upsample(hidden_states, scale=scale)
+ )
+ elif self.downsample is not None:
+ input_tensor = (
+ self.downsample(input_tensor, scale=scale)
+ )
+ hidden_states = (
+ self.downsample(hidden_states, scale=scale)
+ )
+
+ hidden_states = self.conv1(hidden_states)
+
+ if self.time_emb_proj is not None:
+ if not self.skip_time_act:
+ temb = self.nonlinearity(temb)
+ temb = (
+ self.time_emb_proj(temb, scale)[:, :, None, None]
+ )
+
+ if temb is not None and self.time_embedding_norm == "default":
+ hidden_states = hidden_states + temb
+
+ if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+ hidden_states = self.norm2(hidden_states, temb)
+ else:
+ hidden_states = self.norm2(hidden_states)
+
+ if temb is not None and self.time_embedding_norm == "scale_shift":
+ scale, shift = torch.chunk(temb, 2, dim=1)
+ hidden_states = hidden_states * (1 + scale) + shift
+
+ hidden_states = self.nonlinearity(hidden_states)
+
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+
+ if self.conv_shortcut is not None:
+ input_tensor = (
+ self.conv_shortcut(input_tensor)
+ )
+
+ output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+ return output_tensor
+
+
+def get_down_block3d(
+ down_block_type: str,
+ num_layers: int,
+ in_channels: int,
+ out_channels: int,
+ temb_channels: int,
+ add_downsample: bool,
+ downsample_stride: int,
+ resnet_eps: float,
+ resnet_act_fn: str,
+ transformer_layers_per_block: int = 1,
+ num_attention_heads: Optional[int] = None,
+ resnet_groups: Optional[int] = None,
+ cross_attention_dim: Optional[int] = None,
+ downsample_padding: Optional[int] = None,
+ dual_cross_attention: bool = False,
+ use_linear_projection: bool = False,
+ only_cross_attention: bool = False,
+ upcast_attention: bool = False,
+ resnet_time_scale_shift: str = "default",
+ attention_type: str = "default",
+ resnet_skip_time_act: bool = False,
+ resnet_out_scale_factor: float = 1.0,
+ cross_attention_norm: Optional[str] = None,
+ attention_head_dim: Optional[int] = None,
+ downsample_type: Optional[str] = None,
+ dropout: float = 0.0,
+):
+ # If attn head dim is not defined, we default it to the number of heads
+ if attention_head_dim is None:
+ logger.warn(
+ f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+ )
+ attention_head_dim = num_attention_heads
+
+ down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
+ if down_block_type == "DownEncoderBlockCausal3D":
+ return DownEncoderBlockCausal3D(
+ num_layers=num_layers,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ dropout=dropout,
+ add_downsample=add_downsample,
+ downsample_stride=downsample_stride,
+ resnet_eps=resnet_eps,
+ resnet_act_fn=resnet_act_fn,
+ resnet_groups=resnet_groups,
+ downsample_padding=downsample_padding,
+ resnet_time_scale_shift=resnet_time_scale_shift,
+ )
+ raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block3d(
+ up_block_type: str,
+ num_layers: int,
+ in_channels: int,
+ out_channels: int,
+ prev_output_channel: int,
+ temb_channels: int,
+ add_upsample: bool,
+ upsample_scale_factor: Tuple,
+ resnet_eps: float,
+ resnet_act_fn: str,
+ resolution_idx: Optional[int] = None,
+ transformer_layers_per_block: int = 1,
+ num_attention_heads: Optional[int] = None,
+ resnet_groups: Optional[int] = None,
+ cross_attention_dim: Optional[int] = None,
+ dual_cross_attention: bool = False,
+ use_linear_projection: bool = False,
+ only_cross_attention: bool = False,
+ upcast_attention: bool = False,
+ resnet_time_scale_shift: str = "default",
+ attention_type: str = "default",
+ resnet_skip_time_act: bool = False,
+ resnet_out_scale_factor: float = 1.0,
+ cross_attention_norm: Optional[str] = None,
+ attention_head_dim: Optional[int] = None,
+ upsample_type: Optional[str] = None,
+ dropout: float = 0.0,
+) -> nn.Module:
+ # If attn head dim is not defined, we default it to the number of heads
+ if attention_head_dim is None:
+ logger.warn(
+ f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+ )
+ attention_head_dim = num_attention_heads
+
+ up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+ if up_block_type == "UpDecoderBlockCausal3D":
+ return UpDecoderBlockCausal3D(
+ num_layers=num_layers,
+ in_channels=in_channels,
+ out_channels=out_channels,
+ resolution_idx=resolution_idx,
+ dropout=dropout,
+ add_upsample=add_upsample,
+ upsample_scale_factor=upsample_scale_factor,
+ resnet_eps=resnet_eps,
+ resnet_act_fn=resnet_act_fn,
+ resnet_groups=resnet_groups,
+ resnet_time_scale_shift=resnet_time_scale_shift,
+ temb_channels=temb_channels,
+ )
+ raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlockCausal3D(nn.Module):
+ """
+ A 3D UNet mid-block [`UNetMidBlockCausal3D`] with multiple residual blocks and optional attention blocks.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ temb_channels: int,
+ dropout: float = 0.0,
+ num_layers: int = 1,
+ resnet_eps: float = 1e-6,
+ resnet_time_scale_shift: str = "default", # default, spatial
+ resnet_act_fn: str = "swish",
+ resnet_groups: int = 32,
+ attn_groups: Optional[int] = None,
+ resnet_pre_norm: bool = True,
+ add_attention: bool = True,
+ attention_head_dim: int = 1,
+ output_scale_factor: float = 1.0,
+ ):
+ super().__init__()
+ resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+ self.add_attention = add_attention
+
+ if attn_groups is None:
+ attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
+
+ # there is always at least one resnet
+ resnets = [
+ ResnetBlockCausal3D(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ temb_channels=temb_channels,
+ eps=resnet_eps,
+ groups=resnet_groups,
+ dropout=dropout,
+ time_embedding_norm=resnet_time_scale_shift,
+ non_linearity=resnet_act_fn,
+ output_scale_factor=output_scale_factor,
+ pre_norm=resnet_pre_norm,
+ )
+ ]
+ attentions = []
+
+ if attention_head_dim is None:
+ logger.warn(
+ f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
+ )
+ attention_head_dim = in_channels
+
+ for _ in range(num_layers):
+ if self.add_attention:
+ processor = AttnProcessor()
+ 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,
+ processor=processor,
+ )
+ )
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/vae.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..b60e8ef0c860293e1856a79b14f7f4eec5d88a3b
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/hyvideo/vae/vae.py
@@ -0,0 +1,357 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from diffusers.models.attention_processor import SpatialNorm
+from diffusers.utils import BaseOutput, is_torch_version
+from diffusers.utils.torch_utils import randn_tensor
+
+from .unet_causal_3d_blocks import (
+ CausalConv3d,
+ UNetMidBlockCausal3D,
+ get_down_block3d,
+ get_up_block3d,
+)
+
+
+@dataclass
+class DecoderOutput(BaseOutput):
+ r"""
+ Output of decoding method.
+
+ Args:
+ sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ The decoded output sample from the last layer of the model.
+ """
+
+ sample: torch.FloatTensor
+
+
+class EncoderCausal3D(nn.Module):
+ r"""
+ The `EncoderCausal3D` layer of a variational autoencoder that encodes its input into a latent representation.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ down_block_types: Tuple[str, ...] = ("DownEncoderBlockCausal3D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ double_z: bool = True,
+ mid_block_add_attention=True,
+ time_compression_ratio: int = 4,
+ spatial_compression_ratio: int = 8,
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+
+ self.conv_in = CausalConv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1)
+ self.mid_block = None
+ self.down_blocks = nn.ModuleList([])
+
+ # down
+ output_channel = block_out_channels[0]
+ for i, down_block_type in enumerate(down_block_types):
+ input_channel = output_channel
+ output_channel = block_out_channels[i]
+ is_final_block = i == len(block_out_channels) - 1
+ num_spatial_downsample_layers = int(np.log2(spatial_compression_ratio))
+ num_time_downsample_layers = int(np.log2(time_compression_ratio))
+
+ if time_compression_ratio == 4:
+ add_spatial_downsample = bool(i < num_spatial_downsample_layers)
+ add_time_downsample = bool(
+ i >= (len(block_out_channels) - 1 - num_time_downsample_layers)
+ and not is_final_block
+ )
+ else:
+ raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
+
+ downsample_stride_HW = (2, 2) if add_spatial_downsample else (1, 1)
+ downsample_stride_T = (2,) if add_time_downsample else (1,)
+ downsample_stride = tuple(downsample_stride_T + downsample_stride_HW)
+ down_block = get_down_block3d(
+ down_block_type,
+ num_layers=self.layers_per_block,
+ in_channels=input_channel,
+ out_channels=output_channel,
+ add_downsample=bool(add_spatial_downsample or add_time_downsample),
+ downsample_stride=downsample_stride,
+ resnet_eps=1e-6,
+ downsample_padding=0,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ attention_head_dim=output_channel,
+ temb_channels=None,
+ )
+ self.down_blocks.append(down_block)
+
+ # mid
+ self.mid_block = UNetMidBlockCausal3D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default",
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=None,
+ add_attention=mid_block_add_attention,
+ )
+
+ # out
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+
+ conv_out_channels = 2 * out_channels if double_z else out_channels
+ self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3)
+
+ def forward(self, sample: torch.FloatTensor) -> torch.FloatTensor:
+ r"""The forward method of the `EncoderCausal3D` class."""
+ assert len(sample.shape) == 5, "The input tensor should have 5 dimensions"
+
+ sample = self.conv_in(sample)
+
+ # down
+ for down_block in self.down_blocks:
+ sample = down_block(sample)
+
+ # middle
+ sample = self.mid_block(sample)
+
+ # post-process
+ sample = self.conv_norm_out(sample)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+
+class DecoderCausal3D(nn.Module):
+ r"""
+ The `DecoderCausal3D` layer of a variational autoencoder that decodes its latent representation into an output sample.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ up_block_types: Tuple[str, ...] = ("UpDecoderBlockCausal3D",),
+ block_out_channels: Tuple[int, ...] = (64,),
+ layers_per_block: int = 2,
+ norm_num_groups: int = 32,
+ act_fn: str = "silu",
+ norm_type: str = "group", # group, spatial
+ mid_block_add_attention=True,
+ time_compression_ratio: int = 4,
+ spatial_compression_ratio: int = 8,
+ ):
+ super().__init__()
+ self.layers_per_block = layers_per_block
+
+ self.conv_in = CausalConv3d(in_channels, block_out_channels[-1], kernel_size=3, stride=1)
+ self.mid_block = None
+ self.up_blocks = nn.ModuleList([])
+
+ temb_channels = in_channels if norm_type == "spatial" else None
+
+ # mid
+ self.mid_block = UNetMidBlockCausal3D(
+ in_channels=block_out_channels[-1],
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ output_scale_factor=1,
+ resnet_time_scale_shift="default" if norm_type == "group" else norm_type,
+ attention_head_dim=block_out_channels[-1],
+ resnet_groups=norm_num_groups,
+ temb_channels=temb_channels,
+ add_attention=mid_block_add_attention,
+ )
+
+ # up
+ reversed_block_out_channels = list(reversed(block_out_channels))
+ output_channel = reversed_block_out_channels[0]
+ for i, up_block_type in enumerate(up_block_types):
+ prev_output_channel = output_channel
+ output_channel = reversed_block_out_channels[i]
+ is_final_block = i == len(block_out_channels) - 1
+ num_spatial_upsample_layers = int(np.log2(spatial_compression_ratio))
+ num_time_upsample_layers = int(np.log2(time_compression_ratio))
+
+ if time_compression_ratio == 4:
+ add_spatial_upsample = bool(i < num_spatial_upsample_layers)
+ add_time_upsample = bool(
+ i >= len(block_out_channels) - 1 - num_time_upsample_layers
+ and not is_final_block
+ )
+ else:
+ raise ValueError(f"Unsupported time_compression_ratio: {time_compression_ratio}.")
+
+ upsample_scale_factor_HW = (2, 2) if add_spatial_upsample else (1, 1)
+ upsample_scale_factor_T = (2,) if add_time_upsample else (1,)
+ upsample_scale_factor = tuple(upsample_scale_factor_T + upsample_scale_factor_HW)
+ up_block = get_up_block3d(
+ up_block_type,
+ num_layers=self.layers_per_block + 1,
+ in_channels=prev_output_channel,
+ out_channels=output_channel,
+ prev_output_channel=None,
+ add_upsample=bool(add_spatial_upsample or add_time_upsample),
+ upsample_scale_factor=upsample_scale_factor,
+ resnet_eps=1e-6,
+ resnet_act_fn=act_fn,
+ resnet_groups=norm_num_groups,
+ attention_head_dim=output_channel,
+ temb_channels=temb_channels,
+ resnet_time_scale_shift=norm_type,
+ )
+ self.up_blocks.append(up_block)
+ prev_output_channel = output_channel
+
+ # out
+ if norm_type == "spatial":
+ self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels)
+ else:
+ self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+ self.conv_act = nn.SiLU()
+ self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ sample: torch.FloatTensor,
+ latent_embeds: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ r"""The forward method of the `DecoderCausal3D` class."""
+ assert len(sample.shape) == 5, "The input tensor should have 5 dimensions."
+
+ sample = self.conv_in(sample)
+
+ upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+ if self.training and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ if is_torch_version(">=", "1.11.0"):
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block),
+ sample,
+ latent_embeds,
+ use_reentrant=False,
+ )
+ sample = sample.to(upscale_dtype)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(up_block),
+ sample,
+ latent_embeds,
+ use_reentrant=False,
+ )
+ else:
+ # middle
+ sample = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(self.mid_block), sample, latent_embeds
+ )
+ sample = sample.to(upscale_dtype)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = torch.utils.checkpoint.checkpoint(create_custom_forward(up_block), sample, latent_embeds)
+ else:
+ # middle
+ sample = self.mid_block(sample, latent_embeds)
+ sample = sample.to(upscale_dtype)
+
+ # up
+ for up_block in self.up_blocks:
+ sample = up_block(sample, latent_embeds)
+
+ # post-process
+ if latent_embeds is None:
+ sample = self.conv_norm_out(sample)
+ else:
+ sample = self.conv_norm_out(sample, latent_embeds)
+ sample = self.conv_act(sample)
+ sample = self.conv_out(sample)
+
+ return sample
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters: torch.Tensor, deterministic: bool = False):
+ if parameters.ndim == 3:
+ dim = 2 # (B, L, C)
+ elif parameters.ndim == 5 or parameters.ndim == 4:
+ dim = 1 # (B, C, T, H ,W) / (B, C, H, W)
+ else:
+ raise NotImplementedError
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim)
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.deterministic = deterministic
+ self.std = torch.exp(0.5 * self.logvar)
+ self.var = torch.exp(self.logvar)
+ if self.deterministic:
+ self.var = self.std = torch.zeros_like(
+ self.mean, device=self.parameters.device, dtype=self.parameters.dtype
+ )
+
+ def sample(self, generator: Optional[torch.Generator] = None) -> torch.FloatTensor:
+ # make sure sample is on the same device as the parameters and has same dtype
+ sample = randn_tensor(
+ self.mean.shape,
+ generator=generator,
+ device=self.parameters.device,
+ dtype=self.parameters.dtype,
+ )
+ x = self.mean + self.std * sample
+ return x
+
+ def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+ if self.deterministic:
+ return torch.Tensor([0.0])
+ else:
+ reduce_dim = list(range(1, self.mean.ndim))
+ if other is None:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+ dim=reduce_dim,
+ )
+ else:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean - other.mean, 2) / other.var
+ + self.var / other.var
+ - 1.0
+ - self.logvar
+ + other.logvar,
+ dim=reduce_dim,
+ )
+
+ def nll(self, sample: torch.Tensor, dims: Tuple[int, ...] = None) -> torch.Tensor:
+ if dims is None:
+ dims = (1, 2, 3)
+ 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/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/prompts.txt b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/prompts.txt
new file mode 100644
index 0000000000000000000000000000000000000000..ac028cbaf2e7b8e479a0eca36335a24ac720af67
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/prompts.txt
@@ -0,0 +1,6 @@
+realistic style, a lone cowboy rides his horse across an open plain at beautiful sunset, soft light, warm colors.
+extreme close-up with a shallow depth of field of a puddle in a street. reflecting a busy futuristic Tokyo city with bright neon signs, night, lens flare.
+Extreme close-up of chicken and green pepper kebabs grilling on a barbeque with flames. Shallow focus and light smoke. vivid colours.
+Timelapse of the northern lights dancing across the Arctic sky, stars twinkling, snow-covered landscape.
+A panning shot of a serene mountain landscape, slowly revealing snow-capped peaks, granite rocks and a crystal-clear lake reflecting the sky.
+moody shot of a central European alley film noir cinematic black and white high contrast high detail.
\ No newline at end of file
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/requirements.txt b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..a99b9adf08a1ae72ecb322ddd7d4eba67ba4fb13
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/requirements.txt
@@ -0,0 +1,16 @@
+opencv-python==4.9.0.80
+diffusers==0.31.0
+transformers==4.46.3
+tokenizers==0.20.3
+accelerate==1.1.1
+pandas==2.0.3
+numpy==1.24.4
+einops==0.7.0
+tqdm==4.66.2
+loguru==0.7.2
+imageio==2.34.0
+imageio-ffmpeg==0.5.1
+safetensors==0.4.3
+gradio==5.0.0
+deepspeed
+strenum
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/run.sh b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/run.sh
new file mode 100644
index 0000000000000000000000000000000000000000..ec5c416cfbd3cffef21938810188cb70d25b331e
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/run.sh
@@ -0,0 +1,21 @@
+export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"
+export TASK_QUEUE_ENABLE=2
+export HCCL_OP_EXPANSION_MODE="AIV"
+export TOKENIZERS_PARALLELISM=false
+torchrun --nproc_per_node=8 sample_video.py \
+ --model-base /home/weights/HunyuanVideo_weights \
+ --dit-weight /home/weights/HunyuanVideo_weights/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+ --vae-path /home/weights/HunyuanVideo_weights/hunyuan-video-t2v-720p/vae \
+ --text-encoder-path /home/weights/HunyuanVideo_weights/llm \
+ --text-encoder-2-path /home/weights/HunyuanVideo_weights/clip-vit-large-patch14 \
+ --model-resolution "720p" \
+ --video-size 720 1280 \
+ --video-length 129 \
+ --infer-steps 50 \
+ --prompt ./prompts.txt \
+ --seed 42 \
+ --flow-reverse \
+ --num-videos 1 \
+ --save-path ./results \
+ --sp-degree 8 \
+ --algorithm attention_cache
\ No newline at end of file
diff --git a/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/sample_video.py b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/sample_video.py
new file mode 100644
index 0000000000000000000000000000000000000000..aee9c65c7f7513e62ce99e8b215b473c9b7ef2bf
--- /dev/null
+++ b/MindIE/MindIE-Torch/built-in/foundation/HunyuanVideo-Beta/sample_video.py
@@ -0,0 +1,128 @@
+import os
+import time
+from datetime import datetime
+from pathlib import Path
+
+import torch
+import torch_npu
+from loguru import logger
+
+from hyvideo.config import parse_args
+from hyvideo.inference import HunyuanVideoSampler
+from hyvideo.utils.file_utils import save_videos_grid
+
+torch_npu.npu.set_compile_mode(jit_compile=False)
+torch.npu.config.allow_internal_format = False
+
+
+def main():
+ args = parse_args()
+ print(args)
+ models_root_path = Path(args.model_base)
+ if not models_root_path.exists():
+ raise ValueError(f"`models_root` not exists: {models_root_path}")
+
+ # Create save folder to save the samples
+ save_path = args.save_path if args.save_path_suffix == "" else f'{args.save_path}_{args.save_path_suffix}'
+ if not os.path.exists(save_path):
+ os.makedirs(save_path, exist_ok=True)
+
+ # Load models
+ hunyuan_video_sampler = HunyuanVideoSampler.from_pretrained(models_root_path, args=args)
+
+ # Get the updated args
+ args = hunyuan_video_sampler.args
+ if args.algorithm == "dit_cache":
+ from mindiesd import CacheConfig, CacheAgent
+ transformer = hunyuan_video_sampler.pipeline.transformer
+ # double_stream
+ config_double_blocks_cache = CacheConfig(
+ method="dit_block_cache",
+ blocks_count=len(transformer.double_blocks),
+ steps_count=50,
+ step_start=10,
+ step_interval=3,
+ step_end=49,
+ block_start=3,
+ block_end=18
+ )
+ double_blocks_cache = CacheAgent(config_double_blocks_cache)
+ hunyuan_video_sampler.pipeline.transformer.double_blocks_cache = double_blocks_cache
+ # single_stream
+ config_single_blocks_cache = CacheConfig(
+ method="dit_block_cache",
+ blocks_count=len(transformer.single_blocks),
+ steps_count=50,
+ step_start=10,
+ step_interval=3,
+ step_end=49,
+ block_start=5,
+ block_end=35
+ )
+ single_blocks_cache = CacheAgent(config_single_blocks_cache)
+ hunyuan_video_sampler.pipeline.transformer.single_blocks_cache = single_blocks_cache
+ elif args.algorithm == "attention_cache":
+ from mindiesd import CacheConfig, CacheAgent
+ transformer = hunyuan_video_sampler.pipeline.transformer
+ config_double_stream_attention = CacheConfig(
+ method="attention_cache",
+ blocks_count=len(transformer.double_blocks),
+ steps_count=50,
+ step_start=11,
+ step_interval=4,
+ step_end=47
+ )
+ config_single_stream_attention = CacheConfig(
+ method="attention_cache",
+ blocks_count=len(transformer.single_blocks),
+ steps_count=50,
+ step_start=11,
+ step_interval=4,
+ step_end=47
+ )
+ cache_double_stream_attention = CacheAgent(config_double_stream_attention)
+ cache_single_stream_attention = CacheAgent(config_single_stream_attention)
+ for block in transformer.double_blocks:
+ block.attention_cache = cache_double_stream_attention
+ for block in transformer.single_blocks:
+ block.attention_cache = cache_single_stream_attention
+
+ if not isinstance(args.prompt, list):
+ args.prompt = [args.prompt]
+ if len(args.prompt) == 1 and args.prompt[0].endswith('txt'):
+ prompt = open(args.prompt[0], 'r').readlines()
+ args.prompt = [i.strip() for i in prompt]
+
+ for i, prompt in enumerate(args.prompt):
+ # Start sampling
+ outputs = hunyuan_video_sampler.predict(
+ prompt=prompt,
+ height=args.video_size[0],
+ width=args.video_size[1],
+ video_length=args.video_length,
+ seed=args.seed,
+ negative_prompt=args.neg_prompt,
+ infer_steps=args.infer_steps,
+ guidance_scale=args.cfg_scale,
+ num_videos_per_prompt=args.num_videos,
+ flow_shift=args.flow_shift,
+ batch_size=args.batch_size,
+ embedded_guidance_scale=args.embedded_cfg_scale
+ )
+ samples_ls = outputs['samples']
+
+ # Save samples
+ if 'LOCAL_RANK' not in os.environ or int(os.environ['LOCAL_RANK']) == 0:
+ for j, sample in enumerate(samples_ls):
+ sample = sample.unsqueeze(0)
+ time_flag = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d-%H:%M:%S")
+ cur_save_path = f"{save_path}/prompt_{i}_{time_flag}_seed{outputs['seeds'][j]}_{outputs['prompts'][j][:100].replace('/', '')}.mp4"
+ save_videos_grid(sample, cur_save_path, fps=24)
+ logger.info(f'Sample save to: {cur_save_path}')
+
+ # Finalize distributed group
+ hunyuan_video_sampler.finalize()
+
+
+if __name__ == "__main__":
+ main()