
"""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        _lowercase   (				__UpperCAmelCase				,       unittest.TestCase		):
	lowercase_        =	KandinskyPriorPipeline
	lowercase_        =	['prompt']
	lowercase_        =	['prompt', 'negative_prompt']
	lowercase_        =	[
	    'num_images_per_prompt',
	    'generator',
	    'num_inference_steps',
	    'latents',
	    'negative_prompt',
	    'guidance_scale',
	    'output_type',
	    'return_dict',
	]
	lowercase_        =	False




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




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




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




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




	@property
	def 			_UpperCamelCase    ( self					)		->							Optional[int]:
			return 100




	@property
	def 			_UpperCamelCase    ( self					)		->							List[Any]:
			lowerCamelCase     :			List[Any]		=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'					)
			return tokenizer




	@property
	def 			_UpperCamelCase    ( self					)		->							Dict:
			torch.manual_seed(0					)
			lowerCamelCase     :			str		=    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-0_5		, num_attention_heads=4		, num_hidden_layers=5		, pad_token_id=1		, vocab_size=1000		, )
			return CLIPTextModelWithProjection(UpperCAmelCase_					)




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

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

			lowerCamelCase     :			int		=    PriorTransformer(**UpperCAmelCase_					)
			# 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     :			List[Any]		=    nn.Parameter(torch.ones(model.clip_std.shape					)					)
			return model




	@property
	def 			_UpperCamelCase    ( self					)		->							Union[str, Any]:
			torch.manual_seed(0					)
			lowerCamelCase     :			Optional[int]		=    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     :			Any		=    CLIPVisionModelWithProjection(UpperCAmelCase_					)
			return model




	@property
	def 			_UpperCamelCase    ( self					)		->							Optional[int]:
			lowerCamelCase     :			Union[str, Any]		=    CLIPImageProcessor(
			    crop_size=224		, do_center_crop=UpperCAmelCase_		, do_normalize=UpperCAmelCase_		, do_resize=UpperCAmelCase_		, image_mean=[0.48145466, 0.4578275, 0.40821073]		, image_std=[0.26862954, 0.26130258, 0.27577711]		, resample=3		, size=224		, )

			return image_processor




	def 			_UpperCamelCase    ( self					)		->							Dict:
			lowerCamelCase     :			Tuple		=    self.dummy_prior
			lowerCamelCase     :			Any		=    self.dummy_image_encoder
			lowerCamelCase     :			Tuple		=    self.dummy_text_encoder
			lowerCamelCase     :			Any		=    self.dummy_tokenizer
			lowerCamelCase     :			str		=    self.dummy_image_processor

			lowerCamelCase     :			int		=    UnCLIPScheduler(
			    variance_type='fixed_small_log'		, prediction_type='sample'		, num_train_timesteps=1000		, clip_sample=UpperCAmelCase_		, clip_sample_range=10.0		, )

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

			return components




	def 			_UpperCamelCase    ( self		, UpperCAmelCase_		, UpperCAmelCase_=0					)		->							int:
			if str(UpperCAmelCase_					).startswith('mps'					):
					lowerCamelCase     :			str		=    torch.manual_seed(UpperCAmelCase_					)
			else:
					lowerCamelCase     :			int		=    torch.Generator(device=UpperCAmelCase_					).manual_seed(UpperCAmelCase_					)
			lowerCamelCase     :			List[str]		=    {
			    'prompt': 'horse',
			    'generator': generator,
			    'guidance_scale': 4.0,
			    'num_inference_steps': 2,
			    'output_type': 'np',
			}
			return inputs




	def 			_UpperCamelCase    ( self					)		->							str:
			lowerCamelCase     :			Dict		=    'cpu'

			lowerCamelCase     :			Tuple		=    self.get_dummy_components()

			lowerCamelCase     :			Any		=    self.pipeline_class(**UpperCAmelCase_					)
			lowerCamelCase     :			Optional[int]		=    pipe.to(UpperCAmelCase_					)

			pipe.set_progress_bar_config(disable=UpperCAmelCase_					)

			lowerCamelCase     :			int		=    pipe(**self.get_dummy_inputs(UpperCAmelCase_					)					)
			lowerCamelCase     :			List[Any]		=    output.image_embeds

			lowerCamelCase     :			str		=    pipe(
			    **self.get_dummy_inputs(UpperCAmelCase_					)		, return_dict=UpperCAmelCase_		, )[0]

			lowerCamelCase     :			List[Any]		=    image[0, -10:]
			lowerCamelCase     :			int		=    image_from_tuple[0, -10:]

			assert image.shape == (1, 32)

			lowerCamelCase     :			List[Any]		=    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 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase     :			Dict		=    torch_device == 'cpu'
			lowerCamelCase     :			Optional[Any]		=    True
			lowerCamelCase     :			List[str]		=    False

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




	@skip_mps
	def 			_UpperCamelCase    ( self					)		->							Any:
			lowerCamelCase     :			Optional[int]		=    torch_device == 'cpu'
			lowerCamelCase     :			List[Any]		=    False

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








"""simple docstring"""

import torch

from diffusers import KDPMaAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest






class        _lowercase   (				__UpperCAmelCase		):
	lowercase_        =	(KDPMaAncestralDiscreteScheduler,)
	lowercase_        =	1_0




	def 			_UpperCamelCase    ( self		, **UpperCAmelCase_					)		->							Optional[int]:
			lowerCamelCase     :			Any		=    {
			    'num_train_timesteps': 1100,
			    'beta_start': 0.0001,
			    'beta_end': 0.02,
			    'beta_schedule': 'linear',
			}

			config.update(**UpperCAmelCase_					)
			return config




	def 			_UpperCamelCase    ( self					)		->							Any:
			for timesteps in [10, 50, 100, 1000]:
					self.check_over_configs(num_train_timesteps=UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							List[str]:
			for beta_start, beta_end in zip([0.00001, 0.0001, 0.001]		, [0.0002, 0.002, 0.02]					):
					self.check_over_configs(beta_start=UpperCAmelCase_		, beta_end=UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							Optional[int]:
			for schedule in ["linear", "scaled_linear"]:
					self.check_over_configs(beta_schedule=UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							List[Any]:
			if torch_device == "mps":
					return
			lowerCamelCase     :			Any		=    self.scheduler_classes[0]
			lowerCamelCase     :			List[str]		=    self.get_scheduler_config()
			lowerCamelCase     :			Union[str, Any]		=    scheduler_class(**UpperCAmelCase_					)

			scheduler.set_timesteps(self.num_inference_steps					)

			lowerCamelCase     :			Optional[Any]		=    torch.manual_seed(0					)

			lowerCamelCase     :			int		=    self.dummy_model()
			lowerCamelCase     :			Tuple		=    self.dummy_sample_deter * scheduler.init_noise_sigma
			lowerCamelCase     :			Union[str, Any]		=    sample.to(UpperCAmelCase_					)

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

					lowerCamelCase     :			List[str]		=    model(UpperCAmelCase_		, UpperCAmelCase_					)

					lowerCamelCase     :			Tuple		=    scheduler.step(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_		, generator=UpperCAmelCase_					)
					lowerCamelCase     :			str		=    output.prev_sample

			lowerCamelCase     :			Optional[int]		=    torch.sum(torch.abs(UpperCAmelCase_					)					)
			lowerCamelCase     :			Tuple		=    torch.mean(torch.abs(UpperCAmelCase_					)					)

			assert abs(result_sum.item() - 13849.3877					) < 1E-2
			assert abs(result_mean.item() - 18.0331					) < 5E-3




	def 			_UpperCamelCase    ( self					)		->							Optional[int]:
			for prediction_type in ["epsilon", "v_prediction"]:
					self.check_over_configs(prediction_type=UpperCAmelCase_					)




	def 			_UpperCamelCase    ( self					)		->							List[str]:
			if torch_device == "mps":
					return
			lowerCamelCase     :			Dict		=    self.scheduler_classes[0]
			lowerCamelCase     :			Union[str, Any]		=    self.get_scheduler_config(prediction_type='v_prediction'					)
			lowerCamelCase     :			int		=    scheduler_class(**UpperCAmelCase_					)

			scheduler.set_timesteps(self.num_inference_steps					)

			lowerCamelCase     :			int		=    self.dummy_model()
			lowerCamelCase     :			Dict		=    self.dummy_sample_deter * scheduler.init_noise_sigma
			lowerCamelCase     :			Optional[int]		=    sample.to(UpperCAmelCase_					)

			lowerCamelCase     :			Dict		=    torch.manual_seed(0					)

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

					lowerCamelCase     :			Optional[Any]		=    model(UpperCAmelCase_		, UpperCAmelCase_					)

					lowerCamelCase     :			Tuple		=    scheduler.step(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_		, generator=UpperCAmelCase_					)
					lowerCamelCase     :			Optional[Any]		=    output.prev_sample

			lowerCamelCase     :			Union[str, Any]		=    torch.sum(torch.abs(UpperCAmelCase_					)					)
			lowerCamelCase     :			Optional[Any]		=    torch.mean(torch.abs(UpperCAmelCase_					)					)

			assert abs(result_sum.item() - 328.9970					) < 1E-2
			assert abs(result_mean.item() - 0.4284					) < 1E-3




	def 			_UpperCamelCase    ( self					)		->							Any:
			if torch_device == "mps":
					return
			lowerCamelCase     :			Union[str, Any]		=    self.scheduler_classes[0]
			lowerCamelCase     :			Dict		=    self.get_scheduler_config()
			lowerCamelCase     :			Dict		=    scheduler_class(**UpperCAmelCase_					)

			scheduler.set_timesteps(self.num_inference_steps		, device=UpperCAmelCase_					)
			lowerCamelCase     :			int		=    torch.manual_seed(0					)

			lowerCamelCase     :			List[Any]		=    self.dummy_model()
			lowerCamelCase     :			Union[str, Any]		=    self.dummy_sample_deter.to(UpperCAmelCase_					) * scheduler.init_noise_sigma

			for t in scheduler.timesteps:
					lowerCamelCase     :			Optional[int]		=    scheduler.scale_model_input(UpperCAmelCase_		, UpperCAmelCase_					)

					lowerCamelCase     :			str		=    model(UpperCAmelCase_		, UpperCAmelCase_					)

					lowerCamelCase     :			List[str]		=    scheduler.step(UpperCAmelCase_		, UpperCAmelCase_		, UpperCAmelCase_		, generator=UpperCAmelCase_					)
					lowerCamelCase     :			Dict		=    output.prev_sample

			lowerCamelCase     :			List[str]		=    torch.sum(torch.abs(UpperCAmelCase_					)					)
			lowerCamelCase     :			str		=    torch.mean(torch.abs(UpperCAmelCase_					)					)

			assert abs(result_sum.item() - 13849.3818					) < 1E-1
			assert abs(result_mean.item() - 18.0331					) < 1E-3






