
import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




def  __magic_name__ (   __lowerCAmelCase	:						str							)			->   int:
				__lowerCamelCase    					=	[tensor.shape for tensor in tensor_list]
				return all(shape == shapes[0] for shape in shapes[1:]							)

class    lowerCAmelCase__     (	__lowercase  ,	__lowercase  ,	__lowercase  ,	unittest.TestCase			):
		a__       :						str        =		StableDiffusionLatentUpscalePipeline
		a__       :						Union[str, Any]        =		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
		    """height""",
		    """width""",
		    """cross_attention_kwargs""",
		    """negative_prompt_embeds""",
		    """prompt_embeds""",
		}
		a__       :						Dict        =		PipelineTesterMixin.required_optional_params - {"""num_images_per_prompt"""}
		a__       :						Optional[Any]        =		TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
		a__       :						List[Any]        =		frozenset(
		    []			)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
		a__       :						Any        =		frozenset([]			)

		a__       :						str        =		True






		@property
		def  __A       (		self		:							Optional[int]  )     ->							int:
						__lowerCamelCase    					=	1
						__lowerCamelCase    					=	4
						__lowerCamelCase    					=	(16, 16)

						__lowerCamelCase    					=	floats_tensor((batch_size, num_channels) + sizes							,			rng=random.Random(0  )  ).to(SCREAMING_SNAKE_CASE__  )
						return image






		def  __A       (		self		:							List[Any]  )     ->							Tuple:
						torch.manual_seed(0  )
						__lowerCamelCase    					=	UNetaDConditionModel(
						    act_fn='''gelu'''							,			attention_head_dim=8							,			norm_num_groups=SCREAMING_SNAKE_CASE__							,			block_out_channels=[32, 32, 64, 64]							,			time_cond_proj_dim=1_60							,			conv_in_kernel=1							,			conv_out_kernel=1							,			cross_attention_dim=32							,			down_block_types=(
						        '''KDownBlock2D''',
						        '''KCrossAttnDownBlock2D''',
						        '''KCrossAttnDownBlock2D''',
						        '''KCrossAttnDownBlock2D''',
						    )							,			in_channels=8							,			mid_block_type=SCREAMING_SNAKE_CASE__							,			only_cross_attention=SCREAMING_SNAKE_CASE__							,			out_channels=5							,			resnet_time_scale_shift='''scale_shift'''							,			time_embedding_type='''fourier'''							,			timestep_post_act='''gelu'''							,			up_block_types=('''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KCrossAttnUpBlock2D''', '''KUpBlock2D''')							,			)
						__lowerCamelCase    					=	AutoencoderKL(
						    block_out_channels=[32, 32, 64, 64]							,			in_channels=3							,			out_channels=3							,			down_block_types=[
						        '''DownEncoderBlock2D''',
						        '''DownEncoderBlock2D''',
						        '''DownEncoderBlock2D''',
						        '''DownEncoderBlock2D''',
						    ]							,			up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']							,			latent_channels=4							,			)
						__lowerCamelCase    					=	EulerDiscreteScheduler(prediction_type='''sample'''  )
						__lowerCamelCase    					=	CLIPTextConfig(
						    bos_token_id=0							,			eos_token_id=2							,			hidden_size=32							,			intermediate_size=37							,			layer_norm_eps=1e-05							,			num_attention_heads=4							,			num_hidden_layers=5							,			pad_token_id=1							,			vocab_size=10_00							,			hidden_act='''quick_gelu'''							,			projection_dim=5_12							,			)
						__lowerCamelCase    					=	CLIPTextModel(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )

						__lowerCamelCase    					=	{
						    '''unet''': model.eval(),
						    '''vae''': vae.eval(),
						    '''scheduler''': scheduler,
						    '''text_encoder''': text_encoder,
						    '''tokenizer''': tokenizer,
						}

						return components






		def  __A       (		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Optional[int]							,			SCREAMING_SNAKE_CASE__		:							Union[str, Any]=0  )     ->							str:
						if str(SCREAMING_SNAKE_CASE__  ).startswith('''mps'''  ):
										__lowerCamelCase    					=	torch.manual_seed(SCREAMING_SNAKE_CASE__  )
						else:
										__lowerCamelCase    					=	torch.Generator(device=SCREAMING_SNAKE_CASE__  ).manual_seed(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	{
						    '''prompt''': '''A painting of a squirrel eating a burger''',
						    '''image''': self.dummy_image.cpu(),
						    '''generator''': generator,
						    '''num_inference_steps''': 2,
						    '''output_type''': '''numpy''',
						}
						return inputs






		def  __A       (		self		:							Union[str, Any]  )     ->							Optional[Any]:
						__lowerCamelCase    					=	'''cpu'''

						__lowerCamelCase    					=	self.get_dummy_components()
						__lowerCamelCase    					=	self.pipeline_class(**SCREAMING_SNAKE_CASE__  )
						pipe.to(SCREAMING_SNAKE_CASE__  )
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	self.get_dummy_inputs(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	pipe(**SCREAMING_SNAKE_CASE__  ).images
						__lowerCamelCase    					=	image[0, -3:, -3:, -1]

						self.assertEqual(image.shape							,			(1, 2_56, 2_56, 3)  )
						__lowerCamelCase    					=	np.array(
						    [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055]  )
						__lowerCamelCase    					=	np.abs(image_slice.flatten() - expected_slice  ).max()
						self.assertLessEqual(SCREAMING_SNAKE_CASE__							,			1e-3  )






		def  __A       (		self		:							str  )     ->							Optional[Any]:
						super().test_attention_slicing_forward_pass(expected_max_diff=7e-3  )






		def  __A       (		self		:							List[Any]  )     ->							Optional[int]:
						super().test_cpu_offload_forward_pass(expected_max_diff=3e-3  )






		def  __A       (		self		:							Any  )     ->							str:
						super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3  )






		def  __A       (		self		:							Any  )     ->							str:
						super().test_inference_batch_single_identical(expected_max_diff=7e-3  )






		def  __A       (		self		:							Any  )     ->							Any:
						super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3  )






		def  __A       (		self		:							Optional[Any]  )     ->							Optional[Any]:
						super().test_save_load_local(expected_max_difference=3e-3  )






		def  __A       (		self		:							Tuple  )     ->							Dict:
						super().test_save_load_optional_components(expected_max_difference=3e-3  )







		def  __A       (		self		:							Optional[Any]  )     ->							List[Any]:
						__lowerCamelCase    					=	[
						    '''DDIMScheduler''',
						    '''DDPMScheduler''',
						    '''PNDMScheduler''',
						    '''HeunDiscreteScheduler''',
						    '''EulerAncestralDiscreteScheduler''',
						    '''KDPM2DiscreteScheduler''',
						    '''KDPM2AncestralDiscreteScheduler''',
						    '''DPMSolverSDEScheduler''',
						]
						__lowerCamelCase    					=	self.get_dummy_components()
						__lowerCamelCase    					=	self.pipeline_class(**SCREAMING_SNAKE_CASE__  )

						# make sure that PNDM does not need warm-up
						pipe.scheduler.register_to_config(skip_prk_steps=SCREAMING_SNAKE_CASE__  )

						pipe.to(SCREAMING_SNAKE_CASE__  )
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.get_dummy_inputs(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	2

						__lowerCamelCase    					=	[]
						for scheduler_enum in KarrasDiffusionSchedulers:
										if scheduler_enum.name in skip_schedulers:
														# no sigma schedulers are not supported
														# no schedulers
														continue

										__lowerCamelCase    					=	getattr(SCREAMING_SNAKE_CASE__							,			scheduler_enum.name  )
										__lowerCamelCase    					=	scheduler_cls.from_config(pipe.scheduler.config  )
										__lowerCamelCase    					=	pipe(**SCREAMING_SNAKE_CASE__  )[0]
										outputs.append(SCREAMING_SNAKE_CASE__  )

						assert check_same_shape(SCREAMING_SNAKE_CASE__  )





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






		def  __A       (		self		:							List[Any]  )     ->							List[str]:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()






		def  __A       (		self		:							List[Any]  )     ->							Union[str, Any]:
						__lowerCamelCase    					=	torch.manual_seed(33  )

						__lowerCamelCase    					=	StableDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''							,			torch_dtype=torch.floataa  )
						pipe.to('''cuda'''  )

						__lowerCamelCase    					=	StableDiffusionLatentUpscalePipeline.from_pretrained(
						    '''stabilityai/sd-x2-latent-upscaler'''							,			torch_dtype=torch.floataa  )
						upscaler.to('''cuda'''  )

						__lowerCamelCase    					=	'''a photo of an astronaut high resolution, unreal engine, ultra realistic'''

						__lowerCamelCase    					=	pipe(SCREAMING_SNAKE_CASE__							,			generator=SCREAMING_SNAKE_CASE__							,			output_type='''latent'''  ).images

						__lowerCamelCase    					=	upscaler(
						    prompt=SCREAMING_SNAKE_CASE__							,			image=SCREAMING_SNAKE_CASE__							,			num_inference_steps=20							,			guidance_scale=0							,			generator=SCREAMING_SNAKE_CASE__							,			output_type='''np'''							,			).images[0]

						__lowerCamelCase    					=	load_numpy(
						    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'''  )
						assert np.abs((expected_image - image).mean()  ) < 5e-2







		def  __A       (		self		:							Any  )     ->							Union[str, Any]:
						__lowerCamelCase    					=	torch.manual_seed(33  )

						__lowerCamelCase    					=	StableDiffusionLatentUpscalePipeline.from_pretrained(
						    '''stabilityai/sd-x2-latent-upscaler'''							,			torch_dtype=torch.floataa  )
						upscaler.to('''cuda'''  )

						__lowerCamelCase    					=	'''the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'''

						__lowerCamelCase    					=	load_image(
						    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'''  )

						__lowerCamelCase    					=	upscaler(
						    prompt=SCREAMING_SNAKE_CASE__							,			image=SCREAMING_SNAKE_CASE__							,			num_inference_steps=20							,			guidance_scale=0							,			generator=SCREAMING_SNAKE_CASE__							,			output_type='''np'''							,			).images[0]

						__lowerCamelCase    					=	load_numpy(
						    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'''  )
						assert np.abs((expected_image - image).max()  ) < 5e-2





from typing import Optional

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader

class    lowerCAmelCase__     (	__lowercase			):






		def __init__(		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							NestedDataStructureLike[PathLike]							,			SCREAMING_SNAKE_CASE__		:							Optional[NamedSplit] = None							,			SCREAMING_SNAKE_CASE__		:							Optional[Features] = None							,			SCREAMING_SNAKE_CASE__		:							str = None							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							Optional[int] = None							,			**SCREAMING_SNAKE_CASE__		:							str							,			)     ->							Union[str, Any]:
						super().__init__(
						    SCREAMING_SNAKE_CASE__							,			split=SCREAMING_SNAKE_CASE__							,			features=SCREAMING_SNAKE_CASE__							,			cache_dir=SCREAMING_SNAKE_CASE__							,			keep_in_memory=SCREAMING_SNAKE_CASE__							,			streaming=SCREAMING_SNAKE_CASE__							,			num_proc=SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			)
						__lowerCamelCase    					=	path_or_paths if isinstance(SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__  ) else {self.split: path_or_paths}
						__lowerCamelCase    					=	Text(
						    cache_dir=SCREAMING_SNAKE_CASE__							,			data_files=SCREAMING_SNAKE_CASE__							,			features=SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			)







		def  __A       (		self		:							int  )     ->							Dict:
						# Build iterable dataset
						if self.streaming:
										__lowerCamelCase    					=	self.builder.as_streaming_dataset(split=self.split  )
						# Build regular (map-style) dataset
						else:
										__lowerCamelCase    					=	None
										__lowerCamelCase    					=	None
										__lowerCamelCase    					=	None
										__lowerCamelCase    					=	None

										self.builder.download_and_prepare(
										    download_config=SCREAMING_SNAKE_CASE__							,			download_mode=SCREAMING_SNAKE_CASE__							,			verification_mode=SCREAMING_SNAKE_CASE__							,			base_path=SCREAMING_SNAKE_CASE__							,			num_proc=self.num_proc							,			)
										__lowerCamelCase    					=	self.builder.as_dataset(
										    split=self.split							,			verification_mode=SCREAMING_SNAKE_CASE__							,			in_memory=self.keep_in_memory  )
						return dataset



