






import unittest

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

from diffusers import KandinskyVaaPriorPipeline, 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				(_snake_case			,  unittest.TestCase      ):







		'''simple docstring'''



		_snake_case   :   int							=  KandinskyVaaPriorPipeline
		_snake_case   :   str							=  ['''prompt''']
		_snake_case   :   int							=  ['''prompt''', '''negative_prompt''']
		_snake_case   :   Optional[Any]							=  [
		    '''num_images_per_prompt''',
		    '''generator''',
		    '''num_inference_steps''',
		    '''latents''',
		    '''negative_prompt''',
		    '''guidance_scale''',
		    '''output_type''',
		    '''return_dict''',
		]
		_snake_case   :   str							=  False



		@property
		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							return 3_2



		@property
		def  __UpperCAmelCase      (							self       )     ->							Any:
							return 3_2



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



		@property
		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							return self.time_input_dim * 4



		@property
		def  __UpperCAmelCase      (							self       )     ->							Optional[Any]:
							return 1_0_0



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



		@property
		def  __UpperCAmelCase      (							self       )     ->							Any:
							torch.manual_seed(0       )
							UpperCAmelCase_   :				Union[str, Any]						=						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=3_7    ,		layer_norm_eps=1E-05    ,		num_attention_heads=4    ,		num_hidden_layers=5    ,		pad_token_id=1    ,		vocab_size=1_0_0_0    ,		)
							return CLIPTextModelWithProjection(_UpperCamelCase       )



		@property
		def  __UpperCAmelCase      (							self       )     ->							int:
							torch.manual_seed(0       )

							UpperCAmelCase_   :				List[str]						=						{
							    'num_attention_heads': 2,
							    'attention_head_dim': 1_2,
							    'embedding_dim': self.text_embedder_hidden_size,
							    'num_layers': 1,
							}

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



		@property
		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							torch.manual_seed(0       )
							UpperCAmelCase_   :				Dict						=						CLIPVisionConfig(
							    hidden_size=self.text_embedder_hidden_size    ,		image_size=2_2_4    ,		projection_dim=self.text_embedder_hidden_size    ,		intermediate_size=3_7    ,		num_attention_heads=4    ,		num_channels=3    ,		num_hidden_layers=5    ,		patch_size=1_4    ,		)

							UpperCAmelCase_   :				List[str]						=						CLIPVisionModelWithProjection(_UpperCamelCase       )
							return model



		@property
		def  __UpperCAmelCase      (							self       )     ->							Any:
							UpperCAmelCase_   :				List[Any]						=						CLIPImageProcessor(
							    crop_size=2_2_4    ,		do_center_crop=_UpperCamelCase    ,		do_normalize=_UpperCamelCase    ,		do_resize=_UpperCamelCase    ,		image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]    ,		image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]    ,		resample=3    ,		size=2_2_4    ,		)

							return image_processor



		def  __UpperCAmelCase      (							self       )     ->							List[Any]:
							UpperCAmelCase_   :				str						=						self.dummy_prior
							UpperCAmelCase_   :				Optional[int]						=						self.dummy_image_encoder
							UpperCAmelCase_   :				List[str]						=						self.dummy_text_encoder
							UpperCAmelCase_   :				Optional[int]						=						self.dummy_tokenizer
							UpperCAmelCase_   :				Tuple						=						self.dummy_image_processor

							UpperCAmelCase_   :				List[Any]						=						UnCLIPScheduler(
							    variance_type='fixed_small_log'    ,		prediction_type='sample'    ,		num_train_timesteps=1_0_0_0    ,		clip_sample=_UpperCamelCase    ,		clip_sample_range=10.0    ,		)

							UpperCAmelCase_   :				Optional[Any]						=						{
							    '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       )     ->							List[Any]:
							if str(_UpperCamelCase       ).startswith('mps'       ):
												UpperCAmelCase_   :				int						=						torch.manual_seed(_UpperCamelCase       )
							else:
												UpperCAmelCase_   :				Union[str, Any]						=						torch.Generator(device=_UpperCamelCase       ).manual_seed(_UpperCamelCase       )
							UpperCAmelCase_   :				str						=						{
							    'prompt': 'horse',
							    'generator': generator,
							    'guidance_scale': 4.0,
							    'num_inference_steps': 2,
							    'output_type': 'np',
							}
							return inputs



		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				Optional[int]						=						'cpu'

							UpperCAmelCase_   :				List[str]						=						self.get_dummy_components()

							UpperCAmelCase_   :				Tuple						=						self.pipeline_class(**_UpperCamelCase       )
							UpperCAmelCase_   :				str						=						pipe.to(_UpperCamelCase       )

							pipe.set_progress_bar_config(disable=_UpperCamelCase       )

							UpperCAmelCase_   :				Union[str, Any]						=						pipe(**self.get_dummy_inputs(_UpperCamelCase       )       )
							UpperCAmelCase_   :				str						=						output.image_embeds

							UpperCAmelCase_   :				List[Any]						=						pipe(
							    **self.get_dummy_inputs(_UpperCamelCase       )    ,		return_dict=_UpperCamelCase    ,		)[0]

							UpperCAmelCase_   :				int						=						image[0, -1_0:]
							UpperCAmelCase_   :				Union[str, Any]						=						image_from_tuple[0, -1_0:]

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

							UpperCAmelCase_   :				int						=						np.array(
							    [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]       )

							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       )     ->							List[Any]:
							UpperCAmelCase_   :				Optional[Any]						=						torch_device == 'cpu'
							UpperCAmelCase_   :				Optional[Any]						=						True
							UpperCAmelCase_   :				Union[str, Any]						=						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       )     ->							List[str]:
							UpperCAmelCase_   :				int						=						torch_device == 'cpu'
							UpperCAmelCase_   :				int						=						False

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














import os
import re
import shutil
import sys
import tempfile
import unittest

import black


__UpperCAmelCase     							=						os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, 'utils'))

import check_copies  # noqa: E402


# This is the reference code that will be used in the tests.
# If DDPMSchedulerOutput is changed in scheduling_ddpm.py, this code needs to be manually updated.
__UpperCAmelCase     							=						'    \"""\n    Output class for the scheduler\'s step function output.\n\n    Args:\n        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):\n            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the\n            denoising loop.\n        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):\n            The predicted denoised sample (x_{0}) based on the model output from the current timestep.\n            `pred_original_sample` can be used to preview progress or for guidance.\n    \"""\n\n    prev_sample: torch.FloatTensor\n    pred_original_sample: Optional[torch.FloatTensor] = None\n'



class   lowerCamelCase				(unittest.TestCase      ):







		'''simple docstring'''



		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							UpperCAmelCase_   :				str						=						tempfile.mkdtemp()
							os.makedirs(os.path.join(self.diffusers_dir    ,		'schedulers/'       )       )
							UpperCAmelCase_   :				int						=						self.diffusers_dir
							shutil.copy(
							    os.path.join(_UpperCamelCase    ,		'src/diffusers/schedulers/scheduling_ddpm.py'       )    ,		os.path.join(self.diffusers_dir    ,		'schedulers/scheduling_ddpm.py'       )    ,		)



		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				str						=						'src/diffusers'
							shutil.rmtree(self.diffusers_dir       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase=None       )     ->							Optional[Any]:
							UpperCAmelCase_   :				Dict						=						comment + f"\nclass {class_name}(nn.Module):\n" + class_code
							if overwrite_result is not None:
												UpperCAmelCase_   :				Optional[Any]						=						comment + f"\nclass {class_name}(nn.Module):\n" + overwrite_result
							UpperCAmelCase_   :				Any						=						black.Mode(target_versions={black.TargetVersion.PYaa}    ,		line_length=1_1_9       )
							UpperCAmelCase_   :				Optional[Any]						=						black.format_str(_UpperCamelCase    ,		mode=_UpperCamelCase       )
							UpperCAmelCase_   :				Any						=						os.path.join(self.diffusers_dir    ,		'new_code.py'       )
							with open(_UpperCamelCase    ,		'w'    ,		newline='\n'       ) as f:
												f.write(_UpperCamelCase       )
							if overwrite_result is None:
												self.assertTrue(len(check_copies.is_copy_consistent(_UpperCamelCase       )       ) == 0       )
							else:
												check_copies.is_copy_consistent(f.name    ,		overwrite=_UpperCamelCase       )
												with open(_UpperCamelCase    ,		'r'       ) as f:
																	self.assertTrue(f.read()    ,		_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							List[str]:
							UpperCAmelCase_   :				List[Any]						=						check_copies.find_code_in_diffusers('schedulers.scheduling_ddpm.DDPMSchedulerOutput'       )
							self.assertEqual(_UpperCamelCase    ,		_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							str:
							# Base copy consistency
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput'    ,		'DDPMSchedulerOutput'    ,		REFERENCE_CODE + '\n'    ,		)

							# With no empty line at the end
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput'    ,		'DDPMSchedulerOutput'    ,		_UpperCamelCase    ,		)

							# Copy consistency with rename
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test'    ,		'TestSchedulerOutput'    ,		re.sub('DDPM'    ,		'Test'    ,		_UpperCamelCase       )    ,		)

							# Copy consistency with a really long name
							UpperCAmelCase_   :				Optional[int]						=						'TestClassWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason'
							self.check_copy_consistency(
							    f"# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->{long_class_name}"    ,		f"{long_class_name}SchedulerOutput"    ,		re.sub('Bert'    ,		_UpperCamelCase    ,		_UpperCamelCase       )    ,		)

							# Copy consistency with overwrite
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test'    ,		'TestSchedulerOutput'    ,		_UpperCamelCase    ,		overwrite_result=re.sub('DDPM'    ,		'Test'    ,		_UpperCamelCase       )    ,		)






