# Copyright 2024 Marigold authors, PRS ETH Zurich. All rights reserved. # Copyright 2024 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # -------------------------------------------------------------------------- # More information and citation instructions are available on the # -------------------------------------------------------------------------- from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import torch from PIL import Image from tqdm.auto import tqdm from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection from diffusers.image_processor import PipelineImageInput from diffusers.models import ( AutoencoderKL, UNet2DConditionModel, ControlNetModel, ) from diffusers.schedulers import ( DDIMScheduler ) from diffusers.utils import ( BaseOutput, logging, replace_example_docstring, ) from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, logging, scale_lora_layers, unscale_lora_layers from diffusers.utils.torch_utils import randn_tensor from diffusers.pipelines.controlnet import StableDiffusionControlNetPipeline from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.pipelines.marigold.marigold_image_processing import MarigoldImageProcessor from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker import torch.nn.functional as F import pdb logger = logging.get_logger(__name__) # pylint: disable=invalid-name EXAMPLE_DOC_STRING = """ Examples: ```py >>> import diffusers >>> import torch >>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained( ... "prs-eth/marigold-normals-lcm-v0-1", variant="fp16", torch_dtype=torch.float16 ... ).to("cuda") >>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg") >>> normals = pipe(image) >>> vis = pipe.image_processor.visualize_normals(normals.prediction) >>> vis[0].save("einstein_normals.png") ``` """ @dataclass class StableNormalOutput(BaseOutput): """ Output class for Marigold monocular normals prediction pipeline. Args: prediction (`np.ndarray`, `torch.Tensor`): Predicted normals with values in the range [-1, 1]. The shape is always $numimages \times 3 \times height \times width$, regardless of whether the images were passed as a 4D array or a list. uncertainty (`None`, `np.ndarray`, `torch.Tensor`): Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages \times 1 \times height \times width$. latent (`None`, `torch.Tensor`): Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$. """ prediction: Union[np.ndarray, torch.Tensor] latent: Union[None, torch.Tensor] gaus_noise: Union[None, torch.Tensor] from einops import rearrange class DINOv2_Encoder(torch.nn.Module): IMAGENET_DEFAULT_MEAN = [0.485, 0.456, 0.406] IMAGENET_DEFAULT_STD = [0.229, 0.224, 0.225] def __init__( self, model_name = 'dinov2_vitl14', freeze = True, antialias=True, device="cuda", size = 448, ): super(DINOv2_Encoder, self).__init__() self.model = torch.hub.load('facebookresearch/dinov2', model_name) self.model.eval().to(device) self.device = device self.antialias = antialias self.dtype = torch.float32 self.mean = torch.Tensor(self.IMAGENET_DEFAULT_MEAN) self.std = torch.Tensor(self.IMAGENET_DEFAULT_STD) self.size = size if freeze: self.freeze() def freeze(self): for param in self.model.parameters(): param.requires_grad = False @torch.no_grad() def encoder(self, x): ''' x: [b h w c], range from (-1, 1), rbg ''' x = self.preprocess(x).to(self.device, self.dtype) b, c, h, w = x.shape patch_h, patch_w = h // 14, w // 14 embeddings = self.model.forward_features(x)['x_norm_patchtokens'] embeddings = rearrange(embeddings, 'b (h w) c -> b h w c', h = patch_h, w = patch_w) return rearrange(embeddings, 'b h w c -> b c h w') def preprocess(self, x): ''' x ''' # normalize to [0,1], x = torch.nn.functional.interpolate( x, size=(self.size, self.size), mode='bicubic', align_corners=True, antialias=self.antialias, ) x = (x + 1.0) / 2.0 # renormalize according to dino mean = self.mean.view(1, 3, 1, 1).to(x.device) std = self.std.view(1, 3, 1, 1).to(x.device) x = (x - mean) / std return x def to(self, device, dtype=None): if dtype is not None: self.dtype = dtype self.model.to(device, dtype) self.mean.to(device, dtype) self.std.to(device, dtype) else: self.model.to(device) self.mean.to(device) self.std.to(device) return self def __call__(self, x, **kwargs): return self.encoder(x, **kwargs) class StableNormalPipeline(StableDiffusionControlNetPipeline): """ Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io. Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance. This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods implemented for all pipelines (downloading, saving, running on a particular device, etc.). The pipeline also inherits the following loading methods: - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters Args: vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. text_encoder ([`~transformers.CLIPTextModel`]): Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)). tokenizer ([`~transformers.CLIPTokenizer`]): A `CLIPTokenizer` to tokenize text. unet ([`UNet2DConditionModel`]): A `UNet2DConditionModel` to denoise the encoded image latents. controlnet ([`ControlNetModel`] or `List[ControlNetModel]`): Provides additional conditioning to the `unet` during the denoising process. If you set multiple ControlNets as a list, the outputs from each ControlNet are added together to create one combined additional conditioning. scheduler ([`SchedulerMixin`]): A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details about a model's potential harms. feature_extractor ([`~transformers.CLIPImageProcessor`]): A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`. """ model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae" _optional_components = ["safety_checker", "feature_extractor", "image_encoder"] _exclude_from_cpu_offload = ["safety_checker"] _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] def __init__( self, vae: AutoencoderKL, text_encoder: CLIPTextModel, tokenizer: CLIPTokenizer, unet: UNet2DConditionModel, controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel]], dino_controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel]], scheduler: Union[DDIMScheduler], safety_checker: StableDiffusionSafetyChecker, feature_extractor: CLIPImageProcessor, image_encoder: CLIPVisionModelWithProjection = None, requires_safety_checker: bool = True, default_denoising_steps: Optional[int] = 10, default_processing_resolution: Optional[int] = 768, prompt="The normal map", empty_text_embedding=None, ): super().__init__( vae, text_encoder, tokenizer, unet, controlnet, scheduler, safety_checker, feature_extractor, image_encoder, requires_safety_checker, ) self.register_modules( dino_controlnet=dino_controlnet, ) self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor) self.dino_image_processor = lambda x: x / 127.5 -1. self.default_denoising_steps = default_denoising_steps self.default_processing_resolution = default_processing_resolution self.prompt = prompt self.prompt_embeds = None self.empty_text_embedding = empty_text_embedding self.prior = DINOv2_Encoder(size=672) def check_inputs( self, image: PipelineImageInput, num_inference_steps: int, ensemble_size: int, processing_resolution: int, resample_method_input: str, resample_method_output: str, batch_size: int, ensembling_kwargs: Optional[Dict[str, Any]], latents: Optional[torch.Tensor], generator: Optional[Union[torch.Generator, List[torch.Generator]]], output_type: str, output_uncertainty: bool, ) -> int: if num_inference_steps is None: raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.") if num_inference_steps < 1: raise ValueError("`num_inference_steps` must be positive.") if ensemble_size < 1: raise ValueError("`ensemble_size` must be positive.") if ensemble_size == 2: logger.warning( "`ensemble_size` == 2 results are similar to no ensembling (1); " "consider increasing the value to at least 3." ) if ensemble_size == 1 and output_uncertainty: raise ValueError( "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` " "greater than 1." ) if processing_resolution is None: raise ValueError( "`processing_resolution` is not specified and could not be resolved from the model config." ) if processing_resolution < 0: raise ValueError( "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for " "downsampled processing." ) if processing_resolution % self.vae_scale_factor != 0: raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.") if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): raise ValueError( "`resample_method_input` takes string values compatible with PIL library: " "nearest, nearest-exact, bilinear, bicubic, area." ) if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"): raise ValueError( "`resample_method_output` takes string values compatible with PIL library: " "nearest, nearest-exact, bilinear, bicubic, area." ) if batch_size < 1: raise ValueError("`batch_size` must be positive.") if output_type not in ["pt", "np"]: raise ValueError("`output_type` must be one of `pt` or `np`.") if latents is not None and generator is not None: raise ValueError("`latents` and `generator` cannot be used together.") if ensembling_kwargs is not None: if not isinstance(ensembling_kwargs, dict): raise ValueError("`ensembling_kwargs` must be a dictionary.") if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("closest", "mean"): raise ValueError("`ensembling_kwargs['reduction']` can be either `'closest'` or `'mean'`.") # image checks num_images = 0 W, H = None, None if not isinstance(image, list): image = [image] for i, img in enumerate(image): if isinstance(img, np.ndarray) or torch.is_tensor(img): if img.ndim not in (2, 3, 4): raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.") H_i, W_i = img.shape[-2:] N_i = 1 if img.ndim == 4: N_i = img.shape[0] elif isinstance(img, Image.Image): W_i, H_i = img.size N_i = 1 else: raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.") if W is None: W, H = W_i, H_i elif (W, H) != (W_i, H_i): raise ValueError( f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}" ) num_images += N_i # latents checks if latents is not None: if not torch.is_tensor(latents): raise ValueError("`latents` must be a torch.Tensor.") if latents.dim() != 4: raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.") if processing_resolution > 0: max_orig = max(H, W) new_H = H * processing_resolution // max_orig new_W = W * processing_resolution // max_orig if new_H == 0 or new_W == 0: raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]") W, H = new_W, new_H w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w) if latents.shape != shape_expected: raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.") # generator checks if generator is not None: if isinstance(generator, list): if len(generator) != num_images * ensemble_size: raise ValueError( "The number of generators must match the total number of ensemble members for all input images." ) if not all(g.device.type == generator[0].device.type for g in generator): raise ValueError("`generator` device placement is not consistent in the list.") elif not isinstance(generator, torch.Generator): raise ValueError(f"Unsupported generator type: {type(generator)}.") return num_images def progress_bar(self, iterable=None, total=None, desc=None, leave=True): if not hasattr(self, "_progress_bar_config"): self._progress_bar_config = {} elif not isinstance(self._progress_bar_config, dict): raise ValueError( f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}." ) progress_bar_config = dict(**self._progress_bar_config) progress_bar_config["desc"] = progress_bar_config.get("desc", desc) progress_bar_config["leave"] = progress_bar_config.get("leave", leave) if iterable is not None: return tqdm(iterable, **progress_bar_config) elif total is not None: return tqdm(total=total, **progress_bar_config) else: raise ValueError("Either `total` or `iterable` has to be defined.") @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, image: PipelineImageInput, prompt: Union[str, List[str]] = None, negative_prompt: Optional[Union[str, List[str]]] = None, num_inference_steps: Optional[int] = None, ensemble_size: int = 1, processing_resolution: Optional[int] = None, match_input_resolution: bool = True, resample_method_input: str = "bilinear", resample_method_output: str = "bilinear", batch_size: int = 1, ensembling_kwargs: Optional[Dict[str, Any]] = None, latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None, prompt_embeds: Optional[torch.Tensor] = None, negative_prompt_embeds: Optional[torch.Tensor] = None, num_images_per_prompt: Optional[int] = 1, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, controlnet_conditioning_scale: Union[float, List[float]] = 1.0, output_type: str = "np", output_uncertainty: bool = False, output_latent: bool = False, return_dict: bool = True, ): """ Function invoked when calling the pipeline. Args: image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`), `List[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the same width and height. num_inference_steps (`int`, *optional*, defaults to `None`): Number of denoising diffusion steps during inference. The default value `None` results in automatic selection. The number of steps should be at least 10 with the full Marigold models, and between 1 and 4 for Marigold-LCM models. ensemble_size (`int`, defaults to `1`): Number of ensemble predictions. Recommended values are 5 and higher for better precision, or 1 for faster inference. processing_resolution (`int`, *optional*, defaults to `None`): Effective processing resolution. When set to `0`, matches the larger input image dimension. This produces crisper predictions, but may also lead to the overall loss of global context. The default value `None` resolves to the optimal value from the model config. match_input_resolution (`bool`, *optional*, defaults to `True`): When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer side of the output will equal to `processing_resolution`. resample_method_input (`str`, *optional*, defaults to `"bilinear"`): Resampling method used to resize input images to `processing_resolution`. The accepted values are: `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. resample_method_output (`str`, *optional*, defaults to `"bilinear"`): Resampling method used to resize output predictions to match the input resolution. The accepted values are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. batch_size (`int`, *optional*, defaults to `1`): Batch size; only matters when setting `ensemble_size` or passing a tensor of images. ensembling_kwargs (`dict`, *optional*, defaults to `None`) Extra dictionary with arguments for precise ensembling control. The following options are available: - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in every pixel location, can be either `"closest"` or `"mean"`. latents (`torch.Tensor`, *optional*, defaults to `None`): Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`): Random number generator object to ensure reproducibility. output_type (`str`, *optional*, defaults to `"np"`): Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted values are: `"np"` (numpy array) or `"pt"` (torch tensor). output_uncertainty (`bool`, *optional*, defaults to `False`): When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that the `ensemble_size` argument is set to a value above 2. output_latent (`bool`, *optional*, defaults to `False`): When enabled, the output's `latent` field contains the latent codes corresponding to the predictions within the ensemble. These codes can be saved, modified, and used for subsequent calls with the `latents` argument. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple. Examples: Returns: [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`: If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a `tuple` is returned where the first element is the prediction, the second element is the uncertainty (or `None`), and the third is the latent (or `None`). """ # 0. Resolving variables. device = self._execution_device dtype = self.dtype # Model-specific optimal default values leading to fast and reasonable results. if num_inference_steps is None: num_inference_steps = self.default_denoising_steps if processing_resolution is None: processing_resolution = self.default_processing_resolution image, padding, original_resolution = self.image_processor.preprocess( image, processing_resolution, resample_method_input, device, dtype ) # [N,3,PPH,PPW] image_latent, gaus_noise = self.prepare_latents( image, latents, generator, ensemble_size, batch_size ) # [N,4,h,w], [N,4,h,w] # 0. X_start latent obtain predictor = self.x_start_pipeline(image, latents=gaus_noise, processing_resolution=processing_resolution, skip_preprocess=True) x_start_latent = predictor.latent # 1. Check inputs. num_images = self.check_inputs( image, num_inference_steps, ensemble_size, processing_resolution, resample_method_input, resample_method_output, batch_size, ensembling_kwargs, latents, generator, output_type, output_uncertainty, ) # 2. Prepare empty text conditioning. # Model invocation: self.tokenizer, self.text_encoder. if self.empty_text_embedding is None: prompt = "" text_inputs = self.tokenizer( prompt, padding="do_not_pad", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt", ) text_input_ids = text_inputs.input_ids.to(device) self.empty_text_embedding = self.text_encoder(text_input_ids)[0] # [1,2,1024] # 3. prepare prompt if self.prompt_embeds is None: prompt_embeds, negative_prompt_embeds = self.encode_prompt( self.prompt, device, num_images_per_prompt, False, negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=None, lora_scale=None, clip_skip=None, ) self.prompt_embeds = prompt_embeds self.negative_prompt_embeds = negative_prompt_embeds # 5. dino guider features obtaining ## TODO different case-1 dino_features = self.prior(image) dino_features = self.dino_controlnet.dino_controlnet_cond_embedding(dino_features) dino_features = self.match_noisy(dino_features, x_start_latent) del ( image, ) # 7. denoise sampling, using heuritic sampling proposed by Ye. t_start = self.x_start_pipeline.t_start self.scheduler.set_timesteps(num_inference_steps, t_start=t_start,device=device) cond_scale =controlnet_conditioning_scale pred_latent = x_start_latent cur_step = 0 # dino controlnet dino_down_block_res_samples, dino_mid_block_res_sample = self.dino_controlnet( dino_features.detach(), 0, # not depend on time steps encoder_hidden_states=self.prompt_embeds, conditioning_scale=cond_scale, guess_mode=False, return_dict=False, ) assert dino_mid_block_res_sample == None pred_latents = [] last_pred_latent = pred_latent for (t, prev_t) in self.progress_bar(zip(self.scheduler.timesteps,self.scheduler.prev_timesteps), leave=False, desc="Diffusion steps..."): _dino_down_block_res_samples = [dino_down_block_res_sample for dino_down_block_res_sample in dino_down_block_res_samples] # copy, avoid repeat quiery # controlnet down_block_res_samples, mid_block_res_sample = self.controlnet( image_latent.detach(), t, encoder_hidden_states=self.prompt_embeds, conditioning_scale=cond_scale, guess_mode=False, return_dict=False, ) # SG-DRN noise = self.dino_unet_forward( self.unet, pred_latent, t, encoder_hidden_states=self.prompt_embeds, down_block_additional_residuals=down_block_res_samples, mid_block_additional_residual=mid_block_res_sample, dino_down_block_additional_residuals= _dino_down_block_res_samples, return_dict=False, )[0] # [B,4,h,w] pred_latents.append(noise) # ddim steps out = self.scheduler.step( noise, t, prev_t, pred_latent, gaus_noise = gaus_noise, generator=generator, cur_step=cur_step+1 # NOTE that cur_step dirs to next_step )# [B,4,h,w] pred_latent = out.prev_sample cur_step += 1 del ( image_latent, dino_features, ) pred_latent = pred_latents[-1] # using x0 # decoder prediction = self.decode_prediction(pred_latent) prediction = self.image_processor.unpad_image(prediction, padding) # [N*E,3,PH,PW] prediction = self.image_processor.resize_antialias(prediction, original_resolution, resample_method_output, is_aa=False) # [N,3,H,W] if match_input_resolution: prediction = self.image_processor.resize_antialias( prediction, original_resolution, resample_method_output, is_aa=False ) # [N,3,H,W] if match_input_resolution: prediction = self.image_processor.resize_antialias( prediction, original_resolution, resample_method_output, is_aa=False ) # [N,3,H,W] prediction = self.normalize_normals(prediction) # [N,3,H,W] if output_type == "np": prediction = self.image_processor.pt_to_numpy(prediction) # [N,H,W,3] prediction = prediction.clip(min=-1, max=1) # 11. Offload all models self.maybe_free_model_hooks() return StableNormalOutput( prediction=prediction, latent=pred_latent, gaus_noise=gaus_noise ) # Copied from diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents def prepare_latents( self, image: torch.Tensor, latents: Optional[torch.Tensor], generator: Optional[torch.Generator], ensemble_size: int, batch_size: int, ) -> Tuple[torch.Tensor, torch.Tensor]: def retrieve_latents(encoder_output): if hasattr(encoder_output, "latent_dist"): return encoder_output.latent_dist.mode() elif hasattr(encoder_output, "latents"): return encoder_output.latents else: raise AttributeError("Could not access latents of provided encoder_output") image_latent = torch.cat( [ retrieve_latents(self.vae.encode(image[i : i + batch_size])) for i in range(0, image.shape[0], batch_size) ], dim=0, ) # [N,4,h,w] image_latent = image_latent * self.vae.config.scaling_factor image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # [N*E,4,h,w] pred_latent = latents if pred_latent is None: pred_latent = randn_tensor( image_latent.shape, generator=generator, device=image_latent.device, dtype=image_latent.dtype, ) # [N*E,4,h,w] return image_latent, pred_latent def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor: if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels: raise ValueError( f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}." ) prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # [B,3,H,W] return prediction # [B,3,H,W] @staticmethod def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: if normals.dim() != 4 or normals.shape[1] != 3: raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.") norm = torch.norm(normals, dim=1, keepdim=True) normals /= norm.clamp(min=eps) return normals @staticmethod def match_noisy(dino, noisy): _, __, dino_h, dino_w = dino.shape _, __, h, w = noisy.shape if h == dino_h and w == dino_w: return dino else: return F.interpolate(dino, (h, w), mode='bilinear') @staticmethod def dino_unet_forward( self, # NOTE that repurpose to UNet sample: torch.Tensor, timestep: Union[torch.Tensor, float, int], encoder_hidden_states: torch.Tensor, class_labels: Optional[torch.Tensor] = None, timestep_cond: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, cross_attention_kwargs: Optional[Dict[str, Any]] = None, added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, mid_block_additional_residual: Optional[torch.Tensor] = None, dino_down_block_additional_residuals: Optional[torch.Tensor] = None, down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, encoder_attention_mask: Optional[torch.Tensor] = None, return_dict: bool = True, ) -> Union[UNet2DConditionOutput, Tuple]: r""" The [`UNet2DConditionModel`] forward method. Args: sample (`torch.Tensor`): The noisy input tensor with the following shape `(batch, channel, height, width)`. timestep (`torch.Tensor` or `float` or `int`): The number of timesteps to denoise an input. encoder_hidden_states (`torch.Tensor`): The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. class_labels (`torch.Tensor`, *optional*, defaults to `None`): Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed through the `self.time_embedding` layer to obtain the timestep embeddings. attention_mask (`torch.Tensor`, *optional*, defaults to `None`): An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to "discard" tokens. cross_attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). added_cond_kwargs: (`dict`, *optional*): A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks. down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): A tuple of tensors that if specified are added to the residuals of down unet blocks. mid_block_additional_residual: (`torch.Tensor`, *optional*): A tensor that if specified is added to the residual of the middle unet block. down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*): additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s) encoder_attention_mask (`torch.Tensor`): A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to "discard" tokens. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. Returns: [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: If `return_dict` is True, an [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise a `tuple` is returned where the first element is the sample tensor. """ # By default samples have to be AT least a multiple of the overall upsampling factor. # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). # However, the upsampling interpolation output size can be forced to fit any upsampling size # on the fly if necessary. default_overall_up_factor = 2**self.num_upsamplers # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` forward_upsample_size = False upsample_size = None for dim in sample.shape[-2:]: if dim % default_overall_up_factor != 0: # Forward upsample size to force interpolation output size. forward_upsample_size = True break # ensure attention_mask is a bias, and give it a singleton query_tokens dimension # expects mask of shape: # [batch, key_tokens] # adds singleton query_tokens dimension: # [batch, 1, key_tokens] # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) if attention_mask is not None: # assume that mask is expressed as: # (1 = keep, 0 = discard) # convert mask into a bias that can be added to attention scores: # (keep = +0, discard = -10000.0) attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 attention_mask = attention_mask.unsqueeze(1) # convert encoder_attention_mask to a bias the same way we do for attention_mask if encoder_attention_mask is not None: encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0 encoder_attention_mask = encoder_attention_mask.unsqueeze(1) # 0. center input if necessary if self.config.center_input_sample: sample = 2 * sample - 1.0 # 1. time t_emb = self.get_time_embed(sample=sample, timestep=timestep) emb = self.time_embedding(t_emb, timestep_cond) aug_emb = None class_emb = self.get_class_embed(sample=sample, class_labels=class_labels) if class_emb is not None: if self.config.class_embeddings_concat: emb = torch.cat([emb, class_emb], dim=-1) else: emb = emb + class_emb aug_emb = self.get_aug_embed( emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs ) if self.config.addition_embed_type == "image_hint": aug_emb, hint = aug_emb sample = torch.cat([sample, hint], dim=1) emb = emb + aug_emb if aug_emb is not None else emb if self.time_embed_act is not None: emb = self.time_embed_act(emb) encoder_hidden_states = self.process_encoder_hidden_states( encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs ) # 2. pre-process sample = self.conv_in(sample) # 2.5 GLIGEN position net if cross_attention_kwargs is not None and cross_attention_kwargs.get("gligen", None) is not None: cross_attention_kwargs = cross_attention_kwargs.copy() gligen_args = cross_attention_kwargs.pop("gligen") cross_attention_kwargs["gligen"] = {"objs": self.position_net(**gligen_args)} # 3. down # we're popping the `scale` instead of getting it because otherwise `scale` will be propagated # to the internal blocks and will raise deprecation warnings. this will be confusing for our users. if cross_attention_kwargs is not None: cross_attention_kwargs = cross_attention_kwargs.copy() lora_scale = cross_attention_kwargs.pop("scale", 1.0) else: lora_scale = 1.0 if USE_PEFT_BACKEND: # weight the lora layers by setting `lora_scale` for each PEFT layer scale_lora_layers(self, lora_scale) is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets is_adapter = down_intrablock_additional_residuals is not None # maintain backward compatibility for legacy usage, where # T2I-Adapter and ControlNet both use down_block_additional_residuals arg # but can only use one or the other if not is_adapter and mid_block_additional_residual is None and down_block_additional_residuals is not None: deprecate( "T2I should not use down_block_additional_residuals", "1.3.0", "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \ and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \ for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ", standard_warn=False, ) down_intrablock_additional_residuals = down_block_additional_residuals is_adapter = True def residual_downforward( self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, additional_residuals: Optional[torch.Tensor] = None, *args, **kwargs, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: if len(args) > 0 or kwargs.get("scale", None) is not None: deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." deprecate("scale", "1.0.0", deprecation_message) output_states = () for resnet in self.resnets: if self.training and self.gradient_checkpointing: def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs) return custom_forward if is_torch_version(">=", "1.11.0"): hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb, use_reentrant=False ) else: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb ) else: hidden_states = resnet(hidden_states, temb) hidden_states += additional_residuals.pop(0) output_states = output_states + (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) hidden_states += additional_residuals.pop(0) output_states = output_states + (hidden_states,) return hidden_states, output_states def residual_blockforward( self, ## NOTE that repurpose to unet_blocks hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None, encoder_hidden_states: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, cross_attention_kwargs: Optional[Dict[str, Any]] = None, encoder_attention_mask: Optional[torch.Tensor] = None, additional_residuals: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]: if cross_attention_kwargs is not None: if cross_attention_kwargs.get("scale", None) is not None: logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") output_states = () blocks = list(zip(self.resnets, self.attentions)) for i, (resnet, attn) in enumerate(blocks): if self.training and self.gradient_checkpointing: def create_custom_forward(module, return_dict=None): def custom_forward(*inputs): if return_dict is not None: return module(*inputs, return_dict=return_dict) else: return module(*inputs) return custom_forward ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs, ) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] else: hidden_states = resnet(hidden_states, temb) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] hidden_states += additional_residuals.pop(0) output_states = output_states + (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) hidden_states += additional_residuals.pop(0) output_states = output_states + (hidden_states,) return hidden_states, output_states down_intrablock_additional_residuals = dino_down_block_additional_residuals sample += down_intrablock_additional_residuals.pop(0) down_block_res_samples = (sample,) for downsample_block in self.down_blocks: if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: sample, res_samples = residual_blockforward( downsample_block, hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask, additional_residuals = down_intrablock_additional_residuals, ) else: sample, res_samples = residual_downforward( downsample_block, hidden_states=sample, temb=emb, additional_residuals = down_intrablock_additional_residuals, ) down_block_res_samples += res_samples if is_controlnet: new_down_block_res_samples = () for down_block_res_sample, down_block_additional_residual in zip( down_block_res_samples, down_block_additional_residuals ): down_block_res_sample = down_block_res_sample + down_block_additional_residual new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,) down_block_res_samples = new_down_block_res_samples # 4. mid if self.mid_block is not None: if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention: sample = self.mid_block( sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask, ) else: sample = self.mid_block(sample, emb) # To support T2I-Adapter-XL if ( is_adapter and len(down_intrablock_additional_residuals) > 0 and sample.shape == down_intrablock_additional_residuals[0].shape ): sample += down_intrablock_additional_residuals.pop(0) if is_controlnet: sample = sample + mid_block_additional_residual # 5. up for i, upsample_block in enumerate(self.up_blocks): is_final_block = i == len(self.up_blocks) - 1 res_samples = down_block_res_samples[-len(upsample_block.resnets) :] down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)] # if we have not reached the final block and need to forward the # upsample size, we do it here if not is_final_block and forward_upsample_size: upsample_size = down_block_res_samples[-1].shape[2:] if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: sample = upsample_block( hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, upsample_size=upsample_size, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, ) else: sample = upsample_block( hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size, ) # 6. post-process if self.conv_norm_out: sample = self.conv_norm_out(sample) sample = self.conv_act(sample) sample = self.conv_out(sample) if USE_PEFT_BACKEND: # remove `lora_scale` from each PEFT layer unscale_lora_layers(self, lora_scale) if not return_dict: return (sample,) return UNet2DConditionOutput(sample=sample) @staticmethod def ensemble_normals( normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest" ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """ Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`. Args: normals (`torch.Tensor`): Input ensemble normals maps. output_uncertainty (`bool`, *optional*, defaults to `False`): Whether to output uncertainty map. reduction (`str`, *optional*, defaults to `"closest"`): Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and `"mean"`. Returns: A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of uncertainties of shape `(1, 1, H, W)`. """ if normals.dim() != 4 or normals.shape[1] != 3: raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.") if reduction not in ("closest", "mean"): raise ValueError(f"Unrecognized reduction method: {reduction}.") mean_normals = normals.mean(dim=0, keepdim=True) # [1,3,H,W] mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals) # [1,3,H,W] sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True) # [E,1,H,W] sim_cos = sim_cos.clamp(-1, 1) # required to avoid NaN in uncertainty with fp16 uncertainty = None if output_uncertainty: uncertainty = sim_cos.arccos() # [E,1,H,W] uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi # [1,1,H,W] if reduction == "mean": return mean_normals, uncertainty # [1,3,H,W], [1,1,H,W] closest_indices = sim_cos.argmax(dim=0, keepdim=True) # [1,1,H,W] closest_indices = closest_indices.repeat(1, 3, 1, 1) # [1,3,H,W] closest_normals = torch.gather(normals, 0, closest_indices) # [1,3,H,W] return closest_normals, uncertainty # [1,3,H,W], [1,1,H,W] # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, **kwargs, ): """ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. Args: scheduler (`SchedulerMixin`): The scheduler to get timesteps from. num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. timesteps (`List[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. sigmas (`List[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") if timesteps is not None: accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif sigmas is not None: accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if not accept_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigmas schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps