






"""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        lowercase__      (							unittest.TestCase       ):
			_UpperCAmelCase		:Optional[int]             =			StableDiffusionLDMaDPipeline
			_UpperCAmelCase		:Any             =			TEXT_TO_IMAGE_PARAMS
			_UpperCAmelCase		:List[Any]             =			TEXT_TO_IMAGE_BATCH_PARAMS
			_UpperCAmelCase		:Optional[Any]             =			TEXT_TO_IMAGE_IMAGE_PARAMS


			def      UpperCAmelCase__						(				self					:  Any ):
										torch.manual_seed(0 )
										lowerCamelCase_	:							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       ,	)
										lowerCamelCase_	:							Dict									=DDIMScheduler(
										    beta_start=0.00_085       ,	beta_end=0.012       ,	beta_schedule="scaled_linear"       ,	clip_sample=snake_case__       ,	set_alpha_to_one=snake_case__       ,	)
										torch.manual_seed(0 )
										lowerCamelCase_	:							List[str]									=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 )
										lowerCamelCase_	:							Tuple									=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       ,	)
										lowerCamelCase_	:							Union[str, Any]									=CLIPTextModel(snake_case__ )
										lowerCamelCase_	:							Union[str, Any]									=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

										lowerCamelCase_	:							Dict									={
										    "unet": unet,
										    "scheduler": scheduler,
										    "vae": vae,
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "safety_checker": None,
										    "feature_extractor": None,
										}
										return components


			def      UpperCAmelCase__						(				self					:  Optional[int]       ,	snake_case__					:  List[Any]       ,	snake_case__					:  str=0 ):
										if str(snake_case__ ).startswith("mps" ):
																	lowerCamelCase_	:							List[Any]									=torch.manual_seed(snake_case__ )
										else:
																	lowerCamelCase_	:							Any									=torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )
										lowerCamelCase_	:							Optional[int]									={
										    "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      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							Optional[int]									="cpu"  # ensure determinism for the device-dependent torch.Generator

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

										lowerCamelCase_	:							Optional[Any]									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							int									=ldmad_pipe(**snake_case__ )
										lowerCamelCase_							,						lowerCamelCase_	:							int									=output.rgb, output.depth

										lowerCamelCase_	:							Optional[int]									=rgb[0, -3:, -3:, -1]
										lowerCamelCase_	:							Optional[int]									=depth[0, -3:, -1]

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

										lowerCamelCase_	:							str									=np.array(
										    [0.37_338_176, 0.70_247, 0.74_203_193, 0.51_643_604, 0.58_256_793, 0.60_932_136, 0.4_181_095, 0.48_355_877, 0.46_535_262] )
										lowerCamelCase_	:							str									=np.array([103.46_727, 85.812_004, 87.849_236] )

										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      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							List[str]									=self.get_dummy_components()
										lowerCamelCase_	:							Any									=StableDiffusionLDMaDPipeline(**snake_case__ )
										lowerCamelCase_	:							str									=ldmad_pipe.to(snake_case__ )
										ldmad_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Optional[int]									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							Optional[int]									=3 * [inputs["prompt"]]

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

										lowerCamelCase_	:							Optional[Any]									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							Tuple									=3 * [inputs.pop("prompt" )]

										lowerCamelCase_	:							Dict									=ldmad_pipe.tokenizer(
										    snake_case__       ,	padding="max_length"       ,	max_length=ldmad_pipe.tokenizer.model_max_length       ,	truncation=snake_case__       ,	return_tensors="pt"       ,	)
										lowerCamelCase_	:							Any									=text_inputs["input_ids"].to(snake_case__ )

										lowerCamelCase_	:							Union[str, Any]									=ldmad_pipe.text_encoder(snake_case__ )[0]

										lowerCamelCase_	:							Optional[int]									=prompt_embeds

										# forward
										lowerCamelCase_	:							int									=ldmad_pipe(**snake_case__ )
										lowerCamelCase_							,						lowerCamelCase_	:							List[Any]									=output.rgb, output.depth
										lowerCamelCase_	:							int									=rgb_slice_a[0, -3:, -3:, -1]
										lowerCamelCase_	:							List[Any]									=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      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							int									="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase_	:							Union[str, Any]									=self.get_dummy_components()
										lowerCamelCase_	:							Any									=PNDMScheduler(skip_prk_steps=snake_case__ )
										lowerCamelCase_	:							List[str]									=StableDiffusionLDMaDPipeline(**snake_case__ )
										lowerCamelCase_	:							int									=ldmad_pipe.to(snake_case__ )
										ldmad_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Union[str, Any]									=self.get_dummy_inputs(snake_case__ )
										lowerCamelCase_	:							Any									="french fries"
										lowerCamelCase_	:							List[Any]									=ldmad_pipe(**snake_case__       ,	negative_prompt=snake_case__ )

										lowerCamelCase_							,						lowerCamelCase_	:							int									=output.rgb, output.depth
										lowerCamelCase_	:							List[str]									=rgb[0, -3:, -3:, -1]
										lowerCamelCase_	:							List[str]									=depth[0, -3:, -1]

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

										lowerCamelCase_	:							Union[str, Any]									=np.array(
										    [0.37_044, 0.71_811_503, 0.7_223_251, 0.48_603_675, 0.5_638_391, 0.6_364_948, 0.42_833_704, 0.4_901_315, 0.47_926_217] )
										lowerCamelCase_	:							str									=np.array([107.84_738, 84.62_802, 89.962_135] )
										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        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  int ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


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

										lowerCamelCase_	:							Optional[Any]									=self.get_inputs(snake_case__ )
										lowerCamelCase_	:							List[Any]									=ldmad_pipe(**snake_case__ )
										lowerCamelCase_							,						lowerCamelCase_	:							Union[str, Any]									=output.rgb, output.depth
										lowerCamelCase_	:							Tuple									=rgb[0, -3:, -3:, -1].flatten()
										lowerCamelCase_	:							Dict									=rgb[0, -3:, -1].flatten()

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

										lowerCamelCase_	:							Any									=np.array(
										    [0.53_805_465, 0.56_707_305, 0.5_486_515, 0.57_012_236, 0.5_814_511, 0.56_253_487, 0.54_843_014, 0.55_092_263, 0.6_459_706] )
										lowerCamelCase_	:							Optional[Any]									=np.array(
										    [0.9_263_781, 0.6_678_672, 0.5_486_515, 0.92_202_145, 0.67_831_135, 0.56_253_487, 0.9_241_694, 0.7_551_478, 0.6_459_706] )
										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        lowercase__      (							unittest.TestCase       ):


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


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

										lowerCamelCase_	:							List[str]									=self.get_inputs(snake_case__ )
										lowerCamelCase_	:							Optional[Any]									=ldmad_pipe(**snake_case__ )
										lowerCamelCase_							,						lowerCamelCase_	:							int									=output.rgb, output.depth

										lowerCamelCase_	:							List[str]									=0.495_586
										lowerCamelCase_	:							List[str]									=0.33_795_515
										lowerCamelCase_	:							Optional[Any]									=112.48_518
										lowerCamelCase_	:							Optional[Any]									=98.489_746
										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      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							Any									=StableDiffusionLDMaDPipeline.from_pretrained("Intel/ldm3d-4c" ).to(snake_case__ )
										ldmad_pipe.set_progress_bar_config(disable=snake_case__ )

										lowerCamelCase_	:							Tuple									=self.get_inputs(snake_case__ )
										lowerCamelCase_	:							Any									=ldmad_pipe(**snake_case__ )
										lowerCamelCase_							,						lowerCamelCase_	:							Dict									=output.rgb, output.depth

										lowerCamelCase_	:							Dict									=0.4_194_127
										lowerCamelCase_	:							str									=0.35_375_586
										lowerCamelCase_	:							Optional[int]									=0.5_638_502
										lowerCamelCase_	:							Optional[Any]									=0.34_686_103

										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"""



import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows

def 			_snake_case    (					lowerCamelCase__     :   Any  )    ->     Union[str, Any]:  # picklable for multiprocessing
							return i + 1

@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def 			_snake_case    (					)    ->     List[Any]:
							with parallel_backend("spark"  ):
														assert ParallelBackendConfig.backend_name == "spark"

							lowerCamelCase_	:							Optional[Any]									=[1, 2, 3]
							with pytest.raises(lowerCamelCase__  ):
														with parallel_backend("unsupported backend"  ):
																					map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=2  )

							with pytest.raises(lowerCamelCase__  ):
														with parallel_backend("unsupported backend"  ):
																					map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=-1  )

@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize("num_proc"    ,						[2, -1]  )
def 			_snake_case    (					lowerCamelCase__     :   Tuple  )    ->     Optional[Any]:
							lowerCamelCase_	:							str									=[1, 2]
							lowerCamelCase_	:							List[str]									={"a": 1, "b": 2}
							lowerCamelCase_	:							List[str]									={"a": [1, 2], "b": [3, 4]}
							lowerCamelCase_	:							Optional[int]									={"a": {"1": 1}, "b": 2}
							lowerCamelCase_	:							int									={"a": 1, "b": 2, "c": 3, "d": 4}
							lowerCamelCase_	:							Optional[int]									=[2, 3]
							lowerCamelCase_	:							List[Any]									={"a": 2, "b": 3}
							lowerCamelCase_	:							int									={"a": [2, 3], "b": [4, 5]}
							lowerCamelCase_	:							str									={"a": {"1": 2}, "b": 3}
							lowerCamelCase_	:							Dict									={"a": 2, "b": 3, "c": 4, "d": 5}

							with parallel_backend("spark"  ):
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa



