import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp

def 	__lowerCamelCase							(				__a    : List[Any]      ,							__a    : List[Any]      ,							__a    : int      ,							__a    : Optional[int]      ,							__a    : int = 4_096		)      ->       Optional[Any]:
	_lowercase   ,  _lowercase   ,  _lowercase   =key.shape[-3:]
	_lowercase   =value.shape[-1]
	_lowercase   =min(__a      ,							__a		)
	_lowercase   =query / jnp.sqrt(__a		)

	@functools.partial(jax.checkpoint      ,							prevent_cse=__a		)
	def summarize_chunk(__a    : Dict      ,							__a    : Optional[Any]      ,							__a    : Optional[int]		):
		_lowercase   =jnp.einsum("...qhd,...khd->...qhk"      ,							__a      ,							__a      ,							precision=__a		)

		_lowercase   =jnp.max(__a      ,							axis=-1      ,							keepdims=__a		)
		_lowercase   =jax.lax.stop_gradient(__a		)
		_lowercase   =jnp.exp(attn_weights - max_score		)

		_lowercase   =jnp.einsum("...vhf,...qhv->...qhf"      ,							__a      ,							__a      ,							precision=__a		)
		_lowercase   =jnp.einsum("...qhk->...qh"      ,							__a		)

		return (exp_values, exp_weights.sum(axis=-1		), max_score)

	def chunk_scanner(__a    : str		):
		# julienne key array
		_lowercase   =jax.lax.dynamic_slice(
		    operand=__a      ,							start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]      ,							slice_sizes=list(key.shape[:-3]		) + [key_chunk_size, num_heads, k_features]      ,							)

		# julienne value array
		_lowercase   =jax.lax.dynamic_slice(
		    operand=__a      ,							start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]      ,							slice_sizes=list(value.shape[:-3]		) + [key_chunk_size, num_heads, v_features]      ,							)

		return summarize_chunk(__a      ,							__a      ,							__a		)

	_lowercase   ,  _lowercase   ,  _lowercase   =jax.lax.map(f=__a      ,							xs=jnp.arange(0      ,							__a      ,							__a		)		)

	_lowercase   =jnp.max(__a      ,							axis=0      ,							keepdims=__a		)
	_lowercase   =jnp.exp(chunk_max - global_max		)

	chunk_values *= jnp.expand_dims(__a      ,							axis=-1		)
	chunk_weights *= max_diffs

	_lowercase   =chunk_values.sum(axis=0		)
	_lowercase   =jnp.expand_dims(__a      ,							-1		).sum(axis=0		)

	return all_values / all_weights

def 	__lowerCamelCase							(				__a    : Union[str, Any]      ,							__a    : Dict      ,							__a    : int      ,							__a    : Any=jax.lax.Precision.HIGHEST      ,							__a    : int = 1_024      ,							__a    : int = 4_096		)      ->       Any:
	_lowercase   ,  _lowercase   ,  _lowercase   =query.shape[-3:]

	def chunk_scanner(__a    : List[Any]      ,							__a    : Union[str, Any]		):
		# julienne query array
		_lowercase   =jax.lax.dynamic_slice(
		    operand=__a      ,							start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]      ,							slice_sizes=list(query.shape[:-3]		) + [min(__a      ,							__a		), num_heads, q_features]      ,							)

		return (
		    chunk_idx + query_chunk_size,  # unused ignore it
		    _query_chunk_attention(
		        query=__a      ,							key=__a      ,							value=__a      ,							precision=__a      ,							key_chunk_size=__a		),
		)

	_lowercase   ,  _lowercase   =jax.lax.scan(
	    f=__a      ,							init=0      ,							xs=__a      ,							length=math.ceil(num_q / query_chunk_size		)		)

	return jnp.concatenate(__a      ,							axis=-3		)  # fuse the chunked result back
class 	_a  (							nn.Module					):



	"""simple docstring"""



	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					8
	__SCREAMING_SNAKE_CASE              =					64
	__SCREAMING_SNAKE_CASE              =					0.0
	__SCREAMING_SNAKE_CASE              =					False
	__SCREAMING_SNAKE_CASE              =					jnp.floataa


	def __lowerCAmelCase      (   self       ):
		_lowercase   =self.dim_head * self.heads
		_lowercase   =self.dim_head**-0.5

		# Weights were exported with old names {to_q, to_k, to_v, to_out}
		_lowercase   =nn.Dense(lowerCAmelCase_       ,  use_bias=lowerCAmelCase_       ,  dtype=self.dtype       ,  name="to_q"       )
		_lowercase   =nn.Dense(lowerCAmelCase_       ,  use_bias=lowerCAmelCase_       ,  dtype=self.dtype       ,  name="to_k"       )
		_lowercase   =nn.Dense(lowerCAmelCase_       ,  use_bias=lowerCAmelCase_       ,  dtype=self.dtype       ,  name="to_v"       )

		_lowercase   =nn.Dense(self.query_dim       ,  dtype=self.dtype       ,  name="to_out_0"       )
		_lowercase   =nn.Dropout(rate=self.dropout       )


	def __lowerCAmelCase      (   self       ,  lowerCAmelCase_       ):
		_lowercase   ,  _lowercase   ,  _lowercase   =tensor.shape
		_lowercase   =self.heads
		_lowercase   =tensor.reshape(lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  dim // head_size       )
		_lowercase   =jnp.transpose(lowerCAmelCase_       ,  (0, 2, 1, 3)       )
		_lowercase   =tensor.reshape(batch_size * head_size       ,  lowerCAmelCase_       ,  dim // head_size       )
		return tensor


	def __lowerCAmelCase      (   self       ,  lowerCAmelCase_       ):
		_lowercase   ,  _lowercase   ,  _lowercase   =tensor.shape
		_lowercase   =self.heads
		_lowercase   =tensor.reshape(batch_size // head_size       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )
		_lowercase   =jnp.transpose(lowerCAmelCase_       ,  (0, 2, 1, 3)       )
		_lowercase   =tensor.reshape(batch_size // head_size       ,  lowerCAmelCase_       ,  dim * head_size       )
		return tensor


	def __call__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_=None       ,  lowerCAmelCase_=True       ):
		_lowercase   =hidden_states if context is None else context

		_lowercase   =self.query(lowerCAmelCase_       )
		_lowercase   =self.key(lowerCAmelCase_       )
		_lowercase   =self.value(lowerCAmelCase_       )

		_lowercase   =self.reshape_heads_to_batch_dim(lowerCAmelCase_       )
		_lowercase   =self.reshape_heads_to_batch_dim(lowerCAmelCase_       )
		_lowercase   =self.reshape_heads_to_batch_dim(lowerCAmelCase_       )

		if self.use_memory_efficient_attention:
			_lowercase   =query_states.transpose(1       ,  0       ,  2       )
			_lowercase   =key_states.transpose(1       ,  0       ,  2       )
			_lowercase   =value_states.transpose(1       ,  0       ,  2       )

			# this if statement create a chunk size for each layer of the unet
			# the chunk size is equal to the query_length dimension of the deepest layer of the unet

			_lowercase   =query_states.shape[-3]
			if flatten_latent_dim % 64 == 0:
				_lowercase   =int(flatten_latent_dim / 64       )
			elif flatten_latent_dim % 16 == 0:
				_lowercase   =int(flatten_latent_dim / 16       )
			elif flatten_latent_dim % 4 == 0:
				_lowercase   =int(flatten_latent_dim / 4       )
			else:
				_lowercase   =int(lowerCAmelCase_       )

			_lowercase   =jax_memory_efficient_attention(
			    lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  query_chunk_size=lowerCAmelCase_       ,  key_chunk_size=4096 * 4       )

			_lowercase   =hidden_states.transpose(1       ,  0       ,  2       )
		else:
			# compute attentions
			_lowercase   =jnp.einsum("b i d, b j d->b i j"       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )
			_lowercase   =attention_scores * self.scale
			_lowercase   =nn.softmax(lowerCAmelCase_       ,  axis=2       )

			# attend to values
			_lowercase   =jnp.einsum("b i j, b j d -> b i d"       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )

		_lowercase   =self.reshape_batch_dim_to_heads(lowerCAmelCase_       )
		_lowercase   =self.proj_attn(lowerCAmelCase_       )
		return self.dropout_layer(lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )

class 	_a  (							nn.Module					):



	"""simple docstring"""



	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					0.0
	__SCREAMING_SNAKE_CASE              =					False
	__SCREAMING_SNAKE_CASE              =					jnp.floataa
	__SCREAMING_SNAKE_CASE              =					False


	def __lowerCAmelCase      (   self       ):
		# self attention (or cross_attention if only_cross_attention is True)
		_lowercase   =FlaxAttention(
		    self.dim       ,  self.n_heads       ,  self.d_head       ,  self.dropout       ,  self.use_memory_efficient_attention       ,  dtype=self.dtype       )
		# cross attention
		_lowercase   =FlaxAttention(
		    self.dim       ,  self.n_heads       ,  self.d_head       ,  self.dropout       ,  self.use_memory_efficient_attention       ,  dtype=self.dtype       )
		_lowercase   =FlaxFeedForward(dim=self.dim       ,  dropout=self.dropout       ,  dtype=self.dtype       )
		_lowercase   =nn.LayerNorm(epsilon=1e-5       ,  dtype=self.dtype       )
		_lowercase   =nn.LayerNorm(epsilon=1e-5       ,  dtype=self.dtype       )
		_lowercase   =nn.LayerNorm(epsilon=1e-5       ,  dtype=self.dtype       )
		_lowercase   =nn.Dropout(rate=self.dropout       )


	def __call__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_=True       ):
		# self attention
		_lowercase   =hidden_states
		if self.only_cross_attention:
			_lowercase   =self.attna(self.norma(lowerCAmelCase_       )       ,  lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )
		else:
			_lowercase   =self.attna(self.norma(lowerCAmelCase_       )       ,  deterministic=lowerCAmelCase_       )
		_lowercase   =hidden_states + residual

		# cross attention
		_lowercase   =hidden_states
		_lowercase   =self.attna(self.norma(lowerCAmelCase_       )       ,  lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )
		_lowercase   =hidden_states + residual

		# feed forward
		_lowercase   =hidden_states
		_lowercase   =self.ff(self.norma(lowerCAmelCase_       )       ,  deterministic=lowerCAmelCase_       )
		_lowercase   =hidden_states + residual

		return self.dropout_layer(lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )

class 	_a  (							nn.Module					):



	"""simple docstring"""



	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					1
	__SCREAMING_SNAKE_CASE              =					0.0
	__SCREAMING_SNAKE_CASE              =					False
	__SCREAMING_SNAKE_CASE              =					False
	__SCREAMING_SNAKE_CASE              =					jnp.floataa
	__SCREAMING_SNAKE_CASE              =					False


	def __lowerCAmelCase      (   self       ):
		_lowercase   =nn.GroupNorm(num_groups=32       ,  epsilon=1e-5       )

		_lowercase   =self.n_heads * self.d_head
		if self.use_linear_projection:
			_lowercase   =nn.Dense(lowerCAmelCase_       ,  dtype=self.dtype       )
		else:
			_lowercase   =nn.Conv(
			    lowerCAmelCase_       ,  kernel_size=(1, 1)       ,  strides=(1, 1)       ,  padding="VALID"       ,  dtype=self.dtype       ,  )

		_lowercase   =[
		    FlaxBasicTransformerBlock(
		        lowerCAmelCase_       ,  self.n_heads       ,  self.d_head       ,  dropout=self.dropout       ,  only_cross_attention=self.only_cross_attention       ,  dtype=self.dtype       ,  use_memory_efficient_attention=self.use_memory_efficient_attention       ,  )
		    for _ in range(self.depth       )
		]

		if self.use_linear_projection:
			_lowercase   =nn.Dense(lowerCAmelCase_       ,  dtype=self.dtype       )
		else:
			_lowercase   =nn.Conv(
			    lowerCAmelCase_       ,  kernel_size=(1, 1)       ,  strides=(1, 1)       ,  padding="VALID"       ,  dtype=self.dtype       ,  )

		_lowercase   =nn.Dropout(rate=self.dropout       )


	def __call__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_=True       ):
		_lowercase   ,  _lowercase   ,  _lowercase   ,  _lowercase   =hidden_states.shape
		_lowercase   =hidden_states
		_lowercase   =self.norm(lowerCAmelCase_       )
		if self.use_linear_projection:
			_lowercase   =hidden_states.reshape(lowerCAmelCase_       ,  height * width       ,  lowerCAmelCase_       )
			_lowercase   =self.proj_in(lowerCAmelCase_       )
		else:
			_lowercase   =self.proj_in(lowerCAmelCase_       )
			_lowercase   =hidden_states.reshape(lowerCAmelCase_       ,  height * width       ,  lowerCAmelCase_       )

		for transformer_block in self.transformer_blocks:
			_lowercase   =transformer_block(lowerCAmelCase_       ,  lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )

		if self.use_linear_projection:
			_lowercase   =self.proj_out(lowerCAmelCase_       )
			_lowercase   =hidden_states.reshape(lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )
		else:
			_lowercase   =hidden_states.reshape(lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       ,  lowerCAmelCase_       )
			_lowercase   =self.proj_out(lowerCAmelCase_       )

		_lowercase   =hidden_states + residual
		return self.dropout_layer(lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )

class 	_a  (							nn.Module					):



	"""simple docstring"""



	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					0.0
	__SCREAMING_SNAKE_CASE              =					jnp.floataa


	def __lowerCAmelCase      (   self       ):
		# The second linear layer needs to be called
		# net_2 for now to match the index of the Sequential layer
		_lowercase   =FlaxGEGLU(self.dim       ,  self.dropout       ,  self.dtype       )
		_lowercase   =nn.Dense(self.dim       ,  dtype=self.dtype       )


	def __call__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_=True       ):
		_lowercase   =self.net_a(lowerCAmelCase_       ,  deterministic=lowerCAmelCase_       )
		_lowercase   =self.net_a(lowerCAmelCase_       )
		return hidden_states

class 	_a  (							nn.Module					):



	"""simple docstring"""



	__SCREAMING_SNAKE_CASE              =					42
	__SCREAMING_SNAKE_CASE              =					0.0
	__SCREAMING_SNAKE_CASE              =					jnp.floataa


	def __lowerCAmelCase      (   self       ):
		_lowercase   =self.dim * 4
		_lowercase   =nn.Dense(inner_dim * 2       ,  dtype=self.dtype       )
		_lowercase   =nn.Dropout(rate=self.dropout       )


	def __call__(   self       ,  lowerCAmelCase_       ,  lowerCAmelCase_=True       ):
		_lowercase   =self.proj(lowerCAmelCase_       )
		_lowercase   ,  _lowercase   =jnp.split(lowerCAmelCase_       ,  2       ,  axis=2       )
		return self.dropout_layer(hidden_linear * nn.gelu(lowerCAmelCase_       )       ,  deterministic=lowerCAmelCase_       )










"""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  _SCREAMING_SNAKE_CASE       (     SCREAMING_SNAKE_CASE__			,				unittest.TestCase				):



 '''simple docstring'''





 __UpperCAmelCase		   =							KandinskyPipeline
 __UpperCAmelCase		   =							[
     """prompt""",
     """image_embeds""",
     """negative_image_embeds""",
 ]
 __UpperCAmelCase		   =							["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
 __UpperCAmelCase		   =							[
     """generator""",
     """height""",
     """width""",
     """latents""",
     """guidance_scale""",
     """negative_prompt""",
     """num_inference_steps""",
     """return_dict""",
     """guidance_scale""",
     """num_images_per_prompt""",
     """output_type""",
     """return_dict""",
 ]
 __UpperCAmelCase		   =							False






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  return 32






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  return 32






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  return self.time_input_dim






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  return self.time_input_dim * 4






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  return 1_00






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Tuple					      =  XLMRobertaTokenizerFast.from_pretrained("YiYiXu/tiny-random-mclip-base"   )
  return tokenizer






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  torch.manual_seed(0   )
  _UpperCamelCase      :	List[str]					      =  MCLIPConfig(
      numDims=self.cross_attention_dim					,     transformerDimensions=self.text_embedder_hidden_size					,     hidden_size=self.text_embedder_hidden_size					,     intermediate_size=37					,     num_attention_heads=4					,     num_hidden_layers=5					,     vocab_size=10_05					,     )

  _UpperCamelCase      :	Optional[Any]					      =  MultilingualCLIP(lowerCAmelCase__   )
  _UpperCamelCase      :	List[str]					      =  text_encoder.eval()

  return text_encoder






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  torch.manual_seed(0   )

  _UpperCamelCase      :	Dict					      =  {
      "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,
  }

  _UpperCamelCase      :	int					      =  UNetaDConditionModel(**lowerCAmelCase__   )
  return model






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  return {
      "block_out_channels": [32, 64],
      "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": 12,
      "out_channels": 3,
      "up_block_types": [
          "AttnUpDecoderBlock2D",
          "UpDecoderBlock2D",
      ],
      "vq_embed_dim": 4,
  }






 @property
 def lowercase_      (self   ):

  '''simple docstring'''







  torch.manual_seed(0   )
  _UpperCamelCase      :	Dict					      =  VQModel(**self.dummy_movq_kwargs   )
  return model






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	List[str]					      =  self.dummy_text_encoder
  _UpperCamelCase      :	str					      =  self.dummy_tokenizer
  _UpperCamelCase      :	List[str]					      =  self.dummy_unet
  _UpperCamelCase      :	Tuple					      =  self.dummy_movq

  _UpperCamelCase      :	Dict					      =  DDIMScheduler(
      num_train_timesteps=10_00					,     beta_schedule="linear"					,     beta_start=0.0_0085					,     beta_end=0.012					,     clip_sample=lowerCAmelCase__					,     set_alpha_to_one=lowerCAmelCase__					,     steps_offset=1					,     prediction_type="epsilon"					,     thresholding=lowerCAmelCase__					,     )

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






 def lowercase_      (self					,     lowerCAmelCase__					,     lowerCAmelCase__=0   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  floats_tensor((1, self.cross_attention_dim)					,     rng=random.Random(lowerCAmelCase__   )   ).to(lowerCAmelCase__   )
  _UpperCamelCase      :	int					      =  floats_tensor((1, self.cross_attention_dim)					,     rng=random.Random(seed + 1   )   ).to(lowerCAmelCase__   )
  if str(lowerCAmelCase__   ).startswith("mps"   ):
   _UpperCamelCase      :	Optional[Any]					      =  torch.manual_seed(lowerCAmelCase__   )
  else:
   _UpperCamelCase      :	List[Any]					      =  torch.Generator(device=lowerCAmelCase__   ).manual_seed(lowerCAmelCase__   )
  _UpperCamelCase      :	Optional[Any]					      =  {
      "prompt": "horse",
      "image_embeds": image_embeds,
      "negative_image_embeds": negative_image_embeds,
      "generator": generator,
      "height": 64,
      "width": 64,
      "guidance_scale": 4.0,
      "num_inference_steps": 2,
      "output_type": "np",
  }
  return inputs






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Optional[Any]					      =  "cpu"

  _UpperCamelCase      :	Any					      =  self.get_dummy_components()

  _UpperCamelCase      :	Tuple					      =  self.pipeline_class(**lowerCAmelCase__   )
  _UpperCamelCase      :	Optional[int]					      =  pipe.to(lowerCAmelCase__   )

  pipe.set_progress_bar_config(disable=lowerCAmelCase__   )

  _UpperCamelCase      :	Union[str, Any]					      =  pipe(**self.get_dummy_inputs(lowerCAmelCase__   )   )
  _UpperCamelCase      :	Union[str, Any]					      =  output.images

  _UpperCamelCase      :	str					      =  pipe(
      **self.get_dummy_inputs(lowerCAmelCase__   )					,     return_dict=lowerCAmelCase__					,     )[0]

  _UpperCamelCase      :	Tuple					      =  image[0, -3:, -3:, -1]
  _UpperCamelCase      :	Optional[int]					      =  image_from_tuple[0, -3:, -3:, -1]

  assert image.shape == (1, 64, 64, 3)

  _UpperCamelCase      :	Union[str, Any]					      =  np.array(
      [0.32_8663, 1.0, 0.2321_6873, 1.0, 0.9271_7564, 0.463_9046, 0.9689_4777, 0.3171_3378, 0.629_3953]   )

  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  _SCREAMING_SNAKE_CASE       (     unittest.TestCase				):



 '''simple docstring'''





 def lowercase_      (self   ):

  '''simple docstring'''







  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Tuple					      =  load_numpy(
      "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
      "/kandinsky/kandinsky_text2img_cat_fp16.npy"   )

  _UpperCamelCase      :	Optional[int]					      =  KandinskyPriorPipeline.from_pretrained(
      "kandinsky-community/kandinsky-2-1-prior"					,     torch_dtype=torch.floataa   )
  pipe_prior.to(lowerCAmelCase__   )

  _UpperCamelCase      :	List[str]					      =  KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1"					,     torch_dtype=torch.floataa   )
  _UpperCamelCase      :	List[Any]					      =  pipeline.to(lowerCAmelCase__   )
  pipeline.set_progress_bar_config(disable=lowerCAmelCase__   )

  _UpperCamelCase      :	Optional[Any]					      =  "red cat, 4k photo"

  _UpperCamelCase      :	str					      =  torch.Generator(device="cuda"   ).manual_seed(0   )
  _UpperCamelCase      ,	_UpperCamelCase      :	Optional[Any]					      =  pipe_prior(
      lowerCAmelCase__					,     generator=lowerCAmelCase__					,     num_inference_steps=5					,     negative_prompt=""					,     ).to_tuple()

  _UpperCamelCase      :	Optional[Any]					      =  torch.Generator(device="cuda"   ).manual_seed(0   )
  _UpperCamelCase      :	Any					      =  pipeline(
      lowerCAmelCase__					,     image_embeds=lowerCAmelCase__					,     negative_image_embeds=lowerCAmelCase__					,     generator=lowerCAmelCase__					,     num_inference_steps=1_00					,     output_type="np"					,     )

  _UpperCamelCase      :	Optional[Any]					      =  output.images[0]

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

  assert_mean_pixel_difference(lowerCAmelCase__					,     lowerCAmelCase__   )



