




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 		UpperCamelCase__     (lowerCAmelCase__   ,    unittest.TestCase					):

    '''simple docstring'''



    lowerCamelCase_				:     str     					=				KandinskyPipeline
    lowerCamelCase_				:     Any     					=				[
        """prompt""",
        """image_embeds""",
        """negative_image_embeds""",
    ]
    lowerCamelCase_				:     Any     					=				["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
    lowerCamelCase_				:     Optional[int]     					=				[
        """generator""",
        """height""",
        """width""",
        """latents""",
        """guidance_scale""",
        """negative_prompt""",
        """num_inference_steps""",
        """return_dict""",
        """guidance_scale""",
        """num_images_per_prompt""",
        """output_type""",
        """return_dict""",
    ]
    lowerCamelCase_				:     int     					=				False







    @property
    def    _lowercase					(   self					)					->				Any:
           return 32







    @property
    def    _lowercase					(   self					)					->				Optional[int]:
           return 32







    @property
    def    _lowercase					(   self					)					->				Dict:
           return self.time_input_dim







    @property
    def    _lowercase					(   self					)					->				Dict:
           return self.time_input_dim * 4







    @property
    def    _lowercase					(   self					)					->				Union[str, Any]:
           return 100







    @property
    def    _lowercase					(   self					)					->				str:
           lowerCamelCase       :				str      			= XLMRobertaTokenizerFast.from_pretrained("YiYiXu/tiny-random-mclip-base"					)
           return tokenizer







    @property
    def    _lowercase					(   self					)					->				Dict:
           torch.manual_seed(0					)
           lowerCamelCase       :				Tuple      			= 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=1005					,  )

           lowerCamelCase       :				List[str]      			= MultilingualCLIP(UpperCamelCase__					)
           lowerCamelCase       :				Tuple      			= text_encoder.eval()

           return text_encoder







    @property
    def    _lowercase					(   self					)					->				List[Any]:
           torch.manual_seed(0					)

           lowerCamelCase       :				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,
           }

           lowerCamelCase       :				int      			= UNetaDConditionModel(**UpperCamelCase__					)
           return model







    @property
    def    _lowercase					(   self					)					->				Any:
           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					)					->				Any:
           torch.manual_seed(0					)
           lowerCamelCase       :				Union[str, Any]      			= VQModel(**self.dummy_movq_kwargs					)
           return model







    def    _lowercase					(   self					)					->				Tuple:
           lowerCamelCase       :				Optional[int]      			= self.dummy_text_encoder
           lowerCamelCase       :				List[str]      			= self.dummy_tokenizer
           lowerCamelCase       :				Optional[Any]      			= self.dummy_unet
           lowerCamelCase       :				Dict      			= self.dummy_movq

           lowerCamelCase       :				str      			= DDIMScheduler(
               num_train_timesteps=1000					,  beta_schedule="linear"					,  beta_start=0.00085					,  beta_end=0.012					,  clip_sample=UpperCamelCase__					,  set_alpha_to_one=UpperCamelCase__					,  steps_offset=1					,  prediction_type="epsilon"					,  thresholding=UpperCamelCase__					,  )

           lowerCamelCase       :				Optional[int]      			= {
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "unet": unet,
               "scheduler": scheduler,
               "movq": movq,
           }
           return components







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__=0					)					->				Optional[int]:
           lowerCamelCase       :				int      			= floats_tensor((1, self.cross_attention_dim)					,  rng=random.Random(UpperCamelCase__					)					).to(UpperCamelCase__					)
           lowerCamelCase       :				Union[str, Any]      			= floats_tensor((1, self.cross_attention_dim)					,  rng=random.Random(seed + 1					)					).to(UpperCamelCase__					)
           if str(UpperCamelCase__					).startswith("mps"					):
                  lowerCamelCase       :				Any      			= torch.manual_seed(UpperCamelCase__					)
           else:
                  lowerCamelCase       :				Optional[Any]      			= torch.Generator(device=UpperCamelCase__					).manual_seed(UpperCamelCase__					)
           lowerCamelCase       :				Dict      			= {
               "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					)					->				Optional[int]:
           lowerCamelCase       :				Dict      			= "cpu"

           lowerCamelCase       :				Optional[int]      			= self.get_dummy_components()

           lowerCamelCase       :				List[str]      			= self.pipeline_class(**UpperCamelCase__					)
           lowerCamelCase       :				str      			= pipe.to(UpperCamelCase__					)

           pipe.set_progress_bar_config(disable=UpperCamelCase__					)

           lowerCamelCase       :				Optional[int]      			= pipe(**self.get_dummy_inputs(UpperCamelCase__					)					)
           lowerCamelCase       :				Union[str, Any]      			= output.images

           lowerCamelCase       :				Optional[int]      			= pipe(
               **self.get_dummy_inputs(UpperCamelCase__					)					,  return_dict=UpperCamelCase__					,  )[0]

           lowerCamelCase       :				Any      			= image[0, -3:, -3:, -1]
           lowerCamelCase       :				Tuple      			= image_from_tuple[0, -3:, -3:, -1]

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

           lowerCamelCase       :				Optional[int]      			= np.array(
               [0.328663, 1.0, 0.23216873, 1.0, 0.92717564, 0.4639046, 0.96894777, 0.31713378, 0.6293953]					)

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

    '''simple docstring'''



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







    def    _lowercase					(   self					)					->				Any:
           lowerCamelCase       :				Optional[Any]      			= load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/kandinsky/kandinsky_text2img_cat_fp16.npy"					)

           lowerCamelCase       :				Dict      			= KandinskyPriorPipeline.from_pretrained(
               "kandinsky-community/kandinsky-2-1-prior"					,  torch_dtype=torch.floataa					)
           pipe_prior.to(UpperCamelCase__					)

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

           lowerCamelCase       :				Tuple      			= "red cat, 4k photo"

           lowerCamelCase       :				Union[str, Any]      			= torch.Generator(device="cuda"					).manual_seed(0					)
           lowerCamelCase			,				lowerCamelCase       :				Tuple      			= pipe_prior(
               UpperCamelCase__					,  generator=UpperCamelCase__					,  num_inference_steps=5					,  negative_prompt=""					,  ).to_tuple()

           lowerCamelCase       :				Optional[Any]      			= torch.Generator(device="cuda"					).manual_seed(0					)
           lowerCamelCase       :				Tuple      			= pipeline(
               UpperCamelCase__					,  image_embeds=UpperCamelCase__					,  negative_image_embeds=UpperCamelCase__					,  generator=UpperCamelCase__					,  num_inference_steps=100					,  output_type="np"					,  )

           lowerCamelCase       :				Union[str, Any]      			= output.images[0]

           assert image.shape == (512, 512, 3)

           assert_mean_pixel_difference(UpperCamelCase__					,  UpperCamelCase__					)












import types
from typing import List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
from transformers.models.clip.modeling_clip import CLIPTextModelOutput

from diffusers.models import PriorTransformer
from diffusers.pipelines import DiffusionPipeline, StableDiffusionImageVariationPipeline
from diffusers.schedulers import UnCLIPScheduler
from diffusers.utils import logging, randn_tensor


SCREAMING_SNAKE_CASE__	:				List[str]  =       logging.get_logger(__name__)  # pylint: disable=invalid-name

def 					A				(				self			,_SCREAMING_SNAKE_CASE			,_SCREAMING_SNAKE_CASE			,_SCREAMING_SNAKE_CASE			,_SCREAMING_SNAKE_CASE   )     ->			str:
       lowerCamelCase       :				Dict      			= image.to(device=_SCREAMING_SNAKE_CASE   )
       lowerCamelCase       :				str      			= image  # take image as image_embeddings
       lowerCamelCase       :				List[Any]      			= image_embeddings.unsqueeze(1   )

       # duplicate image embeddings for each generation per prompt, using mps friendly method
       lowerCamelCase			,				lowerCamelCase			,				lowerCamelCase       :				str      			= image_embeddings.shape
       lowerCamelCase       :				List[str]      			= image_embeddings.repeat(1			,_SCREAMING_SNAKE_CASE			,1   )
       lowerCamelCase       :				Optional[int]      			= image_embeddings.view(bs_embed * num_images_per_prompt			,_SCREAMING_SNAKE_CASE			,-1   )

       if do_classifier_free_guidance:
              lowerCamelCase       :				Dict      			= torch.zeros_like(_SCREAMING_SNAKE_CASE   )

              # 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       :				Optional[Any]      			= torch.cat([uncond_embeddings, image_embeddings]   )

       return image_embeddings




class 		UpperCamelCase__     (lowerCAmelCase__					):

    '''simple docstring'''



    def __init__(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__ = None					,  )					->				Tuple:
           super().__init__()

           lowerCamelCase       :				str      			= {"image_encoder": None} if decoder_pipe_kwargs is None else decoder_pipe_kwargs

           lowerCamelCase       :				Dict      			= decoder_pipe_kwargs.get("torch_dtype"					,  UpperCamelCase__					) or prior.dtype

           lowerCamelCase       :				Dict      			= StableDiffusionImageVariationPipeline.from_pretrained(
               "lambdalabs/sd-image-variations-diffusers"					,  **UpperCamelCase__					)

           # replace `_encode_image` method
           lowerCamelCase       :				List[str]      			= types.MethodType(_encode_image					,  self.decoder_pipe					)

           self.register_modules(
               prior=UpperCamelCase__					,  tokenizer=UpperCamelCase__					,  text_encoder=UpperCamelCase__					,  prior_scheduler=UpperCamelCase__					,  )







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__ = None					,  UpperCamelCase__ = None					,  )					->				Any:
           if text_model_output is None:
                  lowerCamelCase       :				List[str]      			= len(UpperCamelCase__					) if isinstance(UpperCamelCase__					,  UpperCamelCase__					) else 1
                  # get prompt text embeddings
                  lowerCamelCase       :				Any      			= self.tokenizer(
                      UpperCamelCase__					,  padding="max_length"					,  max_length=self.tokenizer.model_max_length					,  return_tensors="pt"					,  )
                  lowerCamelCase       :				int      			= text_inputs.input_ids
                  lowerCamelCase       :				Optional[Any]      			= text_inputs.attention_mask.bool().to(UpperCamelCase__					)

                  if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
                         lowerCamelCase       :				Optional[int]      			= self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]					)
                         logger.warning(
                             "The following part of your input was truncated because CLIP can only handle sequences up to"
                             F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''					)
                         lowerCamelCase       :				Optional[Any]      			= text_input_ids[:, : self.tokenizer.model_max_length]

                  lowerCamelCase       :				int      			= self.text_encoder(text_input_ids.to(UpperCamelCase__					)					)

                  lowerCamelCase       :				Optional[Any]      			= text_encoder_output.text_embeds
                  lowerCamelCase       :				Optional[int]      			= text_encoder_output.last_hidden_state

           else:
                  lowerCamelCase       :				Dict      			= text_model_output[0].shape[0]
                  lowerCamelCase			,				lowerCamelCase       :				int      			= text_model_output[0], text_model_output[1]
                  lowerCamelCase       :				int      			= text_attention_mask

           lowerCamelCase       :				Dict      			= text_embeddings.repeat_interleave(UpperCamelCase__					,  dim=0					)
           lowerCamelCase       :				int      			= text_encoder_hidden_states.repeat_interleave(UpperCamelCase__					,  dim=0					)
           lowerCamelCase       :				List[str]      			= text_mask.repeat_interleave(UpperCamelCase__					,  dim=0					)

           if do_classifier_free_guidance:
                  lowerCamelCase       :				Dict      			= [""] * batch_size

                  lowerCamelCase       :				List[Any]      			= self.tokenizer(
                      UpperCamelCase__					,  padding="max_length"					,  max_length=self.tokenizer.model_max_length					,  truncation=UpperCamelCase__					,  return_tensors="pt"					,  )
                  lowerCamelCase       :				str      			= uncond_input.attention_mask.bool().to(UpperCamelCase__					)
                  lowerCamelCase       :				int      			= self.text_encoder(uncond_input.input_ids.to(UpperCamelCase__					)					)

                  lowerCamelCase       :				Union[str, Any]      			= uncond_embeddings_text_encoder_output.text_embeds
                  lowerCamelCase       :				List[Any]      			= uncond_embeddings_text_encoder_output.last_hidden_state

                  # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

                  lowerCamelCase       :				Optional[Any]      			= uncond_embeddings.shape[1]
                  lowerCamelCase       :				str      			= uncond_embeddings.repeat(1					,  UpperCamelCase__					)
                  lowerCamelCase       :				Optional[Any]      			= uncond_embeddings.view(batch_size * num_images_per_prompt					,  UpperCamelCase__					)

                  lowerCamelCase       :				Any      			= uncond_text_encoder_hidden_states.shape[1]
                  lowerCamelCase       :				int      			= uncond_text_encoder_hidden_states.repeat(1					,  UpperCamelCase__					,  1					)
                  lowerCamelCase       :				Tuple      			= uncond_text_encoder_hidden_states.view(
                      batch_size * num_images_per_prompt					,  UpperCamelCase__					,  -1					)
                  lowerCamelCase       :				List[str]      			= uncond_text_mask.repeat_interleave(UpperCamelCase__					,  dim=0					)

                  # done duplicates

                  # 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       :				int      			= torch.cat([uncond_embeddings, text_embeddings]					)
                  lowerCamelCase       :				str      			= torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]					)

                  lowerCamelCase       :				Dict      			= torch.cat([uncond_text_mask, text_mask]					)

           return text_embeddings, text_encoder_hidden_states, text_mask







    @property
    def    _lowercase					(   self					)					->				Any:
           if self.device != torch.device("meta"					) or not hasattr(self.prior					,  "_hf_hook"					):
                  return self.device
           for module in self.prior.modules():
                  if (
                      hasattr(UpperCamelCase__					,  "_hf_hook"					)
                      and hasattr(module._hf_hook					,  "execution_device"					)
                      and module._hf_hook.execution_device is not None
                  ):
                         return torch.device(module._hf_hook.execution_device					)
           return self.device







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				Any:
           if latents is None:
                  lowerCamelCase       :				Any      			= randn_tensor(UpperCamelCase__					,  generator=UpperCamelCase__					,  device=UpperCamelCase__					,  dtype=UpperCamelCase__					)
           else:
                  if latents.shape != shape:
                         raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}'''					)
                  lowerCamelCase       :				Any      			= latents.to(UpperCamelCase__					)

           lowerCamelCase       :				List[str]      			= latents * scheduler.init_noise_sigma
           return latents







    def    _lowercase					(   self					,  UpperCamelCase__ = None					)					->				str:
           self.decoder_pipe.to(UpperCamelCase__					)
           super().to(UpperCamelCase__					)







    @torch.no_grad()
    def __call__(   self					,  UpperCamelCase__ = None					,  UpperCamelCase__ = None					,  UpperCamelCase__ = None					,  UpperCamelCase__ = 1					,  UpperCamelCase__ = 25					,  UpperCamelCase__ = None					,  UpperCamelCase__ = None					,  UpperCamelCase__ = None					,  UpperCamelCase__ = None					,  UpperCamelCase__ = 4.0					,  UpperCamelCase__ = 8.0					,  UpperCamelCase__ = 50					,  UpperCamelCase__ = 1					,  UpperCamelCase__ = 0.0					,  UpperCamelCase__ = "pil"					,  UpperCamelCase__ = True					,  )					->				Union[str, Any]:
           if prompt is not None:
                  if isinstance(UpperCamelCase__					,  UpperCamelCase__					):
                         lowerCamelCase       :				Union[str, Any]      			= 1
                  elif isinstance(UpperCamelCase__					,  UpperCamelCase__					):
                         lowerCamelCase       :				int      			= len(UpperCamelCase__					)
                  else:
                         raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(UpperCamelCase__					)}'''					)
           else:
                  lowerCamelCase       :				Any      			= text_model_output[0].shape[0]

           lowerCamelCase       :				str      			= self._execution_device

           lowerCamelCase       :				Optional[Any]      			= batch_size * num_images_per_prompt

           lowerCamelCase       :				int      			= prior_guidance_scale > 1.0 or decoder_guidance_scale > 1.0

           lowerCamelCase			,				lowerCamelCase			,				lowerCamelCase       :				Tuple      			= self._encode_prompt(
               UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)

           # prior

           self.prior_scheduler.set_timesteps(UpperCamelCase__					,  device=UpperCamelCase__					)
           lowerCamelCase       :				Optional[int]      			= self.prior_scheduler.timesteps

           lowerCamelCase       :				List[Any]      			= self.prior.config.embedding_dim

           lowerCamelCase       :				Dict      			= self.prepare_latents(
               (batch_size, embedding_dim)					,  text_embeddings.dtype					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  self.prior_scheduler					,  )

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

                  lowerCamelCase       :				List[Any]      			= self.prior(
                      UpperCamelCase__					,  timestep=UpperCamelCase__					,  proj_embedding=UpperCamelCase__					,  encoder_hidden_states=UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  ).predicted_image_embedding

                  if do_classifier_free_guidance:
                         lowerCamelCase			,				lowerCamelCase       :				List[Any]      			= predicted_image_embedding.chunk(2					)
                         lowerCamelCase       :				Optional[Any]      			= predicted_image_embedding_uncond + prior_guidance_scale * (
                             predicted_image_embedding_text - predicted_image_embedding_uncond
                         )

                  if i + 1 == prior_timesteps_tensor.shape[0]:
                         lowerCamelCase       :				Union[str, Any]      			= None
                  else:
                         lowerCamelCase       :				Dict      			= prior_timesteps_tensor[i + 1]

                  lowerCamelCase       :				Optional[int]      			= self.prior_scheduler.step(
                      UpperCamelCase__					,  timestep=UpperCamelCase__					,  sample=UpperCamelCase__					,  generator=UpperCamelCase__					,  prev_timestep=UpperCamelCase__					,  ).prev_sample

           lowerCamelCase       :				Tuple      			= self.prior.post_process_latents(UpperCamelCase__					)

           lowerCamelCase       :				Any      			= prior_latents

           lowerCamelCase       :				Union[str, Any]      			= self.decoder_pipe(
               image=UpperCamelCase__					,  height=UpperCamelCase__					,  width=UpperCamelCase__					,  num_inference_steps=UpperCamelCase__					,  guidance_scale=UpperCamelCase__					,  generator=UpperCamelCase__					,  output_type=UpperCamelCase__					,  return_dict=UpperCamelCase__					,  num_images_per_prompt=UpperCamelCase__					,  eta=UpperCamelCase__					,  )
           return output






