




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     lowerCamelCase__(					unittest.TestCase):
      UpperCAmelCase__      :	Tuple    		=   StableDiffusionLDMaDPipeline
      UpperCAmelCase__      :	Optional[int]    		=   TEXT_TO_IMAGE_PARAMS
      UpperCAmelCase__      :	Any    		=   TEXT_TO_IMAGE_BATCH_PARAMS
      UpperCAmelCase__      :	Union[str, Any]    		=   TEXT_TO_IMAGE_IMAGE_PARAMS
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             torch.manual_seed(0   )
             __lowerCamelCase   =						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   =						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   =						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   =						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=10_00    ,       )
             __lowerCamelCase   =						CLIPTextModel(__lowerCAmelCase   )
             __lowerCamelCase   =						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )

             __lowerCamelCase   =						{
                 """unet""": unet,
                 """scheduler""": scheduler,
                 """vae""": vae,
                 """text_encoder""": text_encoder,
                 """tokenizer""": tokenizer,
                 """safety_checker""": None,
                 """feature_extractor""": None,
             }
             return components
      def 			lowerCAmelCase__						(				self: Optional[Any]    ,       UpperCamelCase_: Union[str, Any]    ,       UpperCamelCase_: Optional[int]=0   ):







             if str(__lowerCAmelCase   ).startswith("""mps"""   ):
                    __lowerCamelCase   =						torch.manual_seed(__lowerCAmelCase   )
             else:
                    __lowerCamelCase   =						torch.Generator(device=__lowerCAmelCase   ).manual_seed(__lowerCAmelCase   )
             __lowerCamelCase   =						{
                 """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 			lowerCAmelCase__						(				self: List[Any]   ):
             __lowerCamelCase   =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator

             __lowerCamelCase   =						self.get_dummy_components()
             __lowerCamelCase   =						StableDiffusionLDMaDPipeline(**__lowerCAmelCase   )
             __lowerCamelCase   =						ldmad_pipe.to(__lowerCAmelCase   )
             ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

             __lowerCamelCase   =						self.get_dummy_inputs(__lowerCAmelCase   )
             __lowerCamelCase   =						ldmad_pipe(**__lowerCAmelCase   )
             __lowerCamelCase,   __lowerCamelCase   =						output.rgb, output.depth

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

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

             __lowerCamelCase   =						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   =						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 			lowerCAmelCase__						(				self: List[str]   ):
             __lowerCamelCase   =						self.get_dummy_components()
             __lowerCamelCase   =						StableDiffusionLDMaDPipeline(**__lowerCAmelCase   )
             __lowerCamelCase   =						ldmad_pipe.to(__lowerCAmelCase   )
             ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

             __lowerCamelCase   =						self.get_dummy_inputs(__lowerCAmelCase   )
             __lowerCamelCase   =						3 * [inputs["""prompt"""]]

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

             __lowerCamelCase   =						self.get_dummy_inputs(__lowerCAmelCase   )
             __lowerCamelCase   =						3 * [inputs.pop("""prompt"""   )]

             __lowerCamelCase   =						ldmad_pipe.tokenizer(
                 __lowerCAmelCase    ,       padding="""max_length"""    ,       max_length=ldmad_pipe.tokenizer.model_max_length    ,       truncation=__lowerCAmelCase    ,       return_tensors="""pt"""    ,       )
             __lowerCamelCase   =						text_inputs["""input_ids"""].to(__lowerCAmelCase   )

             __lowerCamelCase   =						ldmad_pipe.text_encoder(__lowerCAmelCase   )[0]

             __lowerCamelCase   =						prompt_embeds

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

             __lowerCamelCase   =						self.get_dummy_inputs(__lowerCAmelCase   )
             __lowerCamelCase   =						"""french fries"""
             __lowerCamelCase   =						ldmad_pipe(**__lowerCAmelCase    ,       negative_prompt=__lowerCAmelCase   )

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

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

             __lowerCamelCase   =						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   =						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     lowerCamelCase__(					unittest.TestCase):
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()
      def 			lowerCAmelCase__						(				self: Any    ,       UpperCamelCase_: List[str]    ,       UpperCamelCase_: int="cpu"    ,       UpperCamelCase_: int=torch.floataa    ,       UpperCamelCase_: Optional[Any]=0   ):
             __lowerCamelCase   =						torch.Generator(device=__lowerCAmelCase   ).manual_seed(__lowerCAmelCase   )
             __lowerCamelCase   =						np.random.RandomState(__lowerCAmelCase   ).standard_normal((1, 4, 64, 64)   )
             __lowerCamelCase   =						torch.from_numpy(__lowerCAmelCase   ).to(device=__lowerCAmelCase    ,       dtype=__lowerCAmelCase   )
             __lowerCamelCase   =						{
                 """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 			lowerCAmelCase__						(				self: int   ):
             __lowerCamelCase   =						StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""   )
             __lowerCamelCase   =						ldmad_pipe.to(__lowerCAmelCase   )
             ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

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

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

             __lowerCamelCase   =						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   =						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     lowerCamelCase__(					unittest.TestCase):
      def 			lowerCAmelCase__						(				self: Optional[int]   ):
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()
      def 			lowerCAmelCase__						(				self: Dict    ,       UpperCamelCase_: Tuple    ,       UpperCamelCase_: Any="cpu"    ,       UpperCamelCase_: int=torch.floataa    ,       UpperCamelCase_: Any=0   ):
             __lowerCamelCase   =						torch.Generator(device=__lowerCAmelCase   ).manual_seed(__lowerCAmelCase   )
             __lowerCamelCase   =						np.random.RandomState(__lowerCAmelCase   ).standard_normal((1, 4, 64, 64)   )
             __lowerCamelCase   =						torch.from_numpy(__lowerCAmelCase   ).to(device=__lowerCAmelCase    ,       dtype=__lowerCAmelCase   )
             __lowerCamelCase   =						{
                 """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 			lowerCAmelCase__						(				self: Optional[int]   ):
             __lowerCamelCase   =						StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""   ).to(__lowerCAmelCase   )
             ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

             __lowerCamelCase   =						self.get_inputs(__lowerCAmelCase   )
             __lowerCamelCase   =						ldmad_pipe(**__lowerCAmelCase   )
             __lowerCamelCase,   __lowerCamelCase   =						output.rgb, output.depth

             __lowerCamelCase   =						0.49_5586
             __lowerCamelCase   =						0.3379_5515
             __lowerCamelCase   =						112.4_8518
             __lowerCamelCase   =						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 			lowerCAmelCase__						(				self: List[str]   ):
             __lowerCamelCase   =						StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d-4c"""   ).to(__lowerCAmelCase   )
             ldmad_pipe.set_progress_bar_config(disable=__lowerCAmelCase   )

             __lowerCamelCase   =						self.get_inputs(__lowerCAmelCase   )
             __lowerCamelCase   =						ldmad_pipe(**__lowerCAmelCase   )
             __lowerCamelCase,   __lowerCamelCase   =						output.rgb, output.depth

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

             assert rgb.shape == (1, 5_12, 5_12, 3)
             assert depth.shape == (1, 5_12, 5_12, 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 						=      logging.get_logger(__name__)

_a 						=      {
    "hustvl/yolos-small": "https://huggingface.co/hustvl/yolos-small/resolve/main/config.json",
    # See all YOLOS models at https://huggingface.co/models?filter=yolos
}






class 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						"""yolos"""





			def __init__(    self       ,  __lowerCAmelCase=7_6_8       ,  __lowerCAmelCase=1_2       ,  __lowerCAmelCase=1_2       ,  __lowerCAmelCase=3_0_7_2       ,  __lowerCAmelCase="gelu"       ,  __lowerCAmelCase=0.0       ,  __lowerCAmelCase=0.0       ,  __lowerCAmelCase=0.02       ,  __lowerCAmelCase=1E-12       ,  __lowerCAmelCase=[5_1_2, 8_6_4]       ,  __lowerCAmelCase=1_6       ,  __lowerCAmelCase=3       ,  __lowerCAmelCase=True       ,  __lowerCAmelCase=1_0_0       ,  __lowerCAmelCase=True       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=1       ,  __lowerCAmelCase=5       ,  __lowerCAmelCase=2       ,  __lowerCAmelCase=5       ,  __lowerCAmelCase=2       ,  __lowerCAmelCase=0.1       ,  **__lowerCAmelCase       ,  ):



				'''simple docstring'''







				super().__init__(**__lowerCAmelCase							)

				lowerCamelCase__								=		hidden_size
				lowerCamelCase__								=		num_hidden_layers
				lowerCamelCase__								=		num_attention_heads
				lowerCamelCase__								=		intermediate_size
				lowerCamelCase__								=		hidden_act
				lowerCamelCase__								=		hidden_dropout_prob
				lowerCamelCase__								=		attention_probs_dropout_prob
				lowerCamelCase__								=		initializer_range
				lowerCamelCase__								=		layer_norm_eps
				lowerCamelCase__								=		image_size
				lowerCamelCase__								=		patch_size
				lowerCamelCase__								=		num_channels
				lowerCamelCase__								=		qkv_bias
				lowerCamelCase__								=		num_detection_tokens
				lowerCamelCase__								=		use_mid_position_embeddings
				lowerCamelCase__								=		auxiliary_loss
				# Hungarian matcher
				lowerCamelCase__								=		class_cost
				lowerCamelCase__								=		bbox_cost
				lowerCamelCase__								=		giou_cost
				# Loss coefficients
				lowerCamelCase__								=		bbox_loss_coefficient
				lowerCamelCase__								=		giou_loss_coefficient
				lowerCamelCase__								=		eos_coefficient






class 					__A    (    lowerCAmelCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						version.parse("""1.11""" )





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







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





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return 1E-4





			@property
			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return 1_2



