



'''simple docstring'''






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, 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


enable_full_determinism()

class        SCREAMING_SNAKE_CASE(							unittest.TestCase    ):







	"""simple docstring"""






	lowerCamelCase__    				= StableDiffusionLDMaDPipeline
	lowerCamelCase__    				= TEXT_TO_IMAGE_PARAMS
	lowerCamelCase__    				= TEXT_TO_IMAGE_BATCH_PARAMS
	lowerCamelCase__    				= TEXT_TO_IMAGE_IMAGE_PARAMS







	def   A			(		self    :  int	)    ->	Any:
								torch.manual_seed(0	)
								UpperCAmelCase				: Optional[int]											=	UNetaDConditionModel(
								    block_out_channels=(32, 64)	,    layers_per_block=2	,    sample_size=32	,    in_channels=4	,    out_channels=4	,    down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')	,    up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')	,    cross_attention_dim=32	,    )
								UpperCAmelCase				: Tuple											=	DDIMScheduler(
								    beta_start=0.0_00_85	,    beta_end=0.0_12	,    beta_schedule='''scaled_linear'''	,    clip_sample=__snake_case	,    set_alpha_to_one=__snake_case	,    )
								torch.manual_seed(0	)
								UpperCAmelCase				: Dict											=	AutoencoderKL(
								    block_out_channels=[32, 64]	,    in_channels=6	,    out_channels=6	,    down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']	,    up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']	,    latent_channels=4	,    )
								torch.manual_seed(0	)
								UpperCAmelCase				: Dict											=	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=1000	,    )
								UpperCAmelCase				: Any											=	CLIPTextModel(__snake_case	)
								UpperCAmelCase				: List[Any]											=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)

								UpperCAmelCase				: Union[str, Any]											=	{
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components







	def   A			(		self    :  List[str]	,    __snake_case    :  str	,    __snake_case    :  Dict=0	)    ->	Dict:
								if str(__snake_case	).startswith('''mps'''	):
															UpperCAmelCase				: Optional[Any]											=	torch.manual_seed(__snake_case	)
								else:
															UpperCAmelCase				: Optional[int]											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: str											=	{
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs







	def   A			(		self    :  Union[str, Any]	)    ->	str:
								UpperCAmelCase				: Dict											=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator

								UpperCAmelCase				: int											=	self.get_dummy_components()
								UpperCAmelCase				: Optional[int]											=	StableDiffusionLDMaDPipeline(**__snake_case	)
								UpperCAmelCase				: List[str]											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Any											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: str											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: Union[str, Any]											=	output.rgb, output.depth

								UpperCAmelCase				: List[Any]											=	rgb[0, -3:, -3:, -1]
								UpperCAmelCase				: List[str]											=	depth[0, -3:, -1]

								assert rgb.shape == (1, 64, 64, 3)
								assert depth.shape == (1, 64, 64)

								UpperCAmelCase				: Optional[Any]											=	np.array(
								    [0.37_33_81_76, 0.7_02_47, 0.74_20_31_93, 0.51_64_36_04, 0.58_25_67_93, 0.60_93_21_36, 0.4_18_10_95, 0.48_35_58_77, 0.46_53_52_62]	)
								UpperCAmelCase				: Optional[Any]											=	np.array([1_03.4_67_27, 85.81_20_04, 87.84_92_36]	)

								assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb	).max() < 1E-2
								assert np.abs(image_slice_depth.flatten() - expected_slice_depth	).max() < 1E-2







	def   A			(		self    :  Optional[int]	)    ->	Tuple:
								UpperCAmelCase				: Tuple											=	self.get_dummy_components()
								UpperCAmelCase				: Any											=	StableDiffusionLDMaDPipeline(**__snake_case	)
								UpperCAmelCase				: Any											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: str											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: str											=	3 * [inputs['''prompt''']]

								# forward
								UpperCAmelCase				: Tuple											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: List[Any]											=	output.rgb, output.depth
								UpperCAmelCase				: Optional[Any]											=	rgb_slice_a[0, -3:, -3:, -1]
								UpperCAmelCase				: Optional[Any]											=	depth_slice_a[0, -3:, -1]

								UpperCAmelCase				: Tuple											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: Union[str, Any]											=	3 * [inputs.pop('''prompt'''	)]

								UpperCAmelCase				: Optional[Any]											=	ldmad_pipe.tokenizer(
								    __snake_case	,    padding='''max_length'''	,    max_length=ldmad_pipe.tokenizer.model_max_length	,    truncation=__snake_case	,    return_tensors='''pt'''	,    )
								UpperCAmelCase				: Union[str, Any]											=	text_inputs['''input_ids'''].to(__snake_case	)

								UpperCAmelCase				: Optional[int]											=	ldmad_pipe.text_encoder(__snake_case	)[0]

								UpperCAmelCase				: Optional[int]											=	prompt_embeds

								# forward
								UpperCAmelCase				: int											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: int											=	output.rgb, output.depth
								UpperCAmelCase				: Union[str, Any]											=	rgb_slice_a[0, -3:, -3:, -1]
								UpperCAmelCase				: int											=	depth_slice_a[0, -3:, -1]

								assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()	).max() < 1E-4
								assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()	).max() < 1E-4







	def   A			(		self    :  Optional[Any]	)    ->	Union[str, Any]:
								UpperCAmelCase				: List[Any]											=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase				: Optional[Any]											=	self.get_dummy_components()
								UpperCAmelCase				: Union[str, Any]											=	PNDMScheduler(skip_prk_steps=__snake_case	)
								UpperCAmelCase				: Tuple											=	StableDiffusionLDMaDPipeline(**__snake_case	)
								UpperCAmelCase				: Any											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: int											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	'''french fries'''
								UpperCAmelCase				: Dict											=	ldmad_pipe(**__snake_case	,    negative_prompt=__snake_case	)

								UpperCAmelCase							,	UpperCAmelCase				: Tuple											=	output.rgb, output.depth
								UpperCAmelCase				: List[str]											=	rgb[0, -3:, -3:, -1]
								UpperCAmelCase				: List[Any]											=	depth[0, -3:, -1]

								assert rgb.shape == (1, 64, 64, 3)
								assert depth.shape == (1, 64, 64)

								UpperCAmelCase				: Any											=	np.array(
								    [0.3_70_44, 0.71_81_15_03, 0.7_22_32_51, 0.48_60_36_75, 0.5_63_83_91, 0.6_36_49_48, 0.42_83_37_04, 0.4_90_13_15, 0.47_92_62_17]	)
								UpperCAmelCase				: str											=	np.array([1_07.8_47_38, 84.6_28_02, 89.96_21_35]	)
								assert np.abs(rgb_slice.flatten() - expected_slice_rgb	).max() < 1E-2
								assert np.abs(depth_slice.flatten() - expected_slice_depth	).max() < 1E-2

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







	"""simple docstring"""






	def   A			(		self    :  str	)    ->	Dict:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def   A			(		self    :  List[Any]	,    __snake_case    :  int	,    __snake_case    :  Tuple="cpu"	,    __snake_case    :  Tuple=torch.floataa	,    __snake_case    :  List[Any]=0	)    ->	Union[str, Any]:
								UpperCAmelCase				: List[str]											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: List[str]											=	np.random.RandomState(__snake_case	).standard_normal((1, 4, 64, 64)	)
								UpperCAmelCase				: Optional[Any]											=	torch.from_numpy(__snake_case	).to(device=__snake_case	,    dtype=__snake_case	)
								UpperCAmelCase				: List[str]											=	{
								    '''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   A			(		self    :  str	)    ->	str:
								UpperCAmelCase				: Union[str, Any]											=	StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''	)
								UpperCAmelCase				: Tuple											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Union[str, Any]											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: Optional[int]											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: List[str]											=	output.rgb, output.depth
								UpperCAmelCase				: Union[str, Any]											=	rgb[0, -3:, -3:, -1].flatten()
								UpperCAmelCase				: Any											=	rgb[0, -3:, -1].flatten()

								assert rgb.shape == (1, 512, 512, 3)
								assert depth.shape == (1, 512, 512)

								UpperCAmelCase				: Optional[Any]											=	np.array(
								    [0.53_80_54_65, 0.56_70_73_05, 0.5_48_65_15, 0.57_01_22_36, 0.5_81_45_11, 0.56_25_34_87, 0.54_84_30_14, 0.55_09_22_63, 0.6_45_97_06]	)
								UpperCAmelCase				: str											=	np.array(
								    [0.9_26_37_81, 0.6_67_86_72, 0.5_48_65_15, 0.92_20_21_45, 0.67_83_11_35, 0.56_25_34_87, 0.9_24_16_94, 0.7_55_14_78, 0.6_45_97_06]	)
								assert np.abs(rgb_slice - expected_slice_rgb	).max() < 3E-3
								assert np.abs(depth_slice - expected_slice_depth	).max() < 3E-3

@nightly
@require_torch_gpu
class        SCREAMING_SNAKE_CASE(							unittest.TestCase    ):







	"""simple docstring"""






	def   A			(		self    :  Tuple	)    ->	int:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def   A			(		self    :  str	,    __snake_case    :  Optional[Any]	,    __snake_case    :  Optional[Any]="cpu"	,    __snake_case    :  Dict=torch.floataa	,    __snake_case    :  int=0	)    ->	int:
								UpperCAmelCase				: Tuple											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: Dict											=	np.random.RandomState(__snake_case	).standard_normal((1, 4, 64, 64)	)
								UpperCAmelCase				: int											=	torch.from_numpy(__snake_case	).to(device=__snake_case	,    dtype=__snake_case	)
								UpperCAmelCase				: Optional[int]											=	{
								    '''prompt''': '''a photograph of an astronaut riding a horse''',
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 50,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs







	def   A			(		self    :  Tuple	)    ->	Any:
								UpperCAmelCase				: Any											=	StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''	).to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: List[str]											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: int											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: Dict											=	output.rgb, output.depth

								UpperCAmelCase				: List[Any]											=	0.49_55_86
								UpperCAmelCase				: List[Any]											=	0.33_79_55_15
								UpperCAmelCase				: int											=	1_12.4_85_18
								UpperCAmelCase				: Tuple											=	98.48_97_46
								assert np.abs(expected_rgb_mean - rgb.mean()	) < 1E-3
								assert np.abs(expected_rgb_std - rgb.std()	) < 1E-3
								assert np.abs(expected_depth_mean - depth.mean()	) < 1E-3
								assert np.abs(expected_depth_std - depth.std()	) < 1E-3







	def   A			(		self    :  int	)    ->	Any:
								UpperCAmelCase				: Optional[Any]											=	StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d-4c'''	).to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: List[str]											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: Dict											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: int											=	output.rgb, output.depth

								UpperCAmelCase				: List[Any]											=	0.4_19_41_27
								UpperCAmelCase				: str											=	0.35_37_55_86
								UpperCAmelCase				: Union[str, Any]											=	0.5_63_85_02
								UpperCAmelCase				: Optional[int]											=	0.34_68_61_03

								assert rgb.shape == (1, 512, 512, 3)
								assert depth.shape == (1, 512, 512, 1)
								assert np.abs(expected_rgb_mean - rgb.mean()	) < 1E-3
								assert np.abs(expected_rgb_std - rgb.std()	) < 1E-3
								assert np.abs(expected_depth_mean - depth.mean()	) < 1E-3
								assert np.abs(expected_depth_std - depth.std()	) < 1E-3










'''simple docstring'''






from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate

@dataclass(frozen=A__    )
class        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= field(default="""language-modeling"""							,	metadata={"""include_in_asdict_even_if_is_default""": True}    )

	lowerCamelCase__    				= Features({"""text""": Value("""string"""    )}    )
	lowerCamelCase__    				= Features({}    )
	lowerCamelCase__    				= "text"







	@property
	def   A			(		self    :  Optional[Any]	)    ->	Dict[str, str]:
								return {self.text_column: "text"}





