roi/EditP23
5
1import os2from typing import Any, Dict, Optional3from diffusers.models import AutoencoderKL, UNet2DConditionModel4from diffusers.schedulers import KarrasDiffusionSchedulers5 6import numpy7import torch8import torch.nn as nn9import transformers10from collections import OrderedDict11from PIL import Image12from torchvision import transforms13from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer14 15import diffusers16from diffusers import (17 AutoencoderKL,18 DDPMScheduler,19 DiffusionPipeline,20 EulerAncestralDiscreteScheduler,21 UNet2DConditionModel,22 ImagePipelineOutput,23)24from diffusers.image_processor import VaeImageProcessor25from diffusers.models.attention_processor import (26 Attention,27 AttnProcessor,28 XFormersAttnProcessor,29 AttnProcessor2_0,30)31from diffusers.utils.import_utils import is_xformers_available32 33 34def to_rgb_image(maybe_rgba: Image.Image):35 if maybe_rgba.mode == "RGB":36 return maybe_rgba37 elif maybe_rgba.mode == "RGBA":38 rgba = maybe_rgba39 img = numpy.random.randint(40 127, 128, size=[rgba.size[1], rgba.size[0], 3], dtype=numpy.uint841 )42 img = Image.fromarray(img, "RGB")43 img.paste(rgba, mask=rgba.getchannel("A"))44 return img45 else:46 raise ValueError("Unsupported image type.", maybe_rgba.mode)47 48 49class ReferenceOnlyAttnProc(torch.nn.Module):50 def __init__(self, chained_proc, enabled=False, name=None) -> None:51 super().__init__()52 self.enabled = enabled53 self.chained_proc = chained_proc54 self.name = name55 56 def __call__(57 self,58 attn: Attention,59 hidden_states,60 encoder_hidden_states=None,61 attention_mask=None,62 mode="w",63 ref_dict: dict = None,64 is_cfg_guidance=False,65 ) -> Any:66 if encoder_hidden_states is None:67 encoder_hidden_states = hidden_states68 if self.enabled and is_cfg_guidance:69 res0 = self.chained_proc(70 attn, hidden_states[:1], encoder_hidden_states[:1], attention_mask71 )72 hidden_states = hidden_states[1:]73 encoder_hidden_states = encoder_hidden_states[1:]74 if self.enabled:75 if mode == "w":76 ref_dict[self.name] = encoder_hidden_states77 elif mode == "r":78 encoder_hidden_states = torch.cat(79 [encoder_hidden_states, ref_dict.pop(self.name)], dim=180 )81 elif mode == "m":82 encoder_hidden_states = torch.cat(83 [encoder_hidden_states, ref_dict[self.name]], dim=184 )85 else:86 assert False, mode87 res = self.chained_proc(88 attn, hidden_states, encoder_hidden_states, attention_mask89 )90 if self.enabled and is_cfg_guidance:91 res = torch.cat([res0, res])92 return res93 94 95class RefOnlyNoisedUNet(torch.nn.Module):96 def __init__(97 self,98 unet: UNet2DConditionModel,99 train_sched: DDPMScheduler,100 val_sched: EulerAncestralDiscreteScheduler,101 ) -> None:102 super().__init__()103 self.unet = unet104 self.train_sched = train_sched105 self.val_sched = val_sched106 107 unet_lora_attn_procs = dict()108 for name, _ in unet.attn_processors.items():109 if torch.__version__ >= "2.0":110 default_attn_proc = AttnProcessor2_0()111 elif is_xformers_available():112 default_attn_proc = XFormersAttnProcessor()113 else:114 default_attn_proc = AttnProcessor()115 unet_lora_attn_procs[name] = ReferenceOnlyAttnProc(116 default_attn_proc, enabled=name.endswith("attn1.processor"), name=name117 )118 unet.set_attn_processor(unet_lora_attn_procs)119 120 def __getattr__(self, name: str):121 try:122 return super().__getattr__(name)123 except AttributeError:124 return getattr(self.unet, name)125 126 def forward_cond(127 self,128 noisy_cond_lat,129 timestep,130 encoder_hidden_states,131 class_labels,132 ref_dict,133 is_cfg_guidance,134 **kwargs,135 ):136 if is_cfg_guidance:137 encoder_hidden_states = encoder_hidden_states[1:]138 class_labels = class_labels[1:]139 self.unet(140 noisy_cond_lat,141 timestep,142 encoder_hidden_states=encoder_hidden_states,143 class_labels=class_labels,144 cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict),145 **kwargs,146 )147 148 def forward(149 self,150 sample,151 timestep,152 encoder_hidden_states,153 class_labels=None,154 *args,155 cross_attention_kwargs,156 down_block_res_samples=None,157 mid_block_res_sample=None,158 **kwargs,159 ):160 cond_lat = cross_attention_kwargs["cond_lat"]161 noisy_cond_lat = cross_attention_kwargs.get("noisy_cond_lat", None)162 is_cfg_guidance = cross_attention_kwargs.get("is_cfg_guidance", False)163 noise = torch.randn_like(cond_lat)164 if noisy_cond_lat is None:165 if self.training:166 noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep)167 noisy_cond_lat = self.train_sched.scale_model_input(168 noisy_cond_lat, timestep169 )170 else:171 noisy_cond_lat = self.val_sched.add_noise(172 cond_lat, noise, timestep.reshape(-1)173 )174 noisy_cond_lat = self.val_sched.scale_model_input(175 noisy_cond_lat, timestep.reshape(-1)176 )177 ref_dict = {}178 self.forward_cond(179 noisy_cond_lat,180 timestep,181 encoder_hidden_states,182 class_labels,183 ref_dict,184 is_cfg_guidance,185 **kwargs,186 )187 weight_dtype = self.unet.dtype188 return self.unet(189 sample,190 timestep,191 encoder_hidden_states,192 *args,193 class_labels=class_labels,194 cross_attention_kwargs=dict(195 mode="r", ref_dict=ref_dict, is_cfg_guidance=is_cfg_guidance196 ),197 down_block_additional_residuals=(198 [sample.to(dtype=weight_dtype) for sample in down_block_res_samples]199 if down_block_res_samples is not None200 else None201 ),202 mid_block_additional_residual=(203 mid_block_res_sample.to(dtype=weight_dtype)204 if mid_block_res_sample is not None205 else None206 ),207 **kwargs,208 )209 210 211def scale_latents(latents):212 latents = (latents - 0.22) * 0.75213 return latents214 215 216def unscale_latents(latents):217 latents = latents / 0.75 + 0.22218 return latents219 220 221def scale_image(image):222 image = image * 0.5 / 0.8223 return image224 225 226def unscale_image(image):227 image = image / 0.5 * 0.8228 return image229 230 231class DepthControlUNet(torch.nn.Module):232 def __init__(233 self,234 unet: RefOnlyNoisedUNet,235 controlnet: Optional[diffusers.ControlNetModel] = None,236 conditioning_scale=1.0,237 ) -> None:238 super().__init__()239 self.unet = unet240 if controlnet is None:241 self.controlnet = diffusers.ControlNetModel.from_unet(unet.unet)242 else:243 self.controlnet = controlnet244 DefaultAttnProc = AttnProcessor2_0245 if is_xformers_available():246 DefaultAttnProc = XFormersAttnProcessor247 self.controlnet.set_attn_processor(DefaultAttnProc())248 self.conditioning_scale = conditioning_scale249 250 def __getattr__(self, name: str):251 try:252 return super().__getattr__(name)253 except AttributeError:254 return getattr(self.unet, name)255 256 def forward(257 self,258 sample,259 timestep,260 encoder_hidden_states,261 class_labels=None,262 *args,263 cross_attention_kwargs: dict,264 **kwargs,265 ):266 cross_attention_kwargs = dict(cross_attention_kwargs)267 control_depth = cross_attention_kwargs.pop("control_depth")268 down_block_res_samples, mid_block_res_sample = self.controlnet(269 sample,270 timestep,271 encoder_hidden_states=encoder_hidden_states,272 controlnet_cond=control_depth,273 conditioning_scale=self.conditioning_scale,274 return_dict=False,275 )276 return self.unet(277 sample,278 timestep,279 encoder_hidden_states=encoder_hidden_states,280 down_block_res_samples=down_block_res_samples,281 mid_block_res_sample=mid_block_res_sample,282 cross_attention_kwargs=cross_attention_kwargs,283 )284 285 286class ModuleListDict(torch.nn.Module):287 def __init__(self, procs: dict) -> None:288 super().__init__()289 self.keys = sorted(procs.keys())290 self.values = torch.nn.ModuleList(procs[k] for k in self.keys)291 292 def __getitem__(self, key):293 return self.values[self.keys.index(key)]294 295 296class SuperNet(torch.nn.Module):297 def __init__(self, state_dict: Dict[str, torch.Tensor]):298 super().__init__()299 state_dict = OrderedDict((k, state_dict[k]) for k in sorted(state_dict.keys()))300 self.layers = torch.nn.ModuleList(state_dict.values())301 self.mapping = dict(enumerate(state_dict.keys()))302 self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}303 304 # .processor for unet, .self_attn for text encoder305 self.split_keys = [".processor", ".self_attn"]306 307 # we add a hook to state_dict() and load_state_dict() so that the308 # naming fits with `unet.attn_processors`309 def map_to(module, state_dict, *args, **kwargs):310 new_state_dict = {}311 for key, value in state_dict.items():312 num = int(key.split(".")[1]) # 0 is always "layers"313 new_key = key.replace(f"layers.{num}", module.mapping[num])314 new_state_dict[new_key] = value315 316 return new_state_dict317 318 def remap_key(key, state_dict):319 for k in self.split_keys:320 if k in key:321 return key.split(k)[0] + k322 return key.split(".")[0]323 324 def map_from(module, state_dict, *args, **kwargs):325 all_keys = list(state_dict.keys())326 for key in all_keys:327 replace_key = remap_key(key, state_dict)328 new_key = key.replace(329 replace_key, f"layers.{module.rev_mapping[replace_key]}"330 )331 state_dict[new_key] = state_dict[key]332 del state_dict[key]333 334 self._register_state_dict_hook(map_to)335 self._register_load_state_dict_pre_hook(map_from, with_module=True)336 337 338class Zero123PlusPipeline(diffusers.StableDiffusionPipeline):339 tokenizer: transformers.CLIPTokenizer340 text_encoder: transformers.CLIPTextModel341 vision_encoder: transformers.CLIPVisionModelWithProjection342 343 feature_extractor_clip: transformers.CLIPImageProcessor344 unet: UNet2DConditionModel345 scheduler: diffusers.schedulers.KarrasDiffusionSchedulers346 347 vae: AutoencoderKL348 ramping: nn.Linear349 350 feature_extractor_vae: transformers.CLIPImageProcessor351 352 depth_transforms_multi = transforms.Compose(353 [transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]354 )355 356 def __init__(357 self,358 vae: AutoencoderKL,359 text_encoder: CLIPTextModel,360 tokenizer: CLIPTokenizer,361 unet: UNet2DConditionModel,362 scheduler: KarrasDiffusionSchedulers,363 vision_encoder: transformers.CLIPVisionModelWithProjection,364 feature_extractor_clip: CLIPImageProcessor,365 feature_extractor_vae: CLIPImageProcessor,366 ramping_coefficients: Optional[list] = None,367 safety_checker=None,368 ):369 DiffusionPipeline.__init__(self)370 371 self.register_modules(372 vae=vae,373 text_encoder=text_encoder,374 tokenizer=tokenizer,375 unet=unet,376 scheduler=scheduler,377 safety_checker=None,378 vision_encoder=vision_encoder,379 feature_extractor_clip=feature_extractor_clip,380 feature_extractor_vae=feature_extractor_vae,381 )382 self.register_to_config(ramping_coefficients=ramping_coefficients)383 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)384 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)385 386 def prepare(self):387 train_sched = DDPMScheduler.from_config(self.scheduler.config)388 if isinstance(self.unet, UNet2DConditionModel):389 self.unet = RefOnlyNoisedUNet(self.unet, train_sched, self.scheduler).eval()390 391 def add_controlnet(392 self,393 controlnet: Optional[diffusers.ControlNetModel] = None,394 conditioning_scale=1.0,395 ):396 self.prepare()397 self.unet = DepthControlUNet(self.unet, controlnet, conditioning_scale)398 return SuperNet(OrderedDict([("controlnet", self.unet.controlnet)]))399 400 def encode_condition_image(self, image: torch.Tensor):401 image = self.vae.encode(image).latent_dist.sample()402 return image403 404 def make_condition_lat(405 self,406 local_cond_image,407 num_images_per_prompt: Optional[int] = 1,408 guidance_scale=4.0,409 ):410 local_cond_image = to_rgb_image(local_cond_image)411 local_cond_image_f = self.feature_extractor_vae(412 images=local_cond_image, return_tensors="pt"413 ).pixel_values414 415 image = local_cond_image_f.to(device=self.vae.device, dtype=self.vae.dtype)416 cond_lat = self.encode_condition_image(image)417 if guidance_scale > 1:418 negative_lat = self.encode_condition_image(torch.zeros_like(image))419 cond_lat = torch.cat([negative_lat, cond_lat])420 if num_images_per_prompt > 1:421 bs_embed, *lat_shape = cond_lat.shape422 assert len(lat_shape) == 3423 cond_lat = cond_lat.repeat(1, num_images_per_prompt, 1, 1)424 cond_lat = cond_lat.view(bs_embed * num_images_per_prompt, *lat_shape)425 return cond_lat426 427 @torch.no_grad()428 def __call__(429 self,430 image: Image.Image = None,431 prompt="",432 *args,433 num_images_per_prompt: Optional[int] = 1,434 guidance_scale=4.0,435 depth_image: Image.Image = None,436 output_type: Optional[str] = "pil",437 width=640,438 height=960,439 num_inference_steps=28,440 return_dict=True,441 noisy_cond_lat=None,442 **kwargs,443 ):444 self.prepare()445 if image is None:446 raise ValueError(447 "Inputting embeddings not supported for this pipeline. Please pass an image."448 )449 global_image = image450 local_image = image451 452 assert not isinstance(image, torch.Tensor)453 image = to_rgb_image(image)454 global_image = to_rgb_image(global_image)455 image_2 = self.feature_extractor_clip(456 images=global_image, return_tensors="pt"457 ).pixel_values458 459 if depth_image is not None and hasattr(self.unet, "controlnet"):460 depth_image = to_rgb_image(depth_image)461 depth_image = self.depth_transforms_multi(depth_image).to(462 device=self.unet.controlnet.device, dtype=self.unet.controlnet.dtype463 )464 image_2 = image_2.to(device=self.vae.device, dtype=self.vae.dtype)465 466 encoded = self.vision_encoder(image_2, output_hidden_states=False)467 global_embeds = encoded.image_embeds468 global_embeds = global_embeds.unsqueeze(-2)469 470 if hasattr(self, "encode_prompt"):471 encoder_hidden_states = self.encode_prompt(prompt, self.device, 1, False)[0]472 else:473 encoder_hidden_states = self._encode_prompt(prompt, self.device, 1, False)474 ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)475 encoder_hidden_states = encoder_hidden_states + global_embeds * ramp476 cond_lat = self.make_condition_lat(local_image, num_images_per_prompt, guidance_scale)477 478 cak = dict(cond_lat=cond_lat, noisy_cond_lat=noisy_cond_lat)479 if hasattr(self.unet, "controlnet"):480 cak["control_depth"] = depth_image481 latents: torch.Tensor = (482 super()483 .__call__(484 None,485 *args,486 cross_attention_kwargs=cak,487 guidance_scale=guidance_scale,488 num_images_per_prompt=num_images_per_prompt,489 prompt_embeds=encoder_hidden_states,490 num_inference_steps=num_inference_steps,491 output_type="latent",492 width=width,493 height=height,494 **kwargs,495 )496 .images497 )498 latents = unscale_latents(latents)499 if not output_type == "latent":500 image = unscale_image(501 self.vae.decode(502 latents / self.vae.config.scaling_factor, return_dict=False503 )[0]504 )505 else:506 image = latents507 508 image = self.image_processor.postprocess(image, output_type=output_type)509 if not return_dict:510 return (image,)511 512 return ImagePipelineOutput(images=image)513