





"""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 			_a						(				unittest.TestCase    ):
				UpperCamelCase								=   StableDiffusionLDMaDPipeline
				UpperCamelCase								=   TEXT_TO_IMAGE_PARAMS
				UpperCamelCase								=   TEXT_TO_IMAGE_BATCH_PARAMS
				UpperCamelCase								=   TEXT_TO_IMAGE_IMAGE_PARAMS






				def     snake_case     (			self		:   Union[str, Any]							)     ->				List[str]:




										'''simple docstring'''



										torch.manual_seed(0							)
										_UpperCamelCase      : Optional[Any]         =   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'''),			cross_attention_dim=3_2,			)
										_UpperCamelCase      : Optional[Any]         =   DDIMScheduler(
										    beta_start=0.00_085,			beta_end=0.012,			beta_schedule='''scaled_linear''',			clip_sample=_a,			set_alpha_to_one=_a,			)
										torch.manual_seed(0							)
										_UpperCamelCase      : str         =   AutoencoderKL(
										    block_out_channels=[3_2, 6_4],			in_channels=6,			out_channels=6,			down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],			up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],			latent_channels=4,			)
										torch.manual_seed(0							)
										_UpperCamelCase      : Any         =   CLIPTextConfig(
										    bos_token_id=0,			eos_token_id=2,			hidden_size=3_2,			intermediate_size=3_7,			layer_norm_eps=1e-0_5,			num_attention_heads=4,			num_hidden_layers=5,			pad_token_id=1,			vocab_size=1_0_0_0,			)
										_UpperCamelCase      : List[str]         =   CLIPTextModel(_a							)
										_UpperCamelCase      : List[str]         =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

										_UpperCamelCase      : Tuple         =   {
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										    '''vae''': vae,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''safety_checker''': None,
										    '''feature_extractor''': None,
										}
										return components






				def     snake_case     (			self		:   Union[str, Any],			lowerCAmelCase__		:   int,			lowerCAmelCase__		:   Any=0							)     ->				Optional[Any]:




										'''simple docstring'''



										if str(_a							).startswith('''mps'''							):
																_UpperCamelCase      : Tuple         =   torch.manual_seed(_a							)
										else:
																_UpperCamelCase      : str         =   torch.Generator(device=_a							).manual_seed(_a							)
										_UpperCamelCase      : 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     snake_case     (			self		:   List[str]							)     ->				int:




										'''simple docstring'''



										_UpperCamelCase      : Any         =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator

										_UpperCamelCase      : Union[str, Any]         =   self.get_dummy_components()
										_UpperCamelCase      : str         =   StableDiffusionLDMaDPipeline(**_a							)
										_UpperCamelCase      : int         =   ldmad_pipe.to(_a							)
										ldmad_pipe.set_progress_bar_config(disable=_a							)

										_UpperCamelCase      : Dict         =   self.get_dummy_inputs(_a							)
										_UpperCamelCase      : Any         =   ldmad_pipe(**_a							)
										_UpperCamelCase      : Union[str, Any]         =   output.rgb, output.depth

										_UpperCamelCase      : Tuple         =   rgb[0, -3:, -3:, -1]
										_UpperCamelCase      : int         =   depth[0, -3:, -1]

										assert rgb.shape == (1, 6_4, 6_4, 3)
										assert depth.shape == (1, 6_4, 6_4)

										_UpperCamelCase      : int         =   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]							)
										_UpperCamelCase      : Any         =   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     snake_case     (			self		:   Optional[int]							)     ->				List[str]:




										'''simple docstring'''



										_UpperCamelCase      : Optional[int]         =   self.get_dummy_components()
										_UpperCamelCase      : Optional[Any]         =   StableDiffusionLDMaDPipeline(**_a							)
										_UpperCamelCase      : int         =   ldmad_pipe.to(_a							)
										ldmad_pipe.set_progress_bar_config(disable=_a							)

										_UpperCamelCase      : Union[str, Any]         =   self.get_dummy_inputs(_a							)
										_UpperCamelCase      : Union[str, Any]         =   3 * [inputs['''prompt''']]

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

										_UpperCamelCase      : Any         =   self.get_dummy_inputs(_a							)
										_UpperCamelCase      : List[Any]         =   3 * [inputs.pop('''prompt'''							)]

										_UpperCamelCase      : int         =   ldmad_pipe.tokenizer(
										    _a,			padding='''max_length''',			max_length=ldmad_pipe.tokenizer.model_max_length,			truncation=_a,			return_tensors='''pt''',			)
										_UpperCamelCase      : Tuple         =   text_inputs['''input_ids'''].to(_a							)

										_UpperCamelCase      : Union[str, Any]         =   ldmad_pipe.text_encoder(_a							)[0]

										_UpperCamelCase      : Optional[int]         =   prompt_embeds

										# forward
										_UpperCamelCase      : Dict         =   ldmad_pipe(**_a							)
										_UpperCamelCase      : List[str]         =   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     snake_case     (			self		:   Optional[int]							)     ->				Dict:




										'''simple docstring'''



										_UpperCamelCase      : Optional[Any]         =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
										_UpperCamelCase      : List[str]         =   self.get_dummy_components()
										_UpperCamelCase      : Any         =   PNDMScheduler(skip_prk_steps=_a							)
										_UpperCamelCase      : Any         =   StableDiffusionLDMaDPipeline(**_a							)
										_UpperCamelCase      : Any         =   ldmad_pipe.to(_a							)
										ldmad_pipe.set_progress_bar_config(disable=_a							)

										_UpperCamelCase      : Tuple         =   self.get_dummy_inputs(_a							)
										_UpperCamelCase      : int         =   '''french fries'''
										_UpperCamelCase      : Tuple         =   ldmad_pipe(**_a,			negative_prompt=_a							)

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

										assert rgb.shape == (1, 6_4, 6_4, 3)
										assert depth.shape == (1, 6_4, 6_4)

										_UpperCamelCase      : 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]							)
										_UpperCamelCase      : int         =   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 			_a						(				unittest.TestCase    ):






				def     snake_case     (			self		:   Optional[int]							)     ->				Dict:




										'''simple docstring'''



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






				def     snake_case     (			self		:   Dict,			lowerCAmelCase__		:   List[Any],			lowerCAmelCase__		:   List[Any]="cpu",			lowerCAmelCase__		:   str=torch.floataa,			lowerCAmelCase__		:   Optional[Any]=0							)     ->				Tuple:




										'''simple docstring'''



										_UpperCamelCase      : int         =   torch.Generator(device=_a							).manual_seed(_a							)
										_UpperCamelCase      : Optional[int]         =   np.random.RandomState(_a							).standard_normal((1, 4, 6_4, 6_4)							)
										_UpperCamelCase      : Dict         =   torch.from_numpy(_a							).to(device=_a,			dtype=_a							)
										_UpperCamelCase      : int         =   {
										    '''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     snake_case     (			self		:   Dict							)     ->				Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase      : Tuple         =   StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''							)
										_UpperCamelCase      : Optional[Any]         =   ldmad_pipe.to(_a							)
										ldmad_pipe.set_progress_bar_config(disable=_a							)

										_UpperCamelCase      : int         =   self.get_inputs(_a							)
										_UpperCamelCase      : Optional[int]         =   ldmad_pipe(**_a							)
										_UpperCamelCase      : Any         =   output.rgb, output.depth
										_UpperCamelCase      : Tuple         =   rgb[0, -3:, -3:, -1].flatten()
										_UpperCamelCase      : int         =   rgb[0, -3:, -1].flatten()

										assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
										assert depth.shape == (1, 5_1_2, 5_1_2)

										_UpperCamelCase      : int         =   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]							)
										_UpperCamelCase      : Dict         =   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 			_a						(				unittest.TestCase    ):






				def     snake_case     (			self		:   Union[str, Any]							)     ->				Tuple:




										'''simple docstring'''



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






				def     snake_case     (			self		:   List[str],			lowerCAmelCase__		:   Dict,			lowerCAmelCase__		:   Any="cpu",			lowerCAmelCase__		:   Tuple=torch.floataa,			lowerCAmelCase__		:   Any=0							)     ->				Optional[int]:




										'''simple docstring'''



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






				def     snake_case     (			self		:   int							)     ->				Tuple:




										'''simple docstring'''



										_UpperCamelCase      : Optional[Any]         =   StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''							).to(_a							)
										ldmad_pipe.set_progress_bar_config(disable=_a							)

										_UpperCamelCase      : Any         =   self.get_inputs(_a							)
										_UpperCamelCase      : List[str]         =   ldmad_pipe(**_a							)
										_UpperCamelCase      : Any         =   output.rgb, output.depth

										_UpperCamelCase      : str         =   0.495_586
										_UpperCamelCase      : Union[str, Any]         =   0.33_795_515
										_UpperCamelCase      : Any         =   112.48_518
										_UpperCamelCase      : List[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     snake_case     (			self		:   Optional[Any]							)     ->				int:




										'''simple docstring'''



										_UpperCamelCase      : Optional[Any]         =   StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d-4c'''							).to(_a							)
										ldmad_pipe.set_progress_bar_config(disable=_a							)

										_UpperCamelCase      : List[str]         =   self.get_inputs(_a							)
										_UpperCamelCase      : List[Any]         =   ldmad_pipe(**_a							)
										_UpperCamelCase      : Union[str, Any]         =   output.rgb, output.depth

										_UpperCamelCase      : int         =   0.4_194_127
										_UpperCamelCase      : str         =   0.35_375_586
										_UpperCamelCase      : str         =   0.5_638_502
										_UpperCamelCase      : Tuple         =   0.34_686_103

										assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
										assert depth.shape == (1, 5_1_2, 5_1_2, 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"""




UpperCamelCase_									=[
    """DownloadConfig""",
    """DownloadManager""",
    """DownloadMode""",
    """StreamingDownloadManager""",
]

from .download_config import DownloadConfig
from .download_manager import DownloadManager, DownloadMode
from .streaming_download_manager import StreamingDownloadManager



