"""simple docstring"""


import importlib
import json
import os
from collections import OrderedDict
from typing import Dict, Optional, Union

# Build the list of all feature extractors
from ...configuration_utils import PretrainedConfig
from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
from ...feature_extraction_utils import FeatureExtractionMixin
from ...utils import CONFIG_NAME, FEATURE_EXTRACTOR_NAME, get_file_from_repo, logging
from .auto_factory import _LazyAutoMapping
from .configuration_auto import (
    CONFIG_MAPPING_NAMES,
    AutoConfig,
    model_type_to_module_name,
    replace_list_option_in_docstrings,
)


UpperCAmelCase__  		=						logging.get_logger(__name__)

UpperCAmelCase__  		=						OrderedDict(
    [
        ("""audio-spectrogram-transformer""", """ASTFeatureExtractor"""),
        ("""beit""", """BeitFeatureExtractor"""),
        ("""chinese_clip""", """ChineseCLIPFeatureExtractor"""),
        ("""clap""", """ClapFeatureExtractor"""),
        ("""clip""", """CLIPFeatureExtractor"""),
        ("""clipseg""", """ViTFeatureExtractor"""),
        ("""conditional_detr""", """ConditionalDetrFeatureExtractor"""),
        ("""convnext""", """ConvNextFeatureExtractor"""),
        ("""cvt""", """ConvNextFeatureExtractor"""),
        ("""data2vec-audio""", """Wav2Vec2FeatureExtractor"""),
        ("""data2vec-vision""", """BeitFeatureExtractor"""),
        ("""deformable_detr""", """DeformableDetrFeatureExtractor"""),
        ("""deit""", """DeiTFeatureExtractor"""),
        ("""detr""", """DetrFeatureExtractor"""),
        ("""dinat""", """ViTFeatureExtractor"""),
        ("""donut-swin""", """DonutFeatureExtractor"""),
        ("""dpt""", """DPTFeatureExtractor"""),
        ("""encodec""", """EncodecFeatureExtractor"""),
        ("""flava""", """FlavaFeatureExtractor"""),
        ("""glpn""", """GLPNFeatureExtractor"""),
        ("""groupvit""", """CLIPFeatureExtractor"""),
        ("""hubert""", """Wav2Vec2FeatureExtractor"""),
        ("""imagegpt""", """ImageGPTFeatureExtractor"""),
        ("""layoutlmv2""", """LayoutLMv2FeatureExtractor"""),
        ("""layoutlmv3""", """LayoutLMv3FeatureExtractor"""),
        ("""levit""", """LevitFeatureExtractor"""),
        ("""maskformer""", """MaskFormerFeatureExtractor"""),
        ("""mctct""", """MCTCTFeatureExtractor"""),
        ("""mobilenet_v1""", """MobileNetV1FeatureExtractor"""),
        ("""mobilenet_v2""", """MobileNetV2FeatureExtractor"""),
        ("""mobilevit""", """MobileViTFeatureExtractor"""),
        ("""nat""", """ViTFeatureExtractor"""),
        ("""owlvit""", """OwlViTFeatureExtractor"""),
        ("""perceiver""", """PerceiverFeatureExtractor"""),
        ("""poolformer""", """PoolFormerFeatureExtractor"""),
        ("""regnet""", """ConvNextFeatureExtractor"""),
        ("""resnet""", """ConvNextFeatureExtractor"""),
        ("""segformer""", """SegformerFeatureExtractor"""),
        ("""sew""", """Wav2Vec2FeatureExtractor"""),
        ("""sew-d""", """Wav2Vec2FeatureExtractor"""),
        ("""speech_to_text""", """Speech2TextFeatureExtractor"""),
        ("""speecht5""", """SpeechT5FeatureExtractor"""),
        ("""swiftformer""", """ViTFeatureExtractor"""),
        ("""swin""", """ViTFeatureExtractor"""),
        ("""swinv2""", """ViTFeatureExtractor"""),
        ("""table-transformer""", """DetrFeatureExtractor"""),
        ("""timesformer""", """VideoMAEFeatureExtractor"""),
        ("""tvlt""", """TvltFeatureExtractor"""),
        ("""unispeech""", """Wav2Vec2FeatureExtractor"""),
        ("""unispeech-sat""", """Wav2Vec2FeatureExtractor"""),
        ("""van""", """ConvNextFeatureExtractor"""),
        ("""videomae""", """VideoMAEFeatureExtractor"""),
        ("""vilt""", """ViltFeatureExtractor"""),
        ("""vit""", """ViTFeatureExtractor"""),
        ("""vit_mae""", """ViTFeatureExtractor"""),
        ("""vit_msn""", """ViTFeatureExtractor"""),
        ("""wav2vec2""", """Wav2Vec2FeatureExtractor"""),
        ("""wav2vec2-conformer""", """Wav2Vec2FeatureExtractor"""),
        ("""wavlm""", """Wav2Vec2FeatureExtractor"""),
        ("""whisper""", """WhisperFeatureExtractor"""),
        ("""xclip""", """CLIPFeatureExtractor"""),
        ("""yolos""", """YolosFeatureExtractor"""),
    ]
)

UpperCAmelCase__  		=						_LazyAutoMapping(CONFIG_MAPPING_NAMES, FEATURE_EXTRACTOR_MAPPING_NAMES)






def     __UpperCAmelCase					(			lowercase		):







  """simple docstring"""
  for module_name, extractors in FEATURE_EXTRACTOR_MAPPING_NAMES.items():
    if class_name in extractors:
      _UpperCAmelCase						 =      model_type_to_module_name(lowercase		)

      _UpperCAmelCase						 =      importlib.import_module(f'''.{module_name}'''   ,"""transformers.models"""		)
      try:
        return getattr(lowercase   ,lowercase		)
      except AttributeError:
        continue

  for _, extractor in FEATURE_EXTRACTOR_MAPPING._extra_content.items():
    if getattr(lowercase   ,"""__name__"""   ,lowercase		) == class_name:
      return extractor

    # We did not fine the class, but maybe it's because a dep is missing. In that case, the class will be in the main
    # init and we return the proper dummy to get an appropriate error message.
  _UpperCAmelCase						 =      importlib.import_module("""transformers"""		)
  if hasattr(lowercase   ,lowercase		):
    return getattr(lowercase   ,lowercase		)

  return None






def     __UpperCAmelCase					(			lowercase   ,lowercase = None   ,lowercase = False   ,lowercase = False   ,lowercase = None   ,lowercase = None   ,lowercase = None   ,lowercase = False   ,**lowercase   ,):







  """simple docstring"""
  _UpperCAmelCase						 =      get_file_from_repo(
      lowercase   ,lowercase   ,cache_dir=lowercase   ,force_download=lowercase   ,resume_download=lowercase   ,proxies=lowercase   ,use_auth_token=lowercase   ,revision=lowercase   ,local_files_only=lowercase   ,)
  if resolved_config_file is None:
    logger.info(
        """Could not locate the feature extractor configuration file, will try to use the model config instead."""		)
    return {}

  with open(lowercase   ,encoding="""utf-8"""		) as reader:
    return json.load(lowercase		)


class        a				:


    def __init__(					self	:  str					):
      raise EnvironmentError(
          """AutoFeatureExtractor is designed to be instantiated """
          """using the `AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path)` method."""					)


    @classmethod
    @replace_list_option_in_docstrings(__lowerCAmelCase					)
    def      lowerCAmelCase_  (					cls	:  Optional[Any]   ,				__lowerCAmelCase	:  Optional[Any]   ,				**__lowerCAmelCase	:  Any					):
      _UpperCAmelCase						 =      kwargs.pop("""config"""   ,				__lowerCAmelCase					)
      _UpperCAmelCase						 =      kwargs.pop("""trust_remote_code"""   ,				__lowerCAmelCase					)
      _UpperCAmelCase						 =      True

      _UpperCAmelCase ,				_UpperCAmelCase						 =      FeatureExtractionMixin.get_feature_extractor_dict(__lowerCAmelCase   ,				**__lowerCAmelCase					)
      _UpperCAmelCase						 =      config_dict.get("""feature_extractor_type"""   ,				__lowerCAmelCase					)
      _UpperCAmelCase						 =      None
      if "AutoFeatureExtractor" in config_dict.get("""auto_map"""   ,				{}					):
        _UpperCAmelCase						 =      config_dict["""auto_map"""]["""AutoFeatureExtractor"""]

      # If we don't find the feature extractor class in the feature extractor config, let's try the model config.
      if feature_extractor_class is None and feature_extractor_auto_map is None:
        if not isinstance(__lowerCAmelCase   ,				__lowerCAmelCase					):
          _UpperCAmelCase						 =      AutoConfig.from_pretrained(__lowerCAmelCase   ,				**__lowerCAmelCase					)
        # It could be in `config.feature_extractor_type``
        _UpperCAmelCase						 =      getattr(__lowerCAmelCase   ,				"""feature_extractor_type"""   ,				__lowerCAmelCase					)
        if hasattr(__lowerCAmelCase   ,				"""auto_map"""					) and "AutoFeatureExtractor" in config.auto_map:
          _UpperCAmelCase						 =      config.auto_map["""AutoFeatureExtractor"""]

      if feature_extractor_class is not None:
        _UpperCAmelCase						 =      feature_extractor_class_from_name(__lowerCAmelCase					)

      _UpperCAmelCase						 =      feature_extractor_auto_map is not None
      _UpperCAmelCase						 =      feature_extractor_class is not None or type(__lowerCAmelCase					) in FEATURE_EXTRACTOR_MAPPING
      _UpperCAmelCase						 =      resolve_trust_remote_code(
          __lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase					)

      if has_remote_code and trust_remote_code:
        _UpperCAmelCase						 =      get_class_from_dynamic_module(
            __lowerCAmelCase   ,				__lowerCAmelCase   ,				**__lowerCAmelCase					)
        _UpperCAmelCase						 =      kwargs.pop("""code_revision"""   ,				__lowerCAmelCase					)
        if os.path.isdir(__lowerCAmelCase					):
          feature_extractor_class.register_for_auto_class()
        return feature_extractor_class.from_dict(__lowerCAmelCase   ,				**__lowerCAmelCase					)
      elif feature_extractor_class is not None:
        return feature_extractor_class.from_dict(__lowerCAmelCase   ,				**__lowerCAmelCase					)
      # Last try: we use the FEATURE_EXTRACTOR_MAPPING.
      elif type(__lowerCAmelCase					) in FEATURE_EXTRACTOR_MAPPING:
        _UpperCAmelCase						 =      FEATURE_EXTRACTOR_MAPPING[type(__lowerCAmelCase					)]
        return feature_extractor_class.from_dict(__lowerCAmelCase   ,				**__lowerCAmelCase					)

      raise ValueError(
          f'''Unrecognized feature extractor in {pretrained_model_name_or_path}. Should have a '''
          f'''`feature_extractor_type` key in its {FEATURE_EXTRACTOR_NAME} of {CONFIG_NAME}, or one of the following '''
          f'''`model_type` keys in its {CONFIG_NAME}: {", ".join(c for c in FEATURE_EXTRACTOR_MAPPING_NAMES.keys()					)}'''					)


    @staticmethod
    def      lowerCAmelCase_  (					__lowerCAmelCase	:  Dict   ,				__lowerCAmelCase	:  Optional[Any]					):
      FEATURE_EXTRACTOR_MAPPING.register(__lowerCAmelCase   ,				__lowerCAmelCase					)

"""simple docstring"""


import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class        a				(		lowerCAmelCase_	,							lowerCAmelCase_	,							lowerCAmelCase_	,							unittest.TestCase   ):
    _snake_case     :  Optional[int]     							=							StableUnCLIPPipeline
    _snake_case     :  List[str]     							=							TEXT_TO_IMAGE_PARAMS
    _snake_case     :  Tuple     							=							TEXT_TO_IMAGE_BATCH_PARAMS
    _snake_case     :  List[Any]     							=							TEXT_TO_IMAGE_IMAGE_PARAMS
    _snake_case     :  str     							=							TEXT_TO_IMAGE_IMAGE_PARAMS

    # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
    _snake_case     :  str     							=							False


    def      lowerCAmelCase_  (					self	:  Dict					):
      _UpperCAmelCase						 =      32
      _UpperCAmelCase						 =      embedder_hidden_size

      # prior components

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=__lowerCAmelCase   ,				projection_dim=__lowerCAmelCase   ,				intermediate_size=37   ,				layer_norm_eps=1e-0_5   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1000   ,				)					)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      PriorTransformer(
          num_attention_heads=2   ,				attention_head_dim=12   ,				embedding_dim=__lowerCAmelCase   ,				num_layers=1   ,				)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      DDPMScheduler(
          variance_type="""fixed_small_log"""   ,				prediction_type="""sample"""   ,				num_train_timesteps=1000   ,				clip_sample=__lowerCAmelCase   ,				clip_sample_range=5.0   ,				beta_schedule="""squaredcos_cap_v2"""   ,				)

      # regular denoising components

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      StableUnCLIPImageNormalizer(embedding_dim=__lowerCAmelCase					)
      _UpperCAmelCase						 =      DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""					)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=__lowerCAmelCase   ,				projection_dim=32   ,				intermediate_size=37   ,				layer_norm_eps=1e-0_5   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1000   ,				)					)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      UNetaDConditionModel(
          sample_size=32   ,				in_channels=4   ,				out_channels=4   ,				down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")   ,				up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")   ,				block_out_channels=(32, 64)   ,				attention_head_dim=(2, 4)   ,				class_embed_type="""projection"""   ,				projection_class_embeddings_input_dim=embedder_projection_dim * 2   ,				cross_attention_dim=__lowerCAmelCase   ,				layers_per_block=1   ,				upcast_attention=__lowerCAmelCase   ,				use_linear_projection=__lowerCAmelCase   ,				)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      DDIMScheduler(
          beta_schedule="""scaled_linear"""   ,				beta_start=0.00_085   ,				beta_end=0.012   ,				prediction_type="""v_prediction"""   ,				set_alpha_to_one=__lowerCAmelCase   ,				steps_offset=1   ,				)

      torch.manual_seed(0					)
      _UpperCAmelCase						 =      AutoencoderKL()

      _UpperCAmelCase						 =      {
          # prior components
          """prior_tokenizer""": prior_tokenizer,
          """prior_text_encoder""": prior_text_encoder,
          """prior""": prior,
          """prior_scheduler""": prior_scheduler,
          # image noising components
          """image_normalizer""": image_normalizer,
          """image_noising_scheduler""": image_noising_scheduler,
          # regular denoising components
          """tokenizer""": tokenizer,
          """text_encoder""": text_encoder,
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
      }

      return components


    def      lowerCAmelCase_  (					self	:  List[Any]   ,				__lowerCAmelCase	:  Any   ,				__lowerCAmelCase	:  int=0					):
      if str(__lowerCAmelCase					).startswith("""mps"""					):
        _UpperCAmelCase						 =      torch.manual_seed(__lowerCAmelCase					)
      else:
        _UpperCAmelCase						 =      torch.Generator(device=__lowerCAmelCase					).manual_seed(__lowerCAmelCase					)
      _UpperCAmelCase						 =      {
          """prompt""": """A painting of a squirrel eating a burger""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """prior_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs


    def      lowerCAmelCase_  (					self	:  List[str]					):
      _UpperCAmelCase						 =      torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  Any					):
      _UpperCAmelCase						 =      torch_device in ["""cpu""", """mps"""]

      self._test_inference_batch_single_identical(test_max_difference=__lowerCAmelCase					)



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


    def      lowerCAmelCase_  (					self	:  List[str]					):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()


    def      lowerCAmelCase_  (					self	:  List[str]					):
      _UpperCAmelCase						 =      load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""					)

      _UpperCAmelCase						 =      StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""   ,				torch_dtype=torch.floataa					)
      pipe.to(__lowerCAmelCase					)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase					)
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      _UpperCAmelCase						 =      torch.Generator(device="""cpu"""					).manual_seed(0					)
      _UpperCAmelCase						 =      pipe("""anime turle"""   ,				generator=__lowerCAmelCase   ,				output_type="""np"""					)

      _UpperCAmelCase						 =      output.images[0]

      assert image.shape == (768, 768, 3)

      assert_mean_pixel_difference(__lowerCAmelCase   ,				__lowerCAmelCase					)


    def      lowerCAmelCase_  (					self	:  List[str]					):
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      _UpperCAmelCase						 =      StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""   ,				torch_dtype=torch.floataa					)
      _UpperCAmelCase						 =      pipe.to(__lowerCAmelCase					)
      pipe.set_progress_bar_config(disable=__lowerCAmelCase					)
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      _UpperCAmelCase						 =      pipe(
          """anime turtle"""   ,				prior_num_inference_steps=2   ,				num_inference_steps=2   ,				output_type="""np"""   ,				)

      _UpperCAmelCase						 =      torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9
