



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




 __magic_name__					: List[Any]  =      StableDiffusionLDMaDPipeline
 __magic_name__					: Any  =      TEXT_TO_IMAGE_PARAMS
 __magic_name__					: Optional[int]  =      TEXT_TO_IMAGE_BATCH_PARAMS
 __magic_name__					: Union[str, Any]  =      TEXT_TO_IMAGE_IMAGE_PARAMS




 def       lowercase_  (   self) ->      str:



      """simple docstring"""







      torch.manual_seed(0)
      a_       		=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							,  )
      a_       		=DDIMScheduler(
          beta_start=0.0_0_0_8_5							,  beta_end=0.0_1_2							,  beta_schedule="scaled_linear"							,  clip_sample=lowerCamelCase_							,  set_alpha_to_one=lowerCamelCase_							,  )
      torch.manual_seed(0)
      a_       		=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)
      a_       		=CLIPTextConfig(
          bos_token_id=0							,  eos_token_id=2							,  hidden_size=3_2							,  intermediate_size=3_7							,  layer_norm_eps=1e-05							,  num_attention_heads=4							,  num_hidden_layers=5							,  pad_token_id=1							,  vocab_size=1_0_0_0							,  )
      a_       		=CLIPTextModel(lowerCamelCase_)
      a_       		=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

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




 def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_=0) ->      Dict:



      """simple docstring"""







      if str(lowerCamelCase_).startswith("mps"):
           a_       		=torch.manual_seed(lowerCamelCase_)
      else:
           a_       		=torch.Generator(device=lowerCamelCase_).manual_seed(lowerCamelCase_)
      a_       		={
          "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       lowercase_  (   self) ->      int:



      """simple docstring"""







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

      a_       		=self.get_dummy_components()
      a_       		=StableDiffusionLDMaDPipeline(**lowerCamelCase_)
      a_       		=ldmad_pipe.to(lowerCamelCase_)
      ldmad_pipe.set_progress_bar_config(disable=lowerCamelCase_)

      a_       		=self.get_dummy_inputs(lowerCamelCase_)
      a_       		=ldmad_pipe(**lowerCamelCase_)
      a_		,						a_       		=output.rgb, output.depth

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

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

      a_       		=np.array(
          [0.3_7_3_3_8_1_7_6, 0.7_0_2_4_7, 0.7_4_2_0_3_1_9_3, 0.5_1_6_4_3_6_0_4, 0.5_8_2_5_6_7_9_3, 0.6_0_9_3_2_1_3_6, 0.4_1_8_1_0_9_5, 0.4_8_3_5_5_8_7_7, 0.4_6_5_3_5_2_6_2])
      a_       		=np.array([1_0_3.4_6_7_2_7, 8_5.8_1_2_0_0_4, 8_7.8_4_9_2_3_6])

      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       lowercase_  (   self) ->      Union[str, Any]:



      """simple docstring"""







      a_       		=self.get_dummy_components()
      a_       		=StableDiffusionLDMaDPipeline(**lowerCamelCase_)
      a_       		=ldmad_pipe.to(lowerCamelCase_)
      ldmad_pipe.set_progress_bar_config(disable=lowerCamelCase_)

      a_       		=self.get_dummy_inputs(lowerCamelCase_)
      a_       		=3 * [inputs["prompt"]]

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

      a_       		=self.get_dummy_inputs(lowerCamelCase_)
      a_       		=3 * [inputs.pop("prompt")]

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

      a_       		=ldmad_pipe.text_encoder(lowerCamelCase_)[0]

      a_       		=prompt_embeds

      # forward
      a_       		=ldmad_pipe(**lowerCamelCase_)
      a_		,						a_       		=output.rgb, output.depth
      a_       		=rgb_slice_a[0, -3:, -3:, -1]
      a_       		=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       lowercase_  (   self) ->      Optional[Any]:



      """simple docstring"""







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

      a_       		=self.get_dummy_inputs(lowerCamelCase_)
      a_       		="french fries"
      a_       		=ldmad_pipe(**lowerCamelCase_							,  negative_prompt=lowerCamelCase_)

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

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

      a_       		=np.array(
          [0.3_7_0_4_4, 0.7_1_8_1_1_5_0_3, 0.7_2_2_3_2_5_1, 0.4_8_6_0_3_6_7_5, 0.5_6_3_8_3_9_1, 0.6_3_6_4_9_4_8, 0.4_2_8_3_3_7_0_4, 0.4_9_0_1_3_1_5, 0.4_7_9_2_6_2_1_7])
      a_       		=np.array([1_0_7.8_4_7_3_8, 8_4.6_2_8_0_2, 8_9.9_6_2_1_3_5])
      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       lowercase_  (   self) ->      Union[str, Any]:



      """simple docstring"""







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




 def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_="cpu"							,  lowerCAmelCase_=torch.floataa							,  lowerCAmelCase_=0) ->      int:



      """simple docstring"""







      a_       		=torch.Generator(device=lowerCamelCase_).manual_seed(lowerCamelCase_)
      a_       		=np.random.RandomState(lowerCamelCase_).standard_normal((1, 4, 6_4, 6_4))
      a_       		=torch.from_numpy(lowerCamelCase_).to(device=lowerCamelCase_							,  dtype=lowerCamelCase_)
      a_       		={
          "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       lowercase_  (   self) ->      List[Any]:



      """simple docstring"""







      a_       		=StableDiffusionLDMaDPipeline.from_pretrained("Intel/ldm3d")
      a_       		=ldmad_pipe.to(lowerCamelCase_)
      ldmad_pipe.set_progress_bar_config(disable=lowerCamelCase_)

      a_       		=self.get_inputs(lowerCamelCase_)
      a_       		=ldmad_pipe(**lowerCamelCase_)
      a_		,						a_       		=output.rgb, output.depth
      a_       		=rgb[0, -3:, -3:, -1].flatten()
      a_       		=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)

      a_       		=np.array(
          [0.5_3_8_0_5_4_6_5, 0.5_6_7_0_7_3_0_5, 0.5_4_8_6_5_1_5, 0.5_7_0_1_2_2_3_6, 0.5_8_1_4_5_1_1, 0.5_6_2_5_3_4_8_7, 0.5_4_8_4_3_0_1_4, 0.5_5_0_9_2_2_6_3, 0.6_4_5_9_7_0_6])
      a_       		=np.array(
          [0.9_2_6_3_7_8_1, 0.6_6_7_8_6_7_2, 0.5_4_8_6_5_1_5, 0.9_2_2_0_2_1_4_5, 0.6_7_8_3_1_1_3_5, 0.5_6_2_5_3_4_8_7, 0.9_2_4_1_6_9_4, 0.7_5_5_1_4_7_8, 0.6_4_5_9_7_0_6])
      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       lowercase_  (   self) ->      int:



      """simple docstring"""







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




 def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_="cpu"							,  lowerCAmelCase_=torch.floataa							,  lowerCAmelCase_=0) ->      Dict:



      """simple docstring"""







      a_       		=torch.Generator(device=lowerCamelCase_).manual_seed(lowerCamelCase_)
      a_       		=np.random.RandomState(lowerCamelCase_).standard_normal((1, 4, 6_4, 6_4))
      a_       		=torch.from_numpy(lowerCamelCase_).to(device=lowerCamelCase_							,  dtype=lowerCamelCase_)
      a_       		={
          "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       lowercase_  (   self) ->      Dict:



      """simple docstring"""







      a_       		=StableDiffusionLDMaDPipeline.from_pretrained("Intel/ldm3d").to(lowerCamelCase_)
      ldmad_pipe.set_progress_bar_config(disable=lowerCamelCase_)

      a_       		=self.get_inputs(lowerCamelCase_)
      a_       		=ldmad_pipe(**lowerCamelCase_)
      a_		,						a_       		=output.rgb, output.depth

      a_       		=0.4_9_5_5_8_6
      a_       		=0.3_3_7_9_5_5_1_5
      a_       		=1_1_2.4_8_5_1_8
      a_       		=9_8.4_8_9_7_4_6
      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       lowercase_  (   self) ->      Optional[int]:



      """simple docstring"""







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

      a_       		=self.get_inputs(lowerCamelCase_)
      a_       		=ldmad_pipe(**lowerCamelCase_)
      a_		,						a_       		=output.rgb, output.depth

      a_       		=0.4_1_9_4_1_2_7
      a_       		=0.3_5_3_7_5_5_8_6
      a_       		=0.5_6_3_8_5_0_2
      a_       		=0.3_4_6_8_6_1_0_3

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


from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase							       =			logging.get_logger(__name__)

lowercase							       =			{
    '''google/switch-base-8''': '''https://huggingface.co/google/switch-base-8/blob/main/config.json''',
}







class 	UpperCAmelCase	(		__a):


 '''simple docstring'''




 __magic_name__					: int  =      "switch_transformers"
 __magic_name__					: List[Any]  =      ["past_key_values"]
 __magic_name__					: Union[str, Any]  =      {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}




 def __init__(   self							,  lowerCAmelCase_=3_2_1_2_8							,  lowerCAmelCase_=7_6_8							,  lowerCAmelCase_=6_4							,  lowerCAmelCase_=2_0_4_8							,  lowerCAmelCase_=6_4							,  lowerCAmelCase_=1_2							,  lowerCAmelCase_=3							,  lowerCAmelCase_=1_2							,  lowerCAmelCase_=3							,  lowerCAmelCase_=1_2							,  lowerCAmelCase_=8							,  lowerCAmelCase_=False							,  lowerCAmelCase_=0.0_1							,  lowerCAmelCase_="float32"							,  lowerCAmelCase_=False							,  lowerCAmelCase_=3_2							,  lowerCAmelCase_=1_2_8							,  lowerCAmelCase_=0.1							,  lowerCAmelCase_=1e-6							,  lowerCAmelCase_=0.0_0_1							,  lowerCAmelCase_=0.0_0_1							,  lowerCAmelCase_=1.0							,  lowerCAmelCase_="relu"							,  lowerCAmelCase_=True							,  lowerCAmelCase_=False							,  lowerCAmelCase_=True							,  lowerCAmelCase_=0							,  lowerCAmelCase_=1							,  **lowerCAmelCase_							,  ) ->      Optional[int]:



      """simple docstring"""







      a_       		=vocab_size
      a_       		=d_model
      a_       		=d_kv
      a_       		=d_ff

      a_       		=num_sparse_encoder_layers

      a_       		=num_layers
      a_       		=(
          num_decoder_layers if num_decoder_layers is not None else self.num_layers
      )  # default = symmetry
      a_       		=num_sparse_decoder_layers

      # This tells us, each how many encoder layer we'll have to set a sparse layer.
      if self.num_sparse_encoder_layers > 0:
           a_       		=self.num_layers // self.num_sparse_encoder_layers
      else:
           a_       		=self.num_layers  # HACK: this will create 0 sparse layers

      # This tells us, each how many encoder layer we'll have to set a sparse layer.
      if self.num_sparse_decoder_layers > 0:
           a_       		=self.num_decoder_layers // self.num_sparse_decoder_layers
      else:
           a_       		=self.num_decoder_layers  # HACK: this will create 0 sparse layers

      a_       		=num_heads
      a_       		=num_experts
      a_       		=expert_capacity
      a_       		=router_bias
      a_       		=router_jitter_noise
      if router_dtype not in ["float32", "float16", "bfloat16"]:
           raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}""")
      a_       		=router_dtype

      a_       		=router_ignore_padding_tokens
      a_       		=relative_attention_num_buckets
      a_       		=relative_attention_max_distance

      a_       		=dropout_rate
      a_       		=layer_norm_epsilon
      a_       		=initializer_factor
      a_       		=feed_forward_proj
      a_       		=use_cache
      a_       		=add_router_probs

      a_       		=router_z_loss_coef
      a_       		=router_aux_loss_coef

      a_       		=self.feed_forward_proj.split("-")
      a_       		=act_info[-1]
      a_       		=act_info[0] == "gated"

      if len(lowerCAmelCase_) > 1 and act_info[0] != "gated" or len(lowerCAmelCase_) > 2:
           raise ValueError(
               f"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
               "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
               "'gated-gelu' or 'relu'")

      # for backwards compatibility
      if feed_forward_proj == "gated-gelu":
           a_       		="gelu_new"

      super().__init__(
          pad_token_id=lowerCAmelCase_							,  eos_token_id=lowerCAmelCase_							,  is_encoder_decoder=lowerCAmelCase_							,  **lowerCAmelCase_							,  )


