






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



      A          =    StableDiffusionLDMaDPipeline
      A          =    TEXT_TO_IMAGE_PARAMS
      A          =    TEXT_TO_IMAGE_BATCH_PARAMS
      A          =    TEXT_TO_IMAGE_IMAGE_PARAMS
      def 			snake_case_      (   self		):



          '''simple docstring'''
          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__   :int      							=				DDIMScheduler(
              beta_start=0.0_0085    ,	beta_end=0.012    ,	beta_schedule="scaled_linear"    ,	clip_sample=_lowerCAmelCase    ,	set_alpha_to_one=_lowerCAmelCase    ,	)
          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__   :List[Any]      							=				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=1_000    ,	)
          lowerCAmelCase__   :Optional[Any]      							=				CLIPTextModel(_lowerCAmelCase		)
          lowerCAmelCase__   :str      							=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

          lowerCAmelCase__   :Union[str, Any]      							=				{
              'unet': unet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'safety_checker': None,
              'feature_extractor': None,
          }
          return components
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=0		):



          '''simple docstring'''





          if str(_lowerCAmelCase		).startswith("mps"		):
              lowerCAmelCase__   :List[Any]      							=				torch.manual_seed(_lowerCAmelCase		)
          else:
              lowerCAmelCase__   :int      							=				torch.Generator(device=_lowerCAmelCase		).manual_seed(_lowerCAmelCase		)
          lowerCAmelCase__   :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 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				'cpu'  # ensure determinism for the device-dependent torch.Generator

          lowerCAmelCase__   :Dict      							=				self.get_dummy_components()
          lowerCAmelCase__   :Union[str, Any]      							=				StableDiffusionLDMaDPipeline(**_lowerCAmelCase		)
          lowerCAmelCase__   :str      							=				ldmad_pipe.to(_lowerCAmelCase		)
          ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :int      							=				self.get_dummy_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				ldmad_pipe(**_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				output.rgb, output.depth

          lowerCAmelCase__   :int      							=				rgb[0, -3:, -3:, -1]
          lowerCAmelCase__   :Dict      							=				depth[0, -3:, -1]

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

          lowerCAmelCase__   :Any      							=				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]		)
          lowerCAmelCase__   :Dict      							=				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 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				self.get_dummy_components()
          lowerCAmelCase__   :Any      							=				StableDiffusionLDMaDPipeline(**_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				ldmad_pipe.to(_lowerCAmelCase		)
          ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[Any]      							=				self.get_dummy_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				3 * [inputs['prompt']]

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

          lowerCAmelCase__   :Union[str, Any]      							=				self.get_dummy_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				3 * [inputs.pop("prompt"		)]

          lowerCAmelCase__   :int      							=				ldmad_pipe.tokenizer(
              _lowerCAmelCase    ,	padding="max_length"    ,	max_length=ldmad_pipe.tokenizer.model_max_length    ,	truncation=_lowerCAmelCase    ,	return_tensors="pt"    ,	)
          lowerCAmelCase__   :int      							=				text_inputs['input_ids'].to(_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[Any]      							=				ldmad_pipe.text_encoder(_lowerCAmelCase		)[0]

          lowerCAmelCase__   :int      							=				prompt_embeds

          # forward
          lowerCAmelCase__   :str      							=				ldmad_pipe(**_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				output.rgb, output.depth
          lowerCAmelCase__   :Optional[int]      							=				rgb_slice_a[0, -3:, -3:, -1]
          lowerCAmelCase__   :List[str]      							=				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		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				'cpu'  # ensure determinism for the device-dependent torch.Generator
          lowerCAmelCase__   :Optional[int]      							=				self.get_dummy_components()
          lowerCAmelCase__   :Optional[Any]      							=				PNDMScheduler(skip_prk_steps=_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				StableDiffusionLDMaDPipeline(**_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[Any]      							=				ldmad_pipe.to(_lowerCAmelCase		)
          ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

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

          lowerCAmelCase__   :Any      							=				output.rgb, output.depth
          lowerCAmelCase__   :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__   :List[str]      							=				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]		)
          lowerCAmelCase__   :str      							=				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 			snake_case_      (   self		):



          '''simple docstring'''
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase="cpu"    ,	_lowerCAmelCase=torch.floataa    ,	_lowerCAmelCase=0		):



          '''simple docstring'''
          lowerCAmelCase__   :int      							=				torch.Generator(device=_lowerCAmelCase		).manual_seed(_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				np.random.RandomState(_lowerCAmelCase		).standard_normal((1, 4, 64, 64)		)
          lowerCAmelCase__   :int      							=				torch.from_numpy(_lowerCAmelCase		).to(device=_lowerCAmelCase    ,	dtype=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				{
              '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		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				StableDiffusionLDMaDPipeline.from_pretrained("Intel/ldm3d"		)
          lowerCAmelCase__   :List[Any]      							=				ldmad_pipe.to(_lowerCAmelCase		)
          ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :Tuple      							=				self.get_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :int      							=				ldmad_pipe(**_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				output.rgb, output.depth
          lowerCAmelCase__   :Dict      							=				rgb[0, -3:, -3:, -1].flatten()
          lowerCAmelCase__   :int      							=				rgb[0, -3:, -1].flatten()

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

          lowerCAmelCase__   :List[str]      							=				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]		)
          lowerCAmelCase__   :Any      							=				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 			snake_case_      (   self		):



          '''simple docstring'''
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase="cpu"    ,	_lowerCAmelCase=torch.floataa    ,	_lowerCAmelCase=0		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				torch.Generator(device=_lowerCAmelCase		).manual_seed(_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[Any]      							=				np.random.RandomState(_lowerCAmelCase		).standard_normal((1, 4, 64, 64)		)
          lowerCAmelCase__   :Optional[int]      							=				torch.from_numpy(_lowerCAmelCase		).to(device=_lowerCAmelCase    ,	dtype=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				{
              '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 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				StableDiffusionLDMaDPipeline.from_pretrained("Intel/ldm3d"		).to(_lowerCAmelCase		)
          ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				self.get_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :int      							=				ldmad_pipe(**_lowerCAmelCase		)
          lowerCAmelCase__   :Any      							=				output.rgb, output.depth

          lowerCAmelCase__   :str      							=				0.49_5586
          lowerCAmelCase__   :List[Any]      							=				0.3379_5515
          lowerCAmelCase__   :Optional[Any]      							=				112.4_8518
          lowerCAmelCase__   :int      							=				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 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				StableDiffusionLDMaDPipeline.from_pretrained("Intel/ldm3d-4c"		).to(_lowerCAmelCase		)
          ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				self.get_inputs(_lowerCAmelCase		)
          lowerCAmelCase__   :int      							=				ldmad_pipe(**_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				output.rgb, output.depth

          lowerCAmelCase__   :List[str]      							=				0.419_4127
          lowerCAmelCase__   :int      							=				0.3537_5586
          lowerCAmelCase__   :int      							=				0.563_8502
          lowerCAmelCase__   :Optional[Any]      							=				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











from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_a    :   Tuple        =			logging.get_logger(__name__)

_a    :   Optional[int]        =			{
    """facebook/levit-128S""": """https://huggingface.co/facebook/levit-128S/resolve/main/config.json""",
    # See all LeViT models at https://huggingface.co/models?filter=levit
}







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    '''levit'''
      def __init__(   self    ,	_lowerCAmelCase=224    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=2    ,	_lowerCAmelCase=1    ,	_lowerCAmelCase=16    ,	_lowerCAmelCase=[128, 256, 384]    ,	_lowerCAmelCase=[4, 8, 12]    ,	_lowerCAmelCase=[4, 4, 4]    ,	_lowerCAmelCase=[16, 16, 16]    ,	_lowerCAmelCase=0    ,	_lowerCAmelCase=[2, 2, 2]    ,	_lowerCAmelCase=[2, 2, 2]    ,	_lowerCAmelCase=0.02    ,	**_lowerCAmelCase    ,	):



          '''simple docstring'''
          super().__init__(**_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				image_size
          lowerCAmelCase__   :Dict      							=				num_channels
          lowerCAmelCase__   :List[Any]      							=				kernel_size
          lowerCAmelCase__   :List[Any]      							=				stride
          lowerCAmelCase__   :Tuple      							=				padding
          lowerCAmelCase__   :str      							=				hidden_sizes
          lowerCAmelCase__   :Union[str, Any]      							=				num_attention_heads
          lowerCAmelCase__   :Optional[int]      							=				depths
          lowerCAmelCase__   :Union[str, Any]      							=				key_dim
          lowerCAmelCase__   :str      							=				drop_path_rate
          lowerCAmelCase__   :int      							=				patch_size
          lowerCAmelCase__   :Dict      							=				attention_ratio
          lowerCAmelCase__   :Dict      							=				mlp_ratio
          lowerCAmelCase__   :Optional[int]      							=				initializer_range
          lowerCAmelCase__   :Tuple      							=				[
              ["Subsample", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
              ["Subsample", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
          ]







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    version.parse('''1.11'''		)
      @property
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return OrderedDict(
              [
                  ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
              ]		)



      @property
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return 1e-4



