import os import math from typing import List import torch import torch.nn as nn from diffusers.pipelines.controlnet import MultiControlNetModel from .attention_processor import MaskedIPAttnProcessor2_0 as IPAttnProcessor, AttnProcessor2_0 as AttnProcessor, CNAttnProcessor2_0 as CNAttnProcessor class MSAdapter(torch.nn.Module): def __init__(self, unet, image_proj_model, adapter_modules=None, ckpt_path=None, num_tokens=4, text_tokens=77, max_rn=4, num_dummy_tokens=4, device="cuda", controlnet=None): super().__init__() self.unet = unet self.image_proj_model = image_proj_model self.adapter_modules = adapter_modules self.num_tokens = num_tokens self.num_dummy_tokens = num_dummy_tokens self.text_tokens = text_tokens self.max_rn = max_rn self.device = device self.controlnet = controlnet self.cross_attention_dim = self.unet.config.cross_attention_dim # set attention processor when inference if self.adapter_modules is None: self.set_ms_adapter() # dummy image tokens self.dummy_image_tokens = nn.Parameter(torch.randn(1, self.num_dummy_tokens, self.cross_attention_dim)) if ckpt_path is not None: self.load_from_checkpoint(ckpt_path) def forward(self, noisy_latents, timesteps, encoder_hidden_states, unet_added_cond_kwargs, image_embeds, rf_attention_mask=None, cross_attention_kwargs=None, grounding_kwargs=None): bsz = encoder_hidden_states.shape[0] if grounding_kwargs is None: ip_tokens = self.image_proj_model(image_embeds) # (bsz*rn, num_tokens, cross_attention_dim) else: ip_tokens = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs) # concat multiple images tokens ip_tokens = ip_tokens.view(bsz, -1, ip_tokens.shape[-2], ip_tokens.shape[-1]) # (bsz, rn, num_tokens, cross_attention_dim) total_num_tokens = ip_tokens.shape[-2] # num_tokens ip_tokens = ip_tokens.view(bsz, ip_tokens.shape[-3] * ip_tokens.shape[-2], ip_tokens.shape[-1]) # (bsz, total_num_tokens*rn, cross_attention_dim) dummy_image_tokens = self.dummy_image_tokens.repeat(bsz, 1, 1) ip_tokens = torch.cat([dummy_image_tokens, ip_tokens], dim=1) encoder_hidden_states = torch.cat([encoder_hidden_states, ip_tokens], dim=1) encoder_attention_mask = None if rf_attention_mask is not None: attention_mask = torch.ones((bsz, self.text_tokens)).cuda() rf_attention_mask = torch.repeat_interleave(rf_attention_mask, repeats=total_num_tokens, dim=1) encoder_attention_mask = torch.cat([attention_mask, rf_attention_mask], dim=1) # Predict the noise residual noise_pred = self.unet(noisy_latents, timesteps, encoder_hidden_states, added_cond_kwargs=unet_added_cond_kwargs, encoder_attention_mask=encoder_attention_mask, cross_attention_kwargs=cross_attention_kwargs).sample return noise_pred def save_to_checkpoint(self, output_path: str): if os.path.isdir(output_path): os.makedirs(output_path, exist_ok=True) output_path = os.path.join(output_path, "ms_adapter.bin") state_dict = { "image_proj": self.image_proj_model.state_dict(), "ms_adapter": self.adapter_modules.state_dict(), "dummy_image_tokens": self.dummy_image_tokens, } torch.save(state_dict, output_path) print(f"Successfully saved weights to checkpoint {output_path}") def load_from_checkpoint(self, ckpt_path: str): if os.path.isdir(ckpt_path): ckpt_path = os.path.join(ckpt_path, "ms_adapter.bin") # Calculate original checksums orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()])) orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()])) state_dict = torch.load(ckpt_path, map_location="cpu") # Load state dict for image_proj_model and adapter_modules when using resampler self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True) self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=True) self.load_state_dict({"dummy_image_tokens": state_dict["dummy_image_tokens"]}, strict=False) # Calculate new checksums new_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()])) new_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()])) # Verify if the weights have changed assert orig_ip_proj_sum != new_ip_proj_sum, "Weights of image_proj_model did not change!" assert orig_adapter_sum != new_adapter_sum, "Weights of adapter_modules did not change!" print(f"Successfully loaded weights from checkpoint {ckpt_path}") def set_ms_adapter(self, weight_dtype=torch.float16, cache_attention_maps=True): # set attention processor attn_procs = {} for name in self.unet.attn_processors.keys(): cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim if name.startswith("mid_block"): hidden_size = self.unet.config.block_out_channels[-1] elif name.startswith("up_blocks"): block_id = int(name[len("up_blocks.")]) hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id] elif name.startswith("down_blocks"): block_id = int(name[len("down_blocks.")]) hidden_size = self.unet.config.block_out_channels[block_id] if cross_attention_dim is None: attn_procs[name] = AttnProcessor() else: attn_procs[name] = IPAttnProcessor( hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, num_tokens=self.num_tokens, text_tokens=self.text_tokens, ).to(self.device, dtype=weight_dtype) self.unet.set_attn_processor(attn_procs) self.adapter_modules = torch.nn.ModuleList(self.unet.attn_processors.values()) if self.controlnet is not None: if isinstance(self.controlnet, MultiControlNetModel): for controlnet in self.controlnet.nets: controlnet.set_attn_processor(CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens)) else: self.controlnet.set_attn_processor(CNAttnProcessor(text_tokens=self.text_tokens, num_tokens=self.num_tokens)) @torch.inference_mode() def get_image_embeds(self, processed_images, image_encoder=None, image_proj_type="linear", image_encoder_type="clip", weight_dtype=torch.float16): # get image embeds # processed_images: [bsz, rn, ...] processed_images = processed_images.view(-1, processed_images.shape[-3], processed_images.shape[-2], processed_images.shape[-1]) # (bsz*rn, ...) if image_proj_type == "resampler": image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype), output_hidden_states=True).hidden_states[-2] # (bsz*rn, num_tokens, embedding_dim) else: image_embeds = image_encoder(processed_images.to(self.device, dtype=weight_dtype)).image_embeds # (bsz*rn, embedding_dim) return image_embeds # [bsz*rn, ...] def set_scale(self, scale, subject_scales): for attn_processor in self.pipe.unet.attn_processors.values(): if isinstance(attn_processor, IPAttnProcessor): attn_processor.scale = scale attn_processor.subject_scales = subject_scales def enable_psuedo_attention_mask(self, mask_threshold=0.5, start_step=5): for attn_processor in self.pipe.unet.attn_processors.values(): if isinstance(attn_processor, IPAttnProcessor): attn_processor.mask_threshold = mask_threshold attn_processor.start_step = start_step attn_processor.use_psuedo_attention_mask = True attn_processor.need_text_attention_map = True attn_processor.attention_maps = [] # clear attention maps def generate(self, pipe, pil_images=None, processed_images=None, prompt=None, negative_prompt=None, scale=1.0, num_samples=4, seed=None, guidance_scale=7.5, num_inference_steps=50, image_processor=None, image_encoder=None, image_proj_type="linear", image_encoder_type="clip", weight_dtype=torch.float16, boxes=None, phrases=None, drop_grounding_tokens=None, phrase_idxes=None, eot_idxes=None, height=1024, width=1024, subject_scales=None, mask_threshold=None, start_step=5, **kwargs): # generate images (validation&inference) self.pipe = pipe self.set_scale(scale, subject_scales) if mask_threshold is not None: self.enable_psuedo_attention_mask(mask_threshold, start_step) # pil_images: [[xxx, xxx, xxx], [xxx, xxx, xxx], ...] bsz = len(pil_images) # only support bsz=1 now if processed_images is None: # write in this way to promise it can be extended to batch in the future processed_images = [] for pil_image in pil_images: processed_image = image_processor(images=pil_image, return_tensors="pt").pixel_values processed_images.append(processed_image) processed_images = torch.stack(processed_images, dim=0) num_prompts = bsz if prompt is None: prompt = "best quality, high quality" if negative_prompt is None: negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" # duplicate if not isinstance(prompt, List): prompt = [prompt] * num_prompts if not isinstance(negative_prompt, List): negative_prompt = [negative_prompt] * num_prompts cross_attention_kwargs = None grounding_kwargs = None if boxes is not None: boxes = torch.tensor(boxes).to(self.device, weight_dtype) if phrases is not None: drop_grounding_tokens = drop_grounding_tokens if drop_grounding_tokens is not None else [0]*bsz batch_boxes = boxes.view(bsz*boxes.shape[1], -1) # write in this way to promise it can be extended to batch in the future phrase_input_ids = [] for phrase in phrases: phrase_input_id = pipe.tokenizer(phrase, max_length=pipe.tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt").input_ids phrase_input_ids.append(phrase_input_id) phrase_input_ids = torch.stack(phrase_input_ids) phrase_input_ids = phrase_input_ids.view(-1, phrase_input_ids.shape[-1]) phrase_embeds = pipe.text_encoder(phrase_input_ids.to(self.device)).pooler_output grounding_kwargs = {"boxes": batch_boxes, "phrase_embeds": phrase_embeds, "drop_grounding_tokens": drop_grounding_tokens} else: grounding_kwargs = None boxes = torch.repeat_interleave(boxes, repeats=num_samples, dim=0) uncond_boxes = torch.zeros_like(boxes) boxes = torch.cat([uncond_boxes, boxes], dim=0) cross_attention_kwargs = {"boxes": boxes} if phrase_idxes is not None: phrase_idxes = torch.tensor(phrase_idxes).to(self.device, torch.int) eot_idxes = torch.tensor(eot_idxes).to(self.device, torch.int) phrase_idxes = torch.repeat_interleave(phrase_idxes, repeats=num_samples, dim=0) eot_idxes = torch.repeat_interleave(eot_idxes, repeats=num_samples, dim=0) uncond_phrase_idxes = torch.zeros_like(phrase_idxes) uncond_eot_idxes = torch.zeros_like(eot_idxes) phrase_idxes = torch.cat([uncond_phrase_idxes, phrase_idxes], dim=0) eot_idxes = torch.cat([uncond_eot_idxes, eot_idxes], dim=0) if cross_attention_kwargs is None: cross_attention_kwargs = {"phrase_idxes": phrase_idxes, "eot_idxes": eot_idxes} else: cross_attention_kwargs["phrase_idxes"] = phrase_idxes cross_attention_kwargs["eot_idxes"] = eot_idxes with torch.inference_mode(): image_embeds = self.get_image_embeds(processed_images, image_encoder, image_proj_type=image_proj_type, image_encoder_type=image_encoder_type, weight_dtype=weight_dtype) image_prompt_embeds = self.image_proj_model(image_embeds, grounding_kwargs=grounding_kwargs) image_prompt_embeds = image_prompt_embeds.view(bsz, -1, image_prompt_embeds.shape[-2], image_prompt_embeds.shape[-1]) # (bsz, rn, num_tokens, cross_attention_dim) image_prompt_embeds = image_prompt_embeds.view(bsz, image_prompt_embeds.shape[-3] * image_prompt_embeds.shape[-2], image_prompt_embeds.shape[-1]) # (bsz, total_num_tokens*rn, cross_attention_dim) image_prompt_embeds = torch.cat([self.dummy_image_tokens, image_prompt_embeds], dim=1) uncond_image_prompt_embeds = torch.zeros_like(image_prompt_embeds) bs_embed, seq_len, _ = image_prompt_embeds.shape image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1) image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1) uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) prompt_embeds_, negative_prompt_embeds_, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt( prompt, device=self.device, num_images_per_prompt=num_samples, do_classifier_free_guidance=True, negative_prompt=negative_prompt, ) prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1) negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1) generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None images = pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, generator=generator, cross_attention_kwargs=cross_attention_kwargs, height=height, width=width, **kwargs, ).images return images