



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference



class 						UpperCAmelCase_					( _lowerCAmelCase     ,      unittest.TestCase):
       lowerCamelCase__         =	ShapEPipeline
       lowerCamelCase__         =	['prompt']
       lowerCamelCase__         =	['prompt']
       lowerCamelCase__         =	[
           'num_images_per_prompt',
           'num_inference_steps',
           'generator',
           'latents',
           'guidance_scale',
           'frame_size',
           'output_type',
           'return_dict',
       ]
       lowerCamelCase__         =	False







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              return 32







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              return 32







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              return self.time_input_dim * 4







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              return 8







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              _lowerCAmelCase  :							Tuple				     =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
              return tokenizer







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              torch.manual_seed(0)
              _lowerCAmelCase  :							List[Any]				     =					CLIPTextConfig(
                  bos_token_id=0,    eos_token_id=2,    hidden_size=self.text_embedder_hidden_size,    projection_dim=self.text_embedder_hidden_size,    intermediate_size=37,    layer_norm_eps=1E-05,    num_attention_heads=4,    num_hidden_layers=5,    pad_token_id=1,    vocab_size=1000,    )
              return CLIPTextModelWithProjection(_lowerCAmelCase)







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              torch.manual_seed(0)

              _lowerCAmelCase  :							Optional[int]				     =					{
                  "num_attention_heads": 2,
                  "attention_head_dim": 16,
                  "embedding_dim": self.time_input_dim,
                  "num_embeddings": 32,
                  "embedding_proj_dim": self.text_embedder_hidden_size,
                  "time_embed_dim": self.time_embed_dim,
                  "num_layers": 1,
                  "clip_embed_dim": self.time_input_dim * 2,
                  "additional_embeddings": 0,
                  "time_embed_act_fn": "gelu",
                  "norm_in_type": "layer",
                  "encoder_hid_proj_type": None,
                  "added_emb_type": None,
              }

              _lowerCAmelCase  :							List[str]				     =					PriorTransformer(**_lowerCAmelCase)
              return model







       @property
       def      snake_case__   (     self):

              '''simple docstring'''

              torch.manual_seed(0)

              _lowerCAmelCase  :							Tuple				     =					{
                  "param_shapes": (
                      (self.renderer_dim, 93),
                      (self.renderer_dim, 8),
                      (self.renderer_dim, 8),
                      (self.renderer_dim, 8),
                  ),
                  "d_latent": self.time_input_dim,
                  "d_hidden": self.renderer_dim,
                  "n_output": 12,
                  "background": (
                      0.1,
                      0.1,
                      0.1,
                  ),
              }
              _lowerCAmelCase  :							Dict				     =					ShapERenderer(**_lowerCAmelCase)
              return model







       def      snake_case__   (     self):

              '''simple docstring'''

              _lowerCAmelCase  :							Union[str, Any]				     =					self.dummy_prior
              _lowerCAmelCase  :							Any				     =					self.dummy_text_encoder
              _lowerCAmelCase  :							List[Any]				     =					self.dummy_tokenizer
              _lowerCAmelCase  :							Union[str, Any]				     =					self.dummy_renderer

              _lowerCAmelCase  :							str				     =					HeunDiscreteScheduler(
                  beta_schedule="exp",    num_train_timesteps=1024,    prediction_type="sample",    use_karras_sigmas=_lowerCAmelCase,    clip_sample=_lowerCAmelCase,    clip_sample_range=1.0,    )
              _lowerCAmelCase  :							Dict				     =					{
                  "prior": prior,
                  "text_encoder": text_encoder,
                  "tokenizer": tokenizer,
                  "renderer": renderer,
                  "scheduler": scheduler,
              }

              return components







       def      snake_case__   (     self,    __a,    __a=0):

              '''simple docstring'''

              if str(_lowerCAmelCase).startswith("mps"):
                     _lowerCAmelCase  :							str				     =					torch.manual_seed(_lowerCAmelCase)
              else:
                     _lowerCAmelCase  :							List[Any]				     =					torch.Generator(device=_lowerCAmelCase).manual_seed(_lowerCAmelCase)
              _lowerCAmelCase  :							str				     =					{
                  "prompt": "horse",
                  "generator": generator,
                  "num_inference_steps": 1,
                  "frame_size": 32,
                  "output_type": "np",
              }
              return inputs







       def      snake_case__   (     self):

              '''simple docstring'''

              _lowerCAmelCase  :							List[Any]				     =					"cpu"

              _lowerCAmelCase  :							Tuple				     =					self.get_dummy_components()

              _lowerCAmelCase  :							Dict				     =					self.pipeline_class(**_lowerCAmelCase)
              _lowerCAmelCase  :							List[Any]				     =					pipe.to(_lowerCAmelCase)

              pipe.set_progress_bar_config(disable=_lowerCAmelCase)

              _lowerCAmelCase  :							Dict				     =					pipe(**self.get_dummy_inputs(_lowerCAmelCase))
              _lowerCAmelCase  :							Optional[int]				     =					output.images[0]
              _lowerCAmelCase  :							str				     =					image[0, -3:, -3:, -1]

              assert image.shape == (20, 32, 32, 3)

              _lowerCAmelCase  :							Optional[Any]				     =					np.array(
                  [
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                      0.00_039_216,
                  ])

              assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







       def      snake_case__   (     self):

              '''simple docstring'''

              self._test_inference_batch_consistent(batch_sizes=[1, 2])







       def      snake_case__   (     self):

              '''simple docstring'''

              _lowerCAmelCase  :							Optional[Any]				     =					torch_device == "cpu"
              _lowerCAmelCase  :							Optional[Any]				     =					True

              self._test_inference_batch_single_identical(
                  batch_size=2,    test_max_difference=_lowerCAmelCase,    relax_max_difference=_lowerCAmelCase,    )







       def      snake_case__   (     self):

              '''simple docstring'''

              _lowerCAmelCase  :							List[str]				     =					self.get_dummy_components()
              _lowerCAmelCase  :							Dict				     =					self.pipeline_class(**_lowerCAmelCase)
              _lowerCAmelCase  :							Union[str, Any]				     =					pipe.to(_lowerCAmelCase)
              pipe.set_progress_bar_config(disable=_lowerCAmelCase)

              _lowerCAmelCase  :							Tuple				     =					1
              _lowerCAmelCase  :							List[Any]				     =					2

              _lowerCAmelCase  :							Dict				     =					self.get_dummy_inputs(_lowerCAmelCase)

              for key in inputs.keys():
                     if key in self.batch_params:
                            _lowerCAmelCase  :							Dict				     =					batch_size * [inputs[key]]

              _lowerCAmelCase  :							Dict				     =					pipe(**_lowerCAmelCase,    num_images_per_prompt=_lowerCAmelCase)[0]

              assert images.shape[0] == batch_size * num_images_per_prompt



@slow
@require_torch_gpu
class 						UpperCAmelCase_					( unittest.TestCase):







       def      snake_case__   (     self):

              '''simple docstring'''

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







       def      snake_case__   (     self):

              '''simple docstring'''

              _lowerCAmelCase  :							List[Any]				     =					load_numpy(
                  "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                  "/shap_e/test_shap_e_np_out.npy")
              _lowerCAmelCase  :							int				     =					ShapEPipeline.from_pretrained("openai/shap-e")
              _lowerCAmelCase  :							int				     =					pipe.to(_lowerCAmelCase)
              pipe.set_progress_bar_config(disable=_lowerCAmelCase)

              _lowerCAmelCase  :							Optional[Any]				     =					torch.Generator(device=_lowerCAmelCase).manual_seed(0)

              _lowerCAmelCase  :							Optional[int]				     =					pipe(
                  "a shark",    generator=_lowerCAmelCase,    guidance_scale=15.0,    num_inference_steps=64,    frame_size=64,    output_type="np",    ).images[0]

              assert images.shape == (20, 64, 64, 3)

              assert_mean_pixel_difference(_lowerCAmelCase,    _lowerCAmelCase)





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_snake_case	=						logging.get_logger(__name__)

_snake_case	=						{
    "microsoft/focalnet-tiny": "https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json",
}



class 						UpperCAmelCase_					( a     ,      a):
							lowerCamelCase__         =	'focalnet'







							def __init__(     self,    __a=224,    __a=4,    __a=3,    __a=96,    __a=False,    __a=[192, 384, 768, 768],    __a=[2, 2, 6, 2],    __a=[2, 2, 2, 2],    __a=[3, 3, 3, 3],    __a="gelu",    __a=4.0,    __a=0.0,    __a=0.1,    __a=False,    __a=1E-4,    __a=False,    __a=False,    __a=False,    __a=0.02,    __a=1E-5,    __a=32,    __a=None,    __a=None,    **__a,    ):

														'''simple docstring'''

														super().__init__(**__a)

														_lowerCAmelCase  :							str				     =					image_size
														_lowerCAmelCase  :							List[str]				     =					patch_size
														_lowerCAmelCase  :							List[Any]				     =					num_channels
														_lowerCAmelCase  :							Tuple				     =					embed_dim
														_lowerCAmelCase  :							List[Any]				     =					use_conv_embed
														_lowerCAmelCase  :							Any				     =					hidden_sizes
														_lowerCAmelCase  :							Tuple				     =					depths
														_lowerCAmelCase  :							Dict				     =					focal_levels
														_lowerCAmelCase  :							Optional[Any]				     =					focal_windows
														_lowerCAmelCase  :							str				     =					hidden_act
														_lowerCAmelCase  :							Union[str, Any]				     =					mlp_ratio
														_lowerCAmelCase  :							Any				     =					hidden_dropout_prob
														_lowerCAmelCase  :							Dict				     =					drop_path_rate
														_lowerCAmelCase  :							str				     =					use_layerscale
														_lowerCAmelCase  :							str				     =					layerscale_value
														_lowerCAmelCase  :							Union[str, Any]				     =					use_post_layernorm
														_lowerCAmelCase  :							Optional[int]				     =					use_post_layernorm_in_modulation
														_lowerCAmelCase  :							str				     =					normalize_modulator
														_lowerCAmelCase  :							Any				     =					initializer_range
														_lowerCAmelCase  :							Union[str, Any]				     =					layer_norm_eps
														_lowerCAmelCase  :							Any				     =					encoder_stride
														_lowerCAmelCase  :							List[str]				     =					["stem"] + [f"stage{idx}" for idx in range(1,    len(self.depths) + 1)]
														_lowerCAmelCase     ,					_lowerCAmelCase  :							List[str]				     =					get_aligned_output_features_output_indices(
														    out_features=__a,    out_indices=__a,    stage_names=self.stage_names)
