# Copyright (c) MONAI Consortium # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # ========================================================================= # Adapted from https://github.com/huggingface/diffusers # which has the following license: # https://github.com/huggingface/diffusers/blob/main/LICENSE # # Copyright 2022 UC Berkeley Team and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ========================================================================= from __future__ import annotations from collections.abc import Sequence import torch from torch import nn from monai.networks.blocks import Convolution, SpatialAttentionBlock from monai.networks.blocks.spade_norm import SPADE from monai.networks.nets.diffusion_model_unet import ( DiffusionUnetDownsample, DiffusionUNetResnetBlock, SpatialTransformer, WrappedUpsample, get_down_block, get_mid_block, get_timestep_embedding, zero_module, ) from monai.utils import ensure_tuple_rep __all__ = ["SPADEDiffusionModelUNet"] class SPADEDiffResBlock(nn.Module): """ Residual block with timestep conditioning and SPADE norm. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. temb_channels: number of timestep embedding channels. label_nc: number of semantic channels for SPADE normalisation. out_channels: number of output channels. up: if True, performs upsampling. down: if True, performs downsampling. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. spade_intermediate_channels: number of intermediate channels for SPADE block layer """ def __init__( self, spatial_dims: int, in_channels: int, temb_channels: int, label_nc: int, out_channels: int | None = None, up: bool = False, down: bool = False, norm_num_groups: int = 32, norm_eps: float = 1e-6, spade_intermediate_channels: int = 128, ) -> None: super().__init__() self.spatial_dims = spatial_dims self.channels = in_channels self.emb_channels = temb_channels self.out_channels = out_channels or in_channels self.up = up self.down = down self.norm1 = SPADE( label_nc=label_nc, norm_nc=in_channels, norm="GROUP", norm_params={"num_groups": norm_num_groups, "eps": norm_eps, "affine": True}, hidden_channels=spade_intermediate_channels, kernel_size=3, spatial_dims=spatial_dims, ) self.nonlinearity = nn.SiLU() self.conv1 = Convolution( spatial_dims=spatial_dims, in_channels=in_channels, out_channels=self.out_channels, strides=1, kernel_size=3, padding=1, conv_only=True, ) self.upsample = self.downsample = None if self.up: self.upsample = WrappedUpsample( spatial_dims=spatial_dims, mode="nontrainable", in_channels=in_channels, out_channels=in_channels, interp_mode="nearest", scale_factor=2.0, align_corners=None, ) elif down: self.downsample = DiffusionUnetDownsample(spatial_dims, in_channels, use_conv=False) self.time_emb_proj = nn.Linear(temb_channels, self.out_channels) self.norm2 = SPADE( label_nc=label_nc, norm_nc=self.out_channels, norm="GROUP", norm_params={"num_groups": norm_num_groups, "eps": norm_eps, "affine": True}, hidden_channels=spade_intermediate_channels, kernel_size=3, spatial_dims=spatial_dims, ) self.conv2 = zero_module( Convolution( spatial_dims=spatial_dims, in_channels=self.out_channels, out_channels=self.out_channels, strides=1, kernel_size=3, padding=1, conv_only=True, ) ) self.skip_connection: nn.Module if self.out_channels == in_channels: self.skip_connection = nn.Identity() else: self.skip_connection = Convolution( spatial_dims=spatial_dims, in_channels=in_channels, out_channels=self.out_channels, strides=1, kernel_size=1, padding=0, conv_only=True, ) def forward(self, x: torch.Tensor, emb: torch.Tensor, seg: torch.Tensor) -> torch.Tensor: h = x h = self.norm1(h, seg) h = self.nonlinearity(h) if self.upsample is not None: x = self.upsample(x) h = self.upsample(h) elif self.downsample is not None: x = self.downsample(x) h = self.downsample(h) h = self.conv1(h) if self.spatial_dims == 2: temb = self.time_emb_proj(self.nonlinearity(emb))[:, :, None, None] else: temb = self.time_emb_proj(self.nonlinearity(emb))[:, :, None, None, None] h = h + temb h = self.norm2(h, seg) h = self.nonlinearity(h) h = self.conv2(h) output: torch.Tensor = self.skip_connection(x) + h return output class SPADEUpBlock(nn.Module): """ Unet's up block containing resnet and upsamplers blocks. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. prev_output_channel: number of channels from residual connection. out_channels: number of output channels. temb_channels: number of timestep embedding channels. label_nc: number of semantic channels for SPADE normalisation. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_upsample: if True add downsample block. resblock_updown: if True use residual blocks for upsampling. spade_intermediate_channels: number of intermediate channels for SPADE block layer. """ def __init__( self, spatial_dims: int, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, label_nc: int, num_res_blocks: int = 1, norm_num_groups: int = 32, norm_eps: float = 1e-6, add_upsample: bool = True, resblock_updown: bool = False, spade_intermediate_channels: int = 128, ) -> None: super().__init__() self.resblock_updown = resblock_updown resnets = [] for i in range(num_res_blocks): res_skip_channels = in_channels if (i == num_res_blocks - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( SPADEDiffResBlock( spatial_dims=spatial_dims, in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, label_nc=label_nc, norm_num_groups=norm_num_groups, norm_eps=norm_eps, spade_intermediate_channels=spade_intermediate_channels, ) ) self.resnets = nn.ModuleList(resnets) self.upsampler: nn.Module | None if add_upsample: if resblock_updown: self.upsampler = DiffusionUNetResnetBlock( spatial_dims=spatial_dims, in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, norm_num_groups=norm_num_groups, norm_eps=norm_eps, up=True, ) else: post_conv = Convolution( spatial_dims=spatial_dims, in_channels=out_channels, out_channels=out_channels, strides=1, kernel_size=3, padding=1, conv_only=True, ) self.upsampler = WrappedUpsample( spatial_dims=spatial_dims, mode="nontrainable", in_channels=out_channels, out_channels=out_channels, interp_mode="nearest", scale_factor=2.0, post_conv=post_conv, align_corners=None, ) else: self.upsampler = None def forward( self, hidden_states: torch.Tensor, res_hidden_states_list: list[torch.Tensor], temb: torch.Tensor, seg: torch.Tensor, context: torch.Tensor | None = None, ) -> torch.Tensor: del context for resnet in self.resnets: # pop res hidden states res_hidden_states = res_hidden_states_list[-1] res_hidden_states_list = res_hidden_states_list[:-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) hidden_states = resnet(hidden_states, temb, seg) if self.upsampler is not None: hidden_states = self.upsampler(hidden_states, temb) return hidden_states class SPADEAttnUpBlock(nn.Module): """ Unet's up block containing resnet, upsamplers, and self-attention blocks. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. prev_output_channel: number of channels from residual connection. out_channels: number of output channels. temb_channels: number of timestep embedding channels. label_nc: number of semantic channels for SPADE normalisation num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_upsample: if True add downsample block. resblock_updown: if True use residual blocks for upsampling. num_head_channels: number of channels in each attention head. spade_intermediate_channels: number of intermediate channels for SPADE block layer include_fc: whether to include the final linear layer. Default to True. use_combined_linear: whether to use a single linear layer for qkv projection, default to False. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). """ def __init__( self, spatial_dims: int, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, label_nc: int, num_res_blocks: int = 1, norm_num_groups: int = 32, norm_eps: float = 1e-6, add_upsample: bool = True, resblock_updown: bool = False, num_head_channels: int = 1, spade_intermediate_channels: int = 128, include_fc: bool = True, use_combined_linear: bool = False, use_flash_attention: bool = False, ) -> None: super().__init__() self.resblock_updown = resblock_updown resnets = [] attentions = [] for i in range(num_res_blocks): res_skip_channels = in_channels if (i == num_res_blocks - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( SPADEDiffResBlock( spatial_dims=spatial_dims, in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, label_nc=label_nc, norm_num_groups=norm_num_groups, norm_eps=norm_eps, spade_intermediate_channels=spade_intermediate_channels, ) ) attentions.append( SpatialAttentionBlock( spatial_dims=spatial_dims, num_channels=out_channels, num_head_channels=num_head_channels, norm_num_groups=norm_num_groups, norm_eps=norm_eps, include_fc=include_fc, use_combined_linear=use_combined_linear, use_flash_attention=use_flash_attention, ) ) self.resnets = nn.ModuleList(resnets) self.attentions = nn.ModuleList(attentions) self.upsampler: nn.Module | None if add_upsample: if resblock_updown: self.upsampler = DiffusionUNetResnetBlock( spatial_dims=spatial_dims, in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, norm_num_groups=norm_num_groups, norm_eps=norm_eps, up=True, ) else: post_conv = Convolution( spatial_dims=spatial_dims, in_channels=out_channels, out_channels=out_channels, strides=1, kernel_size=3, padding=1, conv_only=True, ) self.upsampler = WrappedUpsample( spatial_dims=spatial_dims, mode="nontrainable", in_channels=out_channels, out_channels=out_channels, interp_mode="nearest", scale_factor=2.0, post_conv=post_conv, align_corners=None, ) else: self.upsampler = None def forward( self, hidden_states: torch.Tensor, res_hidden_states_list: list[torch.Tensor], temb: torch.Tensor, seg: torch.Tensor, context: torch.Tensor | None = None, ) -> torch.Tensor: del context for resnet, attn in zip(self.resnets, self.attentions): # pop res hidden states res_hidden_states = res_hidden_states_list[-1] res_hidden_states_list = res_hidden_states_list[:-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) hidden_states = resnet(hidden_states, temb, seg) hidden_states = attn(hidden_states).contiguous() if self.upsampler is not None: hidden_states = self.upsampler(hidden_states, temb) return hidden_states class SPADECrossAttnUpBlock(nn.Module): """ Unet's up block containing resnet, upsamplers, and self-attention blocks. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. prev_output_channel: number of channels from residual connection. out_channels: number of output channels. temb_channels: number of timestep embedding channels. label_nc: number of semantic channels for SPADE normalisation. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_upsample: if True add downsample block. resblock_updown: if True use residual blocks for upsampling. num_head_channels: number of channels in each attention head. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. upcast_attention: if True, upcast attention operations to full precision. spade_intermediate_channels: number of intermediate channels for SPADE block layer. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism. (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). """ def __init__( self, spatial_dims: int, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, label_nc: int, num_res_blocks: int = 1, norm_num_groups: int = 32, norm_eps: float = 1e-6, add_upsample: bool = True, resblock_updown: bool = False, num_head_channels: int = 1, transformer_num_layers: int = 1, cross_attention_dim: int | None = None, upcast_attention: bool = False, spade_intermediate_channels: int = 128, include_fc: bool = True, use_combined_linear: bool = False, use_flash_attention: bool = False, ) -> None: super().__init__() self.resblock_updown = resblock_updown resnets = [] attentions = [] for i in range(num_res_blocks): res_skip_channels = in_channels if (i == num_res_blocks - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( SPADEDiffResBlock( spatial_dims=spatial_dims, in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, norm_num_groups=norm_num_groups, norm_eps=norm_eps, label_nc=label_nc, spade_intermediate_channels=spade_intermediate_channels, ) ) attentions.append( SpatialTransformer( spatial_dims=spatial_dims, in_channels=out_channels, num_attention_heads=out_channels // num_head_channels, num_head_channels=num_head_channels, norm_num_groups=norm_num_groups, norm_eps=norm_eps, num_layers=transformer_num_layers, cross_attention_dim=cross_attention_dim, upcast_attention=upcast_attention, include_fc=include_fc, use_combined_linear=use_combined_linear, use_flash_attention=use_flash_attention, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.upsampler: nn.Module | None if add_upsample: if resblock_updown: self.upsampler = DiffusionUNetResnetBlock( spatial_dims=spatial_dims, in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, norm_num_groups=norm_num_groups, norm_eps=norm_eps, up=True, ) else: post_conv = Convolution( spatial_dims=spatial_dims, in_channels=out_channels, out_channels=out_channels, strides=1, kernel_size=3, padding=1, conv_only=True, ) self.upsampler = WrappedUpsample( spatial_dims=spatial_dims, mode="nontrainable", in_channels=out_channels, out_channels=out_channels, interp_mode="nearest", scale_factor=2.0, post_conv=post_conv, align_corners=None, ) else: self.upsampler = None def forward( self, hidden_states: torch.Tensor, res_hidden_states_list: list[torch.Tensor], temb: torch.Tensor, seg: torch.Tensor | None = None, context: torch.Tensor | None = None, ) -> torch.Tensor: for resnet, attn in zip(self.resnets, self.attentions): # pop res hidden states res_hidden_states = res_hidden_states_list[-1] res_hidden_states_list = res_hidden_states_list[:-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) hidden_states = resnet(hidden_states, temb, seg) hidden_states = attn(hidden_states, context=context).contiguous() if self.upsampler is not None: hidden_states = self.upsampler(hidden_states, temb) return hidden_states def get_spade_up_block( spatial_dims: int, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, num_res_blocks: int, norm_num_groups: int, norm_eps: float, add_upsample: bool, resblock_updown: bool, with_attn: bool, with_cross_attn: bool, num_head_channels: int, transformer_num_layers: int, label_nc: int, cross_attention_dim: int | None, upcast_attention: bool = False, spade_intermediate_channels: int = 128, include_fc: bool = True, use_combined_linear: bool = False, use_flash_attention: bool = False, ) -> nn.Module: if with_attn: return SPADEAttnUpBlock( spatial_dims=spatial_dims, in_channels=in_channels, prev_output_channel=prev_output_channel, out_channels=out_channels, temb_channels=temb_channels, label_nc=label_nc, num_res_blocks=num_res_blocks, norm_num_groups=norm_num_groups, norm_eps=norm_eps, add_upsample=add_upsample, resblock_updown=resblock_updown, num_head_channels=num_head_channels, spade_intermediate_channels=spade_intermediate_channels, include_fc=include_fc, use_combined_linear=use_combined_linear, use_flash_attention=use_flash_attention, ) elif with_cross_attn: return SPADECrossAttnUpBlock( spatial_dims=spatial_dims, in_channels=in_channels, prev_output_channel=prev_output_channel, out_channels=out_channels, temb_channels=temb_channels, label_nc=label_nc, num_res_blocks=num_res_blocks, norm_num_groups=norm_num_groups, norm_eps=norm_eps, add_upsample=add_upsample, resblock_updown=resblock_updown, num_head_channels=num_head_channels, transformer_num_layers=transformer_num_layers, cross_attention_dim=cross_attention_dim, upcast_attention=upcast_attention, spade_intermediate_channels=spade_intermediate_channels, use_flash_attention=use_flash_attention, ) else: return SPADEUpBlock( spatial_dims=spatial_dims, in_channels=in_channels, prev_output_channel=prev_output_channel, out_channels=out_channels, temb_channels=temb_channels, label_nc=label_nc, num_res_blocks=num_res_blocks, norm_num_groups=norm_num_groups, norm_eps=norm_eps, add_upsample=add_upsample, resblock_updown=resblock_updown, spade_intermediate_channels=spade_intermediate_channels, ) class SPADEDiffusionModelUNet(nn.Module): """ UNet network with timestep embedding and attention mechanisms for conditioning, with added SPADE normalization for semantic conditioning (Park et.al (2019): https://github.com/NVlabs/SPADE). An example tutorial can be found at https://github.com/Project-MONAI/GenerativeModels/tree/main/tutorials/generative/2d_spade_ldm Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. label_nc: number of semantic channels for SPADE normalisation. num_res_blocks: number of residual blocks (see ResnetBlock) per level. channels: tuple of block output channels. attention_levels: list of levels to add attention. norm_num_groups: number of groups for the normalization. norm_eps: epsilon for the normalization. resblock_updown: if True use residual blocks for up/downsampling. num_head_channels: number of channels in each attention head. with_conditioning: if True add spatial transformers to perform conditioning. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. num_class_embeds: if specified (as an int), then this model will be class-conditional with `num_class_embeds` classes. upcast_attention: if True, upcast attention operations to full precision. spade_intermediate_channels: number of intermediate channels for SPADE block layer. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). """ def __init__( self, spatial_dims: int, in_channels: int, out_channels: int, label_nc: int, num_res_blocks: Sequence[int] | int = (2, 2, 2, 2), channels: Sequence[int] = (32, 64, 64, 64), attention_levels: Sequence[bool] = (False, False, True, True), norm_num_groups: int = 32, norm_eps: float = 1e-6, resblock_updown: bool = False, num_head_channels: int | Sequence[int] = 8, with_conditioning: bool = False, transformer_num_layers: int = 1, cross_attention_dim: int | None = None, num_class_embeds: int | None = None, upcast_attention: bool = False, spade_intermediate_channels: int = 128, include_fc: bool = True, use_combined_linear: bool = False, use_flash_attention: bool = False, ) -> None: super().__init__() if with_conditioning is True and cross_attention_dim is None: raise ValueError( "SPADEDiffusionModelUNet expects dimension of the cross-attention conditioning (cross_attention_dim) " "when using with_conditioning." ) if cross_attention_dim is not None and with_conditioning is False: raise ValueError( "SPADEDiffusionModelUNet expects with_conditioning=True when specifying the cross_attention_dim." ) # All number of channels should be multiple of num_groups if any((out_channel % norm_num_groups) != 0 for out_channel in channels): raise ValueError("SPADEDiffusionModelUNet expects all num_channels being multiple of norm_num_groups") if len(channels) != len(attention_levels): raise ValueError("SPADEDiffusionModelUNet expects num_channels being same size of attention_levels") if isinstance(num_head_channels, int): num_head_channels = ensure_tuple_rep(num_head_channels, len(attention_levels)) if len(num_head_channels) != len(attention_levels): raise ValueError( "num_head_channels should have the same length as attention_levels. For the i levels without attention," " i.e. `attention_level[i]=False`, the num_head_channels[i] will be ignored." ) if isinstance(num_res_blocks, int): num_res_blocks = ensure_tuple_rep(num_res_blocks, len(channels)) if len(num_res_blocks) != len(channels): raise ValueError( "`num_res_blocks` should be a single integer or a tuple of integers with the same length as " "`num_channels`." ) self.in_channels = in_channels self.block_out_channels = channels self.out_channels = out_channels self.num_res_blocks = num_res_blocks self.attention_levels = attention_levels self.num_head_channels = num_head_channels self.with_conditioning = with_conditioning self.label_nc = label_nc # input self.conv_in = Convolution( spatial_dims=spatial_dims, in_channels=in_channels, out_channels=channels[0], strides=1, kernel_size=3, padding=1, conv_only=True, ) # time time_embed_dim = channels[0] * 4 self.time_embed = nn.Sequential( nn.Linear(channels[0], time_embed_dim), nn.SiLU(), nn.Linear(time_embed_dim, time_embed_dim) ) # class embedding self.num_class_embeds = num_class_embeds if num_class_embeds is not None: self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) # down self.down_blocks = nn.ModuleList([]) output_channel = channels[0] for i in range(len(channels)): input_channel = output_channel output_channel = channels[i] is_final_block = i == len(channels) - 1 down_block = get_down_block( spatial_dims=spatial_dims, in_channels=input_channel, out_channels=output_channel, temb_channels=time_embed_dim, num_res_blocks=num_res_blocks[i], norm_num_groups=norm_num_groups, norm_eps=norm_eps, add_downsample=not is_final_block, resblock_updown=resblock_updown, with_attn=(attention_levels[i] and not with_conditioning), with_cross_attn=(attention_levels[i] and with_conditioning), num_head_channels=num_head_channels[i], transformer_num_layers=transformer_num_layers, cross_attention_dim=cross_attention_dim, upcast_attention=upcast_attention, include_fc=include_fc, use_combined_linear=use_combined_linear, use_flash_attention=use_flash_attention, ) self.down_blocks.append(down_block) # mid self.middle_block = get_mid_block( spatial_dims=spatial_dims, in_channels=channels[-1], temb_channels=time_embed_dim, norm_num_groups=norm_num_groups, norm_eps=norm_eps, with_conditioning=with_conditioning, num_head_channels=num_head_channels[-1], transformer_num_layers=transformer_num_layers, cross_attention_dim=cross_attention_dim, upcast_attention=upcast_attention, include_fc=include_fc, use_combined_linear=use_combined_linear, use_flash_attention=use_flash_attention, ) # up self.up_blocks = nn.ModuleList([]) reversed_block_out_channels = list(reversed(channels)) reversed_num_res_blocks = list(reversed(num_res_blocks)) reversed_attention_levels = list(reversed(attention_levels)) reversed_num_head_channels = list(reversed(num_head_channels)) output_channel = reversed_block_out_channels[0] for i in range(len(reversed_block_out_channels)): prev_output_channel = output_channel output_channel = reversed_block_out_channels[i] input_channel = reversed_block_out_channels[min(i + 1, len(channels) - 1)] is_final_block = i == len(channels) - 1 up_block = get_spade_up_block( spatial_dims=spatial_dims, in_channels=input_channel, prev_output_channel=prev_output_channel, out_channels=output_channel, temb_channels=time_embed_dim, num_res_blocks=reversed_num_res_blocks[i] + 1, norm_num_groups=norm_num_groups, norm_eps=norm_eps, add_upsample=not is_final_block, resblock_updown=resblock_updown, with_attn=(reversed_attention_levels[i] and not with_conditioning), with_cross_attn=(reversed_attention_levels[i] and with_conditioning), num_head_channels=reversed_num_head_channels[i], transformer_num_layers=transformer_num_layers, cross_attention_dim=cross_attention_dim, upcast_attention=upcast_attention, label_nc=label_nc, spade_intermediate_channels=spade_intermediate_channels, use_flash_attention=use_flash_attention, ) self.up_blocks.append(up_block) # out self.out = nn.Sequential( nn.GroupNorm(num_groups=norm_num_groups, num_channels=channels[0], eps=norm_eps, affine=True), nn.SiLU(), zero_module( Convolution( spatial_dims=spatial_dims, in_channels=channels[0], out_channels=out_channels, strides=1, kernel_size=3, padding=1, conv_only=True, ) ), ) def forward( self, x: torch.Tensor, timesteps: torch.Tensor, seg: torch.Tensor, context: torch.Tensor | None = None, class_labels: torch.Tensor | None = None, down_block_additional_residuals: tuple[torch.Tensor] | None = None, mid_block_additional_residual: torch.Tensor | None = None, ) -> torch.Tensor: """ Args: x: input tensor (N, C, SpatialDims). timesteps: timestep tensor (N,). seg: Bx[LABEL_NC]x[SPATIAL DIMENSIONS] tensor of segmentations for SPADE norm. context: context tensor (N, 1, ContextDim). class_labels: context tensor (N, ). down_block_additional_residuals: additional residual tensors for down blocks (N, C, FeatureMapsDims). mid_block_additional_residual: additional residual tensor for mid block (N, C, FeatureMapsDims). """ # 1. time t_emb = get_timestep_embedding(timesteps, self.block_out_channels[0]) # timesteps does not contain any weights and will always return f32 tensors # but time_embedding might actually be running in fp16. so we need to cast here. # there might be better ways to encapsulate this. t_emb = t_emb.to(dtype=x.dtype) emb = self.time_embed(t_emb) # 2. class if self.num_class_embeds is not None: if class_labels is None: raise ValueError("class_labels should be provided when num_class_embeds > 0") class_emb = self.class_embedding(class_labels) class_emb = class_emb.to(dtype=x.dtype) emb = emb + class_emb # 3. initial convolution h = self.conv_in(x) # 4. down if context is not None and self.with_conditioning is False: raise ValueError("model should have with_conditioning = True if context is provided") down_block_res_samples: list[torch.Tensor] = [h] for downsample_block in self.down_blocks: h, res_samples = downsample_block(hidden_states=h, temb=emb, context=context) for residual in res_samples: down_block_res_samples.append(residual) # Additional residual conections for Controlnets if down_block_additional_residuals is not None: new_down_block_res_samples: list[torch.Tensor] = [h] for down_block_res_sample, down_block_additional_residual in zip( down_block_res_samples, down_block_additional_residuals ): down_block_res_sample = down_block_res_sample + down_block_additional_residual new_down_block_res_samples.append(down_block_res_sample) down_block_res_samples = new_down_block_res_samples # 5. mid h = self.middle_block(hidden_states=h, temb=emb, context=context) # Additional residual conections for Controlnets if mid_block_additional_residual is not None: h = h + mid_block_additional_residual # 6. up for upsample_block in self.up_blocks: idx: int = -len(upsample_block.resnets) # type: ignore res_samples = down_block_res_samples[idx:] down_block_res_samples = down_block_res_samples[:idx] h = upsample_block(hidden_states=h, res_hidden_states_list=res_samples, seg=seg, temb=emb, context=context) # 7. output block output: torch.Tensor = self.out(h) return output