# 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