Spaces:
Runtime error
Runtime error
| # Copyright 2022 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 | |
| import numpy as np | |
| # limitations under the License. | |
| import torch | |
| from torch import nn | |
| from .attention import AttentionBlock, SpatialTransformer | |
| from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, ResnetBlock2D, Upsample2D | |
| def get_down_block( | |
| down_block_type, | |
| num_layers, | |
| in_channels, | |
| out_channels, | |
| temb_channels, | |
| add_downsample, | |
| resnet_eps, | |
| resnet_act_fn, | |
| attn_num_head_channels, | |
| cross_attention_dim=None, | |
| downsample_padding=None, | |
| ): | |
| down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type | |
| if down_block_type == "DownBlock2D": | |
| return DownBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| add_downsample=add_downsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| downsample_padding=downsample_padding, | |
| ) | |
| elif down_block_type == "AttnDownBlock2D": | |
| return AttnDownBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| add_downsample=add_downsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| downsample_padding=downsample_padding, | |
| attn_num_head_channels=attn_num_head_channels, | |
| ) | |
| elif down_block_type == "CrossAttnDownBlock2D": | |
| if cross_attention_dim is None: | |
| raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D") | |
| return CrossAttnDownBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| add_downsample=add_downsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| downsample_padding=downsample_padding, | |
| cross_attention_dim=cross_attention_dim, | |
| attn_num_head_channels=attn_num_head_channels, | |
| ) | |
| elif down_block_type == "SkipDownBlock2D": | |
| return SkipDownBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| add_downsample=add_downsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| downsample_padding=downsample_padding, | |
| ) | |
| elif down_block_type == "AttnSkipDownBlock2D": | |
| return AttnSkipDownBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| add_downsample=add_downsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| downsample_padding=downsample_padding, | |
| attn_num_head_channels=attn_num_head_channels, | |
| ) | |
| elif down_block_type == "DownEncoderBlock2D": | |
| return DownEncoderBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| add_downsample=add_downsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| downsample_padding=downsample_padding, | |
| ) | |
| def get_up_block( | |
| up_block_type, | |
| num_layers, | |
| in_channels, | |
| out_channels, | |
| prev_output_channel, | |
| temb_channels, | |
| add_upsample, | |
| resnet_eps, | |
| resnet_act_fn, | |
| attn_num_head_channels, | |
| cross_attention_dim=None, | |
| ): | |
| up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type | |
| if up_block_type == "UpBlock2D": | |
| return UpBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| prev_output_channel=prev_output_channel, | |
| temb_channels=temb_channels, | |
| add_upsample=add_upsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| ) | |
| elif up_block_type == "CrossAttnUpBlock2D": | |
| if cross_attention_dim is None: | |
| raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D") | |
| return CrossAttnUpBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| prev_output_channel=prev_output_channel, | |
| temb_channels=temb_channels, | |
| add_upsample=add_upsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| cross_attention_dim=cross_attention_dim, | |
| attn_num_head_channels=attn_num_head_channels, | |
| ) | |
| elif up_block_type == "AttnUpBlock2D": | |
| return AttnUpBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| prev_output_channel=prev_output_channel, | |
| temb_channels=temb_channels, | |
| add_upsample=add_upsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| attn_num_head_channels=attn_num_head_channels, | |
| ) | |
| elif up_block_type == "SkipUpBlock2D": | |
| return SkipUpBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| prev_output_channel=prev_output_channel, | |
| temb_channels=temb_channels, | |
| add_upsample=add_upsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| ) | |
| elif up_block_type == "AttnSkipUpBlock2D": | |
| return AttnSkipUpBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| prev_output_channel=prev_output_channel, | |
| temb_channels=temb_channels, | |
| add_upsample=add_upsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| attn_num_head_channels=attn_num_head_channels, | |
| ) | |
| elif up_block_type == "UpDecoderBlock2D": | |
| return UpDecoderBlock2D( | |
| num_layers=num_layers, | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| add_upsample=add_upsample, | |
| resnet_eps=resnet_eps, | |
| resnet_act_fn=resnet_act_fn, | |
| ) | |
| raise ValueError(f"{up_block_type} does not exist.") | |
| class UNetMidBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| attention_type="default", | |
| output_scale_factor=1.0, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.attention_type = attention_type | |
| resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) | |
| # there is always at least one resnet | |
| resnets = [ | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=in_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ] | |
| attentions = [] | |
| for _ in range(num_layers): | |
| attentions.append( | |
| AttentionBlock( | |
| in_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| num_groups=resnet_groups, | |
| ) | |
| ) | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=in_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| def forward(self, hidden_states, temb=None, encoder_states=None): | |
| hidden_states = self.resnets[0](hidden_states, temb) | |
| for attn, resnet in zip(self.attentions, self.resnets[1:]): | |
| if self.attention_type == "default": | |
| hidden_states = attn(hidden_states) | |
| else: | |
| hidden_states = attn(hidden_states, encoder_states) | |
| hidden_states = resnet(hidden_states, temb) | |
| return hidden_states | |
| class UNetMidBlock2DCrossAttn(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| attention_type="default", | |
| output_scale_factor=1.0, | |
| cross_attention_dim=1280, | |
| **kwargs, | |
| ): | |
| super().__init__() | |
| self.attention_type = attention_type | |
| self.attn_num_head_channels = attn_num_head_channels | |
| resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) | |
| # there is always at least one resnet | |
| resnets = [ | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=in_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ] | |
| attentions = [] | |
| for _ in range(num_layers): | |
| attentions.append( | |
| SpatialTransformer( | |
| in_channels, | |
| attn_num_head_channels, | |
| in_channels // attn_num_head_channels, | |
| depth=1, | |
| context_dim=cross_attention_dim, | |
| ) | |
| ) | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=in_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| def set_attention_slice(self, slice_size): | |
| if slice_size is not None and self.attn_num_head_channels % slice_size != 0: | |
| raise ValueError( | |
| f"Make sure slice_size {slice_size} is a divisor of " | |
| f"the number of heads used in cross_attention {self.attn_num_head_channels}" | |
| ) | |
| if slice_size is not None and slice_size > self.attn_num_head_channels: | |
| raise ValueError( | |
| f"Chunk_size {slice_size} has to be smaller or equal to " | |
| f"the number of heads used in cross_attention {self.attn_num_head_channels}" | |
| ) | |
| for attn in self.attentions: | |
| attn._set_attention_slice(slice_size) | |
| def forward(self, hidden_states, temb=None, encoder_hidden_states=None): | |
| hidden_states = self.resnets[0](hidden_states, temb) | |
| for attn, resnet in zip(self.attentions, self.resnets[1:]): | |
| hidden_states = attn(hidden_states, encoder_hidden_states) | |
| hidden_states = resnet(hidden_states, temb) | |
| return hidden_states | |
| class AttnDownBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| attention_type="default", | |
| output_scale_factor=1.0, | |
| downsample_padding=1, | |
| add_downsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| attentions = [] | |
| self.attention_type = attention_type | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| attentions.append( | |
| AttentionBlock( | |
| out_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_downsample: | |
| self.downsamplers = nn.ModuleList( | |
| [ | |
| Downsample2D( | |
| in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" | |
| ) | |
| ] | |
| ) | |
| else: | |
| self.downsamplers = None | |
| def forward(self, hidden_states, temb=None): | |
| output_states = () | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| hidden_states = resnet(hidden_states, temb) | |
| hidden_states = attn(hidden_states) | |
| output_states += (hidden_states,) | |
| if self.downsamplers is not None: | |
| for downsampler in self.downsamplers: | |
| hidden_states = downsampler(hidden_states) | |
| output_states += (hidden_states,) | |
| return hidden_states, output_states | |
| class CrossAttnDownBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| cross_attention_dim=1280, | |
| attention_type="default", | |
| output_scale_factor=1.0, | |
| downsample_padding=1, | |
| add_downsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| attentions = [] | |
| self.attention_type = attention_type | |
| self.attn_num_head_channels = attn_num_head_channels | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| attentions.append( | |
| SpatialTransformer( | |
| out_channels, | |
| attn_num_head_channels, | |
| out_channels // attn_num_head_channels, | |
| depth=1, | |
| context_dim=cross_attention_dim, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_downsample: | |
| self.downsamplers = nn.ModuleList( | |
| [ | |
| Downsample2D( | |
| in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" | |
| ) | |
| ] | |
| ) | |
| else: | |
| self.downsamplers = None | |
| def set_attention_slice(self, slice_size): | |
| if slice_size is not None and self.attn_num_head_channels % slice_size != 0: | |
| raise ValueError( | |
| f"Make sure slice_size {slice_size} is a divisor of " | |
| f"the number of heads used in cross_attention {self.attn_num_head_channels}" | |
| ) | |
| if slice_size is not None and slice_size > self.attn_num_head_channels: | |
| raise ValueError( | |
| f"Chunk_size {slice_size} has to be smaller or equal to " | |
| f"the number of heads used in cross_attention {self.attn_num_head_channels}" | |
| ) | |
| for attn in self.attentions: | |
| attn._set_attention_slice(slice_size) | |
| def forward(self, hidden_states, temb=None, encoder_hidden_states=None): | |
| output_states = () | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| hidden_states = resnet(hidden_states, temb) | |
| hidden_states = attn(hidden_states, context=encoder_hidden_states) | |
| output_states += (hidden_states,) | |
| if self.downsamplers is not None: | |
| for downsampler in self.downsamplers: | |
| hidden_states = downsampler(hidden_states) | |
| output_states += (hidden_states,) | |
| return hidden_states, output_states | |
| class DownBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| output_scale_factor=1.0, | |
| add_downsample=True, | |
| downsample_padding=1, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_downsample: | |
| self.downsamplers = nn.ModuleList( | |
| [ | |
| Downsample2D( | |
| in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" | |
| ) | |
| ] | |
| ) | |
| else: | |
| self.downsamplers = None | |
| def forward(self, hidden_states, temb=None): | |
| output_states = () | |
| for resnet in self.resnets: | |
| hidden_states = resnet(hidden_states, temb) | |
| output_states += (hidden_states,) | |
| if self.downsamplers is not None: | |
| for downsampler in self.downsamplers: | |
| hidden_states = downsampler(hidden_states) | |
| output_states += (hidden_states,) | |
| return hidden_states, output_states | |
| class DownEncoderBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| output_scale_factor=1.0, | |
| add_downsample=True, | |
| downsample_padding=1, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=None, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_downsample: | |
| self.downsamplers = nn.ModuleList( | |
| [ | |
| Downsample2D( | |
| in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" | |
| ) | |
| ] | |
| ) | |
| else: | |
| self.downsamplers = None | |
| def forward(self, hidden_states): | |
| for resnet in self.resnets: | |
| hidden_states = resnet(hidden_states, temb=None) | |
| if self.downsamplers is not None: | |
| for downsampler in self.downsamplers: | |
| hidden_states = downsampler(hidden_states) | |
| return hidden_states | |
| class AttnDownEncoderBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| output_scale_factor=1.0, | |
| add_downsample=True, | |
| downsample_padding=1, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| attentions = [] | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=None, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| attentions.append( | |
| AttentionBlock( | |
| out_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| num_groups=resnet_groups, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_downsample: | |
| self.downsamplers = nn.ModuleList( | |
| [ | |
| Downsample2D( | |
| in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op" | |
| ) | |
| ] | |
| ) | |
| else: | |
| self.downsamplers = None | |
| def forward(self, hidden_states): | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| hidden_states = resnet(hidden_states, temb=None) | |
| hidden_states = attn(hidden_states) | |
| if self.downsamplers is not None: | |
| for downsampler in self.downsamplers: | |
| hidden_states = downsampler(hidden_states) | |
| return hidden_states | |
| class AttnSkipDownBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| attention_type="default", | |
| output_scale_factor=np.sqrt(2.0), | |
| downsample_padding=1, | |
| add_downsample=True, | |
| ): | |
| super().__init__() | |
| self.attentions = nn.ModuleList([]) | |
| self.resnets = nn.ModuleList([]) | |
| self.attention_type = attention_type | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| self.resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(in_channels // 4, 32), | |
| groups_out=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.attentions.append( | |
| AttentionBlock( | |
| out_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| ) | |
| ) | |
| if add_downsample: | |
| self.resnet_down = ResnetBlock2D( | |
| in_channels=out_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| use_nin_shortcut=True, | |
| down=True, | |
| kernel="fir", | |
| ) | |
| self.downsamplers = nn.ModuleList([FirDownsample2D(in_channels, out_channels=out_channels)]) | |
| self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) | |
| else: | |
| self.resnet_down = None | |
| self.downsamplers = None | |
| self.skip_conv = None | |
| def forward(self, hidden_states, temb=None, skip_sample=None): | |
| output_states = () | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| hidden_states = resnet(hidden_states, temb) | |
| hidden_states = attn(hidden_states) | |
| output_states += (hidden_states,) | |
| if self.downsamplers is not None: | |
| hidden_states = self.resnet_down(hidden_states, temb) | |
| for downsampler in self.downsamplers: | |
| skip_sample = downsampler(skip_sample) | |
| hidden_states = self.skip_conv(skip_sample) + hidden_states | |
| output_states += (hidden_states,) | |
| return hidden_states, output_states, skip_sample | |
| class SkipDownBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_pre_norm: bool = True, | |
| output_scale_factor=np.sqrt(2.0), | |
| add_downsample=True, | |
| downsample_padding=1, | |
| ): | |
| super().__init__() | |
| self.resnets = nn.ModuleList([]) | |
| for i in range(num_layers): | |
| in_channels = in_channels if i == 0 else out_channels | |
| self.resnets.append( | |
| ResnetBlock2D( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(in_channels // 4, 32), | |
| groups_out=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| if add_downsample: | |
| self.resnet_down = ResnetBlock2D( | |
| in_channels=out_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| use_nin_shortcut=True, | |
| down=True, | |
| kernel="fir", | |
| ) | |
| self.downsamplers = nn.ModuleList([FirDownsample2D(in_channels, out_channels=out_channels)]) | |
| self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) | |
| else: | |
| self.resnet_down = None | |
| self.downsamplers = None | |
| self.skip_conv = None | |
| def forward(self, hidden_states, temb=None, skip_sample=None): | |
| output_states = () | |
| for resnet in self.resnets: | |
| hidden_states = resnet(hidden_states, temb) | |
| output_states += (hidden_states,) | |
| if self.downsamplers is not None: | |
| hidden_states = self.resnet_down(hidden_states, temb) | |
| for downsampler in self.downsamplers: | |
| skip_sample = downsampler(skip_sample) | |
| hidden_states = self.skip_conv(skip_sample) + hidden_states | |
| output_states += (hidden_states,) | |
| return hidden_states, output_states, skip_sample | |
| class AttnUpBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| prev_output_channel: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attention_type="default", | |
| attn_num_head_channels=1, | |
| output_scale_factor=1.0, | |
| add_upsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| attentions = [] | |
| self.attention_type = attention_type | |
| for i in range(num_layers): | |
| res_skip_channels = in_channels if (i == num_layers - 1) else out_channels | |
| resnet_in_channels = prev_output_channel if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=resnet_in_channels + res_skip_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| attentions.append( | |
| AttentionBlock( | |
| out_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_upsample: | |
| self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) | |
| else: | |
| self.upsamplers = None | |
| def forward(self, hidden_states, res_hidden_states_tuple, temb=None): | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| # pop res hidden states | |
| res_hidden_states = res_hidden_states_tuple[-1] | |
| res_hidden_states_tuple = res_hidden_states_tuple[:-1] | |
| hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) | |
| hidden_states = resnet(hidden_states, temb) | |
| hidden_states = attn(hidden_states) | |
| if self.upsamplers is not None: | |
| for upsampler in self.upsamplers: | |
| hidden_states = upsampler(hidden_states) | |
| return hidden_states | |
| class CrossAttnUpBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| prev_output_channel: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| cross_attention_dim=1280, | |
| attention_type="default", | |
| output_scale_factor=1.0, | |
| downsample_padding=1, | |
| add_upsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| attentions = [] | |
| self.attention_type = attention_type | |
| self.attn_num_head_channels = attn_num_head_channels | |
| for i in range(num_layers): | |
| res_skip_channels = in_channels if (i == num_layers - 1) else out_channels | |
| resnet_in_channels = prev_output_channel if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=resnet_in_channels + res_skip_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| attentions.append( | |
| SpatialTransformer( | |
| out_channels, | |
| attn_num_head_channels, | |
| out_channels // attn_num_head_channels, | |
| depth=1, | |
| context_dim=cross_attention_dim, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_upsample: | |
| self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) | |
| else: | |
| self.upsamplers = None | |
| def set_attention_slice(self, slice_size): | |
| if slice_size is not None and self.attn_num_head_channels % slice_size != 0: | |
| raise ValueError( | |
| f"Make sure slice_size {slice_size} is a divisor of " | |
| f"the number of heads used in cross_attention {self.attn_num_head_channels}" | |
| ) | |
| if slice_size is not None and slice_size > self.attn_num_head_channels: | |
| raise ValueError( | |
| f"Chunk_size {slice_size} has to be smaller or equal to " | |
| f"the number of heads used in cross_attention {self.attn_num_head_channels}" | |
| ) | |
| for attn in self.attentions: | |
| attn._set_attention_slice(slice_size) | |
| def forward(self, hidden_states, res_hidden_states_tuple, temb=None, encoder_hidden_states=None): | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| # pop res hidden states | |
| res_hidden_states = res_hidden_states_tuple[-1] | |
| res_hidden_states_tuple = res_hidden_states_tuple[:-1] | |
| hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) | |
| hidden_states = resnet(hidden_states, temb) | |
| hidden_states = attn(hidden_states, context=encoder_hidden_states) | |
| if self.upsamplers is not None: | |
| for upsampler in self.upsamplers: | |
| hidden_states = upsampler(hidden_states) | |
| return hidden_states | |
| class UpBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| prev_output_channel: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| output_scale_factor=1.0, | |
| add_upsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| for i in range(num_layers): | |
| res_skip_channels = in_channels if (i == num_layers - 1) else out_channels | |
| resnet_in_channels = prev_output_channel if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=resnet_in_channels + res_skip_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_upsample: | |
| self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) | |
| else: | |
| self.upsamplers = None | |
| def forward(self, hidden_states, res_hidden_states_tuple, temb=None): | |
| for resnet in self.resnets: | |
| # pop res hidden states | |
| res_hidden_states = res_hidden_states_tuple[-1] | |
| res_hidden_states_tuple = res_hidden_states_tuple[:-1] | |
| hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) | |
| hidden_states = resnet(hidden_states, temb) | |
| if self.upsamplers is not None: | |
| for upsampler in self.upsamplers: | |
| hidden_states = upsampler(hidden_states) | |
| return hidden_states | |
| class UpDecoderBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| output_scale_factor=1.0, | |
| add_upsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| for i in range(num_layers): | |
| input_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=input_channels, | |
| out_channels=out_channels, | |
| temb_channels=None, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_upsample: | |
| self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) | |
| else: | |
| self.upsamplers = None | |
| def forward(self, hidden_states): | |
| for resnet in self.resnets: | |
| hidden_states = resnet(hidden_states, temb=None) | |
| if self.upsamplers is not None: | |
| for upsampler in self.upsamplers: | |
| hidden_states = upsampler(hidden_states) | |
| return hidden_states | |
| class AttnUpDecoderBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| out_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_groups: int = 32, | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| output_scale_factor=1.0, | |
| add_upsample=True, | |
| ): | |
| super().__init__() | |
| resnets = [] | |
| attentions = [] | |
| for i in range(num_layers): | |
| input_channels = in_channels if i == 0 else out_channels | |
| resnets.append( | |
| ResnetBlock2D( | |
| in_channels=input_channels, | |
| out_channels=out_channels, | |
| temb_channels=None, | |
| eps=resnet_eps, | |
| groups=resnet_groups, | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| attentions.append( | |
| AttentionBlock( | |
| out_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| num_groups=resnet_groups, | |
| ) | |
| ) | |
| self.attentions = nn.ModuleList(attentions) | |
| self.resnets = nn.ModuleList(resnets) | |
| if add_upsample: | |
| self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) | |
| else: | |
| self.upsamplers = None | |
| def forward(self, hidden_states): | |
| for resnet, attn in zip(self.resnets, self.attentions): | |
| hidden_states = resnet(hidden_states, temb=None) | |
| hidden_states = attn(hidden_states) | |
| if self.upsamplers is not None: | |
| for upsampler in self.upsamplers: | |
| hidden_states = upsampler(hidden_states) | |
| return hidden_states | |
| class AttnSkipUpBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| prev_output_channel: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_pre_norm: bool = True, | |
| attn_num_head_channels=1, | |
| attention_type="default", | |
| output_scale_factor=np.sqrt(2.0), | |
| upsample_padding=1, | |
| add_upsample=True, | |
| ): | |
| super().__init__() | |
| self.attentions = nn.ModuleList([]) | |
| self.resnets = nn.ModuleList([]) | |
| self.attention_type = attention_type | |
| for i in range(num_layers): | |
| res_skip_channels = in_channels if (i == num_layers - 1) else out_channels | |
| resnet_in_channels = prev_output_channel if i == 0 else out_channels | |
| self.resnets.append( | |
| ResnetBlock2D( | |
| in_channels=resnet_in_channels + res_skip_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(resnet_in_channels + res_skip_channels // 4, 32), | |
| groups_out=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.attentions.append( | |
| AttentionBlock( | |
| out_channels, | |
| num_head_channels=attn_num_head_channels, | |
| rescale_output_factor=output_scale_factor, | |
| eps=resnet_eps, | |
| ) | |
| ) | |
| self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) | |
| if add_upsample: | |
| self.resnet_up = ResnetBlock2D( | |
| in_channels=out_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(out_channels // 4, 32), | |
| groups_out=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| use_nin_shortcut=True, | |
| up=True, | |
| kernel="fir", | |
| ) | |
| self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) | |
| self.skip_norm = torch.nn.GroupNorm( | |
| num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True | |
| ) | |
| self.act = nn.SiLU() | |
| else: | |
| self.resnet_up = None | |
| self.skip_conv = None | |
| self.skip_norm = None | |
| self.act = None | |
| def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None): | |
| for resnet in self.resnets: | |
| # pop res hidden states | |
| res_hidden_states = res_hidden_states_tuple[-1] | |
| res_hidden_states_tuple = res_hidden_states_tuple[:-1] | |
| hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) | |
| hidden_states = resnet(hidden_states, temb) | |
| hidden_states = self.attentions[0](hidden_states) | |
| if skip_sample is not None: | |
| skip_sample = self.upsampler(skip_sample) | |
| else: | |
| skip_sample = 0 | |
| if self.resnet_up is not None: | |
| skip_sample_states = self.skip_norm(hidden_states) | |
| skip_sample_states = self.act(skip_sample_states) | |
| skip_sample_states = self.skip_conv(skip_sample_states) | |
| skip_sample = skip_sample + skip_sample_states | |
| hidden_states = self.resnet_up(hidden_states, temb) | |
| return hidden_states, skip_sample | |
| class SkipUpBlock2D(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels: int, | |
| prev_output_channel: int, | |
| out_channels: int, | |
| temb_channels: int, | |
| dropout: float = 0.0, | |
| num_layers: int = 1, | |
| resnet_eps: float = 1e-6, | |
| resnet_time_scale_shift: str = "default", | |
| resnet_act_fn: str = "swish", | |
| resnet_pre_norm: bool = True, | |
| output_scale_factor=np.sqrt(2.0), | |
| add_upsample=True, | |
| upsample_padding=1, | |
| ): | |
| super().__init__() | |
| self.resnets = nn.ModuleList([]) | |
| for i in range(num_layers): | |
| res_skip_channels = in_channels if (i == num_layers - 1) else out_channels | |
| resnet_in_channels = prev_output_channel if i == 0 else out_channels | |
| self.resnets.append( | |
| ResnetBlock2D( | |
| in_channels=resnet_in_channels + res_skip_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min((resnet_in_channels + res_skip_channels) // 4, 32), | |
| groups_out=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| ) | |
| ) | |
| self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels) | |
| if add_upsample: | |
| self.resnet_up = ResnetBlock2D( | |
| in_channels=out_channels, | |
| out_channels=out_channels, | |
| temb_channels=temb_channels, | |
| eps=resnet_eps, | |
| groups=min(out_channels // 4, 32), | |
| groups_out=min(out_channels // 4, 32), | |
| dropout=dropout, | |
| time_embedding_norm=resnet_time_scale_shift, | |
| non_linearity=resnet_act_fn, | |
| output_scale_factor=output_scale_factor, | |
| pre_norm=resnet_pre_norm, | |
| use_nin_shortcut=True, | |
| up=True, | |
| kernel="fir", | |
| ) | |
| self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) | |
| self.skip_norm = torch.nn.GroupNorm( | |
| num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True | |
| ) | |
| self.act = nn.SiLU() | |
| else: | |
| self.resnet_up = None | |
| self.skip_conv = None | |
| self.skip_norm = None | |
| self.act = None | |
| def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=None): | |
| for resnet in self.resnets: | |
| # pop res hidden states | |
| res_hidden_states = res_hidden_states_tuple[-1] | |
| res_hidden_states_tuple = res_hidden_states_tuple[:-1] | |
| hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) | |
| hidden_states = resnet(hidden_states, temb) | |
| if skip_sample is not None: | |
| skip_sample = self.upsampler(skip_sample) | |
| else: | |
| skip_sample = 0 | |
| if self.resnet_up is not None: | |
| skip_sample_states = self.skip_norm(hidden_states) | |
| skip_sample_states = self.act(skip_sample_states) | |
| skip_sample_states = self.skip_conv(skip_sample_states) | |
| skip_sample = skip_sample + skip_sample_states | |
| hidden_states = self.resnet_up(hidden_states, temb) | |
| return hidden_states, skip_sample | |