Anonymous-123/ImageNet-Editing
1
1from abc import abstractmethod2 3import math4 5import numpy as np6import torch as th7import torch.nn as nn8import torch.nn.functional as F9 10from .fp16_util import convert_module_to_f16, convert_module_to_f3211from .nn import (12 checkpoint,13 conv_nd,14 linear,15 avg_pool_nd,16 zero_module,17 normalization,18 timestep_embedding,19)20 21 22class AttentionPool2d(nn.Module):23 """24 Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py25 """26 27 def __init__(28 self,29 spacial_dim: int,30 embed_dim: int,31 num_heads_channels: int,32 output_dim: int = None,33 ):34 super().__init__()35 self.positional_embedding = nn.Parameter(36 th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.537 )38 self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)39 self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)40 self.num_heads = embed_dim // num_heads_channels41 self.attention = QKVAttention(self.num_heads)42 43 def forward(self, x):44 b, c, *_spatial = x.shape45 x = x.reshape(b, c, -1) # NC(HW)46 x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)47 x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)48 x = self.qkv_proj(x)49 x = self.attention(x)50 x = self.c_proj(x)51 return x[:, :, 0]52 53 54class TimestepBlock(nn.Module):55 """56 Any module where forward() takes timestep embeddings as a second argument.57 """58 59 @abstractmethod60 def forward(self, x, emb):61 """62 Apply the module to `x` given `emb` timestep embeddings.63 """64 65 66class TimestepEmbedSequential(nn.Sequential, TimestepBlock):67 """68 A sequential module that passes timestep embeddings to the children that69 support it as an extra input.70 """71 72 def forward(self, x, emb):73 for layer in self:74 if isinstance(layer, TimestepBlock):75 x = layer(x, emb)76 else:77 x = layer(x)78 return x79 80 81class Upsample(nn.Module):82 """83 An upsampling layer with an optional convolution.84 85 :param channels: channels in the inputs and outputs.86 :param use_conv: a bool determining if a convolution is applied.87 :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then88 upsampling occurs in the inner-two dimensions.89 """90 91 def __init__(self, channels, use_conv, dims=2, out_channels=None):92 super().__init__()93 self.channels = channels94 self.out_channels = out_channels or channels95 self.use_conv = use_conv96 self.dims = dims97 if use_conv:98 self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)99 100 def forward(self, x):101 assert x.shape[1] == self.channels102 if self.dims == 3:103 x = F.interpolate(104 x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"105 )106 else:107 x = F.interpolate(x, scale_factor=2, mode="nearest")108 if self.use_conv:109 x = self.conv(x)110 return x111 112 113class Downsample(nn.Module):114 """115 A downsampling layer with an optional convolution.116 117 :param channels: channels in the inputs and outputs.118 :param use_conv: a bool determining if a convolution is applied.119 :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then120 downsampling occurs in the inner-two dimensions.121 """122 123 def __init__(self, channels, use_conv, dims=2, out_channels=None):124 super().__init__()125 self.channels = channels126 self.out_channels = out_channels or channels127 self.use_conv = use_conv128 self.dims = dims129 stride = 2 if dims != 3 else (1, 2, 2)130 if use_conv:131 self.op = conv_nd(132 dims, self.channels, self.out_channels, 3, stride=stride, padding=1133 )134 else:135 assert self.channels == self.out_channels136 self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)137 138 def forward(self, x):139 assert x.shape[1] == self.channels140 return self.op(x)141 142 143class ResBlock(TimestepBlock):144 """145 A residual block that can optionally change the number of channels.146 147 :param channels: the number of input channels.148 :param emb_channels: the number of timestep embedding channels.149 :param dropout: the rate of dropout.150 :param out_channels: if specified, the number of out channels.151 :param use_conv: if True and out_channels is specified, use a spatial152 convolution instead of a smaller 1x1 convolution to change the153 channels in the skip connection.154 :param dims: determines if the signal is 1D, 2D, or 3D.155 :param use_checkpoint: if True, use gradient checkpointing on this module.156 :param up: if True, use this block for upsampling.157 :param down: if True, use this block for downsampling.158 """159 160 def __init__(161 self,162 channels,163 emb_channels,164 dropout,165 out_channels=None,166 use_conv=False,167 use_scale_shift_norm=False,168 dims=2,169 use_checkpoint=False,170 up=False,171 down=False,172 ):173 super().__init__()174 self.channels = channels175 self.emb_channels = emb_channels176 self.dropout = dropout177 self.out_channels = out_channels or channels178 self.use_conv = use_conv179 self.use_checkpoint = use_checkpoint180 self.use_scale_shift_norm = use_scale_shift_norm181 182 self.in_layers = nn.Sequential(183 normalization(channels),184 nn.SiLU(),185 conv_nd(dims, channels, self.out_channels, 3, padding=1),186 )187 188 self.updown = up or down189 190 if up:191 self.h_upd = Upsample(channels, False, dims)192 self.x_upd = Upsample(channels, False, dims)193 elif down:194 self.h_upd = Downsample(channels, False, dims)195 self.x_upd = Downsample(channels, False, dims)196 else:197 self.h_upd = self.x_upd = nn.Identity()198 199 self.emb_layers = nn.Sequential(200 nn.SiLU(),201 linear(202 emb_channels,203 2 * self.out_channels if use_scale_shift_norm else self.out_channels,204 ),205 )206 self.out_layers = nn.Sequential(207 normalization(self.out_channels),208 nn.SiLU(),209 nn.Dropout(p=dropout),210 zero_module(211 conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)212 ),213 )214 215 if self.out_channels == channels:216 self.skip_connection = nn.Identity()217 elif use_conv:218 self.skip_connection = conv_nd(219 dims, channels, self.out_channels, 3, padding=1220 )221 else:222 self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)223 224 def forward(self, x, emb):225 """226 Apply the block to a Tensor, conditioned on a timestep embedding.227 228 :param x: an [N x C x ...] Tensor of features.229 :param emb: an [N x emb_channels] Tensor of timestep embeddings.230 :return: an [N x C x ...] Tensor of outputs.231 """232 return checkpoint(233 self._forward, (x, emb), self.parameters(), self.use_checkpoint234 )235 236 def _forward(self, x, emb):237 if self.updown:238 in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]239 h = in_rest(x)240 h = self.h_upd(h)241 x = self.x_upd(x)242 h = in_conv(h)243 else:244 h = self.in_layers(x)245 emb_out = self.emb_layers(emb).type(h.dtype)246 while len(emb_out.shape) < len(h.shape):247 emb_out = emb_out[..., None]248 if self.use_scale_shift_norm:249 out_norm, out_rest = self.out_layers[0], self.out_layers[1:]250 scale, shift = th.chunk(emb_out, 2, dim=1)251 h = out_norm(h) * (1 + scale) + shift252 h = out_rest(h)253 else:254 h = h + emb_out255 h = self.out_layers(h)256 return self.skip_connection(x) + h257 258 259class AttentionBlock(nn.Module):260 """261 An attention block that allows spatial positions to attend to each other.262 263 Originally ported from here, but adapted to the N-d case.264 https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.265 """266 267 def __init__(268 self,269 channels,270 num_heads=1,271 num_head_channels=-1,272 use_checkpoint=False,273 use_new_attention_order=False,274 ):275 super().__init__()276 self.channels = channels277 if num_head_channels == -1:278 self.num_heads = num_heads279 else:280 assert (281 channels % num_head_channels == 0282 ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"283 self.num_heads = channels // num_head_channels284 self.use_checkpoint = use_checkpoint285 self.norm = normalization(channels)286 self.qkv = conv_nd(1, channels, channels * 3, 1)287 if use_new_attention_order:288 # split qkv before split heads289 self.attention = QKVAttention(self.num_heads)290 else:291 # split heads before split qkv292 self.attention = QKVAttentionLegacy(self.num_heads)293 294 self.proj_out = zero_module(conv_nd(1, channels, channels, 1))295 296 def forward(self, x):297 return checkpoint(self._forward, (x,), self.parameters(), True)298 299 def _forward(self, x):300 b, c, *spatial = x.shape301 x = x.reshape(b, c, -1)302 qkv = self.qkv(self.norm(x))303 h = self.attention(qkv)304 h = self.proj_out(h)305 return (x + h).reshape(b, c, *spatial)306 307 308def count_flops_attn(model, _x, y):309 """310 A counter for the `thop` package to count the operations in an311 attention operation.312 Meant to be used like:313 macs, params = thop.profile(314 model,315 inputs=(inputs, timestamps),316 custom_ops={QKVAttention: QKVAttention.count_flops},317 )318 """319 b, c, *spatial = y[0].shape320 num_spatial = int(np.prod(spatial))321 # We perform two matmuls with the same number of ops.322 # The first computes the weight matrix, the second computes323 # the combination of the value vectors.324 matmul_ops = 2 * b * (num_spatial ** 2) * c325 model.total_ops += th.DoubleTensor([matmul_ops])326 327 328class QKVAttentionLegacy(nn.Module):329 """330 A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping331 """332 333 def __init__(self, n_heads):334 super().__init__()335 self.n_heads = n_heads336 337 def forward(self, qkv):338 """339 Apply QKV attention.340 341 :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.342 :return: an [N x (H * C) x T] tensor after attention.343 """344 bs, width, length = qkv.shape345 assert width % (3 * self.n_heads) == 0346 ch = width // (3 * self.n_heads)347 q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)348 scale = 1 / math.sqrt(math.sqrt(ch))349 weight = th.einsum(350 "bct,bcs->bts", q * scale, k * scale351 ) # More stable with f16 than dividing afterwards352 weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)353 a = th.einsum("bts,bcs->bct", weight, v)354 return a.reshape(bs, -1, length)355 356 @staticmethod357 def count_flops(model, _x, y):358 return count_flops_attn(model, _x, y)359 360 361class QKVAttention(nn.Module):362 """363 A module which performs QKV attention and splits in a different order.364 """365 366 def __init__(self, n_heads):367 super().__init__()368 self.n_heads = n_heads369 370 def forward(self, qkv):371 """372 Apply QKV attention.373 374 :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.375 :return: an [N x (H * C) x T] tensor after attention.376 """377 bs, width, length = qkv.shape378 assert width % (3 * self.n_heads) == 0379 ch = width // (3 * self.n_heads)380 q, k, v = qkv.chunk(3, dim=1)381 scale = 1 / math.sqrt(math.sqrt(ch))382 weight = th.einsum(383 "bct,bcs->bts",384 (q * scale).view(bs * self.n_heads, ch, length),385 (k * scale).view(bs * self.n_heads, ch, length),386 ) # More stable with f16 than dividing afterwards387 weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)388 a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))389 return a.reshape(bs, -1, length)390 391 @staticmethod392 def count_flops(model, _x, y):393 return count_flops_attn(model, _x, y)394 395 396class UNetModel(nn.Module):397 """398 The full UNet model with attention and timestep embedding.399 400 :param in_channels: channels in the input Tensor.401 :param model_channels: base channel count for the model.402 :param out_channels: channels in the output Tensor.403 :param num_res_blocks: number of residual blocks per downsample.404 :param attention_resolutions: a collection of downsample rates at which405 attention will take place. May be a set, list, or tuple.406 For example, if this contains 4, then at 4x downsampling, attention407 will be used.408 :param dropout: the dropout probability.409 :param channel_mult: channel multiplier for each level of the UNet.410 :param conv_resample: if True, use learned convolutions for upsampling and411 downsampling.412 :param dims: determines if the signal is 1D, 2D, or 3D.413 :param num_classes: if specified (as an int), then this model will be414 class-conditional with `num_classes` classes.415 :param use_checkpoint: use gradient checkpointing to reduce memory usage.416 :param num_heads: the number of attention heads in each attention layer.417 :param num_heads_channels: if specified, ignore num_heads and instead use418 a fixed channel width per attention head.419 :param num_heads_upsample: works with num_heads to set a different number420 of heads for upsampling. Deprecated.421 :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.422 :param resblock_updown: use residual blocks for up/downsampling.423 :param use_new_attention_order: use a different attention pattern for potentially424 increased efficiency.425 """426 427 def __init__(428 self,429 image_size,430 in_channels,431 model_channels,432 out_channels,433 num_res_blocks,434 attention_resolutions,435 dropout=0,436 channel_mult=(1, 2, 4, 8),437 conv_resample=True,438 dims=2,439 num_classes=None,440 use_checkpoint=False,441 use_fp16=False,442 num_heads=1,443 num_head_channels=-1,444 num_heads_upsample=-1,445 use_scale_shift_norm=False,446 resblock_updown=False,447 use_new_attention_order=False,448 ):449 super().__init__()450 451 if num_heads_upsample == -1:452 num_heads_upsample = num_heads453 454 self.image_size = image_size455 self.in_channels = in_channels456 self.model_channels = model_channels457 self.out_channels = out_channels458 self.num_res_blocks = num_res_blocks459 self.attention_resolutions = attention_resolutions460 self.dropout = dropout461 self.channel_mult = channel_mult462 self.conv_resample = conv_resample463 self.num_classes = num_classes464 self.use_checkpoint = use_checkpoint465 self.dtype = th.float16 if use_fp16 else th.float32466 self.num_heads = num_heads467 self.num_head_channels = num_head_channels468 self.num_heads_upsample = num_heads_upsample469 470 time_embed_dim = model_channels * 4471 self.time_embed = nn.Sequential(472 linear(model_channels, time_embed_dim),473 nn.SiLU(),474 linear(time_embed_dim, time_embed_dim),475 )476 477 if self.num_classes is not None:478 self.label_emb = nn.Embedding(num_classes, time_embed_dim)479 480 ch = input_ch = int(channel_mult[0] * model_channels)481 self.input_blocks = nn.ModuleList(482 [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]483 )484 self._feature_size = ch485 input_block_chans = [ch]486 ds = 1487 for level, mult in enumerate(channel_mult):488 for _ in range(num_res_blocks):489 layers = [490 ResBlock(491 ch,492 time_embed_dim,493 dropout,494 out_channels=int(mult * model_channels),495 dims=dims,496 use_checkpoint=use_checkpoint,497 use_scale_shift_norm=use_scale_shift_norm,498 )499 ]500 ch = int(mult * model_channels)501 if ds in attention_resolutions:502 layers.append(503 AttentionBlock(504 ch,505 use_checkpoint=use_checkpoint,506 num_heads=num_heads,507 num_head_channels=num_head_channels,508 use_new_attention_order=use_new_attention_order,509 )510 )511 self.input_blocks.append(TimestepEmbedSequential(*layers))512 self._feature_size += ch513 input_block_chans.append(ch)514 if level != len(channel_mult) - 1:515 out_ch = ch516 self.input_blocks.append(517 TimestepEmbedSequential(518 ResBlock(519 ch,520 time_embed_dim,521 dropout,522 out_channels=out_ch,523 dims=dims,524 use_checkpoint=use_checkpoint,525 use_scale_shift_norm=use_scale_shift_norm,526 down=True,527 )528 if resblock_updown529 else Downsample(530 ch, conv_resample, dims=dims, out_channels=out_ch531 )532 )533 )534 ch = out_ch535 input_block_chans.append(ch)536 ds *= 2537 self._feature_size += ch538 539 self.middle_block = TimestepEmbedSequential(540 ResBlock(541 ch,542 time_embed_dim,543 dropout,544 dims=dims,545 use_checkpoint=use_checkpoint,546 use_scale_shift_norm=use_scale_shift_norm,547 ),548 AttentionBlock(549 ch,550 use_checkpoint=use_checkpoint,551 num_heads=num_heads,552 num_head_channels=num_head_channels,553 use_new_attention_order=use_new_attention_order,554 ),555 ResBlock(556 ch,557 time_embed_dim,558 dropout,559 dims=dims,560 use_checkpoint=use_checkpoint,561 use_scale_shift_norm=use_scale_shift_norm,562 ),563 )564 self._feature_size += ch565 566 self.output_blocks = nn.ModuleList([])567 for level, mult in list(enumerate(channel_mult))[::-1]:568 for i in range(num_res_blocks + 1):569 ich = input_block_chans.pop()570 layers = [571 ResBlock(572 ch + ich,573 time_embed_dim,574 dropout,575 out_channels=int(model_channels * mult),576 dims=dims,577 use_checkpoint=use_checkpoint,578 use_scale_shift_norm=use_scale_shift_norm,579 )580 ]581 ch = int(model_channels * mult)582 if ds in attention_resolutions:583 layers.append(584 AttentionBlock(585 ch,586 use_checkpoint=use_checkpoint,587 num_heads=num_heads_upsample,588 num_head_channels=num_head_channels,589 use_new_attention_order=use_new_attention_order,590 )591 )592 if level and i == num_res_blocks:593 out_ch = ch594 layers.append(595 ResBlock(596 ch,597 time_embed_dim,598 dropout,599 out_channels=out_ch,600 dims=dims,601 use_checkpoint=use_checkpoint,602 use_scale_shift_norm=use_scale_shift_norm,603 up=True,604 )605 if resblock_updown606 else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)607 )608 ds //= 2609 self.output_blocks.append(TimestepEmbedSequential(*layers))610 self._feature_size += ch611 612 self.out = nn.Sequential(613 normalization(ch),614 nn.SiLU(),615 zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)),616 )617 618 def convert_to_fp16(self):619 """620 Convert the torso of the model to float16.621 """622 self.input_blocks.apply(convert_module_to_f16)623 self.middle_block.apply(convert_module_to_f16)624 self.output_blocks.apply(convert_module_to_f16)625 626 def convert_to_fp32(self):627 """628 Convert the torso of the model to float32.629 """630 self.input_blocks.apply(convert_module_to_f32)631 self.middle_block.apply(convert_module_to_f32)632 self.output_blocks.apply(convert_module_to_f32)633 634 def forward(self, x, timesteps, y=None):635 """636 Apply the model to an input batch.637 638 :param x: an [N x C x ...] Tensor of inputs.639 :param timesteps: a 1-D batch of timesteps.640 :param y: an [N] Tensor of labels, if class-conditional.641 :return: an [N x C x ...] Tensor of outputs.642 """643 assert (y is not None) == (644 self.num_classes is not None645 ), "must specify y if and only if the model is class-conditional"646 647 hs = []648 emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))649 650 if self.num_classes is not None:651 assert y.shape == (x.shape[0],)652 emb = emb + self.label_emb(y)653 654 h = x.type(self.dtype)655 for module in self.input_blocks:656 h = module(h, emb)657 hs.append(h)658 h = self.middle_block(h, emb)659 for module in self.output_blocks:660 h = th.cat([h, hs.pop()], dim=1)661 h = module(h, emb)662 h = h.type(x.dtype)663 return self.out(h)664 665 666class SuperResModel(UNetModel):667 """668 A UNetModel that performs super-resolution.669 670 Expects an extra kwarg `low_res` to condition on a low-resolution image.671 """672 673 def __init__(self, image_size, in_channels, *args, **kwargs):674 super().__init__(image_size, in_channels * 2, *args, **kwargs)675 676 def forward(self, x, timesteps, low_res=None, **kwargs):677 _, _, new_height, new_width = x.shape678 upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")679 x = th.cat([x, upsampled], dim=1)680 return super().forward(x, timesteps, **kwargs)681 682 683class EncoderUNetModel(nn.Module):684 """685 The half UNet model with attention and timestep embedding.686 687 For usage, see UNet.688 """689 690 def __init__(691 self,692 image_size,693 in_channels,694 model_channels,695 out_channels,696 num_res_blocks,697 attention_resolutions,698 dropout=0,699 channel_mult=(1, 2, 4, 8),700 conv_resample=True,701 dims=2,702 use_checkpoint=False,703 use_fp16=False,704 num_heads=1,705 num_head_channels=-1,706 num_heads_upsample=-1,707 use_scale_shift_norm=False,708 resblock_updown=False,709 use_new_attention_order=False,710 pool="adaptive",711 ):712 super().__init__()713 714 if num_heads_upsample == -1:715 num_heads_upsample = num_heads716 717 self.in_channels = in_channels718 self.model_channels = model_channels719 self.out_channels = out_channels720 self.num_res_blocks = num_res_blocks721 self.attention_resolutions = attention_resolutions722 self.dropout = dropout723 self.channel_mult = channel_mult724 self.conv_resample = conv_resample725 self.use_checkpoint = use_checkpoint726 self.dtype = th.float16 if use_fp16 else th.float32727 self.num_heads = num_heads728 self.num_head_channels = num_head_channels729 self.num_heads_upsample = num_heads_upsample730 731 time_embed_dim = model_channels * 4732 self.time_embed = nn.Sequential(733 linear(model_channels, time_embed_dim),734 nn.SiLU(),735 linear(time_embed_dim, time_embed_dim),736 )737 738 ch = int(channel_mult[0] * model_channels)739 self.input_blocks = nn.ModuleList(740 [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))]741 )742 self._feature_size = ch743 input_block_chans = [ch]744 ds = 1745 for level, mult in enumerate(channel_mult):746 for _ in range(num_res_blocks):747 layers = [748 ResBlock(749 ch,750 time_embed_dim,751 dropout,752 out_channels=int(mult * model_channels),753 dims=dims,754 use_checkpoint=use_checkpoint,755 use_scale_shift_norm=use_scale_shift_norm,756 )757 ]758 ch = int(mult * model_channels)759 if ds in attention_resolutions:760 layers.append(761 AttentionBlock(762 ch,763 use_checkpoint=use_checkpoint,764 num_heads=num_heads,765 num_head_channels=num_head_channels,766 use_new_attention_order=use_new_attention_order,767 )768 )769 self.input_blocks.append(TimestepEmbedSequential(*layers))770 self._feature_size += ch771 input_block_chans.append(ch)772 if level != len(channel_mult) - 1:773 out_ch = ch774 self.input_blocks.append(775 TimestepEmbedSequential(776 ResBlock(777 ch,778 time_embed_dim,779 dropout,780 out_channels=out_ch,781 dims=dims,782 use_checkpoint=use_checkpoint,783 use_scale_shift_norm=use_scale_shift_norm,784 down=True,785 )786 if resblock_updown787 else Downsample(788 ch, conv_resample, dims=dims, out_channels=out_ch789 )790 )791 )792 ch = out_ch793 input_block_chans.append(ch)794 ds *= 2795 self._feature_size += ch796 797 self.middle_block = TimestepEmbedSequential(798 ResBlock(799 ch,800 time_embed_dim,801 dropout,802 dims=dims,803 use_checkpoint=use_checkpoint,804 use_scale_shift_norm=use_scale_shift_norm,805 ),806 AttentionBlock(807 ch,808 use_checkpoint=use_checkpoint,809 num_heads=num_heads,810 num_head_channels=num_head_channels,811 use_new_attention_order=use_new_attention_order,812 ),813 ResBlock(814 ch,815 time_embed_dim,816 dropout,817 dims=dims,818 use_checkpoint=use_checkpoint,819 use_scale_shift_norm=use_scale_shift_norm,820 ),821 )822 self._feature_size += ch823 self.pool = pool824 if pool == "adaptive":825 self.out = nn.Sequential(826 normalization(ch),827 nn.SiLU(),828 nn.AdaptiveAvgPool2d((1, 1)),829 zero_module(conv_nd(dims, ch, out_channels, 1)),830 nn.Flatten(),831 )832 elif pool == "attention":833 assert num_head_channels != -1834 self.out = nn.Sequential(835 normalization(ch),836 nn.SiLU(),837 AttentionPool2d(838 (image_size // ds), ch, num_head_channels, out_channels839 ),840 )841 elif pool == "spatial":842 self.out = nn.Sequential(843 nn.Linear(self._feature_size, 2048),844 nn.ReLU(),845 nn.Linear(2048, self.out_channels),846 )847 elif pool == "spatial_v2":848 self.out = nn.Sequential(849 nn.Linear(self._feature_size, 2048),850 normalization(2048),851 nn.SiLU(),852 nn.Linear(2048, self.out_channels),853 )854 else:855 raise NotImplementedError(f"Unexpected {pool} pooling")856 857 def convert_to_fp16(self):858 """859 Convert the torso of the model to float16.860 """861 self.input_blocks.apply(convert_module_to_f16)862 self.middle_block.apply(convert_module_to_f16)863 864 def convert_to_fp32(self):865 """866 Convert the torso of the model to float32.867 """868 self.input_blocks.apply(convert_module_to_f32)869 self.middle_block.apply(convert_module_to_f32)870 871 def forward(self, x, timesteps):872 """873 Apply the model to an input batch.874 875 :param x: an [N x C x ...] Tensor of inputs.876 :param timesteps: a 1-D batch of timesteps.877 :return: an [N x K] Tensor of outputs.878 """879 emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))880 881 results = []882 h = x.type(self.dtype)883 for module in self.input_blocks:884 h = module(h, emb)885 if self.pool.startswith("spatial"):886 results.append(h.type(x.dtype).mean(dim=(2, 3)))887 h = self.middle_block(h, emb)888 if self.pool.startswith("spatial"):889 results.append(h.type(x.dtype).mean(dim=(2, 3)))890 h = th.cat(results, axis=-1)891 return self.out(h)892 else:893 h = h.type(x.dtype)894 return self.out(h)895