

"""simple docstring"""




import inspect
from typing import List, Optional, Union

import torch
from torch import nn
from torch.nn import functional as F
from torchvision import transforms
from transformers import CLIPImageProcessor, CLIPModel, CLIPTextModel, CLIPTokenizer

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





class 				lowerCamelCase  (nn.Module     ):

			def __init__(   self			:  Optional[Any]      ,			__UpperCAmelCase			:  Optional[Any]      ,			__UpperCAmelCase			:  Optional[Any]=1.0    )				->      List[Any]:
						super().__init__()

						SCREAMING_SNAKE_CASE__									= cut_size
						SCREAMING_SNAKE_CASE__									= cut_power





			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[Any]      ,			__UpperCAmelCase			:  int      ,			__UpperCAmelCase			:  Dict    )				->      List[Any]:
						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__									= pixel_values.shape[2:4]
						SCREAMING_SNAKE_CASE__									= min(__UpperCAmelCase      ,			__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= min(__UpperCAmelCase      ,			__UpperCAmelCase      ,			self.cut_size    )
						SCREAMING_SNAKE_CASE__									= []
						for _ in range(__UpperCAmelCase    ):
									SCREAMING_SNAKE_CASE__									= int(torch.rand([]    ) ** self.cut_power * (max_size - min_size) + min_size    )
									SCREAMING_SNAKE_CASE__									= torch.randint(0      ,			sideX - size + 1      ,			()    )
									SCREAMING_SNAKE_CASE__									= torch.randint(0      ,			sideY - size + 1      ,			()    )
									SCREAMING_SNAKE_CASE__									= pixel_values[:, :, offsety : offsety + size, offsetx : offsetx + size]
									cutouts.append(F.adaptive_avg_poolad(__UpperCAmelCase      ,			self.cut_size    )    )
						return torch.cat(__UpperCAmelCase    )

def 		A    (       snake_case__						,			snake_case__      ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__									= F.normalize(snake_case__						,			dim=-1      )
			SCREAMING_SNAKE_CASE__									= F.normalize(snake_case__						,			dim=-1      )
			return (x - y).norm(dim=-1      ).div(2      ).arcsin().pow(2      ).mul(2      )

def 		A    (       snake_case__						,			snake_case__      ):




			'''simple docstring'''


			for param in model.parameters():
						SCREAMING_SNAKE_CASE__									= value





class 				lowerCamelCase  (A__     ):

			def __init__(   self			:  List[str]      ,			__UpperCAmelCase			:  AutoencoderKL      ,			__UpperCAmelCase			:  CLIPTextModel      ,			__UpperCAmelCase			:  CLIPModel      ,			__UpperCAmelCase			:  CLIPTokenizer      ,			__UpperCAmelCase			:  UNetaDConditionModel      ,			__UpperCAmelCase			:  Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]      ,			__UpperCAmelCase			:  CLIPImageProcessor      ,			)				->      int:
						super().__init__()
						self.register_modules(
						    vae=__UpperCAmelCase      ,			text_encoder=__UpperCAmelCase      ,			clip_model=__UpperCAmelCase      ,			tokenizer=__UpperCAmelCase      ,			unet=__UpperCAmelCase      ,			scheduler=__UpperCAmelCase      ,			feature_extractor=__UpperCAmelCase      ,			)

						SCREAMING_SNAKE_CASE__									= transforms.Normalize(mean=feature_extractor.image_mean      ,			std=feature_extractor.image_std    )
						SCREAMING_SNAKE_CASE__									= (
						    feature_extractor.size
						    if isinstance(feature_extractor.size      ,			__UpperCAmelCase    )
						    else feature_extractor.size["""shortest_edge"""]
						)
						SCREAMING_SNAKE_CASE__									= MakeCutouts(self.cut_out_size    )

						set_requires_grad(self.text_encoder      ,			__UpperCAmelCase    )
						set_requires_grad(self.clip_model      ,			__UpperCAmelCase    )

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

			def    SCREAMING_SNAKE_CASE			(   self			:  int    )				->      int:
						self.enable_attention_slicing(__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  int    )				->      Tuple:
						set_requires_grad(self.vae      ,			__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Dict    )				->      str:
						set_requires_grad(self.vae      ,			__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Tuple:
						set_requires_grad(self.unet      ,			__UpperCAmelCase    )

			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      Any:
						set_requires_grad(self.unet      ,			__UpperCAmelCase    )

			@torch.enable_grad()
			def    SCREAMING_SNAKE_CASE			(   self			:  Union[str, Any]      ,			__UpperCAmelCase			:  Dict      ,			__UpperCAmelCase			:  List[str]      ,			__UpperCAmelCase			:  Tuple      ,			__UpperCAmelCase			:  Tuple      ,			__UpperCAmelCase			:  List[str]      ,			__UpperCAmelCase			:  Optional[int]      ,			__UpperCAmelCase			:  List[Any]      ,			__UpperCAmelCase			:  Any      ,			__UpperCAmelCase			:  List[Any]=True      ,			)				->      Optional[int]:
						SCREAMING_SNAKE_CASE__									= latents.detach().requires_grad_()

						SCREAMING_SNAKE_CASE__									= self.scheduler.scale_model_input(__UpperCAmelCase      ,			__UpperCAmelCase    )

						# predict the noise residual
						SCREAMING_SNAKE_CASE__									= self.unet(__UpperCAmelCase      ,			__UpperCAmelCase      ,			encoder_hidden_states=__UpperCAmelCase    ).sample

						if isinstance(self.scheduler      ,			(PNDMScheduler, DDIMScheduler, DPMSolverMultistepScheduler)    ):
									SCREAMING_SNAKE_CASE__									= self.scheduler.alphas_cumprod[timestep]
									SCREAMING_SNAKE_CASE__									= 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
									SCREAMING_SNAKE_CASE__									= (latents - beta_prod_t ** 0.5 * noise_pred) / alpha_prod_t ** 0.5

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

						SCREAMING_SNAKE_CASE__									= 1 / self.vae.config.scaling_factor * sample
						SCREAMING_SNAKE_CASE__									= self.vae.decode(__UpperCAmelCase    ).sample
						SCREAMING_SNAKE_CASE__									= (image / 2 + 0.5).clamp(0      ,			1    )

						if use_cutouts:
									SCREAMING_SNAKE_CASE__									= self.make_cutouts(__UpperCAmelCase      ,			__UpperCAmelCase    )
						else:
									SCREAMING_SNAKE_CASE__									= transforms.Resize(self.cut_out_size    )(__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= self.normalize(__UpperCAmelCase    ).to(latents.dtype    )

						SCREAMING_SNAKE_CASE__									= self.clip_model.get_image_features(__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= image_embeddings_clip / image_embeddings_clip.norm(p=2      ,			dim=-1      ,			keepdim=__UpperCAmelCase    )

						if use_cutouts:
									SCREAMING_SNAKE_CASE__									= spherical_dist_loss(__UpperCAmelCase      ,			__UpperCAmelCase    )
									SCREAMING_SNAKE_CASE__									= dists.view([num_cutouts, sample.shape[0], -1]    )
									SCREAMING_SNAKE_CASE__									= dists.sum(2    ).mean(0    ).sum() * clip_guidance_scale
						else:
									SCREAMING_SNAKE_CASE__									= spherical_dist_loss(__UpperCAmelCase      ,			__UpperCAmelCase    ).mean() * clip_guidance_scale

						SCREAMING_SNAKE_CASE__									= -torch.autograd.grad(__UpperCAmelCase      ,			__UpperCAmelCase    )[0]

						if isinstance(self.scheduler      ,			__UpperCAmelCase    ):
									SCREAMING_SNAKE_CASE__									= latents.detach() + grads * (sigma**2)
									SCREAMING_SNAKE_CASE__									= noise_pred_original
						else:
									SCREAMING_SNAKE_CASE__									= noise_pred_original - torch.sqrt(__UpperCAmelCase    ) * grads
						return noise_pred, latents





			@torch.no_grad()
			def __call__(   self			:  List[Any]      ,			__UpperCAmelCase			:  Union[str, List[str]]      ,			__UpperCAmelCase			:  Optional[int] = 5_1_2      ,			__UpperCAmelCase			:  Optional[int] = 5_1_2      ,			__UpperCAmelCase			:  Optional[int] = 5_0      ,			__UpperCAmelCase			:  Optional[float] = 7.5      ,			__UpperCAmelCase			:  Optional[int] = 1      ,			__UpperCAmelCase			:  float = 0.0      ,			__UpperCAmelCase			:  Optional[float] = 1_0_0      ,			__UpperCAmelCase			:  Optional[Union[str, List[str]]] = None      ,			__UpperCAmelCase			:  Optional[int] = 4      ,			__UpperCAmelCase			:  Optional[bool] = True      ,			__UpperCAmelCase			:  Optional[torch.Generator] = None      ,			__UpperCAmelCase			:  Optional[torch.FloatTensor] = None      ,			__UpperCAmelCase			:  Optional[str] = "pil"      ,			__UpperCAmelCase			:  bool = True      ,			)				->      Any:
						if isinstance(__UpperCAmelCase      ,			__UpperCAmelCase    ):
									SCREAMING_SNAKE_CASE__									= 1
						elif isinstance(__UpperCAmelCase      ,			__UpperCAmelCase    ):
									SCREAMING_SNAKE_CASE__									= len(__UpperCAmelCase    )
						else:
									raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__UpperCAmelCase    )}"""    )

						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}."""    )

						# get prompt text embeddings
						SCREAMING_SNAKE_CASE__									= self.tokenizer(
						    __UpperCAmelCase      ,			padding="""max_length"""      ,			max_length=self.tokenizer.model_max_length      ,			truncation=__UpperCAmelCase      ,			return_tensors="""pt"""      ,			)
						SCREAMING_SNAKE_CASE__									= self.text_encoder(text_input.input_ids.to(self.device    )    )[0]
						# duplicate text embeddings for each generation per prompt
						SCREAMING_SNAKE_CASE__									= text_embeddings.repeat_interleave(__UpperCAmelCase      ,			dim=0    )

						if clip_guidance_scale > 0:
									if clip_prompt is not None:
												SCREAMING_SNAKE_CASE__									= self.tokenizer(
												    __UpperCAmelCase      ,			padding="""max_length"""      ,			max_length=self.tokenizer.model_max_length      ,			truncation=__UpperCAmelCase      ,			return_tensors="""pt"""      ,			).input_ids.to(self.device    )
									else:
												SCREAMING_SNAKE_CASE__									= text_input.input_ids.to(self.device    )
									SCREAMING_SNAKE_CASE__									= self.clip_model.get_text_features(__UpperCAmelCase    )
									SCREAMING_SNAKE_CASE__									= text_embeddings_clip / text_embeddings_clip.norm(p=2      ,			dim=-1      ,			keepdim=__UpperCAmelCase    )
									# duplicate text embeddings clip for each generation per prompt
									SCREAMING_SNAKE_CASE__									= text_embeddings_clip.repeat_interleave(__UpperCAmelCase      ,			dim=0    )

						# 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.
						SCREAMING_SNAKE_CASE__									= guidance_scale > 1.0
						# get unconditional embeddings for classifier free guidance
						if do_classifier_free_guidance:
									SCREAMING_SNAKE_CASE__									= text_input.input_ids.shape[-1]
									SCREAMING_SNAKE_CASE__									= self.tokenizer([""""""]      ,			padding="""max_length"""      ,			max_length=__UpperCAmelCase      ,			return_tensors="""pt"""    )
									SCREAMING_SNAKE_CASE__									= self.text_encoder(uncond_input.input_ids.to(self.device    )    )[0]
									# duplicate unconditional embeddings for each generation per prompt
									SCREAMING_SNAKE_CASE__									= uncond_embeddings.repeat_interleave(__UpperCAmelCase      ,			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
									SCREAMING_SNAKE_CASE__									= 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`.
						SCREAMING_SNAKE_CASE__									= (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
						SCREAMING_SNAKE_CASE__									= text_embeddings.dtype
						if latents is None:
									if self.device.type == "mps":
												# randn does not work reproducibly on mps
												SCREAMING_SNAKE_CASE__									= torch.randn(__UpperCAmelCase      ,			generator=__UpperCAmelCase      ,			device="""cpu"""      ,			dtype=__UpperCAmelCase    ).to(
												    self.device    )
									else:
												SCREAMING_SNAKE_CASE__									= torch.randn(__UpperCAmelCase      ,			generator=__UpperCAmelCase      ,			device=self.device      ,			dtype=__UpperCAmelCase    )
						else:
									if latents.shape != latents_shape:
												raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""    )
									SCREAMING_SNAKE_CASE__									= latents.to(self.device    )

						# set timesteps
						SCREAMING_SNAKE_CASE__									= """offset""" in set(inspect.signature(self.scheduler.set_timesteps    ).parameters.keys()    )
						SCREAMING_SNAKE_CASE__									= {}
						if accepts_offset:
									SCREAMING_SNAKE_CASE__									= 1

						self.scheduler.set_timesteps(__UpperCAmelCase      ,			**__UpperCAmelCase    )

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

						# scale the initial noise by the standard deviation required by the scheduler
						SCREAMING_SNAKE_CASE__									= 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]
						SCREAMING_SNAKE_CASE__									= """eta""" in set(inspect.signature(self.scheduler.step    ).parameters.keys()    )
						SCREAMING_SNAKE_CASE__									= {}
						if accepts_eta:
									SCREAMING_SNAKE_CASE__									= eta

						# check if the scheduler accepts generator
						SCREAMING_SNAKE_CASE__									= """generator""" in set(inspect.signature(self.scheduler.step    ).parameters.keys()    )
						if accepts_generator:
									SCREAMING_SNAKE_CASE__									= generator

						for i, t in enumerate(self.progress_bar(__UpperCAmelCase    )    ):
									# expand the latents if we are doing classifier free guidance
									SCREAMING_SNAKE_CASE__									= torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents
									SCREAMING_SNAKE_CASE__									= self.scheduler.scale_model_input(__UpperCAmelCase      ,			__UpperCAmelCase    )

									# predict the noise residual
									SCREAMING_SNAKE_CASE__									= self.unet(__UpperCAmelCase      ,			__UpperCAmelCase      ,			encoder_hidden_states=__UpperCAmelCase    ).sample

									# perform classifier free guidance
									if do_classifier_free_guidance:
												SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__									= noise_pred.chunk(2    )
												SCREAMING_SNAKE_CASE__									= noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

									# perform clip guidance
									if clip_guidance_scale > 0:
												SCREAMING_SNAKE_CASE__									= (
												    text_embeddings.chunk(2    )[1] if do_classifier_free_guidance else text_embeddings
												)
												SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__									= self.cond_fn(
												    __UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			)

									# compute the previous noisy sample x_t -> x_t-1
									SCREAMING_SNAKE_CASE__									= self.scheduler.step(__UpperCAmelCase      ,			__UpperCAmelCase      ,			__UpperCAmelCase      ,			**__UpperCAmelCase    ).prev_sample

						# scale and decode the image latents with vae
						SCREAMING_SNAKE_CASE__									= 1 / self.vae.config.scaling_factor * latents
						SCREAMING_SNAKE_CASE__									= self.vae.decode(__UpperCAmelCase    ).sample

						SCREAMING_SNAKE_CASE__									= (image / 2 + 0.5).clamp(0      ,			1    )
						SCREAMING_SNAKE_CASE__									= image.cpu().permute(0      ,			2      ,			3      ,			1    ).numpy()

						if output_type == "pil":
									SCREAMING_SNAKE_CASE__									= self.numpy_to_pil(__UpperCAmelCase    )

						if not return_dict:
									return (image, None)

						return StableDiffusionPipelineOutput(images=__UpperCAmelCase      ,			nsfw_content_detected=__UpperCAmelCase    )





"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class 				lowerCamelCase  (A__     ,unittest.TestCase     ):
			lowerCamelCase__				:						Optional[Any]   						=							KandinskyPipeline
			lowerCamelCase__				:						Union[str, Any]   						=							[
			    'prompt',
			    'image_embeds',
			    'negative_image_embeds',
			]
			lowerCamelCase__				:						Optional[Any]   						=							['prompt', 'negative_prompt', 'image_embeds', 'negative_image_embeds']
			lowerCamelCase__				:						Union[str, Any]   						=							[
			    'generator',
			    'height',
			    'width',
			    'latents',
			    'guidance_scale',
			    'negative_prompt',
			    'num_inference_steps',
			    'return_dict',
			    'guidance_scale',
			    'num_images_per_prompt',
			    'output_type',
			    'return_dict',
			]
			lowerCamelCase__				:						Optional[Any]   						=							False

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      List[str]:
						return 3_2

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      int:
						return 3_2

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  Union[str, Any]    )				->      Any:
						return self.time_input_dim

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      Any:
						return self.time_input_dim * 4

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      Optional[Any]:
						return 1_0_0

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  List[Any]    )				->      List[str]:
						SCREAMING_SNAKE_CASE__									= XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""    )
						return tokenizer

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Dict:
						torch.manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= MCLIPConfig(
						    numDims=self.cross_attention_dim      ,			transformerDimensions=self.text_embedder_hidden_size      ,			hidden_size=self.text_embedder_hidden_size      ,			intermediate_size=3_7      ,			num_attention_heads=4      ,			num_hidden_layers=5      ,			vocab_size=1_0_0_5      ,			)

						SCREAMING_SNAKE_CASE__									= MultilingualCLIP(__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= text_encoder.eval()

						return text_encoder

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  str    )				->      int:
						torch.manual_seed(0    )

						SCREAMING_SNAKE_CASE__									= {
						    """in_channels""": 4,
						    # Out channels is double in channels because predicts mean and variance
						    """out_channels""": 8,
						    """addition_embed_type""": """text_image""",
						    """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
						    """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
						    """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
						    """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
						    """layers_per_block""": 1,
						    """encoder_hid_dim""": self.text_embedder_hidden_size,
						    """encoder_hid_dim_type""": """text_image_proj""",
						    """cross_attention_dim""": self.cross_attention_dim,
						    """attention_head_dim""": 4,
						    """resnet_time_scale_shift""": """scale_shift""",
						    """class_embed_type""": None,
						}

						SCREAMING_SNAKE_CASE__									= UNetaDConditionModel(**__UpperCAmelCase    )
						return model

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  List[Any]    )				->      Union[str, Any]:
						return {
						    "block_out_channels": [3_2, 6_4],
						    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
						    "in_channels": 3,
						    "latent_channels": 4,
						    "layers_per_block": 1,
						    "norm_num_groups": 8,
						    "norm_type": "spatial",
						    "num_vq_embeddings": 1_2,
						    "out_channels": 3,
						    "up_block_types": [
						        "AttnUpDecoderBlock2D",
						        "UpDecoderBlock2D",
						    ],
						    "vq_embed_dim": 4,
						}

			@property
			def    SCREAMING_SNAKE_CASE			(   self			:  int    )				->      List[str]:
						torch.manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= VQModel(**self.dummy_movq_kwargs    )
						return model

			def    SCREAMING_SNAKE_CASE			(   self			:  Any    )				->      Tuple:
						SCREAMING_SNAKE_CASE__									= self.dummy_text_encoder
						SCREAMING_SNAKE_CASE__									= self.dummy_tokenizer
						SCREAMING_SNAKE_CASE__									= self.dummy_unet
						SCREAMING_SNAKE_CASE__									= self.dummy_movq

						SCREAMING_SNAKE_CASE__									= DDIMScheduler(
						    num_train_timesteps=1_0_0_0      ,			beta_schedule="""linear"""      ,			beta_start=0.00_085      ,			beta_end=0.012      ,			clip_sample=__UpperCAmelCase      ,			set_alpha_to_one=__UpperCAmelCase      ,			steps_offset=1      ,			prediction_type="""epsilon"""      ,			thresholding=__UpperCAmelCase      ,			)

						SCREAMING_SNAKE_CASE__									= {
						    """text_encoder""": text_encoder,
						    """tokenizer""": tokenizer,
						    """unet""": unet,
						    """scheduler""": scheduler,
						    """movq""": movq,
						}
						return components

			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[Any]      ,			__UpperCAmelCase			:  Dict      ,			__UpperCAmelCase			:  Optional[Any]=0    )				->      str:
						SCREAMING_SNAKE_CASE__									= floats_tensor((1, self.cross_attention_dim)      ,			rng=random.Random(__UpperCAmelCase    )    ).to(__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= floats_tensor((1, self.cross_attention_dim)      ,			rng=random.Random(seed + 1    )    ).to(__UpperCAmelCase    )
						if str(__UpperCAmelCase    ).startswith("""mps"""    ):
									SCREAMING_SNAKE_CASE__									= torch.manual_seed(__UpperCAmelCase    )
						else:
									SCREAMING_SNAKE_CASE__									= torch.Generator(device=__UpperCAmelCase    ).manual_seed(__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= {
						    """prompt""": """horse""",
						    """image_embeds""": image_embeds,
						    """negative_image_embeds""": negative_image_embeds,
						    """generator""": generator,
						    """height""": 6_4,
						    """width""": 6_4,
						    """guidance_scale""": 4.0,
						    """num_inference_steps""": 2,
						    """output_type""": """np""",
						}
						return inputs





			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      int:
						SCREAMING_SNAKE_CASE__									= """cpu"""

						SCREAMING_SNAKE_CASE__									= self.get_dummy_components()

						SCREAMING_SNAKE_CASE__									= self.pipeline_class(**__UpperCAmelCase    )
						SCREAMING_SNAKE_CASE__									= pipe.to(__UpperCAmelCase    )

						pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= pipe(**self.get_dummy_inputs(__UpperCAmelCase    )    )
						SCREAMING_SNAKE_CASE__									= output.images

						SCREAMING_SNAKE_CASE__									= pipe(
						    **self.get_dummy_inputs(__UpperCAmelCase    )      ,			return_dict=__UpperCAmelCase      ,			)[0]

						SCREAMING_SNAKE_CASE__									= image[0, -3:, -3:, -1]
						SCREAMING_SNAKE_CASE__									= image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 6_4, 6_4, 3)

						SCREAMING_SNAKE_CASE__									= np.array(
						    [0.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]    )

						assert (
						    np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
						), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
						assert (
						    np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1e-2
						), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""





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

			def    SCREAMING_SNAKE_CASE			(   self			:  Optional[Any]    )				->      List[str]:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()





			def    SCREAMING_SNAKE_CASE			(   self			:  List[str]    )				->      Optional[int]:
						SCREAMING_SNAKE_CASE__									= load_numpy(
						    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
						    """/kandinsky/kandinsky_text2img_cat_fp16.npy"""    )

						SCREAMING_SNAKE_CASE__									= KandinskyPriorPipeline.from_pretrained(
						    """kandinsky-community/kandinsky-2-1-prior"""      ,			torch_dtype=torch.floataa    )
						pipe_prior.to(__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= KandinskyPipeline.from_pretrained("""kandinsky-community/kandinsky-2-1"""      ,			torch_dtype=torch.floataa    )
						SCREAMING_SNAKE_CASE__									= pipeline.to(__UpperCAmelCase    )
						pipeline.set_progress_bar_config(disable=__UpperCAmelCase    )

						SCREAMING_SNAKE_CASE__									= """red cat, 4k photo"""

						SCREAMING_SNAKE_CASE__									= torch.Generator(device="""cuda"""    ).manual_seed(0    )
						SCREAMING_SNAKE_CASE__							,						SCREAMING_SNAKE_CASE__									= pipe_prior(
						    __UpperCAmelCase      ,			generator=__UpperCAmelCase      ,			num_inference_steps=5      ,			negative_prompt=""""""      ,			).to_tuple()

						SCREAMING_SNAKE_CASE__									= torch.Generator(device="""cuda"""    ).manual_seed(0    )
						SCREAMING_SNAKE_CASE__									= pipeline(
						    __UpperCAmelCase      ,			image_embeds=__UpperCAmelCase      ,			negative_image_embeds=__UpperCAmelCase      ,			generator=__UpperCAmelCase      ,			num_inference_steps=1_0_0      ,			output_type="""np"""      ,			)

						SCREAMING_SNAKE_CASE__									= output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)

						assert_mean_pixel_difference(__UpperCAmelCase      ,			__UpperCAmelCase    )


