diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/README.md b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..3c3ca6b8d4d8793913f6d2d49dbfb0c5e6638c18 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/README.md @@ -0,0 +1,4 @@ +/.--- +license: apache-2.0 +--- + diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/inference.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..21f85bab8be3e3042140f81f02bbd18abf0231d5 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/inference.py @@ -0,0 +1,144 @@ +import argparse +import logging +import os + +import imageio +import torch +import torch_npu +from diffusers.schedulers import EulerAncestralDiscreteScheduler +from transformers import T5Tokenizer, MT5EncoderModel + +from opensora.models.causalvideovae import ae_stride_config, CausalVAEModelWrapper +from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V +from opensora.sample.pipeline_opensora_sp import OpenSoraPipeline +from utils.parallel_mgr import init_parallel_env, finalize_parallel_env, get_sequence_parallel_rank + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +def load_t2v_checkpoint(model_path): + logger.info(f'load_t2v_checkpoint, {model_path}') + transformer_model = OpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir, + low_cpu_mem_usage=False, device_map=None, + torch_dtype=weight_dtype).to("npu") + transformer_model.eval() + pipeline = OpenSoraPipeline(vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + transformer=transformer_model).to("npu") + if args.algorithm == "dit_cache": + from mindiesd.layers.cache_mgr import CacheManager, DitCacheConfig + config = DitCacheConfig(step_start=20, step_interval=2, block_start=7, num_blocks=21) + cache = CacheManager(config) + pipeline.transformer.cache = cache + return pipeline + + +def run_model_and_save_images(pipeline): + if not isinstance(args.text_prompt, list): + args.text_prompt = [args.text_prompt] + if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'): + text_prompt = open(args.text_prompt[0], 'r').readlines() + args.text_prompt = [i.strip() for i in text_prompt] + + positive_prompt = """ + (masterpiece), (best quality), (ultra-detailed), (unwatermarked), + {}. + emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, + sharp focus, high budget, cinemascope, moody, epic, gorgeous + """ + + negative_prompt = """ + nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, + low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry. + """ + kwargs = {} + if args.algorithm == "sampling_optimize": + kwargs["sampling_optimize"] = True + + for index, prompt in enumerate(args.text_prompt): + videos = pipeline(positive_prompt.format(prompt), + negative_prompt=negative_prompt, + num_frames=args.num_frames, + height=args.height, + width=args.width, + num_inference_steps=args.num_sampling_steps, + guidance_scale=args.guidance_scale, + num_images_per_prompt=1, + mask_feature=True, + max_sequence_length=args.max_sequence_length, + **kwargs + ).images + logger.info(videos.shape) + + if get_sequence_parallel_rank() <= 0: + try: + imageio.mimwrite( + os.path.join( + args.save_img_path, + f'EulerAncestralDiscrete_{index}_gs{args.guidance_scale}_s{args.num_sampling_steps}.mp4' + ), videos[0], + fps=args.fps, quality=6, codec='libx264', + output_params=['-threads', '20']) # highest quality is 10, lowest is 0 + except: + logger.error('Error when saving {}'.format(prompt)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0') + parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256']) + parser.add_argument("--num_frames", type=int, default=93) + parser.add_argument("--height", type=int, default=720) + parser.add_argument("--width", type=int, default=1280) + parser.add_argument('--dtype', type=str, default='bf16', help='Data type used in inference') + parser.add_argument("--cache_dir", type=str, default='./cache_dir') + parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8') + parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8') + parser.add_argument("--text_encoder_name", type=str, default='google/mt5-xxl') + parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v") + parser.add_argument("--guidance_scale", type=float, default=7.5) + parser.add_argument("--num_sampling_steps", type=int, default=50) + parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--max_sequence_length", type=int, default=512) + parser.add_argument("--text_prompt", nargs='+') + parser.add_argument("--tile_overlap_factor", type=float, default=0.25) + parser.add_argument("--algorithm", type=str, default=None, choices=[None, 'dit_cache', 'sampling_optimize']) + args = parser.parse_args() + + if args.dtype not in ['bf16', 'fp16']: + logger.error("Not supported.") + weight_dtype = torch.bfloat16 if args.dtype == 'bf16' else torch.float16 + + local_rank = int(os.getenv('RANK', 0)) + world_size = int(os.getenv('WORLD_SIZE', 1)) + if world_size > 0: + init_parallel_env(enable_sequence_parallelism=True) + + vae = CausalVAEModelWrapper(args.ae_path, dtype=torch.float16).to("npu") + vae.vae.enable_tiling() + vae.vae.tile_overlap_factor = args.tile_overlap_factor + vae.vae.tile_sample_min_size = 256 + vae.vae.tile_latent_min_size = 32 + vae.vae.tile_sample_min_size_t = 29 + vae.vae.tile_latent_min_size_t = 8 + vae.vae_scale_factor = ae_stride_config[args.ae] + vae.eval() + + text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir, + low_cpu_mem_usage=True, torch_dtype=weight_dtype).to("npu") + tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir) + text_encoder.eval() + + scheduler = EulerAncestralDiscreteScheduler() + + if not os.path.exists(args.save_img_path): + os.makedirs(args.save_img_path, exist_ok=True) + + pipeline = load_t2v_checkpoint(args.model_path) + logger.info('load model') + run_model_and_save_images(pipeline) + if world_size > 0: + finalize_parallel_env() diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3eca32c739ed7eaf060421489d9dac7568dffe67 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/__init__.py @@ -0,0 +1 @@ +from .causalvideovae.model import CausalVAEModelWrapper \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bc1357393ee13227aa72acb9111811ee3533826a --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/__init__.py @@ -0,0 +1,32 @@ +from torchvision.transforms import Lambda +from .model.causal_vae import CausalVAEModelWrapper + +ae_stride_config = { + 'CausalVAEModel_D4_2x8x8': [2, 8, 8], + 'CausalVAEModel_D8_2x8x8': [2, 8, 8], + 'CausalVAEModel_D4_4x8x8': [4, 8, 8], + 'CausalVAEModel_D8_4x8x8': [4, 8, 8], +} + + +ae_channel_config = { + 'CausalVAEModel_D4_2x8x8': 4, + 'CausalVAEModel_D8_2x8x8': 8, + 'CausalVAEModel_D4_4x8x8': 4, + 'CausalVAEModel_D8_4x8x8': 8, +} + + +ae_denorm = { + 'CausalVAEModel_D4_2x8x8': lambda x: (x + 1.) / 2., + 'CausalVAEModel_D8_2x8x8': lambda x: (x + 1.) / 2., + 'CausalVAEModel_D4_4x8x8': lambda x: (x + 1.) / 2., + 'CausalVAEModel_D8_4x8x8': lambda x: (x + 1.) / 2., +} + +ae_norm = { + 'CausalVAEModel_D4_2x8x8': Lambda(lambda x: 2. * x - 1.), + 'CausalVAEModel_D8_2x8x8': Lambda(lambda x: 2. * x - 1.), + 'CausalVAEModel_D4_4x8x8': Lambda(lambda x: 2. * x - 1.), + 'CausalVAEModel_D8_4x8x8': Lambda(lambda x: 2. * x - 1.), +} \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..84918b27c2c51943b6fe68caf9ef83518419396d --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/__init__.py @@ -0,0 +1,3 @@ +from .causal_vae import ( + CausalVAEModel, CausalVAEModelWrapper +) \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/causal_vae/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/causal_vae/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..c1aeb800d75346ac7394829e3983eb8a0921db66 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/causal_vae/__init__.py @@ -0,0 +1,29 @@ +from .modeling_causalvae import CausalVAEModel + +from einops import rearrange +from torch import nn + +class CausalVAEModelWrapper(nn.Module): + def __init__(self, model_path, subfolder=None, cache_dir=None, use_ema=False, **kwargs): + super(CausalVAEModelWrapper, self).__init__() + # if os.path.exists(ckpt): + # self.vae = CausalVAEModel.load_from_checkpoint(ckpt) + self.vae = CausalVAEModel.from_pretrained(model_path, subfolder=subfolder, cache_dir=cache_dir, **kwargs) + if use_ema: + self.vae.init_from_ema(model_path) + self.vae = self.vae.ema + def encode(self, x): # b c t h w + # x = self.vae.encode(x).sample() + x = self.vae.encode(x).sample().mul_(0.18215) + return x + def decode(self, x): + # x = self.vae.decode(x) + x = self.vae.decode(x / 0.18215) + x = rearrange(x, 'b c t h w -> b t c h w').contiguous() + return x + + def dtype(self): + return self.vae.dtype + # + # def device(self): + # return self.vae.device \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py new file mode 100755 index 0000000000000000000000000000000000000000..299cc071ae192d1504bc1bc66fe38f829ffdd110 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py @@ -0,0 +1,600 @@ +import os +from copy import deepcopy +from typing import Tuple + +import torch +import torch_npu +import torch.nn as nn +from diffusers.configuration_utils import register_to_config + +from ..modules import Normalize +from ..modules.ops import nonlinearity +from ..modeling_videobase import VideoBaseAE +from ..utils.distrib_utils import DiagonalGaussianDistribution +from ..utils.module_utils import resolve_str_to_obj, Module + + +class Encoder(nn.Module): + def __init__( + self, + z_channels: int, + hidden_size: int, + hidden_size_mult: Tuple[int] = (1, 2, 4, 4), + attn_resolutions: Tuple[int] = (16,), + conv_in: Module = "Conv2d", + conv_out: Module = "CasualConv3d", + attention: Module = "AttnBlock", + resnet_blocks: Tuple[Module] = ( + "ResnetBlock2D", + "ResnetBlock2D", + "ResnetBlock2D", + "ResnetBlock3D", + ), + spatial_downsample: Tuple[Module] = ( + "Downsample", + "Downsample", + "Downsample", + "", + ), + temporal_downsample: Tuple[Module] = ("", "", "TimeDownsampleRes2x", ""), + mid_resnet: Module = "ResnetBlock3D", + dropout: float = 0.0, + resolution: int = 256, + num_res_blocks: int = 2, + double_z: bool = True, + ) -> None: + super().__init__() + assert len(resnet_blocks) == len(hidden_size_mult), print( + hidden_size_mult, resnet_blocks + ) + # ---- Config ---- + self.num_resolutions = len(hidden_size_mult) + self.resolution = resolution + self.num_res_blocks = num_res_blocks + + # ---- In ---- + self.conv_in = resolve_str_to_obj(conv_in)( + 3, hidden_size, kernel_size=3, stride=1, padding=1 + ) + + # ---- Downsample ---- + curr_res = resolution + in_ch_mult = (1,) + tuple(hidden_size_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = hidden_size * in_ch_mult[i_level] + block_out = hidden_size * hidden_size_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + resolve_str_to_obj(resnet_blocks[i_level])( + in_channels=block_in, + out_channels=block_out, + dropout=dropout, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(resolve_str_to_obj(attention)(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if spatial_downsample[i_level]: + down.downsample = resolve_str_to_obj(spatial_downsample[i_level])( + block_in, block_in + ) + curr_res = curr_res // 2 + if temporal_downsample[i_level]: + down.time_downsample = resolve_str_to_obj(temporal_downsample[i_level])( + block_in, block_in + ) + self.down.append(down) + + # ---- Mid ---- + self.mid = nn.Module() + self.mid.block_1 = resolve_str_to_obj(mid_resnet)( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + self.mid.attn_1 = resolve_str_to_obj(attention)(block_in) + self.mid.block_2 = resolve_str_to_obj(mid_resnet)( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + # ---- Out ---- + self.norm_out = Normalize(block_in) + self.conv_out = resolve_str_to_obj(conv_out)( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def forward(self, x): + hs = [self.conv_in(x)] + 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](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if hasattr(self.down[i_level], "downsample"): + hs.append(self.down[i_level].downsample(hs[-1])) + if hasattr(self.down[i_level], "time_downsample"): + hs_down = self.down[i_level].time_downsample(hs[-1]) + hs.append(hs_down) + + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + z_channels: int, + hidden_size: int, + hidden_size_mult: Tuple[int] = (1, 2, 4, 4), + attn_resolutions: Tuple[int] = (16,), + conv_in: Module = "Conv2d", + conv_out: Module = "CasualConv3d", + attention: Module = "AttnBlock", + resnet_blocks: Tuple[Module] = ( + "ResnetBlock3D", + "ResnetBlock3D", + "ResnetBlock3D", + "ResnetBlock3D", + ), + spatial_upsample: Tuple[Module] = ( + "", + "SpatialUpsample2x", + "SpatialUpsample2x", + "SpatialUpsample2x", + ), + temporal_upsample: Tuple[Module] = ("", "", "", "TimeUpsampleRes2x"), + mid_resnet: Module = "ResnetBlock3D", + dropout: float = 0.0, + resolution: int = 256, + num_res_blocks: int = 2, + ): + super().__init__() + # ---- Config ---- + self.num_resolutions = len(hidden_size_mult) + self.resolution = resolution + self.num_res_blocks = num_res_blocks + + # ---- In ---- + block_in = hidden_size * hidden_size_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.conv_in = resolve_str_to_obj(conv_in)( + z_channels, block_in, kernel_size=3, padding=1 + ) + + # ---- Mid ---- + self.mid = nn.Module() + self.mid.block_1 = resolve_str_to_obj(mid_resnet)( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + self.mid.attn_1 = resolve_str_to_obj(attention)(block_in) + self.mid.block_2 = resolve_str_to_obj(mid_resnet)( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + + # ---- Upsample ---- + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = hidden_size * hidden_size_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + resolve_str_to_obj(resnet_blocks[i_level])( + in_channels=block_in, + out_channels=block_out, + dropout=dropout, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(resolve_str_to_obj(attention)(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if spatial_upsample[i_level]: + up.upsample = resolve_str_to_obj(spatial_upsample[i_level])( + block_in, block_in + ) + curr_res = curr_res * 2 + if temporal_upsample[i_level]: + up.time_upsample = resolve_str_to_obj(temporal_upsample[i_level])( + block_in, block_in + ) + self.up.insert(0, up) + + # ---- Out ---- + self.norm_out = Normalize(block_in) + self.conv_out = resolve_str_to_obj(conv_out)( + block_in, 3, kernel_size=3, padding=1 + ) + + def forward(self, z): + h = self.conv_in(z) + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if hasattr(self.up[i_level], "upsample"): + h = self.up[i_level].upsample(h) + if hasattr(self.up[i_level], "time_upsample"): + h = self.up[i_level].time_upsample(h) + + h = self.norm_out(h) + h = nonlinearity(h) + h_dtype = h.dtype + h = self.conv_out(h) + return h + + +class CausalVAEModel(VideoBaseAE): + + @register_to_config + def __init__( + self, + hidden_size: int = 128, + z_channels: int = 4, + hidden_size_mult: Tuple[int] = (1, 2, 4, 4), + attn_resolutions: Tuple[int] = [], + dropout: float = 0.0, + resolution: int = 256, + double_z: bool = True, + embed_dim: int = 4, + num_res_blocks: int = 2, + q_conv: str = "CausalConv3d", + encoder_conv_in: Module = "CausalConv3d", + encoder_conv_out: Module = "CausalConv3d", + encoder_attention: Module = "AttnBlock3D", + encoder_resnet_blocks: Tuple[Module] = ( + "ResnetBlock3D", + "ResnetBlock3D", + "ResnetBlock3D", + "ResnetBlock3D", + ), + encoder_spatial_downsample: Tuple[Module] = ( + "SpatialDownsample2x", + "SpatialDownsample2x", + "SpatialDownsample2x", + "", + ), + encoder_temporal_downsample: Tuple[Module] = ( + "", + "TimeDownsample2x", + "TimeDownsample2x", + "", + ), + encoder_mid_resnet: Module = "ResnetBlock3D", + decoder_conv_in: Module = "CausalConv3d", + decoder_conv_out: Module = "CausalConv3d", + decoder_attention: Module = "AttnBlock3D", + decoder_resnet_blocks: Tuple[Module] = ( + "ResnetBlock3D", + "ResnetBlock3D", + "ResnetBlock3D", + "ResnetBlock3D", + ), + decoder_spatial_upsample: Tuple[Module] = ( + "", + "SpatialUpsample2x", + "SpatialUpsample2x", + "SpatialUpsample2x", + ), + decoder_temporal_upsample: Tuple[Module] = ("", "", "TimeUpsample2x", "TimeUpsample2x"), + decoder_mid_resnet: Module = "ResnetBlock3D", + use_quant_layer: bool = True + ) -> None: + super().__init__() + + self.tile_sample_min_size = 256 + self.tile_sample_min_size_t = 33 + self.tile_latent_min_size = int(self.tile_sample_min_size / (2 ** (len(hidden_size_mult) - 1))) + + # t_down_ratio = [i for i in encoder_temporal_downsample if len(i) > 0] + # self.tile_latent_min_size_t = int((self.tile_sample_min_size_t-1) / (2 ** len(t_down_ratio))) + 1 + self.tile_latent_min_size_t = 16 + self.tile_overlap_factor = 0.125 + self.use_tiling = False + + self.use_quant_layer = use_quant_layer + + self.encoder = Encoder( + z_channels=z_channels, + hidden_size=hidden_size, + hidden_size_mult=hidden_size_mult, + attn_resolutions=attn_resolutions, + conv_in=encoder_conv_in, + conv_out=encoder_conv_out, + attention=encoder_attention, + resnet_blocks=encoder_resnet_blocks, + spatial_downsample=encoder_spatial_downsample, + temporal_downsample=encoder_temporal_downsample, + mid_resnet=encoder_mid_resnet, + dropout=dropout, + resolution=resolution, + num_res_blocks=num_res_blocks, + double_z=double_z, + ) + + self.decoder = Decoder( + z_channels=z_channels, + hidden_size=hidden_size, + hidden_size_mult=hidden_size_mult, + attn_resolutions=attn_resolutions, + conv_in=decoder_conv_in, + conv_out=decoder_conv_out, + attention=decoder_attention, + resnet_blocks=decoder_resnet_blocks, + spatial_upsample=decoder_spatial_upsample, + temporal_upsample=decoder_temporal_upsample, + mid_resnet=decoder_mid_resnet, + dropout=dropout, + resolution=resolution, + num_res_blocks=num_res_blocks, + ) + if self.use_quant_layer: + quant_conv_cls = resolve_str_to_obj(q_conv) + self.quant_conv = quant_conv_cls(2 * z_channels, 2 * embed_dim, 1) + self.post_quant_conv = quant_conv_cls(embed_dim, z_channels, 1) + + def get_encoder(self): + if self.use_quant_layer: + return [self.quant_conv, self.encoder] + return [self.encoder] + + def get_decoder(self): + if self.use_quant_layer: + return [self.post_quant_conv, self.decoder] + return [self.decoder] + + def encode(self, x): + if self.use_tiling and ( + x.shape[-1] > self.tile_sample_min_size + or x.shape[-2] > self.tile_sample_min_size + or x.shape[-3] > self.tile_sample_min_size_t + ): + return self.tiled_encode(x) + h = self.encoder(x) + if self.use_quant_layer: + h = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(h) + return posterior + + def decode(self, z): + if self.use_tiling and ( + z.shape[-1] > self.tile_latent_min_size + or z.shape[-2] > self.tile_latent_min_size + or z.shape[-3] > self.tile_latent_min_size_t + ): + return self.tiled_decode(z) + if self.use_quant_layer: + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward(self, input, sample_posterior=True): + posterior = self.encode(input) + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z) + return dec, posterior + + def on_train_start(self): + self.ema = deepcopy(self) if self.save_ema == True else None + + def get_last_layer(self): + if hasattr(self.decoder.conv_out, "conv"): + return self.decoder.conv_out.conv.weight + else: + return self.decoder.conv_out.weight + + def blend_v( + self, a: torch.Tensor, b: torch.Tensor, blend_extent: int + ) -> torch.Tensor: + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( + 1 - y / blend_extent + ) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h( + self, a: torch.Tensor, b: torch.Tensor, blend_extent: int + ) -> torch.Tensor: + blend_extent = min(a.shape[4], b.shape[4], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( + 1 - x / blend_extent + ) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def tiled_encode(self, x): + t = x.shape[2] + t_chunk_idx = [i for i in range(0, t, self.tile_sample_min_size_t - 1)] + if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0: + t_chunk_start_end = [[0, t]] + else: + t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i + 1] + 1] for i in range(len(t_chunk_idx) - 1)] + if t_chunk_start_end[-1][-1] > t: + t_chunk_start_end[-1][-1] = t + elif t_chunk_start_end[-1][-1] < t: + last_start_end = [t_chunk_idx[-1], t] + t_chunk_start_end.append(last_start_end) + moments = [] + for idx, (start, end) in enumerate(t_chunk_start_end): + chunk_x = x[:, :, start: end] + if idx != 0: + moment = self.tiled_encode2d(chunk_x, return_moments=True)[:, :, 1:] + else: + moment = self.tiled_encode2d(chunk_x, return_moments=True) + moments.append(moment) + moments = torch.cat(moments, dim=2) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def tiled_decode(self, x): + t = x.shape[2] + t_chunk_idx = [i for i in range(0, t, self.tile_latent_min_size_t - 1)] + if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0: + t_chunk_start_end = [[0, t]] + else: + t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i + 1] + 1] for i in range(len(t_chunk_idx) - 1)] + if t_chunk_start_end[-1][-1] > t: + t_chunk_start_end[-1][-1] = t + elif t_chunk_start_end[-1][-1] < t: + last_start_end = [t_chunk_idx[-1], t] + t_chunk_start_end.append(last_start_end) + dec_ = [] + for idx, (start, end) in enumerate(t_chunk_start_end): + chunk_x = x[:, :, start: end] + if idx != 0: + dec = self.tiled_decode2d(chunk_x)[:, :, 1:] + else: + dec = self.tiled_decode2d(chunk_x) + dec_.append(dec) + dec_ = torch.cat(dec_, dim=2) + return dec_ + + def tiled_encode2d(self, x, return_moments=False): + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split the image into 512x512 tiles and encode them separately. + rows = [] + for i in range(0, x.shape[3], overlap_size): + row = [] + for j in range(0, x.shape[4], overlap_size): + tile = x[ + :, + :, + :, + i: i + self.tile_sample_min_size, + j: j + self.tile_sample_min_size, + ] + tile = self.encoder(tile) + if self.use_quant_layer: + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=4)) + + moments = torch.cat(result_rows, dim=3) + posterior = DiagonalGaussianDistribution(moments) + if return_moments: + return moments + return posterior + + def tiled_decode2d(self, z): + + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping 64x64 tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[3], overlap_size): + row = [] + for j in range(0, z.shape[4], overlap_size): + tile = z[ + :, + :, + :, + i: i + self.tile_latent_min_size, + j: j + self.tile_latent_min_size, + ] + if self.use_quant_layer: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=4)) + + dec = torch.cat(result_rows, dim=3) + return dec + + def enable_tiling(self, use_tiling: bool = True): + self.use_tiling = use_tiling + + def disable_tiling(self): + self.enable_tiling(False) + + def init_from_ckpt(self, path, ignore_keys=list()): + sd = torch.load(path, map_location="cpu") + print("init from " + path) + + if "ema_state_dict" in sd and len(sd['ema_state_dict']) > 0 and os.environ.get("NOT_USE_EMA_MODEL", 0) == 0: + print("Load from ema model!") + sd = sd["ema_state_dict"] + sd = {key.replace("module.", ""): value for key, value in sd.items()} + elif "state_dict" in sd: + print("Load from normal model!") + if "gen_model" in sd["state_dict"]: + sd = sd["state_dict"]["gen_model"] + else: + sd = sd["state_dict"] + + keys = list(sd.keys()) + + for k in keys: + for ik in ignore_keys: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + + miss, unexpected = self.load_state_dict(sd, strict=False) + assert len(miss) == 0, f"miss key: {miss}" + if len(unexpected) > 0: + for i in unexpected: + assert 'loss' in i, "unexpected key: {i}" diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modeling_videobase.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modeling_videobase.py new file mode 100755 index 0000000000000000000000000000000000000000..f7686a0f27a0b51a19c2ee9a2220e3149a9b2dcf --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modeling_videobase.py @@ -0,0 +1,50 @@ +import torch +import os +import json +from diffusers.configuration_utils import ConfigMixin +from diffusers.models.modeling_utils import ModelMixin +from typing import Optional, Union +import glob + + +class VideoBaseAE(ModelMixin, ConfigMixin): + config_name = "config.json" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + def encode(self, x: torch.Tensor, *args, **kwargs): + pass + + def decode(self, encoding: torch.Tensor, *args, **kwargs): + pass + + @property + def num_training_steps(self) -> int: + """Total training steps inferred from datamodule and devices.""" + if self.trainer.max_steps: + return self.trainer.max_steps + + limit_batches = self.trainer.limit_train_batches + batches = len(self.train_dataloader()) + batches = min(batches, limit_batches) if isinstance(limit_batches, int) else int(limit_batches * batches) + + num_devices = max(1, self.trainer.num_gpus, self.trainer.num_processes) + if self.trainer.tpu_cores: + num_devices = max(num_devices, self.trainer.tpu_cores) + + effective_accum = self.trainer.accumulate_grad_batches * num_devices + return (batches // effective_accum) * self.trainer.max_epochs + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs): + ckpt_files = glob.glob(os.path.join(pretrained_model_name_or_path, '*.ckpt')) + if ckpt_files: + # Adapt to checkpoint + last_ckpt_file = ckpt_files[-1] + config_file = os.path.join(pretrained_model_name_or_path, cls.config_name) + model = cls.from_config(config_file) + model.init_from_ckpt(last_ckpt_file) + return model + else: + return super().from_pretrained(pretrained_model_name_or_path, **kwargs) \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..7ef41edd27bf3c4b9909a991c11d32d5c84f106e --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/__init__.py @@ -0,0 +1,6 @@ +from .block import Block +from .attention import * +from .conv import * +from .normalize import * +from .resnet_block import * +from .updownsample import * diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/attention.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/attention.py new file mode 100755 index 0000000000000000000000000000000000000000..ac4c96e3b05163fbdf8e214e279b3bc889bd1145 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/attention.py @@ -0,0 +1,230 @@ +import torch.nn as nn +import torch_npu +from .normalize import Normalize +from .conv import CausalConv3d +import torch +from einops import rearrange +from .block import Block +from .ops import video_to_image + +class LinearAttention(Block): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange( + qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3 + ) + k = k.softmax(dim=-1) + context = torch.einsum("bhdn,bhen->bhde", k, v) + out = torch.einsum("bhde,bhdn->bhen", context, q) + out = rearrange( + out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w + ) + return self.to_out(out) + + +class LinAttnBlock(LinearAttention): + """to match AttnBlock usage""" + + def __init__(self, in_channels): + super().__init__(dim=in_channels, heads=1, dim_head=in_channels) + + +class AttnBlock3D(Block): + """Compatible with old versions, there are issues, use with caution.""" + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b, c, t, h, w = q.shape + q = q.reshape(b * t, c, h * w) + q = q.permute(0, 2, 1) # b,hw,c + k = k.reshape(b * t, c, h * w) # b,c,hw + w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b * t, c, h * w) + w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b, c, t, h, w) + + h_ = self.proj_out(h_) + + return x + h_ + +class AttnBlock3DFix(nn.Module): + """ + Thanks to https://github.com/PKU-YuanGroup/Open-Sora-Plan/pull/172. + """ + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + # q: (b c t h w) -> (b t c h w) -> (b*t c h*w) -> (b*t h*w c) + b, c, t, h, w = q.shape + q = q.permute(0, 2, 1, 3, 4) + q = q.reshape(b * t, c, h * w) + q = q.permute(0, 2, 1) + + # k: (b c t h w) -> (b t c h w) -> (b*t c h*w) + k = k.permute(0, 2, 1, 3, 4) + k = k.reshape(b * t, c, h * w) + + # w: (b*t hw hw) + w_ = torch.bmm(q, k) + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + # v: (b c t h w) -> (b t c h w) -> (bt c hw) + # w_: (bt hw hw) -> (bt hw hw) + v = v.permute(0, 2, 1, 3, 4) + v = v.reshape(b * t, c, h * w) + w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + + # h_: (b*t c hw) -> (b t c h w) -> (b c t h w) + h_ = h_.reshape(b, t, c, h, w) + h_ = h_.permute(0, 2, 1, 3 ,4) + + h_ = self.proj_out(h_) + + return x + h_ + + +class AttnBlock(Block): + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.v = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.proj_out = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + + @video_to_image + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b, c, h, w = q.shape + q = q.reshape(b, c, h * w) + q = q.permute(0, 2, 1) # b,hw,c + k = k.reshape(b, c, h * w) # b,c,hw + w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b, c, h * w) + w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b, c, h, w) + + h_ = self.proj_out(h_) + + return x + h_ + + +class TemporalAttnBlock(Block): + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = torch.nn.Conv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = torch.nn.Conv3d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.v = torch.nn.Conv3d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.proj_out = torch.nn.Conv3d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b, c, t, h, w = q.shape + q = rearrange(q, "b c t h w -> (b h w) t c") + k = rearrange(k, "b c t h w -> (b h w) c t") + v = rearrange(v, "b c t h w -> (b h w) c t") + w_ = torch.bmm(q, k) + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + w_ = w_.permute(0, 2, 1) + h_ = torch.bmm(v, w_) + h_ = rearrange(h_, "(b h w) c t -> b c t h w", h=h, w=w) + h_ = self.proj_out(h_) + + return x + h_ + + +def make_attn(in_channels, attn_type="vanilla"): + assert attn_type in ["vanilla", "linear", "none", "vanilla3D"], f"attn_type {attn_type} unknown" + print(f"making attention of type '{attn_type}' with {in_channels} in_channels") + print(attn_type) + if attn_type == "vanilla": + return AttnBlock(in_channels) + elif attn_type == "vanilla3D": + return AttnBlock3D(in_channels) + elif attn_type == "none": + return nn.Identity(in_channels) + else: + return LinAttnBlock(in_channels) \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/block.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/block.py new file mode 100755 index 0000000000000000000000000000000000000000..e93672d20b22cbaaa36b379473149c55f0f44001 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/block.py @@ -0,0 +1,5 @@ +import torch.nn as nn + +class Block(nn.Module): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/conv.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/conv.py new file mode 100755 index 0000000000000000000000000000000000000000..629715b535d184495f33317e5ce3775b3a460a04 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/conv.py @@ -0,0 +1,117 @@ +import torch.nn as nn +from typing import Union, Tuple +import torch.nn.functional as F +import torch +import torch_npu +from .block import Block +from .ops import cast_tuple +from einops import rearrange +from .ops import video_to_image +from torch.utils.checkpoint import checkpoint + +class Conv2d(nn.Conv2d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int]] = 3, + stride: Union[int, Tuple[int]] = 1, + padding: Union[str, int, Tuple[int]] = 0, + dilation: Union[int, Tuple[int]] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + device=None, + dtype=None, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + groups, + bias, + padding_mode, + device, + dtype, + ) + + @video_to_image + def forward(self, x): + return super().forward(x) + + + +class CausalConv3d(nn.Module): + def __init__( + self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], init_method="random", **kwargs + ): + super().__init__() + self.kernel_size = cast_tuple(kernel_size, 3) + self.time_kernel_size = self.kernel_size[0] + self.chan_in = chan_in + self.chan_out = chan_out + stride = kwargs.pop("stride", 1) + padding = kwargs.pop("padding", 0) + padding = list(cast_tuple(padding, 3)) + padding[0] = 0 + stride = cast_tuple(stride, 3) + self.conv = nn.Conv3d(chan_in, chan_out, self.kernel_size, stride=stride, padding=padding) + self.pad = nn.ReplicationPad2d((0, 0, self.time_kernel_size - 1, 0)) + self._init_weights(init_method) + + def _init_weights(self, init_method): + ks = torch.tensor(self.kernel_size) + if init_method == "avg": + assert ( + self.kernel_size[1] == 1 and self.kernel_size[2] == 1 + ), "only support temporal up/down sample" + assert self.chan_in == self.chan_out, "chan_in must be equal to chan_out" + weight = torch.zeros((self.chan_out, self.chan_in, *self.kernel_size)) + + eyes = torch.concat( + [ + torch.eye(self.chan_in).unsqueeze(-1) * 1/3, + torch.eye(self.chan_in).unsqueeze(-1) * 1/3, + torch.eye(self.chan_in).unsqueeze(-1) * 1/3, + ], + dim=-1, + ) + weight[:, :, :, 0, 0] = eyes + + self.conv.weight = nn.Parameter( + weight, + requires_grad=True, + ) + elif init_method == "zero": + self.conv.weight = nn.Parameter( + torch.zeros((self.chan_out, self.chan_in, *self.kernel_size)), + requires_grad=True, + ) + if self.conv.bias is not None: + nn.init.constant_(self.conv.bias, 0) + + def forward(self, x): + x_dtype = x.dtype + # 1 + 16 16 as video, 1 as image + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.time_kernel_size - 1, 1, 1) + ) # b c t h w + x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16 + x = self.conv.to(device=x.device, dtype=torch.float16)(x.to(torch.float16)) + x = x.to(x_dtype) + return torch_npu.npu_format_cast(x, 2) + + +class CausalConv3d_GC(CausalConv3d): + def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int]], init_method="random", **kwargs): + super().__init__(chan_in, chan_out, kernel_size, init_method, **kwargs) + def forward(self, x): + # 1 + 16 16 as video, 1 as image + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.time_kernel_size - 1, 1, 1) + ) # b c t h w + x = torch.concatenate((first_frame_pad, x), dim=2) # 3 + 16 + return checkpoint(self.conv, x) \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/normalize.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/normalize.py new file mode 100755 index 0000000000000000000000000000000000000000..7c8c05f0fa3459214dd077e09a157c588adb637b --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/normalize.py @@ -0,0 +1,101 @@ +import torch +import torch.nn as nn +from .block import Block + +class GroupNorm(Block): + def __init__(self, num_channels, num_groups=32, eps=1e-6, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.norm = torch.nn.GroupNorm( + num_groups=num_groups, num_channels=num_channels, eps=1e-6, affine=True + ) + def forward(self, x): + return self.norm(x) + +def Normalize(in_channels, num_groups=32): + return torch.nn.GroupNorm( + num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) + +class ActNorm(nn.Module): + def __init__(self, num_features, logdet=False, affine=True, + allow_reverse_init=False): + assert affine + super().__init__() + self.logdet = logdet + self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1)) + self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1)) + self.allow_reverse_init = allow_reverse_init + + self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8)) + + def initialize(self, input): + with torch.no_grad(): + flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1) + mean = ( + flatten.mean(1) + .unsqueeze(1) + .unsqueeze(2) + .unsqueeze(3) + .permute(1, 0, 2, 3) + ) + std = ( + flatten.std(1) + .unsqueeze(1) + .unsqueeze(2) + .unsqueeze(3) + .permute(1, 0, 2, 3) + ) + + self.loc.data.copy_(-mean) + self.scale.data.copy_(1 / (std + 1e-6)) + + def forward(self, input, reverse=False): + if reverse: + return self.reverse(input) + if len(input.shape) == 2: + input = input[:,:,None,None] + squeeze = True + else: + squeeze = False + + _, _, height, width = input.shape + + if self.training and self.initialized.item() == 0: + self.initialize(input) + self.initialized.fill_(1) + + h = self.scale * (input + self.loc) + + if squeeze: + h = h.squeeze(-1).squeeze(-1) + + if self.logdet: + log_abs = torch.log(torch.abs(self.scale)) + logdet = height*width*torch.sum(log_abs) + logdet = logdet * torch.ones(input.shape[0]).to(input) + return h, logdet + + return h + + def reverse(self, output): + if self.training and self.initialized.item() == 0: + if not self.allow_reverse_init: + raise RuntimeError( + "Initializing ActNorm in reverse direction is " + "disabled by default. Use allow_reverse_init=True to enable." + ) + else: + self.initialize(output) + self.initialized.fill_(1) + + if len(output.shape) == 2: + output = output[:,:,None,None] + squeeze = True + else: + squeeze = False + + h = output / self.scale - self.loc + + if squeeze: + h = h.squeeze(-1).squeeze(-1) + return h diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/ops.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/ops.py new file mode 100755 index 0000000000000000000000000000000000000000..e298cceebf3f76e0c9afdf75a8dcf942d4515f4c --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/ops.py @@ -0,0 +1,47 @@ +import torch +import torch_npu +import torch.nn as nn +from einops import rearrange + + +def video_to_image(func): + def wrapper(self, x, *args, **kwargs): + if x.dim() == 5: + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = func(self, x, *args, **kwargs) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + return x + + return wrapper + + +def nonlinearity(x): + return nn.functional.silu(x) + + +def cast_tuple(t, length=1): + return t if isinstance(t, tuple) or isinstance(t, list) else ((t,) * length) + + +def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True): + n_dims = len(x.shape) + if src_dim < 0: + src_dim = n_dims + src_dim + if dest_dim < 0: + dest_dim = n_dims + dest_dim + assert 0 <= src_dim < n_dims and 0 <= dest_dim < n_dims + dims = list(range(n_dims)) + del dims[src_dim] + permutation = [] + ctr = 0 + for i in range(n_dims): + if i == dest_dim: + permutation.append(src_dim) + else: + permutation.append(dims[ctr]) + ctr += 1 + x = x.permute(permutation) + if make_contiguous: + x = x.contiguous() + return x diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/resnet_block.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/resnet_block.py new file mode 100755 index 0000000000000000000000000000000000000000..bd225a81af15adfd8278099822f230088ad0324e --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/resnet_block.py @@ -0,0 +1,126 @@ +import torch +import torch.nn as nn +from einops import rearrange, pack, unpack +from .normalize import Normalize +from .ops import nonlinearity, video_to_image +from .conv import CausalConv3d, CausalConv3d_GC +from .block import Block +from torch.utils.checkpoint import checkpoint +import torch_npu + +class ResnetBlock2D(Block): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, + dropout): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = torch.nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = torch.nn.Conv2d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = torch.nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + else: + self.nin_shortcut = torch.nn.Conv2d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + @video_to_image + def forward(self, x): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + x = x + h + return x + +class ResnetBlock3D(Block): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1) + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1) + else: + self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + return x + h + + +class ResnetBlock3D_GC(Block): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1) + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1) + else: + self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0) + + def forward(self, x): + return checkpoint(self._forward, x, use_reentrant=True) + + def _forward(self, x): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + return x + h \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/updownsample.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/updownsample.py new file mode 100755 index 0000000000000000000000000000000000000000..9308b77a977c1bce55e97b232bcf090ef27a4828 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/modules/updownsample.py @@ -0,0 +1,350 @@ +from typing import Union, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from .attention import TemporalAttnBlock +from .block import Block +from .conv import CausalConv3d, CausalConv3d_GC +from .normalize import Normalize +from .ops import cast_tuple, video_to_image +from .resnet_block import ResnetBlock3D + + +class Upsample(Block): + def __init__(self, in_channels, out_channels): + super().__init__() + self.with_conv = True + if self.with_conv: + self.conv = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + + @video_to_image + def forward(self, x): + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample(Block): + def __init__(self, in_channels, out_channels, undown=False): + super().__init__() + self.with_conv = True + self.undown = undown + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + if self.undown: + self.conv = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + else: + self.conv = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=2, + padding=0) + + @video_to_image + def forward(self, x): + if self.with_conv: + if self.undown: + x = self.conv(x) + else: + pad = (0, 1, 0, 1) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + return x + + +class SpatialDownsample2x(Block): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int]] = (3, 3), + stride: Union[int, Tuple[int]] = (2, 2), + **kwargs + ): + super().__init__() + kernel_size = cast_tuple(kernel_size, 2) + stride = cast_tuple(stride, 2) + self.chan_in = chan_in + self.chan_out = chan_out + self.kernel_size = kernel_size + self.conv = CausalConv3d( + self.chan_in, + self.chan_out, + (1,) + self.kernel_size, + stride=(1,) + stride, + padding=0 + ) + + def forward(self, x): + pad = (0, 1, 0, 1, 0, 0) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + return x + + +class SpatialUpsample2x_GC(Block): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int]] = (3, 3), + stride: Union[int, Tuple[int]] = (1, 1), + unup=False, + ): + super().__init__() + self.chan_in = chan_in + self.chan_out = chan_out + self.kernel_size = kernel_size + self.unup = unup + self.conv = CausalConv3d_GC( + self.chan_in, + self.chan_out, + (1,) + self.kernel_size, + stride=(1,) + stride, + padding=1 + ) + + def forward(self, x): + if not self.unup: + t = x.shape[2] + x = rearrange(x, "b c t h w -> b (c t) h w") + x = F.interpolate(x, scale_factor=(2, 2), mode="nearest") + x = rearrange(x, "b (c t) h w -> b c t h w", t=t) + x = self.conv(x) + return x + + +class SpatialUpsample2x(Block): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int]] = (3, 3), + stride: Union[int, Tuple[int]] = (1, 1), + unup=False, + ): + super().__init__() + self.chan_in = chan_in + self.chan_out = chan_out + self.kernel_size = kernel_size + self.unup = unup + self.conv = CausalConv3d( + self.chan_in, + self.chan_out, + (1,) + self.kernel_size, + stride=(1,) + stride, + padding=1 + ) + + def forward(self, x): + if not self.unup: + t = x.shape[2] + x = rearrange(x, "b c t h w -> b (c t) h w") + x = F.interpolate(x, scale_factor=(2, 2), mode="nearest") + x = rearrange(x, "b (c t) h w -> b c t h w", t=t) + x = self.conv(x) + return x + + +class TimeDownsample2x(Block): + def __init__( + self, + chan_in, + chan_out, + kernel_size: int = 3 + ): + super().__init__() + self.kernel_size = kernel_size + self.avg_pool = nn.AvgPool2d((kernel_size, 1), stride=(2, 1)) + self.pad = nn.ReplicationPad3d((0, 0, 0, 0, self.kernel_size - 1, 0)) + + def forward(self, x): + n, c, d, h, w = x.shape + x = self.pad(x) + x = x.view(n * c, -1, h * w) + pooled = self.avg_pool(x) + output = pooled.view(n, c, -1, h, w) + return output + + +class TimeUpsample2x(Block): + def __init__( + self, + chan_in, + chan_out + ): + super().__init__() + + def forward(self, x): + if x.size(2) > 1: + x, x_ = x[:, :, :1], x[:, :, 1:] + x_ = F.interpolate(x_, scale_factor=(2, 1, 1), mode='trilinear') + x = torch.concat([x, x_], dim=2) + return x + + +class TimeDownsampleRes2x(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size: int = 3, + mix_factor: float = 2.0, + ): + super().__init__() + self.kernel_size = cast_tuple(kernel_size, 3) + self.avg_pool = nn.AvgPool2d((kernel_size, 1), stride=(2, 1)) + self.pad = nn.ReplicationPad3d((0, 0, 0, 0, kernel_size - 1, 0)) + self.conv = nn.Conv3d( + in_channels, out_channels, self.kernel_size, stride=(2, 1, 1), padding=(0, 1, 1) + ) + self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor])) + + def forward(self, x): + alpha = torch.sigmoid(self.mix_factor) + n, c, d, h, w = x.shape + x = self.pad(x) + pad_x = x.view(n, c, -1, h, w) + avg_x = self.avg_pool(x.view(n * c, -1, h * w)).view(n, c, -1, h, w).to(x_dtype) + conv_x = self.conv(pad_x) + return alpha * avg_x + (1 - alpha) * conv_x + + +class TimeUpsampleRes2x(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size: int = 3, + mix_factor: float = 2.0, + ): + super().__init__() + self.conv = CausalConv3d( + in_channels, out_channels, kernel_size, padding=1 + ) + self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor])) + + def forward(self, x): + alpha = torch.sigmoid(self.mix_factor) + if x.size(2) > 1: + x, x_ = x[:, :, :1], x[:, :, 1:] + x_ = F.interpolate(x_, scale_factor=(2, 1, 1), mode='trilinear') + x = torch.concat([x, x_], dim=2) + return alpha * x + (1 - alpha) * self.conv(x) + + +class TimeDownsampleResAdv2x(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size: int = 3, + mix_factor: float = 1.5, + ): + super().__init__() + self.kernel_size = cast_tuple(kernel_size, 3) + self.avg_pool = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1)) + self.attn = TemporalAttnBlock(in_channels) + self.res = ResnetBlock3D(in_channels=in_channels, out_channels=in_channels, dropout=0.0) + self.conv = nn.Conv3d( + in_channels, out_channels, self.kernel_size, stride=(2, 1, 1), padding=(0, 1, 1) + ) + self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor])) + + def forward(self, x): + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.kernel_size[0] - 1, 1, 1) + ) + x = torch.concatenate((first_frame_pad, x), dim=2) + alpha = torch.sigmoid(self.mix_factor) + return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(self.attn((self.res(x)))) + + +class TimeUpsampleResAdv2x(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size: int = 3, + mix_factor: float = 1.5, + ): + super().__init__() + self.res = ResnetBlock3D(in_channels=in_channels, out_channels=in_channels, dropout=0.0) + self.attn = TemporalAttnBlock(in_channels) + self.norm = Normalize(in_channels=in_channels) + self.conv = CausalConv3d( + in_channels, out_channels, kernel_size, padding=1 + ) + self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor])) + + def forward(self, x): + if x.size(2) > 1: + x, x_ = x[:, :, :1], x[:, :, 1:] + x_ = F.interpolate(x_, scale_factor=(2, 1, 1), mode='trilinear') + x = torch.concat([x, x_], dim=2) + alpha = torch.sigmoid(self.mix_factor) + return alpha * x + (1 - alpha) * self.conv(self.attn(self.res(x))) + + +class Spatial2xTime2x3DDownsample(Block): + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=2) + + def forward(self, x): + pad = (0, 1, 0, 1, 0, 0) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + return x + + +class Spatial2x3DDownsample(Block): + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=(1, 2, 2)) + + def forward(self, x): + pad = (0, 1, 0, 1, 0, 0) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + return x + + +class Spatial2x3DUpsample(Block): + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1) + + def forward(self, x): + x = F.interpolate(x, scale_factor=(1, 2, 2), mode='trilinear') + return self.conv(x) + + +class Spatial2xTime2x3DUpsample(Block): + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1) + + def forward(self, x): + if x.size(2) > 1: + x, x_ = x[:, :, :1], x[:, :, 1:] + x_ = F.interpolate(x_, scale_factor=(2, 2, 2), mode='trilinear') + x = F.interpolate(x, scale_factor=(1, 2, 2), mode='trilinear') + x = torch.concat([x, x_], dim=2) + else: + x = F.interpolate(x, scale_factor=(1, 2, 2), mode='trilinear') + return self.conv(x) diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/distrib_utils.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/distrib_utils.py new file mode 100755 index 0000000000000000000000000000000000000000..760c0673fe5d8afa663eb1ea5cd7683dbf5dd9f8 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/distrib_utils.py @@ -0,0 +1,42 @@ +import torch +import numpy as np + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) + + def sample(self): + x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) + return x + + def kl(self, other=None): + if self.deterministic: + return torch.Tensor([0.]) + else: + if other is None: + return 0.5 * torch.sum(torch.pow(self.mean, 2) + + self.var - 1.0 - self.logvar, + dim=[1, 2, 3]) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + dim=[1, 2, 3]) + + def nll(self, sample, dims=[1,2,3]): + if self.deterministic: + return torch.Tensor([0.]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims) + + def mode(self): + return self.mean diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/module_utils.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/module_utils.py new file mode 100755 index 0000000000000000000000000000000000000000..3f843e4f3d141cf1226944456213bc7960aa4fc4 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/causalvideovae/model/utils/module_utils.py @@ -0,0 +1,17 @@ +import importlib + +Module = str +MODULES_BASE = "opensora.models.causalvideovae.model.modules." + +def resolve_str_to_obj(str_val, append=True): + if append: + str_val = MODULES_BASE + str_val + module_name, class_name = str_val.rsplit('.', 1) + module = importlib.import_module(module_name) + return getattr(module, class_name) + +def create_instance(module_class_str: str, **kwargs): + module_name, class_name = module_class_str.rsplit('.', 1) + module = importlib.import_module(module_name) + class_ = getattr(module, class_name) + return class_(**kwargs) \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b0dac65653c4ac827e44864913497651c8434874 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/__init__.py @@ -0,0 +1,2 @@ + + \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/__init__.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/modeling_opensora.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/modeling_opensora.py new file mode 100644 index 0000000000000000000000000000000000000000..36b9037b53cfbd9556ea4dc6b3929e2d21d1bc08 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/modeling_opensora.py @@ -0,0 +1,571 @@ +import collections +from typing import Any, Dict, Optional + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.embeddings import PixArtAlphaTextProjection +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import AdaLayerNormSingle +from diffusers.models.transformer_2d import Transformer2DModelOutput +from diffusers.utils import deprecate, logging +from einops import rearrange, repeat +from torch import nn +from torch.nn import functional as F + +from opensora.models.diffusion.opensora.modules import PatchEmbed2D, BasicTransformerBlock, LayerNorm +from utils.parallel_mgr import get_sequence_parallel_state, get_sequence_parallel_size + +logger = logging.get_logger(__name__) + + +def to_2tuple(x): + if isinstance(x, collections.abc.Iterable): + return x + return (x, x) + + +class OpenSoraT2V(ModelMixin, ConfigMixin): + """ + A 2D Transformer model for image-like data. + + Parameters: + num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. + attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. + in_channels (`int`, *optional*): + The number of channels in the input and output (specify if the input is **continuous**). + num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. + sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). + This is fixed during training since it is used to learn a number of position embeddings. + num_vector_embeds (`int`, *optional*): + The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**). + Includes the class for the masked latent pixel. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward. + num_embeds_ada_norm ( `int`, *optional*): + The number of diffusion steps used during training. Pass if at least one of the norm_layers is + `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are + added to the hidden states. + + During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`. + attention_bias (`bool`, *optional*): + Configure if the `TransformerBlocks` attention should contain a bias parameter. + """ + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + out_channels: Optional[int] = None, + num_layers: int = 1, + dropout: float = 0.0, + norm_num_groups: int = 32, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + sample_size: Optional[int] = None, + sample_size_t: Optional[int] = None, + num_vector_embeds: Optional[int] = None, + patch_size: Optional[int] = None, + patch_size_t: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_type: str = "layer_norm", + # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + attention_type: str = "default", + caption_channels: int = None, + interpolation_scale_h: float = None, + interpolation_scale_w: float = None, + interpolation_scale_t: float = None, + use_additional_conditions: Optional[bool] = None, + attention_mode: str = 'xformers', + downsampler: str = None, + use_rope: bool = False, + use_stable_fp32: bool = False, + ): + super().__init__() + + # Validate inputs. + if patch_size is not None: + if norm_type not in ["ada_norm", "ada_norm_zero", "ada_norm_single"]: + raise NotImplementedError( + f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'." + ) + elif norm_type in ["ada_norm", "ada_norm_zero"] and num_embeds_ada_norm is None: + raise ValueError( + f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None." + ) + + # Set some common variables used across the board. + self.use_rope = use_rope + self.use_linear_projection = use_linear_projection + self.interpolation_scale_t = interpolation_scale_t + self.interpolation_scale_h = interpolation_scale_h + self.interpolation_scale_w = interpolation_scale_w + self.downsampler = downsampler + self.caption_channels = caption_channels + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim + self.in_channels = in_channels + self.out_channels = in_channels if out_channels is None else out_channels + self.gradient_checkpointing = False + self.use_additional_conditions = False + self.cache = None + + # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)` + # Define whether input is continuous or discrete depending on configuration + assert in_channels is not None and patch_size is not None + + if norm_type == "layer_norm" and num_embeds_ada_norm is not None: + deprecation_message = ( + f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or" + " incorrectly set to `'layer_norm'`. Make sure to set `norm_type` to `'ada_norm'` in the config." + " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect" + " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it" + " would be very nice if you could open a Pull request for the `transformer/config.json` file" + ) + deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False) + norm_type = "ada_norm" + + # 2. Initialize the right blocks. + # Initialize the output blocks and other projection blocks when necessary. + self._init_patched_inputs(norm_type=norm_type) + + def _init_patched_inputs(self, norm_type): + assert self.config.sample_size_t is not None, "OpenSoraT2V over patched input must provide sample_size_t" + assert self.config.sample_size is not None, "OpenSoraT2V over patched input must provide sample_size" + + self.num_frames = self.config.sample_size_t + self.config.sample_size = to_2tuple(self.config.sample_size) + self.height = self.config.sample_size[0] + self.width = self.config.sample_size[1] + self.patch_size_t = self.config.patch_size_t + self.patch_size = self.config.patch_size + interpolation_scale_t = (( + self.config.sample_size_t - 1) // 16 + 1) if self.config.sample_size_t % 2 == 1 else self.config.sample_size_t / 16 + interpolation_scale_t = ( + self.config.interpolation_scale_t if self.config.interpolation_scale_t is not None else interpolation_scale_t + ) + interpolation_scale = ( + self.config.interpolation_scale_h if self.config.interpolation_scale_h is not None else + self.config.sample_size[0] / 30, + self.config.interpolation_scale_w if self.config.interpolation_scale_w is not None else + self.config.sample_size[1] / 40, + ) + self.pos_embed = PatchEmbed2D( + num_frames=self.config.sample_size_t, + height=self.config.sample_size[0], + width=self.config.sample_size[1], + patch_size_t=self.config.patch_size_t, + patch_size=self.config.patch_size, + in_channels=self.in_channels, + embed_dim=self.inner_dim, + interpolation_scale=interpolation_scale, + interpolation_scale_t=interpolation_scale_t, + use_abs_pos=not self.config.use_rope, + ) + interpolation_scale_thw = (interpolation_scale_t, *interpolation_scale) + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + self.inner_dim, + self.config.num_attention_heads, + self.config.attention_head_dim, + dropout=self.config.dropout, + cross_attention_dim=self.config.cross_attention_dim, + activation_fn=self.config.activation_fn, + num_embeds_ada_norm=self.config.num_embeds_ada_norm, + attention_bias=self.config.attention_bias, + only_cross_attention=self.config.only_cross_attention, + double_self_attention=self.config.double_self_attention, + upcast_attention=self.config.upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=self.config.norm_elementwise_affine, + norm_eps=self.config.norm_eps, + attention_type=self.config.attention_type, + attention_mode=self.config.attention_mode, + downsampler=self.config.downsampler, + use_rope=self.config.use_rope, + interpolation_scale_thw=interpolation_scale_thw, + ) + for _ in range(self.config.num_layers) + ] + ) + + if self.config.norm_type != "ada_norm_single": + self.norm_out = LayerNorm(self.inner_dim, eps=1e-6) + self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) + self.proj_out_2 = nn.Linear( + self.inner_dim, + self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels + ) + elif self.config.norm_type == "ada_norm_single": + self.norm_out = LayerNorm(self.inner_dim, eps=1e-6) + 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, + self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels + ) + + # PixArt-Alpha blocks. + self.adaln_single = None + if self.config.norm_type == "ada_norm_single": + # additional conditions until we find better name + self.adaln_single = AdaLayerNormSingle( + self.inner_dim, use_additional_conditions=self.use_additional_conditions + ) + + self.caption_projection = None + if self.caption_channels is not None: + self.caption_projection = PixArtAlphaTextProjection( + in_features=self.caption_channels, hidden_size=self.inner_dim + ) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states: torch.Tensor, + timestep: Optional[torch.LongTensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + class_labels: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + attention_mask: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + step_id: int = 0, + use_image_num: Optional[int] = 0, + return_dict: bool = True, + ): + """ + The [`Transformer2DModel`] forward method. + + Args: + hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous): + Input `hidden_states`. + encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. If not given, cross-attention defaults to + self-attention. + timestep ( `torch.LongTensor`, *optional*): + Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. + class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): + Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in + `AdaLayerZeroNorm`. + cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + encoder_attention_mask ( `torch.Tensor`, *optional*): + Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: + + * Mask `(batch, sequence_length)` True = keep, False = discard. + * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard. + + If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format + above. This bias will be added to the cross-attention scores. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + batch_size, c, frame, h, w = hidden_states.shape + frame = frame - use_image_num # 21-4=17 + if cross_attention_kwargs is not None: + if cross_attention_kwargs.get("scale", None) is not None: + logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") + + attention_mask_vid, attention_mask_img = None, None + if attention_mask is not None and attention_mask.ndim == 4: + attention_mask = attention_mask.to(self.dtype) + if get_sequence_parallel_state(): + attention_mask_vid = attention_mask[:, :frame * get_sequence_parallel_size()] # b, frame, h, w + attention_mask_img = attention_mask[:, frame * get_sequence_parallel_size():] # b, use_image_num, h, w + else: + attention_mask_vid = attention_mask[:, :frame] # b, frame, h, w + attention_mask_img = attention_mask[:, frame:] # b, use_image_num, h, w + + if attention_mask_vid.numel() > 0: + attention_mask_vid_first_frame = attention_mask_vid[:, :1].repeat(1, self.patch_size_t - 1, 1, 1) + attention_mask_vid = torch.cat([attention_mask_vid_first_frame, attention_mask_vid], dim=1) + attention_mask_vid = attention_mask_vid.unsqueeze(1) # b 1 t h w + attention_mask_vid = F.max_pool3d(attention_mask_vid, + kernel_size=(self.patch_size_t, self.patch_size, self.patch_size), + stride=(self.patch_size_t, self.patch_size, self.patch_size)) + attention_mask_vid = rearrange(attention_mask_vid, 'b 1 t h w -> (b 1) 1 (t h w)') + if attention_mask_img.numel() > 0: + attention_mask_img = F.max_pool2d(attention_mask_img, kernel_size=(self.patch_size, self.patch_size), + stride=(self.patch_size, self.patch_size)) + attention_mask_img = rearrange(attention_mask_img, 'b i h w -> (b i) 1 (h w)') + + attention_mask_vid = (1 - attention_mask_vid.bool().to( + self.dtype)) * -10000.0 if attention_mask_vid.numel() > 0 else None + attention_mask_img = (1 - attention_mask_img.bool().to( + self.dtype)) * -10000.0 if attention_mask_img.numel() > 0 else None + + if frame == 1 and use_image_num == 0 and not get_sequence_parallel_state(): + attention_mask_img = attention_mask_vid + attention_mask_vid = None + # convert encoder_attention_mask to a bias the same way we do for attention_mask + encoder_attention_mask_vid, encoder_attention_mask_img = None, None + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 3: + # b, 1+use_image_num, l -> a video with images + # b, 1, l -> only images + encoder_attention_mask = (1 - encoder_attention_mask.to(self.dtype)) * -10000.0 + in_t = encoder_attention_mask.shape[1] + encoder_attention_mask_vid = encoder_attention_mask[:, :in_t - use_image_num] # b, 1, l + encoder_attention_mask_vid = rearrange(encoder_attention_mask_vid, + 'b 1 l -> (b 1) 1 l') if encoder_attention_mask_vid.numel() > 0 else None + + encoder_attention_mask_img = encoder_attention_mask[:, in_t - use_image_num:] # b, use_image_num, l + encoder_attention_mask_img = rearrange(encoder_attention_mask_img, + 'b i l -> (b i) 1 l') if encoder_attention_mask_img.numel() > 0 else None + + if frame == 1 and use_image_num == 0 and not get_sequence_parallel_state(): + encoder_attention_mask_img = encoder_attention_mask_vid + encoder_attention_mask_vid = None + + if attention_mask_vid is not None: + attention_mask_vid = attention_mask_vid.to(torch.bool) + attention_mask_vid = attention_mask_vid.repeat(1, attention_mask_vid.shape[-1], 1) + encoder_attention_mask_vid = encoder_attention_mask_vid.to(torch.bool) + encoder_attention_mask_vid = encoder_attention_mask_vid.repeat(1, attention_mask_vid.shape[-2], 1) + if attention_mask_img is not None: + attention_mask_img = attention_mask_img.to(torch.bool) + attention_mask_img = attention_mask_img.repeat(1, attention_mask_vid.shape[-1], 1) + encoder_attention_mask_img = encoder_attention_mask_img.to(torch.bool) + encoder_attention_mask_img = encoder_attention_mask_img.repeat(1, attention_mask_vid.shape[-2], 1) + + # 1. Input + frame = ((frame - 1) // self.patch_size_t + 1) if frame % 2 == 1 else frame // self.patch_size_t + height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size + + added_cond_kwargs = {"resolution": None, "aspect_ratio": None} + hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, \ + timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img = self._operate_on_patched_inputs( + hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num + ) + # 2. Blocks + if get_sequence_parallel_state(): + if hidden_states_vid is not None: + hidden_states_vid = rearrange(hidden_states_vid, 'b s h -> s b h', b=batch_size).contiguous() + encoder_hidden_states_vid = rearrange(encoder_hidden_states_vid, 'b s h -> s b h', + b=batch_size).contiguous() + timestep_vid = timestep_vid.view(batch_size, 6, -1).transpose(0, 1).contiguous() + + for i, block in enumerate(self.transformer_blocks): + if hidden_states_vid is not None: + if self.cache: + hidden_states_vid = self.cache(block, step_id, i, + hidden_states_vid, + attention_mask=attention_mask_vid, + encoder_hidden_states=encoder_hidden_states_vid, + encoder_attention_mask=encoder_attention_mask_vid, + timestep=timestep_vid, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + frame=frame, + height=height, + width=width, + ) + else: + hidden_states_vid = block( + hidden_states_vid, + attention_mask=attention_mask_vid, + encoder_hidden_states=encoder_hidden_states_vid, + encoder_attention_mask=encoder_attention_mask_vid, + timestep=timestep_vid, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + frame=frame, + height=height, + width=width, + ) + if hidden_states_img is not None: + if self.cache: + hidden_states_img = self.cache(block, step_id, i, + hidden_states_img, + attention_mask=attention_mask_img, + encoder_hidden_states=encoder_hidden_states_img, + encoder_attention_mask=encoder_attention_mask_img, + timestep=timestep_img, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + frame=1, + height=height, + width=width, + ) + else: + hidden_states_img = block( + hidden_states_img, + attention_mask=attention_mask_img, + encoder_hidden_states=encoder_hidden_states_img, + encoder_attention_mask=encoder_attention_mask_img, + timestep=timestep_img, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + frame=1, + height=height, + width=width, + ) + + if get_sequence_parallel_state(): + if hidden_states_vid is not None: + hidden_states_vid = rearrange(hidden_states_vid, 's b h -> b s h', b=batch_size).contiguous() + + # 3. Output + output_vid, output_img = None, None + if hidden_states_vid is not None: + output_vid = self._get_output_for_patched_inputs( + hidden_states=hidden_states_vid, + timestep=timestep_vid, + class_labels=class_labels, + embedded_timestep=embedded_timestep_vid, + num_frames=frame, + height=height, + width=width, + ) # b c t h w + if hidden_states_img is not None: + output_img = self._get_output_for_patched_inputs( + hidden_states=hidden_states_img, + timestep=timestep_img, + class_labels=class_labels, + embedded_timestep=embedded_timestep_img, + num_frames=1, + height=height, + width=width, + ) # b c 1 h w + if use_image_num != 0: + output_img = rearrange(output_img, '(b i) c 1 h w -> b c i h w', i=use_image_num) + + if output_vid is not None and output_img is not None: + output = torch.cat([output_vid, output_img], dim=2) + elif output_vid is not None: + output = output_vid + elif output_img is not None: + output = output_img + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + def _operate_on_patched_inputs(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, + frame, use_image_num): + hidden_states_vid, hidden_states_img = self.pos_embed(hidden_states.to(self.dtype), frame) + timestep_vid, timestep_img = None, None + embedded_timestep_vid, embedded_timestep_img = None, None + encoder_hidden_states_vid, encoder_hidden_states_img = None, None + + if self.adaln_single is not None: + if self.use_additional_conditions and added_cond_kwargs is None: + raise ValueError( + "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`." + ) + timestep, embedded_timestep = self.adaln_single( + timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype + ) # b 6d, b d + if hidden_states_vid is None: + timestep_img = timestep + embedded_timestep_img = embedded_timestep + else: + timestep_vid = timestep + embedded_timestep_vid = embedded_timestep + if hidden_states_img is not None: + timestep_img = repeat(timestep, 'b d -> (b i) d', i=use_image_num).contiguous() + embedded_timestep_img = repeat(embedded_timestep, 'b d -> (b i) d', i=use_image_num).contiguous() + + if self.caption_projection is not None: + encoder_hidden_states = self.caption_projection( + encoder_hidden_states) # b, 1+use_image_num, l, d or b, 1, l, d + if hidden_states_vid is None: + encoder_hidden_states_img = rearrange(encoder_hidden_states, 'b 1 l d -> (b 1) l d') + else: + encoder_hidden_states_vid = rearrange(encoder_hidden_states[:, :1], 'b 1 l d -> (b 1) l d') + if hidden_states_img is not None: + encoder_hidden_states_img = rearrange(encoder_hidden_states[:, 1:], 'b i l d -> (b i) l d') + + return hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img + + def _get_output_for_patched_inputs( + self, hidden_states, timestep, class_labels, embedded_timestep, num_frames, height=None, width=None + ): + if self.config.norm_type != "ada_norm_single": + conditioning = self.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=self.dtype + ) + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states, scale=(1 + scale), shift=shift) + hidden_states = self.proj_out_2(hidden_states) + elif self.config.norm_type == "ada_norm_single": + shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states, scale=(1 + scale), shift=shift) + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.squeeze(1) + + # unpatchify + if self.adaln_single is None: + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=( + -1, num_frames, height, width, self.patch_size_t, self.patch_size, self.patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nthwopqc->nctohpwq", hidden_states) + output = hidden_states.reshape( + shape=( + -1, self.out_channels, num_frames * self.patch_size_t, height * self.patch_size, + width * self.patch_size) + ) + + return output + + +def OpenSoraT2V_S_122(**kwargs): + return OpenSoraT2V(num_layers=28, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2, + norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1536, **kwargs) + + +def OpenSoraT2V_B_122(**kwargs): + return OpenSoraT2V(num_layers=32, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2, + norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1920, **kwargs) + + +def OpenSoraT2V_L_122(**kwargs): + return OpenSoraT2V(num_layers=40, attention_head_dim=128, num_attention_heads=16, patch_size_t=1, patch_size=2, + norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2048, **kwargs) + + +def OpenSoraT2V_ROPE_L_122(**kwargs): + return OpenSoraT2V(num_layers=32, attention_head_dim=96, num_attention_heads=24, patch_size_t=1, patch_size=2, + norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2304, **kwargs) + + +OpenSora_models = { + "OpenSoraT2V-S/122": OpenSoraT2V_S_122, # 1.1B + "OpenSoraT2V-B/122": OpenSoraT2V_B_122, + "OpenSoraT2V-L/122": OpenSoraT2V_L_122, + "OpenSoraT2V-ROPE-L/122": OpenSoraT2V_ROPE_L_122, +} + +OpenSora_models_class = { + "OpenSoraT2V-S/122": OpenSoraT2V, + "OpenSoraT2V-B/122": OpenSoraT2V, + "OpenSoraT2V-L/122": OpenSoraT2V, + "OpenSoraT2V-ROPE-L/122": OpenSoraT2V, +} diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/modules.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..770ce417769c3d93ab264dd7261c7b2550165115 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/modules.py @@ -0,0 +1,1175 @@ +import math +import re +from typing import Any, Dict, Optional +from typing import Tuple + +import numpy as np +import torch +import torch.nn.functional as F +import torch_npu +from diffusers.models.attention import FeedForward, GatedSelfAttentionDense +from diffusers.models.attention_processor import Attention as Attention_ +from diffusers.models.embeddings import SinusoidalPositionalEmbedding +from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero +from diffusers.utils import deprecate, logging +from diffusers.utils.torch_utils import maybe_allow_in_graph +from einops import rearrange +from torch import nn + +from utils.comm import all_to_all_sbh +from utils.parallel_mgr import get_sequence_parallel_state, get_sequence_parallel_size, get_sequence_parallel_rank +from .rope import PositionGetter3D, RoPE3D + +logger = logging.get_logger(__name__) + + +def get_3d_sincos_pos_embed( + embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, +): + """ + grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or + [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0] + grid_h = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1] + grid_w = np.arange(grid_size[2], dtype=np.float32) / (grid_size[2] / base_size[2]) / interpolation_scale[2] + grid = np.meshgrid(grid_w, grid_h, grid_t) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([3, 1, grid_size[2], grid_size[1], grid_size[0]]) + pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token and extra_tokens > 0: + pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_3d_sincos_pos_embed_from_grid(embed_dim, grid): + if embed_dim % 3 != 0: + raise ValueError("embed_dim must be divisible by 3") + + # use 1/3 of dimensions to encode grid_t/h/w + emb_t = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0]) # (T*H*W, D/3) + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1]) # (T*H*W, D/3) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2]) # (T*H*W, D/3) + + emb = np.concatenate([emb_t, emb_h, emb_w], axis=1) # (T*H*W, D) + return emb + + +def get_2d_sincos_pos_embed( + embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, +): + """ + grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or + [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0] + grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1] + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token and extra_tokens > 0: + pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + if embed_dim % 2 != 0: + raise ValueError("embed_dim must be divisible by 2") + + # use 1/3 of dimensions to encode grid_t/h/w + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) + + emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) + return emb + + +def get_1d_sincos_pos_embed( + embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, +): + """ + grid_size: int of the grid return: pos_embed: [grid_size, embed_dim] or + [1+grid_size, embed_dim] (w/ or w/o cls_token) + """ + + grid = np.arange(grid_size, dtype=np.float32) / (grid_size / base_size) / interpolation_scale + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid) # (H*W, D/2) + if cls_token and extra_tokens > 0: + pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) + """ + if embed_dim % 2 != 0: + raise ValueError("embed_dim must be divisible by 2") + + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000 ** omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +class PatchEmbed2D(nn.Module): + """2D Image to Patch Embedding but with 3D position embedding""" + + def __init__( + self, + num_frames=1, + height=224, + width=224, + patch_size_t=1, + patch_size=16, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + interpolation_scale=(1, 1), + interpolation_scale_t=1, + use_abs_pos=True, + ): + super().__init__() + self.use_abs_pos = use_abs_pos + 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, patch_size), bias=bias + ) + if layer_norm: + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) + else: + self.norm = None + + self.patch_size_t = patch_size_t + self.patch_size = patch_size + + self.height, self.width = height // patch_size, width // patch_size + self.base_size = (height // patch_size, width // patch_size) + self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1]) + pos_embed = get_2d_sincos_pos_embed( + embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale + ) + self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) + + self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t + self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t + self.interpolation_scale_t = interpolation_scale_t + temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t) + self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) + # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0])) + + def forward(self, latent, num_frames): + b, _, _, _, _ = latent.shape + # b c 1 h w + height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size + latent = rearrange(latent, 'b c t h w -> (b t) c h w') + latent = self.proj(latent) + + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C + if self.layer_norm: + latent = self.norm(latent) + + if self.use_abs_pos: + if self.height != height or self.width != width: + pos_embed = get_2d_sincos_pos_embed( + embed_dim=self.pos_embed.shape[-1], + grid_size=(height, width), + base_size=self.base_size, + interpolation_scale=self.interpolation_scale, + ) + pos_embed = torch.from_numpy(pos_embed) + pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) + else: + pos_embed = self.pos_embed + + if self.num_frames != num_frames: + if get_sequence_parallel_state(): + sp_size = get_sequence_parallel_size() + temp_pos_embed = get_1d_sincos_pos_embed( + embed_dim=self.temp_pos_embed.shape[-1], + grid_size=num_frames * sp_size, + base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t, + ) + rank = get_sequence_parallel_rank() % sp_size + st_frame = rank * num_frames + ed_frame = st_frame + num_frames + temp_pos_embed = temp_pos_embed[st_frame: ed_frame] + else: + temp_pos_embed = get_1d_sincos_pos_embed( + embed_dim=self.temp_pos_embed.shape[-1], + grid_size=num_frames, + base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t, + ) + temp_pos_embed = torch.from_numpy(temp_pos_embed) + temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device) + else: + temp_pos_embed = self.temp_pos_embed + + latent = (latent + pos_embed).to(latent.dtype) + + latent = rearrange(latent, '(b t) n c -> b t n c', b=b) + video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:] + + if self.use_abs_pos: + # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh() + temp_pos_embed = temp_pos_embed.unsqueeze(2) + video_latent = (video_latent + temp_pos_embed).to( + video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None + image_latent = (image_latent + temp_pos_embed[:, :1]).to( + image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None + + video_latent = rearrange(video_latent, + 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None + image_latent = rearrange(image_latent, + 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None + + if num_frames == 1 and image_latent is None and not get_sequence_parallel_state(): + image_latent = video_latent + video_latent = None + return video_latent, image_latent + + +class OverlapPatchEmbed3D(nn.Module): + """2D Image to Patch Embedding but with 3D position embedding""" + + def __init__( + self, + num_frames=1, + height=224, + width=224, + patch_size_t=1, + patch_size=16, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + interpolation_scale=(1, 1), + interpolation_scale_t=1, + use_abs_pos=True, + ): + super().__init__() + self.use_abs_pos = use_abs_pos + self.flatten = flatten + self.layer_norm = layer_norm + + self.proj = nn.Conv3d( + in_channels, embed_dim, kernel_size=(patch_size_t, patch_size, patch_size), + stride=(patch_size_t, patch_size, patch_size), bias=bias + ) + if layer_norm: + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) + else: + self.norm = None + + self.patch_size_t = patch_size_t + self.patch_size = patch_size + + self.height, self.width = height // patch_size, width // patch_size + self.base_size = (height // patch_size, width // patch_size) + self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1]) + pos_embed = get_2d_sincos_pos_embed( + embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale + ) + self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) + + self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t + self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t + self.interpolation_scale_t = interpolation_scale_t + temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t) + self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) + + def forward(self, latent, num_frames): + b, _, _, _, _ = latent.shape + # b c 1 h w + height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size + latent = self.proj(latent) + + if self.flatten: + latent = rearrange(latent, 'b c t h w -> (b t) (h w) c ') + if self.layer_norm: + latent = self.norm(latent) + + if self.use_abs_pos: + if self.height != height or self.width != width: + pos_embed = get_2d_sincos_pos_embed( + embed_dim=self.pos_embed.shape[-1], + grid_size=(height, width), + base_size=self.base_size, + interpolation_scale=self.interpolation_scale, + ) + pos_embed = torch.from_numpy(pos_embed) + pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) + else: + pos_embed = self.pos_embed + + if self.num_frames != num_frames: + temp_pos_embed = get_1d_sincos_pos_embed( + embed_dim=self.temp_pos_embed.shape[-1], + grid_size=num_frames, + base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t, + ) + temp_pos_embed = torch.from_numpy(temp_pos_embed) + temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device) + else: + temp_pos_embed = self.temp_pos_embed + + latent = (latent + pos_embed).to(latent.dtype) + + latent = rearrange(latent, '(b t) n c -> b t n c', b=b) + video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:] + + if self.use_abs_pos: + temp_pos_embed = temp_pos_embed.unsqueeze(2) + video_latent = (video_latent + temp_pos_embed).to( + video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None + image_latent = (image_latent + temp_pos_embed[:, :1]).to( + image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None + + video_latent = rearrange(video_latent, + 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None + image_latent = rearrange(image_latent, + 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None + + if num_frames == 1 and image_latent is None: + image_latent = video_latent + video_latent = None + return video_latent, image_latent + + +class OverlapPatchEmbed2D(nn.Module): + """2D Image to Patch Embedding but with 3D position embedding""" + + def __init__( + self, + num_frames=1, + height=224, + width=224, + patch_size_t=1, + patch_size=16, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + interpolation_scale=(1, 1), + interpolation_scale_t=1, + use_abs_pos=True, + ): + super().__init__() + assert patch_size_t == 1 + self.use_abs_pos = use_abs_pos + 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, patch_size), bias=bias + ) + if layer_norm: + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) + else: + self.norm = None + + self.patch_size_t = patch_size_t + self.patch_size = patch_size + + self.height, self.width = height // patch_size, width // patch_size + self.base_size = (height // patch_size, width // patch_size) + self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1]) + pos_embed = get_2d_sincos_pos_embed( + embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale + ) + self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) + + self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t + self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t + self.interpolation_scale_t = interpolation_scale_t + temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t) + self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) + + def forward(self, latent, num_frames): + b, _, _, _, _ = latent.shape + # b c 1 h w + height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size + latent = rearrange(latent, 'b c t h w -> (b t) c h w') + latent = self.proj(latent) + + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C + if self.layer_norm: + latent = self.norm(latent) + + if self.use_abs_pos: + if self.height != height or self.width != width: + pos_embed = get_2d_sincos_pos_embed( + embed_dim=self.pos_embed.shape[-1], + grid_size=(height, width), + base_size=self.base_size, + interpolation_scale=self.interpolation_scale, + ) + pos_embed = torch.from_numpy(pos_embed) + pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) + else: + pos_embed = self.pos_embed + + if self.num_frames != num_frames: + temp_pos_embed = get_1d_sincos_pos_embed( + embed_dim=self.temp_pos_embed.shape[-1], + grid_size=num_frames, + base_size=self.base_size_t, + interpolation_scale=self.interpolation_scale_t, + ) + temp_pos_embed = torch.from_numpy(temp_pos_embed) + temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device) + else: + temp_pos_embed = self.temp_pos_embed + + latent = (latent + pos_embed).to(latent.dtype) + + latent = rearrange(latent, '(b t) n c -> b t n c', b=b) + video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:] + + if self.use_abs_pos: + # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh() + temp_pos_embed = temp_pos_embed.unsqueeze(2) + video_latent = (video_latent + temp_pos_embed).to( + video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None + image_latent = (image_latent + temp_pos_embed[:, :1]).to( + image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None + + video_latent = rearrange(video_latent, + 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None + image_latent = rearrange(image_latent, + 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None + + if num_frames == 1 and image_latent is None: + image_latent = video_latent + video_latent = None + return video_latent, image_latent + + +class Attention(Attention_): + def __init__(self, downsampler, attention_mode, use_rope, interpolation_scale_thw, **kwags): + processor = AttnProcessor2_0(attention_mode=attention_mode, use_rope=use_rope, + interpolation_scale_thw=interpolation_scale_thw) + super().__init__(processor=processor, **kwags) + self.downsampler = None + if downsampler: # downsampler k155_s122 + downsampler_ker_size = list(re.search(r'k(\d{2,3})', downsampler).group(1)) # 122 + down_factor = list(re.search(r's(\d{2,3})', downsampler).group(1)) + downsampler_ker_size = [int(i) for i in downsampler_ker_size] + downsampler_padding = [(i - 1) // 2 for i in downsampler_ker_size] + down_factor = [int(i) for i in down_factor] + + if len(downsampler_ker_size) == 2: + self.downsampler = DownSampler2d(kwags['query_dim'], kwags['query_dim'], + kernel_size=downsampler_ker_size, stride=1, + padding=downsampler_padding, groups=kwags['query_dim'], + down_factor=down_factor, + down_shortcut=True) + elif len(downsampler_ker_size) == 3: + self.downsampler = DownSampler3d(kwags['query_dim'], kwags['query_dim'], + kernel_size=downsampler_ker_size, stride=1, + padding=downsampler_padding, groups=kwags['query_dim'], + down_factor=down_factor, + down_shortcut=True) + + def prepare_attention_mask( + self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3 + ) -> torch.Tensor: + r""" + Prepare the attention mask for the attention computation. + + Args: + attention_mask (`torch.Tensor`): + The attention mask to prepare. + target_length (`int`): + The target length of the attention mask. This is the length of the attention mask after padding. + batch_size (`int`): + The batch size, which is used to repeat the attention mask. + out_dim (`int`, *optional*, defaults to `3`): + The output dimension of the attention mask. Can be either `3` or `4`. + + Returns: + `torch.Tensor`: The prepared attention mask. + """ + head_size = self.heads + if get_sequence_parallel_state(): + head_size = head_size // get_sequence_parallel_size() + if attention_mask is None: + return None + + current_length: int = attention_mask.shape[-1] + if current_length != target_length: + attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) + if out_dim == 3: + if attention_mask.shape[0] < batch_size * head_size: + attention_mask = attention_mask.repeat_interleave(head_size, dim=0) + elif out_dim == 4: + attention_mask = attention_mask.unsqueeze(1) + attention_mask = attention_mask.repeat_interleave(head_size, dim=1) + + return attention_mask + + +class DownSampler3d(nn.Module): + def __init__(self, *args, **kwargs): + ''' Required kwargs: down_factor, downsampler''' + super().__init__() + self.down_factor = kwargs.pop('down_factor') + self.down_shortcut = kwargs.pop('down_shortcut') + self.layer = nn.Conv3d(*args, **kwargs) + + def forward(self, x, attention_mask, t, h, w): + b = x.shape[0] + x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) + + x_dtype = x.dtype + x = self.layer(x) + (x if self.down_shortcut else 0) + + self.t = t // self.down_factor[0] + self.h = h // self.down_factor[1] + self.w = w // self.down_factor[2] + x = rearrange(x, 'b d (t dt) (h dh) (w dw) -> (b dt dh dw) (t h w) d', + t=t // self.down_factor[0], h=h // self.down_factor[1], w=w // self.down_factor[2], + dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2]) + + attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=t, h=h, w=w) + attention_mask = rearrange(attention_mask, 'b 1 (t dt) (h dh) (w dw) -> (b dt dh dw) 1 (t h w)', + t=t // self.down_factor[0], h=h // self.down_factor[1], w=w // self.down_factor[2], + dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2]) + return x, attention_mask + + def reverse(self, x, t, h, w): + x = rearrange(x, '(b dt dh dw) (t h w) d -> b (t dt h dh w dw) d', + t=t, h=h, w=w, + dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2]) + return x + + +class DownSampler2d(nn.Module): + def __init__(self, *args, **kwargs): + ''' Required kwargs: down_factor, downsampler''' + super().__init__() + self.down_factor = kwargs.pop('down_factor') + self.down_shortcut = kwargs.pop('down_shortcut') + self.layer = nn.Conv2d(*args, **kwargs) + + def forward(self, x, attention_mask, t, h, w): + b = x.shape[0] + x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w) + x = self.layer(x) + (x if self.down_shortcut else 0) + + self.t = 1 + self.h = h // self.down_factor[0] + self.w = w // self.down_factor[1] + + x = rearrange(x, 'b d (h dh) (w dw) -> (b dh dw) (h w) d', + h=h // self.down_factor[0], w=w // self.down_factor[1], + dh=self.down_factor[0], dw=self.down_factor[1]) + + attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> (b t) 1 h w', h=h, w=w) + attention_mask = rearrange(attention_mask, 'b 1 (h dh) (w dw) -> (b dh dw) 1 (h w)', + h=h // self.down_factor[0], w=w // self.down_factor[1], + dh=self.down_factor[0], dw=self.down_factor[1]) + return x, attention_mask + + def reverse(self, x, t, h, w): + x = rearrange(x, '(b t dh dw) (h w) d -> b (t h dh w dw) d', + t=t, h=h, w=w, + dh=self.down_factor[0], dw=self.down_factor[1]) + return x + + +class AttnProcessor2_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). + """ + + def __init__(self, attention_mode='xformers', use_rope=False, interpolation_scale_thw=(1, 1, 1)): + self.use_rope = use_rope + self.interpolation_scale_thw = interpolation_scale_thw + if self.use_rope: + self._init_rope(interpolation_scale_thw) + self.attention_mode = attention_mode + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def _init_rope(self, interpolation_scale_thw): + self.rope = RoPE3D(interpolation_scale_thw=interpolation_scale_thw) + self.position_getter = PositionGetter3D() + + def __call__( + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + frame: int = 8, + height: int = 16, + width: int = 16, + *args, + **kwargs, + ) -> torch.FloatTensor: + if len(args) > 0 or kwargs.get("scale", None) is not None: + deprecation_message = ("The `scale` argument is deprecated and will be ignored. Please remove it, " + "as passing it will raise an error in the future. `scale` should directly be " + "passed while calling the underlying pipeline component i.e., " + "via `cross_attention_kwargs`.") + deprecate("scale", "1.0.0", deprecation_message) + + if attn.downsampler is not None: + hidden_states, attention_mask = attn.downsampler(hidden_states, attention_mask, t=frame, h=height, w=width) + frame, height, width = attn.downsampler.t, attn.downsampler.h, attn.downsampler.w + + residual = hidden_states + + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb) + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + if get_sequence_parallel_state(): + sequence_length, batch_size, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + else: + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1]) + + if attn.group_norm is not None: + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + if get_sequence_parallel_state(): + query = query.view(-1, attn.heads, head_dim) # [s // sp, b, h * d] -> [s // sp * b, h, d] + key = key.view(-1, attn.heads, head_dim) + value = value.view(-1, attn.heads, head_dim) + h_size = attn.heads * head_dim + sp_size = get_sequence_parallel_size() + h_size_sp = h_size // sp_size + # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d] + query = all_to_all_sbh(query, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp) + key = all_to_all_sbh(key, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp) + value = all_to_all_sbh(value, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp) + if self.use_rope: + query = query.view(-1, batch_size, attn.heads // sp_size, head_dim) + key = key.view(-1, batch_size, attn.heads // sp_size, head_dim) + # require the shape of (batch_size x nheads x ntokens x dim) + pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, + device=query.device) + query = self.rope(query, pos_thw) + key = self.rope(key, pos_thw) + query = query.view(-1, batch_size, h_size_sp) # SBH + key = key.view(-1, batch_size, h_size_sp) # SBH + value = value.view(-1, batch_size, h_size_sp) # SBH + query = rearrange(query, "s b (n d) -> b n s d", d=head_dim) # BNSD + key = rearrange(key, "s b (n d) -> b n s d", d=head_dim) # BNSD + value = rearrange(value, "s b (n d) -> b n s d", d=head_dim) # BNSD + hidden_states = torch_npu.npu_fusion_attention(query, key, value, + atten_mask=attention_mask, + input_layout="BNSD", + scale=1 / math.sqrt(head_dim), + head_num=attn.heads // sp_size)[0] + hidden_states = rearrange(hidden_states, "b n s d -> s b (n d)") + hidden_states = hidden_states.view(-1, attn.heads // sp_size, head_dim) + + # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d] + hidden_states = all_to_all_sbh(hidden_states, scatter_dim=0, gather_dim=1).view(-1, batch_size, h_size) + else: + query = query.view(batch_size, -1, attn.heads, head_dim) + key = key.view(batch_size, -1, attn.heads, head_dim) + value = value.view(batch_size, -1, attn.heads, head_dim) + if self.use_rope: + # require the shape of (batch_size x nheads x ntokens x dim) + pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) + query = self.rope(query, pos_thw) + key = self.rope(key, pos_thw) + query = rearrange(query, "b s n d -> b n s d") # BNSD + key = rearrange(key, "b s n d -> b n s d") # BNSD + value = rearrange(value, "b s n d -> b n s d") # BNSD + hidden_states = torch_npu.npu_fusion_attention(query, key, value, + atten_mask=attention_mask, + input_layout="BNSD", + scale=1 / math.sqrt(head_dim), + head_num=attn.heads)[0] + hidden_states = rearrange(hidden_states, "b n s d -> b s (n d)") + + # linear proj + hidden_states = attn.to_out[0](hidden_states) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + if input_ndim == 4: + hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + if attn.downsampler is not None: + hidden_states = attn.downsampler.reverse(hidden_states, t=frame, h=height, w=width) + return hidden_states + + +class FeedForward_Conv3d(nn.Module): + def __init__(self, downsampler, dim, hidden_features, bias=True): + super(FeedForward_Conv3d, self).__init__() + + self.bias = bias + + self.project_in = nn.Linear(dim, hidden_features, bias=bias) + + self.dwconv = nn.ModuleList([ + nn.Conv3d(hidden_features, hidden_features, kernel_size=(5, 5, 5), stride=1, padding=(2, 2, 2), dilation=1, + groups=hidden_features, bias=bias), + nn.Conv3d(hidden_features, hidden_features, kernel_size=(3, 3, 3), stride=1, padding=(1, 1, 1), dilation=1, + groups=hidden_features, bias=bias), + nn.Conv3d(hidden_features, hidden_features, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0), dilation=1, + groups=hidden_features, bias=bias) + ]) + + self.project_out = nn.Linear(hidden_features, dim, bias=bias) + + def forward(self, x, t, h, w): + x_dtype = x.dtype + x = self.project_in(x) + x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) + x = F.gelu(x) + out = x + for module in self.dwconv: + out = out + module(x) + out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w) + x = self.project_out(out) + return x + + +class FeedForward_Conv2d(nn.Module): + def __init__(self, downsampler, dim, hidden_features, bias=True): + super(FeedForward_Conv2d, self).__init__() + + self.bias = bias + + self.project_in = nn.Linear(dim, hidden_features, bias=bias) + + self.dwconv = nn.ModuleList([ + nn.Conv2d(hidden_features, hidden_features, kernel_size=(5, 5), stride=1, padding=(2, 2), dilation=1, + groups=hidden_features, bias=bias), + nn.Conv2d(hidden_features, hidden_features, kernel_size=(3, 3), stride=1, padding=(1, 1), dilation=1, + groups=hidden_features, bias=bias), + nn.Conv2d(hidden_features, hidden_features, kernel_size=(1, 1), stride=1, padding=(0, 0), dilation=1, + groups=hidden_features, bias=bias) + ]) + + self.project_out = nn.Linear(hidden_features, dim, bias=bias) + + def forward(self, x, t, h, w): + # import ipdb;ipdb.set_trace() + x = self.project_in(x) + x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w) + x = F.gelu(x) + out = x + for module in self.dwconv: + out = out + module(x) + out = rearrange(out, '(b t) d h w -> b (t h w) d', t=t, h=h, w=w) + x = self.project_out(out) + return x + + +@maybe_allow_in_graph +class BasicTransformerBlock(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, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", + # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + ada_norm_continous_conditioning_embedding_dim: Optional[int] = None, + ada_norm_bias: Optional[int] = None, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = True, + attention_out_bias: bool = True, + attention_mode: str = "xformers", + downsampler: str = None, + use_rope: bool = False, + interpolation_scale_thw: Tuple[int] = (1, 1, 1), + ): + super().__init__() + self.only_cross_attention = only_cross_attention + self.downsampler = downsampler + + # We keep these boolean flags for backward-compatibility. + self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" + self.use_ada_layer_norm_single = norm_type == "ada_norm_single" + self.use_layer_norm = norm_type == "layer_norm" + self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + self.norm_type = norm_type + self.num_embeds_ada_norm = num_embeds_ada_norm + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if norm_type == "ada_norm": + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif norm_type == "ada_norm_zero": + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + elif norm_type == "ada_norm_continuous": + self.norm1 = AdaLayerNormContinuous( + dim, + ada_norm_continous_conditioning_embedding_dim, + norm_elementwise_affine, + norm_eps, + ada_norm_bias, + "rms_norm", + ) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + out_bias=attention_out_bias, + attention_mode=attention_mode, + downsampler=downsampler, + use_rope=use_rope, + interpolation_scale_thw=interpolation_scale_thw, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + if norm_type == "ada_norm": + self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif norm_type == "ada_norm_continuous": + self.norm2 = AdaLayerNormContinuous( + dim, + ada_norm_continous_conditioning_embedding_dim, + norm_elementwise_affine, + norm_eps, + ada_norm_bias, + "rms_norm", + ) + else: + self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) + + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + out_bias=attention_out_bias, + attention_mode=attention_mode, + downsampler=False, + use_rope=False, + interpolation_scale_thw=interpolation_scale_thw, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + if norm_type == "ada_norm_continuous": + self.norm3 = AdaLayerNormContinuous( + dim, + ada_norm_continous_conditioning_embedding_dim, + norm_elementwise_affine, + norm_eps, + ada_norm_bias, + "layer_norm", + ) + + elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm", "ada_norm_continuous"]: + self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) + elif norm_type == "layer_norm_i2vgen": + self.norm3 = None + + if downsampler: + self.ff = FeedForward_Conv2d( + downsampler, + dim, + 2 * dim, + bias=ff_bias, + ) + else: + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + inner_dim=ff_inner_dim, + bias=ff_bias, + ) + + # 4. Fuser + if attention_type == "gated" or attention_type == "gated-text-image": + self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) + + # 5. Scale-shift for PixArt-Alpha. + if norm_type == "ada_norm_single": + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim ** 0.5) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + frame: int = None, + height: int = None, + width: int = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + ) -> torch.FloatTensor: + if cross_attention_kwargs is not None: + if cross_attention_kwargs.get("scale", None) is not None: + logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") + + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + batch_size = hidden_states.shape[0] + + if self.norm_type == "ada_norm": + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.norm_type == "ada_norm_zero": + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]: + norm_hidden_states = self.norm1(hidden_states) + elif self.norm_type == "ada_norm_continuous": + norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"]) + elif self.norm_type == "ada_norm_single": + if get_sequence_parallel_state(): + batch_size = hidden_states.shape[1] + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[:, None] + timestep.reshape(6, batch_size, -1) + ).chunk(6, dim=0) + else: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + else: + raise ValueError("Incorrect norm used") + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + # 1. Prepare GLIGEN inputs + cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, frame=frame, height=height, width=width, + **cross_attention_kwargs, + ) + if self.norm_type == "ada_norm_zero": + attn_output = gate_msa.unsqueeze(1) * attn_output + elif self.norm_type == "ada_norm_single": + attn_output = gate_msa * attn_output + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 1.2 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + + # 3. Cross-Attention + if self.attn2 is not None: + if self.norm_type == "ada_norm": + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]: + norm_hidden_states = self.norm2(hidden_states) + elif self.norm_type == "ada_norm_single": + norm_hidden_states = hidden_states + elif self.norm_type == "ada_norm_continuous": + norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"]) + else: + raise ValueError("Incorrect norm") + + if self.pos_embed is not None and self.norm_type != "ada_norm_single": + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + if self.norm_type == "ada_norm_continuous": + norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"]) + elif not self.norm_type == "ada_norm_single": + norm_hidden_states = self.norm3(hidden_states) + + if self.norm_type == "ada_norm_zero": + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self.norm_type == "ada_norm_single": + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + if self.downsampler: + ff_output = self.ff(norm_hidden_states, t=frame, h=height, w=width) + else: + ff_output = self.ff(norm_hidden_states) + + if self.norm_type == "ada_norm_zero": + ff_output = gate_mlp.unsqueeze(1) * ff_output + elif self.norm_type == "ada_norm_single": + ff_output = gate_mlp * ff_output + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + return hidden_states + + +class LayerNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.hidden_size = hidden_size + self.eps = eps + + def forward(self, x, shift, scale): + return torch_npu.npu_layer_norm_eval( + x, normalized_shape=[self.hidden_size], weight=scale, bias=shift, eps=self.eps + ) diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/rope.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/rope.py new file mode 100644 index 0000000000000000000000000000000000000000..e9d0c7c4eb85e5e1b2ee5bf384d4a12c82c79fec --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/models/diffusion/opensora/rope.py @@ -0,0 +1,88 @@ +import torch +import torch_npu + +from utils.parallel_mgr import get_sequence_parallel_state + + +class PositionGetter3D(object): + """ return positions of patches """ + + def __init__(self, ): + self.cache_positions = {} + + def __call__(self, b, t, h, w, device): + if not (b, t, h, w) in self.cache_positions: + x = torch.arange(w, device=device) + y = torch.arange(h, device=device) + z = torch.arange(t, device=device) + pos = torch.cartesian_prod(z, y, x) + if get_sequence_parallel_state(): + pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, -1, 1).contiguous().expand(3, -1, b).clone() + else: + pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, 1, -1).contiguous().expand(3, b, -1).clone() + poses = (pos[0].contiguous(), pos[1].contiguous(), pos[2].contiguous()) + max_poses = (int(poses[0].max()), int(poses[1].max()), int(poses[2].max())) + + self.cache_positions[b, t, h, w] = (poses, max_poses) + pos = self.cache_positions[b, t, h, w] + + return pos + + +class RoPE3D(torch.nn.Module): + + def __init__(self, freq=10000.0, F0=1.0, interpolation_scale_thw=(1, 1, 1)): + super().__init__() + self.base = freq + self.F0 = F0 + self.interpolation_scale_t = interpolation_scale_thw[0] + self.interpolation_scale_h = interpolation_scale_thw[1] + self.interpolation_scale_w = interpolation_scale_thw[2] + self.cache = {} + + def get_cos_sin(self, D, seq_len, device, dtype, interpolation_scale=1): + 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) / interpolation_scale + 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 + # for (batch_size x ntokens x nheads x dim) + cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :] + sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :] + + return torch_npu.npu_rotary_mul(tokens, cos, sin) + + def forward(self, tokens, positions): + """ + input: + * tokens: batch_size x nheads x ntokens x dim + * positions: batch_size x ntokens x 3 (t, y and x position of each token) + output: + * tokens after appplying RoPE3D (batch_size x nheads x ntokens x x dim) + """ + assert tokens.size(3) % 3 == 0, "number of dimensions should be a multiple of three" + D = tokens.size(3) // 3 + poses, max_poses = positions + assert len(poses) == 3 and poses[0].ndim == 2 # Batch, Seq, 3 + cos_t, sin_t = self.get_cos_sin(D, max_poses[0] + 1, tokens.device, tokens.dtype, self.interpolation_scale_t) + cos_y, sin_y = self.get_cos_sin(D, max_poses[1] + 1, tokens.device, tokens.dtype, self.interpolation_scale_h) + cos_x, sin_x = self.get_cos_sin(D, max_poses[2] + 1, tokens.device, tokens.dtype, self.interpolation_scale_w) + # split features into three along the feature dimension, and apply rope1d on each half + t, y, x = tokens.chunk(3, dim=-1) + t = self.apply_rope1d(t, poses[0], cos_t, sin_t) + y = self.apply_rope1d(y, poses[1], cos_y, sin_y) + x = self.apply_rope1d(x, poses[2], cos_x, sin_x) + tokens = torch.cat((t, y, x), dim=-1) + return tokens diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/sample/pipeline_opensora_sp.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/sample/pipeline_opensora_sp.py new file mode 100644 index 0000000000000000000000000000000000000000..b143cc371dc1a37183cbb77dd0c97e6a645969fb --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/opensora/sample/pipeline_opensora_sp.py @@ -0,0 +1,793 @@ +import html +import inspect +import math +from tqdm import tqdm +import re +import urllib.parse as ul +from typing import Callable, List, Optional, Tuple, Union + +import torch +from diffusers.models import AutoencoderKL, Transformer2DModel +from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput +from diffusers.schedulers import EulerAncestralDiscreteScheduler +from diffusers.utils import ( + BACKENDS_MAPPING, + is_bs4_available, + is_ftfy_available, + logging, +) +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from mindiesd.pipeline.sampling_optm import SamplingOptm +from transformers import T5EncoderModel, T5Tokenizer + +from utils.parallel_mgr import get_sequence_parallel_state, get_sequence_parallel_size, get_sequence_parallel_rank + +logger = logging.get_logger(__name__) + +if is_bs4_available(): + from bs4 import BeautifulSoup + +if is_ftfy_available(): + import ftfy + + +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + **kwargs, +): + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default + timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps` + must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class OpenSoraPipeline(DiffusionPipeline): + r""" + Pipeline for text-to-image generation using PixArt-Sigma. + """ + + bad_punct_regex = re.compile( + r"[" + + "#®•©™&@·º½¾¿¡§~" + + r"\)" + + r"\(" + + r"\]" + + r"\[" + + r"\}" + + r"\{" + + r"\|" + + "\\" + + r"\/" + + r"\*" + + r"]{1,}" + ) # noqa + + _optional_components = ["tokenizer", "text_encoder"] + model_cpu_offload_seq = "text_encoder->transformer->vae" + + def __init__( + self, + tokenizer: T5Tokenizer, + text_encoder: T5EncoderModel, + vae: AutoencoderKL, + transformer: Transformer2DModel, + scheduler: EulerAncestralDiscreteScheduler, + ): + super().__init__() + + self.register_modules( + tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler + ) + + def encode_prompt( + self, + prompt: Union[str, List[str]], + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + num_images_per_prompt: int = 1, + device: Optional[torch.device] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_attention_mask: Optional[torch.FloatTensor] = None, + negative_prompt_attention_mask: Optional[torch.FloatTensor] = None, + clean_caption: bool = False, + max_sequence_length: int = 120, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `List[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `List[str]`, *optional*): + The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` + instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For + PixArt-Alpha, this should be "". + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + whether to use classifier free guidance or not + num_images_per_prompt (`int`, *optional*, defaults to 1): + number of images that should be generated per prompt + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the "" + string. + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt. + """ + + if device is None: + device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('npu') + + 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] + + max_length = max_sequence_length + + if prompt_embeds is None: + prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal( + text_input_ids, untruncated_ids + ): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {max_length} tokens: {removed_text}" + ) + + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.to(device) + + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) + prompt_embeds = prompt_embeds[0] + + if self.text_encoder is not None: + dtype = self.text_encoder.dtype + elif self.transformer is not None: + dtype = self.transformer.dtype + else: + dtype = None + + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + uncond_tokens = [negative_prompt] * batch_size + uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption) + max_length = prompt_embeds.shape[1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_attention_mask=True, + add_special_tokens=True, + return_tensors="pt", + ) + negative_prompt_attention_mask = uncond_input.attention_mask + negative_prompt_attention_mask = negative_prompt_attention_mask.to(device) + + negative_prompt_embeds = self.text_encoder( + uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask + ) + negative_prompt_embeds = negative_prompt_embeds[0] + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1) + else: + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + num_frames, + height, + width, + negative_prompt, + callback_steps, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if num_frames <= 0: + raise ValueError(f"`num_frames` have to be positive but is {num_frames}.") + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + def _text_preprocessing(self, text, clean_caption=False): + if clean_caption and not is_bs4_available(): + logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`")) + logger.warning("Setting `clean_caption` to False...") + clean_caption = False + + if clean_caption and not is_ftfy_available(): + logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`")) + logger.warning("Setting `clean_caption` to False...") + clean_caption = False + + if not isinstance(text, (tuple, list)): + text = [text] + + def process(text: str): + if clean_caption: + text = self._clean_caption(text) + text = self._clean_caption(text) + else: + text = text.lower().strip() + return text + + return [process(t) for t in text] + + # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption + def _clean_caption(self, caption): + caption = str(caption) + caption = ul.unquote_plus(caption) + caption = caption.strip().lower() + caption = re.sub("", "person", caption) + # urls: + caption = re.sub( + r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", + # noqa + "", + caption, + ) # regex for urls + caption = re.sub( + r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", + # noqa + "", + caption, + ) # regex for urls + # html: + caption = BeautifulSoup(caption, features="html.parser").text + + # @ + caption = re.sub(r"@[\w\d]+\b", "", caption) + + # 31C0—31EF CJK Strokes + # 31F0—31FF Katakana Phonetic Extensions + # 3200—32FF Enclosed CJK Letters and Months + # 3300—33FF CJK Compatibility + # 3400—4DBF CJK Unified Ideographs Extension A + # 4DC0—4DFF Yijing Hexagram Symbols + # 4E00—9FFF CJK Unified Ideographs + caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) + caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) + caption = re.sub(r"[\u3200-\u32ff]+", "", caption) + caption = re.sub(r"[\u3300-\u33ff]+", "", caption) + caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) + caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) + # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) + ####################################################### + + # все виды тире / all types of dash --> "-" + caption = re.sub( + r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", + # noqa + "-", + caption, + ) + + # кавычки к одному стандарту + caption = re.sub(r"[`´«»“”¨]", '"', caption) + caption = re.sub(r"[‘’]", "'", caption) + + # " + caption = re.sub(r""?", "", caption) + # & + caption = re.sub(r"&", "", caption) + + # ip adresses: + caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) + + # article ids: + caption = re.sub(r"\d:\d\d\s+$", "", caption) + + # \n + caption = re.sub(r"\\n", " ", caption) + + # "#123" + caption = re.sub(r"#\d{1,3}\b", "", caption) + # "#12345.." + caption = re.sub(r"#\d{5,}\b", "", caption) + # "123456.." + caption = re.sub(r"\b\d{6,}\b", "", caption) + # filenames: + caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption) + + # + caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" + caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" + + caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT + caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " + + # this-is-my-cute-cat / this_is_my_cute_cat + regex2 = re.compile(r"(?:\-|\_)") + if len(re.findall(regex2, caption)) > 3: + caption = re.sub(regex2, " ", caption) + + caption = ftfy.fix_text(caption) + caption = html.unescape(html.unescape(caption)) + + caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 + caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc + caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 + + caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) + caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) + caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) + caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption) + caption = re.sub(r"\bpage\s+\d+\b", "", caption) + + caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a... + + caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) + + caption = re.sub(r"\b\s+\:\s+", r": ", caption) + caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) + caption = re.sub(r"\s+", " ", caption) + + caption.strip() + + caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) + caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) + caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) + caption = re.sub(r"^\.\S+$", "", caption) + return caption.strip() + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents + def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, + latents=None): + shape = ( + batch_size, + num_channels_latents, + (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int( + num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), + math.ceil(int(height) / self.vae.vae_scale_factor[1]), + math.ceil(int(width) / self.vae.vae_scale_factor[2]), + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + + return latents + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + negative_prompt: str = "", + num_inference_steps: int = 20, + timesteps: List[int] = None, + guidance_scale: float = 4.5, + num_images_per_prompt: Optional[int] = 1, + num_frames: Optional[int] = None, + height: Optional[int] = None, + width: Optional[int] = None, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + prompt_attention_mask: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_attention_mask: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + clean_caption: bool = True, + use_resolution_binning: bool = True, + max_sequence_length: int = 300, + **kwargs, + ) -> Union[ImagePipelineOutput, Tuple]: + """ + Function invoked when calling 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. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_inference_steps (`int`, *optional*, defaults to 100): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` + timesteps are used. Must be in descending order. + guidance_scale (`float`, *optional*, defaults to 4.5): + 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_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + height (`int`, *optional*, defaults to self.unet.config.sample_size): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to self.unet.config.sample_size): + The width in pixels of the generated image. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *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 will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be + called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + clean_caption (`bool`, *optional*, defaults to `True`): + Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to + be installed. If the dependencies are not installed, the embeddings will be created from the raw + prompt. + use_resolution_binning (`bool` defaults to `True`): + If set to `True`, the requested height and width are first mapped to the closest resolutions using + `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to + the requested resolution. Useful for generating non-square images. + max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.ImagePipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list with the generated images + """ + # 1. Check inputs. Raise error if not correct + num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0] + height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1] + width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2] + self.check_inputs( + prompt, + num_frames, + height, + width, + negative_prompt, + callback_steps, + prompt_embeds, + negative_prompt_embeds, + prompt_attention_mask, + negative_prompt_attention_mask, + ) + + # 2. Default height and width to transformer + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('npu') + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt, + do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + ) + if do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + + # 5. Prepare latents. + latent_channels = self.transformer.config.in_channels + world_size = get_sequence_parallel_size() + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + latent_channels, + (num_frames + world_size - 1) // world_size if get_sequence_parallel_state() else num_frames, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + if get_sequence_parallel_state(): + prompt_embeds = rearrange(prompt_embeds, 'b (n x) h -> b n x h', n=world_size, + x=prompt_embeds.shape[1] // world_size).contiguous() + rank = get_sequence_parallel_rank() + prompt_embeds = prompt_embeds[:, rank, :, :] + + # 6. Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 6.1 Prepare micro-conditions. + added_cond_kwargs = {"resolution": None, "aspect_ratio": None} + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + def dit_process(latents, prompt_embeds, prompt_attention_mask): + for i, t in enumerate(tqdm(timesteps)): + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + current_timestep = torch.tensor([t], dtype=torch.float64, device=latent_model_input.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + current_timestep = current_timestep.expand(latent_model_input.shape[0]) + + if prompt_embeds.ndim == 3: + prompt_embeds = prompt_embeds.unsqueeze(1) # b l d -> b 1 l d + if prompt_attention_mask.ndim == 2: + prompt_attention_mask = prompt_attention_mask.unsqueeze(1) # b l -> b 1 l + + # prepare attention_mask. + # b c t h w -> b t h w + attention_mask = torch.ones_like(latent_model_input)[:, 0] + if get_sequence_parallel_state(): + attention_mask = attention_mask.repeat(1, world_size, 1, 1) + # predict noise model_output + noise_pred = self.transformer( + latent_model_input, + attention_mask=attention_mask, + encoder_hidden_states=prompt_embeds, + encoder_attention_mask=prompt_attention_mask, + timestep=current_timestep, + step_id=i, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0] + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # learned sigma + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + else: + noise_pred = noise_pred + + # compute previous image: x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + # call the callback, if provided + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + return latents + + if kwargs.get("sampling_optimize", None): + config = {"method": "AdaStep", "skip_thr": 0.015, "max_skip_steps": 4, "decay_ratio": 0.99} + with SamplingOptm(self, + "transformer.forward", + "scheduler.step", + parallel=get_sequence_parallel_state(), + config=config): + latents = dit_process(latents, prompt_embeds, prompt_attention_mask) + else: + latents = dit_process(latents, prompt_embeds, prompt_attention_mask) + + world_size = get_sequence_parallel_size() + if get_sequence_parallel_state(): + latents_shape = list(latents.shape) + full_shape = [latents_shape[0] * world_size] + latents_shape[1:] + all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device) + torch.distributed.all_gather_into_tensor(all_latents, latents) + latents_list = list(all_latents.chunk(world_size, dim=0)) + latents = torch.cat(latents_list, dim=2) + + if not output_type == "latent": + # b t h w c + image = self.decode_latents(latents) + image = image[:, :num_frames, :height, :width] + else: + image = latents + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return ImagePipelineOutput(images=image) + + def decode_latents(self, latents): + video = self.vae.decode(latents.to(self.vae.vae.dtype)) + video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, + 2).contiguous() # b t h w c + return video diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/requirements.txt b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..44097a4d0ea091d5cdc539e92ce3700d8916abc7 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/requirements.txt @@ -0,0 +1,17 @@ +torch==2.1.0 +diffusers==0.28.0 +transformers==4.44.2 +open_clip_torch==2.20.0 +av==12.0.0 +tqdm==4.66.1 +tensorboard==2.11.0 +pre-commit==3.8.0 +mmengine==0.10.4 +ftfy==6.1.3 +accelerate==0.26.1 +bs4 +torchvision==0.16.0 +einops +numpy==1.24.0 +imageio +imageio-ffmpeg diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/run.sh b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..50fb03c198b22d6c8c559779ba1408db98e85a22 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/run.sh @@ -0,0 +1,18 @@ +torchrun --nnodes=1 --nproc_per_node 8 --master_port 29503 \ + inference.py \ + --model_path /path/to/checkpoint-xxx/model_ema \ + --num_frames 93 \ + --height 720 \ + --width 1280 \ + --cache_dir "../cache_dir" \ + --text_encoder_name google/mt5-xxl \ + --text_prompt examples/prompt_list_0.txt \ + --ae CausalVAEModel_D4_4x8x8 \ + --ae_path "/path/to/causalvideovae" \ + --save_img_path "./sample_video_test" \ + --fps 24 \ + --guidance_scale 7.5 \ + --num_sampling_steps 100 \ + --tile_overlap_factor 0.125 \ + --max_sequence_length 512 \ + --dtype bf16 \ No newline at end of file diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/utils/comm.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/utils/comm.py new file mode 100644 index 0000000000000000000000000000000000000000..bac2b5244eefb1ae91d8659431d55e211ca94fe1 --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/utils/comm.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.distributed as dist + +from .parallel_mgr import get_sequence_parallel_size, get_sequence_parallel_group + + +def _all_to_all_func(input_, world_size, process_group, scatter_dim=2, gather_dim=1): + input_list = [t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)] + dist.all_to_all(output_list, input_list, group=process_group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +def split_sequence(input_, process_group: dist.ProcessGroup, dim: int, pad: int): + world_size = dist.get_world_size(process_group) + rank = dist.get_rank(process_group) + if world_size == 1: + return input_ + + if pad > 0: + pad_size = list(input_.shape) + pad_size[dim] = pad + input_ = torch.cat([input_, torch.zeros(pad_size, dtype=input_.dtype, device=input_.device)], dim=dim) + + dim_size = input_.size(dim) + if dim_size % world_size != 0: + raise ValueError( + f"The th{dim} dimensions of input_:{input_.size()} is not divisible by world_size:{world_size}.") + + tensor_list = torch.split(input_, dim_size // world_size, dim=dim) + output = tensor_list[rank].contiguous() + return output + + +def gather_sequence(input_, process_group: dist.ProcessGroup, dim: int, pad: int): + input_ = input_.contiguous() + world_size = dist.get_world_size(process_group) + if world_size == 1: + return input_ + + # all gather + tensor_list = [torch.empty_like(input_) for _ in range(world_size)] + torch.distributed.all_gather(tensor_list, input_, group=process_group) + + # concat + output = torch.cat(tensor_list, dim=dim) + + if pad > 0: + output = output.narrow(dim, 0, output.size(dim) - pad) + + return output + + +# ====== +# Pad +# ====== + +SPTIAL_PAD = 0 +TEMPORAL_PAD = 0 + + +def set_spatial_pad(dim_size: int): + sp_size = get_sequence_parallel_size() + pad = (sp_size - (dim_size % sp_size)) % sp_size + global SPTIAL_PAD + SPTIAL_PAD = pad + + +def get_spatial_pad() -> int: + return SPTIAL_PAD + + +def set_temporal_pad(dim_size: int): + sp_size = get_sequence_parallel_size() + pad = (sp_size - (dim_size % sp_size)) % sp_size + global TEMPORAL_PAD + TEMPORAL_PAD = pad + + +def get_temporal_pad() -> int: + return TEMPORAL_PAD + + +def all_to_all_with_pad( + input_: torch.Tensor, + process_group: dist.ProcessGroup, + **kwargs +): + scatter_dim = kwargs.get("scatter_dim", 2) + gather_dim = kwargs.get("gather_dim", 1) + scatter_pad = kwargs.get("scatter_pad", 0) + gather_pad = kwargs.get("gather_pad", 0) + + if scatter_pad > 0: + pad_shape = list(input_.shape) + pad_shape[scatter_dim] = scatter_pad + pad_tensor = torch.zeros(pad_shape, device=input_.device, dtype=input_.dtype) + input_ = torch.cat([input_, pad_tensor], dim=scatter_dim) + + world_size = dist.get_world_size(process_group) + if input_.shape[scatter_dim] % world_size != 0: + raise ValueError( + f"The scatter_dim:{scatter_dim} of input_:{input_.shape} is not divisible by world_size:{world_size}.") + + input_ = _all_to_all_func(input_, world_size, process_group, scatter_dim, gather_dim) + + if gather_pad > 0: + input_ = input_.narrow(gather_dim, 0, input_.size(gather_dim) - gather_pad) + + return input_ + + +def all_to_all( + tensor: torch.Tensor, + world_size: int, + scatter_dim: int, + gather_dim: int, + process_group: dist.ProcessGroup = None, +): + if process_group is None: + process_group = dist.group.WORLD + return _all_to_all_func(tensor, world_size, process_group, scatter_dim, gather_dim) + + +def all_to_all_sbh( + input_: torch.Tensor, + scatter_dim: int = 1, + gather_dim: int = 0, +): + return single_all_to_all(input_, scatter_dim, gather_dim) + + +def single_all_to_all( + input_: torch.Tensor, + scatter_dim: int, + gather_dim: int, +): + sp_size = get_sequence_parallel_size() + inp_shape = list(input_.shape) + inp_shape[scatter_dim] = inp_shape[scatter_dim] // sp_size + if scatter_dim < 1: + input_t = input_.reshape( + [sp_size, inp_shape[scatter_dim]] + \ + inp_shape[scatter_dim + 1:] + ) + else: + # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them! + input_t = input_.reshape( + [-1, sp_size, inp_shape[scatter_dim]] + \ + inp_shape[scatter_dim + 1:] + ).transpose(0, 1).contiguous() + + output = torch.empty_like(input_t) + + dist.all_to_all_single(output, input_t, group=get_sequence_parallel_group()) + # if scattering the seq-dim, transpose the heads back to the original dimension + if scatter_dim < 1: + output = output.transpose(0, 1).contiguous() + + return output.reshape( + inp_shape[: gather_dim] + [inp_shape[gather_dim] * sp_size, ] + inp_shape[gather_dim + 1:]) diff --git a/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/utils/parallel_mgr.py b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/utils/parallel_mgr.py new file mode 100644 index 0000000000000000000000000000000000000000..817743fe757ff5c97778603477ebd6520f94414f --- /dev/null +++ b/MindIE/MindIE-Torch/built-in/foundation/open_sora_planv1_2/utils/parallel_mgr.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2024 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import os + +import torch.distributed as dist +import torch_npu + + +class ParallelManager(): + def __init__(self, world_size=1, rank=0, group=None): + self.sp_size = world_size + self.sp_group = group + self.enable_sp = world_size > 1 + self.rank = rank + + +PARALLEL_MANAGER = ParallelManager() + + +def set_parallel_manager(world_size, rank, group): + global PARALLEL_MANAGER + PARALLEL_MANAGER = ParallelManager(world_size, rank, group) + + +def get_sequence_parallel_group(): + return PARALLEL_MANAGER.sp_group + + +def get_sequence_parallel_size(): + return PARALLEL_MANAGER.sp_size + + +def get_sequence_parallel_state(): + return PARALLEL_MANAGER.enable_sp + + +def get_sequence_parallel_rank(): + return PARALLEL_MANAGER.rank + + +def init_parallel_env(enable_sequence_parallelism): + rank = int(os.getenv('RANK', 0)) + world_size = int(os.getenv('WORLD_SIZE', 1)) + torch_npu.npu.set_device(rank) + dist.init_process_group( + backend='hccl', init_method='env://', + world_size=world_size, rank=rank + ) + if enable_sequence_parallelism: + set_parallel_manager(world_size, rank, dist.group.WORLD) + + +def finalize_parallel_env(): + if get_sequence_parallel_state and get_sequence_parallel_size() > 1: + dist.destroy_process_group()