
'''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 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   _snake_case		   =							UnCLIPImageVariationPipeline
   _snake_case		   =							IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''guidance_scale'''}
   _snake_case		   =							IMAGE_VARIATION_BATCH_PARAMS

   _snake_case		   =							[
       '''generator''',
       '''return_dict''',
       '''decoder_num_inference_steps''',
       '''super_res_num_inference_steps''',
   ]
   _snake_case		   =							False







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 3_2







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 3_2







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return self.time_input_dim







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return self.time_input_dim * 4







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 1_0_0







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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=3_7 ,   layer_norm_eps=1e-0_5 ,   num_attention_heads=4 ,   num_hidden_layers=5 ,   pad_token_id=1 ,   vocab_size=1_0_0_0 ,   )
      return CLIPTextModelWithProjection(lowerCamelCase__   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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=3_2 ,   intermediate_size=3_7 ,   patch_size=1 ,   )
      return CLIPVisionModelWithProjection(lowerCamelCase__   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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(**lowerCamelCase__   )
      return model







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      torch.manual_seed(0   )

      UpperCamelCase       =  {
          '''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''',
      }

      UpperCamelCase       =  UNetaDConditionModel(**lowerCamelCase__   )
      return model







   @property
   def       UpperCAmelCase				(			self   ):

      '''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       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      torch.manual_seed(0   )

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







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      torch.manual_seed(1   )

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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=1_0_0_0 ,   )

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

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

      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       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__=0 ,   lowerCamelCase__=True   ):

      '''simple docstring'''







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

      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(lowerCamelCase__   )[0]

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  '''cpu'''

      UpperCamelCase       =  self.get_dummy_components()

      UpperCamelCase       =  self.pipeline_class(**lowerCamelCase__   )
      UpperCamelCase       =  pipe.to(lowerCamelCase__   )

      pipe.set_progress_bar_config(disable=lowerCamelCase__   )

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )

      UpperCamelCase       =  pipe(**lowerCamelCase__   )
      UpperCamelCase       =  output.images

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )

      UpperCamelCase       =  pipe(
          **lowerCamelCase__ ,   return_dict=lowerCamelCase__ ,   )[0]

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

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

      UpperCamelCase       =  np.array(
          [
              0.9997,
              0.0002,
              0.9997,
              0.9997,
              0.9969,
              0.0023,
              0.9997,
              0.9969,
              0.9970,
          ]   )

      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       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  '''cpu'''

      UpperCamelCase       =  self.get_dummy_components()

      UpperCamelCase       =  self.pipeline_class(**lowerCamelCase__   )
      UpperCamelCase       =  pipe.to(lowerCamelCase__   )

      pipe.set_progress_bar_config(disable=lowerCamelCase__   )

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )

      UpperCamelCase       =  pipe(**lowerCamelCase__   )
      UpperCamelCase       =  output.images

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )

      UpperCamelCase       =  pipe(
          **lowerCamelCase__ ,   return_dict=lowerCamelCase__ ,   )[0]

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

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

      UpperCamelCase       =  np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]   )

      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       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  '''cpu'''

      UpperCamelCase       =  self.get_dummy_components()

      UpperCamelCase       =  self.pipeline_class(**lowerCamelCase__   )
      UpperCamelCase       =  pipe.to(lowerCamelCase__   )

      pipe.set_progress_bar_config(disable=lowerCamelCase__   )

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )
      UpperCamelCase       =  [
          pipeline_inputs['''image'''],
          pipeline_inputs['''image'''],
      ]

      UpperCamelCase       =  pipe(**lowerCamelCase__   )
      UpperCamelCase       =  output.images

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )
      UpperCamelCase       =  [
          tuple_pipeline_inputs['''image'''],
          tuple_pipeline_inputs['''image'''],
      ]

      UpperCamelCase       =  pipe(
          **lowerCamelCase__ ,   return_dict=lowerCamelCase__ ,   )[0]

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

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

      UpperCamelCase       =  np.array(
          [
              0.9997,
              0.9989,
              0.0008,
              0.0021,
              0.9960,
              0.0018,
              0.0014,
              0.0002,
              0.9933,
          ]   )

      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       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  torch.device('''cpu'''   )


      class 						lowercase__	:






         '''simple docstring'''







         _snake_case		   =							1




      UpperCamelCase       =  self.get_dummy_components()

      UpperCamelCase       =  self.pipeline_class(**lowerCamelCase__   )
      UpperCamelCase       =  pipe.to(lowerCamelCase__   )

      pipe.set_progress_bar_config(disable=lowerCamelCase__   )

      UpperCamelCase       =  torch.Generator(device=lowerCamelCase__   ).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(
          lowerCamelCase__ ,   dtype=lowerCamelCase__ ,   device=lowerCamelCase__ ,   generator=lowerCamelCase__ ,   latents=lowerCamelCase__ ,   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(
          lowerCamelCase__ ,   dtype=lowerCamelCase__ ,   device=lowerCamelCase__ ,   generator=lowerCamelCase__ ,   latents=lowerCamelCase__ ,   scheduler=DummyScheduler()   )

      UpperCamelCase       =  self.get_dummy_inputs(lowerCamelCase__ ,   pil_image=lowerCamelCase__   )

      UpperCamelCase       =  pipe(
          **lowerCamelCase__ ,   decoder_latents=lowerCamelCase__ ,   super_res_latents=lowerCamelCase__   ).images

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

      UpperCamelCase       =  pipe(
          **lowerCamelCase__ ,   decoder_latents=lowerCamelCase__ ,   super_res_latents=lowerCamelCase__ ,   image_embeddings=lowerCamelCase__ ,   ).images

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







   @skip_mps
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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=lowerCamelCase__ ,   expected_max_diff=lowerCamelCase__   )







   @skip_mps
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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=lowerCamelCase__ ,   relax_max_difference=lowerCamelCase__ ,   additional_params_copy_to_batched_inputs=lowerCamelCase__ ,   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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=lowerCamelCase__ ,   additional_params_copy_to_batched_inputs=lowerCamelCase__ ,   )
      else:
         self._test_inference_batch_consistent(
             additional_params_copy_to_batched_inputs=lowerCamelCase__   )







   @skip_mps
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return super().test_dict_tuple_outputs_equivalent()







   @skip_mps
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return super().test_save_load_local()







   @skip_mps
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return super().test_save_load_optional_components()




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






   '''simple docstring'''







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      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(lowerCamelCase__   )
      pipeline.set_progress_bar_config(disable=lowerCamelCase__   )

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

      UpperCamelCase       =  output.images[0]

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

      assert_mean_pixel_difference(lowerCamelCase__ ,   lowerCamelCase__ ,   1_5   )







'''simple docstring'''


import math
from typing import Optional

import numpy as np

from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case_ :  Dict     =					logging.get_logger(__name__)

snake_case_ :  Union[str, Any]     =					{
    'facebook/encodec_24khz': 'https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json',
    'facebook/encodec_48khz': 'https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json',
}


class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   _snake_case		   =							'''encodec'''







   def __init__(			self ,   lowerCamelCase__=[1.5, 3.0, 6.0, 12.0, 24.0] ,   lowerCamelCase__=2_4_0_0_0 ,   lowerCamelCase__=1 ,   lowerCamelCase__=False ,   lowerCamelCase__=None ,   lowerCamelCase__=None ,   lowerCamelCase__=1_2_8 ,   lowerCamelCase__=3_2 ,   lowerCamelCase__=1 ,   lowerCamelCase__=[8, 5, 4, 2] ,   lowerCamelCase__="weight_norm" ,   lowerCamelCase__=7 ,   lowerCamelCase__=7 ,   lowerCamelCase__=3 ,   lowerCamelCase__=2 ,   lowerCamelCase__=True ,   lowerCamelCase__="reflect" ,   lowerCamelCase__=2 ,   lowerCamelCase__=2 ,   lowerCamelCase__=1.0 ,   lowerCamelCase__=1_0_2_4 ,   lowerCamelCase__=None ,   lowerCamelCase__=True ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  target_bandwidths
      UpperCamelCase       =  sampling_rate
      UpperCamelCase       =  audio_channels
      UpperCamelCase       =  normalize
      UpperCamelCase       =  chunk_length_s
      UpperCamelCase       =  overlap
      UpperCamelCase       =  hidden_size
      UpperCamelCase       =  num_filters
      UpperCamelCase       =  num_residual_layers
      UpperCamelCase       =  upsampling_ratios
      UpperCamelCase       =  norm_type
      UpperCamelCase       =  kernel_size
      UpperCamelCase       =  last_kernel_size
      UpperCamelCase       =  residual_kernel_size
      UpperCamelCase       =  dilation_growth_rate
      UpperCamelCase       =  use_causal_conv
      UpperCamelCase       =  pad_mode
      UpperCamelCase       =  compress
      UpperCamelCase       =  num_lstm_layers
      UpperCamelCase       =  trim_right_ratio
      UpperCamelCase       =  codebook_size
      UpperCamelCase       =  codebook_dim if codebook_dim is not None else hidden_size
      UpperCamelCase       =  use_conv_shortcut

      if self.norm_type not in ["weight_norm", "time_group_norm"]:
         raise ValueError(
             f'self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'   )

      super().__init__(**lowerCamelCase__   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      if self.chunk_length_s is None:
         return None
      else:
         return int(self.chunk_length_s * self.sampling_rate   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      if self.chunk_length_s is None or self.overlap is None:
         return None
      else:
         return max(1 ,   int((1.0 - self.overlap) * self.chunk_length   )   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  np.prod(self.upsampling_ratios   )
      return math.ceil(self.sampling_rate / hop_length   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return int(1_0_0_0 * self.target_bandwidths[-1] // (self.frame_rate * 1_0)   )





