"""simple docstring"""


from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


UpperCAmelCase__     						=  logging.get_logger(__name__)


if is_vision_available():
    import PIL


class 				a (						lowerCAmelCase_      ):
    _snake_case   :     int   							=							['pixel_values']


    def __init__(    self		:     Tuple      , __lowerCAmelCase		:     bool = True      , __lowerCAmelCase		:     Dict[str, int] = None      , __lowerCAmelCase		:     PILImageResampling = PILImageResampling.BICUBIC      , __lowerCAmelCase		:     bool = True      , __lowerCAmelCase		:     Dict[str, int] = None      , __lowerCAmelCase		:     bool = True      , __lowerCAmelCase		:     Union[int, float] = 1 / 255      , __lowerCAmelCase		:     bool = True      , __lowerCAmelCase		:     Optional[Union[float, List[float]]] = None      , __lowerCAmelCase		:     Optional[Union[float, List[float]]] = None      , __lowerCAmelCase		:     bool = True      , **__lowerCAmelCase		:     Optional[int]      , ):
      super().__init__(**__lowerCAmelCase    )
      _UpperCAmelCase         =					size if size is not None else {"""shortest_edge""": 224}
      _UpperCAmelCase         =					get_size_dict(__lowerCAmelCase      , default_to_square=__lowerCAmelCase    )
      _UpperCAmelCase         =					crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
      _UpperCAmelCase         =					get_size_dict(__lowerCAmelCase      , default_to_square=__lowerCAmelCase      , param_name="""crop_size"""    )

      _UpperCAmelCase         =					do_resize
      _UpperCAmelCase         =					size
      _UpperCAmelCase         =					resample
      _UpperCAmelCase         =					do_center_crop
      _UpperCAmelCase         =					crop_size
      _UpperCAmelCase         =					do_rescale
      _UpperCAmelCase         =					rescale_factor
      _UpperCAmelCase         =					do_normalize
      _UpperCAmelCase         =					image_mean if image_mean is not None else OPENAI_CLIP_MEAN
      _UpperCAmelCase         =					image_std if image_std is not None else OPENAI_CLIP_STD
      _UpperCAmelCase         =					do_convert_rgb


    def    lowerCAmelCase_     (    self		:     List[str]      , __lowerCAmelCase		:     np.ndarray      , __lowerCAmelCase		:     Dict[str, int]      , __lowerCAmelCase		:     PILImageResampling = PILImageResampling.BICUBIC      , __lowerCAmelCase		:     Optional[Union[str, ChannelDimension]] = None      , **__lowerCAmelCase		:     List[str]      , ):
      _UpperCAmelCase         =					get_size_dict(__lowerCAmelCase      , default_to_square=__lowerCAmelCase    )
      if "shortest_edge" not in size:
        raise ValueError(f'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''    )
      _UpperCAmelCase         =					get_resize_output_image_size(__lowerCAmelCase      , size=size["""shortest_edge"""]      , default_to_square=__lowerCAmelCase    )
      return resize(__lowerCAmelCase      , size=__lowerCAmelCase      , resample=__lowerCAmelCase      , data_format=__lowerCAmelCase      , **__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[Any]      , __lowerCAmelCase		:     np.ndarray      , __lowerCAmelCase		:     Dict[str, int]      , __lowerCAmelCase		:     Optional[Union[str, ChannelDimension]] = None      , **__lowerCAmelCase		:     List[str]      , ):
      _UpperCAmelCase         =					get_size_dict(__lowerCAmelCase    )
      if "height" not in size or "width" not in size:
        raise ValueError(f'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''    )
      return center_crop(__lowerCAmelCase      , size=(size["""height"""], size["""width"""])      , data_format=__lowerCAmelCase      , **__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     str      , __lowerCAmelCase		:     np.ndarray      , __lowerCAmelCase		:     Union[int, float]      , __lowerCAmelCase		:     Optional[Union[str, ChannelDimension]] = None      , **__lowerCAmelCase		:     Optional[Any]      , ):
      return rescale(__lowerCAmelCase      , scale=__lowerCAmelCase      , data_format=__lowerCAmelCase      , **__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Dict      , __lowerCAmelCase		:     np.ndarray      , __lowerCAmelCase		:     Union[float, List[float]]      , __lowerCAmelCase		:     Union[float, List[float]]      , __lowerCAmelCase		:     Optional[Union[str, ChannelDimension]] = None      , **__lowerCAmelCase		:     str      , ):
      return normalize(__lowerCAmelCase      , mean=__lowerCAmelCase      , std=__lowerCAmelCase      , data_format=__lowerCAmelCase      , **__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     int      , __lowerCAmelCase		:     ImageInput      , __lowerCAmelCase		:     bool = None      , __lowerCAmelCase		:     Dict[str, int] = None      , __lowerCAmelCase		:     PILImageResampling = None      , __lowerCAmelCase		:     bool = None      , __lowerCAmelCase		:     int = None      , __lowerCAmelCase		:     bool = None      , __lowerCAmelCase		:     float = None      , __lowerCAmelCase		:     bool = None      , __lowerCAmelCase		:     Optional[Union[float, List[float]]] = None      , __lowerCAmelCase		:     Optional[Union[float, List[float]]] = None      , __lowerCAmelCase		:     bool = None      , __lowerCAmelCase		:     Optional[Union[str, TensorType]] = None      , __lowerCAmelCase		:     Optional[ChannelDimension] = ChannelDimension.FIRST      , **__lowerCAmelCase		:     Any      , ):
      _UpperCAmelCase         =					do_resize if do_resize is not None else self.do_resize
      _UpperCAmelCase         =					size if size is not None else self.size
      _UpperCAmelCase         =					get_size_dict(__lowerCAmelCase      , param_name="""size"""      , default_to_square=__lowerCAmelCase    )
      _UpperCAmelCase         =					resample if resample is not None else self.resample
      _UpperCAmelCase         =					do_center_crop if do_center_crop is not None else self.do_center_crop
      _UpperCAmelCase         =					crop_size if crop_size is not None else self.crop_size
      _UpperCAmelCase         =					get_size_dict(__lowerCAmelCase      , param_name="""crop_size"""      , default_to_square=__lowerCAmelCase    )
      _UpperCAmelCase         =					do_rescale if do_rescale is not None else self.do_rescale
      _UpperCAmelCase         =					rescale_factor if rescale_factor is not None else self.rescale_factor
      _UpperCAmelCase         =					do_normalize if do_normalize is not None else self.do_normalize
      _UpperCAmelCase         =					image_mean if image_mean is not None else self.image_mean
      _UpperCAmelCase         =					image_std if image_std is not None else self.image_std
      _UpperCAmelCase         =					do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

      _UpperCAmelCase         =					make_list_of_images(__lowerCAmelCase    )

      if not valid_images(__lowerCAmelCase    ):
        raise ValueError(
            """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
            """torch.Tensor, tf.Tensor or jax.ndarray."""    )

      if do_resize and size is None:
        raise ValueError("""Size must be specified if do_resize is True."""    )

      if do_center_crop and crop_size is None:
        raise ValueError("""Crop size must be specified if do_center_crop is True."""    )

      if do_rescale and rescale_factor is None:
        raise ValueError("""Rescale factor must be specified if do_rescale is True."""    )

      if do_normalize and (image_mean is None or image_std is None):
        raise ValueError("""Image mean and std must be specified if do_normalize is True."""    )

      # PIL RGBA images are converted to RGB
      if do_convert_rgb:
        _UpperCAmelCase         =					[convert_to_rgb(__lowerCAmelCase    ) for image in images]

      # All transformations expect numpy arrays.
      _UpperCAmelCase         =					[to_numpy_array(__lowerCAmelCase    ) for image in images]

      if do_resize:
        _UpperCAmelCase         =					[self.resize(image=__lowerCAmelCase      , size=__lowerCAmelCase      , resample=__lowerCAmelCase    ) for image in images]

      if do_center_crop:
        _UpperCAmelCase         =					[self.center_crop(image=__lowerCAmelCase      , size=__lowerCAmelCase    ) for image in images]

      if do_rescale:
        _UpperCAmelCase         =					[self.rescale(image=__lowerCAmelCase      , scale=__lowerCAmelCase    ) for image in images]

      if do_normalize:
        _UpperCAmelCase         =					[self.normalize(image=__lowerCAmelCase      , mean=__lowerCAmelCase      , std=__lowerCAmelCase    ) for image in images]

      _UpperCAmelCase         =					[to_channel_dimension_format(__lowerCAmelCase      , __lowerCAmelCase    ) for image in images]

      _UpperCAmelCase         =					{"""pixel_values""": images}
      return BatchFeature(data=__lowerCAmelCase      , tensor_type=__lowerCAmelCase    )

"""simple docstring"""


import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class 				a (						lowerCAmelCase_		,							lowerCAmelCase_		,							lowerCAmelCase_		,							unittest.TestCase      ):
    _snake_case   :     int   							=							StableUnCLIPPipeline
    _snake_case   :     str   							=							TEXT_TO_IMAGE_PARAMS
    _snake_case   :     Any   							=							TEXT_TO_IMAGE_BATCH_PARAMS
    _snake_case   :     Optional[Any]   							=							TEXT_TO_IMAGE_IMAGE_PARAMS
    _snake_case   :     str   							=							TEXT_TO_IMAGE_IMAGE_PARAMS

    # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
    _snake_case   :     str   							=							False


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					32
      _UpperCAmelCase         =					embedder_hidden_size

      # prior components

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0      , eos_token_id=2      , hidden_size=__lowerCAmelCase      , projection_dim=__lowerCAmelCase      , intermediate_size=37      , layer_norm_eps=1e-0_5      , num_attention_heads=4      , num_hidden_layers=5      , pad_token_id=1      , vocab_size=1000      , )    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					PriorTransformer(
          num_attention_heads=2      , attention_head_dim=12      , embedding_dim=__lowerCAmelCase      , num_layers=1      , )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					DDPMScheduler(
          variance_type="""fixed_small_log"""      , prediction_type="""sample"""      , num_train_timesteps=1000      , clip_sample=__lowerCAmelCase      , clip_sample_range=5.0      , beta_schedule="""squaredcos_cap_v2"""      , )

      # regular denoising components

      torch.manual_seed(0    )
      _UpperCAmelCase         =					StableUnCLIPImageNormalizer(embedding_dim=__lowerCAmelCase    )
      _UpperCAmelCase         =					DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0      , eos_token_id=2      , hidden_size=__lowerCAmelCase      , projection_dim=32      , intermediate_size=37      , layer_norm_eps=1e-0_5      , num_attention_heads=4      , num_hidden_layers=5      , pad_token_id=1      , vocab_size=1000      , )    )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					UNetaDConditionModel(
          sample_size=32      , in_channels=4      , out_channels=4      , down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")      , up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")      , block_out_channels=(32, 64)      , attention_head_dim=(2, 4)      , class_embed_type="""projection"""      , projection_class_embeddings_input_dim=embedder_projection_dim * 2      , cross_attention_dim=__lowerCAmelCase      , layers_per_block=1      , upcast_attention=__lowerCAmelCase      , use_linear_projection=__lowerCAmelCase      , )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					DDIMScheduler(
          beta_schedule="""scaled_linear"""      , beta_start=0.00_085      , beta_end=0.012      , prediction_type="""v_prediction"""      , set_alpha_to_one=__lowerCAmelCase      , steps_offset=1      , )

      torch.manual_seed(0    )
      _UpperCAmelCase         =					AutoencoderKL()

      _UpperCAmelCase         =					{
          # prior components
          """prior_tokenizer""": prior_tokenizer,
          """prior_text_encoder""": prior_text_encoder,
          """prior""": prior,
          """prior_scheduler""": prior_scheduler,
          # image noising components
          """image_normalizer""": image_normalizer,
          """image_noising_scheduler""": image_noising_scheduler,
          # regular denoising components
          """tokenizer""": tokenizer,
          """text_encoder""": text_encoder,
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
      }

      return components


    def    lowerCAmelCase_     (    self		:     Optional[int]      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     str=0    ):
      if str(__lowerCAmelCase    ).startswith("""mps"""    ):
        _UpperCAmelCase         =					torch.manual_seed(__lowerCAmelCase    )
      else:
        _UpperCAmelCase         =					torch.Generator(device=__lowerCAmelCase    ).manual_seed(__lowerCAmelCase    )
      _UpperCAmelCase         =					{
          """prompt""": """A painting of a squirrel eating a burger""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """prior_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     List[str]    ):
      _UpperCAmelCase         =					torch_device in ["""cpu""", """mps"""]

      self._test_inference_batch_single_identical(test_max_difference=__lowerCAmelCase    )



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


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


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase         =					load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""    )

      _UpperCAmelCase         =					StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""      , torch_dtype=torch.floataa    )
      pipe.to(__lowerCAmelCase    )
      pipe.set_progress_bar_config(disable=__lowerCAmelCase    )
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      _UpperCAmelCase         =					torch.Generator(device="""cpu"""    ).manual_seed(0    )
      _UpperCAmelCase         =					pipe("""anime turle"""      , generator=__lowerCAmelCase      , output_type="""np"""    )

      _UpperCAmelCase         =					output.images[0]

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

      assert_mean_pixel_difference(__lowerCAmelCase      , __lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Any    ):
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      _UpperCAmelCase         =					StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""      , torch_dtype=torch.floataa    )
      _UpperCAmelCase         =					pipe.to(__lowerCAmelCase    )
      pipe.set_progress_bar_config(disable=__lowerCAmelCase    )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      _UpperCAmelCase         =					pipe(
          """anime turtle"""      , prior_num_inference_steps=2      , num_inference_steps=2      , output_type="""np"""      , )

      _UpperCAmelCase         =					torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9
