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 __snake_case				(						unittest.TestCase		):
   lowerCAmelCase_														=							StableDiffusionLDMaDPipeline
   lowerCAmelCase_														=							TEXT_TO_IMAGE_PARAMS
   lowerCAmelCase_														=							TEXT_TO_IMAGE_BATCH_PARAMS
   lowerCAmelCase_														=							TEXT_TO_IMAGE_IMAGE_PARAMS




   def 	__a						(    self			:				List[Any]      ):



     """simple docstring"""






     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      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       ,	)
     SCREAMING_SNAKE_CASE__										=      DDIMScheduler(
         beta_start=0.0_00_85       ,	beta_end=0.0_12       ,	beta_schedule="""scaled_linear"""       ,	clip_sample=A_       ,	set_alpha_to_one=A_       ,	)
     torch.manual_seed(0      )
     SCREAMING_SNAKE_CASE__										=      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      )
     SCREAMING_SNAKE_CASE__										=      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       ,	)
     SCREAMING_SNAKE_CASE__										=      CLIPTextModel(A_      )
     SCREAMING_SNAKE_CASE__										=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

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




   def 	__a						(    self			:				int       ,	_lowercase			:				Tuple       ,	_lowercase			:				Optional[int]=0      ):



     """simple docstring"""






     if str(A_      ).startswith("""mps"""      ):
       SCREAMING_SNAKE_CASE__										=      torch.manual_seed(A_      )
     else:
       SCREAMING_SNAKE_CASE__										=      torch.Generator(device=A_      ).manual_seed(A_      )
     SCREAMING_SNAKE_CASE__										=      {
         """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			:				Tuple      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      """cpu"""  # ensure determinism for the device-dependent torch.Generator

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

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

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

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

     SCREAMING_SNAKE_CASE__										=      np.array(
         [0.37_33_81_76, 0.7_02_47, 0.74_20_31_93, 0.51_64_36_04, 0.58_25_67_93, 0.60_93_21_36, 0.4_18_10_95, 0.48_35_58_77, 0.46_53_52_62]      )
     SCREAMING_SNAKE_CASE__										=      np.array([1_03.4_67_27, 85.81_20_04, 87.84_92_36]      )

     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			:				Tuple      ):



     """simple docstring"""






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

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(A_      )
     SCREAMING_SNAKE_CASE__										=      3 * [inputs["""prompt"""]]

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

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(A_      )
     SCREAMING_SNAKE_CASE__										=      3 * [inputs.pop("""prompt"""      )]

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

     SCREAMING_SNAKE_CASE__										=      ldmad_pipe.text_encoder(A_      )[0]

     SCREAMING_SNAKE_CASE__										=      prompt_embeds

     # forward
     SCREAMING_SNAKE_CASE__										=      ldmad_pipe(**A_      )
     SCREAMING_SNAKE_CASE__							,     SCREAMING_SNAKE_CASE__										=      output.rgb, output.depth
     SCREAMING_SNAKE_CASE__										=      rgb_slice_a[0, -3:, -3:, -1]
     SCREAMING_SNAKE_CASE__										=      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]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE__										=      self.get_dummy_components()
     SCREAMING_SNAKE_CASE__										=      PNDMScheduler(skip_prk_steps=A_      )
     SCREAMING_SNAKE_CASE__										=      StableDiffusionLDMaDPipeline(**A_      )
     SCREAMING_SNAKE_CASE__										=      ldmad_pipe.to(A_      )
     ldmad_pipe.set_progress_bar_config(disable=A_      )

     SCREAMING_SNAKE_CASE__										=      self.get_dummy_inputs(A_      )
     SCREAMING_SNAKE_CASE__										=      """french fries"""
     SCREAMING_SNAKE_CASE__										=      ldmad_pipe(**A_       ,	negative_prompt=A_      )

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

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

     SCREAMING_SNAKE_CASE__										=      np.array(
         [0.3_70_44, 0.71_81_15_03, 0.7_22_32_51, 0.48_60_36_75, 0.5_63_83_91, 0.6_36_49_48, 0.42_83_37_04, 0.4_90_13_15, 0.47_92_62_17]      )
     SCREAMING_SNAKE_CASE__										=      np.array([1_07.8_47_38, 84.6_28_02, 89.96_21_35]      )
     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 __snake_case				(						unittest.TestCase		):




   def 	__a						(    self			:				Optional[Any]      ):



     """simple docstring"""






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




   def 	__a						(    self			:				Tuple       ,	_lowercase			:				int       ,	_lowercase			:				int="cpu"       ,	_lowercase			:				int=torch.floataa       ,	_lowercase			:				List[Any]=0      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      torch.Generator(device=A_      ).manual_seed(A_      )
     SCREAMING_SNAKE_CASE__										=      np.random.RandomState(A_      ).standard_normal((1, 4, 64, 64)      )
     SCREAMING_SNAKE_CASE__										=      torch.from_numpy(A_      ).to(device=A_       ,	dtype=A_      )
     SCREAMING_SNAKE_CASE__										=      {
         """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			:				List[Any]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""      )
     SCREAMING_SNAKE_CASE__										=      ldmad_pipe.to(A_      )
     ldmad_pipe.set_progress_bar_config(disable=A_      )

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

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

     SCREAMING_SNAKE_CASE__										=      np.array(
         [0.53_80_54_65, 0.56_70_73_05, 0.5_48_65_15, 0.57_01_22_36, 0.5_81_45_11, 0.56_25_34_87, 0.54_84_30_14, 0.55_09_22_63, 0.6_45_97_06]      )
     SCREAMING_SNAKE_CASE__										=      np.array(
         [0.9_26_37_81, 0.6_67_86_72, 0.5_48_65_15, 0.92_20_21_45, 0.67_83_11_35, 0.56_25_34_87, 0.9_24_16_94, 0.7_55_14_78, 0.6_45_97_06]      )
     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 __snake_case				(						unittest.TestCase		):




   def 	__a						(    self			:				Union[str, Any]      ):



     """simple docstring"""






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




   def 	__a						(    self			:				Optional[int]       ,	_lowercase			:				List[Any]       ,	_lowercase			:				Tuple="cpu"       ,	_lowercase			:				int=torch.floataa       ,	_lowercase			:				int=0      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      torch.Generator(device=A_      ).manual_seed(A_      )
     SCREAMING_SNAKE_CASE__										=      np.random.RandomState(A_      ).standard_normal((1, 4, 64, 64)      )
     SCREAMING_SNAKE_CASE__										=      torch.from_numpy(A_      ).to(device=A_       ,	dtype=A_      )
     SCREAMING_SNAKE_CASE__										=      {
         """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			:				Optional[int]      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""      ).to(A_      )
     ldmad_pipe.set_progress_bar_config(disable=A_      )

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

     SCREAMING_SNAKE_CASE__										=      0.49_55_86
     SCREAMING_SNAKE_CASE__										=      0.33_79_55_15
     SCREAMING_SNAKE_CASE__										=      1_12.4_85_18
     SCREAMING_SNAKE_CASE__										=      98.48_97_46
     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			:				int      ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__										=      StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d-4c"""      ).to(A_      )
     ldmad_pipe.set_progress_bar_config(disable=A_      )

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

     SCREAMING_SNAKE_CASE__										=      0.4_19_41_27
     SCREAMING_SNAKE_CASE__										=      0.35_37_55_86
     SCREAMING_SNAKE_CASE__										=      0.5_63_85_02
     SCREAMING_SNAKE_CASE__										=      0.34_68_61_03

     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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


_A		       =						{'configuration_vit_mae': ['VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ViTMAEConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A		       =						[
       'VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST',
       'ViTMAEForPreTraining',
       'ViTMAELayer',
       'ViTMAEModel',
       'ViTMAEPreTrainedModel',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _A		       =						[
       'TFViTMAEForPreTraining',
       'TFViTMAEModel',
       'TFViTMAEPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_vit_mae import VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTMAEConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_vit_mae import (
          VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST,
          ViTMAEForPreTraining,
          ViTMAELayer,
          ViTMAEModel,
          ViTMAEPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_vit_mae import TFViTMAEForPreTraining, TFViTMAEModel, TFViTMAEPreTrainedModel


else:
   import sys

   _A		       =						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





