





"""simple docstring"""






import random
import unittest

import numpy as np
import torch
from PIL import Image
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorEmbaEmbPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 						lowerCAmelCase       (   snake_case      ,						unittest.TestCase				):
       lowerCAmelCase__    			=	KandinskyVaaPriorEmbaEmbPipeline
       lowerCAmelCase__    			=	["""prompt""", """image"""]
       lowerCAmelCase__    			=	["""prompt""", """image"""]
       lowerCAmelCase__    			=	[
           """num_images_per_prompt""",
           """strength""",
           """generator""",
           """num_inference_steps""",
           """latents""",
           """negative_prompt""",
           """guidance_scale""",
           """output_type""",
           """return_dict""",
       ]
       lowerCAmelCase__    			=	False






       @property
       def 				__A			( self  ):
              return 32






       @property
       def 				__A			( self  ):
              return 32






       @property
       def 				__A			( self  ):
              return self.time_input_dim






       @property
       def 				__A			( self  ):
              return self.time_input_dim * 4






       @property
       def 				__A			( self  ):
              return 1_00






       @property
       def 				__A			( self  ):
              _UpperCAmelCase			     =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )
              return tokenizer






       @property
       def 				__A			( self  ):
              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=37     ,	layer_norm_eps=1E-05     ,	num_attention_heads=4     ,	num_hidden_layers=5     ,	pad_token_id=1     ,	vocab_size=10_00     ,	)
              return CLIPTextModelWithProjection(a__  )






       @property
       def 				__A			( self  ):
              torch.manual_seed(0  )

              _UpperCAmelCase			     =  {
                  'num_attention_heads': 2,
                  'attention_head_dim': 12,
                  'embedding_dim': self.text_embedder_hidden_size,
                  'num_layers': 1,
              }

              _UpperCAmelCase			     =  PriorTransformer(**a__  )
              # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
              _UpperCAmelCase			     =  nn.Parameter(torch.ones(model.clip_std.shape  )  )
              return model






       @property
       def 				__A			( self  ):
              torch.manual_seed(0  )
              _UpperCAmelCase			     =  CLIPVisionConfig(
                  hidden_size=self.text_embedder_hidden_size     ,	image_size=2_24     ,	projection_dim=self.text_embedder_hidden_size     ,	intermediate_size=37     ,	num_attention_heads=4     ,	num_channels=3     ,	num_hidden_layers=5     ,	patch_size=14     ,	)

              _UpperCAmelCase			     =  CLIPVisionModelWithProjection(a__  )
              return model






       @property
       def 				__A			( self  ):
              _UpperCAmelCase			     =  CLIPImageProcessor(
                  crop_size=2_24     ,	do_center_crop=a__     ,	do_normalize=a__     ,	do_resize=a__     ,	image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]     ,	image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]     ,	resample=3     ,	size=2_24     ,	)

              return image_processor






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.dummy_prior
              _UpperCAmelCase			     =  self.dummy_image_encoder
              _UpperCAmelCase			     =  self.dummy_text_encoder
              _UpperCAmelCase			     =  self.dummy_tokenizer
              _UpperCAmelCase			     =  self.dummy_image_processor

              _UpperCAmelCase			     =  UnCLIPScheduler(
                  variance_type='fixed_small_log'     ,	prediction_type='sample'     ,	num_train_timesteps=10_00     ,	clip_sample=a__     ,	clip_sample_range=10.0     ,	)

              _UpperCAmelCase			     =  {
                  'prior': prior,
                  'image_encoder': image_encoder,
                  'text_encoder': text_encoder,
                  'tokenizer': tokenizer,
                  'scheduler': scheduler,
                  'image_processor': image_processor,
              }

              return components






       def 				__A			( self     ,	a__     ,	a__=0  ):
              if str(a__  ).startswith('mps'  ):
                     _UpperCAmelCase			     =  torch.manual_seed(a__  )
              else:
                     _UpperCAmelCase			     =  torch.Generator(device=a__  ).manual_seed(a__  )

              _UpperCAmelCase			     =  floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(a__  )  ).to(a__  )
              _UpperCAmelCase			     =  image.cpu().permute(0     ,	2     ,	3     ,	1  )[0]
              _UpperCAmelCase			     =  Image.fromarray(np.uinta(a__  )  ).convert('RGB'  ).resize((2_56, 2_56)  )

              _UpperCAmelCase			     =  {
                  'prompt': 'horse',
                  'image': init_image,
                  'strength': 0.5,
                  'generator': generator,
                  'guidance_scale': 4.0,
                  'num_inference_steps': 2,
                  'output_type': 'np',
              }
              return inputs






       def 				__A			( self  ):
              _UpperCAmelCase			     =  'cpu'

              _UpperCAmelCase			     =  self.get_dummy_components()

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

              pipe.set_progress_bar_config(disable=a__  )

              _UpperCAmelCase			     =  pipe(**self.get_dummy_inputs(a__  )  )
              _UpperCAmelCase			     =  output.image_embeds

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

              _UpperCAmelCase			     =  image[0, -10:]
              _UpperCAmelCase			     =  image_from_tuple[0, -10:]

              assert image.shape == (1, 32)

              _UpperCAmelCase			     =  np.array(
                  [
                      0.1_071_284,
                      1.3_330_271,
                      0.61_260_223,
                      -0.6_691_065,
                      -0.3_846_852,
                      -1.0_303_661,
                      0.22_716_111,
                      0.03_348_901,
                      0.30_040_675,
                      -0.24_805_029,
                  ]  )

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






       @skip_mps
       def 				__A			( self  ):
              _UpperCAmelCase			     =  torch_device == 'cpu'
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  False

              self._test_inference_batch_single_identical(
                  test_max_difference=a__     ,	relax_max_difference=a__     ,	test_mean_pixel_difference=a__     ,	)






       @skip_mps
       def 				__A			( self  ):
              _UpperCAmelCase			     =  torch_device == 'cpu'
              _UpperCAmelCase			     =  False

              self._test_attention_slicing_forward_pass(
                  test_max_difference=a__     ,	test_mean_pixel_difference=a__     ,	)













"""simple docstring"""






import torch

from diffusers import LMSDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class 						lowerCAmelCase       (   snake_case				):
       lowerCAmelCase__    			=	(LMSDiscreteScheduler,)
       lowerCAmelCase__    			=	10






       def 				__A			( self     ,	**a__  ):
              _UpperCAmelCase			     =  {
                  'num_train_timesteps': 11_00,
                  'beta_start': 0.0_001,
                  'beta_end': 0.02,
                  'beta_schedule': 'linear',
              }

              config.update(**a__  )
              return config






       def 				__A			( self  ):
              for timesteps in [10, 50, 1_00, 10_00]:
                     self.check_over_configs(num_train_timesteps=a__  )






       def 				__A			( self  ):
              for beta_start, beta_end in zip([0.00_001, 0.0_001, 0.001]     ,	[0.0_002, 0.002, 0.02]  ):
                     self.check_over_configs(beta_start=a__     ,	beta_end=a__  )






       def 				__A			( self  ):
              for schedule in ["linear", "scaled_linear"]:
                     self.check_over_configs(beta_schedule=a__  )






       def 				__A			( self  ):
              for prediction_type in ["epsilon", "v_prediction"]:
                     self.check_over_configs(prediction_type=a__  )






       def 				__A			( self  ):
              for t in [0, 5_00, 8_00]:
                     self.check_over_forward(time_step=a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.scheduler_classes[0]
              _UpperCAmelCase			     =  self.get_scheduler_config()
              _UpperCAmelCase			     =  scheduler_class(**a__  )

              scheduler.set_timesteps(self.num_inference_steps  )

              _UpperCAmelCase			     =  self.dummy_model()
              _UpperCAmelCase			     =  self.dummy_sample_deter * scheduler.init_noise_sigma

              for i, t in enumerate(scheduler.timesteps  ):
                     _UpperCAmelCase			     =  scheduler.scale_model_input(a__     ,	a__  )

                     _UpperCAmelCase			     =  model(a__     ,	a__  )

                     _UpperCAmelCase			     =  scheduler.step(a__     ,	a__     ,	a__  )
                     _UpperCAmelCase			     =  output.prev_sample

              _UpperCAmelCase			     =  torch.sum(torch.abs(a__  )  )
              _UpperCAmelCase			     =  torch.mean(torch.abs(a__  )  )

              assert abs(result_sum.item() - 1_006.388  ) < 1E-2
              assert abs(result_mean.item() - 1.31  ) < 1E-3






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.scheduler_classes[0]
              _UpperCAmelCase			     =  self.get_scheduler_config(prediction_type='v_prediction'  )
              _UpperCAmelCase			     =  scheduler_class(**a__  )

              scheduler.set_timesteps(self.num_inference_steps  )

              _UpperCAmelCase			     =  self.dummy_model()
              _UpperCAmelCase			     =  self.dummy_sample_deter * scheduler.init_noise_sigma

              for i, t in enumerate(scheduler.timesteps  ):
                     _UpperCAmelCase			     =  scheduler.scale_model_input(a__     ,	a__  )

                     _UpperCAmelCase			     =  model(a__     ,	a__  )

                     _UpperCAmelCase			     =  scheduler.step(a__     ,	a__     ,	a__  )
                     _UpperCAmelCase			     =  output.prev_sample

              _UpperCAmelCase			     =  torch.sum(torch.abs(a__  )  )
              _UpperCAmelCase			     =  torch.mean(torch.abs(a__  )  )

              assert abs(result_sum.item() - 0.0_017  ) < 1E-2
              assert abs(result_mean.item() - 2.2676E-06  ) < 1E-3






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.scheduler_classes[0]
              _UpperCAmelCase			     =  self.get_scheduler_config()
              _UpperCAmelCase			     =  scheduler_class(**a__  )

              scheduler.set_timesteps(self.num_inference_steps     ,	device=a__  )

              _UpperCAmelCase			     =  self.dummy_model()
              _UpperCAmelCase			     =  self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
              _UpperCAmelCase			     =  sample.to(a__  )

              for i, t in enumerate(scheduler.timesteps  ):
                     _UpperCAmelCase			     =  scheduler.scale_model_input(a__     ,	a__  )

                     _UpperCAmelCase			     =  model(a__     ,	a__  )

                     _UpperCAmelCase			     =  scheduler.step(a__     ,	a__     ,	a__  )
                     _UpperCAmelCase			     =  output.prev_sample

              _UpperCAmelCase			     =  torch.sum(torch.abs(a__  )  )
              _UpperCAmelCase			     =  torch.mean(torch.abs(a__  )  )

              assert abs(result_sum.item() - 1_006.388  ) < 1E-2
              assert abs(result_mean.item() - 1.31  ) < 1E-3






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.scheduler_classes[0]
              _UpperCAmelCase			     =  self.get_scheduler_config()
              _UpperCAmelCase			     =  scheduler_class(**a__     ,	use_karras_sigmas=a__  )

              scheduler.set_timesteps(self.num_inference_steps     ,	device=a__  )

              _UpperCAmelCase			     =  self.dummy_model()
              _UpperCAmelCase			     =  self.dummy_sample_deter.to(a__  ) * scheduler.init_noise_sigma
              _UpperCAmelCase			     =  sample.to(a__  )

              for t in scheduler.timesteps:
                     _UpperCAmelCase			     =  scheduler.scale_model_input(a__     ,	a__  )

                     _UpperCAmelCase			     =  model(a__     ,	a__  )

                     _UpperCAmelCase			     =  scheduler.step(a__     ,	a__     ,	a__  )
                     _UpperCAmelCase			     =  output.prev_sample

              _UpperCAmelCase			     =  torch.sum(torch.abs(a__  )  )
              _UpperCAmelCase			     =  torch.mean(torch.abs(a__  )  )

              assert abs(result_sum.item() - 3_812.9_927  ) < 2E-2
              assert abs(result_mean.item() - 4.9_648  ) < 1E-3






