



'''simple docstring'''




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 						_UpperCamelCase  (							lowerCamelCase__   , lowerCamelCase__   , unittest.TestCase   ):


			'''simple docstring'''

			_A				:		str         =		StableDiffusionXLPipeline
			_A				:		Union[str, Any]         =		TEXT_TO_IMAGE_PARAMS
			_A				:		List[Any]         =		TEXT_TO_IMAGE_BATCH_PARAMS
			_A				:		Union[str, Any]         =		TEXT_TO_IMAGE_IMAGE_PARAMS
			_A				:		Optional[int]         =		TEXT_TO_IMAGE_IMAGE_PARAMS

			def UpperCamelCase__ (							self					:			Dict						):





					"""simple docstring"""






					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		UNetaDConditionModel(
					    block_out_channels=(3_2, 6_4)							,					layers_per_block=2							,					sample_size=3_2							,					in_channels=4							,					out_channels=4							,					down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")							,					up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")							,					attention_head_dim=(2, 4)							,					use_linear_projection=lowerCAmelCase__							,					addition_embed_type="""text_time"""							,					addition_time_embed_dim=8							,					transformer_layers_per_block=(1, 2)							,					projection_class_embeddings_input_dim=8_0							,					cross_attention_dim=6_4							,					)
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		EulerDiscreteScheduler(
					    beta_start=0.0_00_85							,					beta_end=0.0_12							,					steps_offset=1							,					beta_schedule="""scaled_linear"""							,					timestep_spacing="""leading"""							,					)
					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		str   	=		AutoencoderKL(
					    block_out_channels=[3_2, 6_4]							,					in_channels=3							,					out_channels=3							,					down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]							,					up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]							,					latent_channels=4							,					sample_size=1_2_8							,					)
					torch.manual_seed(0						)
					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		CLIPTextConfig(
					    bos_token_id=0							,					eos_token_id=2							,					hidden_size=3_2							,					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							,					hidden_act="""gelu"""							,					projection_dim=3_2							,					)
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		CLIPTextModel(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							,					local_files_only=lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		CLIPTextModelWithProjection(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""							,					local_files_only=lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		{
					    """unet""": unet,
					    """scheduler""": scheduler,
					    """vae""": vae,
					    """text_encoder""": text_encoder,
					    """tokenizer""": tokenizer,
					    """text_encoder_2""": text_encoder_a,
					    """tokenizer_2""": tokenizer_a,
					    # "safety_checker": None,
					    # "feature_extractor": None,
					}
					return components

			def UpperCamelCase__ (							self					:			Union[str, Any]							,					lowerCAmelCase__					:			Union[str, Any]							,					lowerCAmelCase__					:			Dict=0						):





					"""simple docstring"""






					if str(lowerCAmelCase__						).startswith("""mps"""						):
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		torch.manual_seed(lowerCAmelCase__						)
					else:
							__SCREAMING_SNAKE_CASE	:		Dict   	=		torch.Generator(device=lowerCAmelCase__						).manual_seed(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		int   	=		{
					    """prompt""": """A painting of a squirrel eating a burger""",
					    """generator""": generator,
					    """num_inference_steps""": 2,
					    """guidance_scale""": 5.0,
					    """output_type""": """numpy""",
					}
					return inputs

			def UpperCamelCase__ (							self					:			Tuple						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Any   	=		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
					__SCREAMING_SNAKE_CASE	:		str   	=		self.get_dummy_components()
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		StableDiffusionXLPipeline(**lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		sd_pipe.to(lowerCAmelCase__						)
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		str   	=		self.get_dummy_inputs(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		sd_pipe(**lowerCAmelCase__						).images
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		image[0, -3:, -3:, -1]

					assert image.shape == (1, 6_4, 6_4, 3)
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		np.array([0.58_73, 0.61_28, 0.47_97, 0.51_22, 0.56_74, 0.46_39, 0.52_27, 0.51_49, 0.47_47]						)

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

			def UpperCamelCase__ (							self					:			List[str]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Dict   	=		self.get_dummy_components()
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		StableDiffusionXLPipeline(**lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Dict   	=		sd_pipe.to(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		sd_pipe.to(lowerCAmelCase__						)
					sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

					# forward without prompt embeds
					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		self.get_dummy_inputs(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		3 * ["""this is a negative prompt"""]
					__SCREAMING_SNAKE_CASE	:		Any   	=		negative_prompt
					__SCREAMING_SNAKE_CASE	:		str   	=		3 * [inputs["""prompt"""]]

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		sd_pipe(**lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		output.images[0, -3:, -3:, -1]

					# forward with prompt embeds
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		self.get_dummy_inputs(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Any   	=		3 * ["""this is a negative prompt"""]
					__SCREAMING_SNAKE_CASE	:		int   	=		3 * [inputs.pop("""prompt"""						)]

					(
					    (
					    __SCREAMING_SNAKE_CASE
					)	,						(
					    __SCREAMING_SNAKE_CASE
					)	,						(
					    __SCREAMING_SNAKE_CASE
					)	,						(
					    __SCREAMING_SNAKE_CASE
					)	,						
					)	:		Any   	=		sd_pipe.encode_prompt(lowerCAmelCase__							,					negative_prompt=lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		Dict   	=		sd_pipe(
					    **lowerCAmelCase__							,					prompt_embeds=lowerCAmelCase__							,					negative_prompt_embeds=lowerCAmelCase__							,					pooled_prompt_embeds=lowerCAmelCase__							,					negative_pooled_prompt_embeds=lowerCAmelCase__							,					)
					__SCREAMING_SNAKE_CASE	:		Any   	=		output.images[0, -3:, -3:, -1]

					# make sure that it's equal
					assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()						).max() < 1E-4

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					super().test_attention_slicing_forward_pass(expected_max_diff=3E-3						)



			def UpperCamelCase__ (							self					:			Optional[int]						):





					"""simple docstring"""






					super().test_inference_batch_single_identical(expected_max_diff=3E-3						)






@slow
@require_torch_gpu
class 						_UpperCamelCase  (							unittest.TestCase   ):


			'''simple docstring'''

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()

			def UpperCamelCase__ (							self					:			Any							,					lowerCAmelCase__					:			List[Any]							,					lowerCAmelCase__					:			str="cpu"							,					lowerCAmelCase__					:			List[str]=torch.floataa							,					lowerCAmelCase__					:			str=0						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Dict   	=		torch.Generator(device=lowerCAmelCase__						).manual_seed(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Dict   	=		np.random.RandomState(lowerCAmelCase__						).standard_normal((1, 4, 6_4, 6_4)						)
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		torch.from_numpy(lowerCAmelCase__						).to(device=lowerCAmelCase__							,					dtype=lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Any   	=		{
					    """prompt""": """a photograph of an astronaut riding a horse""",
					    """latents""": latents,
					    """generator""": generator,
					    """num_inference_steps""": 3,
					    """guidance_scale""": 7.5,
					    """output_type""": """numpy""",
					}
					return inputs



			def UpperCamelCase__ (							self					:			List[str]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""						)
					pipe.to(lowerCAmelCase__						)
					pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.get_inputs(lowerCAmelCase__						)
					__SCREAMING_SNAKE_CASE	:		Dict   	=		pipe(**lowerCAmelCase__						).images
					__SCREAMING_SNAKE_CASE	:		str   	=		image[0, -3:, -3:, -1].flatten()

					assert image.shape == (1, 5_1_2, 5_1_2, 3)
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]						)
					assert np.abs(image_slice - expected_slice						).max() < 7E-3




'''simple docstring'''




import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest






class 						_UpperCamelCase  (							lowerCamelCase__   ):


			'''simple docstring'''

			_A				:		int         =		(VQDiffusionScheduler,)

			def UpperCamelCase__ (							self					:			List[str]							,					**lowerCAmelCase__					:			Dict						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		int   	=		{
					    """num_vec_classes""": 4_0_9_7,
					    """num_train_timesteps""": 1_0_0,
					}

					config.update(**lowerCAmelCase__						)
					return config

			def UpperCamelCase__ (							self					:			Dict							,					lowerCAmelCase__					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Dict   	=		4
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		8
					__SCREAMING_SNAKE_CASE	:		int   	=		8

					__SCREAMING_SNAKE_CASE	:		Dict   	=		torch.randint(0							,					lowerCAmelCase__							,					(batch_size, height * width)						)

					return sample

			@property
			def UpperCamelCase__ (							self					:			str						):





					"""simple docstring"""






					assert False

			def UpperCamelCase__ (							self					:			Dict							,					lowerCAmelCase__					:			str						):





					"""simple docstring"""






					def model(lowerCAmelCase__					:			Optional[Any]							,					lowerCAmelCase__					:			Any							,					*lowerCAmelCase__					:			Optional[int]						):
							__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		sample.shape
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)						)
							__SCREAMING_SNAKE_CASE	:		str   	=		F.log_softmax(logits.double()							,					dim=1						).float()
							return return_value

					return model

			def UpperCamelCase__ (							self					:			Optional[int]						):





					"""simple docstring"""






					for timesteps in [2, 5, 1_0_0, 1_0_0_0]:
							self.check_over_configs(num_train_timesteps=lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					for num_vec_classes in [5, 1_0_0, 1_0_0_0, 4_0_0_0]:
							self.check_over_configs(num_vec_classes=lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			Any						):





					"""simple docstring"""






					for t in [0, 5_0, 9_9]:
							self.check_over_forward(time_step=lowerCAmelCase__						)



			def UpperCamelCase__ (							self					:			List[Any]						):





					"""simple docstring"""






					pass