zparadox/stable-video-diffusion

1from typing import Callable, Iterable, Union2 3import torch4from einops import rearrange, repeat5 6from sgm.modules.diffusionmodules.model import (7    XFORMERS_IS_AVAILABLE,8    AttnBlock,9    Decoder,10    MemoryEfficientAttnBlock,11    ResnetBlock,12)13from sgm.modules.diffusionmodules.openaimodel import ResBlock, timestep_embedding14from sgm.modules.video_attention import VideoTransformerBlock15from sgm.util import partialclass16 17 18class VideoResBlock(ResnetBlock):19    def __init__(20        self,21        out_channels,22        *args,23        dropout=0.0,24        video_kernel_size=3,25        alpha=0.0,26        merge_strategy="learned",27        **kwargs,28    ):29        super().__init__(out_channels=out_channels, dropout=dropout, *args, **kwargs)30        if video_kernel_size is None:31            video_kernel_size = [3, 1, 1]32        self.time_stack = ResBlock(33            channels=out_channels,34            emb_channels=0,35            dropout=dropout,36            dims=3,37            use_scale_shift_norm=False,38            use_conv=False,39            up=False,40            down=False,41            kernel_size=video_kernel_size,42            use_checkpoint=False,43            skip_t_emb=True,44        )45 46        self.merge_strategy = merge_strategy47        if self.merge_strategy == "fixed":48            self.register_buffer("mix_factor", torch.Tensor([alpha]))49        elif self.merge_strategy == "learned":50            self.register_parameter(51                "mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))52            )53        else:54            raise ValueError(f"unknown merge strategy {self.merge_strategy}")55 56    def get_alpha(self, bs):57        if self.merge_strategy == "fixed":58            return self.mix_factor59        elif self.merge_strategy == "learned":60            return torch.sigmoid(self.mix_factor)61        else:62            raise NotImplementedError()63 64    def forward(self, x, temb, skip_video=False, timesteps=None):65        if timesteps is None:66            timesteps = self.timesteps67 68        b, c, h, w = x.shape69 70        x = super().forward(x, temb)71 72        if not skip_video:73            x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)74 75            x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)76 77            x = self.time_stack(x, temb)78 79            alpha = self.get_alpha(bs=b // timesteps)80            x = alpha * x + (1.0 - alpha) * x_mix81 82            x = rearrange(x, "b c t h w -> (b t) c h w")83        return x84 85 86class AE3DConv(torch.nn.Conv2d):87    def __init__(self, in_channels, out_channels, video_kernel_size=3, *args, **kwargs):88        super().__init__(in_channels, out_channels, *args, **kwargs)89        if isinstance(video_kernel_size, Iterable):90            padding = [int(k // 2) for k in video_kernel_size]91        else:92            padding = int(video_kernel_size // 2)93 94        self.time_mix_conv = torch.nn.Conv3d(95            in_channels=out_channels,96            out_channels=out_channels,97            kernel_size=video_kernel_size,98            padding=padding,99        )100 101    def forward(self, input, timesteps, skip_video=False):102        x = super().forward(input)103        if skip_video:104            return x105        x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)106        x = self.time_mix_conv(x)107        return rearrange(x, "b c t h w -> (b t) c h w")108 109 110class VideoBlock(AttnBlock):111    def __init__(112        self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"113    ):114        super().__init__(in_channels)115        # no context, single headed, as in base class116        self.time_mix_block = VideoTransformerBlock(117            dim=in_channels,118            n_heads=1,119            d_head=in_channels,120            checkpoint=False,121            ff_in=True,122            attn_mode="softmax",123        )124 125        time_embed_dim = self.in_channels * 4126        self.video_time_embed = torch.nn.Sequential(127            torch.nn.Linear(self.in_channels, time_embed_dim),128            torch.nn.SiLU(),129            torch.nn.Linear(time_embed_dim, self.in_channels),130        )131 132        self.merge_strategy = merge_strategy133        if self.merge_strategy == "fixed":134            self.register_buffer("mix_factor", torch.Tensor([alpha]))135        elif self.merge_strategy == "learned":136            self.register_parameter(137                "mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))138            )139        else:140            raise ValueError(f"unknown merge strategy {self.merge_strategy}")141 142    def forward(self, x, timesteps, skip_video=False):143        if skip_video:144            return super().forward(x)145 146        x_in = x147        x = self.attention(x)148        h, w = x.shape[2:]149        x = rearrange(x, "b c h w -> b (h w) c")150 151        x_mix = x152        num_frames = torch.arange(timesteps, device=x.device)153        num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)154        num_frames = rearrange(num_frames, "b t -> (b t)")155        t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False)156        emb = self.video_time_embed(t_emb)  # b, n_channels157        emb = emb[:, None, :]158        x_mix = x_mix + emb159 160        alpha = self.get_alpha()161        x_mix = self.time_mix_block(x_mix, timesteps=timesteps)162        x = alpha * x + (1.0 - alpha) * x_mix  # alpha merge163 164        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)165        x = self.proj_out(x)166 167        return x_in + x168 169    def get_alpha(170        self,171    ):172        if self.merge_strategy == "fixed":173            return self.mix_factor174        elif self.merge_strategy == "learned":175            return torch.sigmoid(self.mix_factor)176        else:177            raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}")178 179 180class MemoryEfficientVideoBlock(MemoryEfficientAttnBlock):181    def __init__(182        self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"183    ):184        super().__init__(in_channels)185        # no context, single headed, as in base class186        self.time_mix_block = VideoTransformerBlock(187            dim=in_channels,188            n_heads=1,189            d_head=in_channels,190            checkpoint=False,191            ff_in=True,192            attn_mode="softmax-xformers",193        )194 195        time_embed_dim = self.in_channels * 4196        self.video_time_embed = torch.nn.Sequential(197            torch.nn.Linear(self.in_channels, time_embed_dim),198            torch.nn.SiLU(),199            torch.nn.Linear(time_embed_dim, self.in_channels),200        )201 202        self.merge_strategy = merge_strategy203        if self.merge_strategy == "fixed":204            self.register_buffer("mix_factor", torch.Tensor([alpha]))205        elif self.merge_strategy == "learned":206            self.register_parameter(207                "mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))208            )209        else:210            raise ValueError(f"unknown merge strategy {self.merge_strategy}")211 212    def forward(self, x, timesteps, skip_time_block=False):213        if skip_time_block:214            return super().forward(x)215 216        x_in = x217        x = self.attention(x)218        h, w = x.shape[2:]219        x = rearrange(x, "b c h w -> b (h w) c")220 221        x_mix = x222        num_frames = torch.arange(timesteps, device=x.device)223        num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)224        num_frames = rearrange(num_frames, "b t -> (b t)")225        t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False)226        emb = self.video_time_embed(t_emb)  # b, n_channels227        emb = emb[:, None, :]228        x_mix = x_mix + emb229 230        alpha = self.get_alpha()231        x_mix = self.time_mix_block(x_mix, timesteps=timesteps)232        x = alpha * x + (1.0 - alpha) * x_mix  # alpha merge233 234        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)235        x = self.proj_out(x)236 237        return x_in + x238 239    def get_alpha(240        self,241    ):242        if self.merge_strategy == "fixed":243            return self.mix_factor244        elif self.merge_strategy == "learned":245            return torch.sigmoid(self.mix_factor)246        else:247            raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}")248 249 250def make_time_attn(251    in_channels,252    attn_type="vanilla",253    attn_kwargs=None,254    alpha: float = 0,255    merge_strategy: str = "learned",256):257    assert attn_type in [258        "vanilla",259        "vanilla-xformers",260    ], f"attn_type {attn_type} not supported for spatio-temporal attention"261    print(262        f"making spatial and temporal attention of type '{attn_type}' with {in_channels} in_channels"263    )264    if not XFORMERS_IS_AVAILABLE and attn_type == "vanilla-xformers":265        print(266            f"Attention mode '{attn_type}' is not available. Falling back to vanilla attention. "267            f"This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}"268        )269        attn_type = "vanilla"270 271    if attn_type == "vanilla":272        assert attn_kwargs is None273        return partialclass(274            VideoBlock, in_channels, alpha=alpha, merge_strategy=merge_strategy275        )276    elif attn_type == "vanilla-xformers":277        print(f"building MemoryEfficientAttnBlock with {in_channels} in_channels...")278        return partialclass(279            MemoryEfficientVideoBlock,280            in_channels,281            alpha=alpha,282            merge_strategy=merge_strategy,283        )284    else:285        return NotImplementedError()286 287 288class Conv2DWrapper(torch.nn.Conv2d):289    def forward(self, input: torch.Tensor, **kwargs) -> torch.Tensor:290        return super().forward(input)291 292 293class VideoDecoder(Decoder):294    available_time_modes = ["all", "conv-only", "attn-only"]295 296    def __init__(297        self,298        *args,299        video_kernel_size: Union[int, list] = 3,300        alpha: float = 0.0,301        merge_strategy: str = "learned",302        time_mode: str = "conv-only",303        **kwargs,304    ):305        self.video_kernel_size = video_kernel_size306        self.alpha = alpha307        self.merge_strategy = merge_strategy308        self.time_mode = time_mode309        assert (310            self.time_mode in self.available_time_modes311        ), f"time_mode parameter has to be in {self.available_time_modes}"312        super().__init__(*args, **kwargs)313 314    def get_last_layer(self, skip_time_mix=False, **kwargs):315        if self.time_mode == "attn-only":316            raise NotImplementedError("TODO")317        else:318            return (319                self.conv_out.time_mix_conv.weight320                if not skip_time_mix321                else self.conv_out.weight322            )323 324    def _make_attn(self) -> Callable:325        if self.time_mode not in ["conv-only", "only-last-conv"]:326            return partialclass(327                make_time_attn,328                alpha=self.alpha,329                merge_strategy=self.merge_strategy,330            )331        else:332            return super()._make_attn()333 334    def _make_conv(self) -> Callable:335        if self.time_mode != "attn-only":336            return partialclass(AE3DConv, video_kernel_size=self.video_kernel_size)337        else:338            return Conv2DWrapper339 340    def _make_resblock(self) -> Callable:341        if self.time_mode not in ["attn-only", "only-last-conv"]:342            return partialclass(343                VideoResBlock,344                video_kernel_size=self.video_kernel_size,345                alpha=self.alpha,346                merge_strategy=self.merge_strategy,347            )348        else:349            return super()._make_resblock()350