'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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 				UpperCAmelCase (       _SCREAMING_SNAKE_CASE   ,    unittest.TestCase ):




    '''simple docstring'''







    snake_case__							:   Optional[Any]							=						KandinskyImgaImgPipeline
    snake_case__							:   Dict							=						["prompt", "image_embeds", "negative_image_embeds", "image"]
    snake_case__							:   str							=						[
        "prompt",
        "negative_prompt",
        "image_embeds",
        "negative_image_embeds",
        "image",
    ]
    snake_case__							:   List[str]							=						[
        "generator",
        "height",
        "width",
        "strength",
        "guidance_scale",
        "negative_prompt",
        "num_inference_steps",
        "return_dict",
        "guidance_scale",
        "num_images_per_prompt",
        "output_type",
        "return_dict",
    ]
    snake_case__							:   Tuple							=						False






    @property
    def       _UpperCamelCase    (       self )							-> Any:
         return 32






    @property
    def       _UpperCamelCase    (       self )							-> Tuple:
         return 32






    @property
    def       _UpperCamelCase    (       self )							-> Optional[Any]:
         return self.time_input_dim






    @property
    def       _UpperCamelCase    (       self )							-> Union[str, Any]:
         return self.time_input_dim * 4






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






    @property
    def       _UpperCamelCase    (       self )							-> Tuple:
         SCREAMING_SNAKE_CASE					:			int    	=       XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base' )
         return tokenizer






    @property
    def       _UpperCamelCase    (       self )							-> int:
         torch.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       MCLIPConfig(
             numDims=self.cross_attention_dim    ,							transformerDimensions=self.text_embedder_hidden_size    ,							hidden_size=self.text_embedder_hidden_size    ,							intermediate_size=37    ,							num_attention_heads=4    ,							num_hidden_layers=5    ,							vocab_size=1_005    ,							)

         SCREAMING_SNAKE_CASE					:			int    	=       MultilingualCLIP(lowercase__ )
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       text_encoder.eval()

         return text_encoder






    @property
    def       _UpperCamelCase    (       self )							-> List[Any]:
         torch.manual_seed(0 )

         SCREAMING_SNAKE_CASE					:			List[Any]    	=       {
             'in_channels': 4,
             # Out channels is double in channels because predicts mean and variance
             'out_channels': 8,
             'addition_embed_type': 'text_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': 'text_image_proj',
             'cross_attention_dim': self.cross_attention_dim,
             'attention_head_dim': 4,
             'resnet_time_scale_shift': 'scale_shift',
             'class_embed_type': None,
         }

         SCREAMING_SNAKE_CASE					:			List[str]    	=       UNetaDConditionModel(**lowercase__ )
         return model






    @property
    def       _UpperCamelCase    (       self )							-> Union[str, Any]:
         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       _UpperCamelCase    (       self )							-> Dict:
         torch.manual_seed(0 )
         SCREAMING_SNAKE_CASE					:			Tuple    	=       VQModel(**self.dummy_movq_kwargs )
         return model






    def       _UpperCamelCase    (       self )							-> Optional[Any]:
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       self.dummy_text_encoder
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       self.dummy_tokenizer
         SCREAMING_SNAKE_CASE					:			Dict    	=       self.dummy_unet
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       self.dummy_movq

         SCREAMING_SNAKE_CASE					:			List[Any]    	=       {
             'num_train_timesteps': 1_000,
             '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,
         }

         SCREAMING_SNAKE_CASE					:			Tuple    	=       DDIMScheduler(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       {
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
             'unet': unet,
             'scheduler': scheduler,
             'movq': movq,
         }

         return components






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

         if str(lowercase__ ).startswith('mps' ):
              SCREAMING_SNAKE_CASE					:			List[str]    	=       torch.manual_seed(lowercase__ )
         else:
              SCREAMING_SNAKE_CASE					:			str    	=       torch.Generator(device=lowercase__ ).manual_seed(lowercase__ )
         SCREAMING_SNAKE_CASE					:			Any    	=       {
             'prompt': 'horse',
             '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       _UpperCamelCase    (       self )							-> Dict:
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       'cpu'

         SCREAMING_SNAKE_CASE					:			Tuple    	=       self.get_dummy_components()

         SCREAMING_SNAKE_CASE					:			Dict    	=       self.pipeline_class(**lowercase__ )
         SCREAMING_SNAKE_CASE					:			Tuple    	=       pipe.to(lowercase__ )

         pipe.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			int    	=       pipe(**self.get_dummy_inputs(lowercase__ ) )
         SCREAMING_SNAKE_CASE					:			Any    	=       output.images

         SCREAMING_SNAKE_CASE					:			int    	=       pipe(
             **self.get_dummy_inputs(lowercase__ )    ,							return_dict=lowercase__    ,							)[0]

         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       image[0, -3:, -3:, -1]
         SCREAMING_SNAKE_CASE					:			List[str]    	=       image_from_tuple[0, -3:, -3:, -1]

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

         SCREAMING_SNAKE_CASE					:			str    	=       np.array(
             [0.6_1_4_7_4_9_4_3, 0.6_0_7_3_5_3_9, 0.4_3_3_0_8_5_4_4, 0.5_9_2_8_2_6_9, 0.4_7_4_9_3_5_9_5, 0.4_6_7_5_5_9_7_3, 0.4_6_1_3_8_3_8, 0.4_5_3_6_8_7_9_7, 0.5_0_1_1_9_2_3_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 				UpperCAmelCase (       unittest.TestCase ):




    '''simple docstring'''







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






    def       _UpperCamelCase    (       self )							-> Dict:
         SCREAMING_SNAKE_CASE					:			Dict    	=       load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/kandinsky/kandinsky_img2img_frog.npy' )

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

         SCREAMING_SNAKE_CASE					:			int    	=       KandinskyPriorPipeline.from_pretrained(
             'kandinsky-community/kandinsky-2-1-prior'    ,							torch_dtype=torch.floataa )
         pipe_prior.to(lowercase__ )

         SCREAMING_SNAKE_CASE					:			Dict    	=       KandinskyImgaImgPipeline.from_pretrained(
             'kandinsky-community/kandinsky-2-1'    ,							torch_dtype=torch.floataa )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       pipeline.to(lowercase__ )

         pipeline.set_progress_bar_config(disable=lowercase__ )

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       torch.Generator(device='cpu' ).manual_seed(0 )
         SCREAMING_SNAKE_CASE	,   SCREAMING_SNAKE_CASE					:			List[Any]    	=       pipe_prior(
             lowercase__    ,							generator=lowercase__    ,							num_inference_steps=5    ,							negative_prompt=''    ,							).to_tuple()

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       pipeline(
             lowercase__    ,							image=lowercase__    ,							image_embeds=lowercase__    ,							negative_image_embeds=lowercase__    ,							generator=lowercase__    ,							num_inference_steps=100    ,							height=768    ,							width=768    ,							strength=0.2    ,							output_type='np'    ,							)

         SCREAMING_SNAKE_CASE					:			Dict    	=       output.images[0]

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

         assert_mean_pixel_difference(lowercase__    ,							lowercase__ )




'''simple docstring'''





import torch

from diffusers import DDPMParallelScheduler

from .test_schedulers import SchedulerCommonTest






class 				UpperCAmelCase (       _SCREAMING_SNAKE_CASE ):




    '''simple docstring'''







    snake_case__							:   List[str]							=						(DDPMParallelScheduler,)






    def       _UpperCamelCase    (       self    ,							**lowercase__ )							-> Tuple:
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       {
             'num_train_timesteps': 1_000,
             'beta_start': 0.0_0_0_1,
             'beta_end': 0.0_2,
             'beta_schedule': 'linear',
             'variance_type': 'fixed_small',
             'clip_sample': True,
         }

         config.update(**lowercase__ )
         return config






    def       _UpperCamelCase    (       self )							-> int:
         for timesteps in [1, 5, 100, 1_000]:
              self.check_over_configs(num_train_timesteps=lowercase__ )






    def       _UpperCamelCase    (       self )							-> Optional[Any]:
         for beta_start, beta_end in zip([0.0_0_0_1, 0.0_0_1, 0.0_1, 0.1]    ,							[0.0_0_2, 0.0_2, 0.2, 2] ):
              self.check_over_configs(beta_start=lowercase__    ,							beta_end=lowercase__ )






    def       _UpperCamelCase    (       self )							-> List[Any]:
         for schedule in ["linear", "squaredcos_cap_v2"]:
              self.check_over_configs(beta_schedule=lowercase__ )






    def       _UpperCamelCase    (       self )							-> str:
         for variance in ["fixed_small", "fixed_large", "other"]:
              self.check_over_configs(variance_type=lowercase__ )






    def       _UpperCamelCase    (       self )							-> Optional[Any]:
         for clip_sample in [True, False]:
              self.check_over_configs(clip_sample=lowercase__ )






    def       _UpperCamelCase    (       self )							-> List[str]:
         self.check_over_configs(thresholding=lowercase__ )
         for threshold in [0.5, 1.0, 2.0]:
              for prediction_type in ["epsilon", "sample", "v_prediction"]:
                   self.check_over_configs(
                       thresholding=lowercase__    ,							prediction_type=lowercase__    ,							sample_max_value=lowercase__    ,							)






    def       _UpperCamelCase    (       self )							-> Dict:
         for prediction_type in ["epsilon", "sample", "v_prediction"]:
              self.check_over_configs(prediction_type=lowercase__ )






    def       _UpperCamelCase    (       self )							-> Optional[Any]:
         for t in [0, 500, 999]:
              self.check_over_forward(time_step=lowercase__ )






    def       _UpperCamelCase    (       self )							-> Any:
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			Tuple    	=       scheduler_class(**lowercase__ )

         assert torch.sum(torch.abs(scheduler._get_variance(0 ) - 0.0 ) ) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(487 ) - 0.0_0_9_7_9 ) ) < 1E-5
         assert torch.sum(torch.abs(scheduler._get_variance(999 ) - 0.0_2 ) ) < 1E-5






    def       _UpperCamelCase    (       self )							-> int:
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			Tuple    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       len(lowercase__ )

         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       self.dummy_model()
         SCREAMING_SNAKE_CASE					:			str    	=       self.dummy_sample_deter
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       self.dummy_sample_deter + 0.1
         SCREAMING_SNAKE_CASE					:			int    	=       self.dummy_sample_deter - 0.1

         SCREAMING_SNAKE_CASE					:			List[Any]    	=       samplea.shape[0]
         SCREAMING_SNAKE_CASE					:			str    	=       torch.stack([samplea, samplea, samplea]    ,							dim=0 )
         SCREAMING_SNAKE_CASE					:			List[str]    	=       torch.arange(lowercase__ )[0:3, None].repeat(1    ,							lowercase__ )

         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       model(samples.flatten(0    ,							1 )    ,							timesteps.flatten(0    ,							1 ) )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       scheduler.batch_step_no_noise(lowercase__    ,							timesteps.flatten(0    ,							1 )    ,							samples.flatten(0    ,							1 ) )

         SCREAMING_SNAKE_CASE					:			List[Any]    	=       torch.sum(torch.abs(lowercase__ ) )
         SCREAMING_SNAKE_CASE					:			Any    	=       torch.mean(torch.abs(lowercase__ ) )

         assert abs(result_sum.item() - 1_1_5_3.1_8_3_3 ) < 1E-2
         assert abs(result_mean.item() - 0.5_0_0_5 ) < 1E-3






    def       _UpperCamelCase    (       self )							-> Tuple:
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       len(lowercase__ )

         SCREAMING_SNAKE_CASE					:			Any    	=       self.dummy_model()
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       self.dummy_sample_deter
         SCREAMING_SNAKE_CASE					:			Any    	=       torch.manual_seed(0 )

         for t in reversed(range(lowercase__ ) ):
              # 1. predict noise residual
              SCREAMING_SNAKE_CASE					:			int    	=       model(lowercase__    ,							lowercase__ )

              # 2. predict previous mean of sample x_t-1
              SCREAMING_SNAKE_CASE					:			Dict    	=       scheduler.step(lowercase__    ,							lowercase__    ,							lowercase__    ,							generator=lowercase__ ).prev_sample

              SCREAMING_SNAKE_CASE					:			List[str]    	=       pred_prev_sample

         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       torch.sum(torch.abs(lowercase__ ) )
         SCREAMING_SNAKE_CASE					:			Dict    	=       torch.mean(torch.abs(lowercase__ ) )

         assert abs(result_sum.item() - 2_5_8.9_6_0_6 ) < 1E-2
         assert abs(result_mean.item() - 0.3_3_7_2 ) < 1E-3






    def       _UpperCamelCase    (       self )							-> List[str]:
         SCREAMING_SNAKE_CASE					:			Dict    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			Any    	=       self.get_scheduler_config(prediction_type='v_prediction' )
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			str    	=       len(lowercase__ )

         SCREAMING_SNAKE_CASE					:			Tuple    	=       self.dummy_model()
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.dummy_sample_deter
         SCREAMING_SNAKE_CASE					:			int    	=       torch.manual_seed(0 )

         for t in reversed(range(lowercase__ ) ):
              # 1. predict noise residual
              SCREAMING_SNAKE_CASE					:			Tuple    	=       model(lowercase__    ,							lowercase__ )

              # 2. predict previous mean of sample x_t-1
              SCREAMING_SNAKE_CASE					:			Optional[int]    	=       scheduler.step(lowercase__    ,							lowercase__    ,							lowercase__    ,							generator=lowercase__ ).prev_sample

              SCREAMING_SNAKE_CASE					:			List[Any]    	=       pred_prev_sample

         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       torch.sum(torch.abs(lowercase__ ) )
         SCREAMING_SNAKE_CASE					:			str    	=       torch.mean(torch.abs(lowercase__ ) )

         assert abs(result_sum.item() - 2_0_2.0_2_9_6 ) < 1E-2
         assert abs(result_mean.item() - 0.2_6_3_1 ) < 1E-3






    def       _UpperCamelCase    (       self )							-> List[str]:
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			Optional[Any]    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Tuple    	=       [100, 87, 50, 1, 0]

         scheduler.set_timesteps(timesteps=lowercase__ )

         SCREAMING_SNAKE_CASE					:			str    	=       scheduler.timesteps

         for i, timestep in enumerate(lowercase__ ):
              if i == len(lowercase__ ) - 1:
                   SCREAMING_SNAKE_CASE					:			List[str]    	=       -1
              else:
                   SCREAMING_SNAKE_CASE					:			Any    	=       timesteps[i + 1]

              SCREAMING_SNAKE_CASE					:			str    	=       scheduler.previous_timestep(lowercase__ )
              SCREAMING_SNAKE_CASE					:			str    	=       prev_t.item()

              self.assertEqual(lowercase__    ,							lowercase__ )






    def       _UpperCamelCase    (       self )							-> Optional[int]:
         SCREAMING_SNAKE_CASE					:			Tuple    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			List[Any]    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			int    	=       [100, 87, 50, 51, 0]

         with self.assertRaises(lowercase__    ,							msg='`custom_timesteps` must be in descending order.' ):
              scheduler.set_timesteps(timesteps=lowercase__ )






    def       _UpperCamelCase    (       self )							-> Union[str, Any]:
         SCREAMING_SNAKE_CASE					:			Tuple    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			List[str]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			Union[str, Any]    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Tuple    	=       [100, 87, 50, 1, 0]
         SCREAMING_SNAKE_CASE					:			List[str]    	=       len(lowercase__ )

         with self.assertRaises(lowercase__    ,							msg='Can only pass one of `num_inference_steps` or `custom_timesteps`.' ):
              scheduler.set_timesteps(num_inference_steps=lowercase__    ,							timesteps=lowercase__ )






    def       _UpperCamelCase    (       self )							-> List[str]:
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       self.scheduler_classes[0]
         SCREAMING_SNAKE_CASE					:			Optional[int]    	=       self.get_scheduler_config()
         SCREAMING_SNAKE_CASE					:			Tuple    	=       scheduler_class(**lowercase__ )

         SCREAMING_SNAKE_CASE					:			Any    	=       [scheduler.config.num_train_timesteps]

         with self.assertRaises(
             lowercase__    ,							msg='`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}'    ,							):
              scheduler.set_timesteps(timesteps=lowercase__ )



