





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    Union[str, Any]							=   KandinskyVaaImgaImgPipeline
      A_						:    int							=   ['image_embeds', 'negative_image_embeds', 'image']
      A_						:    Optional[Any]							=   [
          'image_embeds',
          'negative_image_embeds',
          'image',
      ]
      A_						:    Optional[Any]							=   [
          'generator',
          'height',
          'width',
          'strength',
          'guidance_scale',
          'num_inference_steps',
          'return_dict',
          'guidance_scale',
          'num_images_per_prompt',
          'output_type',
          'return_dict',
      ]
      A_						:    str							=   False


      @property
      def        a			(self	:				Any				):





         """simple docstring"""



         return 32


      @property
      def        a			(self	:				int				):





         """simple docstring"""



         return 32


      @property
      def        a			(self	:				List[str]				):





         """simple docstring"""



         return self.time_input_dim


      @property
      def        a			(self	:				List[str]				):





         """simple docstring"""



         return self.time_input_dim * 4


      @property
      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         return 100


      @property
      def        a			(self	:				Dict				):





         """simple docstring"""



         torch.manual_seed(0				)

         __snake_case								=       {
             '''in_channels''': 4,
             # Out channels is double in channels because predicts mean and variance
             '''out_channels''': 8,
             '''addition_embed_type''': '''image''',
             '''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,
             '''encoder_hid_dim''': self.text_embedder_hidden_size,
             '''encoder_hid_dim_type''': '''image_proj''',
             '''cross_attention_dim''': self.cross_attention_dim,
             '''attention_head_dim''': 4,
             '''resnet_time_scale_shift''': '''scale_shift''',
             '''class_embed_type''': None,
         }

         __snake_case								=       UNetaDConditionModel(**a__				)
         return model


      @property
      def        a			(self	:				Any				):





         """simple docstring"""



         return {
             "block_out_channels": [32, 64],
             "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
             "in_channels": 3,
             "latent_channels": 4,
             "layers_per_block": 1,
             "norm_num_groups": 8,
             "norm_type": "spatial",
             "num_vq_embeddings": 12,
             "out_channels": 3,
             "up_block_types": [
                 "AttnUpDecoderBlock2D",
                 "UpDecoderBlock2D",
             ],
             "vq_embed_dim": 4,
         }


      @property
      def        a			(self	:				List[Any]				):





         """simple docstring"""



         torch.manual_seed(0				)
         __snake_case								=       VQModel(**self.dummy_movq_kwargs				)
         return model


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       self.dummy_unet
         __snake_case								=       self.dummy_movq

         __snake_case								=       {
             '''num_train_timesteps''': 1000,
             '''beta_schedule''': '''linear''',
             '''beta_start''': 0.0_0_0_8_5,
             '''beta_end''': 0.0_1_2,
             '''clip_sample''': False,
             '''set_alpha_to_one''': False,
             '''steps_offset''': 0,
             '''prediction_type''': '''epsilon''',
             '''thresholding''': False,
         }

         __snake_case								=       DDIMScheduler(**a__				)

         __snake_case								=       {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }

         return components


      def        a			(self	:				Dict    ,      a__	:				Optional[Any]    ,      a__	:				int=0				):





         """simple docstring"""



         __snake_case								=       floats_tensor((1, self.text_embedder_hidden_size)    ,      rng=random.Random(a__				)				).to(a__				)
         __snake_case								=       floats_tensor((1, self.text_embedder_hidden_size)    ,      rng=random.Random(seed + 1				)				).to(
             a__				)
         # create init_image
         __snake_case								=       floats_tensor((1, 3, 64, 64)    ,      rng=random.Random(a__				)				).to(a__				)
         __snake_case								=       image.cpu().permute(0    ,      2    ,      3    ,      1				)[0]
         __snake_case								=       Image.fromarray(np.uinta(a__				)				).convert('''RGB'''				).resize((256, 256)				)

         if str(a__				).startswith('''mps'''				):
            __snake_case								=       torch.manual_seed(a__				)
         else:
            __snake_case								=       torch.Generator(device=a__				).manual_seed(a__				)
         __snake_case								=       {
             '''image''': init_image,
             '''image_embeds''': image_embeds,
             '''negative_image_embeds''': negative_image_embeds,
             '''generator''': generator,
             '''height''': 64,
             '''width''': 64,
             '''num_inference_steps''': 10,
             '''guidance_scale''': 7.0,
             '''strength''': 0.2,
             '''output_type''': '''np''',
         }
         return inputs


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       '''cpu'''

         __snake_case								=       self.get_dummy_components()

         __snake_case								=       self.pipeline_class(**a__				)
         __snake_case								=       pipe.to(a__				)

         pipe.set_progress_bar_config(disable=a__				)

         __snake_case								=       pipe(**self.get_dummy_inputs(a__				)				)
         __snake_case								=       output.images

         __snake_case								=       pipe(
             **self.get_dummy_inputs(a__				)    ,      return_dict=a__    ,      )[0]

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

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

         __snake_case								=       np.array(
             [0.6_1_9_9_7_7_8, 0.6_3_9_8_4_4_0_6, 0.4_6_1_4_5_7_8_5, 0.6_2_9_4_4_9_8_4, 0.5_6_2_2_2_1_5, 0.4_7_3_0_6_1_3_2, 0.4_7_4_4_1_4_5_6, 0.4_6_0_7_6_0_6, 0.4_8_7_1_9_2_6_3]				)
         assert (
             np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2
         ), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
         assert (
             np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1E-2
         ), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""






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


      def        a			(self	:				List[str]				):





         """simple docstring"""



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


      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinskyv22/kandinskyv22_img2img_frog.npy'''				)

         __snake_case								=       load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''				)
         __snake_case								=       '''A red cartoon frog, 4k'''

         __snake_case								=       KandinskyVaaPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-prior'''    ,      torch_dtype=torch.floataa				)
         pipe_prior.to(a__				)

         __snake_case								=       KandinskyVaaImgaImgPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-decoder'''    ,      torch_dtype=torch.floataa				)
         __snake_case								=       pipeline.to(a__				)

         pipeline.set_progress_bar_config(disable=a__				)

         __snake_case								=       torch.Generator(device='''cpu'''				).manual_seed(0				)
         __snake_case	,     __snake_case								=       pipe_prior(
             a__    ,      generator=a__    ,      num_inference_steps=5    ,      negative_prompt=''''''    ,      ).to_tuple()

         __snake_case								=       pipeline(
             image=a__    ,      image_embeds=a__    ,      negative_image_embeds=a__    ,      generator=a__    ,      num_inference_steps=100    ,      height=768    ,      width=768    ,      strength=0.2    ,      output_type='''np'''    ,      )

         __snake_case								=       output.images[0]

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

         assert_mean_pixel_difference(a__    ,      a__				)








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 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    Optional[int]							=   CycleDiffusionPipeline
      A_						:    List[Any]							=   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
          'negative_prompt',
          'height',
          'width',
          'negative_prompt_embeds',
      }
      A_						:    List[str]							=   PipelineTesterMixin.required_optional_params - {'latents'}
      A_						:    List[str]							=   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'source_prompt'}  )
      A_						:    List[str]							=   IMAGE_TO_IMAGE_IMAGE_PARAMS
      A_						:    Any							=   IMAGE_TO_IMAGE_IMAGE_PARAMS


      def        a			(self	:				Dict				):





         """simple docstring"""



         torch.manual_seed(0				)
         __snake_case								=       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    ,      )
         __snake_case								=       DDIMScheduler(
             beta_start=0.0_0_0_8_5    ,      beta_end=0.0_1_2    ,      beta_schedule='''scaled_linear'''    ,      num_train_timesteps=1000    ,      clip_sample=a__    ,      set_alpha_to_one=a__    ,      )
         torch.manual_seed(0				)
         __snake_case								=       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				)
         __snake_case								=       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=1000    ,      )
         __snake_case								=       CLIPTextModel(a__				)
         __snake_case								=       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

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


      def        a			(self	:				List[str]    ,      a__	:				Tuple    ,      a__	:				Optional[Any]=0				):





         """simple docstring"""



         __snake_case								=       floats_tensor((1, 3, 32, 32)    ,      rng=random.Random(a__				)				).to(a__				)
         __snake_case								=       image / 2 + 0.5
         if str(a__				).startswith('''mps'''				):
            __snake_case								=       torch.manual_seed(a__				)
         else:
            __snake_case								=       torch.Generator(device=a__				).manual_seed(a__				)
         __snake_case								=       {
             '''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        a			(self	:				str				):





         """simple docstring"""



         __snake_case								=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator

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

         __snake_case								=       self.get_dummy_inputs(a__				)
         __snake_case								=       pipe(**a__				)
         __snake_case								=       output.images

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

         assert images.shape == (1, 32, 32, 3)
         __snake_case								=       np.array([0.4_4_5_9, 0.4_9_4_3, 0.4_5_4_4, 0.6_6_4_3, 0.5_4_7_4, 0.4_3_2_7, 0.5_7_0_1, 0.5_9_5_9, 0.5_1_7_9]				)

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


      @unittest.skipIf(torch_device != '''cuda'''    ,      '''This test requires a GPU'''				)
      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       self.get_dummy_components()
         for name, module in components.items():
            if hasattr(a__    ,      '''half'''				):
               __snake_case								=       module.half()
         __snake_case								=       CycleDiffusionPipeline(**a__				)
         __snake_case								=       pipe.to(a__				)
         pipe.set_progress_bar_config(disable=a__				)

         __snake_case								=       self.get_dummy_inputs(a__				)
         __snake_case								=       pipe(**a__				)
         __snake_case								=       output.images

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

         assert images.shape == (1, 32, 32, 3)
         __snake_case								=       np.array([0.3_5_0_6, 0.4_5_4_3, 0.4_4_6, 0.4_5_7_5, 0.5_1_9_5, 0.4_1_5_5, 0.5_2_7_3, 0.5_1_8, 0.4_1_1_6]				)

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


      @skip_mps
      def        a			(self	:				Any				):





         """simple docstring"""



         return super().test_save_load_local()


      @unittest.skip('''non-deterministic pipeline'''				)
      def        a			(self	:				Any				):





         """simple docstring"""



         return super().test_inference_batch_single_identical()


      @skip_mps
      def        a			(self	:				Any				):





         """simple docstring"""



         return super().test_dict_tuple_outputs_equivalent()


      @skip_mps
      def        a			(self	:				str				):





         """simple docstring"""



         return super().test_save_load_optional_components()


      @skip_mps
      def        a			(self	:				Dict				):





         """simple docstring"""



         return super().test_attention_slicing_forward_pass()






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


      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



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


      def        a			(self	:				Tuple				):





         """simple docstring"""



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

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

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

         __snake_case								=       '''A black colored car'''
         __snake_case								=       '''A blue colored car'''

         __snake_case								=       torch.manual_seed(0				)
         __snake_case								=       pipe(
             prompt=a__    ,      source_prompt=a__    ,      image=a__    ,      num_inference_steps=100    ,      eta=0.1    ,      strength=0.8_5    ,      guidance_scale=3    ,      source_guidance_scale=1    ,      generator=a__    ,      output_type='''np'''    ,      )
         __snake_case								=       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        a			(self	:				Tuple				):





         """simple docstring"""



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

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

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

         __snake_case								=       '''A black colored car'''
         __snake_case								=       '''A blue colored car'''

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

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

