



import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class 		UpperCAmelCase						(						snake_case_       ,snake_case_       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     CycleDiffusionPipeline
   SCREAMING_SNAKE_CASE__							=     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
       '''negative_prompt''',
       '''height''',
       '''width''',
       '''negative_prompt_embeds''',
   }
   SCREAMING_SNAKE_CASE__							=     PipelineTesterMixin.required_optional_params - {'''latents'''}
   SCREAMING_SNAKE_CASE__							=     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''source_prompt'''}				)
   SCREAMING_SNAKE_CASE__							=     IMAGE_TO_IMAGE_IMAGE_PARAMS
   SCREAMING_SNAKE_CASE__							=     IMAGE_TO_IMAGE_IMAGE_PARAMS





   def 	__lowerCAmelCase							(	self			):
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  UNetaDConditionModel(
             block_out_channels=(32, 64)       ,     layers_per_block=2       ,     sample_size=32       ,     in_channels=4       ,     out_channels=4       ,     down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')       ,     up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')       ,     cross_attention_dim=32       ,     )
         _lowerCAmelCase 			=  DDIMScheduler(
             beta_start=0.00_085       ,     beta_end=0.012       ,     beta_schedule='''scaled_linear'''       ,     num_train_timesteps=1_000       ,     clip_sample=_lowerCAmelCase       ,     set_alpha_to_one=_lowerCAmelCase       ,     )
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  AutoencoderKL(
             block_out_channels=[32, 64]       ,     in_channels=3       ,     out_channels=3       ,     down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']       ,     up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']       ,     latent_channels=4       ,     )
         torch.manual_seed(0			)
         _lowerCAmelCase 			=  CLIPTextConfig(
             bos_token_id=0       ,     eos_token_id=2       ,     hidden_size=32       ,     intermediate_size=37       ,     layer_norm_eps=1E-05       ,     num_attention_heads=4       ,     num_hidden_layers=5       ,     pad_token_id=1       ,     vocab_size=1_000       ,     )
         _lowerCAmelCase 			=  CLIPTextModel(_lowerCAmelCase			)
         _lowerCAmelCase 			=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

         _lowerCAmelCase 			=  {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''vae''': vae,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''safety_checker''': None,
             '''feature_extractor''': None,
         }
         return components





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase=0			):
         _lowerCAmelCase 			=  floats_tensor((1, 3, 32, 32)       ,     rng=random.Random(_lowerCAmelCase			)			).to(_lowerCAmelCase			)
         _lowerCAmelCase 			=  image / 2 + 0.5
         if str(_lowerCAmelCase			).startswith('''mps'''			):
               _lowerCAmelCase 			=  torch.manual_seed(_lowerCAmelCase			)
         else:
               _lowerCAmelCase 			=  torch.Generator(device=_lowerCAmelCase			).manual_seed(_lowerCAmelCase			)
         _lowerCAmelCase 			=  {
             '''prompt''': '''An astronaut riding an elephant''',
             '''source_prompt''': '''An astronaut riding a horse''',
             '''image''': image,
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''eta''': 0.1,
             '''strength''': 0.8,
             '''guidance_scale''': 3,
             '''source_guidance_scale''': 1,
             '''output_type''': '''numpy''',
         }
         return inputs





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator

         _lowerCAmelCase 			=  self.get_dummy_components()
         _lowerCAmelCase 			=  CycleDiffusionPipeline(**_lowerCAmelCase			)
         _lowerCAmelCase 			=  pipe.to(_lowerCAmelCase			)
         pipe.set_progress_bar_config(disable=_lowerCAmelCase			)

         _lowerCAmelCase 			=  self.get_dummy_inputs(_lowerCAmelCase			)
         _lowerCAmelCase 			=  pipe(**_lowerCAmelCase			)
         _lowerCAmelCase 			=  output.images

         _lowerCAmelCase 			=  images[0, -3:, -3:, -1]

         assert images.shape == (1, 32, 32, 3)
         _lowerCAmelCase 			=  np.array([0.4_459, 0.4_943, 0.4_544, 0.6_643, 0.5_474, 0.4_327, 0.5_701, 0.5_959, 0.5_179]			)

         assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2





   @unittest.skipIf(torch_device != '''cuda'''       ,     '''This test requires a GPU'''			)
   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.get_dummy_components()
         for name, module in components.items():
               if hasattr(_lowerCAmelCase       ,     '''half'''			):
                     _lowerCAmelCase 			=  module.half()
         _lowerCAmelCase 			=  CycleDiffusionPipeline(**_lowerCAmelCase			)
         _lowerCAmelCase 			=  pipe.to(_lowerCAmelCase			)
         pipe.set_progress_bar_config(disable=_lowerCAmelCase			)

         _lowerCAmelCase 			=  self.get_dummy_inputs(_lowerCAmelCase			)
         _lowerCAmelCase 			=  pipe(**_lowerCAmelCase			)
         _lowerCAmelCase 			=  output.images

         _lowerCAmelCase 			=  images[0, -3:, -3:, -1]

         assert images.shape == (1, 32, 32, 3)
         _lowerCAmelCase 			=  np.array([0.3_506, 0.4_543, 0.446, 0.4_575, 0.5_195, 0.4_155, 0.5_273, 0.518, 0.4_116]			)

         assert np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2





   @skip_mps
   def 	__lowerCAmelCase							(	self			):
         return super().test_save_load_local()





   @unittest.skip('''non-deterministic pipeline'''			)
   def 	__lowerCAmelCase							(	self			):
         return super().test_inference_batch_single_identical()





   @skip_mps
   def 	__lowerCAmelCase							(	self			):
         return super().test_dict_tuple_outputs_equivalent()





   @skip_mps
   def 	__lowerCAmelCase							(	self			):
         return super().test_save_load_optional_components()





   @skip_mps
   def 	__lowerCAmelCase							(	self			):
         return super().test_attention_slicing_forward_pass()






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





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





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/cycle-diffusion/black_colored_car.png'''			)
         _lowerCAmelCase 			=  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy'''			)
         _lowerCAmelCase 			=  init_image.resize((512, 512)			)

         _lowerCAmelCase 			=  '''CompVis/stable-diffusion-v1-4'''
         _lowerCAmelCase 			=  DDIMScheduler.from_pretrained(_lowerCAmelCase       ,     subfolder='''scheduler'''			)
         _lowerCAmelCase 			=  CycleDiffusionPipeline.from_pretrained(
             _lowerCAmelCase       ,     scheduler=_lowerCAmelCase       ,     safety_checker=_lowerCAmelCase       ,     torch_dtype=torch.floataa       ,     revision='''fp16'''			)

         pipe.to(_lowerCAmelCase			)
         pipe.set_progress_bar_config(disable=_lowerCAmelCase			)
         pipe.enable_attention_slicing()

         _lowerCAmelCase 			=  '''A black colored car'''
         _lowerCAmelCase 			=  '''A blue colored car'''

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  pipe(
             prompt=_lowerCAmelCase       ,     source_prompt=_lowerCAmelCase       ,     image=_lowerCAmelCase       ,     num_inference_steps=100       ,     eta=0.1       ,     strength=0.85       ,     guidance_scale=3       ,     source_guidance_scale=1       ,     generator=_lowerCAmelCase       ,     output_type='''np'''       ,     )
         _lowerCAmelCase 			=  output.images

         # the values aren't exactly equal, but the images look the same visually
         assert np.abs(image - expected_image			).max() < 5E-1





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/cycle-diffusion/black_colored_car.png'''			)
         _lowerCAmelCase 			=  load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy'''			)
         _lowerCAmelCase 			=  init_image.resize((512, 512)			)

         _lowerCAmelCase 			=  '''CompVis/stable-diffusion-v1-4'''
         _lowerCAmelCase 			=  DDIMScheduler.from_pretrained(_lowerCAmelCase       ,     subfolder='''scheduler'''			)
         _lowerCAmelCase 			=  CycleDiffusionPipeline.from_pretrained(_lowerCAmelCase       ,     scheduler=_lowerCAmelCase       ,     safety_checker=_lowerCAmelCase			)

         pipe.to(_lowerCAmelCase			)
         pipe.set_progress_bar_config(disable=_lowerCAmelCase			)
         pipe.enable_attention_slicing()

         _lowerCAmelCase 			=  '''A black colored car'''
         _lowerCAmelCase 			=  '''A blue colored car'''

         _lowerCAmelCase 			=  torch.manual_seed(0			)
         _lowerCAmelCase 			=  pipe(
             prompt=_lowerCAmelCase       ,     source_prompt=_lowerCAmelCase       ,     image=_lowerCAmelCase       ,     num_inference_steps=100       ,     eta=0.1       ,     strength=0.85       ,     guidance_scale=3       ,     source_guidance_scale=1       ,     generator=_lowerCAmelCase       ,     output_type='''np'''       ,     )
         _lowerCAmelCase 			=  output.images

         assert np.abs(image - expected_image			).max() < 2E-2




from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding





class 		UpperCAmelCase						(						snake_case_				):
   SCREAMING_SNAKE_CASE__							=     '''ClapFeatureExtractor'''
   SCREAMING_SNAKE_CASE__							=     ('''RobertaTokenizer''', '''RobertaTokenizerFast''')





   def __init__(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase			):
         super().__init__(_lowerCAmelCase       ,     _lowerCAmelCase			)





   def __call__(	self       ,     _lowerCAmelCase=None       ,     _lowerCAmelCase=None       ,     _lowerCAmelCase=None       ,     **_lowerCAmelCase			):
         _lowerCAmelCase 			=  kwargs.pop('''sampling_rate'''       ,     _lowerCAmelCase			)

         if text is None and audios is None:
               raise ValueError('''You have to specify either text or audios. Both cannot be none.'''			)

         if text is not None:
               _lowerCAmelCase 			=  self.tokenizer(_lowerCAmelCase       ,     return_tensors=_lowerCAmelCase       ,     **_lowerCAmelCase			)

         if audios is not None:
               _lowerCAmelCase 			=  self.feature_extractor(
                   _lowerCAmelCase       ,     sampling_rate=_lowerCAmelCase       ,     return_tensors=_lowerCAmelCase       ,     **_lowerCAmelCase			)

         if text is not None and audios is not None:
               _lowerCAmelCase 			=  audio_features.input_features
               return encoding
         elif text is not None:
               return encoding
         else:
               return BatchEncoding(data=dict(**_lowerCAmelCase			)       ,     tensor_type=_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self       ,     *_lowerCAmelCase       ,     **_lowerCAmelCase			):
         return self.tokenizer.batch_decode(*_lowerCAmelCase       ,     **_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self       ,     *_lowerCAmelCase       ,     **_lowerCAmelCase			):
         return self.tokenizer.decode(*_lowerCAmelCase       ,     **_lowerCAmelCase			)





   @property
   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.tokenizer.model_input_names
         _lowerCAmelCase 			=  self.feature_extractor.model_input_names
         return list(dict.fromkeys(tokenizer_input_names + feature_extractor_input_names			)			)