

import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class  A						(							_UpperCAmelCase    ,							unittest.TestCase ):







    """simple docstring"""





    lowerCamelCase			=							UnCLIPImageVariationPipeline
    lowerCamelCase			=							IMAGE_VARIATION_PARAMS - {"height", "width", "guidance_scale"}
    lowerCamelCase			=							IMAGE_VARIATION_BATCH_PARAMS

    lowerCamelCase			=							[
        "generator",
        "return_dict",
        "decoder_num_inference_steps",
        "super_res_num_inference_steps",
    ]
    lowerCamelCase			=							False


    @property
    def 						snake_case__     (      self :       Dict		)->    Union[str, Any]:





        '''simple docstring'''







        return 3_2


    @property
    def 						snake_case__     (      self :       str		)->    Tuple:





        '''simple docstring'''







        return 3_2


    @property
    def 						snake_case__     (      self :       Union[str, Any]		)->    Optional[int]:





        '''simple docstring'''







        return self.time_input_dim


    @property
    def 						snake_case__     (      self :       Optional[int]		)->    Union[str, Any]:





        '''simple docstring'''







        return self.time_input_dim * 4


    @property
    def 						snake_case__     (      self :       List[str]		)->    Optional[int]:





        '''simple docstring'''







        return 1_0_0


    @property
    def 						snake_case__     (      self :       Optional[Any]		)->    Any:





        '''simple docstring'''







        A__          =   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)
        return tokenizer


    @property
    def 						snake_case__     (      self :       str		)->    Optional[int]:





        '''simple docstring'''







        torch.manual_seed(0		)
        A__          =   CLIPTextConfig(
            bos_token_id=0,eos_token_id=2,hidden_size=self.text_embedder_hidden_size,projection_dim=self.text_embedder_hidden_size,intermediate_size=3_7,layer_norm_eps=1E-05,num_attention_heads=4,num_hidden_layers=5,pad_token_id=1,vocab_size=1_0_0_0,)
        return CLIPTextModelWithProjection(snake_case__		)


    @property
    def 						snake_case__     (      self :       Any		)->    Dict:





        '''simple docstring'''







        torch.manual_seed(0		)
        A__          =   CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size,projection_dim=self.text_embedder_hidden_size,num_hidden_layers=5,num_attention_heads=4,image_size=3_2,intermediate_size=3_7,patch_size=1,)
        return CLIPVisionModelWithProjection(snake_case__		)


    @property
    def 						snake_case__     (      self :       Any		)->    Dict:





        '''simple docstring'''







        torch.manual_seed(0		)

        A__          =   {
            "clip_embeddings_dim": self.text_embedder_hidden_size,
            "time_embed_dim": self.time_embed_dim,
            "cross_attention_dim": self.cross_attention_dim,
        }

        A__          =   UnCLIPTextProjModel(**snake_case__		)
        return model


    @property
    def 						snake_case__     (      self :       str		)->    str:





        '''simple docstring'''







        torch.manual_seed(0		)

        A__          =   {
            "sample_size": 3_2,
            # RGB in channels
            "in_channels": 3,
            # Out channels is double in channels because predicts mean and variance
            "out_channels": 6,
            "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,
            "cross_attention_dim": self.cross_attention_dim,
            "attention_head_dim": 4,
            "resnet_time_scale_shift": "scale_shift",
            "class_embed_type": "identity",
        }

        A__          =   UNetaDConditionModel(**snake_case__		)
        return model


    @property
    def 						snake_case__     (      self :       Tuple		)->    Optional[int]:





        '''simple docstring'''







        return {
            "sample_size": 6_4,
            "layers_per_block": 1,
            "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
            "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "in_channels": 6,
            "out_channels": 3,
        }


    @property
    def 						snake_case__     (      self :       Union[str, Any]		)->    Optional[int]:





        '''simple docstring'''







        torch.manual_seed(0		)

        A__          =   UNetaDModel(**self.dummy_super_res_kwargs		)
        return model


    @property
    def 						snake_case__     (      self :       int		)->    Dict:





        '''simple docstring'''







        torch.manual_seed(1		)

        A__          =   UNetaDModel(**self.dummy_super_res_kwargs		)
        return model


    def 						snake_case__     (      self :       Dict		)->    Optional[int]:





        '''simple docstring'''







        A__          =   self.dummy_decoder
        A__          =   self.dummy_text_proj
        A__          =   self.dummy_text_encoder
        A__          =   self.dummy_tokenizer
        A__          =   self.dummy_super_res_first
        A__          =   self.dummy_super_res_last

        A__          =   UnCLIPScheduler(
            variance_type='learned_range',prediction_type='epsilon',num_train_timesteps=1_0_0_0,)

        A__          =   UnCLIPScheduler(
            variance_type='fixed_small_log',prediction_type='epsilon',num_train_timesteps=1_0_0_0,)

        A__          =   CLIPImageProcessor(crop_size=3_2,size=3_2		)

        A__          =   self.dummy_image_encoder

        return {
            "decoder": decoder,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "text_proj": text_proj,
            "feature_extractor": feature_extractor,
            "image_encoder": image_encoder,
            "super_res_first": super_res_first,
            "super_res_last": super_res_last,
            "decoder_scheduler": decoder_scheduler,
            "super_res_scheduler": super_res_scheduler,
        }


    def 						snake_case__     (      self :       Dict,lowercase_ :       Dict,lowercase_ :       List[Any]=0,lowercase_ :       int=True		)->    Optional[int]:





        '''simple docstring'''







        A__          =   floats_tensor((1, 3, 3_2, 3_2),rng=random.Random(snake_case__		)		).to(snake_case__		)
        if str(snake_case__		).startswith('mps'		):
            A__          =   torch.manual_seed(snake_case__		)
        else:
            A__          =   torch.Generator(device=snake_case__		).manual_seed(snake_case__		)

        if pil_image:
            A__          =   input_image * 0.5 + 0.5
            A__          =   input_image.clamp(0,1		)
            A__          =   input_image.cpu().permute(0,2,3,1		).float().numpy()
            A__          =   DiffusionPipeline.numpy_to_pil(snake_case__		)[0]

        return {
            "image": input_image,
            "generator": generator,
            "decoder_num_inference_steps": 2,
            "super_res_num_inference_steps": 2,
            "output_type": "np",
        }


    def 						snake_case__     (      self :       Union[str, Any]		)->    Tuple:





        '''simple docstring'''







        A__          =   "cpu"

        A__          =   self.get_dummy_components()

        A__          =   self.pipeline_class(**snake_case__		)
        A__          =   pipe.to(snake_case__		)

        pipe.set_progress_bar_config(disable=snake_case__		)

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)

        A__          =   pipe(**snake_case__		)
        A__          =   output.images

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)

        A__          =   pipe(
            **snake_case__,return_dict=snake_case__,)[0]

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

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

        A__          =   np.array(
            [
                0.9_997,
                0.0_002,
                0.9_997,
                0.9_997,
                0.9_969,
                0.0_023,
                0.9_997,
                0.9_969,
                0.9_970,
            ]		)

        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 						snake_case__     (      self :       Tuple		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   "cpu"

        A__          =   self.get_dummy_components()

        A__          =   self.pipeline_class(**snake_case__		)
        A__          =   pipe.to(snake_case__		)

        pipe.set_progress_bar_config(disable=snake_case__		)

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)

        A__          =   pipe(**snake_case__		)
        A__          =   output.images

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)

        A__          =   pipe(
            **snake_case__,return_dict=snake_case__,)[0]

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

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

        A__          =   np.array([0.9_997, 0.0_003, 0.9_997, 0.9_997, 0.9_970, 0.0_024, 0.9_997, 0.9_971, 0.9_971]		)

        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 						snake_case__     (      self :       List[str]		)->    str:





        '''simple docstring'''







        A__          =   "cpu"

        A__          =   self.get_dummy_components()

        A__          =   self.pipeline_class(**snake_case__		)
        A__          =   pipe.to(snake_case__		)

        pipe.set_progress_bar_config(disable=snake_case__		)

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)
        A__          =   [
            pipeline_inputs["image"],
            pipeline_inputs["image"],
        ]

        A__          =   pipe(**snake_case__		)
        A__          =   output.images

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)
        A__          =   [
            tuple_pipeline_inputs["image"],
            tuple_pipeline_inputs["image"],
        ]

        A__          =   pipe(
            **snake_case__,return_dict=snake_case__,)[0]

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

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

        A__          =   np.array(
            [
                0.9_997,
                0.9_989,
                0.0_008,
                0.0_021,
                0.9_960,
                0.0_018,
                0.0_014,
                0.0_002,
                0.9_933,
            ]		)

        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 						snake_case__     (      self :       Union[str, Any]		)->    Any:





        '''simple docstring'''







        A__          =   torch.device('cpu'		)


        class  A						:







            """simple docstring"""





            lowerCamelCase			=							1




        A__          =   self.get_dummy_components()

        A__          =   self.pipeline_class(**snake_case__		)
        A__          =   pipe.to(snake_case__		)

        pipe.set_progress_bar_config(disable=snake_case__		)

        A__          =   torch.Generator(device=snake_case__		).manual_seed(0		)
        A__          =   pipe.decoder.dtype
        A__          =   1

        A__          =   (
            batch_size,
            pipe.decoder.config.in_channels,
            pipe.decoder.config.sample_size,
            pipe.decoder.config.sample_size,
        )
        A__          =   pipe.prepare_latents(
            snake_case__,dtype=snake_case__,device=snake_case__,generator=snake_case__,latents=snake_case__,scheduler=DummyScheduler()		)

        A__          =   (
            batch_size,
            pipe.super_res_first.config.in_channels // 2,
            pipe.super_res_first.config.sample_size,
            pipe.super_res_first.config.sample_size,
        )
        A__          =   pipe.prepare_latents(
            snake_case__,dtype=snake_case__,device=snake_case__,generator=snake_case__,latents=snake_case__,scheduler=DummyScheduler()		)

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)

        A__          =   pipe(
            **snake_case__,decoder_latents=snake_case__,super_res_latents=snake_case__		).images

        A__          =   self.get_dummy_inputs(snake_case__,pil_image=snake_case__		)
        # Don't pass image, instead pass embedding
        A__          =   pipeline_inputs.pop('image'		)
        A__          =   pipe.image_encoder(snake_case__		).image_embeds

        A__          =   pipe(
            **snake_case__,decoder_latents=snake_case__,super_res_latents=snake_case__,image_embeddings=snake_case__,).images

        # make sure passing text embeddings manually is identical
        assert np.abs(img_out_a - img_out_a		).max() < 1E-4


    @skip_mps
    def 						snake_case__     (      self :       str		)->    Dict:





        '''simple docstring'''







        A__          =   torch_device == "cpu"

        # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
        A__          =   1E-2

        self._test_attention_slicing_forward_pass(
            test_max_difference=snake_case__,expected_max_diff=snake_case__		)


    @skip_mps
    def 						snake_case__     (      self :       Dict		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   torch_device == "cpu"
        A__          =   True
        A__          =   [
            "decoder_num_inference_steps",
            "super_res_num_inference_steps",
        ]

        self._test_inference_batch_single_identical(
            test_max_difference=snake_case__,relax_max_difference=snake_case__,additional_params_copy_to_batched_inputs=snake_case__,)


    def 						snake_case__     (      self :       List[Any]		)->    int:





        '''simple docstring'''







        A__          =   [
            "decoder_num_inference_steps",
            "super_res_num_inference_steps",
        ]

        if torch_device == "mps":
            # TODO: MPS errors with larger batch sizes
            A__          =   [2, 3]
            self._test_inference_batch_consistent(
                batch_sizes=snake_case__,additional_params_copy_to_batched_inputs=snake_case__,)
        else:
            self._test_inference_batch_consistent(
                additional_params_copy_to_batched_inputs=snake_case__		)


    @skip_mps
    def 						snake_case__     (      self :       int		)->    Dict:





        '''simple docstring'''







        return super().test_dict_tuple_outputs_equivalent()


    @skip_mps
    def 						snake_case__     (      self :       Tuple		)->    List[str]:





        '''simple docstring'''







        return super().test_save_load_local()


    @skip_mps
    def 						snake_case__     (      self :       Optional[int]		)->    Optional[Any]:





        '''simple docstring'''







        return super().test_save_load_optional_components()




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







    """simple docstring"""





    def 						snake_case__     (      self :       Optional[int]		)->    Dict:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       List[Any]		)->    List[str]:





        '''simple docstring'''







        A__          =   load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png'		)
        A__          =   load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/unclip/karlo_v1_alpha_cat_variation_fp16.npy'		)

        A__          =   UnCLIPImageVariationPipeline.from_pretrained(
            'kakaobrain/karlo-v1-alpha-image-variations',torch_dtype=torch.floataa		)
        A__          =   pipeline.to(snake_case__		)
        pipeline.set_progress_bar_config(disable=snake_case__		)

        A__          =   torch.Generator(device='cpu'		).manual_seed(0		)
        A__          =   pipeline(
            snake_case__,generator=snake_case__,output_type='np',)

        A__          =   output.images[0]

        assert image.shape == (2_5_6, 2_5_6, 3)

        assert_mean_pixel_difference(snake_case__,snake_case__,1_5		)





"""simple docstring"""







import jax.numpy as jnp

from ...utils import logging
from ..ta.modeling_flax_ta import FlaxTaEncoderModel, FlaxTaForConditionalGeneration, FlaxTaModel
from .configuration_mta import MTaConfig


lowerCAmelCase__      				=   logging.get_logger(__name__)

lowerCAmelCase__      				=   '''T5Config'''
def 					a__			(				SCREAMING_SNAKE_CASE    :    jnp.array				,      SCREAMING_SNAKE_CASE    :    int				,      SCREAMING_SNAKE_CASE    :    int				):


  '''simple docstring'''







  lowerCAmelCase     :     List[str]								=							jnp.zeros_like(SCREAMING_SNAKE_CASE				)
  lowerCAmelCase     :     Optional[int]								=							shifted_input_ids.at[:, 1:].set(input_ids[:, :-1]				)
  lowerCAmelCase     :     List[str]								=							shifted_input_ids.at[:, 0].set(SCREAMING_SNAKE_CASE				)

  lowerCAmelCase     :     str								=							jnp.where(shifted_input_ids == -1_0_0				,      SCREAMING_SNAKE_CASE				,      SCREAMING_SNAKE_CASE				)
  return shifted_input_ids


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				List[Any]   ="mt5"
       a  :				Tuple   =MTaConfig


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				Union[str, Any]   ="mt5"
       a  :				Optional[Any]   =MTaConfig


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       a  :				str   ="mt5"
       a  :				Dict   =MTaConfig





