'''simple docstring'''





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      ):
    a__			:		Union[str, Any]			  =    UnCLIPImageVariationPipeline
    a__			:		Tuple			  =    IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
    a__			:		Union[str, Any]			  =    IMAGE_VARIATION_BATCH_PARAMS

    a__			:		Union[str, Any]			  =    [
        """generator""",
        """return_dict""",
        """decoder_num_inference_steps""",
        """super_res_num_inference_steps""",
    ]
    a__			:		Any			  =    False



    @property
    def    _lowercase		(self       :       int					):
         return 32



    @property
    def    _lowercase		(self       :       Union[str, Any]					):
         return 32



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



    @property
    def    _lowercase		(self       :       Optional[Any]					):
         return self.time_input_dim * 4



    @property
    def    _lowercase		(self       :       Tuple					):
         return 100



    @property
    def    _lowercase		(self       :       List[str]					):
         UpperCAmelCase_			       =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)
         return tokenizer



    @property
    def    _lowercase		(self       :       Optional[int]					):
         torch.manual_seed(0					)
         UpperCAmelCase_			       =  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=37  ,     layer_norm_eps=1E-05  ,     num_attention_heads=4  ,     num_hidden_layers=5  ,     pad_token_id=1  ,     vocab_size=1000  ,     )
         return CLIPTextModelWithProjection(__a					)



    @property
    def    _lowercase		(self       :       Optional[int]					):
         torch.manual_seed(0					)
         UpperCAmelCase_			       =  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=32  ,     intermediate_size=37  ,     patch_size=1  ,     )
         return CLIPVisionModelWithProjection(__a					)



    @property
    def    _lowercase		(self       :       Dict					):
         torch.manual_seed(0					)

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

         UpperCAmelCase_			       =  UnCLIPTextProjModel(**__a					)
         return model



    @property
    def    _lowercase		(self       :       int					):
         torch.manual_seed(0					)

         UpperCAmelCase_			       =  {
             "sample_size": 32,
             # 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",
         }

         UpperCAmelCase_			       =  UNetaDConditionModel(**__a					)
         return model



    @property
    def    _lowercase		(self       :       Union[str, Any]					):
         return {
             "sample_size": 64,
             "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    _lowercase		(self       :       int					):
         torch.manual_seed(0					)

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



    @property
    def    _lowercase		(self       :       Dict					):
         # seeded differently to get different unet than `self.dummy_super_res_first`
         torch.manual_seed(1					)

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



    def    _lowercase		(self       :       Optional[Any]					):
         UpperCAmelCase_			       =  self.dummy_decoder
         UpperCAmelCase_			       =  self.dummy_text_proj
         UpperCAmelCase_			       =  self.dummy_text_encoder
         UpperCAmelCase_			       =  self.dummy_tokenizer
         UpperCAmelCase_			       =  self.dummy_super_res_first
         UpperCAmelCase_			       =  self.dummy_super_res_last

         UpperCAmelCase_			       =  UnCLIPScheduler(
             variance_type="learned_range"  ,     prediction_type="epsilon"  ,     num_train_timesteps=1000  ,     )

         UpperCAmelCase_			       =  UnCLIPScheduler(
             variance_type="fixed_small_log"  ,     prediction_type="epsilon"  ,     num_train_timesteps=1000  ,     )

         UpperCAmelCase_			       =  CLIPImageProcessor(crop_size=32  ,     size=32					)

         UpperCAmelCase_			       =  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    _lowercase		(self       :       Optional[int]  ,     __a       :       Any  ,     __a       :       List[Any]=0  ,     __a       :       str=True					):
         UpperCAmelCase_			       =  floats_tensor((1, 3, 32, 32)  ,     rng=random.Random(__a					)					).to(__a					)
         if str(__a					).startswith("mps"					):
              UpperCAmelCase_			       =  torch.manual_seed(__a					)
         else:
              UpperCAmelCase_			       =  torch.Generator(device=__a					).manual_seed(__a					)

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

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



    def    _lowercase		(self       :       Dict					):
         UpperCAmelCase_			       =  "cpu"

         UpperCAmelCase_			       =  self.get_dummy_components()

         UpperCAmelCase_			       =  self.pipeline_class(**__a					)
         UpperCAmelCase_			       =  pipe.to(__a					)

         pipe.set_progress_bar_config(disable=__a					)

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)

         UpperCAmelCase_			       =  pipe(**__a					)
         UpperCAmelCase_			       =  output.images

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)

         UpperCAmelCase_			       =  pipe(
             **__a  ,     return_dict=__a  ,     )[0]

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

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

         UpperCAmelCase_			       =  np.array(
             [
                 0.99_97,
                 0.00_02,
                 0.99_97,
                 0.99_97,
                 0.99_69,
                 0.00_23,
                 0.99_97,
                 0.99_69,
                 0.99_70,
             ]					)

         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    _lowercase		(self       :       str					):
         UpperCAmelCase_			       =  "cpu"

         UpperCAmelCase_			       =  self.get_dummy_components()

         UpperCAmelCase_			       =  self.pipeline_class(**__a					)
         UpperCAmelCase_			       =  pipe.to(__a					)

         pipe.set_progress_bar_config(disable=__a					)

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)

         UpperCAmelCase_			       =  pipe(**__a					)
         UpperCAmelCase_			       =  output.images

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)

         UpperCAmelCase_			       =  pipe(
             **__a  ,     return_dict=__a  ,     )[0]

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

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

         UpperCAmelCase_			       =  np.array([0.99_97, 0.00_03, 0.99_97, 0.99_97, 0.99_70, 0.00_24, 0.99_97, 0.99_71, 0.99_71]					)

         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    _lowercase		(self       :       List[str]					):
         UpperCAmelCase_			       =  "cpu"

         UpperCAmelCase_			       =  self.get_dummy_components()

         UpperCAmelCase_			       =  self.pipeline_class(**__a					)
         UpperCAmelCase_			       =  pipe.to(__a					)

         pipe.set_progress_bar_config(disable=__a					)

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)
         UpperCAmelCase_			       =  [
             pipeline_inputs["image"],
             pipeline_inputs["image"],
         ]

         UpperCAmelCase_			       =  pipe(**__a					)
         UpperCAmelCase_			       =  output.images

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)
         UpperCAmelCase_			       =  [
             tuple_pipeline_inputs["image"],
             tuple_pipeline_inputs["image"],
         ]

         UpperCAmelCase_			       =  pipe(
             **__a  ,     return_dict=__a  ,     )[0]

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

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

         UpperCAmelCase_			       =  np.array(
             [
                 0.99_97,
                 0.99_89,
                 0.00_08,
                 0.00_21,
                 0.99_60,
                 0.00_18,
                 0.00_14,
                 0.00_02,
                 0.99_33,
             ]					)

         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    _lowercase		(self       :       Tuple					):
         UpperCAmelCase_			       =  torch.device("cpu"					)




         class 			__A :
             a__			:		Optional[Any]			  =    1




         UpperCAmelCase_			       =  self.get_dummy_components()

         UpperCAmelCase_			       =  self.pipeline_class(**__a					)
         UpperCAmelCase_			       =  pipe.to(__a					)

         pipe.set_progress_bar_config(disable=__a					)

         UpperCAmelCase_			       =  torch.Generator(device=__a					).manual_seed(0					)
         UpperCAmelCase_			       =  pipe.decoder.dtype
         UpperCAmelCase_			       =  1

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

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

         UpperCAmelCase_			       =  self.get_dummy_inputs(__a  ,     pil_image=__a					)

         UpperCAmelCase_			       =  pipe(
             **__a  ,     decoder_latents=__a  ,     super_res_latents=__a					).images

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

         UpperCAmelCase_			       =  pipe(
             **__a  ,     decoder_latents=__a  ,     super_res_latents=__a  ,     image_embeddings=__a  ,     ).images

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



    @skip_mps
    def    _lowercase		(self       :       Optional[int]					):
         UpperCAmelCase_			       =  torch_device == "cpu"

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

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



    @skip_mps
    def    _lowercase		(self       :       Optional[Any]					):
         UpperCAmelCase_			       =  torch_device == "cpu"
         UpperCAmelCase_			       =  True
         UpperCAmelCase_			       =  [
             "decoder_num_inference_steps",
             "super_res_num_inference_steps",
         ]

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



    def    _lowercase		(self       :       List[Any]					):
         UpperCAmelCase_			       =  [
             "decoder_num_inference_steps",
             "super_res_num_inference_steps",
         ]

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



    @skip_mps
    def    _lowercase		(self       :       Optional[Any]					):
         return super().test_dict_tuple_outputs_equivalent()



    @skip_mps
    def    _lowercase		(self       :       Optional[Any]					):
         return super().test_save_load_local()






    @skip_mps
    def    _lowercase		(self       :       int					):
         return super().test_save_load_optional_components()




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



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






    def    _lowercase		(self       :       Any					):
         UpperCAmelCase_			       =  load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"					)
         UpperCAmelCase_			       =  load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"					)

         UpperCAmelCase_			       =  UnCLIPImageVariationPipeline.from_pretrained(
             "kakaobrain/karlo-v1-alpha-image-variations"  ,     torch_dtype=torch.floataa					)
         UpperCAmelCase_			       =  pipeline.to(__a					)
         pipeline.set_progress_bar_config(disable=__a					)

         UpperCAmelCase_			       =  torch.Generator(device="cpu"					).manual_seed(0					)
         UpperCAmelCase_			       =  pipeline(
             __a  ,     generator=__a  ,     output_type="np"  ,     )

         UpperCAmelCase_			       =  output.images[0]

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

         assert_mean_pixel_difference(__a  ,     __a  ,     15					)





"""simple docstring"""


import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs






@pytest.mark.parametrize(
    """kwargs, expected""" ,[
        ({"""num_shards""": 0, """max_num_jobs""": 1}, []),
        ({"""num_shards""": 10, """max_num_jobs""": 1}, [range(10		)]),
        ({"""num_shards""": 10, """max_num_jobs""": 10}, [range(lowercase ,i + 1		) for i in range(10		)]),
        ({"""num_shards""": 1, """max_num_jobs""": 10}, [range(1		)]),
        ({"""num_shards""": 10, """max_num_jobs""": 3}, [range(0 ,4		), range(4 ,7		), range(7 ,10		)]),
        ({"""num_shards""": 3, """max_num_jobs""": 10}, [range(0 ,1		), range(1 ,2		), range(2 ,3		)]),
    ] ,)
def      __UpperCAmelCase       (      lowercase ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase    			=   _distribute_shards(**lowercase		)
  assert out == expected






@pytest.mark.parametrize(
    """gen_kwargs, max_num_jobs, expected""" ,[
        ({"""foo""": 0}, 10, [{"""foo""": 0}]),
        ({"""shards""": [0, 1, 2, 3]}, 1, [{"""shards""": [0, 1, 2, 3]}]),
        ({"""shards""": [0, 1, 2, 3]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}, {"""shards""": [2]}, {"""shards""": [3]}]),
        ({"""shards""": [0, 1]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}]),
        ({"""shards""": [0, 1, 2, 3]}, 2, [{"""shards""": [0, 1]}, {"""shards""": [2, 3]}]),
    ] ,)
def      __UpperCAmelCase       (      lowercase ,lowercase ,lowercase		):







  """simple docstring"""
  _UpperCAmelCase    			=   _split_gen_kwargs(lowercase ,lowercase		)
  assert out == expected






@pytest.mark.parametrize(
    """gen_kwargs, expected""" ,[
        ({"""foo""": 0}, 1),
        ({"""shards""": [0]}, 1),
        ({"""shards""": [0, 1, 2, 3]}, 4),
        ({"""shards""": [0, 1, 2, 3], """foo""": 0}, 4),
        ({"""shards""": [0, 1, 2, 3], """other""": (0, 1)}, 4),
        ({"""shards""": [0, 1, 2, 3], """shards2""": [0, 1]}, RuntimeError),
    ] ,)
def      __UpperCAmelCase       (      lowercase ,lowercase		):







  """simple docstring"""
  if expected is RuntimeError:
    with pytest.raises(lowercase		):
      _number_of_shards_in_gen_kwargs(lowercase		)
  else:
    _UpperCAmelCase    			=   _number_of_shards_in_gen_kwargs(lowercase		)
    assert out == expected
