From 4da0e590add7958afc37bdad2e5bb9c77185c8b9 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 31 May 2024 02:41:53 +0300 Subject: [PATCH] Initial ToonCrafter support --- configs/tooncrafter_512_interp.yaml | 103 +++ lvdm/models/autoencoder copy.py | 259 ++++++ lvdm/models/autoencoder.py | 80 +- lvdm/models/autoencoder_dualref.py | 1178 +++++++++++++++++++++++++++ lvdm/models/ddpm3d.py | 37 +- lvdm/modules/attention_svd.py | 759 +++++++++++++++++ lvdm/modules/networks/ae_modules.py | 16 +- nodes.py | 219 ++++- scripts/evaluation/funcs.py | 18 +- 9 files changed, 2640 insertions(+), 29 deletions(-) create mode 100644 configs/tooncrafter_512_interp.yaml create mode 100644 lvdm/models/autoencoder copy.py create mode 100644 lvdm/models/autoencoder_dualref.py create mode 100644 lvdm/modules/attention_svd.py diff --git a/configs/tooncrafter_512_interp.yaml b/configs/tooncrafter_512_interp.yaml new file mode 100644 index 0000000..951175d --- /dev/null +++ b/configs/tooncrafter_512_interp.yaml @@ -0,0 +1,103 @@ +model: + target: .lvdm.models.ddpm3d.LatentVisualDiffusion + params: + rescale_betas_zero_snr: True + parameterization: "v" + linear_start: 0.00085 + linear_end: 0.012 + num_timesteps_cond: 1 + timesteps: 1000 + first_stage_key: video + cond_stage_key: caption + cond_stage_trainable: False + conditioning_key: hybrid + image_size: [40, 64] + channels: 4 + scale_by_std: False + scale_factor: 0.18215 + use_ema: False + uncond_type: 'empty_seq' + use_dynamic_rescale: true + base_scale: 0.7 + fps_condition_type: 'fps' + perframe_ae: True + loop_video: true + unet_config: + target: .lvdm.modules.networks.openaimodel3d.UNetModel + params: + in_channels: 8 + out_channels: 4 + model_channels: 320 + attention_resolutions: + - 4 + - 2 + - 1 + num_res_blocks: 2 + channel_mult: + - 1 + - 2 + - 4 + - 4 + dropout: 0.1 + num_head_channels: 64 + transformer_depth: 1 + context_dim: 1024 + use_linear: true + use_checkpoint: True + temporal_conv: True + temporal_attention: True + temporal_selfatt_only: true + use_relative_position: false + use_causal_attention: False + temporal_length: 16 + addition_attention: true + image_cross_attention: true + default_fs: 24 + fs_condition: true + + first_stage_config: + target: .lvdm.models.autoencoder.AutoencoderKL_Dualref + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: True + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: .lvdm.modules.encoders.condition.FrozenOpenCLIPEmbedder + params: + freeze: true + layer: "penultimate" + + img_cond_stage_config: + target: .lvdm.modules.encoders.condition.FrozenOpenCLIPImageEmbedderV2 + params: + freeze: true + + image_proj_stage_config: + target: .lvdm.modules.encoders.resampler.Resampler + params: + dim: 1024 + depth: 4 + dim_head: 64 + heads: 12 + num_queries: 16 + embedding_dim: 1280 + output_dim: 1024 + ff_mult: 4 + video_length: 16 \ No newline at end of file diff --git a/lvdm/models/autoencoder copy.py b/lvdm/models/autoencoder copy.py new file mode 100644 index 0000000..d49e04e --- /dev/null +++ b/lvdm/models/autoencoder copy.py @@ -0,0 +1,259 @@ +import os +from contextlib import contextmanager +import torch +import numpy as np +from einops import rearrange +import torch.nn.functional as F +import pytorch_lightning as pl +from ...lvdm.modules.networks.ae_modules import Encoder, Decoder +from ...lvdm.distributions import DiagonalGaussianDistribution +from ...utils.utils import instantiate_from_config + +TIMESTEPS=16 +class AutoencoderKL(pl.LightningModule): + def __init__(self, + ddconfig, + lossconfig, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + test=False, + logdir=None, + input_dim=4, + test_args=None, + ): + super().__init__() + self.image_key = image_key + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + self.loss = instantiate_from_config(lossconfig) + assert ddconfig["double_z"] + self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + self.embed_dim = embed_dim + self.input_dim = input_dim + self.test = test + self.test_args = test_args + self.logdir = logdir + 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 + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + if self.test: + self.init_test() + + def init_test(self,): + self.test = True + save_dir = os.path.join(self.logdir, "test") + if 'ckpt' in self.test_args: + ckpt_name = os.path.basename(self.test_args.ckpt).split('.ckpt')[0] + f'_epoch{self._cur_epoch}' + self.root = os.path.join(save_dir, ckpt_name) + else: + self.root = save_dir + if 'test_subdir' in self.test_args: + self.root = os.path.join(save_dir, self.test_args.test_subdir) + + self.root_zs = os.path.join(self.root, "zs") + self.root_dec = os.path.join(self.root, "reconstructions") + self.root_inputs = os.path.join(self.root, "inputs") + os.makedirs(self.root, exist_ok=True) + + if self.test_args.save_z: + os.makedirs(self.root_zs, exist_ok=True) + if self.test_args.save_reconstruction: + os.makedirs(self.root_dec, exist_ok=True) + if self.test_args.save_input: + os.makedirs(self.root_inputs, exist_ok=True) + assert(self.test_args is not None) + self.test_maximum = getattr(self.test_args, 'test_maximum', None) + self.count = 0 + self.eval_metrics = {} + self.decodes = [] + self.save_decode_samples = 2048 + + def init_from_ckpt(self, path, ignore_keys=list()): + sd = torch.load(path, map_location="cpu") + try: + self._cur_epoch = sd['epoch'] + sd = sd["state_dict"] + except: + self._cur_epoch = 'null' + 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) + # self.load_state_dict(sd, strict=True) + print(f"Restored from {path}") + + def encode(self, x, **kwargs): + + h = self.encoder(x) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z, **kwargs): + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward(self, input, sample_posterior=True): + posterior = self.encode(input) + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z) + return dec, posterior + + def get_input(self, batch, k): + x = batch[k] + if x.dim() == 5 and self.input_dim == 4: + b,c,t,h,w = x.shape + self.b = b + self.t = t + x = rearrange(x, 'b c t h w -> (b t) c h w') + + return x + + def training_step(self, batch, batch_idx, optimizer_idx): + inputs = self.get_input(batch, self.image_key) + reconstructions, posterior = self(inputs) + + if optimizer_idx == 0: + # train encoder+decoder+logvar + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log("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=False) + return aeloss + + if optimizer_idx == 1: + # train the discriminator + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + + self.log("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=False) + return discloss + + def validation_step(self, batch, batch_idx): + inputs = self.get_input(batch, self.image_key) + reconstructions, posterior = self(inputs) + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step, + last_layer=self.get_last_layer(), split="val") + + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step, + last_layer=self.get_last_layer(), split="val") + + self.log("val/rec_loss", log_dict_ae["val/rec_loss"]) + 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.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 + + @torch.no_grad() + def log_images(self, batch, only_inputs=False, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + if not only_inputs: + xrec, posterior = 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["samples"] = self.decode(torch.randn_like(posterior.sample())) + log["reconstructions"] = xrec + log["inputs"] = x + 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 IdentityFirstStage(torch.nn.Module): + def __init__(self, *args, vq_interface=False, **kwargs): + self.vq_interface = vq_interface # TODO: Should be true by default but check to not break older stuff + super().__init__() + + def encode(self, x, *args, **kwargs): + return x + + def decode(self, x, *args, **kwargs): + return x + + def quantize(self, x, *args, **kwargs): + if self.vq_interface: + return x, None, [None, None, None] + return x + + def forward(self, x, *args, **kwargs): + return x + +from ...lvdm.models.autoencoder_dualref import VideoDecoder +class AutoencoderKL_Dualref(AutoencoderKL): + def __init__(self, + ddconfig, + lossconfig, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + test=False, + logdir=None, + input_dim=4, + test_args=None, + additional_decode_keys=None, + use_checkpoint=False, + diff_boost_factor=3.0, + ): + super().__init__(ddconfig, lossconfig, embed_dim, ckpt_path, ignore_keys, image_key, colorize_nlabels, monitor, test, logdir, input_dim, test_args, additional_decode_keys, use_checkpoint, diff_boost_factor) + self.decoder = VideoDecoder(**ddconfig) + + def _forward(self, input, sample_posterior=True, **additional_decode_kwargs): + posterior, hidden_states = self.encode(input, return_hidden_states=True) + + hidden_states_first_last = [] + ### use only the first and last hidden states + for hid in hidden_states: + hid = rearrange(hid, '(b t) c h w -> b c t h w', t=TIMESTEPS) + hid_new = torch.cat([hid[:, :, 0:1], hid[:, :, -1:]], dim=2) + hidden_states_first_last.append(hid_new) + + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z, ref_context=hidden_states_first_last, **additional_decode_kwargs) + ## print(input.shape, dec.shape) torch.Size([16, 3, 256, 256]) torch.Size([16, 3, 256, 256]) + return dec, posterior \ No newline at end of file diff --git a/lvdm/models/autoencoder.py b/lvdm/models/autoencoder.py index 0552ea8..baa736e 100644 --- a/lvdm/models/autoencoder.py +++ b/lvdm/models/autoencoder.py @@ -9,7 +9,7 @@ from ...lvdm.modules.networks.ae_modules import Encoder, Decoder from ...lvdm.distributions import DiagonalGaussianDistribution from ...utils.utils import instantiate_from_config - +TIMESTEPS=16 class AutoencoderKL(pl.LightningModule): def __init__(self, ddconfig, @@ -24,6 +24,9 @@ class AutoencoderKL(pl.LightningModule): logdir=None, input_dim=4, test_args=None, + additional_decode_keys=None, + use_checkpoint=False, + diff_boost_factor=3.0, ): super().__init__() self.image_key = image_key @@ -94,25 +97,38 @@ class AutoencoderKL(pl.LightningModule): # self.load_state_dict(sd, strict=True) print(f"Restored from {path}") - def encode(self, x, **kwargs): - - h = self.encoder(x) - moments = self.quant_conv(h) - posterior = DiagonalGaussianDistribution(moments) - return posterior + def encode(self, x, return_hidden_states=False, **kwargs): + if return_hidden_states: + h, hidden = self.encoder(x, return_hidden_states) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior, hidden + else: + h = self.encoder(x) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior def decode(self, z, **kwargs): - z = self.post_quant_conv(z) - dec = self.decoder(z) + if len(kwargs) == 0: ## use the original decoder in AutoencoderKL + z = self.post_quant_conv(z) + dec = self.decoder(z, **kwargs) ##change for SVD decoder by adding **kwargs return dec - def forward(self, input, sample_posterior=True): + def forward(self, input, sample_posterior=True, **additional_decode_kwargs): + input_tuple = (input, ) + forward_temp = partial(self._forward, sample_posterior=sample_posterior, **additional_decode_kwargs) + return checkpoint(forward_temp, input_tuple, self.parameters(), self.use_checkpoint) + + + def _forward(self, input, sample_posterior=True, **additional_decode_kwargs): posterior = self.encode(input) if sample_posterior: z = posterior.sample() else: z = posterior.mode() - dec = self.decode(z) + dec = self.decode(z, **additional_decode_kwargs) + ## print(input.shape, dec.shape) torch.Size([16, 3, 256, 256]) torch.Size([16, 3, 256, 256]) return dec, posterior def get_input(self, batch, k): @@ -216,4 +232,44 @@ class IdentityFirstStage(torch.nn.Module): return x def forward(self, x, *args, **kwargs): - return x \ No newline at end of file + return x + +from ...lvdm.models.autoencoder_dualref import VideoDecoder +class AutoencoderKL_Dualref(AutoencoderKL): + def __init__(self, + ddconfig, + lossconfig, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + test=False, + logdir=None, + input_dim=4, + test_args=None, + additional_decode_keys=None, + use_checkpoint=False, + diff_boost_factor=3.0, + ): + super().__init__(ddconfig, lossconfig, embed_dim, ckpt_path, ignore_keys, image_key, colorize_nlabels, monitor, test, logdir, input_dim, test_args, additional_decode_keys, use_checkpoint, diff_boost_factor) + self.decoder = VideoDecoder(**ddconfig) + + def _forward(self, input, sample_posterior=True, **additional_decode_kwargs): + posterior, hidden_states = self.encode(input, return_hidden_states=True) + + hidden_states_first_last = [] + ### use only the first and last hidden states + for hid in hidden_states: + hid = rearrange(hid, '(b t) c h w -> b c t h w', t=TIMESTEPS) + hid_new = torch.cat([hid[:, :, 0:1], hid[:, :, -1:]], dim=2) + hidden_states_first_last.append(hid_new) + + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z, ref_context=hidden_states_first_last, **additional_decode_kwargs) + ## print(input.shape, dec.shape) torch.Size([16, 3, 256, 256]) torch.Size([16, 3, 256, 256]) + return dec, posterior \ No newline at end of file diff --git a/lvdm/models/autoencoder_dualref.py b/lvdm/models/autoencoder_dualref.py new file mode 100644 index 0000000..55e8db1 --- /dev/null +++ b/lvdm/models/autoencoder_dualref.py @@ -0,0 +1,1178 @@ +#### https://github.com/Stability-AI/generative-models +from einops import rearrange, repeat +import logging +from typing import Any, Callable, Optional, Iterable, Union + +import numpy as np +import torch +import torch.nn as nn +from packaging import version +logpy = logging.getLogger(__name__) + +try: + import xformers + import xformers.ops + + XFORMERS_IS_AVAILABLE = True +except: + XFORMERS_IS_AVAILABLE = False + logpy.warning("no module 'xformers'. Processing without...") + +from ...lvdm.modules.attention_svd import LinearAttention, MemoryEfficientCrossAttention + + +def nonlinearity(x): + # swish + return x * torch.sigmoid(x) + + +def Normalize(in_channels, num_groups=32): + return torch.nn.GroupNorm( + num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) + + +class 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 LinAttnBlock(LinearAttention): + """to match AttnBlock usage""" + + def __init__(self, in_channels): + super().__init__(dim=in_channels, heads=1, dim_head=in_channels) + + +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 attention(self, h_: torch.Tensor) -> torch.Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, h, w = q.shape + q, k, v = map( + lambda x: rearrange(x, "b c h w -> b 1 (h w) c").contiguous(), (q, k, v) + ) + h_ = torch.nn.functional.scaled_dot_product_attention( + q, k, v + ) # scale is dim ** -0.5 per default + # compute attention + + return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) + + def forward(self, x, **kwargs): + h_ = x + h_ = self.attention(h_) + h_ = self.proj_out(h_) + return x + h_ + + +class MemoryEfficientAttnBlock(nn.Module): + """ + Uses xformers efficient implementation, + see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 + Note: this is a single-head self-attention operation + """ + + # + 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 + ) + self.attention_op: Optional[Any] = None + + def attention(self, h_: torch.Tensor) -> torch.Tensor: + 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, k, v = map(lambda x: rearrange(x, "b c h w -> b (h w) c"), (q, k, v)) + + q, k, v = map( + lambda t: t.unsqueeze(3) + .reshape(B, t.shape[1], 1, C) + .permute(0, 2, 1, 3) + .reshape(B * 1, t.shape[1], C) + .contiguous(), + (q, k, v), + ) + out = xformers.ops.memory_efficient_attention( + q, k, v, attn_bias=None, op=self.attention_op + ) + + out = ( + out.unsqueeze(0) + .reshape(B, 1, out.shape[1], C) + .permute(0, 2, 1, 3) + .reshape(B, out.shape[1], C) + ) + return rearrange(out, "b (h w) c -> b c h w", b=B, h=H, w=W, c=C) + + def forward(self, x, **kwargs): + h_ = x + h_ = self.attention(h_) + h_ = self.proj_out(h_) + return x + h_ + + +class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention): + def forward(self, x, context=None, mask=None, **unused_kwargs): + b, c, h, w = x.shape + x = rearrange(x, "b c h w -> b (h w) c") + out = super().forward(x, context=context, mask=mask) + out = rearrange(out, "b (h w) c -> b c h w", h=h, w=w, c=c) + return x + out + + +def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): + print(attn_type) + assert attn_type in [ + "vanilla", + "vanilla-xformers", + "memory-efficient-cross-attn", + "linear", + "none", + "memory-efficient-cross-attn-fusion", + ], f"attn_type {attn_type} unknown" + if ( + version.parse(torch.__version__) < version.parse("2.0.0") + and attn_type != "none" + ): + assert XFORMERS_IS_AVAILABLE, ( + f"We do not support vanilla attention in {torch.__version__} anymore, " + f"as it is too expensive. Please install xformers via e.g. 'pip install xformers==0.0.16'" + ) + # attn_type = "vanilla-xformers" + logpy.info(f"making attention of type '{attn_type}' with {in_channels} in_channels") + if attn_type == "vanilla": + assert attn_kwargs is None + return AttnBlock(in_channels) + elif attn_type == "vanilla-xformers": + logpy.info( + f"building MemoryEfficientAttnBlock with {in_channels} in_channels..." + ) + return MemoryEfficientAttnBlock(in_channels) + elif attn_type == "memory-efficient-cross-attn": + attn_kwargs["query_dim"] = in_channels + return MemoryEfficientCrossAttentionWrapper(**attn_kwargs) + elif attn_type == "memory-efficient-cross-attn-fusion": + attn_kwargs["query_dim"] = in_channels + return MemoryEfficientCrossAttentionWrapperFusion(**attn_kwargs) + elif attn_type == "none": + return nn.Identity(in_channels) + else: + return LinAttnBlock(in_channels) + +class MemoryEfficientCrossAttentionWrapperFusion(MemoryEfficientCrossAttention): + # print('x.shape: ',x.shape, 'context.shape: ',context.shape) ##torch.Size([8, 128, 256, 256]) torch.Size([1, 128, 2, 256, 256]) + def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0, **kwargs): + super().__init__(query_dim, context_dim, heads, dim_head, dropout, **kwargs) + self.norm = Normalize(query_dim) + nn.init.zeros_(self.to_out[0].weight) + nn.init.zeros_(self.to_out[0].bias) + + def forward(self, x, context=None, mask=None): + if self.training: + return checkpoint(self._forward, x, context, mask, use_reentrant=False) + else: + return self._forward(x, context, mask) + + def _forward( + self, + x, + context=None, + mask=None, + ): + bt, c, h, w = x.shape + h_ = self.norm(x) + h_ = rearrange(h_, "b c h w -> b (h w) c") + q = self.to_q(h_) + + + b, c, l, h, w = context.shape + context = rearrange(context, "b c l h w -> (b l) (h w) c") + k = self.to_k(context) + v = self.to_v(context) + k = rearrange(k, "(b l) d c -> b l d c", l=l) + k = torch.cat([k[:, [0] * (bt//b)], k[:, [1]*(bt//b)]], dim=2) + k = rearrange(k, "b l d c -> (b l) d c") + + v = rearrange(v, "(b l) d c -> b l d c", l=l) + v = torch.cat([v[:, [0] * (bt//b)], v[:, [1]*(bt//b)]], dim=2) + v = rearrange(v, "b l d c -> (b l) d c") + + + b, _, _ = q.shape ##actually bt + q, k, v = map( + lambda t: t.unsqueeze(3) + .reshape(b, t.shape[1], self.heads, self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b * self.heads, t.shape[1], self.dim_head) + .contiguous(), + (q, k, v), + ) + + # actually compute the attention, what we cannot get enough of + if version.parse(xformers.__version__) >= version.parse("0.0.21"): + # NOTE: workaround for + # https://github.com/facebookresearch/xformers/issues/845 + max_bs = 32768 + N = q.shape[0] + n_batches = math.ceil(N / max_bs) + out = list() + for i_batch in range(n_batches): + batch = slice(i_batch * max_bs, (i_batch + 1) * max_bs) + out.append( + xformers.ops.memory_efficient_attention( + q[batch], + k[batch], + v[batch], + attn_bias=None, + op=self.attention_op, + ) + ) + out = torch.cat(out, 0) + else: + out = xformers.ops.memory_efficient_attention( + q, k, v, attn_bias=None, op=self.attention_op + ) + + # TODO: Use this directly in the attention operation, as a bias + if exists(mask): + raise NotImplementedError + out = ( + out.unsqueeze(0) + .reshape(b, self.heads, out.shape[1], self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b, out.shape[1], self.heads * self.dim_head) + ) + out = self.to_out(out) + out = rearrange(out, "bt (h w) c -> bt c h w", h=h, w=w, c=c) + return x + out + +class Combiner(nn.Module): + def __init__(self, ch) -> None: + super().__init__() + self.conv = nn.Conv2d(ch,ch,1,padding=0) + + nn.init.zeros_(self.conv.weight) + nn.init.zeros_(self.conv.bias) + + def forward(self, x, context): + if self.training: + return checkpoint(self._forward, x, context, use_reentrant=False) + else: + return self._forward(x, context) + + def _forward(self, x, context): + ## x: b c h w, context: b c 2 h w + b, c, l, h, w = context.shape + bt, c, h, w = x.shape + context = rearrange(context, "b c l h w -> (b l) c h w") + context = self.conv(context) + context = rearrange(context, "(b l) c h w -> b c l h w", l=l) + x = rearrange(x, "(b t) c h w -> b c t h w", t=bt//b) + x[:,:,0] = x[:,:,0] + context[:,:,0] + x[:,:,-1] = x[:,:,-1] + context[:,:,1] + x = rearrange(x, "b c t h w -> (b t) c h w") + return x + + +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, + tanh_out=False, + use_linear_attn=False, + attn_type="vanilla", + attn_level=[2,3], + **ignorekwargs, + ): + super().__init__() + if use_linear_attn: + attn_type = "linear" + 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 + self.tanh_out = tanh_out + self.attn_level = attn_level + # 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) + logpy.info( + "Working with z of shape {} = {} dimensions.".format( + self.z_shape, np.prod(self.z_shape) + ) + ) + + make_attn_cls = self._make_attn() + make_resblock_cls = self._make_resblock() + make_conv_cls = self._make_conv() + # 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 = make_resblock_cls( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + ) + self.mid.attn_1 = make_attn_cls(block_in, attn_type=attn_type) + self.mid.block_2 = make_resblock_cls( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + ) + + # upsampling + self.up = nn.ModuleList() + self.attn_refinement = 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( + make_resblock_cls( + 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(make_attn_cls(block_in, attn_type=attn_type)) + 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 + + if i_level in self.attn_level: + self.attn_refinement.insert(0, make_attn_cls(block_in, attn_type='memory-efficient-cross-attn-fusion', attn_kwargs={})) + else: + self.attn_refinement.insert(0, Combiner(block_in)) + # end + self.norm_out = Normalize(block_in) + self.attn_refinement.append(Combiner(block_in)) + self.conv_out = make_conv_cls( + block_in, out_ch, kernel_size=3, stride=1, padding=1 + ) + + def _make_attn(self) -> Callable: + return make_attn + + def _make_resblock(self) -> Callable: + return ResnetBlock + + def _make_conv(self) -> Callable: + return torch.nn.Conv2d + + def get_last_layer(self, **kwargs): + return self.conv_out.weight + + def forward(self, z, ref_context=None, **kwargs): + ## ref_context: b c 2 h w, 2 means starting and ending frame + # 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, **kwargs) + h = self.mid.attn_1(h, **kwargs) + h = self.mid.block_2(h, temb, **kwargs) + + # 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, **kwargs) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h, **kwargs) + if ref_context: + h = self.attn_refinement[i_level](x=h, context=ref_context[i_level]) + 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) + if ref_context: + # print(h.shape, ref_context[i_level].shape) #torch.Size([8, 128, 256, 256]) torch.Size([1, 128, 2, 256, 256]) + h = self.attn_refinement[-1](x=h, context=ref_context[-1]) + h = self.conv_out(h, **kwargs) + if self.tanh_out: + h = torch.tanh(h) + return h + +##### + + +from abc import abstractmethod +from ...lvdm.models.utils_diffusion import timestep_embedding + +from torch.utils.checkpoint import checkpoint +from ...lvdm.basics import ( + zero_module, + conv_nd, + linear, + normalization, +) +from ...lvdm.modules.networks.openaimodel3d import Upsample, Downsample +class TimestepBlock(nn.Module): + """ + Any module where forward() takes timestep embeddings as a second argument. + """ + + @abstractmethod + def forward(self, x: torch.Tensor, emb: torch.Tensor): + """ + Apply the module to `x` given `emb` timestep embeddings. + """ + +class ResBlock(TimestepBlock): + """ + A residual block that can optionally change the number of channels. + :param channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param use_checkpoint: if True, use gradient checkpointing on this module. + :param up: if True, use this block for upsampling. + :param down: if True, use this block for downsampling. + """ + + def __init__( + self, + channels: int, + emb_channels: int, + dropout: float, + out_channels: Optional[int] = None, + use_conv: bool = False, + use_scale_shift_norm: bool = False, + dims: int = 2, + use_checkpoint: bool = False, + up: bool = False, + down: bool = False, + kernel_size: int = 3, + exchange_temb_dims: bool = False, + skip_t_emb: bool = False, + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.dropout = dropout + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_checkpoint = use_checkpoint + self.use_scale_shift_norm = use_scale_shift_norm + self.exchange_temb_dims = exchange_temb_dims + + if isinstance(kernel_size, Iterable): + padding = [k // 2 for k in kernel_size] + else: + padding = kernel_size // 2 + + self.in_layers = nn.Sequential( + normalization(channels), + nn.SiLU(), + conv_nd(dims, channels, self.out_channels, kernel_size, padding=padding), + ) + + self.updown = up or down + + if up: + self.h_upd = Upsample(channels, False, dims) + self.x_upd = Upsample(channels, False, dims) + elif down: + self.h_upd = Downsample(channels, False, dims) + self.x_upd = Downsample(channels, False, dims) + else: + self.h_upd = self.x_upd = nn.Identity() + + self.skip_t_emb = skip_t_emb + self.emb_out_channels = ( + 2 * self.out_channels if use_scale_shift_norm else self.out_channels + ) + if self.skip_t_emb: + # print(f"Skipping timestep embedding in {self.__class__.__name__}") + assert not self.use_scale_shift_norm + self.emb_layers = None + self.exchange_temb_dims = False + else: + self.emb_layers = nn.Sequential( + nn.SiLU(), + linear( + emb_channels, + self.emb_out_channels, + ), + ) + + self.out_layers = nn.Sequential( + normalization(self.out_channels), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module( + conv_nd( + dims, + self.out_channels, + self.out_channels, + kernel_size, + padding=padding, + ) + ), + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd( + dims, channels, self.out_channels, kernel_size, padding=padding + ) + else: + self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + :param x: an [N x C x ...] Tensor of features. + :param emb: an [N x emb_channels] Tensor of timestep embeddings. + :return: an [N x C x ...] Tensor of outputs. + """ + if self.use_checkpoint: + return checkpoint(self._forward, x, emb, use_reentrant=False) + else: + return self._forward(x, emb) + + def _forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + if self.skip_t_emb: + emb_out = torch.zeros_like(h) + else: + emb_out = self.emb_layers(emb).type(h.dtype) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + if self.use_scale_shift_norm: + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1 + scale) + shift + h = out_rest(h) + else: + if self.exchange_temb_dims: + emb_out = rearrange(emb_out, "b t c ... -> b c t ...") + h = h + emb_out + h = self.out_layers(h) + return self.skip_connection(x) + h +##### + +##### +from ...lvdm.modules.attention_svd import * +class VideoTransformerBlock(nn.Module): + ATTENTION_MODES = { + "softmax": CrossAttention, + "softmax-xformers": MemoryEfficientCrossAttention, + } + + def __init__( + self, + dim, + n_heads, + d_head, + dropout=0.0, + context_dim=None, + gated_ff=True, + checkpoint=True, + timesteps=None, + ff_in=False, + inner_dim=None, + attn_mode="softmax", + disable_self_attn=False, + disable_temporal_crossattention=False, + switch_temporal_ca_to_sa=False, + ): + super().__init__() + + attn_cls = self.ATTENTION_MODES[attn_mode] + + self.ff_in = ff_in or inner_dim is not None + if inner_dim is None: + inner_dim = dim + + assert int(n_heads * d_head) == inner_dim + + self.is_res = inner_dim == dim + + if self.ff_in: + self.norm_in = nn.LayerNorm(dim) + self.ff_in = FeedForward( + dim, dim_out=inner_dim, dropout=dropout, glu=gated_ff + ) + + self.timesteps = timesteps + self.disable_self_attn = disable_self_attn + if self.disable_self_attn: + self.attn1 = attn_cls( + query_dim=inner_dim, + heads=n_heads, + dim_head=d_head, + context_dim=context_dim, + dropout=dropout, + ) # is a cross-attention + else: + self.attn1 = attn_cls( + query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout + ) # is a self-attention + + self.ff = FeedForward(inner_dim, dim_out=dim, dropout=dropout, glu=gated_ff) + + if disable_temporal_crossattention: + if switch_temporal_ca_to_sa: + raise ValueError + else: + self.attn2 = None + else: + self.norm2 = nn.LayerNorm(inner_dim) + if switch_temporal_ca_to_sa: + self.attn2 = attn_cls( + query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout + ) # is a self-attention + else: + self.attn2 = attn_cls( + query_dim=inner_dim, + context_dim=context_dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + ) # is self-attn if context is none + + self.norm1 = nn.LayerNorm(inner_dim) + self.norm3 = nn.LayerNorm(inner_dim) + self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa + + self.checkpoint = checkpoint + if self.checkpoint: + print(f"====>{self.__class__.__name__} is using checkpointing") + else: + print(f"====>{self.__class__.__name__} is NOT using checkpointing") + + def forward( + self, x: torch.Tensor, context: torch.Tensor = None, timesteps: int = None + ) -> torch.Tensor: + if self.checkpoint: + return checkpoint(self._forward, x, context, timesteps, use_reentrant=False) + else: + return self._forward(x, context, timesteps=timesteps) + + def _forward(self, x, context=None, timesteps=None): + assert self.timesteps or timesteps + assert not (self.timesteps and timesteps) or self.timesteps == timesteps + timesteps = self.timesteps or timesteps + B, S, C = x.shape + x = rearrange(x, "(b t) s c -> (b s) t c", t=timesteps) + + if self.ff_in: + x_skip = x + x = self.ff_in(self.norm_in(x)) + if self.is_res: + x += x_skip + + if self.disable_self_attn: + x = self.attn1(self.norm1(x), context=context) + x + else: + x = self.attn1(self.norm1(x)) + x + + if self.attn2 is not None: + if self.switch_temporal_ca_to_sa: + x = self.attn2(self.norm2(x)) + x + else: + x = self.attn2(self.norm2(x), context=context) + x + x_skip = x + x = self.ff(self.norm3(x)) + if self.is_res: + x += x_skip + + x = rearrange( + x, "(b s) t c -> (b t) s c", s=S, b=B // timesteps, c=C, t=timesteps + ) + return x + + def get_last_layer(self): + return self.ff.net[-1].weight + +##### + +##### +import functools +def partialclass(cls, *args, **kwargs): + class NewCls(cls): + __init__ = functools.partialmethod(cls.__init__, *args, **kwargs) + + return NewCls +###### + +class VideoResBlock(ResnetBlock): + def __init__( + self, + out_channels, + *args, + dropout=0.0, + video_kernel_size=3, + alpha=0.0, + merge_strategy="learned", + **kwargs, + ): + super().__init__(out_channels=out_channels, dropout=dropout, *args, **kwargs) + if video_kernel_size is None: + video_kernel_size = [3, 1, 1] + self.time_stack = ResBlock( + channels=out_channels, + emb_channels=0, + dropout=dropout, + dims=3, + use_scale_shift_norm=False, + use_conv=False, + up=False, + down=False, + kernel_size=video_kernel_size, + use_checkpoint=True, + skip_t_emb=True, + ) + + self.merge_strategy = merge_strategy + if self.merge_strategy == "fixed": + self.register_buffer("mix_factor", torch.Tensor([alpha])) + elif self.merge_strategy == "learned": + self.register_parameter( + "mix_factor", torch.nn.Parameter(torch.Tensor([alpha])) + ) + else: + raise ValueError(f"unknown merge strategy {self.merge_strategy}") + + def get_alpha(self, bs): + if self.merge_strategy == "fixed": + return self.mix_factor + elif self.merge_strategy == "learned": + return torch.sigmoid(self.mix_factor) + else: + raise NotImplementedError() + + def forward(self, x, temb, skip_video=False, timesteps=None): + if timesteps is None: + timesteps = self.timesteps + + b, c, h, w = x.shape + + x = super().forward(x, temb) + + if not skip_video: + x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps) + + x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps) + + x = self.time_stack(x, temb) + + alpha = self.get_alpha(bs=b // timesteps) + x = alpha * x + (1.0 - alpha) * x_mix + + x = rearrange(x, "b c t h w -> (b t) c h w") + return x + + +class AE3DConv(torch.nn.Conv2d): + def __init__(self, in_channels, out_channels, video_kernel_size=3, *args, **kwargs): + super().__init__(in_channels, out_channels, *args, **kwargs) + if isinstance(video_kernel_size, Iterable): + padding = [int(k // 2) for k in video_kernel_size] + else: + padding = int(video_kernel_size // 2) + + self.time_mix_conv = torch.nn.Conv3d( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=video_kernel_size, + padding=padding, + ) + + def forward(self, input, timesteps, skip_video=False): + x = super().forward(input) + if skip_video: + return x + x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps) + x = self.time_mix_conv(x) + return rearrange(x, "b c t h w -> (b t) c h w") + + +class VideoBlock(AttnBlock): + def __init__( + self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned" + ): + super().__init__(in_channels) + # no context, single headed, as in base class + self.time_mix_block = VideoTransformerBlock( + dim=in_channels, + n_heads=1, + d_head=in_channels, + checkpoint=True, + ff_in=True, + attn_mode="softmax", + ) + + time_embed_dim = self.in_channels * 4 + self.video_time_embed = torch.nn.Sequential( + torch.nn.Linear(self.in_channels, time_embed_dim), + torch.nn.SiLU(), + torch.nn.Linear(time_embed_dim, self.in_channels), + ) + + self.merge_strategy = merge_strategy + if self.merge_strategy == "fixed": + self.register_buffer("mix_factor", torch.Tensor([alpha])) + elif self.merge_strategy == "learned": + self.register_parameter( + "mix_factor", torch.nn.Parameter(torch.Tensor([alpha])) + ) + else: + raise ValueError(f"unknown merge strategy {self.merge_strategy}") + + def forward(self, x, timesteps, skip_video=False): + if skip_video: + return super().forward(x) + + x_in = x + x = self.attention(x) + h, w = x.shape[2:] + x = rearrange(x, "b c h w -> b (h w) c") + + x_mix = x + num_frames = torch.arange(timesteps, device=x.device) + num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps) + num_frames = rearrange(num_frames, "b t -> (b t)") + t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False) + emb = self.video_time_embed(t_emb) # b, n_channels + emb = emb[:, None, :] + x_mix = x_mix + emb + + alpha = self.get_alpha() + x_mix = self.time_mix_block(x_mix, timesteps=timesteps) + x = alpha * x + (1.0 - alpha) * x_mix # alpha merge + + x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w) + x = self.proj_out(x) + + return x_in + x + + def get_alpha( + self, + ): + if self.merge_strategy == "fixed": + return self.mix_factor + elif self.merge_strategy == "learned": + return torch.sigmoid(self.mix_factor) + else: + raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}") + + +class MemoryEfficientVideoBlock(MemoryEfficientAttnBlock): + def __init__( + self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned" + ): + super().__init__(in_channels) + # no context, single headed, as in base class + self.time_mix_block = VideoTransformerBlock( + dim=in_channels, + n_heads=1, + d_head=in_channels, + checkpoint=True, + ff_in=True, + attn_mode="softmax-xformers", + ) + + time_embed_dim = self.in_channels * 4 + self.video_time_embed = torch.nn.Sequential( + torch.nn.Linear(self.in_channels, time_embed_dim), + torch.nn.SiLU(), + torch.nn.Linear(time_embed_dim, self.in_channels), + ) + + self.merge_strategy = merge_strategy + if self.merge_strategy == "fixed": + self.register_buffer("mix_factor", torch.Tensor([alpha])) + elif self.merge_strategy == "learned": + self.register_parameter( + "mix_factor", torch.nn.Parameter(torch.Tensor([alpha])) + ) + else: + raise ValueError(f"unknown merge strategy {self.merge_strategy}") + + def forward(self, x, timesteps, skip_time_block=False): + if skip_time_block: + return super().forward(x) + + x_in = x + x = self.attention(x) + h, w = x.shape[2:] + x = rearrange(x, "b c h w -> b (h w) c") + + x_mix = x + num_frames = torch.arange(timesteps, device=x.device) + num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps) + num_frames = rearrange(num_frames, "b t -> (b t)") + t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False) + emb = self.video_time_embed(t_emb) # b, n_channels + emb = emb[:, None, :] + x_mix = x_mix + emb + + alpha = self.get_alpha() + x_mix = self.time_mix_block(x_mix, timesteps=timesteps) + x = alpha * x + (1.0 - alpha) * x_mix # alpha merge + + x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w) + x = self.proj_out(x) + + return x_in + x + + def get_alpha( + self, + ): + if self.merge_strategy == "fixed": + return self.mix_factor + elif self.merge_strategy == "learned": + return torch.sigmoid(self.mix_factor) + else: + raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}") + + +def make_time_attn( + in_channels, + attn_type="vanilla", + attn_kwargs=None, + alpha: float = 0, + merge_strategy: str = "learned", +): + assert attn_type in [ + "vanilla", + "vanilla-xformers", + ], f"attn_type {attn_type} not supported for spatio-temporal attention" + print( + f"making spatial and temporal attention of type '{attn_type}' with {in_channels} in_channels" + ) + if not XFORMERS_IS_AVAILABLE and attn_type == "vanilla-xformers": + print( + f"Attention mode '{attn_type}' is not available. Falling back to vanilla attention. " + f"This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}" + ) + attn_type = "vanilla" + + if attn_type == "vanilla": + assert attn_kwargs is None + return partialclass( + VideoBlock, in_channels, alpha=alpha, merge_strategy=merge_strategy + ) + elif attn_type == "vanilla-xformers": + print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...") + return partialclass( + MemoryEfficientVideoBlock, + in_channels, + alpha=alpha, + merge_strategy=merge_strategy, + ) + else: + return NotImplementedError() + + +class Conv2DWrapper(torch.nn.Conv2d): + def forward(self, input: torch.Tensor, **kwargs) -> torch.Tensor: + return super().forward(input) + + +class VideoDecoder(Decoder): + available_time_modes = ["all", "conv-only", "attn-only"] + + def __init__( + self, + *args, + video_kernel_size: Union[int, list] = [3,1,1], + alpha: float = 0.0, + merge_strategy: str = "learned", + time_mode: str = "conv-only", + **kwargs, + ): + self.video_kernel_size = video_kernel_size + self.alpha = alpha + self.merge_strategy = merge_strategy + self.time_mode = time_mode + assert ( + self.time_mode in self.available_time_modes + ), f"time_mode parameter has to be in {self.available_time_modes}" + super().__init__(*args, **kwargs) + + def get_last_layer(self, skip_time_mix=False, **kwargs): + if self.time_mode == "attn-only": + raise NotImplementedError("TODO") + else: + return ( + self.conv_out.time_mix_conv.weight + if not skip_time_mix + else self.conv_out.weight + ) + + def _make_attn(self) -> Callable: + if self.time_mode not in ["conv-only", "only-last-conv"]: + return partialclass( + make_time_attn, + alpha=self.alpha, + merge_strategy=self.merge_strategy, + ) + else: + return super()._make_attn() + + def _make_conv(self) -> Callable: + if self.time_mode != "attn-only": + return partialclass(AE3DConv, video_kernel_size=self.video_kernel_size) + else: + return Conv2DWrapper + + def _make_resblock(self) -> Callable: + if self.time_mode not in ["attn-only", "only-last-conv"]: + return partialclass( + VideoResBlock, + video_kernel_size=self.video_kernel_size, + alpha=self.alpha, + merge_strategy=self.merge_strategy, + ) + else: + return super()._make_resblock() \ No newline at end of file diff --git a/lvdm/models/ddpm3d.py b/lvdm/models/ddpm3d.py index c383188..842680c 100644 --- a/lvdm/models/ddpm3d.py +++ b/lvdm/models/ddpm3d.py @@ -28,7 +28,8 @@ from ...lvdm.common import ( exists, default ) - +import math +from ...lvdm.models.autoencoder_dualref import VideoDecoder __conditioning_keys__ = {'concat': 'c_concat', 'crossattn': 'c_crossattn', 'adm': 'y'} @@ -377,6 +378,10 @@ class LatentDiffusion(DDPM): loop_video=False, fps_condition_type='fs', perframe_ae=False, + # added + logdir=None, + rand_cond_frame=False, + en_and_decode_n_samples_a_time=None, *args, **kwargs): self.num_timesteps_cond = default(num_timesteps_cond, 1) self.scale_by_std = scale_by_std @@ -394,6 +399,11 @@ class LatentDiffusion(DDPM): self.loop_video = loop_video self.fps_condition_type = fps_condition_type self.perframe_ae = perframe_ae + + self.logdir = logdir + self.rand_cond_frame = rand_cond_frame + self.en_and_decode_n_samples_a_time = en_and_decode_n_samples_a_time + try: self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 except: @@ -508,16 +518,27 @@ class LatentDiffusion(DDPM): reshape_back = True else: reshape_back = False - - if not self.perframe_ae: - z = 1. / self.scale_factor * z + + z = 1. / self.scale_factor * z + if not self.perframe_ae: results = self.first_stage_model.decode(z, **kwargs) else: + results = [] - for index in range(z.shape[0]): - frame_z = 1. / self.scale_factor * z[index:index+1,:,:,:] - frame_result = self.first_stage_model.decode(frame_z, **kwargs) - results.append(frame_result) + + n_samples = default(self.en_and_decode_n_samples_a_time, self.temporal_length) + n_rounds = math.ceil(z.shape[0] / n_samples) + with torch.autocast("cuda", enabled=True): + for n in range(n_rounds): + if isinstance(self.first_stage_model.decoder, VideoDecoder): + kwargs.update({"timesteps": len(z[n * n_samples : (n + 1) * n_samples])}) + else: + kwargs = {} + + out = self.first_stage_model.decode( + z[n * n_samples : (n + 1) * n_samples], **kwargs + ) + results.append(out) results = torch.cat(results, dim=0) if reshape_back: diff --git a/lvdm/modules/attention_svd.py b/lvdm/modules/attention_svd.py new file mode 100644 index 0000000..6c5852c --- /dev/null +++ b/lvdm/modules/attention_svd.py @@ -0,0 +1,759 @@ +import logging +import math +from inspect import isfunction +from typing import Any, Optional + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from packaging import version +from torch import nn +from torch.utils.checkpoint import checkpoint + +logpy = logging.getLogger(__name__) + +if version.parse(torch.__version__) >= version.parse("2.0.0"): + SDP_IS_AVAILABLE = True + from torch.backends.cuda import SDPBackend, sdp_kernel + + BACKEND_MAP = { + SDPBackend.MATH: { + "enable_math": True, + "enable_flash": False, + "enable_mem_efficient": False, + }, + SDPBackend.FLASH_ATTENTION: { + "enable_math": False, + "enable_flash": True, + "enable_mem_efficient": False, + }, + SDPBackend.EFFICIENT_ATTENTION: { + "enable_math": False, + "enable_flash": False, + "enable_mem_efficient": True, + }, + None: {"enable_math": True, "enable_flash": True, "enable_mem_efficient": True}, + } +else: + from contextlib import nullcontext + + SDP_IS_AVAILABLE = False + sdp_kernel = nullcontext + BACKEND_MAP = {} + logpy.warn( + f"No SDP backend available, likely because you are running in pytorch " + f"versions < 2.0. In fact, you are using PyTorch {torch.__version__}. " + f"You might want to consider upgrading." + ) + +try: + import xformers + import xformers.ops + + XFORMERS_IS_AVAILABLE = True +except: + XFORMERS_IS_AVAILABLE = False + logpy.warn("no module 'xformers'. Processing without...") + +# from .diffusionmodules.util import mixed_checkpoint as checkpoint + + +def exists(val): + return val is not None + + +def uniq(arr): + return {el: True for el in arr}.keys() + + +def default(val, d): + if exists(val): + return val + return d() if isfunction(d) else d + + +def max_neg_value(t): + return -torch.finfo(t.dtype).max + + +def init_(tensor): + dim = tensor.shape[-1] + std = 1 / math.sqrt(dim) + tensor.uniform_(-std, std) + return tensor + + +# feedforward +class GEGLU(nn.Module): + def __init__(self, dim_in, dim_out): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2) + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0): + super().__init__() + inner_dim = int(dim * mult) + dim_out = default(dim_out, dim) + project_in = ( + nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) + if not glu + else GEGLU(dim, inner_dim) + ) + + self.net = nn.Sequential( + project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out) + ) + + def forward(self, x): + return self.net(x) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def Normalize(in_channels): + return torch.nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + + +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange( + qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3 + ) + k = k.softmax(dim=-1) + context = torch.einsum("bhdn,bhen->bhde", k, v) + out = torch.einsum("bhde,bhdn->bhen", context, q) + out = rearrange( + out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w + ) + return self.to_out(out) + + +class SelfAttention(nn.Module): + ATTENTION_MODES = ("xformers", "torch", "math") + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_scale: Optional[float] = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + attn_mode: str = "torch", + ): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + 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) + assert attn_mode in self.ATTENTION_MODES + self.attn_mode = attn_mode + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, L, C = x.shape + + qkv = self.qkv(x) + if self.attn_mode == "torch": + qkv = rearrange( + qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads + ).float() + q, k, v = qkv[0], qkv[1], qkv[2] # B H L D + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "B H L D -> B L (H D)") + elif self.attn_mode == "xformers": + qkv = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads) + q, k, v = qkv[0], qkv[1], qkv[2] # B L H D + x = xformers.ops.memory_efficient_attention(q, k, v) + x = rearrange(x, "B L H D -> B L (H D)", H=self.num_heads) + elif self.attn_mode == "math": + qkv = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k, v = qkv[0], qkv[1], qkv[2] # B H L D + 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, L, C) + else: + raise NotImplemented + + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SpatialSelfAttention(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 = rearrange(q, "b c h w -> b (h w) c") + k = rearrange(k, "b c h w -> b c (h w)") + w_ = torch.einsum("bij,bjk->bik", q, k) + + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = rearrange(v, "b c h w -> b c (h w)") + w_ = rearrange(w_, "b i j -> b j i") + h_ = torch.einsum("bij,bjk->bik", v, w_) + h_ = rearrange(h_, "b c (h w) -> b c h w", h=h) + h_ = self.proj_out(h_) + + return x + h_ + + +class CrossAttention(nn.Module): + def __init__( + self, + query_dim, + context_dim=None, + heads=8, + dim_head=64, + dropout=0.0, + backend=None, + ): + super().__init__() + inner_dim = dim_head * heads + context_dim = default(context_dim, query_dim) + + self.scale = dim_head**-0.5 + self.heads = heads + + self.to_q = nn.Linear(query_dim, inner_dim, bias=False) + self.to_k = nn.Linear(context_dim, inner_dim, bias=False) + self.to_v = nn.Linear(context_dim, inner_dim, bias=False) + + self.to_out = nn.Sequential( + nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) + ) + self.backend = backend + + def forward( + self, + x, + context=None, + mask=None, + additional_tokens=None, + n_times_crossframe_attn_in_self=0, + ): + h = self.heads + + if additional_tokens is not None: + # get the number of masked tokens at the beginning of the output sequence + n_tokens_to_mask = additional_tokens.shape[1] + # add additional token + x = torch.cat([additional_tokens, x], dim=1) + + q = self.to_q(x) + context = default(context, x) + k = self.to_k(context) + v = self.to_v(context) + + if n_times_crossframe_attn_in_self: + # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439 + assert x.shape[0] % n_times_crossframe_attn_in_self == 0 + n_cp = x.shape[0] // n_times_crossframe_attn_in_self + k = repeat( + k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp + ) + v = repeat( + v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp + ) + + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) + + ## old + """ + sim = einsum('b i d, b j d -> b i j', q, k) * self.scale + del q, k + + if exists(mask): + mask = rearrange(mask, 'b ... -> b (...)') + max_neg_value = -torch.finfo(sim.dtype).max + mask = repeat(mask, 'b j -> (b h) () j', h=h) + sim.masked_fill_(~mask, max_neg_value) + + # attention, what we cannot get enough of + sim = sim.softmax(dim=-1) + + out = einsum('b i j, b j d -> b i d', sim, v) + """ + ## new + with sdp_kernel(**BACKEND_MAP[self.backend]): + # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape) + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask + ) # scale is dim_head ** -0.5 per default + + del q, k, v + out = rearrange(out, "b h n d -> b n (h d)", h=h) + + if additional_tokens is not None: + # remove additional token + out = out[:, n_tokens_to_mask:] + return self.to_out(out) + + +class MemoryEfficientCrossAttention(nn.Module): + # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 + def __init__( + self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, **kwargs + ): + super().__init__() + logpy.debug( + f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, " + f"context_dim is {context_dim} and using {heads} heads with a " + f"dimension of {dim_head}." + ) + inner_dim = dim_head * heads + context_dim = default(context_dim, query_dim) + + self.heads = heads + self.dim_head = dim_head + + self.to_q = nn.Linear(query_dim, inner_dim, bias=False) + self.to_k = nn.Linear(context_dim, inner_dim, bias=False) + self.to_v = nn.Linear(context_dim, inner_dim, bias=False) + + self.to_out = nn.Sequential( + nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) + ) + self.attention_op: Optional[Any] = None + + def forward( + self, + x, + context=None, + mask=None, + additional_tokens=None, + n_times_crossframe_attn_in_self=0, + ): + if additional_tokens is not None: + # get the number of masked tokens at the beginning of the output sequence + n_tokens_to_mask = additional_tokens.shape[1] + # add additional token + x = torch.cat([additional_tokens, x], dim=1) + q = self.to_q(x) + context = default(context, x) + k = self.to_k(context) + v = self.to_v(context) + + if n_times_crossframe_attn_in_self: + # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439 + assert x.shape[0] % n_times_crossframe_attn_in_self == 0 + # n_cp = x.shape[0]//n_times_crossframe_attn_in_self + k = repeat( + k[::n_times_crossframe_attn_in_self], + "b ... -> (b n) ...", + n=n_times_crossframe_attn_in_self, + ) + v = repeat( + v[::n_times_crossframe_attn_in_self], + "b ... -> (b n) ...", + n=n_times_crossframe_attn_in_self, + ) + + b, _, _ = q.shape + q, k, v = map( + lambda t: t.unsqueeze(3) + .reshape(b, t.shape[1], self.heads, self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b * self.heads, t.shape[1], self.dim_head) + .contiguous(), + (q, k, v), + ) + + # actually compute the attention, what we cannot get enough of + if version.parse(xformers.__version__) >= version.parse("0.0.21"): + # NOTE: workaround for + # https://github.com/facebookresearch/xformers/issues/845 + max_bs = 32768 + N = q.shape[0] + n_batches = math.ceil(N / max_bs) + out = list() + for i_batch in range(n_batches): + batch = slice(i_batch * max_bs, (i_batch + 1) * max_bs) + out.append( + xformers.ops.memory_efficient_attention( + q[batch], + k[batch], + v[batch], + attn_bias=None, + op=self.attention_op, + ) + ) + out = torch.cat(out, 0) + else: + out = xformers.ops.memory_efficient_attention( + q, k, v, attn_bias=None, op=self.attention_op + ) + + # TODO: Use this directly in the attention operation, as a bias + if exists(mask): + raise NotImplementedError + out = ( + out.unsqueeze(0) + .reshape(b, self.heads, out.shape[1], self.dim_head) + .permute(0, 2, 1, 3) + .reshape(b, out.shape[1], self.heads * self.dim_head) + ) + if additional_tokens is not None: + # remove additional token + out = out[:, n_tokens_to_mask:] + return self.to_out(out) + + +class BasicTransformerBlock(nn.Module): + ATTENTION_MODES = { + "softmax": CrossAttention, # vanilla attention + "softmax-xformers": MemoryEfficientCrossAttention, # ampere + } + + def __init__( + self, + dim, + n_heads, + d_head, + dropout=0.0, + context_dim=None, + gated_ff=True, + checkpoint=True, + disable_self_attn=False, + attn_mode="softmax", + sdp_backend=None, + ): + super().__init__() + assert attn_mode in self.ATTENTION_MODES + if attn_mode != "softmax" and not XFORMERS_IS_AVAILABLE: + logpy.warn( + f"Attention mode '{attn_mode}' is not available. Falling " + f"back to native attention. This is not a problem in " + f"Pytorch >= 2.0. FYI, you are running with PyTorch " + f"version {torch.__version__}." + ) + attn_mode = "softmax" + elif attn_mode == "softmax" and not SDP_IS_AVAILABLE: + logpy.warn( + "We do not support vanilla attention anymore, as it is too " + "expensive. Sorry." + ) + if not XFORMERS_IS_AVAILABLE: + assert ( + False + ), "Please install xformers via e.g. 'pip install xformers==0.0.16'" + else: + logpy.info("Falling back to xformers efficient attention.") + attn_mode = "softmax-xformers" + attn_cls = self.ATTENTION_MODES[attn_mode] + if version.parse(torch.__version__) >= version.parse("2.0.0"): + assert sdp_backend is None or isinstance(sdp_backend, SDPBackend) + else: + assert sdp_backend is None + self.disable_self_attn = disable_self_attn + self.attn1 = attn_cls( + query_dim=dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + context_dim=context_dim if self.disable_self_attn else None, + backend=sdp_backend, + ) # is a self-attention if not self.disable_self_attn + self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) + self.attn2 = attn_cls( + query_dim=dim, + context_dim=context_dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + backend=sdp_backend, + ) # is self-attn if context is none + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + self.norm3 = nn.LayerNorm(dim) + self.checkpoint = checkpoint + if self.checkpoint: + logpy.debug(f"{self.__class__.__name__} is using checkpointing") + + def forward( + self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0 + ): + kwargs = {"x": x} + + if context is not None: + kwargs.update({"context": context}) + + if additional_tokens is not None: + kwargs.update({"additional_tokens": additional_tokens}) + + if n_times_crossframe_attn_in_self: + kwargs.update( + {"n_times_crossframe_attn_in_self": n_times_crossframe_attn_in_self} + ) + + # return mixed_checkpoint(self._forward, kwargs, self.parameters(), self.checkpoint) + if self.checkpoint: + # inputs = {"x": x, "context": context} + return checkpoint(self._forward, x, context) + # return checkpoint(self._forward, inputs, self.parameters(), self.checkpoint) + else: + return self._forward(**kwargs) + + def _forward( + self, x, context=None, additional_tokens=None, n_times_crossframe_attn_in_self=0 + ): + x = ( + self.attn1( + self.norm1(x), + context=context if self.disable_self_attn else None, + additional_tokens=additional_tokens, + n_times_crossframe_attn_in_self=n_times_crossframe_attn_in_self + if not self.disable_self_attn + else 0, + ) + + x + ) + x = ( + self.attn2( + self.norm2(x), context=context, additional_tokens=additional_tokens + ) + + x + ) + x = self.ff(self.norm3(x)) + x + return x + + +class BasicTransformerSingleLayerBlock(nn.Module): + ATTENTION_MODES = { + "softmax": CrossAttention, # vanilla attention + "softmax-xformers": MemoryEfficientCrossAttention # on the A100s not quite as fast as the above version + # (todo might depend on head_dim, check, falls back to semi-optimized kernels for dim!=[16,32,64,128]) + } + + def __init__( + self, + dim, + n_heads, + d_head, + dropout=0.0, + context_dim=None, + gated_ff=True, + checkpoint=True, + attn_mode="softmax", + ): + super().__init__() + assert attn_mode in self.ATTENTION_MODES + attn_cls = self.ATTENTION_MODES[attn_mode] + self.attn1 = attn_cls( + query_dim=dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + context_dim=context_dim, + ) + self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + self.checkpoint = checkpoint + + def forward(self, x, context=None): + # inputs = {"x": x, "context": context} + # return checkpoint(self._forward, inputs, self.parameters(), self.checkpoint) + return checkpoint(self._forward, x, context) + + def _forward(self, x, context=None): + x = self.attn1(self.norm1(x), context=context) + x + x = self.ff(self.norm2(x)) + x + return x + + +class SpatialTransformer(nn.Module): + """ + Transformer block for image-like data. + First, project the input (aka embedding) + and reshape to b, t, d. + Then apply standard transformer action. + Finally, reshape to image + NEW: use_linear for more efficiency instead of the 1x1 convs + """ + + def __init__( + self, + in_channels, + n_heads, + d_head, + depth=1, + dropout=0.0, + context_dim=None, + disable_self_attn=False, + use_linear=False, + attn_type="softmax", + use_checkpoint=True, + # sdp_backend=SDPBackend.FLASH_ATTENTION + sdp_backend=None, + ): + super().__init__() + logpy.debug( + f"constructing {self.__class__.__name__} of depth {depth} w/ " + f"{in_channels} channels and {n_heads} heads." + ) + + if exists(context_dim) and not isinstance(context_dim, list): + context_dim = [context_dim] + if exists(context_dim) and isinstance(context_dim, list): + if depth != len(context_dim): + logpy.warn( + f"{self.__class__.__name__}: Found context dims " + f"{context_dim} of depth {len(context_dim)}, which does not " + f"match the specified 'depth' of {depth}. Setting context_dim " + f"to {depth * [context_dim[0]]} now." + ) + # depth does not match context dims. + assert all( + map(lambda x: x == context_dim[0], context_dim) + ), "need homogenous context_dim to match depth automatically" + context_dim = depth * [context_dim[0]] + elif context_dim is None: + context_dim = [None] * depth + self.in_channels = in_channels + inner_dim = n_heads * d_head + self.norm = Normalize(in_channels) + if not use_linear: + self.proj_in = nn.Conv2d( + in_channels, inner_dim, kernel_size=1, stride=1, padding=0 + ) + else: + self.proj_in = nn.Linear(in_channels, inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + n_heads, + d_head, + dropout=dropout, + context_dim=context_dim[d], + disable_self_attn=disable_self_attn, + attn_mode=attn_type, + checkpoint=use_checkpoint, + sdp_backend=sdp_backend, + ) + for d in range(depth) + ] + ) + if not use_linear: + self.proj_out = zero_module( + nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0) + ) + else: + # self.proj_out = zero_module(nn.Linear(in_channels, inner_dim)) + self.proj_out = zero_module(nn.Linear(inner_dim, in_channels)) + self.use_linear = use_linear + + def forward(self, x, context=None): + # note: if no context is given, cross-attention defaults to self-attention + if not isinstance(context, list): + context = [context] + b, c, h, w = x.shape + x_in = x + x = self.norm(x) + if not self.use_linear: + x = self.proj_in(x) + x = rearrange(x, "b c h w -> b (h w) c").contiguous() + if self.use_linear: + x = self.proj_in(x) + for i, block in enumerate(self.transformer_blocks): + if i > 0 and len(context) == 1: + i = 0 # use same context for each block + x = block(x, context=context[i]) + if self.use_linear: + x = self.proj_out(x) + x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous() + if not self.use_linear: + x = self.proj_out(x) + return x + x_in + + +class SimpleTransformer(nn.Module): + def __init__( + self, + dim: int, + depth: int, + heads: int, + dim_head: int, + context_dim: Optional[int] = None, + dropout: float = 0.0, + checkpoint: bool = True, + ): + super().__init__() + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + BasicTransformerBlock( + dim, + heads, + dim_head, + dropout=dropout, + context_dim=context_dim, + attn_mode="softmax", + checkpoint=checkpoint, + ) + ) + + def forward( + self, + x: torch.Tensor, + context: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + for layer in self.layers: + x = layer(x, context) + return x \ No newline at end of file diff --git a/lvdm/modules/networks/ae_modules.py b/lvdm/modules/networks/ae_modules.py index 64621fe..55f0e4f 100644 --- a/lvdm/modules/networks/ae_modules.py +++ b/lvdm/modules/networks/ae_modules.py @@ -427,13 +427,17 @@ class Encoder(nn.Module): stride=1, padding=1) - def forward(self, x): + def forward(self, x, return_hidden_states=False): # timestep embedding temb = None # print(f'encoder-input={x.shape}') # downsampling hs = [self.conv_in(x)] + + ## if we return hidden states for decoder usage, we will store them in a list + if return_hidden_states: + hidden_states = [] # print(f'encoder-conv in feat={hs[0].shape}') for i_level in range(self.num_resolutions): for i_block in range(self.num_res_blocks): @@ -442,11 +446,14 @@ class Encoder(nn.Module): if len(self.down[i_level].attn) > 0: h = self.down[i_level].attn[i_block](h) hs.append(h) + if return_hidden_states: + hidden_states.append(h) if i_level != self.num_resolutions-1: # print(f'encoder-downsample (input)={hs[-1].shape}') hs.append(self.down[i_level].downsample(hs[-1])) # print(f'encoder-downsample (output)={hs[-1].shape}') - + if return_hidden_states: + hidden_states.append(hs[0]) # middle h = hs[-1] h = self.mid.block_1(h, temb) @@ -460,7 +467,10 @@ class Encoder(nn.Module): h = nonlinearity(h) h = self.conv_out(h) # print(f'end feat={h.shape}') - return h + if return_hidden_states: + return h, hidden_states + else: + return h class Decoder(nn.Module): diff --git a/nodes.py b/nodes.py index 4dc5e6e..e917011 100644 --- a/nodes.py +++ b/nodes.py @@ -2,7 +2,7 @@ import os from omegaconf import OmegaConf import torch import torch.nn.functional as F -from .scripts.evaluation.funcs import load_model_checkpoint, get_latent_z +from .scripts.evaluation.funcs import load_model_checkpoint, get_latent_z, get_latent_z_with_hidden_states from .utils.utils import instantiate_from_config from einops import repeat import folder_paths @@ -66,7 +66,9 @@ class DynamiCrafterModelLoader: model_path = folder_paths.get_full_path("checkpoints", ckpt_name) ckpt_base_name = os.path.basename(ckpt_name) base_name, _ = os.path.splitext(ckpt_base_name) - if 'interp' in base_name and '512' in base_name: + if 'toon' in base_name and '512' in base_name: + config_file=os.path.join(script_directory, "configs", "tooncrafter_512_interp.yaml") + elif 'interp' in base_name and '512' in base_name: config_file=os.path.join(script_directory, "configs", "dynamicrafter_512_interp_v1.yaml") elif '1024' in base_name: config_file=os.path.join(script_directory, "configs", "dynamicrafter_1024_v1.yaml") @@ -97,7 +99,6 @@ class DynamiCrafterModelLoader: if fp8_unet: self.model.model.diffusion_model = self.model.model.diffusion_model.to(torch.float8_e4m3fn) print(f"Model using dtype: {self.model.dtype}") - print(self.model) return (self.model,) class DynamiCrafterI2V: @@ -298,6 +299,212 @@ class DynamiCrafterI2V: video = F.interpolate(video.permute(0, 3, 1, 2), size=(final_H, final_W), mode="bicubic").permute(0, 2, 3, 1) last_image = video[-1].unsqueeze(0) return (video, last_image) + +class ToonCrafterI2V: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("DCMODEL",), + "image": ("IMAGE",), + "steps": ("INT", {"default": 50, "min": 1, "max": 200, "step": 1}), + "cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 20.0, "step": 0.01}), + "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "frames": ("INT", {"default": 16, "min": 1, "max": 100, "step": 1}), + "prompt": ("STRING", {"multiline": True, "default": "",}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "fs": ("INT", {"default": 10, "min": 2, "max": 100, "step": 1}), + "keep_model_loaded": ("BOOLEAN", {"default": True}), + "vae_dtype": ( + [ + 'fp32', + 'fp16', + 'bf16', + 'auto' + ], { + "default": 'auto' + }), + + }, + "optional": { + "image2": ("IMAGE",), + "mask": ("MASK",), + "frame_window_size": ("INT", {"default": 16, "min": 1, "max": 200, "step": 1}), + "frame_window_stride": ("INT", {"default": 4, "min": 1, "max": 200, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE",) + RETURN_NAMES = ("images", "middle_frames",) + FUNCTION = "process" + CATEGORY = "DynamiCrafterWrapper" + + def process(self, model, image, prompt, cfg, steps, eta, seed, fs, keep_model_loaded, frames, vae_dtype, frame_window_size=16, frame_window_stride=4, mask=None, image2=None): + device = mm.get_torch_device() + mm.unload_all_models() + mm.soft_empty_cache() + + torch.manual_seed(seed) + dtype = model.dtype + if vae_dtype == "auto": + try: + if mm.should_use_bf16(): + model.first_stage_model.to(convert_dtype('bf16')) + else: + model.first_stage_model.to(convert_dtype('fp32')) + except: + raise AttributeError("ComfyUI version too old, can't autodetect properly. Set your dtype manually.") + else: + model.first_stage_model.to(convert_dtype(vae_dtype)) + print(f"VAE using dtype: {model.first_stage_model.dtype}") + + self.model = model + self.model.to(device) + autocast_condition = (dtype != torch.float32) and not comfy.model_management.is_device_mps(device) + with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): + image = image * 2 - 1 + image = image.permute(0, 3, 1, 2).to(dtype).to(device) + + B, C, H, W = image.shape + orig_H, orig_W = H, W + if W % 64 != 0: + W = W - (W % 64) + if H % 64 != 0: + H = H - (H % 64) + if orig_H % 64 != 0 or orig_W % 64 != 0: + image = F.interpolate(image, size=(H, W), mode="bicubic") + + B, C, H, W = image.shape + noise_shape = [B, self.model.model.diffusion_model.out_channels, frames, H // 8, W // 8] + + self.model.first_stage_model.to(device) + + + image2 = image2 * 2 - 1 + image2 = image2.permute(0, 3, 1, 2).to(dtype).to(device) + if image2.shape != image.shape: + image2 = F.interpolate(image, size=(H, W), mode="bicubic") + + videos = image.unsqueeze(2) # bc1hw + videos = repeat(videos, 'b c t h w -> b c (repeat t) h w', repeat=frames//2) + + videos2 = image2.unsqueeze(2) # bc1hw + videos2 = repeat(videos2, 'b c t h w -> b c (repeat t) h w', repeat=frames//2) + + videos = torch.cat([videos, videos2], dim=2) + + z, hs = get_latent_z_with_hidden_states(self.model, videos) + + img_tensor_repeat = torch.zeros_like(z) + img_tensor_repeat[:,:,:1,:,:] = z[:,:,:1,:,:] + img_tensor_repeat[:,:,-1:,:,:] = z[:,:,-1:,:,:] + + + self.model.first_stage_model.to('cpu') + + self.model.cond_stage_model.to(device) + self.model.embedder.to(device) + self.model.image_proj_model.to(device) + + text_emb = self.model.get_learned_conditioning([prompt]) + cond_images = self.model.embedder(image) + img_emb = self.model.image_proj_model(cond_images) + imtext_cond = torch.cat([text_emb, img_emb], dim=1) + del cond_images, img_emb, text_emb + + fs = torch.tensor([fs], dtype=torch.long, device=self.model.device) + cond = {"c_crossattn": [imtext_cond], "c_concat": [img_tensor_repeat]} + + if noise_shape[-1] == 32: + timestep_spacing = "uniform" + guidance_rescale = 0.0 + else: + timestep_spacing = "uniform_trailing" + guidance_rescale = 0.7 + + ## construct unconditional guidance + if cfg != 1.0: + uc_emb = self.model.get_learned_conditioning([""]) + ## process image embedding token + if hasattr(self.model, 'embedder'): + uc_img = torch.zeros(noise_shape[0],3,224,224).to(self.model.device) + ## img: b c h w >> b l c + uc_img = self.model.embedder(uc_img) + uc_img = self.model.image_proj_model(uc_img) + uc_emb = torch.cat([uc_emb, uc_img], dim=1) + if isinstance(cond, dict): + uc = {key:cond[key] for key in cond.keys()} + uc.update({'c_crossattn': [uc_emb]}) + else: + uc = uc_emb + else: + uc = None + + self.model.cond_stage_model.to('cpu') + self.model.embedder.to('cpu') + self.model.image_proj_model.to('cpu') + + if mask is not None: + mask = mask.to(dtype).to(device) + mask = F.interpolate(mask.unsqueeze(0), size=(H // 8, W // 8), mode="nearest").squeeze(0) + mask = (1 - mask) + mask = mask.unsqueeze(1) + B, C, H, W = mask.shape + if B < frames: + mask = mask.unsqueeze(2) + mask = mask.expand(-1, -1, frames, -1, -1) + else: + mask = mask.unsqueeze(0) + mask = mask.permute(0, 2, 1, 3, 4) + mask = torch.where(mask < 1.0, torch.tensor(0.0, device=device, dtype=dtype), torch.tensor(1.0, device=device, dtype=dtype)) + + #inference + ddim_sampler = DDIMSampler(self.model) + samples, _ = ddim_sampler.sample(S=steps, + conditioning=cond, + batch_size=noise_shape[0], + shape=noise_shape[1:], + verbose=True, + unconditional_guidance_scale=cfg, + unconditional_conditioning=uc, + eta=eta, + temporal_length=noise_shape[2], + conditional_guidance_scale_temporal=None, + x_T=None, + fs=fs, + timestep_spacing=timestep_spacing, + guidance_rescale=guidance_rescale, + clean_cond=True, + mask=mask, + x0=img_tensor_repeat.clone() if mask is not None else None, + frame_window_size = frame_window_size, + frame_window_stride = frame_window_stride, + ) + + assert not torch.isnan(samples).any().item(), "Resulting tensor containts NaNs. I'm unsure why this happens, changing step count and/or image dimensions might help." + + ## reconstruct from latent to pixel space + self.model.first_stage_model.to(device) + additional_decode_kwargs = {'ref_context': hs} + decoded_images = self.model.decode_first_stage(samples) #b c t h w + self.model.first_stage_model.to('cpu') + + video = decoded_images.detach().cpu() + video = torch.clamp(video.float(), -1., 1.) + video = (video + 1.0) / 2.0 + video = video.squeeze(0).permute(1, 2, 3, 0) + del decoded_images, samples + + if not keep_model_loaded: + self.model.to('cpu') + mm.soft_empty_cache() + # Ensure the final dimensions are divisible by 2 + final_H = (orig_H // 2) * 2 + final_W = (orig_W // 2) * 2 + + if video.shape[1] != final_H or video.shape[2] != final_W: + video = F.interpolate(video.permute(0, 3, 1, 2), size=(final_H, final_W), mode="bicubic").permute(0, 2, 3, 1) + middle_frames = video[1:-1] + return (video, middle_frames) class DynamiCrafterBatchInterpolation: @classmethod @@ -506,11 +713,13 @@ class DynamiCrafterBatchInterpolation: NODE_CLASS_MAPPINGS = { "DynamiCrafterI2V": DynamiCrafterI2V, "DynamiCrafterModelLoader": DynamiCrafterModelLoader, - "DynamiCrafterBatchInterpolation": DynamiCrafterBatchInterpolation + "DynamiCrafterBatchInterpolation": DynamiCrafterBatchInterpolation, + "ToonCrafterI2V": ToonCrafterI2V } NODE_DISPLAY_NAME_MAPPINGS = { "DynamiCrafterI2V": "DynamiCrafterI2V", "DynamiCrafterModelLoader": "DynamiCrafterModelLoader", - "DynamiCrafterBatchInterpolation": "DynamiCrafterBatchInterpolation" + "DynamiCrafterBatchInterpolation": "DynamiCrafterBatchInterpolation", + "ToonCrafterI2V": "ToonCrafterI2V" } diff --git a/scripts/evaluation/funcs.py b/scripts/evaluation/funcs.py index f1c330d..73f5654 100644 --- a/scripts/evaluation/funcs.py +++ b/scripts/evaluation/funcs.py @@ -52,4 +52,20 @@ def get_latent_z(model, videos): x = rearrange(videos, 'b c t h w -> (b t) c h w') z = model.encode_first_stage(x) z = rearrange(z, '(b t) c h w -> b c t h w', b=b, t=t) - return z \ No newline at end of file + return z + +def get_latent_z_with_hidden_states(model, videos): + b, c, t, h, w = videos.shape + x = rearrange(videos, 'b c t h w -> (b t) c h w') + encoder_posterior, hidden_states = model.first_stage_model.encode(x, return_hidden_states=True) + + hidden_states_first_last = [] + ### use only the first and last hidden states + for hid in hidden_states: + hid = rearrange(hid, '(b t) c h w -> b c t h w', t=t) + hid_new = torch.cat([hid[:, :, 0:1], hid[:, :, -1:]], dim=2) + hidden_states_first_last.append(hid_new) + + z = model.get_first_stage_encoding(encoder_posterior).detach() + z = rearrange(z, '(b t) c h w -> b c t h w', b=b, t=t) + return z, hidden_states_first_last \ No newline at end of file