models/anycontrol.py

import einops
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import cv2

from PIL import Image
from einops import rearrange, repeat
from torchvision.utils import make_grid

from ldm.models.diffusion.ddpm import LatentDiffusion
from ldm.util import log_txt_as_img, instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler

from models.q_formers import load_qformer_model


class AnyControlNet(LatentDiffusion):

 def __init__(self, mode, qformer_config=None, local_control_config=None, global_control_config=None, *args, **kwargs):
 super().__init__(*args, **kwargs)
 assert mode in ['local', 'uni']
 self.mode = mode
 self.qformer_config = qformer_config
 self.local_control_config = local_control_config
 self.global_control_config = global_control_config

 self.model.diffusion_model.requires_grad_(False)
 self.model.diffusion_model.requires_grad_(False)
 self.model.diffusion_model.requires_grad_(False)

 q_former, (vis_processor, txt_processor) = load_qformer_model(qformer_config)
 self.q_former = q_former
 self.qformer_vis_processor = vis_processor
 self.qformer_txt_processor = txt_processor

 self.local_adapter = instantiate_from_config(local_control_config)
 self.local_control_scales = [1.0] * 13
 self.global_adapter = instantiate_from_config(global_control_config) if self.mode == 'uni' else None
 self.clip_embeddings_dim = global_control_config.params.clip_embeddings_dim
 self.color_in_dim = global_control_config.params.color_in_dim

 @torch.no_grad()
 def get_input(self, batch, k, bs=None, *args, **kwargs):
 # latent and text 
 x, c = super().get_input(batch, self.first_stage_key, *args, **kwargs)
 bs = bs or x.size(0)
 shape = self.get_shape(batch, bs)
 
 local_control = self.get_local_conditions_for_vision_encoder(batch, bs)
 local_control = local_control.to(memory_format=torch.contiguous_format).float()

 global_control = {}
 global_conditions = batch['global_conditions'][:bs]
 for key in batch['global_conditions'][0].data.keys():
 global_cond = torch.stack([torch.Tensor(dc.data[key]) for dc in global_conditions])
 global_cond = global_cond.to(self.device).to(memory_format=torch.contiguous_format).float()
 global_control[key] = global_cond

 conditions = dict(
 text=[batch['txt']], 
 c_crossattn=[c], 
 local_control=[local_control], 
 global_control=[global_control],
 ) 
 return x, conditions 

 def apply_model(self, x_noisy, t, cond, local_strength=1.0, content_strength=1.0, color_strength=1.0, *args, **kwargs):
 assert isinstance(cond, dict)
 diffusion_model = self.model.diffusion_model
 cond_txt = torch.cat(cond['c_crossattn'], 1)
 text = cond['text'][0]
 bs = x_noisy.shape[0]

 # extract global control
 if self.mode in ['uni']:
 content_control, color_control = self.global_adapter(
 cond['global_control'][0]['clipembedding'], cond['global_control'][0]['color'])
 else:
 content_control = torch.zeros(bs, self.clip_embeddings_dim).to(self.device).to(memory_format=torch.contiguous_format).float()
 color_control = torch.zeros(bs, self.color_in_dim).to(self.device).to(memory_format=torch.contiguous_format).float()

 # extract local control
 if self.mode in ['local', 'uni']:
 local_features = self.local_adapter.extract_local_features(self.q_former, text, cond['local_control'][0])
 local_control = self.local_adapter(x=x_noisy, timesteps=t, context=cond_txt, local_features=local_features)
 local_control = [c * scale for c, scale in zip(local_control, self.local_control_scales)]

 eps = diffusion_model(
 x=x_noisy, timesteps=t, context=cond_txt, 
 local_control=local_control, local_w=local_strength,
 content_control=content_control, extra_w=content_strength, 
 color_control=color_control, color_w=color_strength)
 return eps

 @torch.no_grad()
 def get_unconditional_conditioning(self, N):
 return self.get_learned_conditioning([""] * N)

 @torch.no_grad()
 def get_unconditional_global_conditioning(self, c):
 if isinstance(c, dict):
 return {k:torch.zeros_like(v) for k,v in c.items()}
 elif isinstance(c, list):
 return [torch.zeros_like(v) for v in c]
 else:
 return torch.zeros_like(c) 

 @torch.no_grad()
 def get_shape(self, batch, N):
 return [dc.data[0].shape[:2] for dc in batch['local_conditions'][:N]]

 @torch.no_grad()
 def get_local_conditions_for_vision_encoder(self, batch, N):
 # return: local_conditions, (bs, num_conds * 3, h, w)
 local_conditions = []
 max_len = max([len(dc.data) for dc in batch['local_conditions'][:N]])
 for dc in batch['local_conditions'][:N]:
 conds = torch.cat([self.qformer_vis_processor['eval'](Image.fromarray(img)).unsqueeze(0) for img in dc.data], dim=1)
 local_conditions.append(conds)
 local_conditions = [F.pad(cond, (0,0,0,0,0,max_len*3-cond.shape[1],0,0)) for cond in local_conditions]
 local_conditions = torch.cat(local_conditions, dim=0).to(self.device) 
 return local_conditions

 @torch.no_grad()
 def get_local_conditions_for_logging(self, batch, N):
 local_conditions = []
 max_len = max([len(dc.data) for dc in batch['local_conditions'][:N]])
 for dc in batch['local_conditions'][:N]:
 conds = torch.stack([torch.Tensor(img).permute(2,0,1) for img in dc.data], dim=0) # (n, c, h, w)
 conds = conds.float() / 255.
 conds = conds * 2.0 - 1.0
 local_conditions.append(conds)
 local_conditions = [F.pad(cond, (0,0,0,0,0,0,0,max_len-cond.shape[0])) for cond in local_conditions]
 local_conditions = torch.stack(local_conditions, dim=0).to(self.device) # (bs, n, c, h, w)
 local_conditions = local_conditions.flatten(1,2)
 return local_conditions

 def clip_batch(self, batch, key, N, flag=True):
 if isinstance(batch, torch.Tensor):
 return batch[:N] 
 elif isinstance(batch, list):
 return batch[:N] 
 batch = batch[key][0] if flag else batch[key]
 if isinstance(batch, torch.Tensor):
 return batch[:N] 
 elif isinstance(batch, list):
 return batch[:N] 
 elif isinstance(batch, dict):
 return {k:self.clip_batch(v,'',N,flag=False) for k,v in batch.items()}
 else:
 raise ValueError(f'Unsupported type {type(batch)}')

 @torch.no_grad()
 def log_images(self, batch, N=4, n_row=2, sample=False, ddim_steps=50, ddim_eta=0.0, 
 plot_denoise_rows=False, plot_diffusion_rows=False, unconditional_guidance_scale=9.0, **kwargs):
 use_ddim = ddim_steps is not None

 log = dict()
 z, c = self.get_input(batch, self.first_stage_key, bs=N)

 shape = self.get_shape(batch, N)
 c_local = self.clip_batch(c, "local_control", N)
 c_global = self.clip_batch(c, "global_control", N)
 c_context = self.clip_batch(c, "c_crossattn", N)
 c_text = self.clip_batch(batch, self.cond_stage_key, N, False)
 
 N = min(z.shape[0], N)
 n_row = min(z.shape[0], n_row)
 log["reconstruction"] = self.decode_first_stage(z)
 log["conditioning"] = log_txt_as_img((512, 512), c_text, size=16)
 log["local_control"] = self.get_local_conditions_for_logging(batch, N)

 if plot_diffusion_rows:
 diffusion_row = list()
 z_start = z[:n_row]
 for t in range(self.num_timesteps):
 if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
 t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
 t = t.to(self.device).long()
 noise = torch.randn_like(z_start)
 z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
 diffusion_row.append(self.decode_first_stage(z_noisy))

 diffusion_row = torch.stack(diffusion_row)
 diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
 diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
 diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
 log["diffusion_row"] = diffusion_grid

 cond_dict = dict(
 local_control=[c_local],
 global_control=[c_global],
 c_crossattn=[c_context],
 text=[c_text],
 shape=[shape],
 )

 if sample:
 samples, z_denoise_row = self.sample_log(cond=cond_dict,
 batch_size=N, ddim=use_ddim,
 ddim_steps=ddim_steps, eta=ddim_eta,
 log_every_t=self.log_every_t * 0.05)
 x_samples = self.decode_first_stage(samples)
 log["samples"] = x_samples
 if plot_denoise_rows:
 if isinstance(z_denoise_row, dict):
 for key in ['pred_x0', 'x_inter']:
 z_denoise_row_key = z_denoise_row[key]
 denoise_grid = self._get_denoise_row_from_list(z_denoise_row_key)
 log[f"denoise_row_{key}"] = denoise_grid
 else:
 denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
 log["denoise_row"] = denoise_grid

 if unconditional_guidance_scale > 1.0:
 uc_context = self.get_unconditional_conditioning(N)
 uc_global = self.get_unconditional_global_conditioning(c_global)
 uc_local = c_local 
 uc_text = c_text 

 uncond_dict = dict(
 local_control=[uc_local],
 global_control=[uc_global],
 c_crossattn=[uc_context],
 text=[uc_text],
 shape=[shape]
 )

 samples_cfg, _ = self.sample_log(cond=cond_dict,
 batch_size=N, ddim=use_ddim,
 ddim_steps=ddim_steps, eta=ddim_eta,
 unconditional_guidance_scale=unconditional_guidance_scale,
 unconditional_conditioning=uncond_dict,
 )
 x_samples_cfg = self.decode_first_stage(samples_cfg)
 log[f"samples_cfg_scale_{unconditional_guidance_scale:.2f}"] = x_samples_cfg

 return log

 @torch.no_grad()
 def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
 ddim_sampler = DDIMSampler(self)
 if cond['shape'] is None:
 h, w = 512, 512
 else:
 h, w = cond["shape"][0][0]
 shape = (self.channels, h // 8, w // 8)
 samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size, shape, cond, verbose=False, **kwargs)
 return samples, intermediates

 def configure_optimizers(self):
 lr = self.learning_rate
 params = list(self.q_former.parameters()) + list(self.local_adapter.parameters())
 if not self.sd_locked:
 params += list(self.model.diffusion_model.output_blocks.parameters())
 params += list(self.model.diffusion_model.out.parameters())

 opt = torch.optim.AdamW(params, lr=lr)
 return opt

 def low_vram_shift(self, is_diffusing):
 if is_diffusing:
 self.model = self.model.cuda()
 self.local_adapter = self.local_adapter.cuda()
 self.first_stage_model = self.first_stage_model.cpu()
 self.cond_stage_model = self.cond_stage_model.cpu()
 else:
 self.model = self.model.cpu()
 self.local_adapter = self.local_adapter.cpu()
 self.first_stage_model = self.first_stage_model.cuda()
 self.cond_stage_model = self.cond_stage_model.cuda()