style-aligned-controlnet

Runtime error

App Files Files Community

style-aligned-controlnet / pipeline_calls.py

ysharma HF staff

Upload 13 files

cfa3979 6 months ago

raw history blame contribute delete

No virus

24.7 kB

	# 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