# -*- coding: utf-8 -*- # Copyright (c) XiMing Xing. All rights reserved. # Author: XiMing Xing # Description: import re from typing import Any, List, Optional, Union, Dict from omegaconf import DictConfig import torch import torch.nn.functional as F from torchvision import transforms from diffusers import StableDiffusionPipeline, UNet2DConditionModel from diffusers import DDIMScheduler from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import ( rescale_noise_cfg, StableDiffusionPipelineOutput) from diffusers.models.attention_processor import LoRAAttnProcessor from diffusers.loaders import AttnProcsLayers from pytorch_svgrender.diffusers_warp import init_StableDiffusion_pipeline, init_diffusers_unet class VectorizedParticleSDSPipeline(torch.nn.Module): def __init__(self, args: DictConfig, diffuser_cfg: DictConfig, guidance_cfg: DictConfig, device: torch.device): super().__init__() self.args = args self.device = device assert guidance_cfg.n_particle >= guidance_cfg.vsd_n_particle assert guidance_cfg.n_particle >= guidance_cfg.phi_n_particle pipe_kwargs = { "device": self.device, "torch_dtype": torch.float32, "local_files_only": not diffuser_cfg.download, "force_download": diffuser_cfg.force_download, "resume_download": diffuser_cfg.resume_download, "ldm_speed_up": args.x.ldm_speed_up, "enable_xformers": args.x.enable_xformers, "gradient_checkpoint": args.x.gradient_checkpoint, "cpu_offload": args.x.cpu_offload, "vae_slicing": False } # load pretrained model self.sd_pipeline = init_StableDiffusion_pipeline( args.x.model_id, custom_pipeline=StableDiffusionPipeline, custom_scheduler=DDIMScheduler, **pipe_kwargs ) # disable grads self.sd_pipeline.vae.requires_grad_(False) self.sd_pipeline.text_encoder.requires_grad_(False) self.sd_pipeline.unet.requires_grad_(False) # set components self.vae = self.sd_pipeline.vae self.unet = self.sd_pipeline.unet self.scheduler = self.sd_pipeline.scheduler self.tokenizer = self.sd_pipeline.tokenizer self.text_encoder = self.sd_pipeline.text_encoder if guidance_cfg.phi_model == 'lora': if guidance_cfg.phi_single: # default, use the single unet # load LoRA model from the pretrained model unet_ = self.unet else: # create a new unet model pipe_kwargs.pop('cpu_offload') pipe_kwargs.pop('vae_slicing') unet_ = init_diffusers_unet(args.x.model_id, **pipe_kwargs) # set correct LoRA layers self.unet_phi, phi_model_layers = self.set_lora_layers(unet_) self.phi_params = list(phi_model_layers.parameters()) self.lora_cross_attention_kwargs = {"scale": guidance_cfg.lora_attn_scale} \ if guidance_cfg.use_attn_scale else {} self.vae_phi = self.vae self.vae_phi.requires_grad_(False) elif guidance_cfg.phi_model == 'unet_simple': self.unet_phi = UNet2DConditionModel( sample_size=64, in_channels=4, out_channels=4, layers_per_block=1, block_out_channels=(128, 256, 384, 512), down_block_types=( "DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", ), up_block_types=( "AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D", ), cross_attention_dim=self.unet.config.cross_attention_dim ).to(device) self.phi_params = list(self.unet_phi.parameters()) self.vae_phi = self.vae # reset lora guidance_cfg.use_attn_scale = False guidance_cfg.lora_attn_scale = False # hyper-params self.phi_single = guidance_cfg.phi_single self.guidance_scale: float = guidance_cfg.guidance_scale self.guidance_scale_lora: float = guidance_cfg.phi_guidance_scale self.grad_clip_val: Union[float, None] = guidance_cfg.grad_clip_val self.vsd_n_particle: int = guidance_cfg.vsd_n_particle self.phi_n_particle: int = guidance_cfg.phi_n_particle self.t_schedule: str = guidance_cfg.t_schedule self.t_range = list(guidance_cfg.t_range) print( f'n_particles: {guidance_cfg.n_particle}, ' f'enhance_particles: {guidance_cfg.particle_aug}, ' f'n_particles of score: {self.vsd_n_particle}, ' f'n_particles of phi_model: {self.phi_n_particle}, \n' f't_range: {self.t_range}, ' f't_schedule: {self.t_schedule}, \n' f'guidance_scale: {self.guidance_scale}, phi_guidance_scale: {self.guidance_scale_lora}.' ) print(f"phi_model: {guidance_cfg.phi_model}, " f"use lora_cross_attn: {guidance_cfg.use_attn_scale}, " f"lora_attn_scale: {guidance_cfg.lora_attn_scale}. \n") # for convenience self.num_train_timesteps = self.scheduler.config.num_train_timesteps self.alphas = self.scheduler.alphas_cumprod.to(self.device) self.text_embeddings = None self.text_embedd_cond, self.text_embedd_uncond = None, None self.text_embeddings_phi = None self.t = None def set_lora_layers(self, unet): # set correct lora layers lora_attn_procs = {} for name in unet.attn_processors.keys(): cross_attention_dim = None if name.endswith("attn1.processor") \ else unet.config.cross_attention_dim if name.startswith("mid_block"): hidden_size = unet.config.block_out_channels[-1] elif name.startswith("up_blocks"): block_id = int(name[len("up_blocks.")]) hidden_size = list(reversed(unet.config.block_out_channels))[block_id] elif name.startswith("down_blocks"): block_id = int(name[len("down_blocks.")]) hidden_size = unet.config.block_out_channels[block_id] lora_attn_procs[name] = LoRAAttnProcessor( hidden_size=hidden_size, cross_attention_dim=cross_attention_dim ).to(self.device) unet.set_attn_processor(lora_attn_procs) lora_layers = AttnProcsLayers(unet.attn_processors) unet.requires_grad_(False) for param in lora_layers.parameters(): param.requires_grad_(True) return unet, lora_layers @torch.no_grad() def encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt=None): # text conditional embed text_inputs = self.tokenizer( prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt", ) prompt_embeds = self.text_encoder(text_inputs.input_ids.to(device))[0] if do_classifier_free_guidance: if negative_prompt is None: uncond_tokens = [""] elif isinstance(negative_prompt, str): uncond_tokens = [negative_prompt] else: uncond_tokens = negative_prompt # unconditional embed uncond_input = self.tokenizer( uncond_tokens, padding="max_length", max_length=prompt_embeds.shape[1], truncation=True, return_tensors="pt", ) negative_prompt_embeds = self.text_encoder(uncond_input.input_ids.to(device))[0] concat_prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) return concat_prompt_embeds, negative_prompt_embeds, prompt_embeds return prompt_embeds, None, None def sampling(self, vae, unet, scheduler, prompt: Union[str, List[str]] = None, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, guidance_scale: float = 7.5, negative_prompt: Optional[Union[str, List[str]]] = None, num_images_per_prompt: Optional[int] = 1, eta: float = 0.0, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, latents: Optional[torch.FloatTensor] = None, output_type: Optional[str] = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, Any]] = None, guidance_rescale: float = 0.0): # 0. Default height and width to unet vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1) height = height or unet.config.sample_size * vae_scale_factor width = width or unet.config.sample_size * vae_scale_factor # 2. Define call parameters if prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = 1 # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 # 3. Encode input prompt prompt_embeds, _, _ = self.encode_prompt( prompt, self.device, do_classifier_free_guidance, negative_prompt, ) # 4. Prepare timesteps scheduler.set_timesteps(num_inference_steps, device=self.device) timesteps = scheduler.timesteps # 5. Prepare latent variables num_channels_latents = unet.config.in_channels latents = self.sd_pipeline.prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, self.device, generator, latents, ) # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.sd_pipeline.prepare_extra_step_kwargs(generator, eta) # 7. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order with self.sd_pipeline.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = scheduler.scale_model_input(latent_model_input, t) # predict the noise residual noise_pred = unet( latent_model_input, t, encoder_hidden_states=prompt_embeds, cross_attention_kwargs=cross_attention_kwargs, return_dict=False, )[0] # perform guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) if do_classifier_free_guidance and guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) # compute the previous noisy sample x_t -> x_t-1 latents = scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] # update progress_bar if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if not output_type == "latent": image = vae.decode(latents / vae.config.scaling_factor, return_dict=False)[0] image, has_nsfw_concept = self.sd_pipeline.run_safety_checker(image, self.device, prompt_embeds.dtype) else: image = latents has_nsfw_concept = None if has_nsfw_concept is None: do_denormalize = [True] * image.shape[0] else: do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] image = self.sd_pipeline.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) # Offload last model to CPU if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: self.final_offload_hook.offload() if not return_dict: return (image, has_nsfw_concept) return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) def sample(self, prompt, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, output_type: Optional[str] = "pil"): return self.sampling(self.vae, self.unet, self.scheduler, prompt=prompt, height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=self.guidance_scale, generator=generator, output_type=output_type) def sample_lora(self, prompt, height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, output_type: Optional[str] = "pil"): return self.sampling(self.vae_phi, self.unet_phi, self.scheduler, prompt=prompt, height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=self.guidance_scale_lora, generator=generator, cross_attention_kwargs=self.lora_cross_attention_kwargs, output_type=output_type) def encode2latent(self, images): images = (2 * images - 1).clamp(-1.0, 1.0) # images: [B, 3, H, W] # encode images latents = self.vae.encode(images).latent_dist.sample() latents = self.vae.config.scaling_factor * latents return latents def get_noise_map(self, noise_pred, guidance_scale=7.5, use_cfg=True): if use_cfg: noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2) noise_map = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond) return noise_map else: return noise_pred def train_phi_model(self, pred_rgb: torch.Tensor, new_timesteps: bool = False, as_latent: bool = False): # interp to 512x512 to be fed into vae. if as_latent: latents = pred_rgb else: pred_rgb_ = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False) # encode image into latents with vae, requires grad! latents = self.encode2latent(pred_rgb_) # get phi particles indices = torch.randperm(latents.size(0)) latents_phi = latents[indices[:self.phi_n_particle]] latents_phi = latents_phi.detach() # get timestep if new_timesteps: t = torch.randint(0, self.num_train_timesteps, (1,), device=self.device) else: t = self.t noise = torch.randn_like(latents_phi) noisy_latents = self.scheduler.add_noise(latents_phi, noise, t) if self.scheduler.config.prediction_type == "epsilon": target = noise elif self.scheduler.config.prediction_type == "v_prediction": target = self.scheduler.get_velocity(latents_phi, noise, t) else: raise ValueError(f"Unknown prediction type {self.scheduler.config.prediction_type}") # predict the noise residual and compute loss noise_pred = self.unet_phi( noisy_latents, t, encoder_hidden_states=self.text_embeddings_phi, cross_attention_kwargs=self.lora_cross_attention_kwargs, ).sample return F.mse_loss(noise_pred, target, reduction="mean") def train_phi_model_refl(self, pred_rgb: torch.Tensor, weight: float = 1, new_timesteps: bool = True): # interp to 512x512 to be fed into vae. pred_rgb_ = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False) # encode image into latents with vae, requires grad! latents = self.encode2latent(pred_rgb_) # get phi particles indices = torch.randperm(latents.size(0)) latents_phi = latents[indices[:self.phi_n_particle]] latents_phi = latents_phi.detach() # get timestep if new_timesteps: t = torch.randint(0, self.num_train_timesteps, (1,), device=self.device) else: t = self.t noise = torch.randn_like(latents_phi) noisy_latents = self.scheduler.add_noise(latents_phi, noise, t) if self.scheduler.config.prediction_type == "epsilon": target = noise elif self.scheduler.config.prediction_type == "v_prediction": target = self.scheduler.get_velocity(latents_phi, noise, t) else: raise ValueError(f"Unknown prediction type {self.scheduler.config.prediction_type}") # predict the noise residual and compute loss noise_pred = self.unet_phi( noisy_latents, t, encoder_hidden_states=self.text_embedd_cond, cross_attention_kwargs=self.lora_cross_attention_kwargs, ).sample rewards = torch.tensor(weight, dtype=torch.float32, device=self.device) return rewards * F.mse_loss(noise_pred, target, reduction="mean") def schedule_timestep(self, step): min_step = int(self.num_train_timesteps * self.t_range[0]) max_step = int(self.num_train_timesteps * self.t_range[1]) if self.t_schedule == 'randint': t = torch.randint(min_step, max_step + 1, [1], dtype=torch.long, device=self.device) elif re.match(r"max_([\d.]+)_(\d+)", self.t_schedule): # Anneal time schedule # e.g: t_schedule == 'max_0.5_200' # [0.02, 0.98] -> [0.02, 0.5] after 200 steps tag, t_val, step_upd = str(self.t_schedule).split('_') t_val, step_upd = float(t_val), int(step_upd) if step >= step_upd: max_step = int(self.num_train_timesteps * t_val) t = torch.randint(min_step, max_step + 1, [1], dtype=torch.long, device=self.device) elif re.match(r"min_([\d.]+)_(\d+)", self.t_schedule): # Anneal time schedule # e.g: t_schedule == 'min_0.5_200' # [0.02, 0.98] -> [0.5, 0.98] after 200 steps tag, t_val, step_upd = str(self.t_schedule).split('_') t_val, step_upd = float(t_val), int(step_upd) if step >= step_upd: min_step = int(self.num_train_timesteps * t_val) t = torch.randint(min_step, max_step + 1, [1], dtype=torch.long, device=self.device) else: raise NotImplementedError(f"{self.t_schedule} is not support.") return t def set_text_embeddings(self, prompt, negative_prompt, do_classifier_free_guidance): if self.text_embeddings is not None: return # encode text prompt text_embeddings, text_embeddings_uncond, text_embeddings_cond = \ self.encode_prompt(prompt, self.device, do_classifier_free_guidance, negative_prompt=negative_prompt) # set pretrained model text embedding text_embeddings_uncond, text_embeddings_cond = text_embeddings.chunk(2) self.text_embedd_uncond, self.text_embedd_cond = text_embeddings_uncond, text_embeddings_cond text_embeddings_unconds = text_embeddings_uncond.repeat_interleave(self.vsd_n_particle, dim=0) text_embeddings_conds = text_embeddings_cond.repeat_interleave(self.vsd_n_particle, dim=0) text_embeddings = torch.cat([text_embeddings_unconds, text_embeddings_conds]) self.text_embeddings = text_embeddings # set phi model text embedding self.text_embeddings_phi = text_embeddings_cond.repeat_interleave(self.phi_n_particle, dim=0) def x_augment(self, x: torch.Tensor, img_size: int = 512): augment_compose = transforms.Compose([ transforms.RandomPerspective(distortion_scale=0.5, p=0.7), transforms.RandomCrop(size=(img_size, img_size), pad_if_needed=True, padding_mode='reflect') ]) return augment_compose(x) def variational_score_distillation(self, pred_rgb: torch.Tensor, step: int, prompt: Union[List, str], negative_prompt: Union[List, str] = None, grad_scale: float = 1.0, enhance_particle: bool = False, im_size: int = 512, as_latent: bool = False): bz = pred_rgb.shape[0] # data enhancement for the input particles pred_rgb = self.x_augment(pred_rgb, im_size) if enhance_particle else pred_rgb # interp to 512x512 to be fed into vae. if as_latent: latents = F.interpolate(pred_rgb, (64, 64), mode='bilinear', align_corners=False) * 2 - 1 else: pred_rgb_ = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False) # encode image into latents with vae, requires grad! # latents = self.encode2latent(pred_rgb_) latent_list = [self.encode2latent(pred_rgb_[i].unsqueeze(0)) for i in range(bz)] latents = torch.cat(latent_list, dim=0) latents = latents.to(self.device) # random sample n_particle_vsd particles from latents latents_vsd = latents[torch.randperm(bz)[:self.vsd_n_particle]] # encode input prompt do_classifier_free_guidance = True self.set_text_embeddings(prompt, negative_prompt, do_classifier_free_guidance) text_embeddings = self.text_embeddings # timestep a.k.a noise level self.t = self.schedule_timestep(step) # predict the noise residual with unet, stop gradient with torch.no_grad(): # add noise noise = torch.randn_like(latents_vsd) latents_noisy = self.scheduler.add_noise(latents_vsd, noise, self.t) # pred noise latent_model_input = torch.cat([latents_noisy] * 2) if do_classifier_free_guidance else latents_noisy # pretrained noise prediction network noise_pred_pretrain = self.unet( latent_model_input, self.t, encoder_hidden_states=text_embeddings, cross_attention_kwargs={'scale': 0.0} if self.phi_single else {} ).sample # use conditional text embeddings in phi_model _, text_embeddings_cond = text_embeddings.chunk(2) # estimated noise prediction network noise_pred_est = self.unet_phi( latents_noisy, self.t, encoder_hidden_states=text_embeddings_cond, cross_attention_kwargs=self.lora_cross_attention_kwargs ).sample # get pretrained score noise_pred_pretrain = self.get_noise_map(noise_pred_pretrain, self.guidance_scale, use_cfg=True) # get estimated score noise_pred_est = self.get_noise_map(noise_pred_est, self.guidance_scale_lora, use_cfg=False) # w(t), sigma_t^2 w = (1 - self.alphas[self.t]) grad = grad_scale * w * (noise_pred_pretrain - noise_pred_est.detach()) grad = torch.nan_to_num(grad) # grad clipping for stable training if self.grad_clip_val is not None and self.grad_clip_val > 0: grad = grad.clamp(-self.grad_clip_val, self.grad_clip_val) # re-parameterization trick: # d(loss)/d(latents) = latents - target = latents - (latents - grad) = grad target = (latents_vsd - grad).detach() loss_vpsd = 0.5 * F.mse_loss(latents_vsd, target, reduction="sum") return loss_vpsd, grad.norm(), latents, self.t