# Copyright 2023 Google LLC # # 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. from __future__ import annotations from typing import Any import torch import numpy as np from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline from diffusers.image_processor import PipelineImageInput from diffusers.utils.torch_utils import is_compiled_module, is_torch_version from transformers import DPTImageProcessor, DPTForDepthEstimation from diffusers import StableDiffusionPanoramaPipeline from PIL import Image import copy T = torch.Tensor TN = T | None def get_depth_map(image: Image, feature_processor: DPTImageProcessor, depth_estimator: DPTForDepthEstimation) -> Image: image = feature_processor(images=image, return_tensors="pt").pixel_values.to("cuda") with torch.no_grad(), torch.autocast("cuda"): depth_map = depth_estimator(image).predicted_depth depth_map = torch.nn.functional.interpolate( depth_map.unsqueeze(1), size=(1024, 1024), mode="bicubic", align_corners=False, ) depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True) depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True) depth_map = (depth_map - depth_min) / (depth_max - depth_min) image = torch.cat([depth_map] * 3, dim=1) image = image.permute(0, 2, 3, 1).cpu().numpy()[0] image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8)) return image def concat_zero_control(control_reisduel: T) -> T: b = control_reisduel.shape[0] // 2 zerso_reisduel = torch.zeros_like(control_reisduel[0:1]) return torch.cat((zerso_reisduel, control_reisduel[:b], zerso_reisduel, control_reisduel[b::])) @torch.no_grad() def controlnet_call( pipeline: StableDiffusionXLControlNetPipeline, prompt: str | list[str] = None, prompt_2: str | list[str] | None = None, image: PipelineImageInput = None, height: int | None = None, width: int | None = None, num_inference_steps: int = 50, guidance_scale: float = 5.0, negative_prompt: str | list[str] | None = None, negative_prompt_2: str | list[str] | None = None, num_images_per_prompt: int = 1, eta: float = 0.0, generator: torch.Generator | None = None, latents: TN = None, prompt_embeds: TN = None, negative_prompt_embeds: TN = None, pooled_prompt_embeds: TN = None, negative_pooled_prompt_embeds: TN = None, cross_attention_kwargs: dict[str, Any] | None = None, controlnet_conditioning_scale: float | list[float] = 1.0, control_guidance_start: float | list[float] = 0.0, control_guidance_end: float | list[float] = 1.0, original_size: tuple[int, int] = None, crops_coords_top_left: tuple[int, int] = (0, 0), target_size: tuple[int, int] | None = None, negative_original_size: tuple[int, int] | None = None, negative_crops_coords_top_left: tuple[int, int] = (0, 0), negative_target_size:tuple[int, int] | None = None, clip_skip: int | None = None, ) -> list[Image]: controlnet = pipeline.controlnet._orig_mod if is_compiled_module(pipeline.controlnet) else pipeline.controlnet # align format for control guidance if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): control_guidance_start = len(control_guidance_end) * [control_guidance_start] elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): control_guidance_end = len(control_guidance_start) * [control_guidance_end] elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): mult = 1 control_guidance_start, control_guidance_end = ( mult * [control_guidance_start], mult * [control_guidance_end], ) # 1. Check inputs. Raise error if not correct pipeline.check_inputs( prompt, prompt_2, image, 1, negative_prompt, negative_prompt_2, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, controlnet_conditioning_scale, control_guidance_start, control_guidance_end, ) pipeline._guidance_scale = guidance_scale # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] device = pipeline._execution_device # 3. Encode input prompt text_encoder_lora_scale = ( cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None ) ( prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, ) = pipeline.encode_prompt( prompt, prompt_2, device, 1, True, negative_prompt, negative_prompt_2, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, lora_scale=text_encoder_lora_scale, clip_skip=clip_skip, ) # 4. Prepare image if isinstance(controlnet, ControlNetModel): image = pipeline.prepare_image( image=image, width=width, height=height, batch_size=1, num_images_per_prompt=1, device=device, dtype=controlnet.dtype, do_classifier_free_guidance=True, guess_mode=False, ) height, width = image.shape[-2:] image = torch.stack([image[0]] * num_images_per_prompt + [image[1]] * num_images_per_prompt) else: assert False # 5. Prepare timesteps pipeline.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = pipeline.scheduler.timesteps # 6. Prepare latent variables num_channels_latents = pipeline.unet.config.in_channels latents = pipeline.prepare_latents( 1 + num_images_per_prompt, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, latents, ) # 6.5 Optionally get Guidance Scale Embedding timestep_cond = None # 7. Prepare extra step kwargs. extra_step_kwargs = pipeline.prepare_extra_step_kwargs(generator, eta) # 7.1 Create tensor stating which controlnets to keep controlnet_keep = [] for i in range(len(timesteps)): keeps = [ 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) for s, e in zip(control_guidance_start, control_guidance_end) ] controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) # 7.2 Prepare added time ids & embeddings if isinstance(image, list): original_size = original_size or image[0].shape[-2:] else: original_size = original_size or image.shape[-2:] target_size = target_size or (height, width) add_text_embeds = pooled_prompt_embeds if pipeline.text_encoder_2 is None: text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) else: text_encoder_projection_dim = pipeline.text_encoder_2.config.projection_dim add_time_ids = pipeline._get_add_time_ids( original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype, text_encoder_projection_dim=text_encoder_projection_dim, ) if negative_original_size is not None and negative_target_size is not None: negative_add_time_ids = pipeline._get_add_time_ids( negative_original_size, negative_crops_coords_top_left, negative_target_size, dtype=prompt_embeds.dtype, text_encoder_projection_dim=text_encoder_projection_dim, ) else: negative_add_time_ids = add_time_ids prompt_embeds = torch.stack([prompt_embeds[0]] + [prompt_embeds[1]] * num_images_per_prompt) negative_prompt_embeds = torch.stack([negative_prompt_embeds[0]] + [negative_prompt_embeds[1]] * num_images_per_prompt) negative_pooled_prompt_embeds = torch.stack([negative_pooled_prompt_embeds[0]] + [negative_pooled_prompt_embeds[1]] * num_images_per_prompt) add_text_embeds = torch.stack([add_text_embeds[0]] + [add_text_embeds[1]] * num_images_per_prompt) prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0) prompt_embeds = prompt_embeds.to(device) add_text_embeds = add_text_embeds.to(device) add_time_ids = add_time_ids.to(device).repeat(1 + num_images_per_prompt, 1) batch_size = num_images_per_prompt + 1 # 8. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * pipeline.scheduler.order is_unet_compiled = is_compiled_module(pipeline.unet) is_controlnet_compiled = is_compiled_module(pipeline.controlnet) is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} controlnet_prompt_embeds = torch.cat((prompt_embeds[1:batch_size], prompt_embeds[1:batch_size])) controlnet_added_cond_kwargs = {key: torch.cat((item[1:batch_size,], item[1:batch_size])) for key, item in added_cond_kwargs.items()} with pipeline.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): # Relevant thread: # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1: torch._inductor.cudagraph_mark_step_begin() # expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t) # controlnet(s) inference control_model_input = torch.cat((latent_model_input[1:batch_size,], latent_model_input[batch_size+1:])) if isinstance(controlnet_keep[i], list): cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] else: controlnet_cond_scale = controlnet_conditioning_scale if isinstance(controlnet_cond_scale, list): controlnet_cond_scale = controlnet_cond_scale[0] cond_scale = controlnet_cond_scale * controlnet_keep[i] if cond_scale > 0: down_block_res_samples, mid_block_res_sample = pipeline.controlnet( control_model_input, t, encoder_hidden_states=controlnet_prompt_embeds, controlnet_cond=image, conditioning_scale=cond_scale, guess_mode=False, added_cond_kwargs=controlnet_added_cond_kwargs, return_dict=False, ) mid_block_res_sample = concat_zero_control(mid_block_res_sample) down_block_res_samples = [concat_zero_control(down_block_res_sample) for down_block_res_sample in down_block_res_samples] else: mid_block_res_sample = down_block_res_samples = None # predict the noise residual noise_pred = pipeline.unet( latent_model_input, t, encoder_hidden_states=prompt_embeds, timestep_cond=timestep_cond, cross_attention_kwargs=cross_attention_kwargs, down_block_additional_residuals=down_block_res_samples, mid_block_additional_residual=mid_block_res_sample, added_cond_kwargs=added_cond_kwargs, return_dict=False, )[0] # perform 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 = pipeline.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipeline.scheduler.order == 0): progress_bar.update() # manually for max memory savings if pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast: pipeline.upcast_vae() latents = latents.to(next(iter(pipeline.vae.post_quant_conv.parameters())).dtype) # make sure the VAE is in float32 mode, as it overflows in float16 needs_upcasting = pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast if needs_upcasting: pipeline.upcast_vae() latents = latents.to(next(iter(pipeline.vae.post_quant_conv.parameters())).dtype) image = pipeline.vae.decode(latents / pipeline.vae.config.scaling_factor, return_dict=False)[0] # cast back to fp16 if needed if needs_upcasting: pipeline.vae.to(dtype=torch.float16) if pipeline.watermark is not None: image = pipeline.watermark.apply_watermark(image) image = pipeline.image_processor.postprocess(image, output_type='pil') # Offload all models pipeline.maybe_free_model_hooks() return image @torch.no_grad() def panorama_call( pipeline: StableDiffusionPanoramaPipeline, prompt: list[str], height: int | None = 512, width: int | None = 2048, num_inference_steps: int = 50, guidance_scale: float = 7.5, view_batch_size: int = 1, negative_prompt: str | list[str] | None = None, num_images_per_prompt: int | None = 1, eta: float = 0.0, generator: torch.Generator | None = None, reference_latent: TN = None, latents: TN = None, prompt_embeds: TN = None, negative_prompt_embeds: TN = None, cross_attention_kwargs: dict[str, Any] | None = None, circular_padding: bool = False, clip_skip: int | None = None, stride=8 ) -> list[Image]: # 0. Default height and width to unet height = height or pipeline.unet.config.sample_size * pipeline.vae_scale_factor width = width or pipeline.unet.config.sample_size * pipeline.vae_scale_factor # 1. Check inputs. Raise error if not correct pipeline.check_inputs( prompt, height, width, 1, negative_prompt, prompt_embeds, negative_prompt_embeds ) device = pipeline._execution_device # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 # 3. Encode input prompt text_encoder_lora_scale = ( cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None ) prompt_embeds, negative_prompt_embeds = pipeline.encode_prompt( prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, lora_scale=text_encoder_lora_scale, clip_skip=clip_skip, ) # 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 # 4. Prepare timesteps pipeline.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = pipeline.scheduler.timesteps # 5. Prepare latent variables num_channels_latents = pipeline.unet.config.in_channels latents = pipeline.prepare_latents( 1, num_channels_latents, height, width, prompt_embeds.dtype, device, generator, latents, ) if reference_latent is None: reference_latent = torch.randn(1, 4, pipeline.unet.config.sample_size, pipeline.unet.config.sample_size, generator=generator) reference_latent = reference_latent.to(device=device, dtype=pipeline.unet.dtype) # 6. Define panorama grid and initialize views for synthesis. # prepare batch grid views = pipeline.get_views(height, width, circular_padding=circular_padding, stride=stride) views_batch = [views[i: i + view_batch_size] for i in range(0, len(views), view_batch_size)] views_scheduler_status = [copy.deepcopy(pipeline.scheduler.__dict__)] * len(views_batch) count = torch.zeros_like(latents) value = torch.zeros_like(latents) # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = pipeline.prepare_extra_step_kwargs(generator, eta) # 8. Denoising loop # Each denoising step also includes refinement of the latents with respect to the # views. num_warmup_steps = len(timesteps) - num_inference_steps * pipeline.scheduler.order negative_prompt_embeds = torch.cat([negative_prompt_embeds[:1], *[negative_prompt_embeds[1:]] * view_batch_size] ) prompt_embeds = torch.cat([prompt_embeds[:1], *[prompt_embeds[1:]] * view_batch_size] ) with pipeline.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): count.zero_() value.zero_() # generate views # Here, we iterate through different spatial crops of the latents and denoise them. These # denoised (latent) crops are then averaged to produce the final latent # for the current timestep via MultiDiffusion. Please see Sec. 4.1 in the # MultiDiffusion paper for more details: https://arxiv.org/abs/2302.08113 # Batch views denoise for j, batch_view in enumerate(views_batch): vb_size = len(batch_view) # get the latents corresponding to the current view coordinates if circular_padding: latents_for_view = [] for h_start, h_end, w_start, w_end in batch_view: if w_end > latents.shape[3]: # Add circular horizontal padding latent_view = torch.cat( ( latents[:, :, h_start:h_end, w_start:], latents[:, :, h_start:h_end, : w_end - latents.shape[3]], ), dim=-1, ) else: latent_view = latents[:, :, h_start:h_end, w_start:w_end] latents_for_view.append(latent_view) latents_for_view = torch.cat(latents_for_view) else: latents_for_view = torch.cat( [ latents[:, :, h_start:h_end, w_start:w_end] for h_start, h_end, w_start, w_end in batch_view ] ) # rematch block's scheduler status pipeline.scheduler.__dict__.update(views_scheduler_status[j]) # expand the latents if we are doing classifier free guidance latent_reference_plus_view = torch.cat((reference_latent, latents_for_view)) latent_model_input = latent_reference_plus_view.repeat(2, 1, 1, 1) prompt_embeds_input = torch.cat([negative_prompt_embeds[: 1 + vb_size], prompt_embeds[: 1 + vb_size]] ) latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t) # predict the noise residual # return noise_pred = pipeline.unet( latent_model_input, t, encoder_hidden_states=prompt_embeds_input, cross_attention_kwargs=cross_attention_kwargs, ).sample # perform 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 latent_reference_plus_view = pipeline.scheduler.step( noise_pred, t, latent_reference_plus_view, **extra_step_kwargs ).prev_sample if j == len(views_batch) - 1: reference_latent = latent_reference_plus_view[:1] latents_denoised_batch = latent_reference_plus_view[1:] # save views scheduler status after sample views_scheduler_status[j] = copy.deepcopy(pipeline.scheduler.__dict__) # extract value from batch for latents_view_denoised, (h_start, h_end, w_start, w_end) in zip( latents_denoised_batch.chunk(vb_size), batch_view ): if circular_padding and w_end > latents.shape[3]: # Case for circular padding value[:, :, h_start:h_end, w_start:] += latents_view_denoised[ :, :, h_start:h_end, : latents.shape[3] - w_start ] value[:, :, h_start:h_end, : w_end - latents.shape[3]] += latents_view_denoised[ :, :, h_start:h_end, latents.shape[3] - w_start: ] count[:, :, h_start:h_end, w_start:] += 1 count[:, :, h_start:h_end, : w_end - latents.shape[3]] += 1 else: value[:, :, h_start:h_end, w_start:w_end] += latents_view_denoised count[:, :, h_start:h_end, w_start:w_end] += 1 # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113 latents = torch.where(count > 0, value / count, value) # call the callback, if provided if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipeline.scheduler.order == 0): progress_bar.update() if circular_padding: image = pipeline.decode_latents_with_padding(latents) else: image = pipeline.vae.decode(latents / pipeline.vae.config.scaling_factor, return_dict=False)[0] reference_image = pipeline.vae.decode(reference_latent / pipeline.vae.config.scaling_factor, return_dict=False)[0] # image, has_nsfw_concept = pipeline.run_safety_checker(image, device, prompt_embeds.dtype) # reference_image, _ = pipeline.run_safety_checker(reference_image, device, prompt_embeds.dtype) image = pipeline.image_processor.postprocess(image, output_type='pil', do_denormalize=[True]) reference_image = pipeline.image_processor.postprocess(reference_image, output_type='pil', do_denormalize=[True]) pipeline.maybe_free_model_hooks() return reference_image + image