Dynamatrix/DiffBIR-OpenXLab
0
1from typing import Mapping, Any2import copy3from collections import OrderedDict4 5import einops6import torch7import torch as th8import torch.nn as nn9 10from ldm.modules.diffusionmodules.util import (11 conv_nd,12 linear,13 zero_module,14 timestep_embedding,15)16from ldm.modules.attention import SpatialTransformer17from ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock, UNetModel18from ldm.models.diffusion.ddpm import LatentDiffusion19from ldm.util import log_txt_as_img, exists, instantiate_from_config20from ldm.modules.distributions.distributions import DiagonalGaussianDistribution21from utils.common import frozen_module22from .spaced_sampler import SpacedSampler23 24 25class ControlledUnetModel(UNetModel):26 def forward(self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs):27 hs = []28 with torch.no_grad():29 t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)30 emb = self.time_embed(t_emb)31 h = x.type(self.dtype)32 for module in self.input_blocks:33 h = module(h, emb, context)34 hs.append(h)35 h = self.middle_block(h, emb, context)36 37 if control is not None:38 h += control.pop()39 40 for i, module in enumerate(self.output_blocks):41 if only_mid_control or control is None:42 h = torch.cat([h, hs.pop()], dim=1)43 else:44 h = torch.cat([h, hs.pop() + control.pop()], dim=1)45 h = module(h, emb, context)46 47 h = h.type(x.dtype)48 return self.out(h)49 50 51class ControlNet(nn.Module):52 def __init__(53 self,54 image_size,55 in_channels,56 model_channels,57 hint_channels,58 num_res_blocks,59 attention_resolutions,60 dropout=0,61 channel_mult=(1, 2, 4, 8),62 conv_resample=True,63 dims=2,64 use_checkpoint=False,65 use_fp16=False,66 num_heads=-1,67 num_head_channels=-1,68 num_heads_upsample=-1,69 use_scale_shift_norm=False,70 resblock_updown=False,71 use_new_attention_order=False,72 use_spatial_transformer=False, # custom transformer support73 transformer_depth=1, # custom transformer support74 context_dim=None, # custom transformer support75 n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model76 legacy=True,77 disable_self_attentions=None,78 num_attention_blocks=None,79 disable_middle_self_attn=False,80 use_linear_in_transformer=False,81 ):82 super().__init__()83 if use_spatial_transformer:84 assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'85 86 if context_dim is not None:87 assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'88 from omegaconf.listconfig import ListConfig89 if type(context_dim) == ListConfig:90 context_dim = list(context_dim)91 92 if num_heads_upsample == -1:93 num_heads_upsample = num_heads94 95 if num_heads == -1:96 assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'97 98 if num_head_channels == -1:99 assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'100 101 self.dims = dims102 self.image_size = image_size103 self.in_channels = in_channels104 self.model_channels = model_channels105 if isinstance(num_res_blocks, int):106 self.num_res_blocks = len(channel_mult) * [num_res_blocks]107 else:108 if len(num_res_blocks) != len(channel_mult):109 raise ValueError("provide num_res_blocks either as an int (globally constant) or "110 "as a list/tuple (per-level) with the same length as channel_mult")111 self.num_res_blocks = num_res_blocks112 if disable_self_attentions is not None:113 # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not114 assert len(disable_self_attentions) == len(channel_mult)115 if num_attention_blocks is not None:116 assert len(num_attention_blocks) == len(self.num_res_blocks)117 assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))118 print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "119 f"This option has LESS priority than attention_resolutions {attention_resolutions}, "120 f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "121 f"attention will still not be set.")122 123 self.attention_resolutions = attention_resolutions124 self.dropout = dropout125 self.channel_mult = channel_mult126 self.conv_resample = conv_resample127 self.use_checkpoint = use_checkpoint128 self.dtype = th.float16 if use_fp16 else th.float32129 self.num_heads = num_heads130 self.num_head_channels = num_head_channels131 self.num_heads_upsample = num_heads_upsample132 self.predict_codebook_ids = n_embed is not None133 134 time_embed_dim = model_channels * 4135 self.time_embed = nn.Sequential(136 linear(model_channels, time_embed_dim),137 nn.SiLU(),138 linear(time_embed_dim, time_embed_dim),139 )140 141 self.input_blocks = nn.ModuleList(142 [143 TimestepEmbedSequential(144 conv_nd(dims, in_channels + hint_channels, model_channels, 3, padding=1)145 )146 ]147 )148 self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels)])149 150 self._feature_size = model_channels151 input_block_chans = [model_channels]152 ch = model_channels153 ds = 1154 for level, mult in enumerate(channel_mult):155 for nr in range(self.num_res_blocks[level]):156 layers = [157 ResBlock(158 ch,159 time_embed_dim,160 dropout,161 out_channels=mult * model_channels,162 dims=dims,163 use_checkpoint=use_checkpoint,164 use_scale_shift_norm=use_scale_shift_norm,165 )166 ]167 ch = mult * model_channels168 if ds in attention_resolutions:169 if num_head_channels == -1:170 dim_head = ch // num_heads171 else:172 num_heads = ch // num_head_channels173 dim_head = num_head_channels174 if legacy:175 # num_heads = 1176 dim_head = ch // num_heads if use_spatial_transformer else num_head_channels177 if exists(disable_self_attentions):178 disabled_sa = disable_self_attentions[level]179 else:180 disabled_sa = False181 182 if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:183 layers.append(184 AttentionBlock(185 ch,186 use_checkpoint=use_checkpoint,187 num_heads=num_heads,188 num_head_channels=dim_head,189 use_new_attention_order=use_new_attention_order,190 ) if not use_spatial_transformer else SpatialTransformer(191 ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,192 disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,193 use_checkpoint=use_checkpoint194 )195 )196 self.input_blocks.append(TimestepEmbedSequential(*layers))197 self.zero_convs.append(self.make_zero_conv(ch))198 self._feature_size += ch199 input_block_chans.append(ch)200 if level != len(channel_mult) - 1:201 out_ch = ch202 self.input_blocks.append(203 TimestepEmbedSequential(204 ResBlock(205 ch,206 time_embed_dim,207 dropout,208 out_channels=out_ch,209 dims=dims,210 use_checkpoint=use_checkpoint,211 use_scale_shift_norm=use_scale_shift_norm,212 down=True,213 )214 if resblock_updown215 else Downsample(216 ch, conv_resample, dims=dims, out_channels=out_ch217 )218 )219 )220 ch = out_ch221 input_block_chans.append(ch)222 self.zero_convs.append(self.make_zero_conv(ch))223 ds *= 2224 self._feature_size += ch225 226 if num_head_channels == -1:227 dim_head = ch // num_heads228 else:229 num_heads = ch // num_head_channels230 dim_head = num_head_channels231 if legacy:232 # num_heads = 1233 dim_head = ch // num_heads if use_spatial_transformer else num_head_channels234 self.middle_block = TimestepEmbedSequential(235 ResBlock(236 ch,237 time_embed_dim,238 dropout,239 dims=dims,240 use_checkpoint=use_checkpoint,241 use_scale_shift_norm=use_scale_shift_norm,242 ),243 AttentionBlock(244 ch,245 use_checkpoint=use_checkpoint,246 num_heads=num_heads,247 num_head_channels=dim_head,248 use_new_attention_order=use_new_attention_order,249 ) if not use_spatial_transformer else SpatialTransformer( # always uses a self-attn250 ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim,251 disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,252 use_checkpoint=use_checkpoint253 ),254 ResBlock(255 ch,256 time_embed_dim,257 dropout,258 dims=dims,259 use_checkpoint=use_checkpoint,260 use_scale_shift_norm=use_scale_shift_norm,261 ),262 )263 self.middle_block_out = self.make_zero_conv(ch)264 self._feature_size += ch265 266 def make_zero_conv(self, channels):267 return TimestepEmbedSequential(zero_module(conv_nd(self.dims, channels, channels, 1, padding=0)))268 269 def forward(self, x, hint, timesteps, context, **kwargs):270 t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)271 emb = self.time_embed(t_emb)272 x = torch.cat((x, hint), dim=1)273 outs = []274 275 h = x.type(self.dtype)276 for module, zero_conv in zip(self.input_blocks, self.zero_convs):277 h = module(h, emb, context)278 outs.append(zero_conv(h, emb, context))279 280 h = self.middle_block(h, emb, context)281 outs.append(self.middle_block_out(h, emb, context))282 283 return outs284 285 286class ControlLDM(LatentDiffusion):287 288 def __init__(289 self,290 control_stage_config: Mapping[str, Any],291 control_key: str,292 sd_locked: bool,293 only_mid_control: bool,294 learning_rate: float,295 preprocess_config,296 *args,297 **kwargs298 ) -> "ControlLDM":299 super().__init__(*args, **kwargs)300 # instantiate control module301 self.control_model: ControlNet = instantiate_from_config(control_stage_config)302 self.control_key = control_key303 self.sd_locked = sd_locked304 self.only_mid_control = only_mid_control305 self.learning_rate = learning_rate306 self.control_scales = [1.0] * 13307 308 # instantiate preprocess module (SwinIR)309 self.preprocess_model = instantiate_from_config(preprocess_config)310 frozen_module(self.preprocess_model)311 312 # instantiate condition encoder, since our condition encoder has the same 313 # structure with AE encoder, we just make a copy of AE encoder. please314 # note that AE encoder's parameters has not been initialized here.315 self.cond_encoder = nn.Sequential(OrderedDict([316 ("encoder", copy.deepcopy(self.first_stage_model.encoder)), # cond_encoder.encoder317 ("quant_conv", copy.deepcopy(self.first_stage_model.quant_conv)) # cond_encoder.quant_conv318 ]))319 frozen_module(self.cond_encoder)320 321 def apply_condition_encoder(self, control):322 c_latent_meanvar = self.cond_encoder(control * 2 - 1)323 c_latent = DiagonalGaussianDistribution(c_latent_meanvar).mode() # only use mode324 c_latent = c_latent * self.scale_factor325 return c_latent326 327 @torch.no_grad()328 def get_input(self, batch, k, bs=None, *args, **kwargs):329 x, c = super().get_input(batch, self.first_stage_key, *args, **kwargs)330 control = batch[self.control_key]331 if bs is not None:332 control = control[:bs]333 control = control.to(self.device)334 control = einops.rearrange(control, 'b h w c -> b c h w')335 control = control.to(memory_format=torch.contiguous_format).float()336 lq = control337 # apply preprocess model338 control = self.preprocess_model(control)339 # apply condition encoder340 c_latent = self.apply_condition_encoder(control)341 return x, dict(c_crossattn=[c], c_latent=[c_latent], lq=[lq], c_concat=[control])342 343 def apply_model(self, x_noisy, t, cond, *args, **kwargs):344 assert isinstance(cond, dict)345 diffusion_model = self.model.diffusion_model346 347 cond_txt = torch.cat(cond['c_crossattn'], 1)348 349 if cond['c_latent'] is None:350 eps = diffusion_model(x=x_noisy, timesteps=t, context=cond_txt, control=None, only_mid_control=self.only_mid_control)351 else:352 control = self.control_model(353 x=x_noisy, hint=torch.cat(cond['c_latent'], 1),354 timesteps=t, context=cond_txt355 )356 control = [c * scale for c, scale in zip(control, self.control_scales)]357 eps = diffusion_model(x=x_noisy, timesteps=t, context=cond_txt, control=control, only_mid_control=self.only_mid_control)358 359 return eps360 361 @torch.no_grad()362 def get_unconditional_conditioning(self, N):363 return self.get_learned_conditioning([""] * N)364 365 @torch.no_grad()366 def log_images(self, batch, sample_steps=50):367 log = dict()368 z, c = self.get_input(batch, self.first_stage_key)369 c_lq = c["lq"][0]370 c_latent = c["c_latent"][0]371 c_cat, c = c["c_concat"][0], c["c_crossattn"][0]372 373 log["hq"] = (self.decode_first_stage(z) + 1) / 2374 log["control"] = c_cat375 log["decoded_control"] = (self.decode_first_stage(c_latent) + 1) / 2376 log["lq"] = c_lq377 log["text"] = (log_txt_as_img((512, 512), batch[self.cond_stage_key], size=16) + 1) / 2378 379 samples = self.sample_log(380 # TODO: remove c_concat from cond381 cond={"c_concat": [c_cat], "c_crossattn": [c], "c_latent": [c_latent]},382 steps=sample_steps383 )384 x_samples = self.decode_first_stage(samples)385 log["samples"] = (x_samples + 1) / 2386 387 return log388 389 @torch.no_grad()390 def sample_log(self, cond, steps):391 sampler = SpacedSampler(self)392 b, c, h, w = cond["c_concat"][0].shape393 shape = (b, self.channels, h // 8, w // 8)394 samples = sampler.sample(395 steps, shape, cond, unconditional_guidance_scale=1.0,396 unconditional_conditioning=None397 )398 return samples399 400 def configure_optimizers(self):401 lr = self.learning_rate402 params = list(self.control_model.parameters())403 if not self.sd_locked:404 params += list(self.model.diffusion_model.output_blocks.parameters())405 params += list(self.model.diffusion_model.out.parameters())406 opt = torch.optim.AdamW(params, lr=lr)407 return opt408 409 def validation_step(self, batch, batch_idx):410 # TODO: 411 pass412