diff --git a/README.md b/README.md index 30562c9b9..7d96030d9 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ Project stages: - [x] Support Dynamic mask input inspired [FiT](https://github.com/whlzy/FiT). - [x] Add class-conditioning on embeddings. - [ ] Incorporating [Latte](https://github.com/Vchitect/Latte) as main codebase. -- [ ] Add VAE model, which is borrowed from [Stable Diffusion](https://github.com/CompVis/latent-diffusion). +- [x] Add VAE model, which is borrowed from [Stable Diffusion](https://github.com/CompVis/latent-diffusion). - [ ] Joint dynamic mask input with VAE. - [ ] Make the codebase ready for the cluster training. Add SLURM scripts. - [ ] Add sampling script. @@ -67,7 +67,7 @@ Project stages: - [ ] Add flash-attention2 from pytorch. - [ ] Add xformers. - [ ] Add accelerate to automatically manage training, e.g. mixed precision training. - - [ ] Add gradient checkpoint. + - [x] Add gradient checkpoint. - [ ] Train using the deepspeed engine. #### Control model with more condition @@ -114,7 +114,7 @@ conda create -n opensora python=3.8 -y conda activate opensora pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117 pip install -r requirements.txt -cd src/sora/modules/ae/vqvae/videogpt/ +cd VideoGPT pip install -e . cd .. ``` diff --git a/opensora/dataset/__init__.py b/opensora/dataset/__init__.py new file mode 100644 index 000000000..565ff71d3 --- /dev/null +++ b/opensora/dataset/__init__.py @@ -0,0 +1,38 @@ +from torchvision.transforms import Compose +from torchvision.transforms._transforms_video import ToTensorVideo, RandomHorizontalFlipVideo + +from opensora.models.ae import vae, vqvae, videovae, videovqvae +from .feature_datasets import LandscopeFeatures +from torchvision import transforms +from torchvision.transforms import Lambda + +from .transform import LongSideScale, TemporalRandomCrop +from .ucf101 import UCF101 + + +def getdataset(args): + temporal_sample = TemporalRandomCrop(args.num_frames * args.sample_rate) # 16 x + if args.ae in videovqvae: + norm_fun = Lambda(lambda x: x - 0.5) + elif args.ae in videovae: + raise NotImplementedError + elif args.ae in vae: + norm_fun = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True) + elif args.ae in vqvae: + norm_fun = Lambda(lambda x: 2. * x - 1.) + else: + raise NotImplementedError + + if args.dataset == 'landscope_feature': + temporal_sample = TemporalRandomCrop(args.num_frames) # 16 1 + return LandscopeFeatures(args, temporal_sample=temporal_sample) + elif args.dataset == 'ucf101': + transform = Compose( + [ + ToTensorVideo(), # cthw + LongSideScale(size=args.max_image_size), + RandomHorizontalFlipVideo(p=0.5), + norm_fun, + ] + ) + return UCF101(args, transform=transform, temporal_sample=temporal_sample) \ No newline at end of file diff --git a/opensora/dataset/feature_datasets.py b/opensora/dataset/feature_datasets.py new file mode 100644 index 000000000..7f8f251e9 --- /dev/null +++ b/opensora/dataset/feature_datasets.py @@ -0,0 +1,49 @@ +import os +import torch +import random +import torch.utils.data as data + +import numpy as np +from glob import glob +from PIL import Image +from tqdm import tqdm + + +class LandscopeFeatures(data.Dataset): + def __init__(self, args, temporal_sample): + + self.args = args + self.data_path = args.data_path + self.temporal_sample = temporal_sample + + data_all = list(glob(self.data_path)) + self.num_frames = self.args.num_frames + print('Building dataset...') + self.data_all = [i for i in tqdm(data_all) if self.num_frames < np.load(i).shape[1]] + print(f'Total {len(self.data_all)} to train, origin dataset have {len(data_all)} videos.') + + + + def __getitem__(self, index): + try: + # import ipdb;ipdb.set_trace() + npy_path = self.data_all[index] + vframes = np.load(npy_path)[0] # f, c, h, w + + total_frames = len(vframes) + + # Sampling video frames + start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) + assert end_frame_ind - start_frame_ind >= self.target_video_len + frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.target_video_len, dtype=int) + video_clip = vframes[frame_indice] # + video_clip = torch.from_numpy(video_clip) + return video_clip, 1 + except Exception as e: + print(f'Error with {e}', npy_path) + return self.__getitem__(random.randint(0, self.__len__()-1)) + + + def __len__(self): + return len(self.data_all) + diff --git a/opensora/dataset/transform.py b/opensora/dataset/transform.py new file mode 100644 index 000000000..829b3fa65 --- /dev/null +++ b/opensora/dataset/transform.py @@ -0,0 +1,53 @@ +import math +import random + +import torch +from torchvision import transforms +from torchvision.transforms import Lambda + + + +class TemporalRandomCrop(object): + """Temporally crop the given frame indices at a random location. + + Args: + size (int): Desired length of frames will be seen in the model. + """ + + def __init__(self, size): + self.size = size + + def __call__(self, total_frames): + rand_end = max(0, total_frames - self.size - 1) + begin_index = random.randint(0, rand_end) + end_index = min(begin_index + self.size, total_frames) + return begin_index, end_index + + +class LongSideScale(torch.nn.Module): + """ + ``nn.Module`` wrapper for ``pytorchvideo.transforms.functional.short_side_scale``. + """ + + def __init__( + self, size: int, interpolation: str = "bilinear" + ): + super().__init__() + self._size = size + self._interpolation = interpolation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (torch.Tensor): video tensor with shape (C, T, H, W). + """ + assert len(x.shape) == 4 + assert x.dtype == torch.float32 + c, t, h, w = x.shape + if w < h: + new_w = int(math.floor((float(w) / h) * self._size)) + new_h = self._size + else: + new_w = self._size + new_h = int(math.floor((float(h) / w) * self._size)) + return torch.nn.functional.interpolate(x, size=(new_h, new_w), mode=self._interpolation, align_corners=False) diff --git a/opensora/dataset/ucf101.py b/opensora/dataset/ucf101.py new file mode 100644 index 000000000..2be67cf10 --- /dev/null +++ b/opensora/dataset/ucf101.py @@ -0,0 +1,105 @@ +import math +import os + +import decord +import numpy as np +import torch +import torchvision +from decord import VideoReader, cpu +from torch.utils.data import Dataset +from torchvision.transforms import Compose, Lambda, ToTensor +from torchvision.transforms._transforms_video import NormalizeVideo, RandomCropVideo, RandomHorizontalFlipVideo +from pytorchvideo.transforms import ApplyTransformToKey, ShortSideScale, UniformTemporalSubsample +from torch.nn import functional as F +import random + + +class DecordInit(object): + """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader.""" + + def __init__(self, num_threads=1): + self.num_threads = num_threads + self.ctx = decord.cpu(0) + + def __call__(self, filename): + """Perform the Decord initialization. + Args: + results (dict): The resulting dict to be modified and passed + to the next transform in pipeline. + """ + reader = decord.VideoReader(filename, + ctx=self.ctx, + num_threads=self.num_threads) + return reader + + def __repr__(self): + repr_str = (f'{self.__class__.__name__}(' + f'sr={self.sr},' + f'num_threads={self.num_threads})') + return repr_str + +class UCF101(Dataset): + def __init__(self, args, transform, temporal_sample): + self.data_path = args.data_path + self.num_frames = args.num_frames + self.transform = transform + self.temporal_sample = temporal_sample + self.v_decoder = DecordInit() + + self.classes = sorted(os.listdir(self.data_path)) + self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)} + self.samples = self._make_dataset() + + + def _make_dataset(self): + dataset = [] + for class_name in self.classes: + class_path = os.path.join(self.data_path, class_name) + for fname in os.listdir(class_path): + if fname.endswith('.avi'): + item = (os.path.join(class_path, fname), self.class_to_idx[class_name]) + dataset.append(item) + return dataset + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + video_path, label = self.samples[idx] + + try: + video = self.tv_read(video_path) + video = self.transform(video) # C T H W + return video, label + except Exception as e: + print(f'Error with {e}, {video_path}') + return self.__getitem__(random.randint(0, self.__len__()-1)) + + def tv_read(self, path): + vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW') + total_frames = len(vframes) + + # Sampling video frames + start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) + assert end_frame_ind - start_frame_ind >= self.num_frames + frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int) + # frame_indice = np.linspace(0, 30, self.target_video_len, dtype=int) + # print(frame_indice) + video = vframes[frame_indice] # + video = video.permute(0, 2, 3, 1) # (T, C, H, W) -> (T H W C) + + return video + + def decord_read(self, path): + decord_vr = self.v_decoder(path) + total_frames = len(decord_vr) + # Sampling video frames + start_frame_ind, end_frame_ind = self.temporal_sample(total_frames) + assert end_frame_ind - start_frame_ind >= self.target_video_len + frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.target_video_len, dtype=int) + + video_data = decord_vr.get_batch(frame_indice).asnumpy() + video_data = torch.from_numpy(video_data) + # video_data = video_data.permute(0, 3, 1, 2) # (T, H, W, C) -> (T C H W) + return video_data + diff --git a/opensora/models/ae/__init__.py b/opensora/models/ae/__init__.py new file mode 100644 index 000000000..59a6c0e35 --- /dev/null +++ b/opensora/models/ae/__init__.py @@ -0,0 +1,10 @@ +from .imagebase import imagebase_ae, vae, vqvae +from .videobase import videobase_ae, videovae, videovqvae + + +def getae(args): + ae = imagebase_ae.get(args.ae, None) + if ae is None: + ae = videobase_ae.get(args.ae, None) + assert ae is not None + return ae(args.ae) \ No newline at end of file diff --git a/opensora/models/ae/imagebase/__init__.py b/opensora/models/ae/imagebase/__init__.py new file mode 100644 index 000000000..b1b07abf8 --- /dev/null +++ b/opensora/models/ae/imagebase/__init__.py @@ -0,0 +1,14 @@ +from .vae.vae import HFVAEWrapper +from .vae.vae import SDVAEWrapper +from .vqvae.vqvae import SDVQVAEWrapper + +vae = ['stabilityai/sd-vae-ft-mse', 'stabilityai/sd-vae-ft-ema'] +vqvae = ['vqgan_imagenet_f16_1024', 'vqgan_imagenet_f16_16384', 'vqgan_gumbel_f8'] + +imagebase_ae = { + 'stabilityai/sd-vae-ft-mse': HFVAEWrapper, + 'stabilityai/sd-vae-ft-ema': HFVAEWrapper, + 'vqgan_imagenet_f16_1024': SDVQVAEWrapper, + 'vqgan_imagenet_f16_16384': SDVQVAEWrapper, + 'vqgan_gumbel_f8': SDVQVAEWrapper, +} \ No newline at end of file diff --git a/src/sora/models/interpolation/networks/__init__.py b/opensora/models/ae/imagebase/vae/__init__.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/__init__.py rename to opensora/models/ae/imagebase/vae/__init__.py diff --git a/opensora/models/ae/imagebase/vae/vae.py b/opensora/models/ae/imagebase/vae/vae.py new file mode 100644 index 000000000..2a4fac05a --- /dev/null +++ b/opensora/models/ae/imagebase/vae/vae.py @@ -0,0 +1,26 @@ + +from torch import nn +from diffusers.models import AutoencoderKL + + +class HFVAEWrapper(nn.Module): + def __init__(self, hfvae='mse'): + super(HFVAEWrapper, self).__init__() + self.vae = AutoencoderKL.from_pretrained(f'stabilityai/sd-vae-ft-{hfvae}', cache_dir='cache_dir') + def encode(self, x): # b c h w + x = self.vae.encode(x).latent_dist.sample().mul_(0.18215) + return x + def decode(self, x): + x = self.vae.decode(x / 0.18215).sample + return x + +class SDVAEWrapper(nn.Module): + def __init__(self): + super(SDVAEWrapper, self).__init__() + raise NotImplementedError + + def encode(self, x): # b c h w + raise NotImplementedError + + def decode(self, x): + raise NotImplementedError \ No newline at end of file diff --git a/src/sora/models/interpolation/networks/blocks/__init__.py b/opensora/models/ae/imagebase/vqvae/__init__.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/blocks/__init__.py rename to opensora/models/ae/imagebase/vqvae/__init__.py diff --git a/opensora/models/ae/imagebase/vqvae/model.py b/opensora/models/ae/imagebase/vqvae/model.py new file mode 100644 index 000000000..7c9a757f7 --- /dev/null +++ b/opensora/models/ae/imagebase/vqvae/model.py @@ -0,0 +1,775 @@ +# pytorch_diffusion + derived encoder decoder +import math +import torch +import torch.nn as nn +import numpy as np + + +def get_timestep_embedding(timesteps, embedding_dim): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: + From Fairseq. + Build sinusoidal embeddings. + This matches the implementation in tensor2tensor, but differs slightly + from the description in Section 3.5 of "Attention Is All You Need". + """ + assert len(timesteps.shape) == 1 + + half_dim = embedding_dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) + emb = emb.to(device=timesteps.device) + emb = timesteps.float()[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: # zero pad + emb = torch.nn.functional.pad(emb, (0,1,0,0)) + return emb + + +def nonlinearity(x): + # swish + return x*torch.sigmoid(x) + + +def Normalize(in_channels): + return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + +class Upsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=3, + stride=1, + padding=1) + + 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(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=3, + stride=2, + padding=0) + + def forward(self, x): + if self.with_conv: + 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 ResnetBlock(nn.Module): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, + dropout, temb_channels=512): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.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) + if temb_channels > 0: + self.temb_proj = torch.nn.Linear(temb_channels, + out_channels) + 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) + + def forward(self, x, temb): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None] + + 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 AttnBlock(nn.Module): + 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) + + + 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 Model(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, use_timestep=True): + super().__init__() + self.ch = ch + self.temb_ch = self.ch*4 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + self.use_timestep = use_timestep + if self.use_timestep: + # timestep embedding + self.temb = nn.Module() + self.temb.dense = nn.ModuleList([ + torch.nn.Linear(self.ch, + self.temb_ch), + torch.nn.Linear(self.temb_ch, + self.temb_ch), + ]) + + # downsampling + self.conv_in = torch.nn.Conv2d(in_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + skip_in = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + if i_block == self.num_res_blocks: + skip_in = ch*in_ch_mult[i_level] + block.append(ResnetBlock(in_channels=block_in+skip_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1) + + + def forward(self, x, t=None): + #assert x.shape[2] == x.shape[3] == self.resolution + + if self.use_timestep: + # timestep embedding + assert t is not None + temb = get_timestep_embedding(t, self.ch) + temb = self.temb.dense[0](temb) + temb = nonlinearity(temb) + temb = self.temb.dense[1](temb) + else: + temb = None + + # downsampling + 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], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + 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]( + torch.cat([h, hs.pop()], dim=1), temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Encoder(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, z_channels, double_z=True, **ignore_kwargs): + super().__init__() + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + # downsampling + self.conv_in = torch.nn.Conv2d(in_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + 2*z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1) + + + def forward(self, x): + #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution) + + # timestep embedding + temb = None + + # downsampling + 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], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, z_channels, give_pre_end=False, **ignorekwargs): + super().__init__() + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.give_pre_end = give_pre_end + + # compute in_ch_mult, block_in and curr_res at lowest res + in_ch_mult = (1,)+tuple(ch_mult) + block_in = ch*ch_mult[self.num_resolutions-1] + curr_res = resolution // 2**(self.num_resolutions-1) + self.z_shape = (1,z_channels,curr_res,curr_res) + print("Working with z of shape {} = {} dimensions.".format( + self.z_shape, np.prod(self.z_shape))) + + # z to block_in + self.conv_in = torch.nn.Conv2d(z_channels, + block_in, + kernel_size=3, + stride=1, + padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, z): + #assert z.shape[1:] == self.z_shape[1:] + self.last_z_shape = z.shape + + # timestep embedding + temb = None + + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + 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, temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class VUNet(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, + in_channels, c_channels, + resolution, z_channels, use_timestep=False, **ignore_kwargs): + super().__init__() + self.ch = ch + self.temb_ch = self.ch*4 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + + self.use_timestep = use_timestep + if self.use_timestep: + # timestep embedding + self.temb = nn.Module() + self.temb.dense = nn.ModuleList([ + torch.nn.Linear(self.ch, + self.temb_ch), + torch.nn.Linear(self.temb_ch, + self.temb_ch), + ]) + + # downsampling + self.conv_in = torch.nn.Conv2d(c_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv) + curr_res = curr_res // 2 + self.down.append(down) + + self.z_in = torch.nn.Conv2d(z_channels, + block_in, + kernel_size=1, + stride=1, + padding=0) + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=2*block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + skip_in = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + if i_block == self.num_res_blocks: + skip_in = ch*in_ch_mult[i_level] + block.append(ResnetBlock(in_channels=block_in+skip_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1) + + + def forward(self, x, z): + #assert x.shape[2] == x.shape[3] == self.resolution + + if self.use_timestep: + # timestep embedding + assert t is not None + temb = get_timestep_embedding(t, self.ch) + temb = self.temb.dense[0](temb) + temb = nonlinearity(temb) + temb = self.temb.dense[1](temb) + else: + temb = None + + # downsampling + 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], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + z = self.z_in(z) + h = torch.cat((h,z),dim=1) + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + 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]( + torch.cat([h, hs.pop()], dim=1), temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class SimpleDecoder(nn.Module): + def __init__(self, in_channels, out_channels, *args, **kwargs): + super().__init__() + self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1), + ResnetBlock(in_channels=in_channels, + out_channels=2 * in_channels, + temb_channels=0, dropout=0.0), + ResnetBlock(in_channels=2 * in_channels, + out_channels=4 * in_channels, + temb_channels=0, dropout=0.0), + ResnetBlock(in_channels=4 * in_channels, + out_channels=2 * in_channels, + temb_channels=0, dropout=0.0), + nn.Conv2d(2*in_channels, in_channels, 1), + Upsample(in_channels, with_conv=True)]) + # end + self.norm_out = Normalize(in_channels) + self.conv_out = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + for i, layer in enumerate(self.model): + if i in [1,2,3]: + x = layer(x, None) + else: + x = layer(x) + + h = self.norm_out(x) + h = nonlinearity(h) + x = self.conv_out(h) + return x + + +class UpsampleDecoder(nn.Module): + def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution, + ch_mult=(2,2), dropout=0.0): + super().__init__() + # upsampling + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + block_in = in_channels + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.res_blocks = nn.ModuleList() + self.upsample_blocks = nn.ModuleList() + for i_level in range(self.num_resolutions): + res_block = [] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + res_block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + self.res_blocks.append(nn.ModuleList(res_block)) + if i_level != self.num_resolutions - 1: + self.upsample_blocks.append(Upsample(block_in, True)) + curr_res = curr_res * 2 + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + # upsampling + h = x + for k, i_level in enumerate(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.res_blocks[i_level][i_block](h, None) + if i_level != self.num_resolutions - 1: + h = self.upsample_blocks[k](h) + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h \ No newline at end of file diff --git a/opensora/models/ae/imagebase/vqvae/quantize.py b/opensora/models/ae/imagebase/vqvae/quantize.py new file mode 100644 index 000000000..dd258efe9 --- /dev/null +++ b/opensora/models/ae/imagebase/vqvae/quantize.py @@ -0,0 +1,447 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from torch import einsum +from einops import rearrange + + +class VectorQuantizer(nn.Module): + """ + see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py + ____________________________________________ + Discretization bottleneck part of the VQ-VAE. + Inputs: + - n_e : number of embeddings + - e_dim : dimension of embedding + - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2 + _____________________________________________ + """ + + # NOTE: this class contains a bug regarding beta; see VectorQuantizer2 for + # a fix and use legacy=False to apply that fix. VectorQuantizer2 can be + # used wherever VectorQuantizer has been used before and is additionally + # more efficient. + def __init__(self, n_e, e_dim, beta): + super(VectorQuantizer, self).__init__() + self.n_e = n_e + self.e_dim = e_dim + self.beta = beta + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + + def forward(self, z): + """ + Inputs the output of the encoder network z and maps it to a discrete + one-hot vector that is the index of the closest embedding vector e_j + z (continuous) -> z_q (discrete) + z.shape = (batch, channel, height, width) + quantization pipeline: + 1. get encoder input (B,C,H,W) + 2. flatten input to (B*H*W,C) + """ + # reshape z -> (batch, height, width, channel) and flatten + z = z.permute(0, 2, 3, 1).contiguous() + z_flattened = z.view(-1, self.e_dim) + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + + d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \ + torch.matmul(z_flattened, self.embedding.weight.t()) + + ## could possible replace this here + # #\start... + # find closest encodings + min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1) + + min_encodings = torch.zeros( + min_encoding_indices.shape[0], self.n_e).to(z) + min_encodings.scatter_(1, min_encoding_indices, 1) + + # dtype min encodings: torch.float32 + # min_encodings shape: torch.Size([2048, 512]) + # min_encoding_indices.shape: torch.Size([2048, 1]) + + # get quantized latent vectors + z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape) + # .........\end + + # with: + # .........\start + # min_encoding_indices = torch.argmin(d, dim=1) + # z_q = self.embedding(min_encoding_indices) + # ......\end......... (TODO) + + # compute loss for embedding + loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * \ + torch.mean((z_q - z.detach()) ** 2) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # perplexity + e_mean = torch.mean(min_encodings, dim=0) + perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10))) + + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + return z_q, loss, (perplexity, min_encodings, min_encoding_indices) + + def get_codebook_entry(self, indices, shape): + # shape specifying (batch, height, width, channel) + # TODO: check for more easy handling with nn.Embedding + min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices) + min_encodings.scatter_(1, indices[:, None], 1) + + # get quantized latent vectors + z_q = torch.matmul(min_encodings.float(), self.embedding.weight) + + if shape is not None: + z_q = z_q.view(shape) + + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + return z_q + + +class GumbelQuantize(nn.Module): + """ + credit to @karpathy: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py (thanks!) + Gumbel Softmax trick quantizer + Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016 + https://arxiv.org/abs/1611.01144 + """ + + def __init__(self, num_hiddens, embedding_dim, n_embed, straight_through=True, + kl_weight=5e-4, temp_init=1.0, use_vqinterface=True, + remap=None, unknown_index="random"): + super().__init__() + + self.embedding_dim = embedding_dim + self.n_embed = n_embed + + self.straight_through = straight_through + self.temperature = temp_init + self.kl_weight = kl_weight + + self.proj = nn.Conv2d(num_hiddens, n_embed, 1) + self.embed = nn.Embedding(n_embed, embedding_dim) + + self.use_vqinterface = use_vqinterface + + self.remap = remap + if self.remap is not None: + self.register_buffer("used", torch.tensor(np.load(self.remap))) + self.re_embed = self.used.shape[0] + self.unknown_index = unknown_index # "random" or "extra" or integer + if self.unknown_index == "extra": + self.unknown_index = self.re_embed + self.re_embed = self.re_embed + 1 + print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. " + f"Using {self.unknown_index} for unknown indices.") + else: + self.re_embed = n_embed + + def remap_to_used(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + match = (inds[:, :, None] == used[None, None, ...]).long() + new = match.argmax(-1) + unknown = match.sum(2) < 1 + if self.unknown_index == "random": + new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) + else: + new[unknown] = self.unknown_index + return new.reshape(ishape) + + def unmap_to_all(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + if self.re_embed > self.used.shape[0]: # extra token + inds[inds >= self.used.shape[0]] = 0 # simply set to zero + back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) + return back.reshape(ishape) + + def forward(self, z, temp=None, return_logits=False): + # force hard = True when we are in eval mode, as we must quantize. actually, always true seems to work + hard = self.straight_through if self.training else True + temp = self.temperature if temp is None else temp + + logits = self.proj(z) + if self.remap is not None: + # continue only with used logits + full_zeros = torch.zeros_like(logits) + logits = logits[:, self.used, ...] + + soft_one_hot = F.gumbel_softmax(logits, tau=temp, dim=1, hard=hard) + if self.remap is not None: + # go back to all entries but unused set to zero + full_zeros[:, self.used, ...] = soft_one_hot + soft_one_hot = full_zeros + z_q = einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight) + + # + kl divergence to the prior loss + qy = F.softmax(logits, dim=1) + diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.n_embed + 1e-10), dim=1).mean() + + ind = soft_one_hot.argmax(dim=1) + if self.remap is not None: + ind = self.remap_to_used(ind) + if self.use_vqinterface: + if return_logits: + return z_q, diff, (None, None, ind), logits + return z_q, diff, (None, None, ind) + return z_q, diff, ind + + def get_codebook_entry(self, indices, shape): + b, h, w, c = shape + assert b * h * w == indices.shape[0] + indices = rearrange(indices, '(b h w) -> b h w', b=b, h=h, w=w) + if self.remap is not None: + indices = self.unmap_to_all(indices) + one_hot = F.one_hot(indices, num_classes=self.n_embed).permute(0, 3, 1, 2).float() + z_q = einsum('b n h w, n d -> b d h w', one_hot, self.embed.weight) + return z_q + + +class VectorQuantizer2(nn.Module): + """ + Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly + avoids costly matrix multiplications and allows for post-hoc remapping of indices. + """ + + # NOTE: due to a bug the beta term was applied to the wrong term. for + # backwards compatibility we use the buggy version by default, but you can + # specify legacy=False to fix it. + def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random", + sane_index_shape=False, legacy=True): + super().__init__() + self.n_e = n_e + self.e_dim = e_dim + self.beta = beta + self.legacy = legacy + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + + self.remap = remap + if self.remap is not None: + self.register_buffer("used", torch.tensor(np.load(self.remap))) + self.re_embed = self.used.shape[0] + self.unknown_index = unknown_index # "random" or "extra" or integer + if self.unknown_index == "extra": + self.unknown_index = self.re_embed + self.re_embed = self.re_embed + 1 + print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " + f"Using {self.unknown_index} for unknown indices.") + else: + self.re_embed = n_e + + self.sane_index_shape = sane_index_shape + + def remap_to_used(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + match = (inds[:, :, None] == used[None, None, ...]).long() + new = match.argmax(-1) + unknown = match.sum(2) < 1 + if self.unknown_index == "random": + new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) + else: + new[unknown] = self.unknown_index + return new.reshape(ishape) + + def unmap_to_all(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + if self.re_embed > self.used.shape[0]: # extra token + inds[inds >= self.used.shape[0]] = 0 # simply set to zero + back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) + return back.reshape(ishape) + + def forward(self, z, temp=None, rescale_logits=False, return_logits=False): + assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" + assert rescale_logits == False, "Only for interface compatible with Gumbel" + assert return_logits == False, "Only for interface compatible with Gumbel" + # reshape z -> (batch, height, width, channel) and flatten + z = rearrange(z, 'b c h w -> b h w c').contiguous() + z_flattened = z.view(-1, self.e_dim) + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + + d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \ + torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n')) + + min_encoding_indices = torch.argmin(d, dim=1) + z_q = self.embedding(min_encoding_indices).view(z.shape) + perplexity = None + min_encodings = None + + # compute loss for embedding + if not self.legacy: + loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + \ + torch.mean((z_q - z.detach()) ** 2) + else: + loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * \ + torch.mean((z_q - z.detach()) ** 2) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # reshape back to match original input shape + z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous() + + if self.remap is not None: + min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis + min_encoding_indices = self.remap_to_used(min_encoding_indices) + min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten + + if self.sane_index_shape: + min_encoding_indices = min_encoding_indices.reshape( + z_q.shape[0], z_q.shape[2], z_q.shape[3]) + + return z_q, loss, (perplexity, min_encodings, min_encoding_indices) + + def get_codebook_entry(self, indices, shape): + # shape specifying (batch, height, width, channel) + if self.remap is not None: + indices = indices.reshape(shape[0], -1) # add batch axis + indices = self.unmap_to_all(indices) + indices = indices.reshape(-1) # flatten again + + # get quantized latent vectors + z_q = self.embedding(indices) + + if shape is not None: + z_q = z_q.view(shape) + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + return z_q + + +class EmbeddingEMA(nn.Module): + def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5): + super().__init__() + self.decay = decay + self.eps = eps + weight = torch.randn(num_tokens, codebook_dim) + self.weight = nn.Parameter(weight, requires_grad=False) + self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad=False) + self.embed_avg = nn.Parameter(weight.clone(), requires_grad=False) + self.update = True + + def forward(self, embed_id): + return F.embedding(embed_id, self.weight) + + def cluster_size_ema_update(self, new_cluster_size): + self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay) + + def embed_avg_ema_update(self, new_embed_avg): + self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay) + + def weight_update(self, num_tokens): + n = self.cluster_size.sum() + smoothed_cluster_size = ( + (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n + ) + # normalize embedding average with smoothed cluster size + embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1) + self.weight.data.copy_(embed_normalized) + + +class EMAVectorQuantizer(nn.Module): + def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5, + remap=None, unknown_index="random"): + super().__init__() + self.codebook_dim = codebook_dim + self.num_tokens = num_tokens + self.beta = beta + self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps) + + self.remap = remap + if self.remap is not None: + self.register_buffer("used", torch.tensor(np.load(self.remap))) + self.re_embed = self.used.shape[0] + self.unknown_index = unknown_index # "random" or "extra" or integer + if self.unknown_index == "extra": + self.unknown_index = self.re_embed + self.re_embed = self.re_embed + 1 + print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. " + f"Using {self.unknown_index} for unknown indices.") + else: + self.re_embed = n_embed + + def remap_to_used(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + match = (inds[:, :, None] == used[None, None, ...]).long() + new = match.argmax(-1) + unknown = match.sum(2) < 1 + if self.unknown_index == "random": + new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) + else: + new[unknown] = self.unknown_index + return new.reshape(ishape) + + def unmap_to_all(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + if self.re_embed > self.used.shape[0]: # extra token + inds[inds >= self.used.shape[0]] = 0 # simply set to zero + back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) + return back.reshape(ishape) + + def forward(self, z): + # reshape z -> (batch, height, width, channel) and flatten + # z, 'b c h w -> b h w c' + z = rearrange(z, 'b c h w -> b h w c') + z_flattened = z.reshape(-1, self.codebook_dim) + + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \ + self.embedding.weight.pow(2).sum(dim=1) - 2 * \ + torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n' + + encoding_indices = torch.argmin(d, dim=1) + + z_q = self.embedding(encoding_indices).view(z.shape) + encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype) + avg_probs = torch.mean(encodings, dim=0) + perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) + + if self.training and self.embedding.update: + # EMA cluster size + encodings_sum = encodings.sum(0) + self.embedding.cluster_size_ema_update(encodings_sum) + # EMA embedding average + embed_sum = encodings.transpose(0, 1) @ z_flattened + self.embedding.embed_avg_ema_update(embed_sum) + # normalize embed_avg and update weight + self.embedding.weight_update(self.num_tokens) + + # compute loss for embedding + loss = self.beta * F.mse_loss(z_q.detach(), z) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # reshape back to match original input shape + # z_q, 'b h w c -> b c h w' + z_q = rearrange(z_q, 'b h w c -> b c h w') + return z_q, loss, (perplexity, encodings, encoding_indices) \ No newline at end of file diff --git a/opensora/models/ae/imagebase/vqvae/vqgan.py b/opensora/models/ae/imagebase/vqvae/vqgan.py new file mode 100644 index 000000000..9e9125be1 --- /dev/null +++ b/opensora/models/ae/imagebase/vqvae/vqgan.py @@ -0,0 +1,419 @@ +import torch +import torch.nn.functional as F +import pytorch_lightning as pl +import argparse, os, sys, datetime, glob, importlib + +from .model import Encoder, Decoder +from .quantize import VectorQuantizer2 as VectorQuantizer +from .quantize import GumbelQuantize +from .quantize import EMAVectorQuantizer + + + +def get_obj_from_str(string, reload=False): + module, cls = string.rsplit(".", 1) + if reload: + module_imp = importlib.import_module(module) + importlib.reload(module_imp) + return getattr(importlib.import_module(module, package=None), cls) + + +def instantiate_from_config(config): + if not "target" in config: + raise KeyError("Expected key `target` to instantiate.") + return get_obj_from_str(config["target"])(**config.get("params", dict())) + + +class VQModel(pl.LightningModule): + def __init__(self, + ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + remap=None, + sane_index_shape=False, # tell vector quantizer to return indices as bhw + ): + super().__init__() + self.image_key = image_key + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + self.loss = instantiate_from_config(lossconfig) + self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25, + remap=remap, sane_index_shape=sane_index_shape) + self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + self.image_key = image_key + if colorize_nlabels is not None: + assert type(colorize_nlabels)==int + self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) + if monitor is not None: + self.monitor = monitor + + def init_from_ckpt(self, path, ignore_keys=list()): + sd = torch.load(path, map_location="cpu")["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] + self.load_state_dict(sd, strict=False) + print(f"Restored from {path}") + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + quant, emb_loss, info = self.quantize(h) + return quant, emb_loss, info + + def decode(self, quant): + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + + def decode_code(self, code_b): + quant_b = self.quantize.embed_code(code_b) + dec = self.decode(quant_b) + return dec + + def forward(self, input): + quant, diff, _ = self.encode(input) + dec = self.decode(quant) + return dec, diff + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format) + return x.float() + + def training_step(self, batch, batch_idx, optimizer_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + + if optimizer_idx == 0: + # autoencode + aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + + self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return aeloss + + if optimizer_idx == 1: + # discriminator + discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return discloss + + def validation_step(self, batch, batch_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step, + last_layer=self.get_last_layer(), split="val") + + discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step, + last_layer=self.get_last_layer(), split="val") + rec_loss = log_dict_ae["val/rec_loss"] + self.log("val/rec_loss", rec_loss, + prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True) + self.log("val/aeloss", aeloss, + prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True) + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + return self.log_dict + + def configure_optimizers(self): + lr = self.learning_rate + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quantize.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), + lr=lr, betas=(0.5, 0.9)) + return [opt_ae, opt_disc], [] + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def log_images(self, batch, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + xrec, _ = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["inputs"] = x + log["reconstructions"] = xrec + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.*(x-x.min())/(x.max()-x.min()) - 1. + return x + + +class VQSegmentationModel(VQModel): + def __init__(self, n_labels, *args, **kwargs): + super().__init__(*args, **kwargs) + self.register_buffer("colorize", torch.randn(3, n_labels, 1, 1)) + + def configure_optimizers(self): + lr = self.learning_rate + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quantize.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + return opt_ae + + def training_step(self, batch, batch_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="train") + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return aeloss + + def validation_step(self, batch, batch_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="val") + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + total_loss = log_dict_ae["val/total_loss"] + self.log("val/total_loss", total_loss, + prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True) + return aeloss + + @torch.no_grad() + def log_images(self, batch, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + xrec, _ = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + # convert logits to indices + xrec = torch.argmax(xrec, dim=1, keepdim=True) + xrec = F.one_hot(xrec, num_classes=x.shape[1]) + xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float() + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["inputs"] = x + log["reconstructions"] = xrec + return log + + +class VQNoDiscModel(VQModel): + def __init__(self, + ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None + ): + super().__init__(ddconfig=ddconfig, lossconfig=lossconfig, n_embed=n_embed, embed_dim=embed_dim, + ckpt_path=ckpt_path, ignore_keys=ignore_keys, image_key=image_key, + colorize_nlabels=colorize_nlabels) + + def training_step(self, batch, batch_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + # autoencode + aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="train") + output = pl.TrainResult(minimize=aeloss) + output.log("train/aeloss", aeloss, + prog_bar=True, logger=True, on_step=True, on_epoch=True) + output.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return output + + def validation_step(self, batch, batch_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="val") + rec_loss = log_dict_ae["val/rec_loss"] + output = pl.EvalResult(checkpoint_on=rec_loss) + output.log("val/rec_loss", rec_loss, + prog_bar=True, logger=True, on_step=True, on_epoch=True) + output.log("val/aeloss", aeloss, + prog_bar=True, logger=True, on_step=True, on_epoch=True) + output.log_dict(log_dict_ae) + + return output + + def configure_optimizers(self): + optimizer = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quantize.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=self.learning_rate, betas=(0.5, 0.9)) + return optimizer + + +class GumbelVQ(VQModel): + def __init__(self, + ddconfig, + lossconfig, + n_embed, + embed_dim, + temperature_scheduler_config, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + kl_weight=1e-8, + remap=None, + ): + + z_channels = ddconfig["z_channels"] + super().__init__(ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=ignore_keys, + image_key=image_key, + colorize_nlabels=colorize_nlabels, + monitor=monitor, + ) + + self.loss.n_classes = n_embed + self.vocab_size = n_embed + + self.quantize = GumbelQuantize(z_channels, embed_dim, + n_embed=n_embed, + kl_weight=kl_weight, temp_init=1.0, + remap=remap) + + self.temperature_scheduler = instantiate_from_config(temperature_scheduler_config) # annealing of temp + + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + + def temperature_scheduling(self): + self.quantize.temperature = self.temperature_scheduler(self.global_step) + + def encode_to_prequant(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + return h + + def decode_code(self, code_b): + raise NotImplementedError + + def training_step(self, batch, batch_idx, optimizer_idx): + self.temperature_scheduling() + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x) + + if optimizer_idx == 0: + # autoencode + aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + self.log("temperature", self.quantize.temperature, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return aeloss + + if optimizer_idx == 1: + # discriminator + discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return discloss + + def validation_step(self, batch, batch_idx): + x = self.get_input(batch, self.image_key) + xrec, qloss = self(x, return_pred_indices=True) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step, + last_layer=self.get_last_layer(), split="val") + + discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step, + last_layer=self.get_last_layer(), split="val") + rec_loss = log_dict_ae["val/rec_loss"] + self.log("val/rec_loss", rec_loss, + prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) + self.log("val/aeloss", aeloss, + prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + return self.log_dict + + def log_images(self, batch, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + # encode + h = self.encoder(x) + h = self.quant_conv(h) + quant, _, _ = self.quantize(h) + # decode + x_rec = self.decode(quant) + log["inputs"] = x + log["reconstructions"] = x_rec + return log + + +class EMAVQ(VQModel): + def __init__(self, + ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + remap=None, + sane_index_shape=False, # tell vector quantizer to return indices as bhw + ): + super().__init__(ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=ignore_keys, + image_key=image_key, + colorize_nlabels=colorize_nlabels, + monitor=monitor, + ) + self.quantize = EMAVectorQuantizer(n_embed=n_embed, + embedding_dim=embed_dim, + beta=0.25, + remap=remap) + def configure_optimizers(self): + lr = self.learning_rate + #Remove self.quantize from parameter list since it is updated via EMA + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), + lr=lr, betas=(0.5, 0.9)) + return [opt_ae, opt_disc], [] \ No newline at end of file diff --git a/opensora/models/ae/imagebase/vqvae/vqvae.py b/opensora/models/ae/imagebase/vqvae/vqvae.py new file mode 100644 index 000000000..c43f8c42b --- /dev/null +++ b/opensora/models/ae/imagebase/vqvae/vqvae.py @@ -0,0 +1,34 @@ +from torch import nn +import yaml +import torch +from omegaconf import OmegaConf +from .vqgan import VQModel, GumbelVQ + +def load_config(config_path, display=False): + config = OmegaConf.load(config_path) + if display: + print(yaml.dump(OmegaConf.to_container(config))) + return config + + +def load_vqgan(config, ckpt_path=None, is_gumbel=False): + if is_gumbel: + model = GumbelVQ(**config.model.params) + else: + model = VQModel(**config.model.params) + if ckpt_path is not None: + sd = torch.load(ckpt_path, map_location="cpu")["state_dict"] + missing, unexpected = model.load_state_dict(sd, strict=False) + return model.eval() + + +class SDVQVAEWrapper(nn.Module): + def __init__(self, name): + super(SDVQVAEWrapper, self).__init__() + raise NotImplementedError + + def encode(self, x): # b c h w + raise NotImplementedError + + def decode(self, x): + raise NotImplementedError diff --git a/opensora/models/ae/videobase/__init__.py b/opensora/models/ae/videobase/__init__.py new file mode 100644 index 000000000..e70ea192c --- /dev/null +++ b/opensora/models/ae/videobase/__init__.py @@ -0,0 +1,11 @@ +from .vqvae.videogpt import VideoGPTVQVAEWrapper + +videovqvae = ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4'] +videovae = [] + +videobase_ae = { + 'bair_stride4x2x2': VideoGPTVQVAEWrapper, + 'ucf101_stride4x4x4': VideoGPTVQVAEWrapper, + 'kinetics_stride4x4x4': VideoGPTVQVAEWrapper, + 'kinetics_stride2x4x4': VideoGPTVQVAEWrapper, +} \ No newline at end of file diff --git a/src/sora/models/captioner/placeholder b/opensora/models/ae/videobase/vae/placeholder similarity index 100% rename from src/sora/models/captioner/placeholder rename to opensora/models/ae/videobase/vae/placeholder diff --git a/opensora/models/ae/videobase/vqvae/videogpt/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/__init__.py new file mode 100644 index 000000000..0fe3d520c --- /dev/null +++ b/opensora/models/ae/videobase/vqvae/videogpt/__init__.py @@ -0,0 +1,16 @@ +from torch import nn + +from .videogpt import load_vqvae + + +class VideoGPTVQVAEWrapper(nn.Module): + def __init__(self, ckpt='kinetics_stride4x4x4'): + super(VideoGPTVQVAEWrapper, self).__init__() + self.vqvae = load_vqvae(ckpt, root='./cache_dir') + def encode(self, x): # b c h w + x = self.vqvae.pre_vq_conv(self.vqvae.encoder(x)) + return x + def decode(self, x): + vq_output = self.vqvae.codebook(x) + x = self.vqvae.decoder(self.vqvae.post_vq_conv(vq_output['embeddings'])) + return x \ No newline at end of file diff --git a/src/sora/modules/ae/vqvae/videogpt/rec_video.py b/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/rec_video.py rename to opensora/models/ae/videobase/vqvae/videogpt/rec_video.py diff --git a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/requires.txt b/opensora/models/ae/videobase/vqvae/videogpt/requirements.txt similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/requires.txt rename to opensora/models/ae/videobase/vqvae/videogpt/requirements.txt index 21b4a9949..fde2a9f3d 100644 --- a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/requires.txt +++ b/opensora/models/ae/videobase/vqvae/videogpt/requirements.txt @@ -1,9 +1,9 @@ -av -gdown -gradio -h5py -imageio -moviepy -requests -scikit-video tqdm +requests +h5py +scikit-video +av +gradio +moviepy +imageio +gdown diff --git a/src/sora/models/interpolation/utils/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/__init__.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/utils/__init__.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/__init__.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/compute_fvd.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/compute_fvd.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/__init__.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/__init__.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/__init__.py diff --git a/src/sora/models/super_resolution/placeholder b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/__init__.py similarity index 100% rename from src/sora/models/super_resolution/placeholder rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/__init__.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh diff --git a/src/sora/modules/ae/vae/placeholder b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/__init__.py similarity index 100% rename from src/sora/modules/ae/vae/placeholder rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/__init__.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/sample_videogpt.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/sample_videogpt.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/train_videogpt.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/train_videogpt.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py diff --git a/src/sora/modules/ae/vqvae/videogpt/scripts/train_vqvae.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/scripts/train_vqvae.py rename to opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py diff --git a/src/sora/modules/ae/vqvae/videogpt/setup.py b/opensora/models/ae/videobase/vqvae/videogpt/setup.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/setup.py rename to opensora/models/ae/videobase/vqvae/videogpt/setup.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/__init__.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/attention.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/attention.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/data.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/data.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/download.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/download.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py diff --git a/src/sora/modules/diffusion/unet/placeholder b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/__init__.py similarity index 100% rename from src/sora/modules/diffusion/unet/placeholder rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/__init__.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/fvd.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/fvd.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/gpt.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/gpt.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/resnet.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/resnet.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/utils.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/utils.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py diff --git a/src/sora/modules/ae/vqvae/videogpt/videogpt/vqvae.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py similarity index 100% rename from src/sora/modules/ae/vqvae/videogpt/videogpt/vqvae.py rename to opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py diff --git a/src/sora/utils/placeholder b/opensora/models/captioner/placeholder similarity index 100% rename from src/sora/utils/placeholder rename to opensora/models/captioner/placeholder diff --git a/src/sora/modules/diffusion/dit/diffusion/__init__.py b/opensora/models/diffusion/dit/diffusion/__init__.py similarity index 100% rename from src/sora/modules/diffusion/dit/diffusion/__init__.py rename to opensora/models/diffusion/dit/diffusion/__init__.py diff --git a/src/sora/modules/diffusion/dit/diffusion/diffusion_utils.py b/opensora/models/diffusion/dit/diffusion/diffusion_utils.py similarity index 100% rename from src/sora/modules/diffusion/dit/diffusion/diffusion_utils.py rename to opensora/models/diffusion/dit/diffusion/diffusion_utils.py diff --git a/src/sora/modules/diffusion/dit/diffusion/gaussian_diffusion.py b/opensora/models/diffusion/dit/diffusion/gaussian_diffusion.py similarity index 100% rename from src/sora/modules/diffusion/dit/diffusion/gaussian_diffusion.py rename to opensora/models/diffusion/dit/diffusion/gaussian_diffusion.py diff --git a/src/sora/modules/diffusion/dit/diffusion/respace.py b/opensora/models/diffusion/dit/diffusion/respace.py similarity index 100% rename from src/sora/modules/diffusion/dit/diffusion/respace.py rename to opensora/models/diffusion/dit/diffusion/respace.py diff --git a/src/sora/modules/diffusion/dit/diffusion/timestep_sampler.py b/opensora/models/diffusion/dit/diffusion/timestep_sampler.py similarity index 100% rename from src/sora/modules/diffusion/dit/diffusion/timestep_sampler.py rename to opensora/models/diffusion/dit/diffusion/timestep_sampler.py diff --git a/src/sora/modules/diffusion/dit/models.py b/opensora/models/diffusion/dit/models.py similarity index 94% rename from src/sora/modules/diffusion/dit/models.py rename to opensora/models/diffusion/dit/models.py index 2e22359f2..c92f970be 100644 --- a/src/sora/modules/diffusion/dit/models.py +++ b/opensora/models/diffusion/dit/models.py @@ -52,28 +52,28 @@ class Attention(nn.Module): q, k, v = qkv.unbind(0) q, k = self.q_norm(q), self.k_norm(k) - if self.fused_attn: - x = F.scaled_dot_product_attention( - q, k, v, - attn_mask=attention_mask, - dropout_p=self.attn_drop.p if self.training else 0., - ) - else: - q = q * self.scale - attn = q @ k.transpose(-2, -1) - if attention_mask is not None: - attention_mask = attention_mask.flatten(1).unsqueeze(-1) # bs t h w -> bs thw 1 - attention_mask = attention_mask @ attention_mask.transpose(1, 2) # bs thw 1 @ bs 1 thw = bs thw thw - attention_mask = attention_mask.unsqueeze(1) - assert attn.shape[0] == attention_mask.shape[0] - attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(attn.dtype).min) - attn = attn + attention_mask - attn = attn.softmax(dim=-1) - if torch.any(torch.isnan(attn)): - print('torch.any(torch.isnan(attn))') - attn = attn.masked_fill(torch.isnan(attn), float(0.)) - attn = self.attn_drop(attn) - x = attn @ v + # if self.fused_attn: + # x = F.scaled_dot_product_attention( + # q, k, v, + # attn_mask=attention_mask, + # dropout_p=self.attn_drop.p if self.training else 0., + # ) + # else: + q = q * self.scale + attn = q @ k.transpose(-2, -1) + if attention_mask is not None: + attention_mask = attention_mask.flatten(1).unsqueeze(-1) # bs t h w -> bs thw 1 + attention_mask = attention_mask @ attention_mask.transpose(1, 2) # bs thw 1 @ bs 1 thw = bs thw thw + attention_mask = attention_mask.unsqueeze(1) + assert attn.shape[0] == attention_mask.shape[0] + attention_mask = attention_mask.masked_fill(attention_mask == 0, torch.finfo(attn.dtype).min) + attn = attn + attention_mask + attn = attn.softmax(dim=-1) + if torch.any(torch.isnan(attn)): + print('torch.any(torch.isnan(attn))') + attn = attn.masked_fill(torch.isnan(attn), float(0.)) + attn = self.attn_drop(attn) + x = attn @ v x = x.transpose(1, 2).reshape(B, N, C) x = self.proj(x) diff --git a/opensora/models/diffusion/latte/diffusion/__init__.py b/opensora/models/diffusion/latte/diffusion/__init__.py new file mode 100644 index 000000000..8a09d8c46 --- /dev/null +++ b/opensora/models/diffusion/latte/diffusion/__init__.py @@ -0,0 +1,47 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +from . import gaussian_diffusion as gd +from .respace import SpacedDiffusion, space_timesteps + + +def create_diffusion( + timestep_respacing, + noise_schedule="linear", + use_kl=False, + sigma_small=False, + predict_xstart=False, + learn_sigma=True, + # learn_sigma=False, + rescale_learned_sigmas=False, + diffusion_steps=1000 +): + betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) + if use_kl: + loss_type = gd.LossType.RESCALED_KL + elif rescale_learned_sigmas: + loss_type = gd.LossType.RESCALED_MSE + else: + loss_type = gd.LossType.MSE + if timestep_respacing is None or timestep_respacing == "": + timestep_respacing = [diffusion_steps] + return SpacedDiffusion( + use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), + betas=betas, + model_mean_type=( + gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X + ), + model_var_type=( + ( + gd.ModelVarType.FIXED_LARGE + if not sigma_small + else gd.ModelVarType.FIXED_SMALL + ) + if not learn_sigma + else gd.ModelVarType.LEARNED_RANGE + ), + loss_type=loss_type + # rescale_timesteps=rescale_timesteps, + ) diff --git a/opensora/models/diffusion/latte/diffusion/diffusion_utils.py b/opensora/models/diffusion/latte/diffusion/diffusion_utils.py new file mode 100644 index 000000000..e493a6a3e --- /dev/null +++ b/opensora/models/diffusion/latte/diffusion/diffusion_utils.py @@ -0,0 +1,88 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +import torch as th +import numpy as np + + +def normal_kl(mean1, logvar1, mean2, logvar2): + """ + Compute the KL divergence between two gaussians. + Shapes are automatically broadcasted, so batches can be compared to + scalars, among other use cases. + """ + tensor = None + for obj in (mean1, logvar1, mean2, logvar2): + if isinstance(obj, th.Tensor): + tensor = obj + break + assert tensor is not None, "at least one argument must be a Tensor" + + # Force variances to be Tensors. Broadcasting helps convert scalars to + # Tensors, but it does not work for th.exp(). + logvar1, logvar2 = [ + x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) + for x in (logvar1, logvar2) + ] + + return 0.5 * ( + -1.0 + + logvar2 + - logvar1 + + th.exp(logvar1 - logvar2) + + ((mean1 - mean2) ** 2) * th.exp(-logvar2) + ) + + +def approx_standard_normal_cdf(x): + """ + A fast approximation of the cumulative distribution function of the + standard normal. + """ + return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) + + +def continuous_gaussian_log_likelihood(x, *, means, log_scales): + """ + Compute the log-likelihood of a continuous Gaussian distribution. + :param x: the targets + :param means: the Gaussian mean Tensor. + :param log_scales: the Gaussian log stddev Tensor. + :return: a tensor like x of log probabilities (in nats). + """ + centered_x = x - means + inv_stdv = th.exp(-log_scales) + normalized_x = centered_x * inv_stdv + log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x) + return log_probs + + +def discretized_gaussian_log_likelihood(x, *, means, log_scales): + """ + Compute the log-likelihood of a Gaussian distribution discretizing to a + given image. + :param x: the target images. It is assumed that this was uint8 values, + rescaled to the range [-1, 1]. + :param means: the Gaussian mean Tensor. + :param log_scales: the Gaussian log stddev Tensor. + :return: a tensor like x of log probabilities (in nats). + """ + assert x.shape == means.shape == log_scales.shape + centered_x = x - means + inv_stdv = th.exp(-log_scales) + plus_in = inv_stdv * (centered_x + 1.0 / 255.0) + cdf_plus = approx_standard_normal_cdf(plus_in) + min_in = inv_stdv * (centered_x - 1.0 / 255.0) + cdf_min = approx_standard_normal_cdf(min_in) + log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) + log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) + cdf_delta = cdf_plus - cdf_min + log_probs = th.where( + x < -0.999, + log_cdf_plus, + th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), + ) + assert log_probs.shape == x.shape + return log_probs diff --git a/opensora/models/diffusion/latte/diffusion/gaussian_diffusion.py b/opensora/models/diffusion/latte/diffusion/gaussian_diffusion.py new file mode 100644 index 000000000..7fc3d43a7 --- /dev/null +++ b/opensora/models/diffusion/latte/diffusion/gaussian_diffusion.py @@ -0,0 +1,881 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + + +import math + +import numpy as np +import torch as th +import enum + +from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl + + +def mean_flat(tensor): + """ + Take the mean over all non-batch dimensions. + """ + return tensor.mean(dim=list(range(1, len(tensor.shape)))) + + +class ModelMeanType(enum.Enum): + """ + Which type of output the model predicts. + """ + + PREVIOUS_X = enum.auto() # the model predicts x_{t-1} + START_X = enum.auto() # the model predicts x_0 + EPSILON = enum.auto() # the model predicts epsilon + + +class ModelVarType(enum.Enum): + """ + What is used as the model's output variance. + The LEARNED_RANGE option has been added to allow the model to predict + values between FIXED_SMALL and FIXED_LARGE, making its job easier. + """ + + LEARNED = enum.auto() + FIXED_SMALL = enum.auto() + FIXED_LARGE = enum.auto() + LEARNED_RANGE = enum.auto() + + +class LossType(enum.Enum): + MSE = enum.auto() # use raw MSE loss (and KL when learning variances) + RESCALED_MSE = ( + enum.auto() + ) # use raw MSE loss (with RESCALED_KL when learning variances) + KL = enum.auto() # use the variational lower-bound + RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB + + def is_vb(self): + return self == LossType.KL or self == LossType.RESCALED_KL + + +def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + warmup_time = int(num_diffusion_timesteps * warmup_frac) + betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) + return betas + + +def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): + """ + This is the deprecated API for creating beta schedules. + See get_named_beta_schedule() for the new library of schedules. + """ + if beta_schedule == "quad": + betas = ( + np.linspace( + beta_start ** 0.5, + beta_end ** 0.5, + num_diffusion_timesteps, + dtype=np.float64, + ) + ** 2 + ) + elif beta_schedule == "linear": + betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "warmup10": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) + elif beta_schedule == "warmup50": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) + elif beta_schedule == "const": + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1 + betas = 1.0 / np.linspace( + num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 + ) + else: + raise NotImplementedError(beta_schedule) + assert betas.shape == (num_diffusion_timesteps,) + return betas + + +def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): + """ + Get a pre-defined beta schedule for the given name. + The beta schedule library consists of beta schedules which remain similar + in the limit of num_diffusion_timesteps. + Beta schedules may be added, but should not be removed or changed once + they are committed to maintain backwards compatibility. + """ + if schedule_name == "linear": + # Linear schedule from Ho et al, extended to work for any number of + # diffusion steps. + scale = 1000 / num_diffusion_timesteps + return get_beta_schedule( + "linear", + beta_start=scale * 0.0001, + beta_end=scale * 0.02, + num_diffusion_timesteps=num_diffusion_timesteps, + ) + elif schedule_name == "squaredcos_cap_v2": + return betas_for_alpha_bar( + num_diffusion_timesteps, + lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, + ) + else: + raise NotImplementedError(f"unknown beta schedule: {schedule_name}") + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, + which defines the cumulative product of (1-beta) over time from t = [0,1]. + :param num_diffusion_timesteps: the number of betas to produce. + :param alpha_bar: a lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + :param max_beta: the maximum beta to use; use values lower than 1 to + prevent singularities. + """ + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +class GaussianDiffusion: + """ + Utilities for training and sampling diffusion models. + Original ported from this codebase: + https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 + :param betas: a 1-D numpy array of betas for each diffusion timestep, + starting at T and going to 1. + """ + + def __init__( + self, + *, + betas, + model_mean_type, + model_var_type, + loss_type + ): + + self.model_mean_type = model_mean_type + self.model_var_type = model_var_type + self.loss_type = loss_type + + # Use float64 for accuracy. + betas = np.array(betas, dtype=np.float64) + self.betas = betas + assert len(betas.shape) == 1, "betas must be 1-D" + assert (betas > 0).all() and (betas <= 1).all() + + self.num_timesteps = int(betas.shape[0]) + + alphas = 1.0 - betas + self.alphas_cumprod = np.cumprod(alphas, axis=0) + self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) + self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) + assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) + self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) + self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) + self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) + self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + self.posterior_variance = ( + betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain + self.posterior_log_variance_clipped = np.log( + np.append(self.posterior_variance[1], self.posterior_variance[1:]) + ) if len(self.posterior_variance) > 1 else np.array([]) + + self.posterior_mean_coef1 = ( + betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + self.posterior_mean_coef2 = ( + (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod) + ) + + def q_mean_variance(self, x_start, t): + """ + Get the distribution q(x_t | x_0). + :param x_start: the [N x C x ...] tensor of noiseless inputs. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ + mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) + log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) + return mean, variance, log_variance + + def q_sample(self, x_start, t, noise=None): + """ + Diffuse the data for a given number of diffusion steps. + In other words, sample from q(x_t | x_0). + :param x_start: the initial data batch. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :param noise: if specified, the split-out normal noise. + :return: A noisy version of x_start. + """ + if noise is None: + noise = th.randn_like(x_start) + assert noise.shape == x_start.shape + return ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + def q_posterior_mean_variance(self, x_start, x_t, t): + """ + Compute the mean and variance of the diffusion posterior: + q(x_{t-1} | x_t, x_0) + """ + assert x_start.shape == x_t.shape + posterior_mean = ( + _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x_t.shape + ) + assert ( + posterior_mean.shape[0] + == posterior_variance.shape[0] + == posterior_log_variance_clipped.shape[0] + == x_start.shape[0] + ) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): + """ + Apply the model to get p(x_{t-1} | x_t), as well as a prediction of + the initial x, x_0. + :param model: the model, which takes a signal and a batch of timesteps + as input. + :param x: the [N x C x ...] tensor at time t. + :param t: a 1-D Tensor of timesteps. + :param clip_denoised: if True, clip the denoised signal into [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. Applies before + clip_denoised. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict with the following keys: + - 'mean': the model mean output. + - 'variance': the model variance output. + - 'log_variance': the log of 'variance'. + - 'pred_xstart': the prediction for x_0. + """ + if model_kwargs is None: + model_kwargs = {} + + B, F, C = x.shape[:3] + assert t.shape == (B,) + model_output = model(x, t, **model_kwargs) + # try: + # model_output = model_output.sample # for tav unet + # except: + # model_output = model(x, t, **model_kwargs) + if isinstance(model_output, tuple): + model_output, extra = model_output + else: + extra = None + + if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: + assert model_output.shape == (B, F, C * 2, *x.shape[3:]) + model_output, model_var_values = th.split(model_output, C, dim=2) + min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape) + max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) + # The model_var_values is [-1, 1] for [min_var, max_var]. + frac = (model_var_values + 1) / 2 + model_log_variance = frac * max_log + (1 - frac) * min_log + model_variance = th.exp(model_log_variance) + else: + model_variance, model_log_variance = { + # for fixedlarge, we set the initial (log-)variance like so + # to get a better decoder log likelihood. + ModelVarType.FIXED_LARGE: ( + np.append(self.posterior_variance[1], self.betas[1:]), + np.log(np.append(self.posterior_variance[1], self.betas[1:])), + ), + ModelVarType.FIXED_SMALL: ( + self.posterior_variance, + self.posterior_log_variance_clipped, + ), + }[self.model_var_type] + model_variance = _extract_into_tensor(model_variance, t, x.shape) + model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) + + def process_xstart(x): + if denoised_fn is not None: + x = denoised_fn(x) + if clip_denoised: + return x.clamp(-1, 1) + return x + + if self.model_mean_type == ModelMeanType.START_X: + pred_xstart = process_xstart(model_output) + else: + pred_xstart = process_xstart( + self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) + ) + model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t) + + assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape + return { + "mean": model_mean, + "variance": model_variance, + "log_variance": model_log_variance, + "pred_xstart": pred_xstart, + "extra": extra, + } + + def _predict_xstart_from_eps(self, x_t, t, eps): + assert x_t.shape == eps.shape + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t + - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps + ) + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute the mean for the previous step, given a function cond_fn that + computes the gradient of a conditional log probability with respect to + x. In particular, cond_fn computes grad(log(p(y|x))), and we want to + condition on y. + This uses the conditioning strategy from Sohl-Dickstein et al. (2015). + """ + gradient = cond_fn(x, t, **model_kwargs) + new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() + return new_mean + + def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute what the p_mean_variance output would have been, should the + model's score function be conditioned by cond_fn. + See condition_mean() for details on cond_fn. + Unlike condition_mean(), this instead uses the conditioning strategy + from Song et al (2020). + """ + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + + eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) + eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs) + + out = p_mean_var.copy() + out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) + out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t) + return out + + def p_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + ): + """ + Sample x_{t-1} from the model at the given timestep. + :param model: the model to sample from. + :param x: the current tensor at x_{t-1}. + :param t: the value of t, starting at 0 for the first diffusion step. + :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - 'sample': a random sample from the model. + - 'pred_xstart': a prediction of x_0. + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + noise = th.randn_like(x) + nonzero_mask = ( + (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + if cond_fn is not None: + out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs) + sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"]} + + def p_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + """ + Generate samples from the model. + :param model: the model module. + :param shape: the shape of the samples, (N, C, H, W). + :param noise: if specified, the noise from the encoder to sample. + Should be of the same shape as `shape`. + :param clip_denoised: if True, clip x_start predictions to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param device: if specified, the device to create the samples on. + If not specified, use a model parameter's device. + :param progress: if True, show a tqdm progress bar. + :return: a non-differentiable batch of samples. + """ + final = None + for sample in self.p_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + ): + final = sample + return final["sample"] + + def p_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + """ + Generate samples from the model and yield intermediate samples from + each timestep of diffusion. + Arguments are the same as p_sample_loop(). + Returns a generator over dicts, where each dict is the return value of + p_sample(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + if noise is not None: + img = noise + else: + img = th.randn(*shape, device=device) + indices = list(range(self.num_timesteps))[::-1] + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + with th.no_grad(): + out = self.p_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + ) + yield out + img = out["sample"] + + def ddim_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + eta=0.0, + ): + """ + Sample x_{t-1} from the model using DDIM. + Same usage as p_sample(). + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + if cond_fn is not None: + out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) + + # Usually our model outputs epsilon, but we re-derive it + # in case we used x_start or x_prev prediction. + eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) + + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) + sigma = ( + eta + * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) + * th.sqrt(1 - alpha_bar / alpha_bar_prev) + ) + # Equation 12. + noise = th.randn_like(x) + mean_pred = ( + out["pred_xstart"] * th.sqrt(alpha_bar_prev) + + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps + ) + nonzero_mask = ( + (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + sample = mean_pred + nonzero_mask * sigma * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"]} + + def ddim_reverse_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + eta=0.0, + ): + """ + Sample x_{t+1} from the model using DDIM reverse ODE. + """ + assert eta == 0.0, "Reverse ODE only for deterministic path" + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + if cond_fn is not None: + out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) + # Usually our model outputs epsilon, but we re-derive it + # in case we used x_start or x_prev prediction. + eps = ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x + - out["pred_xstart"] + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) + alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) + + # Equation 12. reversed + mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps + + return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} + + def ddim_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + """ + Generate samples from the model using DDIM. + Same usage as p_sample_loop(). + """ + final = None + for sample in self.ddim_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + eta=eta, + ): + final = sample + return final["sample"] + + def ddim_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + """ + Use DDIM to sample from the model and yield intermediate samples from + each timestep of DDIM. + Same usage as p_sample_loop_progressive(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + if noise is not None: + img = noise + else: + img = th.randn(*shape, device=device) + indices = list(range(self.num_timesteps))[::-1] + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + with th.no_grad(): + out = self.ddim_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + eta=eta, + ) + yield out + img = out["sample"] + + def _vb_terms_bpd( + self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None + ): + """ + Get a term for the variational lower-bound. + The resulting units are bits (rather than nats, as one might expect). + This allows for comparison to other papers. + :return: a dict with the following keys: + - 'output': a shape [N] tensor of NLLs or KLs. + - 'pred_xstart': the x_0 predictions. + """ + true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + ) + out = self.p_mean_variance( + model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs + ) + kl = normal_kl( + true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] + ) + kl = mean_flat(kl) / np.log(2.0) + + decoder_nll = -discretized_gaussian_log_likelihood( + x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] + ) + assert decoder_nll.shape == x_start.shape + decoder_nll = mean_flat(decoder_nll) / np.log(2.0) + + # At the first timestep return the decoder NLL, + # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) + output = th.where((t == 0), decoder_nll, kl) + return {"output": output, "pred_xstart": out["pred_xstart"]} + + def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): + """ + Compute training losses for a single timestep. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param t: a batch of timestep indices. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param noise: if specified, the specific Gaussian noise to try to remove. + :return: a dict with the key "loss" containing a tensor of shape [N]. + Some mean or variance settings may also have other keys. + """ + if model_kwargs is None: + model_kwargs = {} + if noise is None: + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start, t, noise=noise) + + terms = {} + + if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: + terms["loss"] = self._vb_terms_bpd( + model=model, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + model_kwargs=model_kwargs, + )["output"] + if self.loss_type == LossType.RESCALED_KL: + terms["loss"] *= self.num_timesteps + elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: + model_output = model(x_t, t, **model_kwargs) + # try: + # model_output = model(x_t, t, **model_kwargs).sample # for tav unet + # except: + # model_output = model(x_t, t, **model_kwargs) + + if self.model_var_type in [ + ModelVarType.LEARNED, + ModelVarType.LEARNED_RANGE, + ]: + B, F, C = x_t.shape[:3] + assert model_output.shape == (B, F, C * 2, *x_t.shape[3:]) + model_output, model_var_values = th.split(model_output, C, dim=2) + # Learn the variance using the variational bound, but don't let + # it affect our mean prediction. + frozen_out = th.cat([model_output.detach(), model_var_values], dim=2) + terms["vb"] = self._vb_terms_bpd( + model=lambda *args, r=frozen_out: r, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + )["output"] + if self.loss_type == LossType.RESCALED_MSE: + # Divide by 1000 for equivalence with initial implementation. + # Without a factor of 1/1000, the VB term hurts the MSE term. + terms["vb"] *= self.num_timesteps / 1000.0 + + target = { + ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + )[0], + ModelMeanType.START_X: x_start, + ModelMeanType.EPSILON: noise, + }[self.model_mean_type] + assert model_output.shape == target.shape == x_start.shape + terms["mse"] = mean_flat((target - model_output) ** 2) + if "vb" in terms: + terms["loss"] = terms["mse"] + terms["vb"] + else: + terms["loss"] = terms["mse"] + else: + raise NotImplementedError(self.loss_type) + + return terms + + def _prior_bpd(self, x_start): + """ + Get the prior KL term for the variational lower-bound, measured in + bits-per-dim. + This term can't be optimized, as it only depends on the encoder. + :param x_start: the [N x C x ...] tensor of inputs. + :return: a batch of [N] KL values (in bits), one per batch element. + """ + batch_size = x_start.shape[0] + t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) + qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) + kl_prior = normal_kl( + mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 + ) + return mean_flat(kl_prior) / np.log(2.0) + + def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): + """ + Compute the entire variational lower-bound, measured in bits-per-dim, + as well as other related quantities. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param clip_denoised: if True, clip denoised samples. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - total_bpd: the total variational lower-bound, per batch element. + - prior_bpd: the prior term in the lower-bound. + - vb: an [N x T] tensor of terms in the lower-bound. + - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. + - mse: an [N x T] tensor of epsilon MSEs for each timestep. + """ + device = x_start.device + batch_size = x_start.shape[0] + + vb = [] + xstart_mse = [] + mse = [] + for t in list(range(self.num_timesteps))[::-1]: + t_batch = th.tensor([t] * batch_size, device=device) + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) + # Calculate VLB term at the current timestep + with th.no_grad(): + out = self._vb_terms_bpd( + model, + x_start=x_start, + x_t=x_t, + t=t_batch, + clip_denoised=clip_denoised, + model_kwargs=model_kwargs, + ) + vb.append(out["output"]) + xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) + eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) + mse.append(mean_flat((eps - noise) ** 2)) + + vb = th.stack(vb, dim=1) + xstart_mse = th.stack(xstart_mse, dim=1) + mse = th.stack(mse, dim=1) + + prior_bpd = self._prior_bpd(x_start) + total_bpd = vb.sum(dim=1) + prior_bpd + return { + "total_bpd": total_bpd, + "prior_bpd": prior_bpd, + "vb": vb, + "xstart_mse": xstart_mse, + "mse": mse, + } + + +def _extract_into_tensor(arr, timesteps, broadcast_shape): + """ + Extract values from a 1-D numpy array for a batch of indices. + :param arr: the 1-D numpy array. + :param timesteps: a tensor of indices into the array to extract. + :param broadcast_shape: a larger shape of K dimensions with the batch + dimension equal to the length of timesteps. + :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. + """ + res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() + while len(res.shape) < len(broadcast_shape): + res = res[..., None] + return res + th.zeros(broadcast_shape, device=timesteps.device) diff --git a/opensora/models/diffusion/latte/diffusion/respace.py b/opensora/models/diffusion/latte/diffusion/respace.py new file mode 100644 index 000000000..31c009238 --- /dev/null +++ b/opensora/models/diffusion/latte/diffusion/respace.py @@ -0,0 +1,130 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py +import torch +import numpy as np +import torch as th + +from .gaussian_diffusion import GaussianDiffusion + + +def space_timesteps(num_timesteps, section_counts): + """ + Create a list of timesteps to use from an original diffusion process, + given the number of timesteps we want to take from equally-sized portions + of the original process. + For example, if there's 300 timesteps and the section counts are [10,15,20] + then the first 100 timesteps are strided to be 10 timesteps, the second 100 + are strided to be 15 timesteps, and the final 100 are strided to be 20. + If the stride is a string starting with "ddim", then the fixed striding + from the DDIM paper is used, and only one section is allowed. + :param num_timesteps: the number of diffusion steps in the original + process to divide up. + :param section_counts: either a list of numbers, or a string containing + comma-separated numbers, indicating the step count + per section. As a special case, use "ddimN" where N + is a number of steps to use the striding from the + DDIM paper. + :return: a set of diffusion steps from the original process to use. + """ + if isinstance(section_counts, str): + if section_counts.startswith("ddim"): + desired_count = int(section_counts[len("ddim") :]) + for i in range(1, num_timesteps): + if len(range(0, num_timesteps, i)) == desired_count: + return set(range(0, num_timesteps, i)) + raise ValueError( + f"cannot create exactly {num_timesteps} steps with an integer stride" + ) + section_counts = [int(x) for x in section_counts.split(",")] + size_per = num_timesteps // len(section_counts) + extra = num_timesteps % len(section_counts) + start_idx = 0 + all_steps = [] + for i, section_count in enumerate(section_counts): + size = size_per + (1 if i < extra else 0) + if size < section_count: + raise ValueError( + f"cannot divide section of {size} steps into {section_count}" + ) + if section_count <= 1: + frac_stride = 1 + else: + frac_stride = (size - 1) / (section_count - 1) + cur_idx = 0.0 + taken_steps = [] + for _ in range(section_count): + taken_steps.append(start_idx + round(cur_idx)) + cur_idx += frac_stride + all_steps += taken_steps + start_idx += size + return set(all_steps) + + +class SpacedDiffusion(GaussianDiffusion): + """ + A diffusion process which can skip steps in a base diffusion process. + :param use_timesteps: a collection (sequence or set) of timesteps from the + original diffusion process to retain. + :param kwargs: the kwargs to create the base diffusion process. + """ + + def __init__(self, use_timesteps, **kwargs): + self.use_timesteps = set(use_timesteps) + self.timestep_map = [] + self.original_num_steps = len(kwargs["betas"]) + + base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa + last_alpha_cumprod = 1.0 + new_betas = [] + for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): + if i in self.use_timesteps: + new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) + last_alpha_cumprod = alpha_cumprod + self.timestep_map.append(i) + kwargs["betas"] = np.array(new_betas) + super().__init__(**kwargs) + + def p_mean_variance( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) + + # @torch.compile + def training_losses( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().training_losses(self._wrap_model(model), *args, **kwargs) + + def condition_mean(self, cond_fn, *args, **kwargs): + return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) + + def condition_score(self, cond_fn, *args, **kwargs): + return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) + + def _wrap_model(self, model): + if isinstance(model, _WrappedModel): + return model + return _WrappedModel( + model, self.timestep_map, self.original_num_steps + ) + + def _scale_timesteps(self, t): + # Scaling is done by the wrapped model. + return t + + +class _WrappedModel: + def __init__(self, model, timestep_map, original_num_steps): + self.model = model + self.timestep_map = timestep_map + # self.rescale_timesteps = rescale_timesteps + self.original_num_steps = original_num_steps + + def __call__(self, x, ts, **kwargs): + map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) + new_ts = map_tensor[ts] + # if self.rescale_timesteps: + # new_ts = new_ts.float() * (1000.0 / self.original_num_steps) + return self.model(x, new_ts, **kwargs) diff --git a/opensora/models/diffusion/latte/diffusion/timestep_sampler.py b/opensora/models/diffusion/latte/diffusion/timestep_sampler.py new file mode 100644 index 000000000..a3f369847 --- /dev/null +++ b/opensora/models/diffusion/latte/diffusion/timestep_sampler.py @@ -0,0 +1,150 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +from abc import ABC, abstractmethod + +import numpy as np +import torch as th +import torch.distributed as dist + + +def create_named_schedule_sampler(name, diffusion): + """ + Create a ScheduleSampler from a library of pre-defined samplers. + :param name: the name of the sampler. + :param diffusion: the diffusion object to sample for. + """ + if name == "uniform": + return UniformSampler(diffusion) + elif name == "loss-second-moment": + return LossSecondMomentResampler(diffusion) + else: + raise NotImplementedError(f"unknown schedule sampler: {name}") + + +class ScheduleSampler(ABC): + """ + A distribution over timesteps in the diffusion process, intended to reduce + variance of the objective. + By default, samplers perform unbiased importance sampling, in which the + objective's mean is unchanged. + However, subclasses may override sample() to change how the resampled + terms are reweighted, allowing for actual changes in the objective. + """ + + @abstractmethod + def weights(self): + """ + Get a numpy array of weights, one per diffusion step. + The weights needn't be normalized, but must be positive. + """ + + def sample(self, batch_size, device): + """ + Importance-sample timesteps for a batch. + :param batch_size: the number of timesteps. + :param device: the torch device to save to. + :return: a tuple (timesteps, weights): + - timesteps: a tensor of timestep indices. + - weights: a tensor of weights to scale the resulting losses. + """ + w = self.weights() + p = w / np.sum(w) + indices_np = np.random.choice(len(p), size=(batch_size,), p=p) + indices = th.from_numpy(indices_np).long().to(device) + weights_np = 1 / (len(p) * p[indices_np]) + weights = th.from_numpy(weights_np).float().to(device) + return indices, weights + + +class UniformSampler(ScheduleSampler): + def __init__(self, diffusion): + self.diffusion = diffusion + self._weights = np.ones([diffusion.num_timesteps]) + + def weights(self): + return self._weights + + +class LossAwareSampler(ScheduleSampler): + def update_with_local_losses(self, local_ts, local_losses): + """ + Update the reweighting using losses from a model. + Call this method from each rank with a batch of timesteps and the + corresponding losses for each of those timesteps. + This method will perform synchronization to make sure all of the ranks + maintain the exact same reweighting. + :param local_ts: an integer Tensor of timesteps. + :param local_losses: a 1D Tensor of losses. + """ + batch_sizes = [ + th.tensor([0], dtype=th.int32, device=local_ts.device) + for _ in range(dist.get_world_size()) + ] + dist.all_gather( + batch_sizes, + th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), + ) + + # Pad all_gather batches to be the maximum batch size. + batch_sizes = [x.item() for x in batch_sizes] + max_bs = max(batch_sizes) + + timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] + loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] + dist.all_gather(timestep_batches, local_ts) + dist.all_gather(loss_batches, local_losses) + timesteps = [ + x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] + ] + losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] + self.update_with_all_losses(timesteps, losses) + + @abstractmethod + def update_with_all_losses(self, ts, losses): + """ + Update the reweighting using losses from a model. + Sub-classes should override this method to update the reweighting + using losses from the model. + This method directly updates the reweighting without synchronizing + between workers. It is called by update_with_local_losses from all + ranks with identical arguments. Thus, it should have deterministic + behavior to maintain state across workers. + :param ts: a list of int timesteps. + :param losses: a list of float losses, one per timestep. + """ + + +class LossSecondMomentResampler(LossAwareSampler): + def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): + self.diffusion = diffusion + self.history_per_term = history_per_term + self.uniform_prob = uniform_prob + self._loss_history = np.zeros( + [diffusion.num_timesteps, history_per_term], dtype=np.float64 + ) + self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) + + def weights(self): + if not self._warmed_up(): + return np.ones([self.diffusion.num_timesteps], dtype=np.float64) + weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) + weights /= np.sum(weights) + weights *= 1 - self.uniform_prob + weights += self.uniform_prob / len(weights) + return weights + + def update_with_all_losses(self, ts, losses): + for t, loss in zip(ts, losses): + if self._loss_counts[t] == self.history_per_term: + # Shift out the oldest loss term. + self._loss_history[t, :-1] = self._loss_history[t, 1:] + self._loss_history[t, -1] = loss + else: + self._loss_history[t, self._loss_counts[t]] = loss + self._loss_counts[t] += 1 + + def _warmed_up(self): + return (self._loss_counts == self.history_per_term).all() diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py new file mode 100644 index 000000000..5b6bf66b0 --- /dev/null +++ b/opensora/models/diffusion/latte/latte.py @@ -0,0 +1,539 @@ +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# GLIDE: https://github.com/openai/glide-text2im +# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py +# -------------------------------------------------------- +import math +import torch +import torch.nn as nn +import numpy as np + +from einops import rearrange, repeat +from timm.models.vision_transformer import Mlp, PatchEmbed + +import torch.utils.checkpoint as cp + +# for i in sys.path: +# print(i) + +# the xformers lib allows less memory, faster training and inference +try: + import xformers + import xformers.ops +except: + XFORMERS_IS_AVAILBLE = False + +# from timm.models.layers.helpers import to_2tuple +# from timm.models.layers.trace_utils import _assert + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + +################################################################################# +# Attention Layers from TIMM # +################################################################################# + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_lora=False, attention_mode='math'): + super().__init__() + assert dim % num_heads == 0, 'dim should be divisible by num_heads' + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim ** -0.5 + self.attention_mode = attention_mode + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous() + q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) + + if self.attention_mode == 'xformers': # cause loss nan while using with amp + x = xformers.ops.memory_efficient_attention(q, k, v).reshape(B, N, C) + + elif self.attention_mode == 'flash': + # cause loss nan while using with amp + # Optionally use the context manager to ensure one of the fused kerenels is run + with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False): + x = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C) # require pytorch 2.0 + + elif self.attention_mode == 'math': + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + + else: + raise NotImplemented + + x = self.proj(x) + x = self.proj_drop(x) + return x + + +################################################################################# +# Embedding Layers for Timesteps and Class Labels # +################################################################################# + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t, use_fp16=False): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + if use_fp16: + t_freq = t_freq.to(dtype=torch.float16) + t_emb = self.mlp(t_freq) + return t_emb + + +class LabelEmbedder(nn.Module): + """ + Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance. + """ + def __init__(self, num_classes, hidden_size, dropout_prob): + super().__init__() + use_cfg_embedding = dropout_prob > 0 + self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size) + self.num_classes = num_classes + self.dropout_prob = dropout_prob + + def token_drop(self, labels, force_drop_ids=None): + """ + Drops labels to enable classifier-free guidance. + """ + if force_drop_ids is None: + drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob + else: + drop_ids = force_drop_ids == 1 + labels = torch.where(drop_ids, self.num_classes, labels) + return labels + + def forward(self, labels, train, force_drop_ids=None): + use_dropout = self.dropout_prob > 0 + if (train and use_dropout) or (force_drop_ids is not None): + labels = self.token_drop(labels, force_drop_ids) + embeddings = self.embedding_table(labels) + return embeddings + + +################################################################################# +# Core Latte Model # +################################################################################# + +class TransformerBlock(nn.Module): + """ + A Latte tansformer block with adaptive layer norm zero (adaLN-Zero) conditioning. + """ + def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + mlp_hidden_dim = int(hidden_size * mlp_ratio) + approx_gelu = lambda: nn.GELU(approximate="tanh") + self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 6 * hidden_size, bias=True) + ) + + def forward(self, x, c): + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) + x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa)) + x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) + return x + + +class FinalLayer(nn.Module): + """ + The final layer of Latte. + """ + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True) + ) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class Latte(nn.Module): + """ + Diffusion model with a Transformer backbone. + """ + def __init__( + self, + input_size=32, + patch_size=2, + in_channels=4, + hidden_size=1152, + depth=28, + num_heads=16, + mlp_ratio=4.0, + num_frames=16, + class_dropout_prob=0.1, + num_classes=1000, + learn_sigma=True, + extras=1, + attention_mode='math', + ): + super().__init__() + self.learn_sigma = learn_sigma + self.in_channels = in_channels + self.out_channels = in_channels * 2 if learn_sigma else in_channels + self.patch_size = patch_size + self.num_heads = num_heads + self.extras = extras + self.num_frames = num_frames + self.gradient_checkpointing = False + + self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True) + self.t_embedder = TimestepEmbedder(hidden_size) + + if self.extras == 2: + self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob) + if self.extras == 78: # timestep + text_embedding + self.text_embedding_projection = nn.Sequential( + nn.SiLU(), + nn.Linear(77 * 768, hidden_size, bias=True) + ) + + num_patches = self.x_embedder.num_patches + # Will use fixed sin-cos embedding: + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False) + self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False) + self.hidden_size = hidden_size + + self.blocks = nn.ModuleList([ + TransformerBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode) for _ in range(depth) + ]) + + self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels) + self.initialize_weights() + + def initialize_weights(self): + # Initialize transformer layers: + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + self.apply(_basic_init) + + # Initialize (and freeze) pos_embed by sin-cos embedding: + pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5)) + self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) + + temp_embed = get_1d_sincos_temp_embed(self.temp_embed.shape[-1], self.temp_embed.shape[-2]) + self.temp_embed.data.copy_(torch.from_numpy(temp_embed).float().unsqueeze(0)) + + # Initialize patch_embed like nn.Linear (instead of nn.Conv2d): + w = self.x_embedder.proj.weight.data + nn.init.xavier_uniform_(w.view([w.shape[0], -1])) + nn.init.constant_(self.x_embedder.proj.bias, 0) + + if self.extras == 2: + # Initialize label embedding table: + nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02) + + # Initialize timestep embedding MLP: + nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) + + # Zero-out adaLN modulation layers in Latte blocks: + for block in self.blocks: + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + # Zero-out output layers: + nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.final_layer.linear.weight, 0) + nn.init.constant_(self.final_layer.linear.bias, 0) + + def unpatchify(self, x): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = self.out_channels + p = self.x_embedder.patch_size[0] + h = w = int(x.shape[1] ** 0.5) + assert h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], h, w, p, p, c)) + x = torch.einsum('nhwpqc->nchpwq', x) + imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p)) + return imgs + + def ckpt_wrapper(self, module): + def ckpt_forward(*inputs): + outputs = module(*inputs) + return outputs + return ckpt_forward + # @torch.cuda.amp.autocast() + # @torch.compile + def forward(self, + x, + t, + y=None, + text_embedding=None, + use_fp16=False): + """ + Forward pass of Latte. + x: (N, F, C, H, W) tensor of video inputs + t: (N,) tensor of diffusion timesteps + y: (N,) tensor of class labels + """ + if use_fp16: + x = x.to(dtype=torch.float16) + + batches, frames, channels, high, weight = x.shape + x = rearrange(x, 'b f c h w -> (b f) c h w') + x = self.x_embedder(x) + self.pos_embed + t = self.t_embedder(t, use_fp16=use_fp16) + timestep_spatial = repeat(t, 'n d -> (n c) d', c=self.temp_embed.shape[1]) + timestep_temp = repeat(t, 'n d -> (n c) d', c=self.pos_embed.shape[1]) + + if self.extras == 2: + y = self.y_embedder(y, self.training) + y_spatial = repeat(y, 'n d -> (n c) d', c=self.temp_embed.shape[1]) + y_temp = repeat(y, 'n d -> (n c) d', c=self.pos_embed.shape[1]) + elif self.extras == 78: + text_embedding = self.text_embedding_projection(text_embedding.reshape(batches, -1)) + text_embedding_spatial = repeat(text_embedding, 'n d -> (n c) d', c=self.temp_embed.shape[1]) + text_embedding_temp = repeat(text_embedding, 'n d -> (n c) d', c=self.pos_embed.shape[1]) + + for i in range(0, len(self.blocks), 2): + spatial_block, temp_block = self.blocks[i:i+2] + if self.extras == 2: + c = timestep_spatial + y_spatial + elif self.extras == 78: + c = timestep_spatial + text_embedding_spatial + else: + c = timestep_spatial + if self.gradient_checkpointing and self.training: + x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(spatial_block), x, c) + else: + x = spatial_block(x, c) + + x = rearrange(x, '(b f) t d -> (b t) f d', b=batches) + # Add Time Embedding + if i == 0: + x = x + self.temp_embed + + if self.extras == 2: + c = timestep_temp + y_temp + elif self.extras == 78: + c = timestep_temp + text_embedding_temp + else: + c = timestep_temp + + if self.gradient_checkpointing and self.training: + x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(temp_block), x, c) + else: + x = temp_block(x, c) + x = rearrange(x, '(b t) f d -> (b f) t d', b=batches) + + if self.extras == 2: + c = timestep_spatial + y_spatial + else: + c = timestep_spatial + x = self.final_layer(x, c) + x = self.unpatchify(x) + x = rearrange(x, '(b f) c h w -> b f c h w', b=batches) + return x + + def forward_with_cfg(self, x, t, y=None, cfg_scale=7.0, use_fp16=False, text_embedding=None): + """ + Forward pass of Latte, but also batches the unconditional forward pass for classifier-free guidance. + """ + # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb + half = x[: len(x) // 2] + combined = torch.cat([half, half], dim=0) + if use_fp16: + combined = combined.to(dtype=torch.float16) + model_out = self.forward(combined, t, y=y, use_fp16=use_fp16, text_embedding=text_embedding) + # For exact reproducibility reasons, we apply classifier-free guidance on only + # three channels by default. The standard approach to cfg applies it to all channels. + # This can be done by uncommenting the following line and commenting-out the line following that. + # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] + # eps, rest = model_out[:, :3], model_out[:, 3:] + eps, rest = model_out[:, :, :4, ...], model_out[:, :, 4:, ...] + cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) + half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) + eps = torch.cat([half_eps, half_eps], dim=0) + return torch.cat([eps, rest], dim=2) + + +################################################################################# +# Sine/Cosine Positional Embedding Functions # +################################################################################# +# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py + +def get_1d_sincos_temp_embed(embed_dim, length): + pos = torch.arange(0, length).unsqueeze(1) + return get_1d_sincos_pos_embed_from_grid(embed_dim, pos) + +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): + """ + 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, dtype=np.float32) + grid_w = np.arange(grid_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size, grid_size]) + 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): + assert embed_dim % 2 == 0 + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) + + emb = np.concatenate([emb_h, emb_w], axis=1) + return emb + + +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) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2. + omega = 1. / 10000**omega + + pos = pos.reshape(-1) + out = np.einsum('m,d->md', pos, omega) + + emb_sin = np.sin(out) + emb_cos = np.cos(out) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) + return emb + + +################################################################################# +# Latte Configs # +################################################################################# + +def Latte_XL_2(**kwargs): + return Latte(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs) + +def Latte_XL_4(**kwargs): + return Latte(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs) + +def Latte_XL_8(**kwargs): + return Latte(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs) + +def Latte_L_2(**kwargs): + return Latte(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs) + +def Latte_L_4(**kwargs): + return Latte(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs) + +def Latte_L_8(**kwargs): + return Latte(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs) + +def Latte_B_2(**kwargs): + return Latte(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs) + +def Latte_B_4(**kwargs): + return Latte(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs) + +def Latte_B_8(**kwargs): + return Latte(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs) + +def Latte_S_2(**kwargs): + return Latte(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs) + +def Latte_S_4(**kwargs): + return Latte(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs) + +def Latte_S_8(**kwargs): + return Latte(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs) + + +Latte_models = { + 'Latte-XL/2': Latte_XL_2, 'Latte-XL/4': Latte_XL_4, 'Latte-XL/8': Latte_XL_8, + 'Latte-L/2': Latte_L_2, 'Latte-L/4': Latte_L_4, 'Latte-L/8': Latte_L_8, + 'Latte-B/2': Latte_B_2, 'Latte-B/4': Latte_B_4, 'Latte-B/8': Latte_B_8, + 'Latte-S/2': Latte_S_2, 'Latte-S/4': Latte_S_4, 'Latte-S/8': Latte_S_8, +} + +if __name__ == '__main__': + + import torch + + device = "cuda" if torch.cuda.is_available() else "cpu" + + img = torch.randn(3, 16, 4, 32, 32).to(device) + t = torch.tensor([1, 2, 3]).to(device) + y = torch.tensor([1, 2, 3]).to(device) + network = Latte_XL_2().to(device) + from thop import profile + flops, params = profile(network, inputs=(img, t)) + print('FLOPs = ' + str(flops/1000**3) + 'G') + print('Params = ' + str(params/1000**2) + 'M') + # y_embeder = LabelEmbedder(num_classes=101, hidden_size=768, dropout_prob=0.5).to(device) + # lora.mark_only_lora_as_trainable(network) + # out = y_embeder(y, True) + # out = network(img, t, y) + # print(out.shape) \ No newline at end of file diff --git a/opensora/models/diffusion/unet/placeholder b/opensora/models/diffusion/unet/placeholder new file mode 100644 index 000000000..e69de29bb diff --git a/src/sora/models/interpolation/cfgs/AMT-G.yaml b/opensora/models/frame_interpolation/cfgs/AMT-G.yaml old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/cfgs/AMT-G.yaml rename to opensora/models/frame_interpolation/cfgs/AMT-G.yaml diff --git a/src/sora/models/interpolation/interpolation.py b/opensora/models/frame_interpolation/interpolation.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/interpolation.py rename to opensora/models/frame_interpolation/interpolation.py diff --git a/src/sora/models/interpolation/networks/AMT-G.py b/opensora/models/frame_interpolation/networks/AMT-G.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/AMT-G.py rename to opensora/models/frame_interpolation/networks/AMT-G.py diff --git a/opensora/models/frame_interpolation/networks/__init__.py b/opensora/models/frame_interpolation/networks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/opensora/models/frame_interpolation/networks/blocks/__init__.py b/opensora/models/frame_interpolation/networks/blocks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/sora/models/interpolation/networks/blocks/feat_enc.py b/opensora/models/frame_interpolation/networks/blocks/feat_enc.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/blocks/feat_enc.py rename to opensora/models/frame_interpolation/networks/blocks/feat_enc.py diff --git a/src/sora/models/interpolation/networks/blocks/ifrnet.py b/opensora/models/frame_interpolation/networks/blocks/ifrnet.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/blocks/ifrnet.py rename to opensora/models/frame_interpolation/networks/blocks/ifrnet.py diff --git a/src/sora/models/interpolation/networks/blocks/multi_flow.py b/opensora/models/frame_interpolation/networks/blocks/multi_flow.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/blocks/multi_flow.py rename to opensora/models/frame_interpolation/networks/blocks/multi_flow.py diff --git a/src/sora/models/interpolation/networks/blocks/raft.py b/opensora/models/frame_interpolation/networks/blocks/raft.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/networks/blocks/raft.py rename to opensora/models/frame_interpolation/networks/blocks/raft.py diff --git a/opensora/models/frame_interpolation/utils/__init__.py b/opensora/models/frame_interpolation/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/sora/models/interpolation/utils/build_utils.py b/opensora/models/frame_interpolation/utils/build_utils.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/utils/build_utils.py rename to opensora/models/frame_interpolation/utils/build_utils.py diff --git a/src/sora/models/interpolation/utils/dist_utils.py b/opensora/models/frame_interpolation/utils/dist_utils.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/utils/dist_utils.py rename to opensora/models/frame_interpolation/utils/dist_utils.py diff --git a/src/sora/models/interpolation/utils/flow_utils.py b/opensora/models/frame_interpolation/utils/flow_utils.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/utils/flow_utils.py rename to opensora/models/frame_interpolation/utils/flow_utils.py diff --git a/src/sora/models/interpolation/utils/utils.py b/opensora/models/frame_interpolation/utils/utils.py old mode 100755 new mode 100644 similarity index 100% rename from src/sora/models/interpolation/utils/utils.py rename to opensora/models/frame_interpolation/utils/utils.py diff --git a/opensora/models/super_resolution/placeholder b/opensora/models/super_resolution/placeholder new file mode 100644 index 000000000..e69de29bb diff --git a/opensora/sample/placeholder b/opensora/sample/placeholder new file mode 100644 index 000000000..e69de29bb diff --git a/src/train.py b/opensora/train/train.py similarity index 91% rename from src/train.py rename to opensora/train/train.py index 9e919cb6a..625355d07 100644 --- a/src/train.py +++ b/opensora/train/train.py @@ -24,11 +24,11 @@ from time import time import argparse import logging import os - -from sora.modules.ae.vqvae.videogpt.videogpt import load_vqvae as load_autoencoder -from sora.dataset.videodata import Collate, UCF101ClassConditionedDataset as Dataset -from sora.modules.diffusion.dit.models import DiT_models as Diffusion_models -from sora.modules.diffusion.dit.diffusion import create_diffusion +from opensora.dataset import getdataset +from opensora.models.ae import getae +from opensora.utils.dataset_utils import Collate +from opensora.models.diffusion.dit.models import DiT_models as Diffusion_models +from opensora.models.diffusion.dit.diffusion import create_diffusion ################################################################################# # Training Helper Functions # @@ -114,9 +114,13 @@ def main(args): logger = create_logger(None) # Create model: - vae_stride_t, vae_stride_h, vae_stride_w = [int(i) for i in args.vae[-5:].split('x')] + vae_stride_t, vae_stride_h, vae_stride_w = [int(i) for i in args.ae[-5:].split('x')] + args.vae_stride_t, args.vae_stride_h, args.vae_stride_w = vae_stride_t, vae_stride_h, vae_stride_w + args.vae_stride = args.vae_stride_h patch_size = args.model[-3:] patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + args.patch_size = patch_size_h + args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w assert args.num_frames % vae_stride_t == 0, "Image size must be divisible by 8 (for the VAE encoder)." assert args.max_image_size % vae_stride_h == 0, "Image size must be divisible by 8 (for the VAE encoder)." assert vae_stride_h == vae_stride_w, "Support now." @@ -150,8 +154,7 @@ def main(args): requires_grad(ema, False) model = DDP(model.to(device), device_ids=[rank]) diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule - # vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device) - vae = load_autoencoder(args.vae, root='./').to(device) + ae = getae(args).to(device) logger.info(f"{model}") logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}") for n, p in model.named_parameters(): @@ -162,8 +165,7 @@ def main(args): opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0) # Setup data: - dataset = Dataset(args.data_path, args.sample_rate, args.num_frames, args.max_image_size, - dynamic_frames=args.dynamic_frames) + dataset = getdataset(args) sampler = DistributedSampler( dataset, num_replicas=dist.get_world_size(), @@ -179,7 +181,7 @@ def main(args): num_workers=args.num_workers, pin_memory=True, drop_last=True, - collate_fn=Collate(args.max_image_size, vae_stride_h, patch_size_h, patch_size_t, args.num_frames) + collate_fn=Collate(args) ) logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})") @@ -203,8 +205,7 @@ def main(args): y = y.to(device) attn_mask = attn_mask.to(device) with torch.no_grad(): - # Map input images to latent space + normalize latents: - x = vae.pre_vq_conv(vae.encoder(x)) + x = ae.encode(x) # Map input images to latent space + normalize latents t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) model_kwargs = dict(y=y, attention_mask=attn_mask) loss_dict = diffusion.training_losses(model, x, t, model_kwargs) @@ -261,6 +262,7 @@ def main(args): if __name__ == "__main__": # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters). parser = argparse.ArgumentParser() + parser.add_argument("--dataset", type=str, required=True) parser.add_argument("--data-path", type=str, required=True) parser.add_argument("--results-dir", type=str, default="results") parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122") @@ -274,8 +276,8 @@ if __name__ == "__main__": # -------------------------------------- - parser.add_argument("--vae", type=str, choices=['bair_stride4x2x2', 'ucf101_stride4x4x4', - 'kinetics_stride4x4x4', 'kinetics_stride2x4x4'], + parser.add_argument("--ae", type=str, choices=['bair_stride4x2x2', 'ucf101_stride4x4x4', + 'kinetics_stride4x4x4', 'kinetics_stride2x4x4'], default="ucf101_stride4x4x4") parser.add_argument("--pt-ckpt", type=str, default=None) parser.add_argument("--sample-rate", type=int, default=4) diff --git a/opensora/utils/convert_video_offline.py b/opensora/utils/convert_video_offline.py new file mode 100644 index 000000000..94f8791e4 --- /dev/null +++ b/opensora/utils/convert_video_offline.py @@ -0,0 +1,208 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +A minimal training script for DiT using PyTorch DDP. +""" +import sys + +import imageio +import torch + +# the first flag below was False when we tested this script but True makes A100 training a lot faster: +from einops import rearrange +from tqdm import tqdm + +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from torchvision.datasets import ImageFolder +from torchvision import transforms +import numpy as np +from collections import OrderedDict +from PIL import Image +from copy import deepcopy +from glob import glob +from time import time +import argparse +import logging +import os + +from datasets import video_transforms, UCF101 +from diffusers.models import AutoencoderKL + + +################################################################################# +# Training Helper Functions # +################################################################################# + + +def center_crop_arr(pil_image, image_size): + """ + Center cropping implementation from ADM. + https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126 + """ + while min(*pil_image.size) >= 2 * image_size: + pil_image = pil_image.resize( + tuple(x // 2 for x in pil_image.size), resample=Image.BOX + ) + + scale = image_size / min(*pil_image.size) + pil_image = pil_image.resize( + tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC + ) + + arr = np.array(pil_image) + crop_y = (arr.shape[0] - image_size) // 2 + crop_x = (arr.shape[1] - image_size) // 2 + return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]) + + +################################################################################# +# Training Loop # +################################################################################# + +def main(args): + """ + Trains a new DiT model. + """ + assert torch.cuda.is_available(), "Training currently requires at least one GPU." + + device = torch.device('cuda:0') + seed = args.global_seed + torch.manual_seed(seed) + torch.cuda.set_device(device) + # args.data_path = "E:/UCF-101" + args.frame_interval = 3 + args.num_frames = 16 + + # Setup a feature folder: + os.makedirs(args.features_path, exist_ok=True) + os.makedirs(os.path.join(args.features_path, args.features_name), exist_ok=True) + # os.makedirs(os.path.join(args.features_path, 'imagenet256_labels'), exist_ok=True) + + # Create model: + assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)." + latent_size = args.image_size // 8 + vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}", cache_dir='cache_dir').to(device) + + # Setup data: + + temporal_sample = video_transforms.TemporalRandomCrop(args.num_frames * args.frame_interval) # 16 1 + transform_ucf101 = transforms.Compose([ + video_transforms.ToTensorVideo(), # TCHW + # video_transforms.RandomHorizontalFlipVideo(), + video_transforms.UCFCenterCropVideo(args.image_size), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True) + ]) + + dataset = UCF101(args, transform=transform_ucf101, temporal_sample=temporal_sample) + + + loader = DataLoader( + dataset, + batch_size=1, + shuffle=False, + num_workers=args.num_workers, + pin_memory=True, + drop_last=True + ) + + train_steps = 0 + for x in tqdm(loader): + # x = x.to(device) + x = x['video'].to(device, non_blocking=True) + # y = y.to(device) + with torch.no_grad(): + # Map input images to latent space + normalize latents: + # x = vae.encode(x).latent_dist.sample().mul_(0.18215) + b, _, _, _, _ = x.shape + x = rearrange(x, 'b f c h w -> (b f) c h w').contiguous() + B = x.shape[0] + n_sample = 64 + n_round = B // n_sample + x_enc = [] + print(x.shape) + for i in range(n_round+1): + if i*n_sample == x.shape[0]: + break + x_ = vae.encode(x[i*n_sample:(i+1)*n_sample]).latent_dist.sample().mul_(0.18215) + x_enc.append(x_) + x = torch.cat(x_enc, dim=0) + x = rearrange(x, '(b f) c h w -> b f c h w', b=b).contiguous() + + + + # b, f, c, h, w = x.shape + # samples = rearrange(x, 'b f c h w -> (b f) c h w') + # samples_dec = [] + # for i in tqdm(range(n_round+1)): + # samples_ = vae.decode(samples[i*n_sample:(i+1)*n_sample] / 0.18215).sample + # samples_dec.append(samples_) + # samples = torch.cat(samples_dec, dim=0) + # + # samples = rearrange(samples, '(b f) c h w -> b f c h w', b=b) + # + # video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous() + # + # imageio.mimwrite('1.mp4', video_, fps=8, quality=9) + # sys.exit() + + x = x.detach().cpu().numpy() # (1, 4, 32, 32) + np.save(f'{args.features_path}/{args.features_name}/{train_steps}.npy', x) + + # y = y.detach().cpu().numpy() # (1,) + # np.save(f'{args.features_path}/imagenet256_labels/{train_steps}.npy', y) + + train_steps += 1 + # print(train_steps) + + +if __name__ == "__main__": + # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters). + # parser = argparse.ArgumentParser() + # parser.add_argument("--data-path", type=str) + # parser.add_argument("--features-path", type=str, default="features") + # parser.add_argument("--features-name", type=str, default="512video") + # parser.add_argument("--results-dir", type=str, default="results") + # parser.add_argument("--image-size", type=int, choices=[128, 256, 512], default=512) + # parser.add_argument("--num-classes", type=int, default=1000) + # parser.add_argument("--epochs", type=int, default=1400) + # parser.add_argument("--global-batch-size", type=int, default=256) + # parser.add_argument("--global-seed", type=int, default=0) + # parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="mse") # Choice doesn't affect training + # parser.add_argument("--num-workers", type=int, default=0) + # parser.add_argument("--log-every", type=int, default=100) + # parser.add_argument("--ckpt-every", type=int, default=50_000) + # args = parser.parse_args() + # main(args) + device = torch.device('cuda:0') + x = torch.from_numpy(np.load('1.npy')).to(device) + vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-mse", cache_dir='cache_dir').to(device) + vae.eval() + b, f, c, h, w = x.shape + + n_sample = 1 + n_round = b*f // n_sample + samples = rearrange(x, 'b f c h w -> (b f) c h w') + samples_dec = [] + + for i in tqdm(range(n_round+1)): + with torch.no_grad(): + samples_ = vae.decode(samples[i*n_sample:(i+1)*n_sample] / 0.18215).sample + samples_dec.append(samples_.cpu()) + samples = torch.cat(samples_dec, dim=0) + + samples = rearrange(samples, '(b f) c h w -> b f c h w', b=b) + + video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).permute(0, 2, 3, 1).contiguous() + + imageio.mimwrite('1.mp4', video_, fps=10, quality=9) + + diff --git a/opensora/utils/dataset_utils.py b/opensora/utils/dataset_utils.py new file mode 100644 index 000000000..835cfcd9f --- /dev/null +++ b/opensora/utils/dataset_utils.py @@ -0,0 +1,63 @@ +import math +from torch.nn import functional as F +import torch + + +def pad_to_multiple(number, ds_stride): + remainder = number % ds_stride + if remainder == 0: + return number + else: + padding = ds_stride - remainder + return number + padding + +class Collate: + def __init__(self, args): + self.max_image_size = args.max_image_size + self.vae_stride = args.vae_stride + self.patch_size = args.patch_size + self.patch_size_t = args.patch_size_t + self.num_frames = args.num_frames + + def __call__(self, batch): + batch_tubes, labels = tuple(zip(*batch)) + labels = torch.as_tensor(labels).to(torch.long) + ds_stride = self.vae_stride * self.patch_size + t_ds_stride = self.vae_stride * self.patch_size_t + + # pad to max multiple of ds_stride + batch_input_size = [i.shape for i in batch_tubes] + max_t, max_h, max_w = max([i[1] for i in batch_input_size]), \ + max([i[2] for i in batch_input_size]), \ + max([i[3] for i in batch_input_size]) + pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t, t_ds_stride), \ + pad_to_multiple(max_h, ds_stride), \ + pad_to_multiple(max_w, ds_stride) + each_pad_t_h_w = [[pad_max_t - i.shape[1], + pad_max_h - i.shape[2], + pad_max_w - i.shape[3]] for i in batch_tubes] + pad_batch_tubes = [F.pad(im, + (0, pad_w, + 0, pad_h, + 0, pad_t), value=0) for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)] + pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0) + + # make attention_mask + max_tube_size = [pad_max_t, pad_max_h, pad_max_w] + max_latent_size = [max_tube_size[0] // self.vae_stride, + max_tube_size[1] // self.vae_stride, + max_tube_size[2] // self.vae_stride] + max_patchify_latent_size = [max_latent_size[0] // self.patch_size_t, + max_latent_size[1] // self.patch_size, + max_latent_size[2] // self.patch_size] + valid_patchify_latent_size = [[int(math.ceil(i[1] / t_ds_stride)), + int(math.ceil(i[2] / ds_stride)), + int(math.ceil(i[3] / ds_stride))] for i in batch_input_size] + attention_mask = [F.pad(torch.ones(i), + (0, max_patchify_latent_size[2] - i[2], + 0, max_patchify_latent_size[1] - i[1], + 0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size] + attention_mask = torch.stack(attention_mask) + + return pad_batch_tubes, labels, attention_mask + diff --git a/opensora/utils/movie2clip.py b/opensora/utils/movie2clip.py new file mode 100644 index 000000000..2da7a1d18 --- /dev/null +++ b/opensora/utils/movie2clip.py @@ -0,0 +1,66 @@ +import cv2 +from tqdm import tqdm + + +def crop_video(input_file, start_and_end): + # 打开视频文件 + video = cv2.VideoCapture(input_file) + + # 获取视频的帧速率和总帧数 + fps = video.get(cv2.CAP_PROP_FPS) + total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + + + current_frame = 0 + # 设置当前帧的索引 + for s, e in tqdm(start_and_end): + # 计算裁剪的起始帧和结束帧 + start_frame = int(s * fps) + end_frame = int(e * fps) + # 设置视频编解码器并创建输出视频对象 + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + output = cv2.VideoWriter(f'clip/{s}-{e}.mp4', fourcc, fps, (int(video.get(3)), int(video.get(4)))) + + + # 读取并写入视频帧，直到达到结束帧 + while current_frame < end_frame: + ret, frame = video.read() + if not ret: + break + if current_frame >= start_frame: + output.write(frame) + current_frame += 1 + + # 释放视频对象和输出对象 + video.release() + output.release() + + print("视频裁剪完成！") + +# 使用示例 +input_file = r"C:\Users\ABin\Desktop\S01E03 热带宁静Tropical Serenity_超清 4K.mp4" +start_and_end = [ + [24, 71], + [76, 178], + [182, 265], + [268, 315], + [322, 366], + [372, 467], + [473, 535], + [550, 610], + [618, 684], + [689, 792], + [800, 880], + [884, 945], + [951, 1002], + [1009, 1116], + [1127, 1196], + [1205, 1301], + [1306, 1381], + [1386, 1421], + [1430, 1505], + [1571, 1628], + [1636, 1724], + [1731, 1784], +] +crop_video(input_file, start_and_end) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..456552755 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "opensora" +version = "1.0.0" +description = "Reproduce OpenAI's Sora." +readme = "README.md" +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", +] +dependencies = [ + "torch==2.0.1", "torchvision==0.15.2", + "transformers==4.32.0", "albumentations==1.4.0", "av==11.0.0", "decord==0.6.0", "einops==0.3.0", "fastapi==0.110.0", + "accelerate==0.21.0", "gdown==5.1.0", "h5py==3.10.0", "idna==3.6", 'imageio==2.34.0', "matplotlib==3.7.5", "numpy==1.24.4", + "omegaconf==2.1.1", "opencv-python==4.9.0.80", "opencv-python-headless==4.9.0.80", "pandas==2.0.3", "pillow==10.2.0", + "pydub==0.25.1", "pytorch-lightning==1.4.2", "pytorchvideo==0.1.5", "PyYAML==6.0.1", "regex==2023.12.25", + "requests==2.31.0", "scikit-learn==1.3.2", "scipy==1.10.1", "six==1.16.0", "tensorboard==2.14.0", "test-tube==0.7.5", + "timm==0.9.16", "torchdiffeq==0.2.3", "torchmetrics==0.5.0", "tqdm==4.66.2", "urllib3==2.2.1", "uvicorn==0.27.1", + "diffusers==0.24.0", "scikit-video==1.1.11" +] + + +[project.urls] +"Homepage" = "https://github.com/PKU-YuanGroup/Open-Sora-Plan" +"Bug Tracker" = "https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues" + +[tool.setuptools.packages.find] +exclude = ["assets*", "docker*", "docs", "scripts*"] + +[tool.wheel] +exclude = ["assets*", "docker*", "docs", "scripts*"] diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 768271f62..000000000 --- a/requirements.txt +++ /dev/null @@ -1,35 +0,0 @@ -albumentations==1.4.0 -av==11.0.0 -decord==0.6.0 -einops==0.3.0 -fastapi==0.110.0 -gdown==5.1.0 -h5py==3.10.0 -idna==3.6 -imageio==2.34.0 -matplotlib==3.7.5 -numpy==1.24.4 -omegaconf==2.1.1 -opencv-python==4.9.0.80 -opencv-python-headless==4.9.0.80 -pandas==2.0.3 -pillow==10.2.0 -pydub==0.25.1 -pytorch-lightning==1.4.2 -pytorchvideo==0.1.5 -PyYAML==6.0.1 -regex==2023.12.25 -requests==2.31.0 -scikit-learn==1.3.2 -scipy==1.10.1 -six==1.16.0 -tensorboard==2.14.0 -test-tube==0.7.5 -timm==0.9.16 -torchdiffeq==0.2.3 -torchmetrics==0.5.0 -tqdm==4.66.2 -transformers==4.32.0 -urllib3==2.2.1 -uvicorn==0.27.1 -wandb==0.16.3 diff --git a/scripts/slurm/placeholder b/scripts/slurm/placeholder new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/train.sh b/scripts/train.sh index 31a35e23e..071b92c17 100644 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -1,7 +1,7 @@ -torchrun --nproc_per_node=8 src/train.py \ - --model DiT-XL/122 \ - --vae ucf101_stride4x4x4 \ +torchrun --nproc_per_node=1 --master_port=29501 opensora/train/train.py \ + --model DiT-XL/122 --dataset ucf101\ + --ae ucf101_stride4x4x4 \ --data-path /remote-home/yeyang/UCF-101 --num-classes 101 \ --sample-rate 2 --num-frames 8 --max-image-size 128 --clip-grad-norm 1 \ - --epochs 14000 --global-batch-size 256 --lr 1e-4 \ - --ckpt-every 1000 --log-every 1000 \ No newline at end of file + --epochs 14000 --global-batch-size 2 --lr 1e-4 \ + --ckpt-every 1000 --log-every 50 \ No newline at end of file diff --git a/src/sora/dataset/videodata.py b/src/sora/dataset/videodata.py deleted file mode 100644 index 8580e4003..000000000 --- a/src/sora/dataset/videodata.py +++ /dev/null @@ -1,174 +0,0 @@ -import math -import os - -import numpy as np -import torch -from decord import VideoReader, cpu -from torch.utils.data import Dataset -from torchvision.transforms import Compose, Lambda, ToTensor -from torchvision.transforms._transforms_video import NormalizeVideo, RandomCropVideo, RandomHorizontalFlipVideo -from pytorchvideo.transforms import ApplyTransformToKey, ShortSideScale, UniformTemporalSubsample -from torch.nn import functional as F -import random - - - -class LongSideScale(torch.nn.Module): - """ - ``nn.Module`` wrapper for ``pytorchvideo.transforms.functional.short_side_scale``. - """ - - def __init__( - self, size: int, interpolation: str = "bilinear" - ): - super().__init__() - self._size = size - self._interpolation = interpolation - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Args: - x (torch.Tensor): video tensor with shape (C, T, H, W). - """ - assert len(x.shape) == 4 - assert x.dtype == torch.float32 - c, t, h, w = x.shape - if w < h: - new_w = int(math.floor((float(w) / h) * self._size)) - new_h = self._size - else: - new_w = self._size - new_h = int(math.floor((float(h) / w) * self._size)) - return torch.nn.functional.interpolate(x, size=(new_h, new_w), mode=self._interpolation, align_corners=False) - -class UCF101ClassConditionedDataset(Dataset): - def __init__(self, root_dir, sample_rate, num_frames, max_image_size, dynamic_frames=False): - self.root_dir = root_dir - - self.classes = sorted(os.listdir(root_dir)) - self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)} - self.samples = self._make_dataset() - - self.sample_rate = sample_rate - self.num_frames = num_frames - self.sample_frames_len = self.sample_rate * self.num_frames - self.max_image_size = max_image_size - self.dynamic_frames = dynamic_frames - self.transform = Compose( - [ - Lambda(lambda x: ((x / 255.0) - 0.5)), - LongSideScale(size=self.max_image_size), - # RandomHorizontalFlipVideo(p=0.5), - ] - ) - - def _make_dataset(self): - dataset = [] - for class_name in self.classes: - class_path = os.path.join(self.root_dir, class_name) - for fname in os.listdir(class_path): - if fname.endswith('.avi'): - item = (os.path.join(class_path, fname), self.class_to_idx[class_name]) - dataset.append(item) - return dataset - - def __len__(self): - return len(self.samples) - - def __getitem__(self, idx): - video_path, label = self.samples[idx] - - try: - video_data = self.read_video(video_path) - video_outputs = self.transform(video_data) - # video_outputs = torch.rand(3, 16, 128, 128) - return video_outputs, label - except Exception as e: - print(f'Error with {e}, {video_path}') - return self.__getitem__(random.randint(0, self.__len__()-1)) - - - def read_video(self, video_path): - decord_vr = VideoReader(video_path, ctx=cpu(0)) - - total_frames = len(decord_vr) - if total_frames > self.sample_frames_len: - s = random.randint(0, total_frames - self.sample_frames_len - 1) - e = s + self.sample_frames_len - num_frames = self.num_frames - else: - s = 0 - e = total_frames - num_frames = int(total_frames / self.sample_frames_len * self.num_frames) - print(f'sample_frames_len {self.sample_frames_len}, only can {num_frames*self.sample_rate}', video_path, total_frames) - - # random drop to dynamic input frames - frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int) - if self.dynamic_frames and total_frames > self.sample_frames_len: # actually only second-half is dynamic, because num_frames are rare... - cut_idx = random.randint(num_frames // 2, num_frames) - frame_id_list = frame_id_list[:cut_idx] - - video_data = decord_vr.get_batch(frame_id_list).asnumpy() - video_data = torch.from_numpy(video_data) - video_data = video_data.permute(3, 0, 1, 2) # (T, H, W, C) -> (C, T, H, W) - return video_data - -def pad_to_multiple(number, ds_stride): - remainder = number % ds_stride - if remainder == 0: - return number - else: - padding = ds_stride - remainder - return number + padding - -class Collate: - def __init__(self, max_image_size, vae_stride, patch_size, patch_size_t, num_frames): - self.max_image_size = max_image_size - self.vae_stride = vae_stride - self.patch_size = patch_size - self.patch_size_t = patch_size_t - self.num_frames = num_frames - - def __call__(self, batch): - batch_tubes, labels = tuple(zip(*batch)) - labels = torch.as_tensor(labels).to(torch.long) - ds_stride = self.vae_stride * self.patch_size - t_ds_stride = self.vae_stride * self.patch_size_t - - # pad to max multiple of ds_stride - batch_input_size = [i.shape for i in batch_tubes] - max_t, max_h, max_w = max([i[1] for i in batch_input_size]), \ - max([i[2] for i in batch_input_size]), \ - max([i[3] for i in batch_input_size]) - pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t, t_ds_stride), \ - pad_to_multiple(max_h, ds_stride), \ - pad_to_multiple(max_w, ds_stride) - each_pad_t_h_w = [[pad_max_t - i.shape[1], - pad_max_h - i.shape[2], - pad_max_w - i.shape[3]] for i in batch_tubes] - pad_batch_tubes = [F.pad(im, - (0, pad_w, - 0, pad_h, - 0, pad_t), value=0) for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)] - pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0) - - # make attention_mask - max_tube_size = [pad_max_t, pad_max_h, pad_max_w] - max_latent_size = [max_tube_size[0] // self.vae_stride, - max_tube_size[1] // self.vae_stride, - max_tube_size[2] // self.vae_stride] - max_patchify_latent_size = [max_latent_size[0] // self.patch_size_t, - max_latent_size[1] // self.patch_size, - max_latent_size[2] // self.patch_size] - valid_patchify_latent_size = [[int(math.ceil(i[1] / t_ds_stride)), - int(math.ceil(i[2] / ds_stride)), - int(math.ceil(i[3] / ds_stride))] for i in batch_input_size] - attention_mask = [F.pad(torch.ones(i), - (0, max_patchify_latent_size[2] - i[2], - 0, max_patchify_latent_size[1] - i[1], - 0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size] - attention_mask = torch.stack(attention_mask) - - return pad_batch_tubes, labels, attention_mask - - diff --git a/src/sora/models/interpolation/environment.yaml b/src/sora/models/interpolation/environment.yaml deleted file mode 100755 index cd402d0bc..000000000 --- a/src/sora/models/interpolation/environment.yaml +++ /dev/null @@ -1,19 +0,0 @@ -name: amt -channels: - - pytorch - - conda-forge - - defaults -dependencies: - - python=3.8.5 - - pip=20.3 - - cudatoolkit=11.3 - - pytorch=1.11.0 - - torchvision=0.12.0 - - numpy=1.21.5 - - pip: - - opencv-python==4.1.2.30 - - imageio==2.19.3 - - omegaconf==2.3.0 - - Pillow==9.4.0 - - tqdm==4.64.1 - - wandb==0.12.21 \ No newline at end of file diff --git a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/PKG-INFO b/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/PKG-INFO deleted file mode 100644 index 3618bcc11..000000000 --- a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/PKG-INFO +++ /dev/null @@ -1,10 +0,0 @@ -Metadata-Version: 1.0 -Name: VideoGPT -Version: 1.0 -Summary: PyTorch package for VideoGPT -Home-page: http://github.com/wilson1yan/VideoGPT -Author: Wilson Yan -Author-email: wilson1.yan@berkeley.edu -License: BSD -Description: UNKNOWN -Platform: UNKNOWN diff --git a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/SOURCES.txt b/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/SOURCES.txt deleted file mode 100644 index 8a243c829..000000000 --- a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/SOURCES.txt +++ /dev/null @@ -1,15 +0,0 @@ -setup.py -VideoGPT.egg-info/PKG-INFO -VideoGPT.egg-info/SOURCES.txt -VideoGPT.egg-info/dependency_links.txt -VideoGPT.egg-info/requires.txt -VideoGPT.egg-info/top_level.txt -VideoGPT.egg-info/zip-safe -videogpt/__init__.py -videogpt/attention.py -videogpt/data.py -videogpt/download.py -videogpt/gpt.py -videogpt/resnet.py -videogpt/utils.py -videogpt/vqvae.py \ No newline at end of file diff --git a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/dependency_links.txt b/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/dependency_links.txt deleted file mode 100644 index 8b1378917..000000000 --- a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/top_level.txt b/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/top_level.txt deleted file mode 100644 index 4c03b8dcf..000000000 --- a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -videogpt diff --git a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/zip-safe b/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/zip-safe deleted file mode 100644 index 8b1378917..000000000 --- a/src/sora/modules/ae/vqvae/videogpt/VideoGPT.egg-info/zip-safe +++ /dev/null @@ -1 +0,0 @@ - diff --git a/src/sora/modules/ae/vqvae/videogpt/requirements.txt b/src/sora/modules/ae/vqvae/videogpt/requirements.txt deleted file mode 100644 index 47320c044..000000000 --- a/src/sora/modules/ae/vqvae/videogpt/requirements.txt +++ /dev/null @@ -1,9 +0,0 @@ -tqdm -requests -h5py -scikit-video -av -gradio -moviepy -imageio -gdown