






"""simple docstring"""



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  str ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			@property
			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Any									=1
										lowerCamelCase_	:							str									=3
										lowerCamelCase_	:							Any									=(32, 32)

										lowerCamelCase_	:							List[Any]									=floats_tensor((batch_size, num_channels) + sizes       ,	rng=random.Random(0 ) ).to(snake_case__ )
										return image


			@property
			def      UpperCAmelCase__						(				self					:  Tuple ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Union[str, Any]									=UNetaDConditionModel(
										    block_out_channels=(32, 32, 64)       ,	layers_per_block=2       ,	sample_size=32       ,	in_channels=7       ,	out_channels=4       ,	down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")       ,	up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")       ,	cross_attention_dim=32       ,	attention_head_dim=8       ,	use_linear_projection=snake_case__       ,	only_cross_attention=(True, True, False)       ,	num_class_embeds=100       ,	)
										return model


			@property
			def      UpperCAmelCase__						(				self					:  List[Any] ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							List[str]									=AutoencoderKL(
										    block_out_channels=[32, 32, 64]       ,	in_channels=3       ,	out_channels=3       ,	down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]       ,	up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]       ,	latent_channels=4       ,	)
										return model


			@property
			def      UpperCAmelCase__						(				self					:  int ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							Dict									=CLIPTextConfig(
										    bos_token_id=0       ,	eos_token_id=2       ,	hidden_size=32       ,	intermediate_size=37       ,	layer_norm_eps=1E-05       ,	num_attention_heads=4       ,	num_hidden_layers=5       ,	pad_token_id=1       ,	vocab_size=1000       ,	hidden_act="gelu"       ,	projection_dim=512       ,	)
										return CLIPTextModel(snake_case__ )


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Any									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							Optional[int]									=self.dummy_cond_unet_upscale
										lowerCamelCase_	:							Dict									=DDPMScheduler()
										lowerCamelCase_	:							Tuple									=DDIMScheduler(prediction_type="v_prediction" )
										lowerCamelCase_	:							str									=self.dummy_vae
										lowerCamelCase_	:							Any									=self.dummy_text_encoder
										lowerCamelCase_	:							Optional[int]									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

										lowerCamelCase_	:							Union[str, Any]									=self.dummy_image.cpu().permute(0       ,	2       ,	3       ,	1 )[0]
										lowerCamelCase_	:							Any									=Image.fromarray(np.uinta(snake_case__ ) ).convert("RGB" ).resize((64, 64) )

										# make sure here that pndm scheduler skips prk
										lowerCamelCase_	:							str									=StableDiffusionUpscalePipeline(
										    unet=snake_case__       ,	low_res_scheduler=snake_case__       ,	scheduler=snake_case__       ,	vae=snake_case__       ,	text_encoder=snake_case__       ,	tokenizer=snake_case__       ,	max_noise_level=350       ,	)
										lowerCamelCase_	:							Any									=sd_pipe.to(snake_case__ )
										sd_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							List[Any]									="A painting of a squirrel eating a burger"
										lowerCamelCase_	:							Dict									=torch.Generator(device=snake_case__ ).manual_seed(0 )
										lowerCamelCase_	:							Optional[Any]									=sd_pipe(
										    [prompt]       ,	image=snake_case__       ,	generator=snake_case__       ,	guidance_scale=6.0       ,	noise_level=20       ,	num_inference_steps=2       ,	output_type="np"       ,	)

										lowerCamelCase_	:							List[str]									=output.images

										lowerCamelCase_	:							int									=torch.Generator(device=snake_case__ ).manual_seed(0 )
										lowerCamelCase_	:							int									=sd_pipe(
										    [prompt]       ,	image=snake_case__       ,	generator=snake_case__       ,	guidance_scale=6.0       ,	noise_level=20       ,	num_inference_steps=2       ,	output_type="np"       ,	return_dict=snake_case__       ,	)[0]

										lowerCamelCase_	:							str									=image[0, -3:, -3:, -1]
										lowerCamelCase_	:							List[str]									=image_from_tuple[0, -3:, -3:, -1]

										lowerCamelCase_	:							Union[str, Any]									=low_res_image.size[0] * 4
										assert image.shape == (1, expected_height_width, expected_height_width, 3)
										lowerCamelCase_	:							int									=np.array([0.3_113, 0.3_910, 0.4_272, 0.4_859, 0.5_061, 0.4_652, 0.5_362, 0.5_715, 0.5_661] )

										assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2


			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							Tuple									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							Any									=self.dummy_cond_unet_upscale
										lowerCamelCase_	:							int									=DDPMScheduler()
										lowerCamelCase_	:							str									=DDIMScheduler(prediction_type="v_prediction" )
										lowerCamelCase_	:							Union[str, Any]									=self.dummy_vae
										lowerCamelCase_	:							List[Any]									=self.dummy_text_encoder
										lowerCamelCase_	:							Optional[int]									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

										lowerCamelCase_	:							List[str]									=self.dummy_image.cpu().permute(0       ,	2       ,	3       ,	1 )[0]
										lowerCamelCase_	:							Any									=Image.fromarray(np.uinta(snake_case__ ) ).convert("RGB" ).resize((64, 64) )

										# make sure here that pndm scheduler skips prk
										lowerCamelCase_	:							Union[str, Any]									=StableDiffusionUpscalePipeline(
										    unet=snake_case__       ,	low_res_scheduler=snake_case__       ,	scheduler=snake_case__       ,	vae=snake_case__       ,	text_encoder=snake_case__       ,	tokenizer=snake_case__       ,	max_noise_level=350       ,	)
										lowerCamelCase_	:							Optional[Any]									=sd_pipe.to(snake_case__ )
										sd_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Any									="A painting of a squirrel eating a burger"
										lowerCamelCase_	:							Union[str, Any]									=sd_pipe(
										    2 * [prompt]       ,	image=2 * [low_res_image]       ,	guidance_scale=6.0       ,	noise_level=20       ,	num_inference_steps=2       ,	output_type="np"       ,	)
										lowerCamelCase_	:							int									=output.images
										assert image.shape[0] == 2

										lowerCamelCase_	:							List[str]									=torch.Generator(device=snake_case__ ).manual_seed(0 )
										lowerCamelCase_	:							Optional[Any]									=sd_pipe(
										    [prompt]       ,	image=snake_case__       ,	generator=snake_case__       ,	num_images_per_prompt=2       ,	guidance_scale=6.0       ,	noise_level=20       ,	num_inference_steps=2       ,	output_type="np"       ,	)
										lowerCamelCase_	:							Union[str, Any]									=output.images
										assert image.shape[0] == 2


			@unittest.skipIf(torch_device != "cuda"       ,	"This test requires a GPU" )
			def      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							List[Any]									=self.dummy_cond_unet_upscale
										lowerCamelCase_	:							List[str]									=DDPMScheduler()
										lowerCamelCase_	:							int									=DDIMScheduler(prediction_type="v_prediction" )
										lowerCamelCase_	:							Optional[int]									=self.dummy_vae
										lowerCamelCase_	:							List[Any]									=self.dummy_text_encoder
										lowerCamelCase_	:							Union[str, Any]									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

										lowerCamelCase_	:							List[str]									=self.dummy_image.cpu().permute(0       ,	2       ,	3       ,	1 )[0]
										lowerCamelCase_	:							Union[str, Any]									=Image.fromarray(np.uinta(snake_case__ ) ).convert("RGB" ).resize((64, 64) )

										# put models in fp16, except vae as it overflows in fp16
										lowerCamelCase_	:							Tuple									=unet.half()
										lowerCamelCase_	:							str									=text_encoder.half()

										# make sure here that pndm scheduler skips prk
										lowerCamelCase_	:							int									=StableDiffusionUpscalePipeline(
										    unet=snake_case__       ,	low_res_scheduler=snake_case__       ,	scheduler=snake_case__       ,	vae=snake_case__       ,	text_encoder=snake_case__       ,	tokenizer=snake_case__       ,	max_noise_level=350       ,	)
										lowerCamelCase_	:							Optional[Any]									=sd_pipe.to(snake_case__ )
										sd_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Tuple									="A painting of a squirrel eating a burger"
										lowerCamelCase_	:							Optional[int]									=torch.manual_seed(0 )
										lowerCamelCase_	:							int									=sd_pipe(
										    [prompt]       ,	image=snake_case__       ,	generator=snake_case__       ,	num_inference_steps=2       ,	output_type="np"       ,	).images

										lowerCamelCase_	:							Dict									=low_res_image.size[0] * 4
										assert image.shape == (1, expected_height_width, expected_height_width, 3)




@slow
@require_torch_gpu
class        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  int ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							List[Any]									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/sd2-upscale/low_res_cat.png" )
										lowerCamelCase_	:							List[str]									=load_numpy(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
										    "/upsampled_cat.npy" )

										lowerCamelCase_	:							List[str]									="stabilityai/stable-diffusion-x4-upscaler"
										lowerCamelCase_	:							Dict									=StableDiffusionUpscalePipeline.from_pretrained(snake_case__ )
										pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )
										pipe.enable_attention_slicing()

										lowerCamelCase_	:							List[str]									="a cat sitting on a park bench"

										lowerCamelCase_	:							Optional[Any]									=torch.manual_seed(0 )
										lowerCamelCase_	:							List[str]									=pipe(
										    prompt=snake_case__       ,	image=snake_case__       ,	generator=snake_case__       ,	output_type="np"       ,	)
										lowerCamelCase_	:							List[str]									=output.images[0]

										assert image.shape == (512, 512, 3)
										assert np.abs(expected_image - image ).max() < 1E-3


			def      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							Union[str, Any]									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/sd2-upscale/low_res_cat.png" )
										lowerCamelCase_	:							int									=load_numpy(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
										    "/upsampled_cat_fp16.npy" )

										lowerCamelCase_	:							int									="stabilityai/stable-diffusion-x4-upscaler"
										lowerCamelCase_	:							Tuple									=StableDiffusionUpscalePipeline.from_pretrained(
										    snake_case__       ,	torch_dtype=torch.floataa       ,	)
										pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )
										pipe.enable_attention_slicing()

										lowerCamelCase_	:							List[Any]									="a cat sitting on a park bench"

										lowerCamelCase_	:							Tuple									=torch.manual_seed(0 )
										lowerCamelCase_	:							Tuple									=pipe(
										    prompt=snake_case__       ,	image=snake_case__       ,	generator=snake_case__       ,	output_type="np"       ,	)
										lowerCamelCase_	:							List[str]									=output.images[0]

										assert image.shape == (512, 512, 3)
										assert np.abs(expected_image - image ).max() < 5E-1


			def      UpperCAmelCase__						(				self					:  Any ):
										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()
										torch.cuda.reset_peak_memory_stats()

										lowerCamelCase_	:							Optional[Any]									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/sd2-upscale/low_res_cat.png" )

										lowerCamelCase_	:							List[str]									="stabilityai/stable-diffusion-x4-upscaler"
										lowerCamelCase_	:							str									=StableDiffusionUpscalePipeline.from_pretrained(
										    snake_case__       ,	torch_dtype=torch.floataa       ,	)
										pipe.to(snake_case__ )
										pipe.set_progress_bar_config(disable=snake_case__ )
										pipe.enable_attention_slicing(1 )
										pipe.enable_sequential_cpu_offload()

										lowerCamelCase_	:							Optional[Any]									="a cat sitting on a park bench"

										lowerCamelCase_	:							List[Any]									=torch.manual_seed(0 )
										lowerCamelCase_	:							Dict									=pipe(
										    prompt=snake_case__       ,	image=snake_case__       ,	generator=snake_case__       ,	num_inference_steps=5       ,	output_type="np"       ,	)

										lowerCamelCase_	:							str									=torch.cuda.max_memory_allocated()
										# make sure that less than 2.9 GB is allocated
										assert mem_bytes < 2.9 * 10**9











"""simple docstring"""



import inspect
from typing import Optional, Union

import numpy as np
import PIL
import torch
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPFeatureExtractor, CLIPModel, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    DPMSolverMultistepScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput
from diffusers.utils import (
    PIL_INTERPOLATION,
    randn_tensor,
)

def 			_snake_case    (					lowerCamelCase__     :   int    ,						lowerCamelCase__     :   Union[str, Any]    ,						lowerCamelCase__     :   int  )    ->     List[Any]:
							if isinstance(lowerCamelCase__    ,						torch.Tensor  ):
														return image
							elif isinstance(lowerCamelCase__    ,						PIL.Image.Image  ):
														lowerCamelCase_	:							int									=[image]

							if isinstance(image[0]    ,						PIL.Image.Image  ):
														lowerCamelCase_	:							Tuple									=[np.array(i.resize((w, h)    ,						resample=PIL_INTERPOLATION["lanczos"]  )  )[None, :] for i in image]
														lowerCamelCase_	:							Any									=np.concatenate(lowerCamelCase__    ,						axis=0  )
														lowerCamelCase_	:							Any									=np.array(lowerCamelCase__  ).astype(np.floataa  ) / 255.0
														lowerCamelCase_	:							Any									=image.transpose(0    ,						3    ,						1    ,						2  )
														lowerCamelCase_	:							Dict									=2.0 * image - 1.0
														lowerCamelCase_	:							List[Any]									=torch.from_numpy(lowerCamelCase__  )
							elif isinstance(image[0]    ,						torch.Tensor  ):
														lowerCamelCase_	:							Tuple									=torch.cat(lowerCamelCase__    ,						dim=0  )
							return image

def 			_snake_case    (					lowerCamelCase__     :   Union[str, Any]    ,						lowerCamelCase__     :   Dict    ,						lowerCamelCase__     :   Dict    ,						lowerCamelCase__     :   Optional[int]=0.9995  )    ->     List[str]:
							if not isinstance(lowerCamelCase__    ,						np.ndarray  ):
														lowerCamelCase_	:							List[Any]									=True
														lowerCamelCase_	:							Tuple									=va.device
														lowerCamelCase_	:							Optional[int]									=va.cpu().numpy()
														lowerCamelCase_	:							List[Any]									=va.cpu().numpy()

							lowerCamelCase_	:							Any									=np.sum(va * va / (np.linalg.norm(lowerCamelCase__  ) * np.linalg.norm(lowerCamelCase__  ))  )
							if np.abs(lowerCamelCase__  ) > DOT_THRESHOLD:
														lowerCamelCase_	:							str									=(1 - t) * va + t * va
							else:
														lowerCamelCase_	:							Tuple									=np.arccos(lowerCamelCase__  )
														lowerCamelCase_	:							Any									=np.sin(lowerCamelCase__  )
														lowerCamelCase_	:							Optional[int]									=theta_a * t
														lowerCamelCase_	:							Tuple									=np.sin(lowerCamelCase__  )
														lowerCamelCase_	:							List[str]									=np.sin(theta_a - theta_t  ) / sin_theta_a
														lowerCamelCase_	:							Union[str, Any]									=sin_theta_t / sin_theta_a
														lowerCamelCase_	:							str									=sa * va + sa * va

							if inputs_are_torch:
														lowerCamelCase_	:							Any									=torch.from_numpy(lowerCamelCase__  ).to(lowerCamelCase__  )

							return va

def 			_snake_case    (					lowerCamelCase__     :   Tuple    ,						lowerCamelCase__     :   Dict  )    ->     List[str]:
							lowerCamelCase_	:							List[Any]									=F.normalize(lowerCamelCase__    ,						dim=-1  )
							lowerCamelCase_	:							Any									=F.normalize(lowerCamelCase__    ,						dim=-1  )
							return (x - y).norm(dim=-1  ).div(2  ).arcsin().pow(2  ).mul(2  )

def 			_snake_case    (					lowerCamelCase__     :   str    ,						lowerCamelCase__     :   List[Any]  )    ->     Any:
							for param in model.parameters():
														lowerCamelCase_	:							Any									=value







class        lowercase__      (							snake_case__       ):


			def __init__(				self					:  Union[str, Any]       ,	snake_case__					:  AutoencoderKL       ,	snake_case__					:  CLIPTextModel       ,	snake_case__					:  CLIPModel       ,	snake_case__					:  CLIPTokenizer       ,	snake_case__					:  UNetaDConditionModel       ,	snake_case__					:  Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]       ,	snake_case__					:  CLIPFeatureExtractor       ,	snake_case__					:  Optional[Any]=None       ,	snake_case__					:  int=None       ,	snake_case__					:  str=None       ,	):
										super().__init__()
										self.register_modules(
										    vae=snake_case__       ,	text_encoder=snake_case__       ,	clip_model=snake_case__       ,	tokenizer=snake_case__       ,	unet=snake_case__       ,	scheduler=snake_case__       ,	feature_extractor=snake_case__       ,	coca_model=snake_case__       ,	coca_tokenizer=snake_case__       ,	coca_transform=snake_case__       ,	)
										lowerCamelCase_	:							Optional[Any]									=(
										    feature_extractor.size
										    if isinstance(feature_extractor.size       ,	snake_case__ )
										    else feature_extractor.size["shortest_edge"]
										)
										lowerCamelCase_	:							Union[str, Any]									=transforms.Normalize(mean=feature_extractor.image_mean       ,	std=feature_extractor.image_std )
										set_requires_grad(self.text_encoder       ,	snake_case__ )
										set_requires_grad(self.clip_model       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  Optional[Any]       ,	snake_case__					:  Optional[Union[str, int]] = "auto" ):
										if slice_size == "auto":
																	# half the attention head size is usually a good trade-off between
																	# speed and memory
																	lowerCamelCase_	:							List[Any]									=self.unet.config.attention_head_dim // 2
										self.unet.set_attention_slice(snake_case__ )


			def      UpperCAmelCase__						(				self					:  int ):
										self.enable_attention_slicing(snake_case__ )


			def      UpperCAmelCase__						(				self					:  str ):
										set_requires_grad(self.vae       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  List[str] ):
										set_requires_grad(self.vae       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  List[Any] ):
										set_requires_grad(self.unet       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  Tuple ):
										set_requires_grad(self.unet       ,	snake_case__ )


			def      UpperCAmelCase__						(				self					:  str       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  List[str]       ,	snake_case__					:  Tuple ):
										# get the original timestep using init_timestep
										lowerCamelCase_	:							Optional[int]									=min(int(num_inference_steps * strength )       ,	snake_case__ )

										lowerCamelCase_	:							Dict									=max(num_inference_steps - init_timestep       ,	0 )
										lowerCamelCase_	:							Union[str, Any]									=self.scheduler.timesteps[t_start:]

										return timesteps, num_inference_steps - t_start


			def      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  List[str]       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  List[Any]       ,	snake_case__					:  List[Any]=None ):
										if not isinstance(snake_case__       ,	torch.Tensor ):
																	raise ValueError(F"""`image` has to be of type `torch.Tensor` but is {type(snake_case__ )}""" )

										lowerCamelCase_	:							List[str]									=image.to(device=snake_case__       ,	dtype=snake_case__ )

										if isinstance(snake_case__       ,	snake_case__ ):
																	lowerCamelCase_	:							Optional[int]									=[
																	    self.vae.encode(image[i : i + 1] ).latent_dist.sample(generator[i] ) for i in range(snake_case__ )
																	]
																	lowerCamelCase_	:							Tuple									=torch.cat(snake_case__       ,	dim=0 )
										else:
																	lowerCamelCase_	:							Any									=self.vae.encode(snake_case__ ).latent_dist.sample(snake_case__ )

										# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
										lowerCamelCase_	:							Optional[Any]									=0.18_215 * init_latents
										lowerCamelCase_	:							Optional[Any]									=init_latents.repeat_interleave(snake_case__       ,	dim=0 )

										lowerCamelCase_	:							List[Any]									=randn_tensor(init_latents.shape       ,	generator=snake_case__       ,	device=snake_case__       ,	dtype=snake_case__ )

										# get latents
										lowerCamelCase_	:							Optional[int]									=self.scheduler.add_noise(snake_case__       ,	snake_case__       ,	snake_case__ )
										lowerCamelCase_	:							int									=init_latents

										return latents


			def      UpperCAmelCase__						(				self					:  Dict       ,	snake_case__					:  List[Any] ):
										lowerCamelCase_	:							Optional[int]									=self.coca_transform(snake_case__ ).unsqueeze(0 )
										with torch.no_grad(), torch.cuda.amp.autocast():
																	lowerCamelCase_	:							Tuple									=self.coca_model.generate(transformed_image.to(device=self.device       ,	dtype=self.coca_model.dtype ) )
										lowerCamelCase_	:							str									=self.coca_tokenizer.decode(generated[0].cpu().numpy() )
										return generated.split("<end_of_text>" )[0].replace("<start_of_text>"       ,	"" ).rstrip(" .," )


			def      UpperCAmelCase__						(				self					:  str       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  Optional[int] ):
										lowerCamelCase_	:							int									=self.feature_extractor.preprocess(snake_case__ )
										lowerCamelCase_	:							Tuple									=torch.from_numpy(clip_image_input["pixel_values"][0] ).unsqueeze(0 ).to(self.device ).half()
										lowerCamelCase_	:							Optional[int]									=self.clip_model.get_image_features(snake_case__ )
										lowerCamelCase_	:							List[str]									=image_embeddings_clip / image_embeddings_clip.norm(p=2       ,	dim=-1       ,	keepdim=snake_case__ )
										lowerCamelCase_	:							str									=image_embeddings_clip.repeat_interleave(snake_case__       ,	dim=0 )
										return image_embeddings_clip


			@torch.enable_grad()
			def      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  Dict       ,	snake_case__					:  Any       ,	snake_case__					:  str       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  List[Any]       ,	snake_case__					:  str       ,	):
										lowerCamelCase_	:							Any									=latents.detach().requires_grad_()

										lowerCamelCase_	:							Union[str, Any]									=self.scheduler.scale_model_input(snake_case__       ,	snake_case__ )

										# predict the noise residual
										lowerCamelCase_	:							Any									=self.unet(snake_case__       ,	snake_case__       ,	encoder_hidden_states=snake_case__ ).sample

										if isinstance(self.scheduler       ,	(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler) ):
																	lowerCamelCase_	:							Any									=self.scheduler.alphas_cumprod[timestep]
																	lowerCamelCase_	:							str									=1 - alpha_prod_t
																	# compute predicted original sample from predicted noise also called
																	# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
																	lowerCamelCase_	:							Dict									=(latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

																	lowerCamelCase_	:							Tuple									=torch.sqrt(snake_case__ )
																	lowerCamelCase_	:							Optional[int]									=pred_original_sample * (fac) + latents * (1 - fac)
										elif isinstance(self.scheduler       ,	snake_case__ ):
																	lowerCamelCase_	:							Optional[Any]									=self.scheduler.sigmas[index]
																	lowerCamelCase_	:							Optional[int]									=latents - sigma * noise_pred
										else:
																	raise ValueError(F"""scheduler type {type(self.scheduler )} not supported""" )

										# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
										lowerCamelCase_	:							str									=1 / 0.18_215 * sample
										lowerCamelCase_	:							List[Any]									=self.vae.decode(snake_case__ ).sample
										lowerCamelCase_	:							List[Any]									=(image / 2 + 0.5).clamp(0       ,	1 )

										lowerCamelCase_	:							Any									=transforms.Resize(self.feature_extractor_size )(snake_case__ )
										lowerCamelCase_	:							Any									=self.normalize(snake_case__ ).to(latents.dtype )

										lowerCamelCase_	:							Any									=self.clip_model.get_image_features(snake_case__ )
										lowerCamelCase_	:							Optional[int]									=image_embeddings_clip / image_embeddings_clip.norm(p=2       ,	dim=-1       ,	keepdim=snake_case__ )

										lowerCamelCase_	:							Union[str, Any]									=spherical_dist_loss(snake_case__       ,	snake_case__ ).mean() * clip_guidance_scale

										lowerCamelCase_	:							Union[str, Any]									=-torch.autograd.grad(snake_case__       ,	snake_case__ )[0]

										if isinstance(self.scheduler       ,	snake_case__ ):
																	lowerCamelCase_	:							int									=latents.detach() + grads * (sigma**2)
																	lowerCamelCase_	:							Tuple									=noise_pred_original
										else:
																	lowerCamelCase_	:							Union[str, Any]									=noise_pred_original - torch.sqrt(snake_case__ ) * grads
										return noise_pred, latents


			@torch.no_grad()
			def __call__(				self					:  Any       ,	snake_case__					:  Union[torch.FloatTensor, PIL.Image.Image]       ,	snake_case__					:  Union[torch.FloatTensor, PIL.Image.Image]       ,	snake_case__					:  Optional[str] = None       ,	snake_case__					:  Optional[str] = None       ,	snake_case__					:  Optional[int] = 512       ,	snake_case__					:  Optional[int] = 512       ,	snake_case__					:  float = 0.6       ,	snake_case__					:  Optional[int] = 50       ,	snake_case__					:  Optional[float] = 7.5       ,	snake_case__					:  Optional[int] = 1       ,	snake_case__					:  float = 0.0       ,	snake_case__					:  Optional[float] = 100       ,	snake_case__					:  Optional[torch.Generator] = None       ,	snake_case__					:  Optional[str] = "pil"       ,	snake_case__					:  bool = True       ,	snake_case__					:  float = 0.8       ,	snake_case__					:  float = 0.1       ,	snake_case__					:  float = 0.1       ,	):
										if isinstance(snake_case__       ,	snake_case__ ) and len(snake_case__ ) != batch_size:
																	raise ValueError(F"""You have passed {batch_size} batch_size, but only {len(snake_case__ )} generators.""" )

										if height % 8 != 0 or width % 8 != 0:
																	raise ValueError(F"""`height` and `width` have to be divisible by 8 but are {height} and {width}.""" )

										if isinstance(snake_case__       ,	torch.Generator ) and batch_size > 1:
																	lowerCamelCase_	:							List[str]									=[generator] + [None] * (batch_size - 1)

										lowerCamelCase_	:							Any									=[
										    ("model", self.coca_model is None),
										    ("tokenizer", self.coca_tokenizer is None),
										    ("transform", self.coca_transform is None),
										]
										lowerCamelCase_	:							Optional[Any]									=[x[0] for x in coca_is_none if x[1]]
										lowerCamelCase_	:							Any									=", ".join(snake_case__ )
										# generate prompts with coca model if prompt is None
										if content_prompt is None:
																	if len(snake_case__ ):
																								raise ValueError(
																								    F"""Content prompt is None and CoCa [{coca_is_none_str}] is None."""
																								    F"""Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.""" )
																	lowerCamelCase_	:							Union[str, Any]									=self.get_image_description(snake_case__ )
										if style_prompt is None:
																	if len(snake_case__ ):
																								raise ValueError(
																								    F"""Style prompt is None and CoCa [{coca_is_none_str}] is None."""
																								    F""" Set prompt or pass Coca [{coca_is_none_str}] to DiffusionPipeline.""" )
																	lowerCamelCase_	:							str									=self.get_image_description(snake_case__ )

										# get prompt text embeddings for content and style
										lowerCamelCase_	:							str									=self.tokenizer(
										    snake_case__       ,	padding="max_length"       ,	max_length=self.tokenizer.model_max_length       ,	truncation=snake_case__       ,	return_tensors="pt"       ,	)
										lowerCamelCase_	:							List[str]									=self.text_encoder(content_text_input.input_ids.to(self.device ) )[0]

										lowerCamelCase_	:							int									=self.tokenizer(
										    snake_case__       ,	padding="max_length"       ,	max_length=self.tokenizer.model_max_length       ,	truncation=snake_case__       ,	return_tensors="pt"       ,	)
										lowerCamelCase_	:							Optional[int]									=self.text_encoder(style_text_input.input_ids.to(self.device ) )[0]

										lowerCamelCase_	:							Dict									=slerp(snake_case__       ,	snake_case__       ,	snake_case__ )

										# duplicate text embeddings for each generation per prompt
										lowerCamelCase_	:							str									=text_embeddings.repeat_interleave(snake_case__       ,	dim=0 )

										# set timesteps
										lowerCamelCase_	:							List[Any]									="offset" in set(inspect.signature(self.scheduler.set_timesteps ).parameters.keys() )
										lowerCamelCase_	:							List[Any]									={}
										if accepts_offset:
																	lowerCamelCase_	:							Optional[int]									=1

										self.scheduler.set_timesteps(snake_case__       ,	**snake_case__ )
										# Some schedulers like PNDM have timesteps as arrays
										# It's more optimized to move all timesteps to correct device beforehand
										self.scheduler.timesteps.to(self.device )

										lowerCamelCase_							,						lowerCamelCase_	:							Optional[int]									=self.get_timesteps(snake_case__       ,	snake_case__       ,	self.device )
										lowerCamelCase_	:							Union[str, Any]									=timesteps[:1].repeat(snake_case__ )

										# Preprocess image
										lowerCamelCase_	:							str									=preprocess(snake_case__       ,	snake_case__       ,	snake_case__ )
										lowerCamelCase_	:							int									=self.prepare_latents(
										    snake_case__       ,	snake_case__       ,	snake_case__       ,	text_embeddings.dtype       ,	self.device       ,	snake_case__ )

										lowerCamelCase_	:							Dict									=preprocess(snake_case__       ,	snake_case__       ,	snake_case__ )
										lowerCamelCase_	:							Tuple									=self.prepare_latents(
										    snake_case__       ,	snake_case__       ,	snake_case__       ,	text_embeddings.dtype       ,	self.device       ,	snake_case__ )

										lowerCamelCase_	:							int									=slerp(snake_case__       ,	snake_case__       ,	snake_case__ )

										if clip_guidance_scale > 0:
																	lowerCamelCase_	:							List[Any]									=self.get_clip_image_embeddings(snake_case__       ,	snake_case__ )
																	lowerCamelCase_	:							Optional[Any]									=self.get_clip_image_embeddings(snake_case__       ,	snake_case__ )
																	lowerCamelCase_	:							List[Any]									=slerp(
																	    snake_case__       ,	snake_case__       ,	snake_case__ )

										# 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.
										lowerCamelCase_	:							str									=guidance_scale > 1.0
										# get unconditional embeddings for classifier free guidance
										if do_classifier_free_guidance:
																	lowerCamelCase_	:							List[Any]									=content_text_input.input_ids.shape[-1]
																	lowerCamelCase_	:							Union[str, Any]									=self.tokenizer([""]       ,	padding="max_length"       ,	max_length=snake_case__       ,	return_tensors="pt" )
																	lowerCamelCase_	:							List[Any]									=self.text_encoder(uncond_input.input_ids.to(self.device ) )[0]
																	# duplicate unconditional embeddings for each generation per prompt
																	lowerCamelCase_	:							List[str]									=uncond_embeddings.repeat_interleave(snake_case__       ,	dim=0 )

																	# 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
																	lowerCamelCase_	:							List[str]									=torch.cat([uncond_embeddings, text_embeddings] )

										# get the initial random noise unless the user supplied it

										# Unlike in other pipelines, latents need to be generated in the target device
										# for 1-to-1 results reproducibility with the CompVis implementation.
										# However this currently doesn't work in `mps`.
										lowerCamelCase_	:							Optional[Any]									=(batch_size, self.unet.config.in_channels, height // 8, width // 8)
										lowerCamelCase_	:							List[str]									=text_embeddings.dtype
										if latents is None:
																	if self.device.type == "mps":
																								# randn does not work reproducibly on mps
																								lowerCamelCase_	:							Union[str, Any]									=torch.randn(snake_case__       ,	generator=snake_case__       ,	device="cpu"       ,	dtype=snake_case__ ).to(
																								    self.device )
																	else:
																								lowerCamelCase_	:							Tuple									=torch.randn(snake_case__       ,	generator=snake_case__       ,	device=self.device       ,	dtype=snake_case__ )
										else:
																	if latents.shape != latents_shape:
																								raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}""" )
																	lowerCamelCase_	:							Any									=latents.to(self.device )

										# scale the initial noise by the standard deviation required by the scheduler
										lowerCamelCase_	:							Optional[int]									=latents * self.scheduler.init_noise_sigma

										# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
										# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
										# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
										# and should be between [0, 1]
										lowerCamelCase_	:							List[Any]									="eta" in set(inspect.signature(self.scheduler.step ).parameters.keys() )
										lowerCamelCase_	:							Dict									={}
										if accepts_eta:
																	lowerCamelCase_	:							Tuple									=eta

										# check if the scheduler accepts generator
										lowerCamelCase_	:							Tuple									="generator" in set(inspect.signature(self.scheduler.step ).parameters.keys() )
										if accepts_generator:
																	lowerCamelCase_	:							Optional[int]									=generator

										with self.progress_bar(total=snake_case__ ):
																	for i, t in enumerate(snake_case__ ):
																								# expand the latents if we are doing classifier free guidance
																								lowerCamelCase_	:							List[str]									=torch.cat([latents] * 2 ) if do_classifier_free_guidance else latents
																								lowerCamelCase_	:							Optional[Any]									=self.scheduler.scale_model_input(snake_case__       ,	snake_case__ )

																								# predict the noise residual
																								lowerCamelCase_	:							Tuple									=self.unet(snake_case__       ,	snake_case__       ,	encoder_hidden_states=snake_case__ ).sample

																								# perform classifier free guidance
																								if do_classifier_free_guidance:
																															lowerCamelCase_							,						lowerCamelCase_	:							Dict									=noise_pred.chunk(2 )
																															lowerCamelCase_	:							Optional[Any]									=noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

																								# perform clip guidance
																								if clip_guidance_scale > 0:
																															lowerCamelCase_	:							Optional[int]									=(
																															    text_embeddings.chunk(2 )[1] if do_classifier_free_guidance else text_embeddings
																															)
																															lowerCamelCase_							,						lowerCamelCase_	:							Union[str, Any]									=self.cond_fn(
																															    snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	)

																								# compute the previous noisy sample x_t -> x_t-1
																								lowerCamelCase_	:							Dict									=self.scheduler.step(snake_case__       ,	snake_case__       ,	snake_case__       ,	**snake_case__ ).prev_sample

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
										lowerCamelCase_	:							str									=1 / 0.18_215 * latents
										lowerCamelCase_	:							List[str]									=self.vae.decode(snake_case__ ).sample

										lowerCamelCase_	:							str									=(image / 2 + 0.5).clamp(0       ,	1 )
										lowerCamelCase_	:							List[Any]									=image.cpu().permute(0       ,	2       ,	3       ,	1 ).numpy()

										if output_type == "pil":
																	lowerCamelCase_	:							List[str]									=self.numpy_to_pil(snake_case__ )

										if not return_dict:
																	return (image, None)

										return StableDiffusionPipelineOutput(images=snake_case__       ,	nsfw_content_detected=snake_case__ )



