diff --git "a/utils/pipeline.py" "b/utils/pipeline.py" new file mode 100644--- /dev/null +++ "b/utils/pipeline.py" @@ -0,0 +1,3828 @@ +import inspect, math +from typing import Callable, List, Optional, Union +from dataclasses import dataclass +import PIL +from PIL import Image +import numpy as np +import torch +import kornia +import torch.distributed as dist +from tqdm import tqdm +from diffusers.utils import is_accelerate_available +from packaging import version +from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection,CLIPFeatureExtractor + +from diffusers.configuration_utils import FrozenDict +from diffusers.models import AutoencoderKL, UNet2DConditionModel +from diffusers import DiffusionPipeline +from diffusers.schedulers import ( + DDIMScheduler, + DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) +from diffusers.utils import deprecate, logging, BaseOutput + +from einops import rearrange + +from ref_encoder.latent_controlnet import ControlNetModel +from ref_encoder.reference_control import ReferenceAttentionControl +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin +from diffusers.models.modeling_utils import ModelMixin + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class CCProjection(ModelMixin, ConfigMixin): + def __init__(self, in_channel=772, out_channel=768): + super().__init__() + self.in_channel = in_channel + self.out_channel = out_channel + self.projection = torch.nn.Linear(in_channel, out_channel) + + def forward(self, x): + return self.projection(x) + +@dataclass +class PipelineOutput(BaseOutput): + samples: Union[torch.Tensor, np.ndarray] + + +class Hair3dPipeline(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: ControlNetModel, + # cc_projection: CCProjection, + image_encoder: CLIPVisionModelWithProjection, + feature_extractor: CLIPFeatureExtractor, + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + # cc_projection=cc_projection, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError("Please install accelerate via `pip install accelerate`") + + device = torch.device(f"cuda:{gpu_id}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + @property + def _execution_device(self): + if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): + return self.device + for module in self.unet.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def CLIP_preprocess(self, x): + dtype = x.dtype + # following openai's implementation + # TODO HF OpenAI CLIP preprocessing issue https://github.com/huggingface/transformers/issues/22505#issuecomment-1650170741 + # follow openai preprocessing to keep exact same, input tensor [-1, 1], otherwise the preprocessing will be different, https://github.com/huggingface/transformers/pull/22608 + if isinstance(x, torch.Tensor): + if x.min() < -1.0 or x.max() > 1.0: + raise ValueError("Expected input tensor to have values in the range [-1, 1]") + x = kornia.geometry.resize(x.to(torch.float32), (224, 224), interpolation='bicubic', align_corners=True, antialias=False).to(dtype=dtype) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, torch.Tensor([0.48145466, 0.4578275, 0.40821073]), + torch.Tensor([0.26862954, 0.26130258, 0.27577711])) + return x + + # from image_variation + def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + if isinstance(image, torch.Tensor): + # Batch single image + if image.ndim == 3: + assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + image = image.unsqueeze(0) + + assert image.ndim == 4, "Image must have 4 dimensions" + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError("Image should be in [-1, 1] range") + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + image = [np.array(i.convert("RGB"))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + image = image.to(device=device, dtype=dtype) + + image = self.CLIP_preprocess(image) + # if not isinstance(image, torch.Tensor): + # # 0-255 + # print("Warning: image is processed by hf's preprocess, which is different from openai original's.") + # image = self.feature_extractor(images=image, return_tensors="pt").pixel_values + image_embeddings = self.image_encoder(image).image_embeds.to(dtype=dtype) + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + + return image_embeddings + + def _encode_pose(self, pose, device, num_images_per_prompt, do_classifier_free_guidance): + + dtype = next(self.unet.parameters()).dtype + if isinstance(pose, torch.Tensor): + pose_embeddings = pose.unsqueeze(0).to(device=device, dtype=dtype) + pose_embeddings = pose_embeddings.unsqueeze(0).to(device=device, dtype=dtype) + else: + if isinstance(pose[0], list): + pose = torch.Tensor(pose) + else: + pose = torch.Tensor([pose]) + x, y, z = pose[:,0].unsqueeze(1), pose[:,1].unsqueeze(1), pose[:,2].unsqueeze(1) + pose_embeddings = torch.cat([torch.deg2rad(x), + torch.sin(torch.deg2rad(y)), + torch.cos(torch.deg2rad(y)), + z], dim=-1).unsqueeze(1).to(device=device, dtype=dtype) # B, 1, 4 + # duplicate pose embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = pose_embeddings.shape + pose_embeddings = pose_embeddings.repeat(1, num_images_per_prompt, 1) + pose_embeddings = pose_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(pose_embeddings) + pose_embeddings = torch.cat([negative_prompt_embeds, pose_embeddings]) + return pose_embeddings + + def _encode_image_with_pose(self, image, pose, device, num_images_per_prompt, do_classifier_free_guidance): + img_prompt_embeds = self._encode_image(image, device, num_images_per_prompt, False) + pose_prompt_embeds = self._encode_pose(pose, device, num_images_per_prompt, False) + prompt_embeds = torch.cat([img_prompt_embeds, pose_prompt_embeds], dim=-1) + prompt_embeds = self.cc_projection(prompt_embeds) + # follow 0123, add negative prompt, after projection + if do_classifier_free_guidance: + negative_prompt = torch.zeros_like(prompt_embeds) + prompt_embeds = torch.cat([negative_prompt, prompt_embeds]) + return prompt_embeds + + def _encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt): + if isinstance(prompt, torch.Tensor): + batch_size = prompt.shape[0] + text_input_ids = prompt + else: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, + untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents).sample + + print("image", torch.max(image), torch.min(image)) + + image = (image / 2 + 0.5).clamp(0, 1).permute(0, 2, 3, 1) + image = image.cpu().squeeze(0).float().numpy() + + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs(self, prompt, height, width, callback_steps): + if not isinstance(prompt, str) and not isinstance(prompt, list) and not isinstance(prompt, torch.Tensor): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + self.vae.encode(image[i: i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) + + image_latents = self.vae.config.scaling_factor * image_latents + + return image_latents + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, + clip_length=16): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + noise = latents.clone() + latents = latents * self.scheduler.init_noise_sigma + return latents, noise + + def prepare_condition(self, condition, device, dtype, do_classifier_free_guidance): + if isinstance(condition, torch.Tensor): + # suppose input is [-1, 1] + condition = condition + elif isinstance(condition, np.ndarray): + # suppose input is [0, 255] + condition = self.images2latents(condition, dtype).cuda() + if do_classifier_free_guidance: + condition_pad = torch.ones_like(condition) * -1 + condition = torch.cat([condition_pad, condition]) + return condition + + @torch.no_grad() + def images2latents(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + def images2latents_new(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 255.0 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + @torch.no_grad() + def encode_single_image_latents(self, images, mask, dtype): + device = self._execution_device + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device) + latents = self.vae.encode(images[None, :])['latent_dist'].mean * 0.18215 + + images = images.unsqueeze(0) + + mask = torch.from_numpy(mask).float().to(dtype).to(device) / 255.0 + if mask.ndim == 2: + mask = mask[None, None, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + + mask = F.interpolate(mask, size=latents.shape[-2:], mode='nearest') + return latents, images, mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + 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, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + controlnet_condition: list = None, + controlnet_conditioning_scale: Optional[float] = 1.0, + init_latents: Optional[torch.FloatTensor] = None, + num_actual_inference_steps: Optional[int] = None, + reference_encoder=None, + ref_image=None, + t2i=False, + style_fidelity=1.0, + prompt_img = None, + poses = None, + **kwargs, + ): + controlnet = self.controlnet + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + # batch_size = 1 if isinstance(prompt, str) else len(prompt) + batch_size = 1 + if latents is not None: + batch_size = latents.shape[0] + if isinstance(prompt, list): + batch_size = len(prompt) + + device = self._execution_device + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + if not isinstance(prompt, torch.Tensor): + prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size + if negative_prompt is not None: + negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size + text_embeddings = self._encode_prompt( + prompt, device, do_classifier_free_guidance, negative_prompt + ) + text_embeddings = torch.cat([text_embeddings]) + + reference_control_writer = ReferenceAttentionControl(reference_encoder, do_classifier_free_guidance=True, + style_fidelity=style_fidelity, + mode='write', fusion_blocks='full') + reference_control_reader = ReferenceAttentionControl(self.unet, do_classifier_free_guidance=True, mode='read', + style_fidelity=style_fidelity, + fusion_blocks='full') + + is_dist_initialized = kwargs.get("dist", False) + rank = kwargs.get("rank", 0) + + # Prepare control_img + control = self.prepare_condition( + condition=controlnet_condition, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + # for b in range(control.size(0)): + # max_value = torch.max(control[b]) + # min_value = torch.min(control[b]) + # control[b] = (control[b] - min_value) / (max_value - min_value) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + if isinstance(latents, tuple): + latents, noise = latents + + latents_dtype = latents.dtype + + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # For img2img setting + if num_actual_inference_steps is None: + num_actual_inference_steps = num_inference_steps + + if isinstance(ref_image, str): + ref_image_latents = self.images2latents(np.array(Image.open(ref_image).resize((width, height))), + latents_dtype).cuda() + elif isinstance(ref_image, np.ndarray): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + elif isinstance(ref_image, torch.Tensor): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + + + ref_padding_latents = torch.ones_like(ref_image_latents) * -1 + ref_image_latents = torch.cat([ref_padding_latents, ref_image_latents]) if do_classifier_free_guidance else ref_image_latents + # prompt_embeds = self._encode_image_with_pose(prompt_img, poses, device, 1, do_classifier_free_guidance) + + # Denoising loop + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), disable=(rank != 0)): + if num_actual_inference_steps is not None and i < num_inference_steps - num_actual_inference_steps: + continue + + # writer + ref_latents_input = ref_image_latents + reference_encoder( + ref_latents_input, + t, + # encoder_hidden_states=prompt_embeds, + encoder_hidden_states=text_embeddings, + return_dict=False, + ) + + reference_control_reader.update(reference_control_writer) + # prepare latents + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + if t2i: + pass + + else: + # controlnet + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + controlnet_cond=control, + return_dict=False, + ) + down_block_res_samples = [sample * controlnet_conditioning_scale for sample in down_block_res_samples] + mid_block_res_sample = mid_block_res_sample * controlnet_conditioning_scale + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # clean the reader + reference_control_reader.clear() + + # 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if is_dist_initialized: + dist.broadcast(latents, 0) + dist.barrier() + + reference_control_writer.clear() + + samples = self.decode_latents(latents) + if is_dist_initialized: + dist.barrier() + + # Convert to tensor + if output_type == "tensor": + samples = torch.from_numpy(samples) + + if not return_dict: + return samples + + return PipelineOutput(samples=samples) + +class Hair3dPipeline_controlnet_simple(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: ControlNetModel, + cc_projection: CCProjection, + image_encoder: CLIPVisionModelWithProjection, + feature_extractor: CLIPFeatureExtractor, + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + cc_projection=cc_projection, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError("Please install accelerate via `pip install accelerate`") + + device = torch.device(f"cuda:{gpu_id}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + @property + def _execution_device(self): + if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): + return self.device + for module in self.unet.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def CLIP_preprocess(self, x): + dtype = x.dtype + # following openai's implementation + # TODO HF OpenAI CLIP preprocessing issue https://github.com/huggingface/transformers/issues/22505#issuecomment-1650170741 + # follow openai preprocessing to keep exact same, input tensor [-1, 1], otherwise the preprocessing will be different, https://github.com/huggingface/transformers/pull/22608 + if isinstance(x, torch.Tensor): + if x.min() < -1.0 or x.max() > 1.0: + raise ValueError("Expected input tensor to have values in the range [-1, 1]") + x = kornia.geometry.resize(x.to(torch.float32), (224, 224), interpolation='bicubic', align_corners=True, antialias=False).to(dtype=dtype) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, torch.Tensor([0.48145466, 0.4578275, 0.40821073]), + torch.Tensor([0.26862954, 0.26130258, 0.27577711])) + return x + + # from image_variation + def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + if isinstance(image, torch.Tensor): + # Batch single image + if image.ndim == 3: + assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + image = image.unsqueeze(0) + + assert image.ndim == 4, "Image must have 4 dimensions" + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError("Image should be in [-1, 1] range") + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + image = [np.array(i.convert("RGB"))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + image = image.to(device=device, dtype=dtype) + + image = self.CLIP_preprocess(image) + # if not isinstance(image, torch.Tensor): + # # 0-255 + # print("Warning: image is processed by hf's preprocess, which is different from openai original's.") + # image = self.feature_extractor(images=image, return_tensors="pt").pixel_values + image_embeddings = self.image_encoder(image).image_embeds.to(dtype=dtype) + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + + return image_embeddings + + def _encode_pose(self, pose, device, num_images_per_prompt, do_classifier_free_guidance): + + dtype = next(self.unet.parameters()).dtype + if isinstance(pose, torch.Tensor): + pose_embeddings = pose.unsqueeze(0).to(device=device, dtype=dtype) + #pose_embeddings = pose_embeddings.unsqueeze(0).to(device=device, dtype=dtype) + else: + if isinstance(pose[0], list): + pose = torch.Tensor(pose) + else: + pose = torch.Tensor([pose]) + x, y, z = pose[:,0].unsqueeze(1), pose[:,1].unsqueeze(1), pose[:,2].unsqueeze(1) + pose_embeddings = torch.cat([torch.deg2rad(x), + torch.sin(torch.deg2rad(y)), + torch.cos(torch.deg2rad(y)), + z], dim=-1).unsqueeze(1).to(device=device, dtype=dtype) # B, 1, 4 + # duplicate pose embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = pose_embeddings.shape + # pose_embeddings = pose_embeddings.repeat(1, num_images_per_prompt, 1) + # pose_embeddings = pose_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(pose_embeddings) + pose_embeddings = torch.cat([negative_prompt_embeds, pose_embeddings]) + return pose_embeddings + + def _encode_image_with_pose(self, image, pose, device, num_images_per_prompt, do_classifier_free_guidance): + img_prompt_embeds = self._encode_image(image, device, num_images_per_prompt, False) + pose_prompt_embeds = self._encode_pose(pose, device, num_images_per_prompt, False) + prompt_embeds = torch.cat([img_prompt_embeds, pose_prompt_embeds], dim=-1) + prompt_embeds = self.cc_projection(prompt_embeds) + prompt_embeds = img_prompt_embeds + # follow 0123, add negative prompt, after projection + if do_classifier_free_guidance: + negative_prompt = torch.zeros_like(prompt_embeds) + prompt_embeds = torch.cat([negative_prompt, prompt_embeds]) + return prompt_embeds + + def _encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt): + if isinstance(prompt, torch.Tensor): + batch_size = prompt.shape[0] + text_input_ids = prompt + else: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, + untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents).sample + + print("image", torch.max(image), torch.min(image)) + + image = (image / 2 + 0.5).clamp(0, 1).permute(0, 2, 3, 1) + image = image.cpu().squeeze(0).float().numpy() + + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs(self, prompt, height, width, callback_steps): + if not isinstance(prompt, str) and not isinstance(prompt, list) and not isinstance(prompt, torch.Tensor): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + self.vae.encode(image[i: i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) + + image_latents = self.vae.config.scaling_factor * image_latents + + return image_latents + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, + clip_length=16): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + noise = latents.clone() + latents = latents * self.scheduler.init_noise_sigma + return latents, noise + + def prepare_condition(self, condition, device, dtype, do_classifier_free_guidance): + if isinstance(condition, torch.Tensor): + # suppose input is [-1, 1] + condition = condition + elif isinstance(condition, np.ndarray): + # suppose input is [0, 255] + condition = self.images2latents(condition, dtype).cuda() + if do_classifier_free_guidance: + condition_pad = torch.ones_like(condition) * -1 + condition = torch.cat([condition_pad, condition]) + return condition + + @torch.no_grad() + def images2latents(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + def images2latents_new(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 255.0 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + @torch.no_grad() + def encode_single_image_latents(self, images, mask, dtype): + device = self._execution_device + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device) + latents = self.vae.encode(images[None, :])['latent_dist'].mean * 0.18215 + + images = images.unsqueeze(0) + + mask = torch.from_numpy(mask).float().to(dtype).to(device) / 255.0 + if mask.ndim == 2: + mask = mask[None, None, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + + mask = F.interpolate(mask, size=latents.shape[-2:], mode='nearest') + return latents, images, mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + 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, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + controlnet_condition: list = None, + controlnet_conditioning_scale: Optional[float] = 1.0, + init_latents: Optional[torch.FloatTensor] = None, + num_actual_inference_steps: Optional[int] = None, + reference_encoder=None, + ref_image=None, + t2i=False, + style_fidelity=1.0, + prompt_img = None, + poses = None, + **kwargs, + ): + controlnet = self.controlnet + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + # batch_size = 1 if isinstance(prompt, str) else len(prompt) + batch_size = 1 + if latents is not None: + batch_size = latents.shape[0] + if isinstance(prompt, list): + batch_size = len(prompt) + + device = self._execution_device + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + if not isinstance(prompt, torch.Tensor): + prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size + if negative_prompt is not None: + negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size + text_embeddings = self._encode_prompt( + prompt, device, do_classifier_free_guidance, negative_prompt + ) + text_embeddings = torch.cat([text_embeddings]) + + # reference_control_writer = ReferenceAttentionControl(reference_encoder, do_classifier_free_guidance=True, + # style_fidelity=style_fidelity, + # mode='write', fusion_blocks='full') + # reference_control_reader = ReferenceAttentionControl(self.unet, do_classifier_free_guidance=True, mode='read', + # style_fidelity=style_fidelity, + # fusion_blocks='full') + + is_dist_initialized = kwargs.get("dist", False) + rank = kwargs.get("rank", 0) + + # Prepare control_img + control = self.prepare_condition( + condition=controlnet_condition, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + # for b in range(control.size(0)): + # max_value = torch.max(control[b]) + # min_value = torch.min(control[b]) + # control[b] = (control[b] - min_value) / (max_value - min_value) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + if isinstance(latents, tuple): + latents, noise = latents + + latents_dtype = latents.dtype + + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # For img2img setting + if num_actual_inference_steps is None: + num_actual_inference_steps = num_inference_steps + + if isinstance(ref_image, str): + ref_image_latents = self.images2latents(np.array(Image.open(ref_image).resize((width, height))), + latents_dtype).cuda() + elif isinstance(ref_image, np.ndarray): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + elif isinstance(ref_image, torch.Tensor): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + + + ref_padding_latents = torch.ones_like(ref_image_latents) * -1 + ref_image_latents = torch.cat([ref_padding_latents, ref_image_latents]) if do_classifier_free_guidance else ref_image_latents + prompt_embeds = self._encode_image_with_pose(prompt_img, poses, device, 1, do_classifier_free_guidance) + + # Denoising loop + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), disable=(rank != 0)): + if num_actual_inference_steps is not None and i < num_inference_steps - num_actual_inference_steps: + continue + + # writer + # ref_latents_input = ref_image_latents + # reference_encoder( + # ref_latents_input, + # t, + # encoder_hidden_states=text_embeddings, + # return_dict=False, + # ) + + # reference_control_reader.update(reference_control_writer) + # prepare latents + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + if t2i: + pass + + else: + # controlnet + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + controlnet_cond=control, + return_dict=False, + ) + down_block_res_samples = [sample * controlnet_conditioning_scale for sample in down_block_res_samples] + mid_block_res_sample = mid_block_res_sample * controlnet_conditioning_scale + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # clean the reader + # reference_control_reader.clear() + + # 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if is_dist_initialized: + dist.broadcast(latents, 0) + dist.barrier() + + #reference_control_writer.clear() + + samples = self.decode_latents(latents) + if is_dist_initialized: + dist.barrier() + + # Convert to tensor + if output_type == "tensor": + samples = torch.from_numpy(samples) + + if not return_dict: + return samples + + return PipelineOutput(samples=samples) + + +class Hair3dPipeline_controlnet(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: ControlNetModel, + cc_projection: CCProjection, + image_encoder: CLIPVisionModelWithProjection, + feature_extractor: CLIPFeatureExtractor, + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + cc_projection=cc_projection, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError("Please install accelerate via `pip install accelerate`") + + device = torch.device(f"cuda:{gpu_id}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + @property + def _execution_device(self): + if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): + return self.device + for module in self.unet.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def CLIP_preprocess(self, x): + dtype = x.dtype + # following openai's implementation + # TODO HF OpenAI CLIP preprocessing issue https://github.com/huggingface/transformers/issues/22505#issuecomment-1650170741 + # follow openai preprocessing to keep exact same, input tensor [-1, 1], otherwise the preprocessing will be different, https://github.com/huggingface/transformers/pull/22608 + if isinstance(x, torch.Tensor): + if x.min() < -1.0 or x.max() > 1.0: + raise ValueError("Expected input tensor to have values in the range [-1, 1]") + x = kornia.geometry.resize(x.to(torch.float32), (224, 224), interpolation='bicubic', align_corners=True, antialias=False).to(dtype=dtype) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, torch.Tensor([0.48145466, 0.4578275, 0.40821073]), + torch.Tensor([0.26862954, 0.26130258, 0.27577711])) + return x + + # from image_variation + def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + if isinstance(image, torch.Tensor): + # Batch single image + if image.ndim == 3: + assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + image = image.unsqueeze(0) + + assert image.ndim == 4, "Image must have 4 dimensions" + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError("Image should be in [-1, 1] range") + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + image = [np.array(i.convert("RGB"))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + image = image.to(device=device, dtype=dtype) + + image = self.CLIP_preprocess(image) + # if not isinstance(image, torch.Tensor): + # # 0-255 + # print("Warning: image is processed by hf's preprocess, which is different from openai original's.") + # image = self.feature_extractor(images=image, return_tensors="pt").pixel_values + image_embeddings = self.image_encoder(image).image_embeds.to(dtype=dtype) + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + + return image_embeddings + + def _encode_pose(self, pose, device, num_images_per_prompt, do_classifier_free_guidance): + + dtype = next(self.unet.parameters()).dtype + if isinstance(pose, torch.Tensor): + pose_embeddings = pose.unsqueeze(0).to(device=device, dtype=dtype) + #pose_embeddings = pose_embeddings.unsqueeze(0).to(device=device, dtype=dtype) + else: + if isinstance(pose[0], list): + pose = torch.Tensor(pose) + else: + pose = torch.Tensor([pose]) + x, y, z = pose[:,0].unsqueeze(1), pose[:,1].unsqueeze(1), pose[:,2].unsqueeze(1) + pose_embeddings = torch.cat([torch.deg2rad(x), + torch.sin(torch.deg2rad(y)), + torch.cos(torch.deg2rad(y)), + z], dim=-1).unsqueeze(1).to(device=device, dtype=dtype) # B, 1, 4 + # duplicate pose embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = pose_embeddings.shape + # pose_embeddings = pose_embeddings.repeat(1, num_images_per_prompt, 1) + # pose_embeddings = pose_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(pose_embeddings) + pose_embeddings = torch.cat([negative_prompt_embeds, pose_embeddings]) + return pose_embeddings + + def _encode_image_with_pose(self, image, pose, device, num_images_per_prompt, do_classifier_free_guidance): + img_prompt_embeds = self._encode_image(image, device, num_images_per_prompt, False) + pose_prompt_embeds = self._encode_pose(pose, device, num_images_per_prompt, False) + prompt_embeds = torch.cat([img_prompt_embeds, pose_prompt_embeds], dim=-1) + prompt_embeds = self.cc_projection(prompt_embeds) + prompt_embeds = img_prompt_embeds + # follow 0123, add negative prompt, after projection + if do_classifier_free_guidance: + negative_prompt = torch.zeros_like(prompt_embeds) + prompt_embeds = torch.cat([negative_prompt, prompt_embeds]) + return prompt_embeds + + def _encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt): + if isinstance(prompt, torch.Tensor): + batch_size = prompt.shape[0] + text_input_ids = prompt + else: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, + untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents).sample + + print("image", torch.max(image), torch.min(image)) + + image = (image / 2 + 0.5).clamp(0, 1).permute(0, 2, 3, 1) + image = image.cpu().squeeze(0).float().numpy() + + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs(self, prompt, height, width, callback_steps): + if not isinstance(prompt, str) and not isinstance(prompt, list) and not isinstance(prompt, torch.Tensor): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + self.vae.encode(image[i: i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) + + image_latents = self.vae.config.scaling_factor * image_latents + + return image_latents + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, + clip_length=16): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + noise = latents.clone() + latents = latents * self.scheduler.init_noise_sigma + return latents, noise + + def prepare_condition(self, condition, device, dtype, do_classifier_free_guidance): + if isinstance(condition, torch.Tensor): + # suppose input is [-1, 1] + condition = condition + elif isinstance(condition, np.ndarray): + # suppose input is [0, 255] + condition = self.images2latents(condition, dtype).cuda() + condition = condition/self.vae.config.scaling_factor + if do_classifier_free_guidance: + condition_pad = torch.ones_like(condition) * -1 + condition = torch.cat([condition_pad, condition]) + return condition + + @torch.no_grad() + def images2latents(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + def images2latents_new(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 255.0 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + @torch.no_grad() + def encode_single_image_latents(self, images, mask, dtype): + device = self._execution_device + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device) + latents = self.vae.encode(images[None, :])['latent_dist'].mean * 0.18215 + + images = images.unsqueeze(0) + + mask = torch.from_numpy(mask).float().to(dtype).to(device) / 255.0 + if mask.ndim == 2: + mask = mask[None, None, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + + mask = F.interpolate(mask, size=latents.shape[-2:], mode='nearest') + return latents, images, mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + 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, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + controlnet_condition: list = None, + controlnet_conditioning_scale: Optional[float] = 1.0, + init_latents: Optional[torch.FloatTensor] = None, + num_actual_inference_steps: Optional[int] = None, + reference_encoder=None, + ref_image=None, + t2i=False, + style_fidelity=1.0, + prompt_img = None, + poses = None, + **kwargs, + ): + controlnet = self.controlnet + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + # batch_size = 1 if isinstance(prompt, str) else len(prompt) + batch_size = 1 + if latents is not None: + batch_size = latents.shape[0] + if isinstance(prompt, list): + batch_size = len(prompt) + + device = self._execution_device + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + if not isinstance(prompt, torch.Tensor): + prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size + if negative_prompt is not None: + negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size + text_embeddings = self._encode_prompt( + prompt, device, do_classifier_free_guidance, negative_prompt + ) + text_embeddings = torch.cat([text_embeddings]) + + reference_control_writer = ReferenceAttentionControl(reference_encoder, do_classifier_free_guidance=True, + style_fidelity=style_fidelity, + mode='write', fusion_blocks='full') + reference_control_reader = ReferenceAttentionControl(self.unet, do_classifier_free_guidance=True, mode='read', + style_fidelity=style_fidelity, + fusion_blocks='full') + + is_dist_initialized = kwargs.get("dist", False) + rank = kwargs.get("rank", 0) + + # Prepare control_img + control = self.prepare_condition( + condition=controlnet_condition, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + # for b in range(control.size(0)): + # max_value = torch.max(control[b]) + # min_value = torch.min(control[b]) + # control[b] = (control[b] - min_value) / (max_value - min_value) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + if isinstance(latents, tuple): + latents, noise = latents + + latents_dtype = latents.dtype + + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # For img2img setting + if num_actual_inference_steps is None: + num_actual_inference_steps = num_inference_steps + + if isinstance(ref_image, str): + ref_image_latents = self.images2latents(np.array(Image.open(ref_image).resize((width, height))), + latents_dtype).cuda() + elif isinstance(ref_image, np.ndarray): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + elif isinstance(ref_image, torch.Tensor): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + + + ref_padding_latents = torch.ones_like(ref_image_latents) * -1 + ref_image_latents = torch.cat([ref_padding_latents, ref_image_latents]) if do_classifier_free_guidance else ref_image_latents + prompt_embeds = self._encode_image_with_pose(prompt_img, poses, device, 1, do_classifier_free_guidance) + + # Denoising loop + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), disable=(rank != 0)): + if num_actual_inference_steps is not None and i < num_inference_steps - num_actual_inference_steps: + continue + + # writer + ref_latents_input = ref_image_latents + reference_encoder( + ref_latents_input, + t, + encoder_hidden_states=text_embeddings, + return_dict=False, + ) + + reference_control_reader.update(reference_control_writer) + # prepare latents + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + if t2i: + pass + + else: + # controlnet + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + controlnet_cond=control, + return_dict=False, + ) + down_block_res_samples = [sample * controlnet_conditioning_scale for sample in down_block_res_samples] + mid_block_res_sample = mid_block_res_sample * controlnet_conditioning_scale + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # clean the reader + reference_control_reader.clear() + + # 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if is_dist_initialized: + dist.broadcast(latents, 0) + dist.barrier() + + reference_control_writer.clear() + + samples = self.decode_latents(latents) + if is_dist_initialized: + dist.barrier() + + # Convert to tensor + if output_type == "tensor": + samples = torch.from_numpy(samples) + + if not return_dict: + return samples + + return PipelineOutput(samples=samples) + + + + + +class Hair3dPipeline_hair_encoder(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + # controlnet: ControlNetModel, + # cc_projection: CCProjection, + image_encoder: CLIPVisionModelWithProjection, + feature_extractor: CLIPFeatureExtractor, + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + # controlnet=controlnet, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + # cc_projection=cc_projection, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError("Please install accelerate via `pip install accelerate`") + + device = torch.device(f"cuda:{gpu_id}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + @property + def _execution_device(self): + if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): + return self.device + for module in self.unet.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def CLIP_preprocess(self, x): + dtype = x.dtype + # following openai's implementation + # TODO HF OpenAI CLIP preprocessing issue https://github.com/huggingface/transformers/issues/22505#issuecomment-1650170741 + # follow openai preprocessing to keep exact same, input tensor [-1, 1], otherwise the preprocessing will be different, https://github.com/huggingface/transformers/pull/22608 + if isinstance(x, torch.Tensor): + if x.min() < -1.0 or x.max() > 1.0: + raise ValueError("Expected input tensor to have values in the range [-1, 1]") + x = kornia.geometry.resize(x.to(torch.float32), (224, 224), interpolation='bicubic', align_corners=True, antialias=False).to(dtype=dtype) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, torch.Tensor([0.48145466, 0.4578275, 0.40821073]), + torch.Tensor([0.26862954, 0.26130258, 0.27577711])) + return x + + # from image_variation + def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + if isinstance(image, torch.Tensor): + # Batch single image + if image.ndim == 3: + assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + image = image.unsqueeze(0) + + assert image.ndim == 4, "Image must have 4 dimensions" + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError("Image should be in [-1, 1] range") + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + image = [np.array(i.convert("RGB"))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + image = image.to(device=device, dtype=dtype) + + image = self.CLIP_preprocess(image) + # if not isinstance(image, torch.Tensor): + # # 0-255 + # print("Warning: image is processed by hf's preprocess, which is different from openai original's.") + # image = self.feature_extractor(images=image, return_tensors="pt").pixel_values + image_embeddings = self.image_encoder(image).image_embeds.to(dtype=dtype) + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + + return image_embeddings + + def _encode_pose(self, pose, device, num_images_per_prompt, do_classifier_free_guidance): + + dtype = next(self.unet.parameters()).dtype + if isinstance(pose, torch.Tensor): + pose_embeddings = pose.unsqueeze(0).to(device=device, dtype=dtype) + pose_embeddings = pose_embeddings.unsqueeze(0).to(device=device, dtype=dtype) + else: + if isinstance(pose[0], list): + pose = torch.Tensor(pose) + else: + pose = torch.Tensor([pose]) + x, y, z = pose[:,0].unsqueeze(1), pose[:,1].unsqueeze(1), pose[:,2].unsqueeze(1) + pose_embeddings = torch.cat([torch.deg2rad(x), + torch.sin(torch.deg2rad(y)), + torch.cos(torch.deg2rad(y)), + z], dim=-1).unsqueeze(1).to(device=device, dtype=dtype) # B, 1, 4 + # duplicate pose embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = pose_embeddings.shape + pose_embeddings = pose_embeddings.repeat(1, num_images_per_prompt, 1) + pose_embeddings = pose_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(pose_embeddings) + pose_embeddings = torch.cat([negative_prompt_embeds, pose_embeddings]) + return pose_embeddings + + def _encode_image_with_pose(self, image, pose, device, num_images_per_prompt, do_classifier_free_guidance): + img_prompt_embeds = self._encode_image(image, device, num_images_per_prompt, False) + pose_prompt_embeds = self._encode_pose(pose, device, num_images_per_prompt, False) + prompt_embeds = torch.cat([img_prompt_embeds, pose_prompt_embeds], dim=-1) + prompt_embeds = self.cc_projection(prompt_embeds) + # follow 0123, add negative prompt, after projection + if do_classifier_free_guidance: + negative_prompt = torch.zeros_like(prompt_embeds) + prompt_embeds = torch.cat([negative_prompt, prompt_embeds]) + return prompt_embeds + + def _encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt): + if isinstance(prompt, torch.Tensor): + batch_size = prompt.shape[0] + text_input_ids = prompt + else: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, + untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents).sample + + print("image", torch.max(image), torch.min(image)) + + image = (image / 2 + 0.5).clamp(0, 1).permute(0, 2, 3, 1) + image = image.cpu().squeeze(0).float().numpy() + + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs(self, prompt, height, width, callback_steps): + if not isinstance(prompt, str) and not isinstance(prompt, list) and not isinstance(prompt, torch.Tensor): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + self.vae.encode(image[i: i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) + + image_latents = self.vae.config.scaling_factor * image_latents + + return image_latents + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, + clip_length=16): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + noise = latents.clone() + latents = latents * self.scheduler.init_noise_sigma + return latents, noise + + def prepare_condition(self, condition, device, dtype, do_classifier_free_guidance): + if isinstance(condition, torch.Tensor): + # suppose input is [-1, 1] + condition = condition + elif isinstance(condition, np.ndarray): + # suppose input is [0, 255] + condition = self.images2latents(condition, dtype).cuda() + if do_classifier_free_guidance: + condition_pad = torch.ones_like(condition) * -1 + condition = torch.cat([condition_pad, condition]) + return condition + + @torch.no_grad() + def images2latents(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + def images2latents_new(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 255.0 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + @torch.no_grad() + def encode_single_image_latents(self, images, mask, dtype): + device = self._execution_device + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device) + latents = self.vae.encode(images[None, :])['latent_dist'].mean * 0.18215 + + images = images.unsqueeze(0) + + mask = torch.from_numpy(mask).float().to(dtype).to(device) / 255.0 + if mask.ndim == 2: + mask = mask[None, None, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + + mask = F.interpolate(mask, size=latents.shape[-2:], mode='nearest') + return latents, images, mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + 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, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + controlnet_condition: list = None, + controlnet_conditioning_scale: Optional[float] = 1.0, + init_latents: Optional[torch.FloatTensor] = None, + num_actual_inference_steps: Optional[int] = None, + reference_encoder=None, + ref_image=None, + t2i=False, + style_fidelity=1.0, + prompt_img = None, + poses = None, + **kwargs, + ): + controlnet = self.controlnet + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + # batch_size = 1 if isinstance(prompt, str) else len(prompt) + batch_size = 1 + if latents is not None: + batch_size = latents.shape[0] + if isinstance(prompt, list): + batch_size = len(prompt) + + device = self._execution_device + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + if not isinstance(prompt, torch.Tensor): + prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size + if negative_prompt is not None: + negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size + text_embeddings = self._encode_prompt( + prompt, device, do_classifier_free_guidance, negative_prompt + ) + text_embeddings = torch.cat([text_embeddings]) + + reference_control_writer = ReferenceAttentionControl(reference_encoder, do_classifier_free_guidance=True, + style_fidelity=style_fidelity, + mode='write', fusion_blocks='full') + reference_control_reader = ReferenceAttentionControl(self.unet, do_classifier_free_guidance=True, mode='read', + style_fidelity=style_fidelity, + fusion_blocks='full') + + is_dist_initialized = kwargs.get("dist", False) + rank = kwargs.get("rank", 0) + + # Prepare control_img + control = self.prepare_condition( + condition=controlnet_condition, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + # for b in range(control.size(0)): + # max_value = torch.max(control[b]) + # min_value = torch.min(control[b]) + # control[b] = (control[b] - min_value) / (max_value - min_value) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + if isinstance(latents, tuple): + latents, noise = latents + + latents_dtype = latents.dtype + + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # For img2img setting + if num_actual_inference_steps is None: + num_actual_inference_steps = num_inference_steps + + if isinstance(ref_image, str): + ref_image_latents = self.images2latents(np.array(Image.open(ref_image).resize((width, height))), + latents_dtype).cuda() + elif isinstance(ref_image, np.ndarray): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + elif isinstance(ref_image, torch.Tensor): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + + + ref_padding_latents = torch.ones_like(ref_image_latents) * -1 + ref_image_latents = torch.cat([ref_padding_latents, ref_image_latents]) if do_classifier_free_guidance else ref_image_latents + # prompt_embeds = self._encode_image_with_pose(prompt_img, poses, device, 1, do_classifier_free_guidance) + + # Denoising loop + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), disable=(rank != 0)): + if num_actual_inference_steps is not None and i < num_inference_steps - num_actual_inference_steps: + continue + + # writer + ref_latents_input = ref_image_latents + reference_encoder( + ref_latents_input, + t, + # encoder_hidden_states=prompt_embeds, + encoder_hidden_states=text_embeddings, + return_dict=False, + ) + + reference_control_reader.update(reference_control_writer) + # prepare latents + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + if t2i: + pass + + else: + # controlnet + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + controlnet_cond=control, + return_dict=False, + ) + down_block_res_samples = [sample * controlnet_conditioning_scale for sample in down_block_res_samples] + mid_block_res_sample = mid_block_res_sample * controlnet_conditioning_scale + + # predict the noise residual + #latent_model_input = torch.cat([latent_model_input, control], dim=1) + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # clean the reader + reference_control_reader.clear() + + # 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if is_dist_initialized: + dist.broadcast(latents, 0) + dist.barrier() + + reference_control_writer.clear() + + samples = self.decode_latents(latents) + if is_dist_initialized: + dist.barrier() + + # Convert to tensor + if output_type == "tensor": + samples = torch.from_numpy(samples) + + if not return_dict: + return samples + + return PipelineOutput(samples=samples) + + + +class Hair3dPipeline_controlnet_sv3d(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: ControlNetModel, + cc_projection: CCProjection, + image_encoder: CLIPVisionModelWithProjection, + feature_extractor: CLIPFeatureExtractor, + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + cc_projection=cc_projection, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError("Please install accelerate via `pip install accelerate`") + + device = torch.device(f"cuda:{gpu_id}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + @property + def _execution_device(self): + if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): + return self.device + for module in self.unet.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def _get_add_time_ids( + self, + noise_aug_strength: torch.tensor, + polars_rad: torch.tensor, + azimuths_rad: torch.tensor, + #dtype: torch.dtype, + # batch_size: int, + # num_videos_per_prompt: int, + do_classifier_free_guidance: bool, + ): + cond_aug = noise_aug_strength.repeat(polars_rad.shape[0], 1) + cond_aug = cond_aug.to(polars_rad.device) + # polars_rad = torch.tensor(polars_rad, dtype=dtype) + # azimuths_rad = torch.tensor(azimuths_rad, dtype=dtype) + + if do_classifier_free_guidance: + cond_aug = torch.cat([cond_aug, cond_aug]) + polars_rad = torch.cat([polars_rad, polars_rad]) + azimuths_rad = torch.cat([azimuths_rad, azimuths_rad]) + + add_time_ids = [cond_aug, polars_rad, azimuths_rad] + + return add_time_ids + + def CLIP_preprocess(self, x): + dtype = x.dtype + # following openai's implementation + # TODO HF OpenAI CLIP preprocessing issue https://github.com/huggingface/transformers/issues/22505#issuecomment-1650170741 + # follow openai preprocessing to keep exact same, input tensor [-1, 1], otherwise the preprocessing will be different, https://github.com/huggingface/transformers/pull/22608 + if isinstance(x, torch.Tensor): + if x.min() < -1.0 or x.max() > 1.0: + raise ValueError("Expected input tensor to have values in the range [-1, 1]") + x = kornia.geometry.resize(x.to(torch.float32), (224, 224), interpolation='bicubic', align_corners=True, antialias=False).to(dtype=dtype) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, torch.Tensor([0.48145466, 0.4578275, 0.40821073]), + torch.Tensor([0.26862954, 0.26130258, 0.27577711])) + return x + + # from image_variation + def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + if isinstance(image, torch.Tensor): + # Batch single image + if image.ndim == 3: + assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + image = image.unsqueeze(0) + + assert image.ndim == 4, "Image must have 4 dimensions" + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError("Image should be in [-1, 1] range") + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + image = [np.array(i.convert("RGB"))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + image = image.to(device=device, dtype=dtype) + + image = self.CLIP_preprocess(image) + # if not isinstance(image, torch.Tensor): + # # 0-255 + # print("Warning: image is processed by hf's preprocess, which is different from openai original's.") + # image = self.feature_extractor(images=image, return_tensors="pt").pixel_values + image_embeddings = self.image_encoder(image).image_embeds.to(dtype=dtype) + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + + return image_embeddings + + def _encode_pose(self, pose, device, num_images_per_prompt, do_classifier_free_guidance): + + dtype = next(self.unet.parameters()).dtype + if isinstance(pose, torch.Tensor): + pose_embeddings = pose.unsqueeze(0).to(device=device, dtype=dtype) + #pose_embeddings = pose_embeddings.unsqueeze(0).to(device=device, dtype=dtype) + else: + if isinstance(pose[0], list): + pose = torch.Tensor(pose) + else: + pose = torch.Tensor([pose]) + x, y, z = pose[:,0].unsqueeze(1), pose[:,1].unsqueeze(1), pose[:,2].unsqueeze(1) + pose_embeddings = torch.cat([torch.deg2rad(x), + torch.sin(torch.deg2rad(y)), + torch.cos(torch.deg2rad(y)), + z], dim=-1).unsqueeze(1).to(device=device, dtype=dtype) # B, 1, 4 + # duplicate pose embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = pose_embeddings.shape + # pose_embeddings = pose_embeddings.repeat(1, num_images_per_prompt, 1) + # pose_embeddings = pose_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(pose_embeddings) + pose_embeddings = torch.cat([negative_prompt_embeds, pose_embeddings]) + return pose_embeddings + + def _encode_image_with_pose(self, image, pose, device, num_images_per_prompt, do_classifier_free_guidance): + img_prompt_embeds = self._encode_image(image, device, num_images_per_prompt, False) + pose_prompt_embeds = self._encode_pose(pose, device, num_images_per_prompt, False) + prompt_embeds = torch.cat([img_prompt_embeds, pose_prompt_embeds], dim=-1) + prompt_embeds = self.cc_projection(prompt_embeds) + prompt_embeds = img_prompt_embeds + # follow 0123, add negative prompt, after projection + if do_classifier_free_guidance: + negative_prompt = torch.zeros_like(prompt_embeds) + prompt_embeds = torch.cat([negative_prompt, prompt_embeds]) + return prompt_embeds + + def _encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt): + if isinstance(prompt, torch.Tensor): + batch_size = prompt.shape[0] + text_input_ids = prompt + else: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, + untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents).sample + + print("image", torch.max(image), torch.min(image)) + + image = (image / 2 + 0.5).clamp(0, 1).permute(0, 2, 3, 1) + image = image.cpu().squeeze(0).float().numpy() + + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs(self, prompt, height, width, callback_steps): + if not isinstance(prompt, str) and not isinstance(prompt, list) and not isinstance(prompt, torch.Tensor): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + self.vae.encode(image[i: i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) + + image_latents = self.vae.config.scaling_factor * image_latents + + return image_latents + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, + clip_length=16): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + noise = latents.clone() + latents = latents * self.scheduler.init_noise_sigma + return latents, noise + + def prepare_condition(self, condition, device, dtype, do_classifier_free_guidance): + if isinstance(condition, torch.Tensor): + # suppose input is [-1, 1] + condition = condition + elif isinstance(condition, np.ndarray): + # suppose input is [0, 255] + condition = self.images2latents(condition, dtype).cuda() + if do_classifier_free_guidance: + condition_pad = torch.ones_like(condition) * -1 + condition = torch.cat([condition_pad, condition]) + return condition + + @torch.no_grad() + def images2latents(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + def images2latents_new(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 255.0 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + @torch.no_grad() + def encode_single_image_latents(self, images, mask, dtype): + device = self._execution_device + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device) + latents = self.vae.encode(images[None, :])['latent_dist'].mean * 0.18215 + + images = images.unsqueeze(0) + + mask = torch.from_numpy(mask).float().to(dtype).to(device) / 255.0 + if mask.ndim == 2: + mask = mask[None, None, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + + mask = F.interpolate(mask, size=latents.shape[-2:], mode='nearest') + return latents, images, mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + 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, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + controlnet_condition: list = None, + controlnet_conditioning_scale: Optional[float] = 1.0, + init_latents: Optional[torch.FloatTensor] = None, + num_actual_inference_steps: Optional[int] = None, + reference_encoder=None, + ref_image=None, + t2i=False, + style_fidelity=1.0, + prompt_img = None, + poses = None, + x = None, + y = None, + controlnet_ablation = False, + hair_encoder_add_xy = True, + controlnet_encoder_add_xy = True, + **kwargs, + ): + controlnet = self.controlnet + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + # batch_size = 1 if isinstance(prompt, str) else len(prompt) + batch_size = 1 + if latents is not None: + batch_size = latents.shape[0] + if isinstance(prompt, list): + batch_size = len(prompt) + + device = self._execution_device + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + if not isinstance(prompt, torch.Tensor): + prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size + if negative_prompt is not None: + negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size + text_embeddings = self._encode_prompt( + prompt, device, do_classifier_free_guidance, negative_prompt + ) + text_embeddings = torch.cat([text_embeddings]) + + reference_control_writer = ReferenceAttentionControl(reference_encoder, do_classifier_free_guidance=True, + style_fidelity=style_fidelity, + mode='write', fusion_blocks='full') + reference_control_reader = ReferenceAttentionControl(self.unet, do_classifier_free_guidance=True, mode='read', + style_fidelity=style_fidelity, + fusion_blocks='full') + + is_dist_initialized = kwargs.get("dist", False) + rank = kwargs.get("rank", 0) + + # Prepare control_img + if controlnet_ablation: + control = controlnet_condition + control = torch.from_numpy(control).float().to(controlnet.dtype) / 127.5 - 1 + control = rearrange(control, "h w c -> c h w").to(device)[None, :] + else: + control = self.prepare_condition( + condition=controlnet_condition, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + # for b in range(control.size(0)): + # max_value = torch.max(control[b]) + # min_value = torch.min(control[b]) + # control[b] = (control[b] - min_value) / (max_value - min_value) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + if isinstance(latents, tuple): + latents, noise = latents + + latents_dtype = latents.dtype + + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # For img2img setting + if num_actual_inference_steps is None: + num_actual_inference_steps = num_inference_steps + + if isinstance(ref_image, str): + ref_image_latents = self.images2latents(np.array(Image.open(ref_image).resize((width, height))), + latents_dtype).cuda() + elif isinstance(ref_image, np.ndarray): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + elif isinstance(ref_image, torch.Tensor): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + + + ref_padding_latents = torch.ones_like(ref_image_latents) * -1 + ref_image_latents = torch.cat([ref_padding_latents, ref_image_latents]) if do_classifier_free_guidance else ref_image_latents + prompt_embeds = self._encode_image_with_pose(prompt_img, poses, device, 1, do_classifier_free_guidance) + noise_aug_strength = 1e-5 + noise_aug_strength = torch.tensor(noise_aug_strength, dtype=torch.float32).unsqueeze(0) + noise_aug_strength = noise_aug_strength.to(device) + if (x is not None) and (y is not None): + x = x.to(device) + y = y.to(device) + add_time_ids = self._get_add_time_ids(noise_aug_strength, x, y, do_classifier_free_guidance) + else: + add_time_ids = None + + + # Denoising loop + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), disable=(rank != 0)): + if num_actual_inference_steps is not None and i < num_inference_steps - num_actual_inference_steps: + continue + + # writer + ref_latents_input = ref_image_latents + + if hair_encoder_add_xy: + reference_encoder( + ref_latents_input, + t, + encoder_hidden_states=text_embeddings, + return_dict=False, + # add_time_ids = add_time_ids, + ) + else: + reference_encoder( + ref_latents_input, + t, + encoder_hidden_states=text_embeddings, + return_dict=False, + add_time_ids = None, + ) + + # reference_control_reader.update(reference_control_writer) + + # prepare latents + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + if t2i: + pass + + else: + # controlnet + if controlnet_encoder_add_xy: + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + controlnet_cond=control, + return_dict=False, + add_time_ids = add_time_ids, + ) + else: + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + controlnet_cond=control, + return_dict=False, + add_time_ids = None, + ) + + down_block_res_samples = [sample * controlnet_conditioning_scale for sample in down_block_res_samples] + mid_block_res_sample = mid_block_res_sample * controlnet_conditioning_scale + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # clean the reader + # reference_control_reader.clear() + + # 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if is_dist_initialized: + dist.broadcast(latents, 0) + dist.barrier() + + # reference_control_writer.clear() + + samples = self.decode_latents(latents) + if is_dist_initialized: + dist.barrier() + + # Convert to tensor + if output_type == "tensor": + samples = torch.from_numpy(samples) + + if not return_dict: + return samples + + return PipelineOutput(samples=samples) + + + +class Hair3dPipeline_controlnet_only_sv3d(DiffusionPipeline): + _optional_components = [] + + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + unet: UNet2DConditionModel, + controlnet: ControlNetModel, + cc_projection: CCProjection, + image_encoder: CLIPVisionModelWithProjection, + feature_extractor: CLIPFeatureExtractor, + scheduler: Union[ + DDIMScheduler, + PNDMScheduler, + LMSDiscreteScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, + ], + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + controlnet=controlnet, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + cc_projection=cc_projection, + ) + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + + def enable_vae_slicing(self): + self.vae.enable_slicing() + + def disable_vae_slicing(self): + self.vae.disable_slicing() + + def enable_sequential_cpu_offload(self, gpu_id=0): + if is_accelerate_available(): + from accelerate import cpu_offload + else: + raise ImportError("Please install accelerate via `pip install accelerate`") + + device = torch.device(f"cuda:{gpu_id}") + + for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]: + if cpu_offloaded_model is not None: + cpu_offload(cpu_offloaded_model, device) + + @property + def _execution_device(self): + if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"): + return self.device + for module in self.unet.modules(): + if ( + hasattr(module, "_hf_hook") + and hasattr(module._hf_hook, "execution_device") + and module._hf_hook.execution_device is not None + ): + return torch.device(module._hf_hook.execution_device) + return self.device + + def _get_add_time_ids( + self, + noise_aug_strength: torch.tensor, + polars_rad: torch.tensor, + azimuths_rad: torch.tensor, + #dtype: torch.dtype, + # batch_size: int, + # num_videos_per_prompt: int, + do_classifier_free_guidance: bool, + ): + cond_aug = noise_aug_strength.repeat(polars_rad.shape[0], 1) + cond_aug = cond_aug.to(polars_rad.device) + # polars_rad = torch.tensor(polars_rad, dtype=dtype) + # azimuths_rad = torch.tensor(azimuths_rad, dtype=dtype) + + if do_classifier_free_guidance: + cond_aug = torch.cat([cond_aug, cond_aug]) + polars_rad = torch.cat([polars_rad, polars_rad]) + azimuths_rad = torch.cat([azimuths_rad, azimuths_rad]) + + add_time_ids = [cond_aug, polars_rad, azimuths_rad] + + return add_time_ids + + def CLIP_preprocess(self, x): + dtype = x.dtype + # following openai's implementation + # TODO HF OpenAI CLIP preprocessing issue https://github.com/huggingface/transformers/issues/22505#issuecomment-1650170741 + # follow openai preprocessing to keep exact same, input tensor [-1, 1], otherwise the preprocessing will be different, https://github.com/huggingface/transformers/pull/22608 + if isinstance(x, torch.Tensor): + if x.min() < -1.0 or x.max() > 1.0: + raise ValueError("Expected input tensor to have values in the range [-1, 1]") + x = kornia.geometry.resize(x.to(torch.float32), (224, 224), interpolation='bicubic', align_corners=True, antialias=False).to(dtype=dtype) + x = (x + 1.) / 2. + # renormalize according to clip + x = kornia.enhance.normalize(x, torch.Tensor([0.48145466, 0.4578275, 0.40821073]), + torch.Tensor([0.26862954, 0.26130258, 0.27577711])) + return x + + # from image_variation + def _encode_image(self, image, device, num_images_per_prompt, do_classifier_free_guidance): + dtype = next(self.image_encoder.parameters()).dtype + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + if isinstance(image, torch.Tensor): + # Batch single image + if image.ndim == 3: + assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + image = image.unsqueeze(0) + + assert image.ndim == 4, "Image must have 4 dimensions" + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError("Image should be in [-1, 1] range") + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + image = [np.array(i.convert("RGB"))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + image = image.to(device=device, dtype=dtype) + + image = self.CLIP_preprocess(image) + # if not isinstance(image, torch.Tensor): + # # 0-255 + # print("Warning: image is processed by hf's preprocess, which is different from openai original's.") + # image = self.feature_extractor(images=image, return_tensors="pt").pixel_values + image_embeddings = self.image_encoder(image).image_embeds.to(dtype=dtype) + image_embeddings = image_embeddings.unsqueeze(1) + + # duplicate image embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = image_embeddings.shape + image_embeddings = image_embeddings.repeat(1, num_images_per_prompt, 1) + image_embeddings = image_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(image_embeddings) + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + image_embeddings = torch.cat([negative_prompt_embeds, image_embeddings]) + + return image_embeddings + + def _encode_pose(self, pose, device, num_images_per_prompt, do_classifier_free_guidance): + + dtype = next(self.unet.parameters()).dtype + if isinstance(pose, torch.Tensor): + pose_embeddings = pose.unsqueeze(0).to(device=device, dtype=dtype) + #pose_embeddings = pose_embeddings.unsqueeze(0).to(device=device, dtype=dtype) + else: + if isinstance(pose[0], list): + pose = torch.Tensor(pose) + else: + pose = torch.Tensor([pose]) + x, y, z = pose[:,0].unsqueeze(1), pose[:,1].unsqueeze(1), pose[:,2].unsqueeze(1) + pose_embeddings = torch.cat([torch.deg2rad(x), + torch.sin(torch.deg2rad(y)), + torch.cos(torch.deg2rad(y)), + z], dim=-1).unsqueeze(1).to(device=device, dtype=dtype) # B, 1, 4 + # duplicate pose embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = pose_embeddings.shape + # pose_embeddings = pose_embeddings.repeat(1, num_images_per_prompt, 1) + # pose_embeddings = pose_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1) + if do_classifier_free_guidance: + negative_prompt_embeds = torch.zeros_like(pose_embeddings) + pose_embeddings = torch.cat([negative_prompt_embeds, pose_embeddings]) + return pose_embeddings + + def _encode_image_with_pose(self, image, pose, device, num_images_per_prompt, do_classifier_free_guidance): + img_prompt_embeds = self._encode_image(image, device, num_images_per_prompt, False) + pose_prompt_embeds = self._encode_pose(pose, device, num_images_per_prompt, False) + prompt_embeds = torch.cat([img_prompt_embeds, pose_prompt_embeds], dim=-1) + prompt_embeds = self.cc_projection(prompt_embeds) + prompt_embeds = img_prompt_embeds + # follow 0123, add negative prompt, after projection + if do_classifier_free_guidance: + negative_prompt = torch.zeros_like(prompt_embeds) + prompt_embeds = torch.cat([negative_prompt, prompt_embeds]) + return prompt_embeds + + def _encode_prompt(self, prompt, device, do_classifier_free_guidance, negative_prompt): + if isinstance(prompt, torch.Tensor): + batch_size = prompt.shape[0] + text_input_ids = prompt + else: + batch_size = 1 if isinstance(prompt, str) else len(prompt) + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, + untruncated_ids): + removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1: -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {self.tokenizer.model_max_length} tokens: {removed_text}" + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = text_inputs.attention_mask.to(device) + else: + attention_mask = None + + text_embeddings = self.text_encoder( + text_input_ids.to(device), + attention_mask=attention_mask, + ) + text_embeddings = text_embeddings[0] + + # duplicate text embeddings for each generation per prompt, using mps friendly method + bs_embed, seq_len, _ = text_embeddings.shape + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance: + uncond_tokens: List[str] + if negative_prompt is None: + uncond_tokens = [""] * batch_size + elif type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + elif isinstance(negative_prompt, str): + uncond_tokens = [negative_prompt] + elif batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + else: + uncond_tokens = negative_prompt + + max_length = text_input_ids.shape[-1] + uncond_input = self.tokenizer( + uncond_tokens, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="pt", + ) + + if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask: + attention_mask = uncond_input.attention_mask.to(device) + else: + attention_mask = None + + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(device), + attention_mask=attention_mask, + ) + uncond_embeddings = uncond_embeddings[0] + + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = uncond_embeddings.shape[1] + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def decode_latents(self, latents): + + latents = 1 / 0.18215 * latents + image = self.vae.decode(latents).sample + + print("image", torch.max(image), torch.min(image)) + + image = (image / 2 + 0.5).clamp(0, 1).permute(0, 2, 3, 1) + image = image.cpu().squeeze(0).float().numpy() + + return image + + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs(self, prompt, height, width, callback_steps): + if not isinstance(prompt, str) and not isinstance(prompt, list) and not isinstance(prompt, torch.Tensor): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if height % 8 != 0 or width % 8 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.") + + if (callback_steps is None) or ( + callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0) + ): + raise ValueError( + f"`callback_steps` has to be a positive integer but is {callback_steps} of type" + f" {type(callback_steps)}." + ) + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + self.vae.encode(image[i: i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = self.vae.encode(image).latent_dist.sample(generator=generator) + + image_latents = self.vae.config.scaling_factor * image_latents + + return image_latents + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, + clip_length=16): + shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + rand_device = "cpu" if device.type == "mps" else device + + if isinstance(generator, list): + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device) + + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + + # scale the initial noise by the standard deviation required by the scheduler + noise = latents.clone() + latents = latents * self.scheduler.init_noise_sigma + return latents, noise + + def prepare_condition(self, condition, device, dtype, do_classifier_free_guidance): + if isinstance(condition, torch.Tensor): + # suppose input is [-1, 1] + condition = condition + elif isinstance(condition, np.ndarray): + # suppose input is [0, 255] + condition = self.images2latents(condition, dtype).cuda() + if do_classifier_free_guidance: + condition_pad = torch.ones_like(condition) * -1 + condition = torch.cat([condition_pad, condition]) + return condition + + @torch.no_grad() + def images2latents(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + def images2latents_new(self, images, dtype): + """ + Convert RGB image to VAE latents + """ + device = self._execution_device + if isinstance(images, torch.Tensor): + # suppose input is [-1, 1] + images = images.to(dtype) + if images.ndim == 3: + images = images.unsqueeze(0) + elif isinstance(images, np.ndarray): + # suppose input is [0, 255] + images = torch.from_numpy(images).float().to(dtype) / 255.0 + images = rearrange(images, "h w c -> c h w").to(device)[None, :] + latents = self.vae.encode(images)['latent_dist'].mean * self.vae.config.scaling_factor + return latents + + @torch.no_grad() + def encode_single_image_latents(self, images, mask, dtype): + device = self._execution_device + images = torch.from_numpy(images).float().to(dtype) / 127.5 - 1 + images = rearrange(images, "h w c -> c h w").to(device) + latents = self.vae.encode(images[None, :])['latent_dist'].mean * 0.18215 + + images = images.unsqueeze(0) + + mask = torch.from_numpy(mask).float().to(dtype).to(device) / 255.0 + if mask.ndim == 2: + mask = mask[None, None, :] + elif mask.ndim == 3: + mask = mask[:, None, :, :] + + mask = F.interpolate(mask, size=latents.shape[-2:], mode='nearest') + return latents, images, mask + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]], + 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, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: Optional[int] = 1, + controlnet_condition: list = None, + controlnet_conditioning_scale: Optional[float] = 1.0, + init_latents: Optional[torch.FloatTensor] = None, + num_actual_inference_steps: Optional[int] = None, + reference_encoder=None, + ref_image=None, + t2i=False, + style_fidelity=1.0, + prompt_img = None, + poses = None, + x = None, + y = None, + **kwargs, + ): + controlnet = self.controlnet + + # Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # Check inputs. Raise error if not correct + self.check_inputs(prompt, height, width, callback_steps) + + # Define call parameters + # batch_size = 1 if isinstance(prompt, str) else len(prompt) + batch_size = 1 + if latents is not None: + batch_size = latents.shape[0] + if isinstance(prompt, list): + batch_size = len(prompt) + + device = self._execution_device + do_classifier_free_guidance = guidance_scale > 1.0 + + # Encode input prompt + if not isinstance(prompt, torch.Tensor): + prompt = prompt if isinstance(prompt, list) else [prompt] * batch_size + if negative_prompt is not None: + negative_prompt = negative_prompt if isinstance(negative_prompt, list) else [negative_prompt] * batch_size + text_embeddings = self._encode_prompt( + prompt, device, do_classifier_free_guidance, negative_prompt + ) + text_embeddings = torch.cat([text_embeddings]) + + reference_control_writer = ReferenceAttentionControl(reference_encoder, do_classifier_free_guidance=True, + style_fidelity=style_fidelity, + mode='write', fusion_blocks='full') + reference_control_reader = ReferenceAttentionControl(self.unet, do_classifier_free_guidance=True, mode='read', + style_fidelity=style_fidelity, + fusion_blocks='full') + + is_dist_initialized = kwargs.get("dist", False) + rank = kwargs.get("rank", 0) + + # Prepare control_img + control = self.prepare_condition( + condition=controlnet_condition, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + # for b in range(control.size(0)): + # max_value = torch.max(control[b]) + # min_value = torch.min(control[b]) + # control[b] = (control[b] - min_value) / (max_value - min_value) + + # Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size, + num_channels_latents, + height, + width, + text_embeddings.dtype, + device, + generator, + latents, + ) + if isinstance(latents, tuple): + latents, noise = latents + + latents_dtype = latents.dtype + + # Prepare extra step kwargs. + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # For img2img setting + if num_actual_inference_steps is None: + num_actual_inference_steps = num_inference_steps + + if isinstance(ref_image, str): + ref_image_latents = self.images2latents(np.array(Image.open(ref_image).resize((width, height))), + latents_dtype).cuda() + elif isinstance(ref_image, np.ndarray): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + elif isinstance(ref_image, torch.Tensor): + ref_image_latents = self.images2latents(ref_image, latents_dtype).cuda() + + + ref_padding_latents = torch.ones_like(ref_image_latents) * -1 + ref_image_latents = torch.cat([ref_padding_latents, ref_image_latents]) if do_classifier_free_guidance else ref_image_latents + prompt_embeds = self._encode_image_with_pose(prompt_img, poses, device, 1, do_classifier_free_guidance) + noise_aug_strength = 1e-5 + noise_aug_strength = torch.tensor(noise_aug_strength, dtype=torch.float32).unsqueeze(0) + noise_aug_strength = noise_aug_strength.to(device) + if (x is not None) and (y is not None): + x = x.to(device) + y = y.to(device) + add_time_ids = self._get_add_time_ids(noise_aug_strength, x, y, do_classifier_free_guidance) + else: + add_time_ids = None + + + # Denoising loop + for i, t in tqdm(enumerate(timesteps), total=len(timesteps), disable=(rank != 0)): + if num_actual_inference_steps is not None and i < num_inference_steps - num_actual_inference_steps: + continue + + # writer + # prepare latents + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + + if t2i: + pass + + else: + # controlnet + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + controlnet_cond=control, + return_dict=False, + add_time_ids = add_time_ids, + ) + down_block_res_samples = [sample * controlnet_conditioning_scale for sample in down_block_res_samples] + mid_block_res_sample = mid_block_res_sample * controlnet_conditioning_scale + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=text_embeddings, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + 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) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + if is_dist_initialized: + dist.broadcast(latents, 0) + dist.barrier() + + + samples = self.decode_latents(latents) + if is_dist_initialized: + dist.barrier() + + # Convert to tensor + if output_type == "tensor": + samples = torch.from_numpy(samples) + + if not return_dict: + return samples + + return PipelineOutput(samples=samples)