


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     UpperCAmelCase__       (					unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		List[str]      =  StableDiffusionLDMaDPipeline
    UpperCAmelCase__						:		Tuple      =  TEXT_TO_IMAGE_PARAMS
    UpperCAmelCase__						:		Optional[Any]      =  TEXT_TO_IMAGE_BATCH_PARAMS
    UpperCAmelCase__						:		str      =  TEXT_TO_IMAGE_IMAGE_PARAMS






    def _a					(			self    )  ->       Union[str, Any]:
         torch.manual_seed(0    )
         __UpperCamelCase    							=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    							=DDIMScheduler(
             beta_start=0.0_0085					,   beta_end=0.012					,   beta_schedule='scaled_linear'					,   clip_sample=A_					,   set_alpha_to_one=A_					,   )
         torch.manual_seed(0    )
         __UpperCamelCase    							=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    							=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    							=CLIPTextModel(A_    )
         __UpperCamelCase    							=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

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






    def _a					(			self					,   A_					,   A_=0    )  ->       List[str]:
         if str(A_    ).startswith('mps'    ):
              __UpperCamelCase    							=torch.manual_seed(A_    )
         else:
              __UpperCamelCase    							=torch.Generator(device=A_    ).manual_seed(A_    )
         __UpperCamelCase    							={
             '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    )  ->       Dict:
         __UpperCamelCase    							='cpu'  # ensure determinism for the device-dependent torch.Generator

         __UpperCamelCase    							=self.get_dummy_components()
         __UpperCamelCase    							=StableDiffusionLDMaDPipeline(**A_    )
         __UpperCamelCase    							=ldmad_pipe.to(A_    )
         ldmad_pipe.set_progress_bar_config(disable=A_    )

         __UpperCamelCase    							=self.get_dummy_inputs(A_    )
         __UpperCamelCase    							=ldmad_pipe(**A_    )
         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth

         __UpperCamelCase    							=rgb[0, -3:, -3:, -1]
         __UpperCamelCase    							=depth[0, -3:, -1]

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

         __UpperCamelCase    							=np.array(
             [0.3733_8176, 0.7_0247, 0.7420_3193, 0.5164_3604, 0.5825_6793, 0.6093_2136, 0.418_1095, 0.4835_5877, 0.4653_5262]    )
         __UpperCamelCase    							=np.array([103.4_6727, 85.81_2004, 87.84_9236]    )

         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    )  ->       str:
         __UpperCamelCase    							=self.get_dummy_components()
         __UpperCamelCase    							=StableDiffusionLDMaDPipeline(**A_    )
         __UpperCamelCase    							=ldmad_pipe.to(A_    )
         ldmad_pipe.set_progress_bar_config(disable=A_    )

         __UpperCamelCase    							=self.get_dummy_inputs(A_    )
         __UpperCamelCase    							=3 * [inputs['prompt']]

         # forward
         __UpperCamelCase    							=ldmad_pipe(**A_    )
         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth
         __UpperCamelCase    							=rgb_slice_a[0, -3:, -3:, -1]
         __UpperCamelCase    							=depth_slice_a[0, -3:, -1]

         __UpperCamelCase    							=self.get_dummy_inputs(A_    )
         __UpperCamelCase    							=3 * [inputs.pop('prompt'    )]

         __UpperCamelCase    							=ldmad_pipe.tokenizer(
             A_					,   padding='max_length'					,   max_length=ldmad_pipe.tokenizer.model_max_length					,   truncation=A_					,   return_tensors='pt'					,   )
         __UpperCamelCase    							=text_inputs['input_ids'].to(A_    )

         __UpperCamelCase    							=ldmad_pipe.text_encoder(A_    )[0]

         __UpperCamelCase    							=prompt_embeds

         # forward
         __UpperCamelCase    							=ldmad_pipe(**A_    )
         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth
         __UpperCamelCase    							=rgb_slice_a[0, -3:, -3:, -1]
         __UpperCamelCase    							=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]:
         __UpperCamelCase    							='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase    							=self.get_dummy_components()
         __UpperCamelCase    							=PNDMScheduler(skip_prk_steps=A_    )
         __UpperCamelCase    							=StableDiffusionLDMaDPipeline(**A_    )
         __UpperCamelCase    							=ldmad_pipe.to(A_    )
         ldmad_pipe.set_progress_bar_config(disable=A_    )

         __UpperCamelCase    							=self.get_dummy_inputs(A_    )
         __UpperCamelCase    							='french fries'
         __UpperCamelCase    							=ldmad_pipe(**A_					,   negative_prompt=A_    )

         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth
         __UpperCamelCase    							=rgb[0, -3:, -3:, -1]
         __UpperCamelCase    							=depth[0, -3:, -1]

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

         __UpperCamelCase    							=np.array(
             [0.3_7044, 0.7181_1503, 0.722_3251, 0.4860_3675, 0.563_8391, 0.636_4948, 0.4283_3704, 0.490_1315, 0.4792_6217]    )
         __UpperCamelCase    							=np.array([107.8_4738, 84.6_2802, 89.96_2135]    )
         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     UpperCAmelCase__       (					unittest.TestCase		):



    """simple docstring"""






    def _a					(			self    )  ->       Any:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






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

         __UpperCamelCase    							=self.get_inputs(A_    )
         __UpperCamelCase    							=ldmad_pipe(**A_    )
         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth
         __UpperCamelCase    							=rgb[0, -3:, -3:, -1].flatten()
         __UpperCamelCase    							=rgb[0, -3:, -1].flatten()

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

         __UpperCamelCase    							=np.array(
             [0.5380_5465, 0.5670_7305, 0.548_6515, 0.5701_2236, 0.581_4511, 0.5625_3487, 0.5484_3014, 0.5509_2263, 0.645_9706]    )
         __UpperCamelCase    							=np.array(
             [0.926_3781, 0.667_8672, 0.548_6515, 0.9220_2145, 0.6783_1135, 0.5625_3487, 0.924_1694, 0.755_1478, 0.645_9706]    )
         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     UpperCAmelCase__       (					unittest.TestCase		):



    """simple docstring"""






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






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

         __UpperCamelCase    							=self.get_inputs(A_    )
         __UpperCamelCase    							=ldmad_pipe(**A_    )
         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth

         __UpperCamelCase    							=0.49_5586
         __UpperCamelCase    							=0.3379_5515
         __UpperCamelCase    							=112.4_8518
         __UpperCamelCase    							=98.48_9746
         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    )  ->       Optional[Any]:
         __UpperCamelCase    							=StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'    ).to(A_    )
         ldmad_pipe.set_progress_bar_config(disable=A_    )

         __UpperCamelCase    							=self.get_inputs(A_    )
         __UpperCamelCase    							=ldmad_pipe(**A_    )
         __UpperCamelCase							,      __UpperCamelCase    							=output.rgb, output.depth

         __UpperCamelCase    							=0.419_4127
         __UpperCamelCase    							=0.3537_5586
         __UpperCamelCase    							=0.563_8502
         __UpperCamelCase    							=0.3468_6103

         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







import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


__a  			=					logging.get_logger(__name__)
class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	a			:Union[str, Any]         =	'upernet'
	def __init__(       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Dict=None       ,							SCREAMING_SNAKE_CASE_  :			str=5_1_2       ,							SCREAMING_SNAKE_CASE_  :			Tuple=0.02       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=[1, 2, 3, 6]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]=True       ,							SCREAMING_SNAKE_CASE_  :			Tuple=0.4       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]=3_8_4       ,							SCREAMING_SNAKE_CASE_  :			str=2_5_6       ,							SCREAMING_SNAKE_CASE_  :			Dict=1       ,							SCREAMING_SNAKE_CASE_  :			Tuple=False       ,							SCREAMING_SNAKE_CASE_  :			str=2_5_5       ,							**SCREAMING_SNAKE_CASE_  :			str       ,							)       ->     int:
					super().__init__(**SCREAMING_SNAKE_CASE_       )

					if backbone_config is None:
									logger.info('''`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'''       )
									lowercase_       =					CONFIG_MAPPING['''resnet'''](out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']       )
					elif isinstance(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ):
									lowercase_       =					backbone_config.get('''model_type'''       )
									lowercase_       =					CONFIG_MAPPING[backbone_model_type]
									lowercase_       =					config_class.from_dict(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					backbone_config
					lowercase_       =					hidden_size
					lowercase_       =					initializer_range
					lowercase_       =					pool_scales
					lowercase_       =					use_auxiliary_head
					lowercase_       =					auxiliary_loss_weight
					lowercase_       =					auxiliary_in_channels
					lowercase_       =					auxiliary_channels
					lowercase_       =					auxiliary_num_convs
					lowercase_       =					auxiliary_concat_input
					lowercase_       =					loss_ignore_index




	def 						_lowercase  (       self  :			List[str]       )       ->     List[str]:
					lowercase_       =					copy.deepcopy(self.__dict__       )
					lowercase_       =					self.backbone_config.to_dict()
					lowercase_       =					self.__class__.model_type
					return output




