



"""simple docstring"""





import unittest

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

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

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class       __lowerCamelCase	(   __lowercase   ,							unittest.TestCase    ):
   __UpperCamelCase									=   KandinskyPriorPipeline
   __UpperCamelCase									=   ['prompt']
   __UpperCamelCase									=   ['prompt', 'negative_prompt']
   __UpperCamelCase									=   [
       'num_images_per_prompt',
       'generator',
       'num_inference_steps',
       'latents',
       'negative_prompt',
       'guidance_scale',
       'output_type',
       'return_dict',
   ]
   __UpperCamelCase									=   False







   @property
   def 							A__						(self				):




      '''simple docstring'''




      return 32







   @property
   def 							A__						(self				):




      '''simple docstring'''




      return 32







   @property
   def 							A__						(self				):




      '''simple docstring'''




      return self.time_input_dim







   @property
   def 							A__						(self				):




      '''simple docstring'''




      return self.time_input_dim * 4







   @property
   def 							A__						(self				):




      '''simple docstring'''




      return 100







   @property
   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)
      return tokenizer







   @property
   def 							A__						(self				):




      '''simple docstring'''




      torch.manual_seed(0				)
      _lowerCAmelCase					     =		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=1_000       ,	)
      return CLIPTextModelWithProjection(lowerCamelCase				)







   @property
   def 							A__						(self				):




      '''simple docstring'''




      torch.manual_seed(0				)

      _lowerCAmelCase					     =		{
          """num_attention_heads""": 2,
          """attention_head_dim""": 12,
          """embedding_dim""": self.text_embedder_hidden_size,
          """num_layers""": 1,
      }

      _lowerCAmelCase					     =		PriorTransformer(**lowerCamelCase				)
      # 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
      _lowerCAmelCase					     =		nn.Parameter(torch.ones(model.clip_std.shape				)				)
      return model







   @property
   def 							A__						(self				):




      '''simple docstring'''




      torch.manual_seed(0				)
      _lowerCAmelCase					     =		CLIPVisionConfig(
          hidden_size=self.text_embedder_hidden_size       ,	image_size=224       ,	projection_dim=self.text_embedder_hidden_size       ,	intermediate_size=37       ,	num_attention_heads=4       ,	num_channels=3       ,	num_hidden_layers=5       ,	patch_size=14       ,	)

      _lowerCAmelCase					     =		CLIPVisionModelWithProjection(lowerCamelCase				)
      return model







   @property
   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		CLIPImageProcessor(
          crop_size=224       ,	do_center_crop=lowerCamelCase       ,	do_normalize=lowerCamelCase       ,	do_resize=lowerCamelCase       ,	image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]       ,	image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]       ,	resample=3       ,	size=224       ,	)

      return image_processor







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.dummy_prior
      _lowerCAmelCase					     =		self.dummy_image_encoder
      _lowerCAmelCase					     =		self.dummy_text_encoder
      _lowerCAmelCase					     =		self.dummy_tokenizer
      _lowerCAmelCase					     =		self.dummy_image_processor

      _lowerCAmelCase					     =		UnCLIPScheduler(
          variance_type="""fixed_small_log"""       ,	prediction_type="""sample"""       ,	num_train_timesteps=1_000       ,	clip_sample=lowerCamelCase       ,	clip_sample_range=10.0       ,	)

      _lowerCAmelCase					     =		{
          """prior""": prior,
          """image_encoder""": image_encoder,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """scheduler""": scheduler,
          """image_processor""": image_processor,
      }

      return components







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase=0				):




      '''simple docstring'''




      if str(lowerCamelCase				).startswith("""mps"""				):
         _lowerCAmelCase					     =		torch.manual_seed(lowerCamelCase				)
      else:
         _lowerCAmelCase					     =		torch.Generator(device=lowerCamelCase				).manual_seed(lowerCamelCase				)
      _lowerCAmelCase					     =		{
          """prompt""": """horse""",
          """generator""": generator,
          """guidance_scale""": 4.0,
          """num_inference_steps""": 2,
          """output_type""": """np""",
      }
      return inputs







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		"""cpu"""

      _lowerCAmelCase					     =		self.get_dummy_components()

      _lowerCAmelCase					     =		self.pipeline_class(**lowerCamelCase				)
      _lowerCAmelCase					     =		pipe.to(lowerCamelCase				)

      pipe.set_progress_bar_config(disable=lowerCamelCase				)

      _lowerCAmelCase					     =		pipe(**self.get_dummy_inputs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		output.image_embeds

      _lowerCAmelCase					     =		pipe(
          **self.get_dummy_inputs(lowerCamelCase				)       ,	return_dict=lowerCamelCase       ,	)[0]

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

      assert image.shape == (1, 32)

      _lowerCAmelCase					     =		np.array(
          [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]				)

      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				):




      '''simple docstring'''




      _lowerCAmelCase					     =		torch_device == """cpu"""
      _lowerCAmelCase					     =		True
      _lowerCAmelCase					     =		False

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







   @skip_mps
   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		torch_device == """cpu"""
      _lowerCAmelCase					     =		False

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




"""simple docstring"""





import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class       __lowerCamelCase	(   __lowercase    ):
   __UpperCamelCase									=   (HeunDiscreteScheduler,)
   __UpperCamelCase									=   10







   def 							A__						(self       ,	**lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		{
          """num_train_timesteps""": 1_100,
          """beta_start""": 0.0001,
          """beta_end""": 0.02,
          """beta_schedule""": """linear""",
      }

      config.update(**lowerCamelCase				)
      return config







   def 							A__						(self				):




      '''simple docstring'''




      for timesteps in [10, 50, 100, 1_000]:
         self.check_over_configs(num_train_timesteps=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]       ,	[0.0002, 0.002, 0.02]				):
         self.check_over_configs(beta_start=lowerCamelCase       ,	beta_end=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for schedule in ["linear", "scaled_linear", "exp"]:
         self.check_over_configs(beta_schedule=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for clip_sample_range in [1.0, 2.0, 3.0]:
         self.check_over_configs(clip_sample_range=lowerCamelCase       ,	clip_sample=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for prediction_type in ["epsilon", "v_prediction", "sample"]:
         self.check_over_configs(prediction_type=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config(prediction_type="""v_prediction"""				)
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 4.6934e-07				) < 1e-2
         assert abs(result_mean.item() - 6.1112e-10				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 4.693428650170972e-07				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

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

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if str(lowerCamelCase				).startswith("""cpu"""				):
         # The following sum varies between 148 and 156 on mps. Why?
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      elif str(lowerCamelCase				).startswith("""mps"""				):
         # Larger tolerance on mps
         assert abs(result_mean.item() - 0.0002				) < 1e-2
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase       ,	use_karras_sigmas=lowerCamelCase				)

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

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      assert abs(result_sum.item() - 0.0_0015				) < 1e-2
      assert abs(result_mean.item() - 1.9869554535034695e-07				) < 1e-2