diff --git a/Step-Video-TI2V/.gitignore b/Step-Video-TI2V/.gitignore new file mode 100755 index 0000000..2aa424c --- /dev/null +++ b/Step-Video-TI2V/.gitignore @@ -0,0 +1,7 @@ +tests +debug +dev +*.egg-info +__pycache__ +lib +results \ No newline at end of file diff --git a/Step-Video-TI2V/LICENSE b/Step-Video-TI2V/LICENSE new file mode 100644 index 0000000..0e58920 --- /dev/null +++ b/Step-Video-TI2V/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 stepfun-ai + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Step-Video-TI2V/README.md b/Step-Video-TI2V/README.md new file mode 100644 index 0000000..cc7edb5 --- /dev/null +++ b/Step-Video-TI2V/README.md @@ -0,0 +1,39 @@ +# Step-Video-Ti2V + +

+ +

+
+   +   +   +
+ +
+   +
+ +## ๐Ÿ”ฅ๐Ÿ”ฅ๐Ÿ”ฅ News!! +* Feb 17, 2025: ๐Ÿ‘‹ We release the inference code and model weights of Step-Video-Ti2V. [Download](https://huggingface.co/stepfun-ai/stepvideo-ti2v) +* Feb 17, 2025: ๐ŸŽ‰ We have made our technical report available as open source. [Read](https://arxiv.org/abs/2502.10248) + + +```bash +git clone https://github.com/stepfun-ai/Step-Video-TI2V.git +conda create -n stepvideo python=3.10 +conda activate stepvideo + +cd StepFun-StepVideo +pip install -e . +``` + +### ๐Ÿš€ 4.2. Inference Scripts +```bash +python api/call_remote_server.py --model_dir where_you_download_dir & ## We assume you have more than 4 GPUs available. This command will return the URL for both the caption API and the VAE API. Please use the returned URL in the following command. + +parallel=4 # or parallel=8 +url='127.0.0.1' +model_dir=where_you_download_dir + +torchrun --nproc_per_node $parallel run_parallel --model_dir $model_dir --vae_url $url --caption_url $url --ulysses_degree $parallel --prompt "็ฌ‘่ตทๆฅ" --infer_steps 15 --cfg_scale 5.0 --time_shift 17.0 +``` diff --git a/Step-Video-TI2V/api/call_remote_server.py b/Step-Video-TI2V/api/call_remote_server.py new file mode 100644 index 0000000..52e2535 --- /dev/null +++ b/Step-Video-TI2V/api/call_remote_server.py @@ -0,0 +1,221 @@ +import torch +import os +from flask import Flask, Response, jsonify, request, Blueprint +from flask_restful import Api, Resource +import pickle +import argparse +import threading +import argparse + + +device = f'cuda:{torch.cuda.device_count()-1}' +torch.cuda.set_device(device) +dtype = torch.bfloat16 + +def parsed_args(): + parser = argparse.ArgumentParser(description="StepVideo API Functions") + parser.add_argument('--model_dir', type=str) + parser.add_argument('--clip_dir', type=str, default='hunyuan_clip') + parser.add_argument('--llm_dir', type=str, default='step_llm') + parser.add_argument('--vae_dir', type=str, default='vae') + parser.add_argument('--port', type=str, default='8080') + args = parser.parse_args() + return args + + + +class StepVaePipeline(Resource): + def __init__(self, vae_dir, version=2): + self.vae = self.build_vae(vae_dir, version) + self.scale_factor = 1.0 + + def build_vae(self, vae_dir, version=2): + from stepvideo.vae.vae import AutoencoderKL + (model_name, z_channels) = ("vae_v2.safetensors", 64) if version == 2 else ("vae.safetensors", 16) + model_path = os.path.join(vae_dir, model_name) + + model = AutoencoderKL( + z_channels=z_channels, + model_path=model_path, + version=version, + ).to(dtype).to(device).eval() + print("Inintialized vae...") + return model + + def decode(self, samples, *args, **kwargs): + with torch.no_grad(): + try: + dtype = next(self.vae.parameters()).dtype + device = next(self.vae.parameters()).device + samples = self.vae.decode(samples.to(dtype).to(device) / self.scale_factor) + if hasattr(samples,'sample'): + samples = samples.sample + return samples + except Exception as err: + print(f"vae decode error: {err}") + torch.cuda.empty_cache() + return None + + def encode(self, videos, *args, **kwargs): + with torch.no_grad(): + try: + dtype = next(self.vae.parameters()).dtype + device = next(self.vae.parameters()).device + latents = self.vae.encode(videos.to(dtype).to(device))*self.scale_factor + if hasattr(latents,'sample'): + latents = latents.sample + return latents + except Exception as err: + print(f"vae encode error: {err}") + torch.cuda.empty_cache() + return None + +lock = threading.Lock() +class VAEapi(Resource): + def __init__(self, vae_pipeline): + self.vae_pipeline = vae_pipeline + + def get(self): + with lock: + try: + feature = pickle.loads(request.get_data()) + feature['api'] = 'vae' + + feature = {k:v for k, v in feature.items() if v is not None} + video_latents = self.vae_pipeline.decode(**feature) + response = pickle.dumps(video_latents) + + except Exception as e: + print("Caught Exception: ", e) + return Response(e) + + return Response(response) + + +class VAEEncodeapi(Resource): + def __init__(self, vae_pipeline): + self.vae_pipeline = vae_pipeline + + def get(self): + with lock: + try: + feature = pickle.loads(request.get_data()) + feature['api'] = 'vae-encode' + + feature = {k:v for k, v in feature.items() if v is not None} + video_latents = self.vae_pipeline.encode(**feature) + response = pickle.dumps(video_latents) + + except Exception as e: + print("Caught Exception: ", e) + return Response(e) + + return Response(response) + + +class CaptionPipeline(Resource): + def __init__(self, llm_dir, clip_dir): + self.text_encoder = self.build_llm(llm_dir) + self.clip = self.build_clip(clip_dir) + + def build_llm(self, model_dir): + from stepvideo.text_encoder.stepllm import STEP1TextEncoder + text_encoder = STEP1TextEncoder(model_dir, max_length=320).to(dtype).to(device).eval() + print("Inintialized text encoder...") + return text_encoder + + def build_clip(self, model_dir): + from stepvideo.text_encoder.clip import HunyuanClip + clip = HunyuanClip(model_dir, max_length=77).to(device).eval() + print("Inintialized clip encoder...") + return clip + + def embedding(self, prompts, *args, **kwargs): + with torch.no_grad(): + try: + y, y_mask = self.text_encoder(prompts) + + clip_embedding, _ = self.clip(prompts) + + len_clip = clip_embedding.shape[1] + y_mask = torch.nn.functional.pad(y_mask, (len_clip, 0), value=1) ## pad attention_mask with clip's length + + data = { + 'y': y.detach().cpu(), + 'y_mask': y_mask.detach().cpu(), + 'clip_embedding': clip_embedding.to(torch.bfloat16).detach().cpu() + } + + return data + except Exception as err: + print(f"{err}") + return None + + + +lock = threading.Lock() +class Captionapi(Resource): + def __init__(self, caption_pipeline): + self.caption_pipeline = caption_pipeline + + def get(self): + with lock: + try: + feature = pickle.loads(request.get_data()) + feature['api'] = 'caption' + + feature = {k:v for k, v in feature.items() if v is not None} + embeddings = self.caption_pipeline.embedding(**feature) + response = pickle.dumps(embeddings) + + except Exception as e: + print("Caught Exception: ", e) + return Response(e) + + return Response(response) + + + + +class RemoteServer(object): + def __init__(self, args) -> None: + self.app = Flask(__name__) + root = Blueprint("root", __name__) + self.app.register_blueprint(root) + api = Api(self.app) + + self.vae_pipeline = StepVaePipeline( + vae_dir=os.path.join(args.model_dir, args.vae_dir) + ) + api.add_resource( + VAEapi, + "/vae-api", + resource_class_args=[self.vae_pipeline], + ) + + api.add_resource( + VAEEncodeapi, + "/vae-encode-api", + resource_class_args=[self.vae_pipeline], + ) + + self.caption_pipeline = CaptionPipeline( + llm_dir=os.path.join(args.model_dir, args.llm_dir), + clip_dir=os.path.join(args.model_dir, args.clip_dir) + ) + api.add_resource( + Captionapi, + "/caption-api", + resource_class_args=[self.caption_pipeline], + ) + + + def run(self, host="0.0.0.0", port=8080): + self.app.run(host, port=port, threaded=True, debug=False) + + +if __name__ == "__main__": + args = parsed_args() + flask_server = RemoteServer(args) + flask_server.run(host="0.0.0.0", port=args.port) + \ No newline at end of file diff --git a/Step-Video-TI2V/assets/demo.png b/Step-Video-TI2V/assets/demo.png new file mode 100644 index 0000000..a2e9a13 Binary files /dev/null and b/Step-Video-TI2V/assets/demo.png differ diff --git a/Step-Video-TI2V/run_parallel.py b/Step-Video-TI2V/run_parallel.py new file mode 100755 index 0000000..67f23ac --- /dev/null +++ b/Step-Video-TI2V/run_parallel.py @@ -0,0 +1,42 @@ +from stepvideo.diffusion.video_pipeline import StepVideoPipeline +import torch.distributed as dist +import torch +from stepvideo.config import parse_args +from stepvideo.parallel import initialize_parall_group, get_parallel_group +from stepvideo.utils import setup_seed + + +if __name__ == "__main__": + args = parse_args() + initialize_parall_group(ring_degree=args.ring_degree, ulysses_degree=args.ulysses_degree) + + local_rank = get_parallel_group().local_rank + device = torch.device(f"cuda:{local_rank}") + + setup_seed(args.seed) + + pipeline = StepVideoPipeline.from_pretrained(args.model_dir).to(dtype=torch.bfloat16, device="cpu") + + pipeline.transformer = pipeline.transformer.to(device) + pipeline.setup_api( + vae_url = args.vae_url, + caption_url = args.caption_url, + ) + + + prompt = args.prompt + videos = pipeline( + prompt=prompt, + first_image=args.first_image_path, + num_frames=args.num_frames, + height=args.height, + width=args.width, + num_inference_steps = args.infer_steps, + guidance_scale=args.cfg_scale, + time_shift=args.time_shift, + pos_magic=args.pos_magic, + neg_magic=args.neg_magic, + output_file_name=args.output_file_name + ) + + dist.destroy_process_group() \ No newline at end of file diff --git a/Step-Video-TI2V/setup.py b/Step-Video-TI2V/setup.py new file mode 100644 index 0000000..c1c66c6 --- /dev/null +++ b/Step-Video-TI2V/setup.py @@ -0,0 +1,57 @@ +from setuptools import find_packages, setup +import subprocess + + +def get_cuda_version(): + try: + nvcc_version = subprocess.check_output(["nvcc", "--version"]).decode("utf-8") + version_line = [line for line in nvcc_version.split("\n") if "release" in line][ + 0 + ] + cuda_version = version_line.split(" ")[-2].replace(",", "") + return "cu" + cuda_version.replace(".", "") + except Exception as e: + return "no_cuda" + + +if __name__ == "__main__": + with open("README.md", "r") as f: + long_description = f.read() + fp = open("stepvideo/__version__.py", "r").read() + version = eval(fp.strip().split()[-1]) + + setup( + name="stepvideo", + author="Step-Video Team", + packages=find_packages(), + install_requires=[ + "torchvision==0.20", + "torch==2.5.0", + "accelerate>=1.0.0", + "transformers>=4.39.1", + "diffusers>=0.31.0", + "sentencepiece>=0.1.99", + "imageio>=2.37.0", + "optimus==2.1", + "numpy", + "einops", + "aiohttp", + "asyncio", + "flask", + "flask_restful", + "ffmpeg-python", + "requests", + "xfuser==0.4.2rc2" + ], + url="", + description="A 30B DiT based text to video and image generation model", + long_description=long_description, + long_description_content_type="text/markdown", + version=version, + classifiers=[ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", + ], + include_package_data=True, + python_requires=">=3.10", + ) \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/__init__.py b/Step-Video-TI2V/stepvideo/__init__.py new file mode 100644 index 0000000..e919871 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/__init__.py @@ -0,0 +1,7 @@ +import os + +os.environ["NCCL_DEBUG"] = "ERROR" + +from .diffusion.scheduler import * +from .diffusion.video_pipeline import * +from .modules.model import * \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/__version__.py b/Step-Video-TI2V/stepvideo/__version__.py new file mode 100644 index 0000000..a68927d --- /dev/null +++ b/Step-Video-TI2V/stepvideo/__version__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/config.py b/Step-Video-TI2V/stepvideo/config.py new file mode 100644 index 0000000..01a1173 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/config.py @@ -0,0 +1,195 @@ +import argparse + +def parse_args(namespace=None): + parser = argparse.ArgumentParser(description="StepVideo inference script") + + parser = add_extra_models_args(parser) + parser = add_denoise_schedule_args(parser) + parser = add_inference_args(parser) + parser = add_parallel_args(parser) + + args = parser.parse_args(namespace=namespace) + + return args + + + +def add_extra_models_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group( + title="Extra models args, including vae, text encoders and tokenizers)" + ) + + group.add_argument( + "--vae_url", + type=str, + default='127.0.0.1', + help="vae url.", + ) + group.add_argument( + "--caption_url", + type=str, + default='127.0.0.1', + help="caption url.", + ) + + return parser + + +def add_denoise_schedule_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Denoise schedule args") + + # Flow Matching + group.add_argument( + "--time_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.", + ) + + return parser + + +def add_inference_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Inference args") + + # ======================== Model loads ======================== + group.add_argument( + "--model_dir", + type=str, + default="./ckpts", + help="Root path of all the models, including t2v models and extra models.", + ) + group.add_argument( + "--model_resolution", + type=str, + default="540p", + choices=["540p"], + help="Root path of all the models, including t2v models and extra models.", + ) + 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( + "--save_path", + type=str, + default="./results", + help="Path to save the generated samples.", + ) + group.add_argument( + "--output_file_name", + type=str, + default="out", + help="Name to save the generated 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( + "--num_frames", + type=int, + default=102, + help="How many frames to sample from a video. ", + ) + group.add_argument( + "--height", + type=int, + default=544, + help="The height of video sample", + ) + group.add_argument( + "--width", + type=int, + default=992, + help="The width of video sample", + ) + # --- prompt --- + group.add_argument( + "--prompt", + type=str, + default=None, + help="Prompt for sampling during evaluation.", + ) + group.add_argument( + "--first_image_path", + type=str, + default='./assets/demo.png', + help="The reference image path for image-to-video task.", + ) + group.add_argument("--seed", type=int, default=1234, help="Seed for evaluation.") + + # Classifier-Free Guidance + group.add_argument( + "--pos_magic", type=str, default="่ถ…้ซ˜ๆธ…ใ€HDR ่ง†้ข‘ใ€็Žฏๅขƒๅ…‰ใ€ๆœๆฏ”ๅ…จๆ™ฏๅฃฐใ€็”ป้ข็จณๅฎšใ€ๆต็•…ๅŠจไฝœใ€้€ผ็œŸ็š„็ป†่Š‚ใ€ไธ“ไธš็บงๆž„ๅ›พใ€่ถ…็Žฐๅฎžไธปไน‰ใ€่‡ช็„ถใ€็”ŸๅŠจใ€่ถ…็ป†่Š‚ใ€ๆธ…ๆ™ฐใ€‚", help="Positive magic prompt for sampling." + ) + group.add_argument( + "--neg_magic", type=str, default="็”ป้ขๆš—ใ€ไฝŽๅˆ†่พจ็Ž‡ใ€ไธ่‰ฏๆ‰‹ใ€ๆ–‡ๆœฌใ€็ผบๅฐ‘ๆ‰‹ๆŒ‡ใ€ๅคšไฝ™็š„ๆ‰‹ๆŒ‡ใ€่ฃๅ‰ชใ€ไฝŽ่ดจ้‡ใ€้ข—็ฒ’็Šถใ€็ญพๅใ€ๆฐดๅฐใ€็”จๆˆทๅใ€ๆจก็ณŠใ€‚", help="Negative magic prompt for sampling." + ) + group.add_argument( + "--cfg_scale", type=float, default=9.0, help="Classifier free guidance scale." + ) + + + return parser + + +def add_parallel_args(parser: argparse.ArgumentParser): + group = parser.add_argument_group(title="Parallel args") + + # ======================== Model loads ======================== + group.add_argument( + "--ulysses_degree", + type=int, + default=8, + help="Ulysses degree.", + ) + group.add_argument( + "--ring_degree", + type=int, + default=1, + help="Ulysses degree.", + ) + + group.add_argument( + "--tensor_parallel_degree", + type=int, + default=1, + help="Tensor parallel degree.", + ) + return parser diff --git a/Step-Video-TI2V/stepvideo/diffusion/scheduler.py b/Step-Video-TI2V/stepvideo/diffusion/scheduler.py new file mode 100644 index 0000000..7996a1d --- /dev/null +++ b/Step-Video-TI2V/stepvideo/diffusion/scheduler.py @@ -0,0 +1,236 @@ +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import numpy as np +import torch + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput, logging +from diffusers.schedulers.scheduling_utils import SchedulerMixin + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + 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, + reverse: bool = False, + solver: str = "euler", + device: Union[str, torch.device] = 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.device = device + + 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, + time_shift: float = 13.0, + device: Union[str, torch.device] = 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. + """ + device = device or self.device + self.num_inference_steps = num_inference_steps + + sigmas = torch.linspace(1, 0, num_inference_steps + 1, device=device) + sigmas = self.sd3_time_shift(sigmas, time_shift) + + if not self.config.reverse: + sigmas = 1 - sigmas + + self.sigmas = sigmas + self.timesteps = sigmas[:-1] + + # 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, time_shift: float = 13.0): + return (time_shift * t) / (1 + (time_shift - 1) * t) + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + return_dict: bool = False, + ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + n_tokens (`int`, *optional*): + Number of tokens in the input sequence. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] + + if self.config.solver == "euler": + prev_sample = sample + model_output.to(torch.float32) * dt + else: + raise ValueError( + f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}" + ) + + # upon completion increase step index by one + self._step_index += 1 + + if not return_dict: + return prev_sample + + return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) + + def __len__(self): + return self.config.num_train_timesteps diff --git a/Step-Video-TI2V/stepvideo/diffusion/video_pipeline.py b/Step-Video-TI2V/stepvideo/diffusion/video_pipeline.py new file mode 100755 index 0000000..824e0bd --- /dev/null +++ b/Step-Video-TI2V/stepvideo/diffusion/video_pipeline.py @@ -0,0 +1,423 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. + +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from dataclasses import dataclass + +import numpy as np +import pickle +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.utils import BaseOutput +import asyncio + +from stepvideo.modules.model import StepVideoModel +from stepvideo.diffusion.scheduler import FlowMatchDiscreteScheduler +from stepvideo.utils import VideoProcessor +from torchvision import transforms +from PIL import Image as PILImage + +import os + + +def call_api_gen(url, api, port=8080): + url =f"http://{url}:{port}/{api}-api" + import aiohttp + async def _fn(samples, *args, **kwargs): + if api=='vae': + data = { + "samples": samples, + } + elif api=='vae-encode': + data = { + "videos": samples, + } + elif api == 'caption': + data = { + "prompts": samples, + } + else: + raise Exception(f"Not supported api: {api}...") + + async with aiohttp.ClientSession() as sess: + data_bytes = pickle.dumps(data) + async with sess.get(url, data=data_bytes, timeout=12000) as response: + result = bytearray() + while not response.content.at_eof(): + chunk = await response.content.read(1024) + result += chunk + response_data = pickle.loads(result) + return response_data + + return _fn + + + + +@dataclass +class StepVideoPipelineOutput(BaseOutput): + video: Union[torch.Tensor, np.ndarray] + + +class StepVideoPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-video generation using StepVideo. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + transformer ([`StepVideoModel`]): + Conditional Transformer to denoise the encoded image latents. + scheduler ([`FlowMatchDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae_url: + remote vae server's url. + caption_url: + remote caption (stepllm and clip) server's url. + """ + + def __init__( + self, + transformer: StepVideoModel, + scheduler: FlowMatchDiscreteScheduler, + vae_url: str = '127.0.0.1', + caption_url: str = '127.0.0.1', + save_path: str = './results', + name_suffix: str = 'comfyui', + ): + super().__init__() + + self.register_modules( + transformer=transformer, + scheduler=scheduler, + ) + + self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 8 + self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16 + self.video_processor = VideoProcessor(save_path, name_suffix) + + self.vae_url = vae_url + self.caption_url = caption_url + self.setup_api(self.vae_url, self.caption_url) + + def setup_api(self, vae_url, caption_url): + self.vae_url = vae_url + self.caption_url = caption_url + self.caption = call_api_gen(caption_url, 'caption') + self.vae = call_api_gen(vae_url, 'vae') + self.vae_encode = call_api_gen(vae_url, 'vae-encode') + return self + + def encode_prompt( + self, + prompt: str, + neg_magic: str = '', + pos_magic: str = '', + ): + device = self._execution_device + prompts = [prompt+pos_magic] + bs = len(prompts) + prompts += [neg_magic]*bs + + data = asyncio.run(self.caption(prompts)) + prompt_embeds, prompt_attention_mask, clip_embedding = data['y'].to(device), data['y_mask'].to(device), data['clip_embedding'].to(device) + + return prompt_embeds, clip_embedding, prompt_attention_mask + + def decode_vae(self, samples): + samples = asyncio.run(self.vae(samples.cpu())) + return samples + + def encode_vae(self, img): + latents = asyncio.run(self.vae_encode(img)) + return latents + + def check_inputs(self, num_frames, width, height): + num_frames = max(num_frames//17*17, 1) + width = max(width//16*16, 16) + height = max(height//16*16, 16) + return num_frames, width, height + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: 64, + height: int = 544, + width: int = 992, + num_frames: int = 204, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_frames, width, height = self.check_inputs(num_frames, width, height) + shape = ( + batch_size, + max(num_frames//17*3, 1), + num_channels_latents, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) # b,f,c,h,w + 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 generator is None: + generator = torch.Generator(device=self._execution_device) + + latents = torch.randn(shape, generator=generator, device=device, dtype=dtype) + return latents + + + def resize_to_desired_aspect_ratio(self, video, aspect_size): + ## video is in shape [f, c, h, w] + height, width = video.shape[-2:] + + aspect_ratio = [w/h for h, w in aspect_size] + # # resize + aspect_ratio_fact = width / height + bucket_idx = np.argmin(np.abs(aspect_ratio_fact - np.array(aspect_ratio))) + aspect_ratio = aspect_ratio[bucket_idx] + target_size_height, target_size_width = aspect_size[bucket_idx] + + if aspect_ratio_fact < aspect_ratio: + scale = target_size_width / width + else: + scale = target_size_height / height + + width_scale = int(round(width * scale)) + height_scale = int(round(height * scale)) + + + # # crop + delta_h = height_scale - target_size_height + delta_w = width_scale - target_size_width + assert delta_w>=0 + assert delta_h>=0 + assert not all( + [delta_h, delta_w] + ) + top = delta_h//2 + left = delta_w//2 + + ## resize image and crop + resize_crop_transform = transforms.Compose([ + transforms.Resize((height_scale, width_scale)), + lambda x: transforms.functional.crop(x, top, left, target_size_height, target_size_width), + ]) + + video = torch.stack([resize_crop_transform(frame.contiguous()) for frame in video], dim=0) + return video + + + def prepare_condition_hidden_states( + self, + img: Union[str, PILImage.Image, torch.Tensor]=None, + batch_size: int = 1, + num_channels_latents: int = 64, + height: int = 544, + width: int = 992, + num_frames: int = 204, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None + ): + if isinstance(img, str): + assert os.path.exists(img) + img = PILImage.open(img) + + if isinstance(img, PILImage.Image): + img_tensor = transforms.ToTensor()(img.convert('RGB'))*2-1 + else: + img_tensor = img + + num_frames, width, height = self.check_inputs(num_frames, width, height) + + img_tensor = self.resize_to_desired_aspect_ratio(img_tensor[None], aspect_size=[(height, width)])[None] + + img_emb = self.encode_vae(img_tensor).repeat(batch_size, 1,1,1,1).to(device) + + padding_tensor = torch.zeros((batch_size, max(num_frames//17*3, 1)-1, num_channels_latents, int(height) // self.vae_scale_factor_spatial, int(width) // self.vae_scale_factor_spatial,), device=device) + condition_hidden_states = torch.cat([img_emb, padding_tensor], dim=1) + + condition_hidden_states = condition_hidden_states.repeat(2, 1,1,1,1) ## for CFG + return condition_hidden_states.to(dtype) + + @torch.inference_mode() + def __call__( + self, + prompt: Union[str, List[str]] = None, + height: int = 544, + width: int = 992, + num_frames: int = 102, + num_inference_steps: int = 50, + guidance_scale: float = 9.0, + time_shift: float = 13.0, + neg_magic: str = "", + pos_magic: str = "", + num_videos_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + first_image: Union[str, PILImage.Image, torch.Tensor] = None, + motion_score: float = 2.0, + output_type: Optional[str] = "mp4", + output_file_name: Optional[str] = "", + return_dict: bool = True, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + height (`int`, defaults to `544`): + The height in pixels of the generated image. + width (`int`, defaults to `992`): + The width in pixels of the generated image. + num_frames (`int`, defaults to `204`): + The number of frames in the generated video. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, defaults to `9.0`): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + 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`. + first_image (`str`, `PIL.Image`, `torch.Tensor`): + A path for the reference image + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + output_file_name(`str`, *optional*`): + The output mp4 file name. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`StepVideoPipelineOutput`] instead of a plain tuple. + + Examples: + + Returns: + [`~StepVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`StepVideoPipelineOutput`] 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. + """ + + # 1. Check inputs. Raise error if not correct + device = self._execution_device + + # 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] + + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + prompt_embeds, prompt_embeds_2, prompt_attention_mask = self.encode_prompt( + prompt=prompt, + neg_magic=neg_magic, + pos_magic=pos_magic, + ) + + transformer_dtype = self.transformer.dtype + prompt_embeds = prompt_embeds.to(transformer_dtype) + prompt_attention_mask = prompt_attention_mask.to(transformer_dtype) + prompt_embeds_2 = prompt_embeds_2.to(transformer_dtype) + + # 4. Prepare timesteps + self.scheduler.set_timesteps( + num_inference_steps=num_inference_steps, + time_shift=time_shift, + device=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, + num_frames, + torch.bfloat16, + device, + generator, + latents, + ) + condition_hidden_states = self.prepare_condition_hidden_states( + first_image, + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + dtype=torch.bfloat16, + device=device) + + # 7. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(self.scheduler.timesteps): + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(transformer_dtype) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + encoder_hidden_states_2=prompt_embeds_2, + condition_hidden_states=condition_hidden_states, + motion_score=motion_score, + return_dict=False, + ) + # perform guidance + if do_classifier_free_guidance: + noise_pred_text, noise_pred_uncond = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step( + model_output=noise_pred, + timestep=t, + sample=latents + ) + + progress_bar.update() + + if not torch.distributed.is_initialized() or int(torch.distributed.get_rank())==0: + if not output_type == "latent": + video = self.decode_vae(latents) + video = self.video_processor.postprocess_video(video, output_file_name=output_file_name, output_type=output_type) + else: + video = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, ) + + return StepVideoPipelineOutput(video=video) + + + \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/modules/attentions.py b/Step-Video-TI2V/stepvideo/modules/attentions.py new file mode 100755 index 0000000..c6ef95c --- /dev/null +++ b/Step-Video-TI2V/stepvideo/modules/attentions.py @@ -0,0 +1,62 @@ +import torch +import torch.nn as nn +from einops import rearrange + +try: + from xfuser.core.long_ctx_attention import xFuserLongContextAttention +except ImportError: + xFuserLongContextAttention = None + + +class Attention(nn.Module): + def __init__(self): + super().__init__() + + def attn_processor(self, attn_type): + if attn_type == 'torch': + return self.torch_attn_func + elif attn_type == 'parallel': + return self.parallel_attn_func + else: + raise Exception('Not supported attention type...') + + def torch_attn_func( + self, + q, + k, + v, + attn_mask=None, + causal=False, + drop_rate=0.0, + **kwargs + ): + + if attn_mask is not None and attn_mask.dtype != torch.bool: + attn_mask = attn_mask.to(q.dtype) + + if attn_mask is not None and attn_mask.ndim == 3: ## no head + n_heads = q.shape[2] + attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) + + q, k, v = map(lambda x: rearrange(x, 'b s h d -> b h s d'), (q, k, v)) + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal + ) + x = rearrange(x, 'b h s d -> b s h d') + return x + + def parallel_attn_func( + self, + q, + k, + v, + causal=False, + **kwargs + ): + assert xFuserLongContextAttention is not None; 'to use sequence parallel attention, xFuserLongContextAttention should be imported...' + hybrid_seq_parallel_attn = xFuserLongContextAttention() + x = hybrid_seq_parallel_attn( + None, q,k,v, causal=causal + ) + return x + diff --git a/Step-Video-TI2V/stepvideo/modules/blocks.py b/Step-Video-TI2V/stepvideo/modules/blocks.py new file mode 100755 index 0000000..004b3f6 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/modules/blocks.py @@ -0,0 +1,314 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import torch +import torch.nn as nn +from typing import Optional +from einops import rearrange +from stepvideo.modules.rope import RoPE3D +from stepvideo.modules.attentions import Attention +from stepvideo.modules.normalization import RMSNorm + + +class SelfAttention(Attention): + def __init__(self, hidden_dim, head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type='torch'): + super().__init__() + self.head_dim = head_dim + self.n_heads = hidden_dim // head_dim + + self.wqkv = nn.Linear(hidden_dim, hidden_dim*3, bias=bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) + + self.with_rope = with_rope + self.with_qk_norm = with_qk_norm + if self.with_qk_norm: + self.q_norm = RMSNorm(head_dim, elementwise_affine=True) + self.k_norm = RMSNorm(head_dim, elementwise_affine=True) + + if self.with_rope: + self.rope_3d = RoPE3D(freq=1e4, F0=1.0, scaling_factor=1.0) + self.rope_ch_split = [64, 32, 32] + + self.core_attention = self.attn_processor(attn_type=attn_type) + self.parallel = attn_type=='parallel' + + def apply_rope3d(self, x, fhw_positions, rope_ch_split, parallel=True): + x = self.rope_3d(x, fhw_positions, rope_ch_split, parallel) + return x + + def forward( + self, + x, + cu_seqlens=None, + max_seqlen=None, + rope_positions=None, + attn_mask=None + ): + xqkv = self.wqkv(x) + xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3*self.head_dim) + + xq, xk, xv = torch.split(xqkv, [self.head_dim]*3, dim=-1) ## seq_len, n, dim + + if self.with_qk_norm: + xq = self.q_norm(xq) + xk = self.k_norm(xk) + + if self.with_rope: + xq = self.apply_rope3d(xq, rope_positions, self.rope_ch_split, parallel=self.parallel) + xk = self.apply_rope3d(xk, rope_positions, self.rope_ch_split, parallel=self.parallel) + + output = self.core_attention( + xq, + xk, + xv, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + attn_mask=attn_mask + ) + output = rearrange(output, 'b s h d -> b s (h d)') + output = self.wo(output) + + return output + + +class CrossAttention(Attention): + def __init__(self, hidden_dim, head_dim, bias=False, with_qk_norm=True, attn_type='torch'): + super().__init__() + self.head_dim = head_dim + self.n_heads = hidden_dim // head_dim + + self.wq = nn.Linear(hidden_dim, hidden_dim, bias=bias) + self.wkv = nn.Linear(hidden_dim, hidden_dim*2, bias=bias) + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=bias) + + self.with_qk_norm = with_qk_norm + if self.with_qk_norm: + self.q_norm = RMSNorm(head_dim, elementwise_affine=True) + self.k_norm = RMSNorm(head_dim, elementwise_affine=True) + + self.core_attention = self.attn_processor(attn_type=attn_type) + + def forward( + self, + x: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attn_mask=None + ): + xq = self.wq(x) + xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim) + + xkv = self.wkv(encoder_hidden_states) + xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2*self.head_dim) + + xk, xv = torch.split(xkv, [self.head_dim]*2, dim=-1) ## seq_len, n, dim + + if self.with_qk_norm: + xq = self.q_norm(xq) + xk = self.k_norm(xk) + + output = self.core_attention( + xq, + xk, + xv, + attn_mask=attn_mask + ) + + output = rearrange(output, 'b s h d -> b s (h d)') + output = self.wo(output) + + return output + + +class GELU(nn.Module): + r""" + GELU activation function with tanh approximation support with `approximate="tanh"`. + + Parameters: + dim_in (`int`): The number of channels in the input. + dim_out (`int`): The number of channels in the output. + approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. + bias (`bool`, defaults to True): Whether to use a bias in the linear layer. + """ + + def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out, bias=bias) + self.approximate = approximate + + def gelu(self, gate: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(gate, approximate=self.approximate) + + def forward(self, hidden_states): + hidden_states = self.proj(hidden_states) + hidden_states = self.gelu(hidden_states) + return hidden_states + + +class FeedForward(nn.Module): + def __init__( + self, + dim: int, + inner_dim: Optional[int] = None, + dim_out: Optional[int] = None, + mult: int = 4, + bias: bool = False, + ): + super().__init__() + inner_dim = dim*mult if inner_dim is None else inner_dim + dim_out = dim if dim_out is None else dim_out + self.net = nn.ModuleList([ + GELU(dim, inner_dim, approximate="tanh", bias=bias), + nn.Identity(), + nn.Linear(inner_dim, dim_out, bias=bias) + ]) + + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + for module in self.net: + hidden_states = module(hidden_states) + return hidden_states + + +def modulate(x, scale, shift): + x = x * (1 + scale) + shift + return x + +def gate(x, gate): + x = gate * x + return x + + +class StepVideoTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__( + self, + dim: int, + attention_head_dim: int, + norm_eps: float = 1e-5, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = False, + attention_type: str = 'parallel' + ): + super().__init__() + self.dim = dim + self.norm1 = nn.LayerNorm(dim, eps=norm_eps) + self.attn1 = SelfAttention(dim, attention_head_dim, bias=False, with_rope=True, with_qk_norm=True, attn_type=attention_type) + + self.norm2 = nn.LayerNorm(dim, eps=norm_eps) + self.attn2 = CrossAttention(dim, attention_head_dim, bias=False, with_qk_norm=True, attn_type='torch') + + self.ff = FeedForward(dim=dim, inner_dim=ff_inner_dim, dim_out=dim, bias=ff_bias) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) /dim**0.5) + + @torch.no_grad() + def forward( + self, + q: torch.Tensor, + kv: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + attn_mask = None, + rope_positions: list = None, + ) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + torch.clone(chunk) for chunk in (self.scale_shift_table[None] + timestep.reshape(-1, 6, self.dim)).chunk(6, dim=1) + ) + + scale_shift_q = modulate(self.norm1(q), scale_msa, shift_msa) + + attn_q = self.attn1( + scale_shift_q, + rope_positions=rope_positions + ) + + q = gate(attn_q, gate_msa) + q + + attn_q = self.attn2( + q, + kv, + attn_mask + ) + + q = attn_q + q + + scale_shift_q = modulate(self.norm2(q), scale_mlp, shift_mlp) + + ff_output = self.ff(scale_shift_q) + + q = gate(ff_output, gate_mlp) + q + + return q + + +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding""" + + def __init__( + self, + patch_size=64, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + ): + super().__init__() + + self.flatten = flatten + self.layer_norm = layer_norm + + self.proj = nn.Conv2d( + in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias + ) + + def forward(self, latent): + latent = self.proj(latent).to(latent.dtype) + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC + if self.layer_norm: + latent = self.norm(latent) + + return latent + \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/modules/model.py b/Step-Video-TI2V/stepvideo/modules/model.py new file mode 100755 index 0000000..c3167ac --- /dev/null +++ b/Step-Video-TI2V/stepvideo/modules/model.py @@ -0,0 +1,220 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +from typing import Any, Dict, Optional, Union +import torch +from torch import nn +import os +from einops import rearrange, repeat +from stepvideo.modules.blocks import StepVideoTransformerBlock, PatchEmbed + +from stepvideo.utils import with_empty_init +from stepvideo.parallel import parallel_forward + +from stepvideo.modules.normalization import ( + PixArtAlphaTextProjection, + AdaLayerNormSingle + ) + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + + + +class StepVideoModel(ModelMixin, ConfigMixin): + _no_split_modules = ["StepVideoTransformerBlock", "PatchEmbed"] + + @with_empty_init + @register_to_config + def __init__( + self, + num_attention_heads: int = 48, + attention_head_dim: int = 128, + in_channels: int = 64, + out_channels: Optional[int] = 64, + num_layers: int = 48, + dropout: float = 0.0, + patch_size: int = 1, + norm_type: str = "ada_norm_single", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + use_additional_conditions: Optional[bool] = False, + caption_channels: Optional[int]|list|tuple = [6144, 1024], + attention_type: Optional[str] = "parallel", + ): + super().__init__() + + # Set some common variables used across the board. + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.out_channels = in_channels if out_channels is None else out_channels + + self.use_additional_conditions = use_additional_conditions + + self.pos_embed = PatchEmbed( + patch_size=patch_size, + in_channels=self.config.in_channels if not use_additional_conditions else self.config.in_channels*2, + embed_dim=self.inner_dim, + ) + + self.transformer_blocks = nn.ModuleList( + [ + StepVideoTransformerBlock( + dim=self.inner_dim, + attention_head_dim=self.config.attention_head_dim, + attention_type=attention_type + ) + for _ in range(self.config.num_layers) + ] + ) + + # 3. Output blocks. + self.norm_out = nn.LayerNorm(self.inner_dim, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels) + self.patch_size = patch_size + + self.adaln_single = AdaLayerNormSingle( + self.inner_dim, use_additional_conditions=self.use_additional_conditions + ) + + if isinstance(self.config.caption_channels, int): + caption_channel = self.config.caption_channels + else: + caption_channel, clip_channel = self.config.caption_channels + self.clip_projection = nn.Linear(clip_channel, self.inner_dim) + + self.caption_norm = nn.LayerNorm(caption_channel, eps=norm_eps, elementwise_affine=norm_elementwise_affine) + + self.caption_projection = PixArtAlphaTextProjection( + in_features=caption_channel, hidden_size=self.inner_dim + ) + + self.parallel = attention_type=='parallel' + + def patchfy(self, hidden_states, condition_hidden_states=None): + if condition_hidden_states is not None: + hidden_states = torch.cat([hidden_states, condition_hidden_states], dim=2) + + hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w') + hidden_states = self.pos_embed(hidden_states) + return hidden_states + + def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states, q_seqlen): + kv_seqlens = encoder_attention_mask.sum(dim=1).int() + mask = torch.zeros([len(kv_seqlens), q_seqlen, max(kv_seqlens)], dtype=torch.bool, device=encoder_attention_mask.device) + encoder_hidden_states = encoder_hidden_states[:,: max(kv_seqlens)] + for i, kv_len in enumerate(kv_seqlens): + mask[i, :, :kv_len] = 1 + return encoder_hidden_states, mask + + + @parallel_forward + def block_forward( + self, + hidden_states, + encoder_hidden_states=None, + timestep=None, + rope_positions=None, + attn_mask=None, + parallel=True + ): + + for i, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep=timestep, + attn_mask=attn_mask, + rope_positions=rope_positions + ) + + return hidden_states + + + @torch.inference_mode() + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_hidden_states_2: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + fps: torch.Tensor=None, + condition_hidden_states: torch.Tensor = None, + motion_score: torch.Tensor=None, + return_dict: bool = True, + ): + assert hidden_states.ndim==5; "hidden_states's shape should be (bsz, f, ch, h ,w)" + + bsz, frame, _, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + + hidden_states = self.patchfy(hidden_states, condition_hidden_states) + len_frame = hidden_states.shape[1] + + if self.use_additional_conditions: + added_cond_kwargs = { + "motion_score": torch.tensor([motion_score], device=hidden_states.device, dtype=hidden_states.dtype).repeat(bsz), + } + else: + added_cond_kwargs = {} + + timestep, embedded_timestep = self.adaln_single( + timestep, added_cond_kwargs=added_cond_kwargs + ) + + encoder_hidden_states = self.caption_projection(self.caption_norm(encoder_hidden_states)) + + if encoder_hidden_states_2 is not None and hasattr(self, 'clip_projection'): + clip_embedding = self.clip_projection(encoder_hidden_states_2) + encoder_hidden_states = torch.cat([clip_embedding, encoder_hidden_states], dim=1) + + hidden_states = rearrange(hidden_states, '(b f) l d-> b (f l) d', b=bsz, f=frame, l=len_frame).contiguous() + encoder_hidden_states, attn_mask = self.prepare_attn_mask(encoder_attention_mask, encoder_hidden_states, q_seqlen=frame*len_frame) + + hidden_states = self.block_forward( + hidden_states, + encoder_hidden_states, + timestep=timestep, + rope_positions=[frame, height, width], + attn_mask=attn_mask, + parallel=self.parallel + ) + + hidden_states = rearrange(hidden_states, 'b (f l) d -> (b f) l d', b=bsz, f=frame, l=len_frame) + + embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d', f=frame).contiguous() + + shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) + # Modulation + hidden_states = hidden_states * (1 + scale) + shift + hidden_states = self.proj_out(hidden_states) + + # unpatchify + hidden_states = hidden_states.reshape( + shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) + ) + + hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q') + output = hidden_states.reshape( + shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) + ) + + output = rearrange(output, '(b f) c h w -> b f c h w', f=frame) + + if return_dict: + return {'x': output} + return output + + \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/modules/normalization.py b/Step-Video-TI2V/stepvideo/modules/normalization.py new file mode 100755 index 0000000..1cdedbc --- /dev/null +++ b/Step-Video-TI2V/stepvideo/modules/normalization.py @@ -0,0 +1,309 @@ +from typing import Any, Dict, Optional, Union, Tuple +import torch +import torch.nn as nn +import math + + +class RMSNorm(nn.Module): + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x): + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +ACTIVATION_FUNCTIONS = { + "swish": nn.SiLU(), + "silu": nn.SiLU(), + "mish": nn.Mish(), + "gelu": nn.GELU(), + "relu": nn.ReLU(), +} + + +def get_activation(act_fn: str) -> nn.Module: + """Helper function to get activation function from string. + + Args: + act_fn (str): Name of activation function. + + Returns: + nn.Module: Activation function. + """ + + act_fn = act_fn.lower() + if act_fn in ACTIVATION_FUNCTIONS: + return ACTIVATION_FUNCTIONS[act_fn] + else: + raise ValueError(f"Unsupported activation function: {act_fn}") + + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the + embeddings. :return: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + + +class Timesteps(nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + + def forward(self, timesteps): + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + ) + return t_emb + + + +class TimestepEmbedding(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + sample_proj_bias=True + ): + super().__init__() + linear_cls = nn.Linear + + self.linear_1 = linear_cls( + in_channels, + time_embed_dim, + bias=sample_proj_bias, + ) + + if cond_proj_dim is not None: + self.cond_proj = linear_cls( + cond_proj_dim, + in_channels, + bias=False, + ) + else: + self.cond_proj = None + + self.act = get_activation(act_fn) + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + + self.linear_2 = linear_cls( + time_embed_dim, + time_embed_dim_out, + bias=sample_proj_bias, + ) + + if post_act_fn is None: + self.post_act = None + else: + self.post_act = get_activation(post_act_fn) + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + + + +class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): + def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.use_additional_conditions = use_additional_conditions + if self.use_additional_conditions: + self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.motion_score_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timestep, motion_score=None): + hidden_dtype = next(self.timestep_embedder.parameters()).dtype + + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + if self.use_additional_conditions: + batch_size = timestep.shape[0] + motion_score_emb = self.additional_condition_proj(motion_score.flatten()).to(hidden_dtype) + motion_score_emb = self.motion_score_embedder(motion_score_emb).reshape(batch_size, -1) + conditioning = timesteps_emb + motion_score_emb + + else: + conditioning = timesteps_emb + + return conditioning + + + +class AdaLayerNormSingle(nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + + As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + use_additional_conditions (`bool`): To use additional conditions for normalization or not. + """ + def __init__(self, embedding_dim: int, use_additional_conditions: bool = False, time_step_rescale=1000): + super().__init__() + + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim, size_emb_dim=embedding_dim // 2, use_additional_conditions=use_additional_conditions + ) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + self.time_step_rescale = time_step_rescale ## timestep usually in [0, 1], we rescale it to [0,1000] for stability + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + ): + embedded_timestep = self.emb(timestep*self.time_step_rescale, **added_cond_kwargs) + + out = self.linear(self.silu(embedded_timestep)) + + return out, embedded_timestep + + +class PixArtAlphaTextProjection(nn.Module): + """ + Projects caption 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_features, hidden_size): + super().__init__() + self.linear_1 = nn.Linear( + in_features, + hidden_size, + bias=True, + ) + self.act_1 = nn.GELU(approximate="tanh") + self.linear_2 = nn.Linear( + hidden_size, + hidden_size, + bias=True, + ) + + 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 + diff --git a/Step-Video-TI2V/stepvideo/modules/rope.py b/Step-Video-TI2V/stepvideo/modules/rope.py new file mode 100755 index 0000000..ca03f1c --- /dev/null +++ b/Step-Video-TI2V/stepvideo/modules/rope.py @@ -0,0 +1,90 @@ +import torch +from stepvideo.parallel import get_sequence_parallel_world_size, get_sequence_parallel_rank + + +class RoPE1D: + def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): + self.base = freq + self.F0 = F0 + self.scaling_factor = scaling_factor + self.cache = {} + + def get_cos_sin(self, D, seq_len, device, dtype): + if (D, seq_len, device, dtype) not in self.cache: + inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D)) + t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) + freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype) + freqs = torch.cat((freqs, freqs), dim=-1) + cos = freqs.cos() # (Seq, Dim) + sin = freqs.sin() + self.cache[D, seq_len, device, dtype] = (cos, sin) + return self.cache[D, seq_len, device, dtype] + + @staticmethod + def rotate_half(x): + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + def apply_rope1d(self, tokens, pos1d, cos, sin): + assert pos1d.ndim == 2 + cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :] + sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :] + return (tokens * cos) + (self.rotate_half(tokens) * sin) + + def __call__(self, tokens, positions): + """ + input: + * tokens: batch_size x ntokens x nheads x dim + * positions: batch_size x ntokens (t position of each token) + output: + * tokens after appplying RoPE2D (batch_size x ntokens x nheads x dim) + """ + D = tokens.size(3) + assert positions.ndim == 2 # Batch, Seq + cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype) + tokens = self.apply_rope1d(tokens, positions, cos, sin) + return tokens + + + +class RoPE3D(RoPE1D): + def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): + super(RoPE3D, self).__init__(freq, F0, scaling_factor) + self.position_cache = {} + + def get_mesh_3d(self, rope_positions, bsz): + f, h, w = rope_positions + + if f"{f}-{h}-{w}" not in self.position_cache: + x = torch.arange(f, device='cpu') + y = torch.arange(h, device='cpu') + z = torch.arange(w, device='cpu') + self.position_cache[f"{f}-{h}-{w}"] = torch.cartesian_prod(x, y, z).view(1, f*h*w, 3).expand(bsz, -1, 3) + return self.position_cache[f"{f}-{h}-{w}"] + + def __call__(self, tokens, rope_positions, ch_split, parallel=False): + """ + input: + * tokens: batch_size x ntokens x nheads x dim + * rope_positions: list of (f, h, w) + output: + * tokens after appplying RoPE2D (batch_size x ntokens x nheads x dim) + """ + assert sum(ch_split) == tokens.size(-1); + + mesh_grid = self.get_mesh_3d(rope_positions, bsz=tokens.shape[0]) + out = [] + for i, (D, x) in enumerate(zip(ch_split, torch.split(tokens, ch_split, dim=-1))): + cos, sin = self.get_cos_sin(D, int(mesh_grid.max()) + 1, tokens.device, tokens.dtype) + + if parallel: + mesh = torch.chunk(mesh_grid[:, :, i], get_sequence_parallel_world_size(),dim=1)[get_sequence_parallel_rank()].clone() + else: + mesh = mesh_grid[:, :, i].clone() + x = self.apply_rope1d(x, mesh.to(tokens.device), cos, sin) + out.append(x) + + tokens = torch.cat(out, dim=-1) + return tokens + + diff --git a/Step-Video-TI2V/stepvideo/parallel.py b/Step-Video-TI2V/stepvideo/parallel.py new file mode 100644 index 0000000..ed785bd --- /dev/null +++ b/Step-Video-TI2V/stepvideo/parallel.py @@ -0,0 +1,46 @@ +import torch.distributed as dist +import xfuser +import torch + + +def initialize_parall_group(ring_degree, ulysses_degree): + dist.init_process_group("nccl") + xfuser.core.distributed.init_distributed_environment( + rank=dist.get_rank(), + world_size=dist.get_world_size() + ) + + xfuser.core.distributed.initialize_model_parallel( + sequence_parallel_degree=ulysses_degree, + ring_degree=ring_degree, + ulysses_degree=ulysses_degree, + ) + torch.cuda.set_device(dist.get_rank()) + +def get_parallel_group(): + return xfuser.core.distributed.get_world_group() + +def get_sequence_parallel_world_size(): + return xfuser.core.distributed.parallel_state.get_sequence_parallel_world_size() + +def get_sequence_parallel_rank(): + return xfuser.core.distributed.parallel_state.get_sequence_parallel_rank() + +def get_sp_group(): + return xfuser.core.distributed.parallel_state.get_sp_group() + + + +def parallel_forward(fn_): + def wrapTheFunction(_, hidden_states, *args, **kwargs): + if kwargs['parallel']: + hidden_states = torch.chunk(hidden_states, get_sequence_parallel_world_size(), dim=-2)[get_sequence_parallel_rank()] + kwargs['attn_mask'] = torch.chunk(kwargs['attn_mask'], get_sequence_parallel_world_size(), dim=-2)[get_sequence_parallel_rank()] + output = fn_(_, hidden_states, *args, **kwargs) + + if kwargs['parallel']: + output = get_sp_group().all_gather(output.contiguous(), dim=-2) + + return output + + return wrapTheFunction \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/text_encoder/__init__.py b/Step-Video-TI2V/stepvideo/text_encoder/__init__.py new file mode 100644 index 0000000..a114abd --- /dev/null +++ b/Step-Video-TI2V/stepvideo/text_encoder/__init__.py @@ -0,0 +1,20 @@ +import torch +from stepvideo.config import parse_args +import os + + +accepted_version = { + '2.2': 'liboptimus_ths-torch2.2-cu121.cpython-310-x86_64-linux-gnu.so', + '2.3': 'liboptimus_ths-torch2.3-cu121.cpython-310-x86_64-linux-gnu.so', + '2.5': 'liboptimus_ths-torch2.5-cu124.cpython-310-x86_64-linux-gnu.so', +} + +try: + args = parse_args() + version = '.'.join(torch.__version__.split('.')[:2]) + if version in accepted_version: + torch.ops.load_library(os.path.join(args.model_dir, f'lib/{accepted_version[version]}')) + else: + raise ValueError("Not supported torch version for liboptimus") +except Exception as err: + print(err) diff --git a/Step-Video-TI2V/stepvideo/text_encoder/clip.py b/Step-Video-TI2V/stepvideo/text_encoder/clip.py new file mode 100755 index 0000000..1e5574b --- /dev/null +++ b/Step-Video-TI2V/stepvideo/text_encoder/clip.py @@ -0,0 +1,35 @@ +import torch +import torch.nn as nn +from transformers import BertModel, BertTokenizer +import os + + +class HunyuanClip(nn.Module): + """ + Hunyuan clip code copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py + hunyuan's clip used BertModel and BertTokenizer, so we copy it. + """ + def __init__(self, model_dir, max_length=77): + super(HunyuanClip, self).__init__() + + self.max_length = max_length + self.tokenizer = BertTokenizer.from_pretrained(os.path.join(model_dir, 'tokenizer')) + self.text_encoder = BertModel.from_pretrained(os.path.join(model_dir, 'clip_text_encoder')) + + @torch.no_grad + def forward(self, prompts, with_mask=True): + self.device = next(self.text_encoder.parameters()).device + text_inputs = self.tokenizer( + prompts, + padding="max_length", + max_length=self.max_length, + truncation=True, + return_attention_mask=True, + return_tensors="pt", + ) + prompt_embeds = self.text_encoder( + text_inputs.input_ids.to(self.device), + attention_mask=text_inputs.attention_mask.to(self.device) if with_mask else None, + ) + return prompt_embeds.last_hidden_state, prompt_embeds.pooler_output + \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/text_encoder/flashattention.py b/Step-Video-TI2V/stepvideo/text_encoder/flashattention.py new file mode 100755 index 0000000..5fe543f --- /dev/null +++ b/Step-Video-TI2V/stepvideo/text_encoder/flashattention.py @@ -0,0 +1,37 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import torch + +def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=True, + return_attn_probs=False, tp_group_rank=0, tp_group_size=1): + softmax_scale = q.size(-1) ** (-0.5) if softmax_scale is None else softmax_scale + return torch.ops.Optimus.fwd(q, k, v, None, dropout_p, softmax_scale, causal, return_attn_probs, None, tp_group_rank, tp_group_size)[0] + + +class FlashSelfAttention(torch.nn.Module): + def __init__( + self, + attention_dropout=0.0, + ): + super().__init__() + self.dropout_p = attention_dropout + + + def forward(self, q, k, v, cu_seqlens=None, max_seq_len=None): + if cu_seqlens is None: + output = flash_attn_func(q, k, v, dropout_p=self.dropout_p) + else: + raise ValueError('cu_seqlens is not supported!') + + return output + diff --git a/Step-Video-TI2V/stepvideo/text_encoder/stepllm.py b/Step-Video-TI2V/stepvideo/text_encoder/stepllm.py new file mode 100755 index 0000000..951dbd9 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/text_encoder/stepllm.py @@ -0,0 +1,303 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import os +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from stepvideo.text_encoder.flashattention import FlashSelfAttention +from stepvideo.modules.normalization import RMSNorm +from stepvideo.text_encoder.tokenizer import LLaMaEmbedding, Wrapped_StepChatTokenizer +from stepvideo.utils import with_empty_init +from safetensors.torch import load_file +from transformers.modeling_utils import PretrainedConfig, PreTrainedModel +from einops import rearrange +import json + + + +def safediv(n, d): + q, r = divmod(n, d) + assert r == 0 + return q + + +class MultiQueryAttention(nn.Module): + def __init__(self, cfg, layer_id=None): + super().__init__() + + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.max_seq_len = cfg.seq_length + self.use_flash_attention = cfg.use_flash_attn + assert self.use_flash_attention, 'FlashAttention is required!' + + self.n_groups = cfg.num_attention_groups + self.tp_size = 1 + self.n_local_heads = cfg.num_attention_heads + self.n_local_groups = self.n_groups + + self.wqkv = nn.Linear( + cfg.hidden_size, + cfg.hidden_size + self.head_dim * 2 * self.n_groups, + bias=False, + ) + self.wo = nn.Linear( + cfg.hidden_size, + cfg.hidden_size, + bias=False, + ) + + assert self.use_flash_attention, 'non-Flash attention not supported yet.' + self.core_attention = FlashSelfAttention(attention_dropout=cfg.attention_dropout) + + self.layer_id = layer_id + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], + max_seq_len: Optional[torch.Tensor], + ): + seqlen, bsz, dim = x.shape + xqkv = self.wqkv(x) + + xq, xkv = torch.split( + xqkv, + (dim // self.tp_size, + self.head_dim*2*self.n_groups // self.tp_size + ), + dim=-1, + ) + + # gather on 1st dimension + xq = xq.view(seqlen, bsz, self.n_local_heads, self.head_dim) + xkv = xkv.view(seqlen, bsz, self.n_local_groups, 2 * self.head_dim) + xk, xv = xkv.chunk(2, -1) + + # rotary embedding + flash attn + xq = rearrange(xq, "s b h d -> b s h d") + xk = rearrange(xk, "s b h d -> b s h d") + xv = rearrange(xv, "s b h d -> b s h d") + + q_per_kv = self.n_local_heads // self.n_local_groups + if q_per_kv > 1: + b, s, h, d = xk.size() + if h == 1: + xk = xk.expand(b, s, q_per_kv, d) + xv = xv.expand(b, s, q_per_kv, d) + else: + ''' To cover the cases where h > 1, we have + the following implementation, which is equivalent to: + xk = xk.repeat_interleave(q_per_kv, dim=-2) + xv = xv.repeat_interleave(q_per_kv, dim=-2) + but can avoid calling aten::item() that involves cpu. + ''' + idx = torch.arange(q_per_kv * h, device=xk.device).reshape(q_per_kv, -1).permute(1, 0).flatten() + xk = torch.index_select(xk.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + xv = torch.index_select(xv.repeat(1, 1, q_per_kv, 1), 2, idx).contiguous() + + if self.use_flash_attention: + output = self.core_attention(xq, xk, xv, + cu_seqlens=cu_seqlens, + max_seq_len=max_seq_len) + # reduce-scatter only support first dimention now + output = rearrange(output, "b s h d -> s b (h d)").contiguous() + else: + xq, xk, xv = [ + rearrange(x, "b s ... -> s b ...").contiguous() + for x in (xq, xk, xv) + ] + output = self.core_attention(xq, xk, xv, mask) + output = self.wo(output) + return output + + + +class FeedForward(nn.Module): + def __init__( + self, + cfg, + dim: int, + hidden_dim: int, + layer_id: int, + multiple_of: int=256, + ): + super().__init__() + + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + self.swiglu = swiglu + + self.w1 = nn.Linear( + dim, + 2 * hidden_dim, + bias=False, + ) + self.w2 = nn.Linear( + hidden_dim, + dim, + bias=False, + ) + + def forward(self, x): + x = self.swiglu(self.w1(x)) + output = self.w2(x) + return output + + + +class TransformerBlock(nn.Module): + def __init__( + self, cfg, layer_id: int + ): + super().__init__() + + self.n_heads = cfg.num_attention_heads + self.dim = cfg.hidden_size + self.head_dim = cfg.hidden_size // cfg.num_attention_heads + self.attention = MultiQueryAttention( + cfg, + layer_id=layer_id, + ) + + self.feed_forward = FeedForward( + cfg, + dim=cfg.hidden_size, + hidden_dim=cfg.ffn_hidden_size, + layer_id=layer_id, + ) + self.layer_id = layer_id + self.attention_norm = RMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + self.ffn_norm = RMSNorm( + cfg.hidden_size, + eps=cfg.layernorm_epsilon, + ) + + def forward( + self, + x: torch.Tensor, + mask: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], + max_seq_len: Optional[torch.Tensor], + ): + residual = self.attention.forward( + self.attention_norm(x), mask, + cu_seqlens, max_seq_len + ) + h = x + residual + ffn_res = self.feed_forward.forward(self.ffn_norm(h)) + out = h + ffn_res + return out + + +class Transformer(nn.Module): + def __init__( + self, + config, + max_seq_size=8192, + ): + super().__init__() + self.num_layers = config.num_layers + self.layers = self._build_layers(config) + + def _build_layers(self, config): + layers = torch.nn.ModuleList() + for layer_id in range(self.num_layers): + layers.append( + TransformerBlock( + config, + layer_id=layer_id + 1 , + ) + ) + return layers + + def forward( + self, + hidden_states, + attention_mask, + cu_seqlens=None, + max_seq_len=None, + ): + + if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor): + max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu") + + for lid, layer in enumerate(self.layers): + hidden_states = layer( + hidden_states, + attention_mask, + cu_seqlens, + max_seq_len, + ) + return hidden_states + + +class Step1Model(PreTrainedModel): + config_class=PretrainedConfig + @with_empty_init + def __init__( + self, + config, + ): + super().__init__(config) + self.tok_embeddings = LLaMaEmbedding(config) + self.transformer = Transformer(config) + + def forward( + self, + input_ids=None, + attention_mask=None, + ): + + hidden_states = self.tok_embeddings(input_ids) + + hidden_states = self.transformer( + hidden_states, + attention_mask, + ) + return hidden_states + + + +class STEP1TextEncoder(torch.nn.Module): + def __init__(self, model_dir, max_length=320): + super(STEP1TextEncoder, self).__init__() + self.max_length = max_length + self.text_tokenizer = Wrapped_StepChatTokenizer(os.path.join(model_dir, 'step1_chat_tokenizer.model')) + text_encoder = Step1Model.from_pretrained(model_dir) + self.text_encoder = text_encoder.eval().to(torch.bfloat16) + + @torch.no_grad + def forward(self, prompts, with_mask=True, max_length=None): + self.device = next(self.text_encoder.parameters()).device + with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16): + if type(prompts) is str: + prompts = [prompts] + + txt_tokens = self.text_tokenizer( + prompts, max_length=max_length or self.max_length, padding="max_length", truncation=True, return_tensors="pt" + ) + y = self.text_encoder( + txt_tokens.input_ids.to(self.device), + attention_mask=txt_tokens.attention_mask.to(self.device) if with_mask else None + ) + y_mask = txt_tokens.attention_mask + return y.transpose(0,1), y_mask + diff --git a/Step-Video-TI2V/stepvideo/text_encoder/tokenizer.py b/Step-Video-TI2V/stepvideo/text_encoder/tokenizer.py new file mode 100755 index 0000000..6fa7b48 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/text_encoder/tokenizer.py @@ -0,0 +1,206 @@ +# Copyright 2025 StepFun Inc. All Rights Reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# ============================================================================== +import torch.nn as nn +import torch +from typing import List + + +class LLaMaEmbedding(nn.Module): + """Language model embeddings. + + Arguments: + hidden_size: hidden size + vocab_size: vocabulary size + max_sequence_length: maximum size of sequence. This + is used for positional embedding + embedding_dropout_prob: dropout probability for embeddings + init_method: weight initialization method + num_tokentypes: size of the token-type embeddings. 0 value + will ignore this embedding + """ + + def __init__(self, + cfg, + ): + super().__init__() + self.hidden_size = cfg.hidden_size + self.params_dtype = cfg.params_dtype + self.fp32_residual_connection = cfg.fp32_residual_connection + self.embedding_weights_in_fp32 = cfg.embedding_weights_in_fp32 + self.word_embeddings = torch.nn.Embedding( + cfg.padded_vocab_size, self.hidden_size, + ) + self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) + + def forward(self, input_ids): + # Embeddings. + if self.embedding_weights_in_fp32: + self.word_embeddings = self.word_embeddings.to(torch.float32) + embeddings = self.word_embeddings(input_ids) + if self.embedding_weights_in_fp32: + embeddings = embeddings.to(self.params_dtype) + self.word_embeddings = self.word_embeddings.to(self.params_dtype) + + # Data format change to avoid explicit tranposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + + # Dropout. + embeddings = self.embedding_dropout(embeddings) + + return embeddings + + + +class StepChatTokenizer: + """Step Chat Tokenizer""" + + def __init__( + self, model_file, name="StepChatTokenizer", + bot_token="<|BOT|>", # Begin of Turn + eot_token="<|EOT|>", # End of Turn + call_start_token="<|CALL_START|>", # Call Start + call_end_token="<|CALL_END|>", # Call End + think_start_token="<|THINK_START|>", # Think Start + think_end_token="<|THINK_END|>", # Think End + mask_start_token="<|MASK_1e69f|>", # Mask start + mask_end_token="<|UNMASK_1e69f|>", # Mask end + ): + import sentencepiece + + self._tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) + + self._vocab = {} + self._inv_vocab = {} + + self._special_tokens = {} + self._inv_special_tokens = {} + + self._t5_tokens = [] + + for idx in range(self._tokenizer.get_piece_size()): + text = self._tokenizer.id_to_piece(idx) + self._inv_vocab[idx] = text + self._vocab[text] = idx + + if self._tokenizer.is_control(idx) or self._tokenizer.is_unknown(idx): + self._special_tokens[text] = idx + self._inv_special_tokens[idx] = text + + self._unk_id = self._tokenizer.unk_id() + self._bos_id = self._tokenizer.bos_id() + self._eos_id = self._tokenizer.eos_id() + + for token in [ + bot_token, eot_token, call_start_token, call_end_token, + think_start_token, think_end_token + ]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + assert token in self._special_tokens, f"Token '{token}' is not a special token" + + for token in [mask_start_token, mask_end_token]: + assert token in self._vocab, f"Token '{token}' not found in tokenizer" + + self._bot_id = self._tokenizer.piece_to_id(bot_token) + self._eot_id = self._tokenizer.piece_to_id(eot_token) + self._call_start_id = self._tokenizer.piece_to_id(call_start_token) + self._call_end_id = self._tokenizer.piece_to_id(call_end_token) + self._think_start_id = self._tokenizer.piece_to_id(think_start_token) + self._think_end_id = self._tokenizer.piece_to_id(think_end_token) + self._mask_start_id = self._tokenizer.piece_to_id(mask_start_token) + self._mask_end_id = self._tokenizer.piece_to_id(mask_end_token) + + self._underline_id = self._tokenizer.piece_to_id("\u2581") + + @property + def vocab(self): + return self._vocab + + @property + def inv_vocab(self): + return self._inv_vocab + + @property + def vocab_size(self): + return self._tokenizer.vocab_size() + + def tokenize(self, text: str) -> List[int]: + return self._tokenizer.encode_as_ids(text) + + def detokenize(self, token_ids: List[int]) -> str: + return self._tokenizer.decode_ids(token_ids) + + +class Tokens: + def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len) -> None: + self.input_ids = input_ids + self.attention_mask = attention_mask + self.cu_input_ids = cu_input_ids + self.cu_seqlens = cu_seqlens + self.max_seq_len = max_seq_len + def to(self, device): + self.input_ids = self.input_ids.to(device) + self.attention_mask = self.attention_mask.to(device) + self.cu_input_ids = self.cu_input_ids.to(device) + self.cu_seqlens = self.cu_seqlens.to(device) + return self + +class Wrapped_StepChatTokenizer(StepChatTokenizer): + def __call__(self, text, max_length=320, padding="max_length", truncation=True, return_tensors="pt"): + # [bos, ..., eos, pad, pad, ..., pad] + self.BOS = 1 + self.EOS = 2 + self.PAD = 2 + out_tokens = [] + attn_mask = [] + if len(text) == 0: + part_tokens = [self.BOS] + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1]*valid_size+[0]*(max_length-valid_size)) + else: + for part in text: + part_tokens = self.tokenize(part) + part_tokens = part_tokens[:(max_length - 2)] # leave 2 space for bos and eos + part_tokens = [self.BOS] + part_tokens + [self.EOS] + valid_size = len(part_tokens) + if len(part_tokens) < max_length: + part_tokens += [self.PAD] * (max_length - valid_size) + out_tokens.append(part_tokens) + attn_mask.append([1]*valid_size+[0]*(max_length-valid_size)) + + out_tokens = torch.tensor(out_tokens, dtype=torch.long) + attn_mask = torch.tensor(attn_mask, dtype=torch.long) + + # padding y based on tp size + padded_len = 0 + padded_flag = True if padded_len > 0 else False + if padded_flag: + pad_tokens = torch.tensor([[self.PAD] * max_length], device=out_tokens.device) + pad_attn_mask = torch.tensor([[1]*padded_len+[0]*(max_length-padded_len)], device=attn_mask.device) + out_tokens = torch.cat([out_tokens, pad_tokens], dim=0) + attn_mask = torch.cat([attn_mask, pad_attn_mask], dim=0) + + # cu_seqlens + cu_out_tokens = out_tokens.masked_select(attn_mask != 0).unsqueeze(0) + seqlen = attn_mask.sum(dim=1).tolist() + cu_seqlens = torch.cumsum(torch.tensor([0]+seqlen), 0).to(device=out_tokens.device,dtype=torch.int32) + max_seq_len = max(seqlen) + return Tokens(out_tokens, cu_out_tokens, attn_mask, cu_seqlens, max_seq_len) + + \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/utils/__init__.py b/Step-Video-TI2V/stepvideo/utils/__init__.py new file mode 100755 index 0000000..a77226b --- /dev/null +++ b/Step-Video-TI2V/stepvideo/utils/__init__.py @@ -0,0 +1,2 @@ +from .utils import * +from .video_process import * \ No newline at end of file diff --git a/Step-Video-TI2V/stepvideo/utils/utils.py b/Step-Video-TI2V/stepvideo/utils/utils.py new file mode 100755 index 0000000..b0f1dc7 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/utils/utils.py @@ -0,0 +1,64 @@ +import numpy as np +import random +import torch +from functools import wraps +import torch.utils._device + + +def setup_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +class EmptyInitOnDevice(torch.overrides.TorchFunctionMode): + def __init__(self, device=None): + self.device = device + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if getattr(func, '__module__', None) == 'torch.nn.init': + if 'tensor' in kwargs: + return kwargs['tensor'] + else: + return args[0] + if self.device is not None and func in torch.utils._device._device_constructors() and kwargs.get('device') is None: + kwargs['device'] = self.device + return func(*args, **kwargs) + + +def with_empty_init(func): + @wraps(func) + def wrapper(*args, **kwargs): + with EmptyInitOnDevice('cpu'): + return func(*args, **kwargs) + return wrapper + + + + +def culens2mask( + cu_seqlens=None, + cu_seqlens_kv=None, + max_seqlen=None, + max_seqlen_kv=None, + is_causal=False +): + assert len(cu_seqlens) == len(cu_seqlens_kv); "q k v should have same bsz..." + bsz = len(cu_seqlens) - 1 + seqlens = cu_seqlens[1:]-cu_seqlens[:-1] + seqlens_kv = cu_seqlens_kv[1:]-cu_seqlens_kv[:-1] + + attn_mask = torch.zeros(bsz, max_seqlen, max_seqlen_kv, dtype=torch.bool) + for i, (seq_len, seq_len_kv) in enumerate(zip(seqlens, seqlens_kv)): + if is_causal: + attn_mask[i, :seq_len, :seq_len_kv] = torch.triu(torch.ones(seq_len, seq_len_kv), diagonal=1).bool() + else: + attn_mask[i, :seq_len, :seq_len_kv] = torch.ones([seq_len, seq_len_kv], dtype=torch.bool) + + return attn_mask + + diff --git a/Step-Video-TI2V/stepvideo/utils/video_process.py b/Step-Video-TI2V/stepvideo/utils/video_process.py new file mode 100644 index 0000000..7780317 --- /dev/null +++ b/Step-Video-TI2V/stepvideo/utils/video_process.py @@ -0,0 +1,50 @@ +import numpy as np +import datetime +import torch +import os +import imageio + + +class VideoProcessor: + def __init__(self, save_path: str='./results', name_suffix: str=''): + self.save_path = save_path + os.makedirs(self.save_path, exist_ok=True) + self.name_suffix = name_suffix + + def crop2standard540p(self, vid_array): + _, height, width, _ = vid_array.shape + height_center = height//2 + width_center = width//2 + if width_center>height_center: ## horizon mode + return vid_array[:, height_center-270:height_center+270, width_center-480:width_center+480] + elif width_center= 2**30) or (out_numel >= 2**30): + assert conv_layer.stride[0] == 1, "time split asks time stride = 1" + + B,T,H,W,C = x.shape + K = conv_layer.kernel_size[0] + + chunks = 4 + chunk_size = T // chunks + + if residual is None: + out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual, only_return_output=True) + else: + out_nhwc = residual + + assert B == 1 + outs = [] + for i in range(chunks): + if i == chunks-1: + xi = x[:1,chunk_size*i:] + out_nhwci = out_nhwc[:1,chunk_size*i:] + else: + xi = x[:1,chunk_size*i:chunk_size*(i+1)+K-1] + out_nhwci = out_nhwc[:1,chunk_size*i:chunk_size*(i+1)] + if residual is not None: + if i == chunks-1: + ri = residual[:1,chunk_size*i:] + else: + ri = residual[:1,chunk_size*i:chunk_size*(i+1)] + else: + ri = None + out_nhwci.copy_(base_conv3d(xi, conv_layer, channel_last=True, residual=ri)) + else: + out_nhwc = base_conv3d(x, conv_layer, channel_last=True, residual=residual) + return out_nhwc + + + +class Upsample2D(nn.Module): + def __init__(self, + channels, + use_conv=False, + use_conv_transpose=False, + out_channels=None): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + + if use_conv: + self.conv = nn.Conv2d(self.channels, self.out_channels, 3, padding=1) + else: + assert "Not Supported" + self.conv = nn.ConvTranspose2d(channels, self.out_channels, 4, 2, 1) + + def forward(self, x, output_size=None): + assert x.shape[-1] == self.channels + + if self.use_conv_transpose: + return self.conv(x) + + if output_size is None: + x = F.interpolate( + x.permute(0,3,1,2).to(memory_format=torch.channels_last), + scale_factor=2.0, mode='nearest').permute(0,2,3,1).contiguous() + else: + x = F.interpolate( + x.permute(0,3,1,2).to(memory_format=torch.channels_last), + size=output_size, mode='nearest').permute(0,2,3,1).contiguous() + + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + +class Downsample2D(nn.Module): + def __init__(self, channels, use_conv=False, out_channels=None, padding=1): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + + if use_conv: + self.conv = nn.Conv2d(self.channels, self.out_channels, 3, stride=stride, padding=padding) + else: + assert self.channels == self.out_channels + self.conv = nn.AvgPool2d(kernel_size=stride, stride=stride) + + def forward(self, x): + assert x.shape[-1] == self.channels + if self.use_conv and self.padding == 0: + pad = (0, 0, 0, 1, 0, 1) + x = F.pad(x, pad, mode="constant", value=0) + + assert x.shape[-1] == self.channels + # x = self.conv(x) + x = base_conv2d(x, self.conv, channel_last=True) + return x + + + +class CausalConv(nn.Module): + def __init__(self, + chan_in, + chan_out, + kernel_size, + **kwargs + ): + super().__init__() + + if isinstance(kernel_size, int): + kernel_size = kernel_size if isinstance(kernel_size, tuple) else ((kernel_size,) * 3) + time_kernel_size, height_kernel_size, width_kernel_size = kernel_size + + self.dilation = kwargs.pop('dilation', 1) + self.stride = kwargs.pop('stride', 1) + if isinstance(self.stride, int): + self.stride = (self.stride, 1, 1) + time_pad = self.dilation * (time_kernel_size - 1) + max((1 - self.stride[0]), 0) + height_pad = height_kernel_size // 2 + width_pad = width_kernel_size // 2 + self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) + self.time_uncausal_padding = (width_pad, width_pad, height_pad, height_pad, 0, 0) + + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None): + x = nn.functional.pad(x, + self.time_causal_padding if is_init else self.time_uncausal_padding) + + x = self.conv(x) + if residual is not None: + x.add_(residual) + return x + + +class ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert out_channels * factor**3 % in_channels == 0 + self.repeats = out_channels * factor**3 // in_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view(x.size(0), self.out_channels, self.factor, self.factor, self.factor, x.size(2), x.size(3), x.size(4)) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(x.size(0), self.out_channels, x.size(2)*self.factor, x.size(4)*self.factor, x.size(6)*self.factor) + x = x[:, :, self.factor - 1:, :, :] + return x + +class ConvPixelShuffleUpSampleLayer3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ): + super().__init__() + self.factor = factor + out_ratio = factor**3 + self.conv = CausalConv( + in_channels, + out_channels * out_ratio, + kernel_size=kernel_size + ) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_shuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_shuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + batch_size, channels, depth, height, width = x.size() + new_channels = channels // (factor ** 3) + new_depth = depth * factor + new_height = height * factor + new_width = width * factor + + x = x.view(batch_size, new_channels, factor, factor, factor, depth, height, width) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view(batch_size, new_channels, new_depth, new_height, new_width) + x = x[:, :, factor - 1:, :, :] + return x + +class ConvPixelUnshuffleDownSampleLayer3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + factor: int, + ): + super().__init__() + self.factor = factor + out_ratio = factor**3 + assert out_channels % out_ratio == 0 + self.conv = CausalConv( + in_channels, + out_channels // out_ratio, + kernel_size=kernel_size + ) + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + x = self.conv(x, is_init) + x = self.pixel_unshuffle_3d(x, self.factor) + return x + + @staticmethod + def pixel_unshuffle_3d(x: torch.Tensor, factor: int) -> torch.Tensor: + pad = (0, 0, 0, 0, factor-1, 0) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // factor, factor, H // factor, factor, W // factor, factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * factor**3, D // factor, H // factor, W // factor) + return x + +class PixelUnshuffleChannelAveragingDownSampleLayer3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert in_channels * factor**3 % out_channels == 0 + self.group_size = in_channels * factor**3 // out_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + pad = (0, 0, 0, 0, self.factor-1, 0) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // self.factor, self.factor, H // self.factor, self.factor, W // self.factor, self.factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor) + x = x.view(B, self.out_channels, self.group_size, D // self.factor, H // self.factor, W // self.factor) + x = x.mean(dim=2) + return x + + def __init__( + self, + in_channels: int, + out_channels: int, + factor: int, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor = factor + assert in_channels * factor**3 % out_channels == 0 + self.group_size = in_channels * factor**3 // out_channels + + def forward(self, x: torch.Tensor, is_init=True) -> torch.Tensor: + pad = (0, 0, 0, 0, self.factor-1, 0) # (left, right, top, bottom, front, back) + x = F.pad(x, pad) + B, C, D, H, W = x.shape + x = x.view(B, C, D // self.factor, self.factor, H // self.factor, self.factor, W // self.factor, self.factor) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * self.factor**3, D // self.factor, H // self.factor, W // self.factor) + x = x.view(B, self.out_channels, self.group_size, D // self.factor, H // self.factor, W // self.factor) + x = x.mean(dim=2) + return x + + + + +def base_group_norm_with_zero_pad(x, norm_layer, act_silu=True, pad_size=2): + out_shape = list(x.shape) + out_shape[1] += pad_size + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + out[:, pad_size:] = base_group_norm(x, norm_layer, act_silu=act_silu, channel_last=True) + out[:, :pad_size] = 0 + return out + + +class CausalConvChannelLast(CausalConv): + def __init__(self, + chan_in, + chan_out, + kernel_size, + **kwargs + ): + super().__init__( + chan_in, chan_out, kernel_size, **kwargs) + + self.time_causal_padding = (0, 0) + self.time_causal_padding + self.time_uncausal_padding = (0, 0) + self.time_uncausal_padding + + def forward(self, x, is_init=True, residual=None): + if self.is_first_run: + self.is_first_run = False + # self.conv.weight = nn.Parameter(self.conv.weight.permute(0,2,3,4,1).contiguous()) + + x = nn.functional.pad(x, + self.time_causal_padding if is_init else self.time_uncausal_padding) + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + +class CausalConvAfterNorm(CausalConv): + def __init__(self, + chan_in, + chan_out, + kernel_size, + **kwargs + ): + super().__init__( + chan_in, chan_out, kernel_size, **kwargs) + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, padding=(0, 1, 1), **kwargs) + else: + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride=self.stride, dilation=self.dilation, **kwargs) + self.is_first_run = True + + def forward(self, x, is_init=True, residual=None): + if self.is_first_run: + self.is_first_run = False + + if self.time_causal_padding == (1, 1, 1, 1, 2, 0): + pass + else: + x = nn.functional.pad(x, self.time_causal_padding).contiguous() + + x = base_conv3d_channel_last(x, self.conv, residual=residual) + return x + +class AttnBlock(nn.Module): + def __init__(self, + in_channels + ): + super().__init__() + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) + + def attention(self, x, is_init=True): + x = base_group_norm(x, self.norm, act_silu=False, channel_last=True) + q = self.q(x, is_init) + k = self.k(x, is_init) + v = self.v(x, is_init) + + b, t, h, w, c = q.shape + q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v)) + x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) + x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) + + return x + + def forward(self, x): + x = x.permute(0,2,3,4,1).contiguous() + h = self.attention(x) + x = self.proj_out(h, residual=x) + x = x.permute(0,4,1,2,3) + return x + +class Resnet3DBlock(nn.Module): + def __init__(self, + in_channels, + out_channels=None, + temb_channels=512, + conv_shortcut=False, + ): + super().__init__() + + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels) + self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) + if temb_channels > 0: + self.temb_proj = nn.Linear(temb_channels, out_channels) + + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels) + self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3) + + assert conv_shortcut is False + self.use_conv_shortcut = conv_shortcut + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) + else: + self.nin_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=1) + + def forward(self, x, temb=None, is_init=True): + x = x.permute(0,2,3,4,1).contiguous() + + h = base_group_norm_with_zero_pad(x, self.norm1, act_silu=True, pad_size=2) + h = self.conv1(h) + if temb is not None: + h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] + + x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x + + h = base_group_norm_with_zero_pad(h, self.norm2, act_silu=True, pad_size=2) + x = self.conv2(h, residual=x) + + x = x.permute(0,4,1,2,3) + return x + + +class Downsample3D(nn.Module): + def __init__(self, + in_channels, + with_conv, + stride + ): + super().__init__() + + self.with_conv = with_conv + if with_conv: + self.conv = CausalConv(in_channels, in_channels, kernel_size=3, stride=stride) + + def forward(self, x, is_init=True): + if self.with_conv: + x = self.conv(x, is_init) + else: + x = nn.functional.avg_pool3d(x, kernel_size=2, stride=2) + return x + +class VideoEncoder(nn.Module): + def __init__(self, + ch=32, + ch_mult=(4, 8, 16, 16), + num_res_blocks=2, + in_channels=3, + z_channels=16, + double_z=True, + down_sampling_layer=[1, 2], + resamp_with_conv=True, + version=1, + ): + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = CausalConv(in_channels, ch, kernel_size=3) + self.down_sampling_layer = down_sampling_layer + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + if i_level in self.down_sampling_layer: + down.downsample = Downsample3D(block_in, resamp_with_conv, stride=(2, 2, 2)) + else: + down.downsample = Downsample2D(block_in, resamp_with_conv, padding=0) #DIFF + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.version = version + if version == 2: + channels = 4 * z_channels * 2 ** 3 + self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D(block_in, channels, kernel_size=3, factor=2) + self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D(block_in, channels, 2) + self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D(channels, 2 * z_channels if double_z else z_channels, 1) + self.conv_out = CausalConvChannelLast(channels, 2 * z_channels if double_z else z_channels, kernel_size=3) + else: + self.conv_out = CausalConvAfterNorm(block_in, 2 * z_channels if double_z else z_channels, kernel_size=3) + + @torch.inference_mode() + def forward(self, x, video_frame_num, is_init=True): + # timestep embedding + temb = None + + t = video_frame_num + + # downsampling + h = self.conv_in(x, is_init) + + # make it real channel last, but behave like normal layout + h = h.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3) + + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](h, temb, is_init) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + + if i_level != self.num_resolutions - 1: + if isinstance(self.down[i_level].downsample, Downsample2D): + _, _, t, _, _ = h.shape + h = rearrange(h, "b c t h w -> (b t) h w c", t=t) + h = self.down[i_level].downsample(h) + h = rearrange(h, "(b t) h w c -> b c t h w", t=t) + else: + h = self.down[i_level].downsample(h, is_init) + + h = self.mid.block_1(h, temb, is_init) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb, is_init) + + h = h.permute(0,2,3,4,1).contiguous() # b c l h w -> b l h w c + if self.version == 2: + h = base_group_norm(h, self.norm_out, act_silu=True, channel_last=True) + h = h.permute(0,4,1,2,3).contiguous() + shortcut = self.shortcut_pathify(h, is_init) + h = self.conv_patchify(h, is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_out(h, is_init).permute(0,2,3,4,1) + h = self.conv_out(h.permute(0,2,3,4,1).contiguous(), is_init) + h = h.add_(shortcut) + else: + h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) + h = self.conv_out(h, is_init) + h = h.permute(0,4,1,2,3) # b l h w c -> b c l h w + + h = rearrange(h, "b c t h w -> b t c h w") + return h + + +class Res3DBlockUpsample(nn.Module): + def __init__(self, + input_filters, + num_filters, + down_sampling_stride, + down_sampling=False + ): + super().__init__() + + self.input_filters = input_filters + self.num_filters = num_filters + + self.act_ = nn.SiLU(inplace=True) + + self.conv1 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3]) + self.norm1 = nn.GroupNorm(32, num_filters) + + self.conv2 = CausalConvChannelLast(num_filters, num_filters, kernel_size=[3, 3, 3]) + self.norm2 = nn.GroupNorm(32, num_filters) + + self.down_sampling = down_sampling + if down_sampling: + self.down_sampling_stride = down_sampling_stride + else: + self.down_sampling_stride = [1, 1, 1] + + if num_filters != input_filters or down_sampling: + self.conv3 = CausalConvChannelLast(input_filters, num_filters, kernel_size=[1, 1, 1], stride=self.down_sampling_stride) + self.norm3 = nn.GroupNorm(32, num_filters) + + def forward(self, x, is_init=False): + x = x.permute(0,2,3,4,1).contiguous() + + residual = x + + h = self.conv1(x, is_init) + h = base_group_norm(h, self.norm1, act_silu=True, channel_last=True) + + h = self.conv2(h, is_init) + h = base_group_norm(h, self.norm2, act_silu=False, channel_last=True) + + if self.down_sampling or self.num_filters != self.input_filters: + x = self.conv3(x, is_init) + x = base_group_norm(x, self.norm3, act_silu=False, channel_last=True) + + h.add_(x) + h = self.act_(h) + if residual is not None: + h.add_(residual) + + h = h.permute(0,4,1,2,3) + return h + +class Upsample3D(nn.Module): + def __init__(self, + in_channels, + scale_factor=2 + ): + super().__init__() + + self.scale_factor = scale_factor + self.conv3d = Res3DBlockUpsample(input_filters=in_channels, + num_filters=in_channels, + down_sampling_stride=(1, 1, 1), + down_sampling=False) + + def forward(self, x, is_init=True, is_split=True): + b, c, t, h, w = x.shape + + # x = x.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3).to(memory_format=torch.channels_last_3d) + if is_split: + split_size = c // 8 + x_slices = torch.split(x, split_size, dim=1) + x = [nn.functional.interpolate(x, scale_factor=self.scale_factor) for x in x_slices] + x = torch.cat(x, dim=1) + else: + x = nn.functional.interpolate(x, scale_factor=self.scale_factor) + + x = self.conv3d(x, is_init) + return x + +class VideoDecoder(nn.Module): + def __init__(self, + ch=128, + z_channels=16, + out_channels=3, + ch_mult=(1, 2, 4, 4), + num_res_blocks=2, + temporal_up_layers=[2, 3], + temporal_downsample=4, + resamp_with_conv=True, + version=1, + ): + super().__init__() + + temb_ch = 0 + + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.temporal_downsample = temporal_downsample + + block_in = ch * ch_mult[self.num_resolutions - 1] + self.version = version + if version == 2: + channels = 4 * z_channels * 2 ** 3 + self.conv_in = CausalConv(z_channels, channels, kernel_size=3) + self.shortcut_in = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(z_channels, channels, 1) + self.conv_unpatchify = ConvPixelShuffleUpSampleLayer3D(channels, block_in, kernel_size=3, factor=2) + self.shortcut_unpathify = ChannelDuplicatingPixelUnshuffleUpSampleLayer3D(channels, block_in, 2) + else: + self.conv_in = CausalConv(z_channels, block_in, kernel_size=3) + + # middle + self.mid = nn.Module() + self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) + + # upsampling + self.up_id = len(temporal_up_layers) + self.video_frame_num = 1 + self.cur_video_frame_num = self.video_frame_num // 2 ** self.up_id + 1 + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + if i_level in temporal_up_layers: + up.upsample = Upsample3D(block_in) + self.cur_video_frame_num = self.cur_video_frame_num * 2 + else: + up.upsample = Upsample2D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) + self.conv_out = CausalConvAfterNorm(block_in, out_channels, kernel_size=3) + + @torch.inference_mode() + def forward(self, z, is_init=True): + z = rearrange(z, "b t c h w -> b c t h w") + + h = self.conv_in(z, is_init=is_init) + if self.version == 2: + shortcut = self.shortcut_in(z, is_init=is_init) + h = h.add_(shortcut) + shortcut = self.shortcut_unpathify(h, is_init=is_init) + h = self.conv_unpatchify(h, is_init=is_init) + h = h.add_(shortcut) + + temb = None + + h = h.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3) + h = self.mid.block_1(h, temb, is_init=is_init) + h = self.mid.attn_1(h) + h = h.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3) + h = self.mid.block_2(h, temb, is_init=is_init) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = h.permute(0,2,3,4,1).contiguous().permute(0,4,1,2,3) + h = self.up[i_level].block[i_block](h, temb, is_init=is_init) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + if isinstance(self.up[i_level].upsample, Upsample2D): + B = h.size(0) + h = h.permute(0,2,3,4,1).flatten(0,1) + h = self.up[i_level].upsample(h) + h = h.unflatten(0, (B, -1)).permute(0,4,1,2,3) + else: + h = self.up[i_level].upsample(h, is_init=is_init) + + # end + h = h.permute(0,2,3,4,1) # b c l h w -> b l h w c + h = base_group_norm_with_zero_pad(h, self.norm_out, act_silu=True, pad_size=2) + h = self.conv_out(h) + h = h.permute(0,4,1,2,3) + + if is_init: + h = h[:, :, (self.temporal_downsample - 1):] + return h + + + +def rms_norm(input, normalized_shape, eps=1e-6): + dtype = input.dtype + input = input.to(torch.float32) + variance = input.pow(2).flatten(-len(normalized_shape)).mean(-1)[(...,) + (None,) * len(normalized_shape)] + input = input * torch.rsqrt(variance + eps) + return input.to(dtype) + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters, deterministic=False, rms_norm_mean=False, only_return_mean=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=-3) #N,[X],C,H,W + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + self.deterministic = deterministic + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, + device=self.parameters.device, + dtype=self.parameters.dtype) + if rms_norm_mean: + self.mean = rms_norm(self.mean, self.mean.size()[1:]) + self.only_return_mean = only_return_mean + + def sample(self, generator=None): + # make sure sample is on the same device + # as the parameters and has same dtype + sample = torch.randn( + self.mean.shape, generator=generator, device=self.parameters.device) + sample = sample.to(dtype=self.parameters.dtype) + x = self.mean + self.std * sample + if self.only_return_mean: + return self.mean + else: + return x + +class AutoencoderKL(nn.Module): + @with_empty_init + def __init__(self, + in_channels=3, + out_channels=3, + z_channels=16, + num_res_blocks=2, + model_path=None, + weight_dict={}, + world_size=1, + version=1, + ): + super().__init__() + + self.frame_len = 17 + self.latent_len = 3 if version == 2 else 5 + + base_group_norm.spatial = True if version == 2 else False + + self.encoder = VideoEncoder( + in_channels=in_channels, + z_channels=z_channels, + num_res_blocks=num_res_blocks, + version=version, + ) + + self.decoder = VideoDecoder( + z_channels=z_channels, + out_channels=out_channels, + num_res_blocks=num_res_blocks, + version=version, + ) + + if model_path is not None: + weight_dict = self.init_from_ckpt(model_path) + if len(weight_dict) != 0: + self.load_from_dict(weight_dict) + self.convert_channel_last() + + self.world_size = world_size + + def init_from_ckpt(self, model_path): + from safetensors import safe_open + p = {} + with safe_open(model_path, framework="pt", device="cpu") as f: + for k in f.keys(): + tensor = f.get_tensor(k) + if k.startswith("decoder.conv_out."): + k = k.replace("decoder.conv_out.", "decoder.conv_out.conv.") + p[k] = tensor + return p + + def load_from_dict(self, p): + self.load_state_dict(p) + + def convert_channel_last(self): + #Conv2d NCHW->NHWC + pass + + def naive_encode(self, x, is_init_image=True): + b, l, c, h, w = x.size() + x = rearrange(x, 'b l c h w -> b c l h w').contiguous() + z = self.encoder(x, l, True) # ไธ‹้‡‡ๆ ท[1, 4, 8, 16, 16] + return z + + @torch.inference_mode() + def encode(self, x): + # b (nc cf) c h w -> (b nc) cf c h w -> encode -> (b nc) cf c h w -> b (nc cf) c h w + chunks = list(x.split(self.frame_len, dim=1)) + for i in range(len(chunks)): + chunks[i] = self.naive_encode(chunks[i], True) + z = torch.cat(chunks, dim=1) + + posterior = DiagonalGaussianDistribution(z) + return posterior.sample() + + def decode_naive(self, z, is_init=True): + z = z.to(next(self.decoder.parameters()).dtype) + dec = self.decoder(z, is_init) + return dec + + @torch.inference_mode() + def decode(self, z): + # b (nc cf) c h w -> (b nc) cf c h w -> decode -> (b nc) c cf h w -> b (nc cf) c h w + chunks = list(z.split(self.latent_len, dim=1)) + + if self.world_size > 1: + chunks_total_num = len(chunks) + max_num_per_rank = (chunks_total_num + self.world_size - 1) // self.world_size + rank = torch.distributed.get_rank() + chunks_ = chunks[max_num_per_rank * rank : max_num_per_rank * (rank + 1)] + if len(chunks_) < max_num_per_rank: + chunks_.extend(chunks[:max_num_per_rank-len(chunks_)]) + chunks = chunks_ + + for i in range(len(chunks)): + chunks[i] = self.decode_naive(chunks[i], True).permute(0,2,1,3,4) + x = torch.cat(chunks, dim=1) + + if self.world_size > 1: + x_ = torch.empty([x.size(0), (self.world_size * max_num_per_rank) * self.frame_len, *x.shape[2:]], dtype=x.dtype, device=x.device) + torch.distributed.all_gather_into_tensor(x_, x) + x = x_[:, : chunks_total_num * self.frame_len] + + x = self.mix(x) + return x + + def mix(self, x): + remain_scale = 0.6 + mix_scale = 1. - remain_scale + front = slice(self.frame_len - 1, x.size(1) - 1, self.frame_len) + back = slice(self.frame_len, x.size(1), self.frame_len) + x[:, back] = x[:, back] * remain_scale + x[:, front] * mix_scale + x[:, front] = x[:, front] * remain_scale + x[:, back] * mix_scale + return x