| import functools
|
| from typing import Iterable, Union
|
|
|
| import torch
|
| from einops import rearrange, repeat
|
|
|
| import comfy.ops
|
| ops = comfy.ops.disable_weight_init
|
|
|
| from .diffusionmodules.model import (
|
| AttnBlock,
|
| Decoder,
|
| ResnetBlock,
|
| )
|
| from .diffusionmodules.openaimodel import ResBlock, timestep_embedding
|
| from .attention import BasicTransformerBlock
|
|
|
| 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=False,
|
| 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):
|
| b, c, h, w = x.shape
|
| if timesteps is None:
|
| timesteps = b
|
|
|
| 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).to(x.device)
|
| 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(ops.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 = ops.Conv3d(
|
| in_channels=out_channels,
|
| out_channels=out_channels,
|
| kernel_size=video_kernel_size,
|
| padding=padding,
|
| )
|
|
|
| def forward(self, input, timesteps=None, skip_video=False):
|
| if timesteps is None:
|
| timesteps = input.shape[0]
|
| 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 AttnVideoBlock(AttnBlock):
|
| def __init__(
|
| self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"
|
| ):
|
| super().__init__(in_channels)
|
|
|
| self.time_mix_block = BasicTransformerBlock(
|
| dim=in_channels,
|
| n_heads=1,
|
| d_head=in_channels,
|
| checkpoint=False,
|
| ff_in=True,
|
| )
|
|
|
| time_embed_dim = self.in_channels * 4
|
| self.video_time_embed = torch.nn.Sequential(
|
| ops.Linear(self.in_channels, time_embed_dim),
|
| torch.nn.SiLU(),
|
| ops.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=None, skip_time_block=False):
|
| if skip_time_block:
|
| return super().forward(x)
|
|
|
| if timesteps is None:
|
| timesteps = x.shape[0]
|
|
|
| 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)
|
| emb = emb[:, None, :]
|
| x_mix = x_mix + emb
|
|
|
| alpha = self.get_alpha().to(x.device)
|
| x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
|
| x = alpha * x + (1.0 - alpha) * x_mix
|
|
|
| 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",
|
| conv_op=ops.Conv2d,
|
| ):
|
| return partialclass(
|
| AttnVideoBlock, in_channels, alpha=alpha, merge_strategy=merge_strategy
|
| )
|
|
|
|
|
| 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,
|
| 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}"
|
|
|
| if self.time_mode != "attn-only":
|
| kwargs["conv_out_op"] = partialclass(AE3DConv, video_kernel_size=self.video_kernel_size)
|
| if self.time_mode not in ["conv-only", "only-last-conv"]:
|
| kwargs["attn_op"] = partialclass(make_time_attn, alpha=self.alpha, merge_strategy=self.merge_strategy)
|
| if self.time_mode not in ["attn-only", "only-last-conv"]:
|
| kwargs["resnet_op"] = partialclass(VideoResBlock, video_kernel_size=self.video_kernel_size, alpha=self.alpha, merge_strategy=self.merge_strategy)
|
|
|
| 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
|
| )
|
|
|
