Spaces:
Sleeping
Sleeping
| from abc import abstractmethod | |
| from functools import partial | |
| import math | |
| from typing import Iterable | |
| import numpy as np | |
| import torch as th | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from ldm.modules.diffusionmodules.util import ( | |
| checkpoint, | |
| conv_nd, | |
| linear, | |
| avg_pool_nd, | |
| zero_module, | |
| normalization, | |
| timestep_embedding, | |
| ) | |
| from ldm.modules.attention import SpatialTransformer | |
| from ldm.modules.diffusionmodules.openaimodel import convert_module_to_f16, convert_module_to_f32, AttentionPool2d, \ | |
| TimestepBlock, TimestepEmbedSequential, Upsample, TransposedUpsample, Downsample, ResBlock, AttentionBlock, count_flops_attn, \ | |
| QKVAttentionLegacy, QKVAttention | |
| class UNetModel(nn.Module): | |
| """ | |
| The full UNet model with attention and timestep embedding. | |
| :param in_channels: channels in the input Tensor. | |
| :param model_channels: base channel count for the model. | |
| :param out_channels: channels in the output Tensor. | |
| :param num_res_blocks: number of residual blocks per downsample. | |
| :param attention_resolutions: a collection of downsample rates at which | |
| attention will take place. May be a set, list, or tuple. | |
| For example, if this contains 4, then at 4x downsampling, attention | |
| will be used. | |
| :param dropout: the dropout probability. | |
| :param channel_mult: channel multiplier for each level of the UNet. | |
| :param conv_resample: if True, use learned convolutions for upsampling and | |
| downsampling. | |
| :param dims: determines if the signal is 1D, 2D, or 3D. | |
| :param num_classes: if specified (as an int), then this model will be | |
| class-conditional with `num_classes` classes. | |
| :param use_checkpoint: use gradient checkpointing to reduce memory usage. | |
| :param num_heads: the number of attention heads in each attention layer. | |
| :param num_heads_channels: if specified, ignore num_heads and instead use | |
| a fixed channel width per attention head. | |
| :param num_heads_upsample: works with num_heads to set a different number | |
| of heads for upsampling. Deprecated. | |
| :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. | |
| :param resblock_updown: use residual blocks for up/downsampling. | |
| :param use_new_attention_order: use a different attention pattern for potentially | |
| increased efficiency. | |
| """ | |
| def __init__( | |
| self, | |
| image_size, | |
| in_channels, | |
| model_channels, | |
| out_channels, | |
| num_res_blocks, | |
| attention_resolutions, | |
| dropout=0, | |
| channel_mult=(1, 2, 4, 8), | |
| conv_resample=True, | |
| dims=2, | |
| num_classes=None, | |
| use_checkpoint=False, | |
| use_fp16=False, | |
| num_heads=-1, | |
| num_head_channels=-1, | |
| num_heads_upsample=-1, | |
| use_scale_shift_norm=False, | |
| resblock_updown=False, | |
| use_new_attention_order=False, | |
| use_spatial_transformer=False, # custom transformer support | |
| transformer_depth=1, # custom transformer support | |
| context_dim=None, # custom transformer support | |
| n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model | |
| legacy=True, | |
| use_context_project=False, # custom text to audio support | |
| use_context_attn=True # custom text to audio support | |
| ): | |
| super().__init__() | |
| if use_spatial_transformer: | |
| assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' | |
| if context_dim is not None and not use_context_project: | |
| assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' | |
| from omegaconf.listconfig import ListConfig | |
| if type(context_dim) == ListConfig: | |
| context_dim = list(context_dim) | |
| if num_heads_upsample == -1: | |
| num_heads_upsample = num_heads | |
| if num_heads == -1: | |
| assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' | |
| if num_head_channels == -1: | |
| assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' | |
| self.image_size = image_size | |
| self.in_channels = in_channels | |
| self.model_channels = model_channels | |
| self.out_channels = out_channels | |
| self.num_res_blocks = num_res_blocks | |
| self.attention_resolutions = attention_resolutions | |
| self.dropout = dropout | |
| self.channel_mult = channel_mult | |
| self.conv_resample = conv_resample | |
| self.num_classes = num_classes | |
| self.use_checkpoint = use_checkpoint | |
| self.dtype = th.float16 if use_fp16 else th.float32 | |
| self.num_heads = num_heads | |
| self.num_head_channels = num_head_channels | |
| self.num_heads_upsample = num_heads_upsample | |
| self.predict_codebook_ids = n_embed is not None | |
| time_embed_dim = model_channels * 4 | |
| self.time_embed = nn.Sequential( | |
| linear(model_channels, time_embed_dim), | |
| nn.SiLU(), | |
| linear(time_embed_dim, time_embed_dim), | |
| ) | |
| if self.num_classes is not None: | |
| self.label_emb = nn.Embedding(num_classes, time_embed_dim) | |
| self.input_blocks = nn.ModuleList( | |
| [ | |
| TimestepEmbedSequential( | |
| conv_nd(dims, in_channels, model_channels, 3, padding=1) | |
| ) | |
| ] | |
| ) | |
| self._feature_size = model_channels | |
| input_block_chans = [model_channels] | |
| ch = model_channels | |
| ds = 1 | |
| for level, mult in enumerate(channel_mult): | |
| for _ in range(num_res_blocks): | |
| layers = [ | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=mult * model_channels, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| ) | |
| ] | |
| ch = mult * model_channels | |
| if ds in attention_resolutions: | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| layers.append( | |
| AttentionBlock( | |
| ch, | |
| use_checkpoint=use_checkpoint, | |
| num_heads=num_heads, | |
| num_head_channels=dim_head, | |
| use_new_attention_order=use_new_attention_order, | |
| ) if not use_spatial_transformer else SpatialTransformer( | |
| ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim | |
| ) | |
| ) | |
| self.input_blocks.append(TimestepEmbedSequential(*layers)) | |
| self._feature_size += ch | |
| input_block_chans.append(ch) | |
| if level != len(channel_mult) - 1: | |
| out_ch = ch | |
| self.input_blocks.append( | |
| TimestepEmbedSequential( | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=out_ch, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| down=True, | |
| ) | |
| if resblock_updown | |
| else Downsample( | |
| ch, conv_resample, dims=dims, out_channels=out_ch | |
| ) | |
| ) | |
| ) | |
| ch = out_ch | |
| input_block_chans.append(ch) | |
| ds *= 2 | |
| self._feature_size += ch | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| self.middle_block = TimestepEmbedSequential( | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| ), | |
| AttentionBlock( | |
| ch, | |
| use_checkpoint=use_checkpoint, | |
| num_heads=num_heads, | |
| num_head_channels=dim_head, | |
| use_new_attention_order=use_new_attention_order, | |
| ) if not use_spatial_transformer else SpatialTransformer( | |
| ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim | |
| ), | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| ), | |
| ) | |
| self._feature_size += ch | |
| self.output_blocks = nn.ModuleList([]) | |
| for level, mult in list(enumerate(channel_mult))[::-1]: | |
| for i in range(num_res_blocks + 1): | |
| ich = input_block_chans.pop() | |
| layers = [ | |
| ResBlock( | |
| ch + ich, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=model_channels * mult, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| ) | |
| ] | |
| ch = model_channels * mult | |
| if ds in attention_resolutions: | |
| if num_head_channels == -1: | |
| dim_head = ch // num_heads | |
| else: | |
| num_heads = ch // num_head_channels | |
| dim_head = num_head_channels | |
| if legacy: | |
| #num_heads = 1 | |
| dim_head = ch // num_heads if use_spatial_transformer else num_head_channels | |
| layers.append( | |
| AttentionBlock( | |
| ch, | |
| use_checkpoint=use_checkpoint, | |
| num_heads=num_heads_upsample, | |
| num_head_channels=dim_head, | |
| use_new_attention_order=use_new_attention_order, | |
| ) if not use_spatial_transformer else SpatialTransformer( | |
| ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim | |
| ) | |
| ) | |
| if level and i == num_res_blocks: | |
| out_ch = ch | |
| layers.append( | |
| ResBlock( | |
| ch, | |
| time_embed_dim, | |
| dropout, | |
| out_channels=out_ch, | |
| dims=dims, | |
| use_checkpoint=use_checkpoint, | |
| use_scale_shift_norm=use_scale_shift_norm, | |
| up=True, | |
| ) | |
| if resblock_updown | |
| else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) | |
| ) | |
| ds //= 2 | |
| self.output_blocks.append(TimestepEmbedSequential(*layers)) | |
| self._feature_size += ch | |
| self.out = nn.Sequential( | |
| normalization(ch), | |
| nn.SiLU(), | |
| zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)), | |
| ) | |
| if self.predict_codebook_ids: | |
| self.id_predictor = nn.Sequential( | |
| normalization(ch), | |
| conv_nd(dims, model_channels, n_embed, 1), | |
| #nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits | |
| ) | |
| self.use_context_project = use_context_project | |
| if use_context_project: | |
| self.context_project = linear(context_dim, time_embed_dim) | |
| self.use_context_attn = use_context_attn | |
| def convert_to_fp16(self): | |
| """ | |
| Convert the torso of the model to float16. | |
| """ | |
| self.input_blocks.apply(convert_module_to_f16) | |
| self.middle_block.apply(convert_module_to_f16) | |
| self.output_blocks.apply(convert_module_to_f16) | |
| def convert_to_fp32(self): | |
| """ | |
| Convert the torso of the model to float32. | |
| """ | |
| self.input_blocks.apply(convert_module_to_f32) | |
| self.middle_block.apply(convert_module_to_f32) | |
| self.output_blocks.apply(convert_module_to_f32) | |
| def forward(self, x, timesteps=None, context=None, y=None,**kwargs): | |
| """ | |
| Apply the model to an input batch. | |
| :param x: an [N x C x ...] Tensor of inputs. | |
| :param timesteps: a 1-D batch of timesteps. | |
| :param context: conditioning plugged in via crossattn | |
| :param y: an [N] Tensor of labels, if class-conditional. | |
| :return: an [N x C x ...] Tensor of outputs. | |
| """ | |
| assert (y is not None) == ( | |
| self.num_classes is not None | |
| ), "must specify y if and only if the model is class-conditional" | |
| hs = [] | |
| t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) | |
| emb = self.time_embed(t_emb) | |
| if self.num_classes is not None: | |
| assert y.shape == (x.shape[0],) | |
| emb = emb + self.label_emb(y) | |
| # For text-to-audio using global CLIP | |
| if self.use_context_project: | |
| context = self.context_project(context) | |
| emb = emb + context.squeeze(1) | |
| h = x.type(self.dtype) | |
| for module in self.input_blocks: | |
| h = module(h, emb, context if self.use_context_attn else None) | |
| hs.append(h) | |
| h = self.middle_block(h, emb, context if self.use_context_attn else None) | |
| for module in self.output_blocks: | |
| h = th.cat([h, hs.pop()], dim=1) | |
| h = module(h, emb, context if self.use_context_attn else None) | |
| h = h.type(x.dtype) | |
| if self.predict_codebook_ids: | |
| return self.id_predictor(h) | |
| else: | |
| return self.out(h) | |