LTT/PRM
24
1import os2import numpy as np3import torch4import torch.nn as nn5import torch.nn.functional as F6import pytorch_lightning as pl7from tqdm import tqdm8from torchvision.transforms import v29from torchvision.utils import make_grid, save_image10from einops import rearrange11 12from src.utils.train_util import instantiate_from_config13from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, DDPMScheduler, UNet2DConditionModel14from .pipeline import RefOnlyNoisedUNet15 16 17def scale_latents(latents):18 latents = (latents - 0.22) * 0.7519 return latents20 21 22def unscale_latents(latents):23 latents = latents / 0.75 + 0.2224 return latents25 26 27def scale_image(image):28 image = image * 0.5 / 0.829 return image30 31 32def unscale_image(image):33 image = image / 0.5 * 0.834 return image35 36 37def extract_into_tensor(a, t, x_shape):38 b, *_ = t.shape39 out = a.gather(-1, t)40 return out.reshape(b, *((1,) * (len(x_shape) - 1)))41 42 43class MVDiffusion(pl.LightningModule):44 def __init__(45 self,46 stable_diffusion_config,47 drop_cond_prob=0.1,48 ):49 super(MVDiffusion, self).__init__()50 51 self.drop_cond_prob = drop_cond_prob52 53 self.register_schedule()54 55 # init modules56 pipeline = DiffusionPipeline.from_pretrained(**stable_diffusion_config)57 pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(58 pipeline.scheduler.config, timestep_spacing='trailing'59 )60 self.pipeline = pipeline61 62 train_sched = DDPMScheduler.from_config(self.pipeline.scheduler.config)63 if isinstance(self.pipeline.unet, UNet2DConditionModel):64 self.pipeline.unet = RefOnlyNoisedUNet(self.pipeline.unet, train_sched, self.pipeline.scheduler)65 66 self.train_scheduler = train_sched # use ddpm scheduler during training67 68 self.unet = pipeline.unet69 70 # validation output buffer71 self.validation_step_outputs = []72 73 def register_schedule(self):74 self.num_timesteps = 100075 76 # replace scaled_linear schedule with linear schedule as Zero123++77 beta_start = 0.0008578 beta_end = 0.012079 betas = torch.linspace(beta_start, beta_end, 1000, dtype=torch.float32)80 81 alphas = 1. - betas82 alphas_cumprod = torch.cumprod(alphas, dim=0)83 alphas_cumprod_prev = torch.cat([torch.ones(1, dtype=torch.float64), alphas_cumprod[:-1]], 0)84 85 self.register_buffer('betas', betas.float())86 self.register_buffer('alphas_cumprod', alphas_cumprod.float())87 self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float())88 89 # calculations for diffusion q(x_t | x_{t-1}) and others90 self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod).float())91 self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1 - alphas_cumprod).float())92 93 self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod).float())94 self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1).float())95 96 def on_fit_start(self):97 device = torch.device(f'cuda:{self.global_rank}')98 self.pipeline.to(device)99 if self.global_rank == 0:100 os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True)101 os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True)102 103 def prepare_batch_data(self, batch):104 # prepare stable diffusion input105 cond_imgs = batch['cond_imgs'] # (B, C, H, W)106 cond_imgs = cond_imgs.to(self.device)107 108 # random resize the condition image109 cond_size = np.random.randint(128, 513)110 cond_imgs = v2.functional.resize(cond_imgs, cond_size, interpolation=3, antialias=True).clamp(0, 1)111 112 target_imgs = batch['target_imgs'] # (B, 6, C, H, W)113 target_imgs = v2.functional.resize(target_imgs, 320, interpolation=3, antialias=True).clamp(0, 1)114 target_imgs = rearrange(target_imgs, 'b (x y) c h w -> b c (x h) (y w)', x=3, y=2) # (B, C, 3H, 2W)115 target_imgs = target_imgs.to(self.device)116 117 return cond_imgs, target_imgs118 119 @torch.no_grad()120 def forward_vision_encoder(self, images):121 dtype = next(self.pipeline.vision_encoder.parameters()).dtype122 image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]123 image_pt = self.pipeline.feature_extractor_clip(images=image_pil, return_tensors="pt").pixel_values124 image_pt = image_pt.to(device=self.device, dtype=dtype)125 global_embeds = self.pipeline.vision_encoder(image_pt, output_hidden_states=False).image_embeds126 global_embeds = global_embeds.unsqueeze(-2)127 128 encoder_hidden_states = self.pipeline._encode_prompt("", self.device, 1, False)[0]129 ramp = global_embeds.new_tensor(self.pipeline.config.ramping_coefficients).unsqueeze(-1)130 encoder_hidden_states = encoder_hidden_states + global_embeds * ramp131 132 return encoder_hidden_states133 134 @torch.no_grad()135 def encode_condition_image(self, images):136 dtype = next(self.pipeline.vae.parameters()).dtype137 image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]138 image_pt = self.pipeline.feature_extractor_vae(images=image_pil, return_tensors="pt").pixel_values139 image_pt = image_pt.to(device=self.device, dtype=dtype)140 latents = self.pipeline.vae.encode(image_pt).latent_dist.sample()141 return latents142 143 @torch.no_grad()144 def encode_target_images(self, images):145 dtype = next(self.pipeline.vae.parameters()).dtype146 # equals to scaling images to [-1, 1] first and then call scale_image147 images = (images - 0.5) / 0.8 # [-0.625, 0.625]148 posterior = self.pipeline.vae.encode(images.to(dtype)).latent_dist149 latents = posterior.sample() * self.pipeline.vae.config.scaling_factor150 latents = scale_latents(latents)151 return latents152 153 def forward_unet(self, latents, t, prompt_embeds, cond_latents):154 dtype = next(self.pipeline.unet.parameters()).dtype155 latents = latents.to(dtype)156 prompt_embeds = prompt_embeds.to(dtype)157 cond_latents = cond_latents.to(dtype)158 cross_attention_kwargs = dict(cond_lat=cond_latents)159 pred_noise = self.pipeline.unet(160 latents,161 t,162 encoder_hidden_states=prompt_embeds,163 cross_attention_kwargs=cross_attention_kwargs,164 return_dict=False,165 )[0]166 return pred_noise167 168 def predict_start_from_z_and_v(self, x_t, t, v):169 return (170 extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -171 extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v172 )173 174 def get_v(self, x, noise, t):175 return (176 extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -177 extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x178 )179 180 def training_step(self, batch, batch_idx):181 # get input182 cond_imgs, target_imgs = self.prepare_batch_data(batch)183 184 # sample random timestep185 B = cond_imgs.shape[0]186 187 t = torch.randint(0, self.num_timesteps, size=(B,)).long().to(self.device)188 189 # classifier-free guidance190 if np.random.rand() < self.drop_cond_prob:191 prompt_embeds = self.pipeline._encode_prompt([""]*B, self.device, 1, False)192 cond_latents = self.encode_condition_image(torch.zeros_like(cond_imgs))193 else:194 prompt_embeds = self.forward_vision_encoder(cond_imgs)195 cond_latents = self.encode_condition_image(cond_imgs)196 197 latents = self.encode_target_images(target_imgs)198 noise = torch.randn_like(latents)199 latents_noisy = self.train_scheduler.add_noise(latents, noise, t)200 201 v_pred = self.forward_unet(latents_noisy, t, prompt_embeds, cond_latents)202 v_target = self.get_v(latents, noise, t)203 204 loss, loss_dict = self.compute_loss(v_pred, v_target)205 206 # logging207 self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True)208 self.log("global_step", self.global_step, prog_bar=True, logger=True, on_step=True, on_epoch=False)209 lr = self.optimizers().param_groups[0]['lr']210 self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)211 212 if self.global_step % 500 == 0 and self.global_rank == 0:213 with torch.no_grad():214 latents_pred = self.predict_start_from_z_and_v(latents_noisy, t, v_pred)215 216 latents = unscale_latents(latents_pred)217 images = unscale_image(self.pipeline.vae.decode(latents / self.pipeline.vae.config.scaling_factor, return_dict=False)[0]) # [-1, 1]218 images = (images * 0.5 + 0.5).clamp(0, 1)219 images = torch.cat([target_imgs, images], dim=-2)220 221 grid = make_grid(images, nrow=images.shape[0], normalize=True, value_range=(0, 1))222 save_image(grid, os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png'))223 224 return loss225 226 def compute_loss(self, noise_pred, noise_gt):227 loss = F.mse_loss(noise_pred, noise_gt)228 229 prefix = 'train'230 loss_dict = {}231 loss_dict.update({f'{prefix}/loss': loss})232 233 return loss, loss_dict234 235 @torch.no_grad()236 def validation_step(self, batch, batch_idx):237 # get input238 cond_imgs, target_imgs = self.prepare_batch_data(batch)239 240 images_pil = [v2.functional.to_pil_image(cond_imgs[i]) for i in range(cond_imgs.shape[0])]241 242 outputs = []243 for cond_img in images_pil:244 latent = self.pipeline(cond_img, num_inference_steps=75, output_type='latent').images245 image = unscale_image(self.pipeline.vae.decode(latent / self.pipeline.vae.config.scaling_factor, return_dict=False)[0]) # [-1, 1]246 image = (image * 0.5 + 0.5).clamp(0, 1)247 outputs.append(image)248 outputs = torch.cat(outputs, dim=0).to(self.device)249 images = torch.cat([target_imgs, outputs], dim=-2)250 251 self.validation_step_outputs.append(images)252 253 @torch.no_grad()254 def on_validation_epoch_end(self):255 images = torch.cat(self.validation_step_outputs, dim=0)256 257 all_images = self.all_gather(images)258 all_images = rearrange(all_images, 'r b c h w -> (r b) c h w')259 260 if self.global_rank == 0:261 grid = make_grid(all_images, nrow=8, normalize=True, value_range=(0, 1))262 save_image(grid, os.path.join(self.logdir, 'images_val', f'val_{self.global_step:07d}.png'))263 264 self.validation_step_outputs.clear() # free memory265 266 def configure_optimizers(self):267 lr = self.learning_rate268 269 optimizer = torch.optim.AdamW(self.unet.parameters(), lr=lr)270 scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 3000, eta_min=lr/4)271 272 return {'optimizer': optimizer, 'lr_scheduler': scheduler}273