




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   __a(					_a    ,						unittest.TestCase       ):






    """simple docstring"""





    lowerCAmelCase        =				KandinskyVaaImgaImgPipeline
    lowerCAmelCase        =				['''image_embeds''', '''negative_image_embeds''', '''image''']
    lowerCAmelCase        =				[
        '''image_embeds''',
        '''negative_image_embeds''',
        '''image''',
    ]
    lowerCAmelCase        =				[
        '''generator''',
        '''height''',
        '''width''',
        '''strength''',
        '''guidance_scale''',
        '''num_inference_steps''',
        '''return_dict''',
        '''guidance_scale''',
        '''num_images_per_prompt''',
        '''output_type''',
        '''return_dict''',
    ]
    lowerCAmelCase        =				False





    @property
    def 							a__				(				self  )			->		Optional[Any]:
         return 32





    @property
    def 							a__				(				self  )			->		str:
         return 32





    @property
    def 							a__				(				self  )			->		Tuple:
         return self.time_input_dim





    @property
    def 							a__				(				self  )			->		str:
         return self.time_input_dim * 4





    @property
    def 							a__				(				self  )			->		Any:
         return 100





    @property
    def 							a__				(				self  )			->		Tuple:
         torch.manual_seed(0  )

         UpperCAmelCase_      :							Union[str, Any]       =      {
             '''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,
         }

         UpperCAmelCase_      :							int       =      UNetaDConditionModel(**_SCREAMING_SNAKE_CASE  )
         return model





    @property
    def 							a__				(				self  )			->		Optional[int]:
         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[str]:
         torch.manual_seed(0  )
         UpperCAmelCase_      :							Tuple       =      VQModel(**self.dummy_movq_kwargs  )
         return model





    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							Dict       =      self.dummy_unet
         UpperCAmelCase_      :							str       =      self.dummy_movq

         UpperCAmelCase_      :							List[str]       =      {
             '''num_train_timesteps''': 1_000,
             '''beta_schedule''': '''linear''',
             '''beta_start''': 0.0_00_85,
             '''beta_end''': 0.0_12,
             '''clip_sample''': False,
             '''set_alpha_to_one''': False,
             '''steps_offset''': 0,
             '''prediction_type''': '''epsilon''',
             '''thresholding''': False,
         }

         UpperCAmelCase_      :							Optional[Any]       =      DDIMScheduler(**_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							int       =      {
             '''unet''': unet,
             '''scheduler''': scheduler,
             '''movq''': movq,
         }

         return components





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE=0  )			->		List[str]:
         UpperCAmelCase_      :							Optional[Any]       =      floats_tensor((1, self.text_embedder_hidden_size)       ,rng=random.Random(_SCREAMING_SNAKE_CASE  )  ).to(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Union[str, Any]       =      floats_tensor((1, self.text_embedder_hidden_size)       ,rng=random.Random(seed + 1  )  ).to(
             _SCREAMING_SNAKE_CASE  )
         # create init_image
         UpperCAmelCase_      :							Optional[int]       =      floats_tensor((1, 3, 64, 64)       ,rng=random.Random(_SCREAMING_SNAKE_CASE  )  ).to(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      image.cpu().permute(0       ,2       ,3       ,1  )[0]
         UpperCAmelCase_      :							Optional[Any]       =      Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE  )  ).convert('''RGB'''  ).resize((256, 256)  )

         if str(_SCREAMING_SNAKE_CASE  ).startswith('''mps'''  ):
              UpperCAmelCase_      :							str       =      torch.manual_seed(_SCREAMING_SNAKE_CASE  )
         else:
              UpperCAmelCase_      :							int       =      torch.Generator(device=_SCREAMING_SNAKE_CASE  ).manual_seed(_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      {
             '''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  )			->		Optional[Any]:
         UpperCAmelCase_      :							Optional[Any]       =      '''cpu'''

         UpperCAmelCase_      :							Optional[Any]       =      self.get_dummy_components()

         UpperCAmelCase_      :							Dict       =      self.pipeline_class(**_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Any       =      pipe.to(_SCREAMING_SNAKE_CASE  )

         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							List[str]       =      pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE  )  )
         UpperCAmelCase_      :							Optional[int]       =      output.images

         UpperCAmelCase_      :							Tuple       =      pipe(
             **self.get_dummy_inputs(_SCREAMING_SNAKE_CASE  )       ,return_dict=_SCREAMING_SNAKE_CASE       ,)[0]

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

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

         UpperCAmelCase_      :							List[Any]       =      np.array(
             [0.6_19_97_78, 0.63_98_44_06, 0.46_14_57_85, 0.62_94_49_84, 0.5_62_22_15, 0.47_30_61_32, 0.47_44_14_56, 0.4_60_76_06, 0.48_71_92_63]  )
         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   __a(					unittest.TestCase       ):






    """simple docstring"""





    def 							a__				(				self  )			->		Union[str, Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





    def 							a__				(				self  )			->		Optional[int]:
         UpperCAmelCase_      :							List[str]       =      load_numpy(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/kandinskyv22/kandinskyv22_img2img_frog.npy'''  )

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

         UpperCAmelCase_      :							int       =      KandinskyVaaPriorPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-prior'''       ,torch_dtype=torch.floataa  )
         pipe_prior.to(_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Tuple       =      KandinskyVaaImgaImgPipeline.from_pretrained(
             '''kandinsky-community/kandinsky-2-2-decoder'''       ,torch_dtype=torch.floataa  )
         UpperCAmelCase_      :							Any       =      pipeline.to(_SCREAMING_SNAKE_CASE  )

         pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Any       =      torch.Generator(device='''cpu'''  ).manual_seed(0  )
         UpperCAmelCase_,   UpperCAmelCase_      :							int       =      pipe_prior(
             _SCREAMING_SNAKE_CASE       ,generator=_SCREAMING_SNAKE_CASE       ,num_inference_steps=5       ,negative_prompt=''''''       ,).to_tuple()

         UpperCAmelCase_      :							str       =      pipeline(
             image=_SCREAMING_SNAKE_CASE       ,image_embeds=_SCREAMING_SNAKE_CASE       ,negative_image_embeds=_SCREAMING_SNAKE_CASE       ,generator=_SCREAMING_SNAKE_CASE       ,num_inference_steps=100       ,height=768       ,width=768       ,strength=0.2       ,output_type='''np'''       ,)

         UpperCAmelCase_      :							List[str]       =      output.images[0]

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

         assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )





import tempfile

import torch

from diffusers import IPNDMScheduler

from .test_schedulers import SchedulerCommonTest

class   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				(IPNDMScheduler,)
    lowerCAmelCase        =				(('''num_inference_steps''', 50),)





    def 							a__				(				self       ,**_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							Tuple       =      {'''num_train_timesteps''': 1_000}
         config.update(**_SCREAMING_SNAKE_CASE  )
         return config





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE=0       ,**_SCREAMING_SNAKE_CASE  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[int]       =      dict(self.forward_default_kwargs  )
         UpperCAmelCase_      :							int       =      kwargs.pop('''num_inference_steps'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							str       =      self.dummy_sample
         UpperCAmelCase_      :							Optional[Any]       =      0.1 * sample
         UpperCAmelCase_      :							Optional[int]       =      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

         for scheduler_class in self.scheduler_classes:
              UpperCAmelCase_      :							Tuple       =      self.get_scheduler_config(**_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Optional[int]       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )
              scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )
              # copy over dummy past residuals
              UpperCAmelCase_      :							Optional[Any]       =      dummy_past_residuals[:]

              if time_step is None:
                   UpperCAmelCase_      :							Tuple       =      scheduler.timesteps[len(scheduler.timesteps  ) // 2]

              with tempfile.TemporaryDirectory() as tmpdirname:
                   scheduler.save_config(_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							List[Any]       =      scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE  )
                   new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )
                   # copy over dummy past residuals
                   UpperCAmelCase_      :							Optional[int]       =      dummy_past_residuals[:]

              UpperCAmelCase_      :							List[Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							str       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"

              UpperCAmelCase_      :							Union[str, Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Optional[int]       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"





    def 							a__				(				self  )			->		List[Any]:
         pass





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE=0       ,**_SCREAMING_SNAKE_CASE  )			->		List[Any]:
         UpperCAmelCase_      :							List[str]       =      dict(self.forward_default_kwargs  )
         UpperCAmelCase_      :							Tuple       =      kwargs.pop('''num_inference_steps'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      self.dummy_sample
         UpperCAmelCase_      :							Tuple       =      0.1 * sample
         UpperCAmelCase_      :							int       =      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

         for scheduler_class in self.scheduler_classes:
              UpperCAmelCase_      :							str       =      self.get_scheduler_config()
              UpperCAmelCase_      :							str       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )
              scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )

              # copy over dummy past residuals (must be after setting timesteps)
              UpperCAmelCase_      :							Tuple       =      dummy_past_residuals[:]

              if time_step is None:
                   UpperCAmelCase_      :							Dict       =      scheduler.timesteps[len(scheduler.timesteps  ) // 2]

              with tempfile.TemporaryDirectory() as tmpdirname:
                   scheduler.save_config(_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							Optional[int]       =      scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE  )
                   # copy over dummy past residuals
                   new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )

                   # copy over dummy past residual (must be after setting timesteps)
                   UpperCAmelCase_      :							str       =      dummy_past_residuals[:]

              UpperCAmelCase_      :							Tuple       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Optional[int]       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"

              UpperCAmelCase_      :							str       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							List[str]       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"





    def 							a__				(				self       ,**_SCREAMING_SNAKE_CASE  )			->		List[Any]:
         UpperCAmelCase_      :							Tuple       =      self.scheduler_classes[0]
         UpperCAmelCase_      :							List[Any]       =      self.get_scheduler_config(**_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Tuple       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Any       =      10
         UpperCAmelCase_      :							Tuple       =      self.dummy_model()
         UpperCAmelCase_      :							List[str]       =      self.dummy_sample_deter
         scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )

         for i, t in enumerate(scheduler.timesteps  ):
              UpperCAmelCase_      :							int       =      model(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Union[str, Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ).prev_sample

         for i, t in enumerate(scheduler.timesteps  ):
              UpperCAmelCase_      :							Optional[int]       =      model(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Dict       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ).prev_sample

         return sample





    def 							a__				(				self  )			->		str:
         UpperCAmelCase_      :							List[Any]       =      dict(self.forward_default_kwargs  )

         UpperCAmelCase_      :							Any       =      kwargs.pop('''num_inference_steps'''       ,_SCREAMING_SNAKE_CASE  )

         for scheduler_class in self.scheduler_classes:
              UpperCAmelCase_      :							int       =      self.get_scheduler_config()
              UpperCAmelCase_      :							List[Any]       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )

              UpperCAmelCase_      :							Optional[int]       =      self.dummy_sample
              UpperCAmelCase_      :							List[Any]       =      0.1 * sample

              if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE       ,'''set_timesteps'''  ):
                   scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )
              elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE       ,'''set_timesteps'''  ):
                   UpperCAmelCase_      :							Union[str, Any]       =      num_inference_steps

              # copy over dummy past residuals (must be done after set_timesteps)
              UpperCAmelCase_      :							str       =      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
              UpperCAmelCase_      :							int       =      dummy_past_residuals[:]

              UpperCAmelCase_      :							List[Any]       =      scheduler.timesteps[5]
              UpperCAmelCase_      :							Dict       =      scheduler.timesteps[6]

              UpperCAmelCase_      :							List[Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Dict       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              self.assertEqual(output_a.shape       ,sample.shape  )
              self.assertEqual(output_a.shape       ,output_a.shape  )

              UpperCAmelCase_      :							List[Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Dict       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              self.assertEqual(output_a.shape       ,sample.shape  )
              self.assertEqual(output_a.shape       ,output_a.shape  )





    def 							a__				(				self  )			->		Any:
         for timesteps in [100, 1_000]:
              self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE       ,time_step=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Union[str, Any]:
         for t, num_inference_steps in zip([1, 5, 10]       ,[10, 50, 100]  ):
              self.check_over_forward(num_inference_steps=_SCREAMING_SNAKE_CASE       ,time_step=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[int]       =      self.full_loop()
         UpperCAmelCase_      :							str       =      torch.mean(torch.abs(_SCREAMING_SNAKE_CASE  )  )

         assert abs(result_mean.item() - 2_540_529  ) < 10