




import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device



def     lowerCAmelCase		(		_lowerCAmelCase				: nn.Module						):




 """simple docstring"""







 UpperCAmelCase__							=    {}
 for name in unet.attn_processors.keys():
  UpperCAmelCase__							=    None if name.endswith("attn1.processor"						) else unet.config.cross_attention_dim
  if name.startswith("mid_block"						):
   UpperCAmelCase__							=    unet.config.block_out_channels[-1]
  elif name.startswith("up_blocks"						):
   UpperCAmelCase__							=    int(name[len("up_blocks."						)]						)
   UpperCAmelCase__							=    list(reversed(unet.config.block_out_channels						)						)[block_id]
  elif name.startswith("down_blocks"						):
   UpperCAmelCase__							=    int(name[len("down_blocks."						)]						)
   UpperCAmelCase__							=    unet.config.block_out_channels[block_id]
  UpperCAmelCase__							=    (
      LoRAAttnProcessor2_0 if hasattr(_lowerCAmelCase    ,				"scaled_dot_product_attention"						) else LoRAAttnProcessor
  )
  UpperCAmelCase__							=    lora_attn_processor_class(
      hidden_size=_lowerCAmelCase    ,				cross_attention_dim=_lowerCAmelCase						)
 UpperCAmelCase__							=    AttnProcsLayers(_lowerCAmelCase						)
 return lora_attn_procs, unet_lora_layers



def     lowerCAmelCase		(		_lowerCAmelCase				: nn.Module						):




 """simple docstring"""







 UpperCAmelCase__							=    {}
 UpperCAmelCase__							=    (
     LoRAAttnProcessor2_0 if hasattr(_lowerCAmelCase    ,				"scaled_dot_product_attention"						) else LoRAAttnProcessor
 )
 for name, module in text_encoder.named_modules():
  if name.endswith(_lowerCAmelCase						):
   UpperCAmelCase__							=    lora_attn_processor_class(
       hidden_size=module.out_proj.out_features    ,				cross_attention_dim=_lowerCAmelCase						)
 return text_lora_attn_procs



def     lowerCAmelCase		(		_lowerCAmelCase				: nn.Module						):




 """simple docstring"""







 UpperCAmelCase__							=    create_text_encoder_lora_attn_procs(_lowerCAmelCase						)
 UpperCAmelCase__							=    AttnProcsLayers(_lowerCAmelCase						)
 return text_encoder_lora_layers



def     lowerCAmelCase		(		_lowerCAmelCase				: str    ,				_lowerCAmelCase				: int=False						):




 """simple docstring"""







 for _, attn_proc in text_lora_attn_procs.items():
  # set up.weights
  for layer_name, layer_module in attn_proc.named_modules():
   if layer_name.endswith("_lora"						):
    UpperCAmelCase__							=    (
        torch.randn_like(layer_module.up.weight						)
        if randn_weight
        else torch.zeros_like(layer_module.up.weight						)
    )
    UpperCAmelCase__							=    torch.nn.Parameter(_lowerCAmelCase						)


class       _UpperCamelCase   (					unittest.TestCase		):
      def 						UpperCAmelCase_    (      self :Dict   )  ->      Tuple:
       torch.manual_seed(0   )
       UpperCAmelCase__							=    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			,   )
       UpperCAmelCase__							=    DDIMScheduler(
           beta_start=0.0_00_85			,   beta_end=0.0_12			,   beta_schedule="scaled_linear"			,   clip_sample=lowerCamelCase			,   set_alpha_to_one=lowerCamelCase			,   steps_offset=1			,   )
       torch.manual_seed(0   )
       UpperCAmelCase__							=    AutoencoderKL(
           block_out_channels=[32, 64]			,   in_channels=3			,   out_channels=3			,   down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]			,   up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]			,   latent_channels=4			,   )
       UpperCAmelCase__							=    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=1000			,   )
       UpperCAmelCase__							=    CLIPTextModel(lowerCamelCase   )
       UpperCAmelCase__							=    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )

       UpperCAmelCase__			,			UpperCAmelCase__							=    create_unet_lora_layers(lowerCamelCase   )
       UpperCAmelCase__							=    create_text_encoder_lora_layers(lowerCamelCase   )

       UpperCAmelCase__							=    {
           "unet": unet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "safety_checker": None,
           "feature_extractor": None,
       }
       UpperCAmelCase__							=    {
           "unet_lora_layers": unet_lora_layers,
           "text_encoder_lora_layers": text_encoder_lora_layers,
           "unet_lora_attn_procs": unet_lora_attn_procs,
       }
       return pipeline_components, lora_components
      def 						UpperCAmelCase_    (      self :Any   )  ->      str:
       UpperCAmelCase__							=    1
       UpperCAmelCase__							=    10
       UpperCAmelCase__							=    4
       UpperCAmelCase__							=    (32, 32)

       UpperCAmelCase__							=    torch.manual_seed(0   )
       UpperCAmelCase__							=    floats_tensor((batch_size, num_channels) + sizes   )
       UpperCAmelCase__							=    torch.randint(1			,   lowerCamelCase			,   size=(batch_size, sequence_length)			,   generator=lowerCamelCase   )

       UpperCAmelCase__							=    {
           "prompt": "A painting of a squirrel eating a burger",
           "generator": generator,
           "num_inference_steps": 2,
           "guidance_scale": 6.0,
           "output_type": "numpy",
       }

       return noise, input_ids, pipeline_inputs

       # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      int:
       UpperCAmelCase__							=    77

       UpperCAmelCase__							=    torch.randint(2			,   56			,   size=(1, max_seq_length)			,   generator=torch.manual_seed(0   )   )

       UpperCAmelCase__							=    {}
       UpperCAmelCase__							=    inputs
       return prepared_inputs
      def 						UpperCAmelCase_    (      self :str			,   lowerCamelCase :Union[str, Any]   )  ->      Any:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       LoraLoaderMixin.save_lora_weights(
           save_directory=lowerCamelCase			,   unet_lora_layers=lora_components["unet_lora_layers"]			,   text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,   )
       self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase			,   "pytorch_lora_weights.bin"   )   )   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      Any:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
       UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_inputs()

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
        LoraLoaderMixin.save_lora_weights(
            save_directory=lowerCamelCase			,   unet_lora_layers=lora_components["unet_lora_layers"]			,   text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,   )
        self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase			,   "pytorch_lora_weights.bin"   )   )   )
        sd_pipe.load_lora_weights(lowerCamelCase   )

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase   )			,   torch.from_numpy(lowerCamelCase   )   )   )
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      str:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
       UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_inputs()

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
        LoraLoaderMixin.save_lora_weights(
            save_directory=lowerCamelCase			,   unet_lora_layers=lora_components["unet_lora_layers"]			,   text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,   safe_serialization=lowerCamelCase			,   )
        self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase			,   "pytorch_lora_weights.safetensors"   )   )   )
        sd_pipe.load_lora_weights(lowerCamelCase   )

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase   )			,   torch.from_numpy(lowerCamelCase   )   )   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      Any:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    lora_components["unet_lora_attn_procs"]
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
       UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_inputs()

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
        UpperCAmelCase__							=    sd_pipe.unet
        unet.set_attn_processor(lowerCamelCase   )
        unet.save_attn_procs(lowerCamelCase   )
        self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase			,   "pytorch_lora_weights.bin"   )   )   )
        sd_pipe.load_lora_weights(lowerCamelCase   )

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase   )			,   torch.from_numpy(lowerCamelCase   )   )   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      Any:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )

       UpperCAmelCase__							=    self.get_dummy_tokens()

       # inference without lora
       UpperCAmelCase__							=    pipe.text_encoder(**lowerCamelCase   )[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with zeroed out up.weights
       UpperCAmelCase__							=    create_text_encoder_lora_attn_procs(pipe.text_encoder   )
       set_lora_up_weights(lowerCamelCase			,   randn_weight=lowerCamelCase   )

       # monkey patch
       pipe._modify_text_encoder(lowerCamelCase   )

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       UpperCAmelCase__							=    pipe.text_encoder(**lowerCamelCase   )[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert torch.allclose(
           lowerCamelCase			,   lowerCamelCase   ), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

       # create lora_attn_procs with randn up.weights
       UpperCAmelCase__							=    create_text_encoder_lora_attn_procs(pipe.text_encoder   )
       set_lora_up_weights(lowerCamelCase			,   randn_weight=lowerCamelCase   )

       # monkey patch
       pipe._modify_text_encoder(lowerCamelCase   )

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       UpperCAmelCase__							=    pipe.text_encoder(**lowerCamelCase   )[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert not torch.allclose(
           lowerCamelCase			,   lowerCamelCase   ), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Tuple:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )

       UpperCAmelCase__							=    self.get_dummy_tokens()

       # inference without lora
       UpperCAmelCase__							=    pipe.text_encoder(**lowerCamelCase   )[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with randn up.weights
       UpperCAmelCase__							=    create_text_encoder_lora_attn_procs(pipe.text_encoder   )
       set_lora_up_weights(lowerCamelCase			,   randn_weight=lowerCamelCase   )

       # monkey patch
       pipe._modify_text_encoder(lowerCamelCase   )

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       UpperCAmelCase__							=    pipe.text_encoder(**lowerCamelCase   )[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert not torch.allclose(
           lowerCamelCase			,   lowerCamelCase   ), "lora outputs should be different to without lora outputs"

       # remove monkey patch
       pipe._remove_text_encoder_monkey_patch()

       # inference with removed lora
       UpperCAmelCase__							=    pipe.text_encoder(**lowerCamelCase   )[0]
       assert outputs_without_lora_removed.shape == (1, 77, 32)

       assert torch.allclose(
           lowerCamelCase			,   lowerCamelCase   ), "remove lora monkey patch should restore the original outputs"
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      List[Any]:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
       UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_inputs()

       with tempfile.TemporaryDirectory() as tmpdirname:
        LoraLoaderMixin.save_lora_weights(
            save_directory=lowerCamelCase			,   unet_lora_layers=lora_components["unet_lora_layers"]			,   text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,   )
        self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase			,   "pytorch_lora_weights.bin"   )   )   )
        sd_pipe.load_lora_weights(lowerCamelCase   )

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    lora_images[0, -3:, -3:, -1]

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase			,   cross_attention_kwargs={"scale": 0.5}   ).images
       UpperCAmelCase__							=    lora_images_with_scale[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(
           torch.allclose(torch.from_numpy(lowerCamelCase   )			,   torch.from_numpy(lowerCamelCase   )   )   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      List[str]:
       with tempfile.TemporaryDirectory() as tmpdirname:
        self.create_lora_weight_file(lowerCamelCase   )

        UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
        UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
        UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

        # check if vanilla attention processors are used
        for _, module in sd_pipe.unet.named_modules():
         if isinstance(lowerCamelCase			,   lowerCamelCase   ):
          self.assertIsInstance(module.processor			,   (AttnProcessor, AttnProcessor2_0)   )

            # load LoRA weight file
        sd_pipe.load_lora_weights(lowerCamelCase   )

        # check if lora attention processors are used
        for _, module in sd_pipe.unet.named_modules():
         if isinstance(lowerCamelCase			,   lowerCamelCase   ):
          UpperCAmelCase__							=    (
              LoRAAttnProcessor2_0 if hasattr(lowerCamelCase			,   "scaled_dot_product_attention"   ) else LoRAAttnProcessor
          )
          self.assertIsInstance(module.processor			,   lowerCamelCase   )
      @unittest.skipIf(torch_device != "cuda"			,   "This test is supposed to run on GPU"   )
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      str:
       with tempfile.TemporaryDirectory() as tmpdirname:
        self.create_lora_weight_file(lowerCamelCase   )

        UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
        UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
        UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

        # check if xFormers attention processors are used
        for _, module in sd_pipe.unet.named_modules():
         if isinstance(lowerCamelCase			,   lowerCamelCase   ):
          self.assertIsInstance(module.processor			,   lowerCamelCase   )

            # load LoRA weight file
        sd_pipe.load_lora_weights(lowerCamelCase   )

        # check if lora attention processors are used
        for _, module in sd_pipe.unet.named_modules():
         if isinstance(lowerCamelCase			,   lowerCamelCase   ):
          self.assertIsInstance(module.processor			,   lowerCamelCase   )







      @unittest.skipIf(torch_device != "cuda"			,   "This test is supposed to run on GPU"   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      List[Any]:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_components()
       UpperCAmelCase__							=    StableDiffusionPipeline(**lowerCamelCase   )
       UpperCAmelCase__							=    sd_pipe.to(lowerCamelCase   )
       sd_pipe.set_progress_bar_config(disable=lowerCamelCase   )

       UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    self.get_dummy_inputs()

       # enable XFormers
       sd_pipe.enable_xformers_memory_efficient_attention()

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
        LoraLoaderMixin.save_lora_weights(
            save_directory=lowerCamelCase			,   unet_lora_layers=lora_components["unet_lora_layers"]			,   text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,   )
        self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase			,   "pytorch_lora_weights.bin"   )   )   )
        sd_pipe.load_lora_weights(lowerCamelCase   )

       UpperCAmelCase__							=    sd_pipe(**lowerCamelCase   ).images
       UpperCAmelCase__							=    lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(lowerCamelCase   )			,   torch.from_numpy(lowerCamelCase   )   )   )











import argparse
import glob
import inspect
import os
import re
import unittest

from get_test_info import get_test_classes

from tests.test_pipeline_mixin import pipeline_test_mapping


_lowerCAmelCase    : int		= {}
for task, _ in pipeline_test_mapping.items():
  _lowerCAmelCase    : List[str]		= {"pt": None, "tf": None}


# DO **NOT** add item to this set (unless the reason is approved)
_lowerCAmelCase    : Optional[int]		= {
    "tests/models/esm/test_modeling_esmfold.py",  # The pipeline test mapping is added to `test_modeling_esm.py`
}



def     lowerCAmelCase		(		_lowerCAmelCase				: Union[str, Any]						):




 """simple docstring"""







 if "ModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
  return "pt"
 elif "TFModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
  return "tf"
 elif "FlaxModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
  return "flax"
 else:
  return None



def     lowerCAmelCase		(		_lowerCAmelCase				: str    ,				_lowerCAmelCase				: Optional[Any]						):




 """simple docstring"""







 if PIPELINE_TEST_MAPPING[task].get(_lowerCAmelCase    ,				_lowerCAmelCase						) is not None:
  return PIPELINE_TEST_MAPPING[task][framework]

 UpperCAmelCase__							=    pipeline_test_mapping[task]["test"]
 UpperCAmelCase__							=    None

 if framework == "pt":
  UpperCAmelCase__							=    getattr(_lowerCAmelCase    ,				"model_mapping"    ,				_lowerCAmelCase						)
 elif framework == "tf":
  UpperCAmelCase__							=    getattr(_lowerCAmelCase    ,				"tf_model_mapping"    ,				_lowerCAmelCase						)

 if mapping is not None:
  UpperCAmelCase__							=    dict(mapping.items()						)

 # cache the results
 UpperCAmelCase__							=    mapping
 return mapping



def     lowerCAmelCase		(		_lowerCAmelCase				: Tuple    ,				_lowerCAmelCase				: List[str]						):




 """simple docstring"""







 UpperCAmelCase__							=    get_framework(_lowerCAmelCase						)
 if framework is None:
  return None
 UpperCAmelCase__							=    get_mapping_for_task(_lowerCAmelCase    ,				_lowerCAmelCase						)
 if mapping is None:
  return None

 UpperCAmelCase__							=    list({model_class.config_class for model_class in test_class.all_model_classes}						)
 if len(_lowerCAmelCase						) != 1:
  raise ValueError("There should be exactly one configuration class from `test_class.all_model_classes`."						)

 # This could be a list/tuple of model classes, but it's rare.
 UpperCAmelCase__							=    mapping.get(config_classes[0]    ,				_lowerCAmelCase						)
 if isinstance(_lowerCAmelCase    ,				(tuple, list)						):
  UpperCAmelCase__							=    sorted(_lowerCAmelCase    ,				key=lambda _lowerCAmelCase						: x.__name__						)

 return model_class



def     lowerCAmelCase		(		_lowerCAmelCase				: str						):




 """simple docstring"""







 UpperCAmelCase__							=    [(task, get_model_for_pipeline_test(_lowerCAmelCase    ,				_lowerCAmelCase						)) for task in pipeline_test_mapping]
 UpperCAmelCase__							=    sorted([(task, model) for task, model in mapping if model is not None]    ,				key=lambda _lowerCAmelCase						: x[0]						)

 return dict(_lowerCAmelCase						)



def     lowerCAmelCase		(		_lowerCAmelCase				: Optional[Any]						):




 """simple docstring"""







 UpperCAmelCase__							=    get_framework(_lowerCAmelCase						)
 if framework == "pt":
  UpperCAmelCase__							=    "torch"
 UpperCAmelCase__							=    "{}"

 UpperCAmelCase__							=    get_pipeline_model_mapping(_lowerCAmelCase						)
 if len(_lowerCAmelCase						) == 0:
  return ""

 UpperCAmelCase__							=    []
 for task, model_classes in mapping.items():
  if isinstance(_lowerCAmelCase    ,				(tuple, list)						):
   # A list/tuple of model classes
   UpperCAmelCase__							=    "(" + ", ".join([x.__name__ for x in model_classes]						) + ")"
  else:
   # A single model class
   UpperCAmelCase__							=    model_classes.__name__
  texts.append(F'''"{task}": {value}'''						)
 UpperCAmelCase__							=    "{" + ", ".join(_lowerCAmelCase						) + "}"
 UpperCAmelCase__							=    F'''pipeline_model_mapping = {text} if is_{framework}_available() else {default_value}'''

 return text



def     lowerCAmelCase		(		_lowerCAmelCase				: Any						):




 """simple docstring"""







 UpperCAmelCase__							=    {"ModelTesterMixin", "TFModelTesterMixin", "FlaxModelTesterMixin"}
 if not issubclass(_lowerCAmelCase    ,				unittest.TestCase						):
  return False
 return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__]						)						) > 0



def     lowerCAmelCase		(		_lowerCAmelCase				: Union[str, Any]						):




 """simple docstring"""







 UpperCAmelCase__							=    [x for x in get_test_classes(_lowerCAmelCase						) if is_valid_test_class(_lowerCAmelCase						)]

 UpperCAmelCase__							=    None
 for test_class in test_classes:
  # If a test class has defined `pipeline_model_mapping`, let's take it
  if getattr(_lowerCAmelCase    ,				"pipeline_model_mapping"    ,				_lowerCAmelCase						) is not None:
   UpperCAmelCase__							=    test_class
   break
    # Take the test class with the shortest name (just a heuristic)
 if target_test_class is None and len(_lowerCAmelCase						) > 0:
  UpperCAmelCase__							=    sorted(_lowerCAmelCase    ,				key=lambda _lowerCAmelCase						: (len(x.__name__						), x.__name__)						)[0]

 return target_test_class



def     lowerCAmelCase		(		_lowerCAmelCase				: int    ,				_lowerCAmelCase				: Optional[int]    ,				_lowerCAmelCase				: str						):




 """simple docstring"""







 UpperCAmelCase__							=    start_idx
 for idx, line in enumerate(lines[start_idx:]						):
  UpperCAmelCase__							=    len(_lowerCAmelCase						) - len(line.lstrip()						)
  if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
   UpperCAmelCase__							=    start_idx + idx
  elif idx > 0 and indent <= indent_level:
   # Outside the definition block of `pipeline_model_mapping`
   break

 return end_idx



def     lowerCAmelCase		(		_lowerCAmelCase				: int    ,				_lowerCAmelCase				: str=False						):




 """simple docstring"""







 if getattr(_lowerCAmelCase    ,				"pipeline_model_mapping"    ,				_lowerCAmelCase						) is not None:
  if not overwrite:
   return "", -1

 UpperCAmelCase__							=    get_pipeline_model_mapping_string(_lowerCAmelCase						)
 if len(_lowerCAmelCase						) == 0:
  return "", -1
 UpperCAmelCase__							=    line_to_add + "\n"

 # The code defined the class `test_class`
 UpperCAmelCase__			,			UpperCAmelCase__							=    inspect.getsourcelines(_lowerCAmelCase						)
 # `inspect` gives the code for an object, including decorator(s) if any.
 # We (only) need the exact line of the class definition.
 for idx, line in enumerate(_lowerCAmelCase						):
  if line.lstrip().startswith("class "						):
   UpperCAmelCase__							=    class_lines[idx:]
   class_start_line_no += idx
   break
 UpperCAmelCase__							=    class_start_line_no + len(_lowerCAmelCase						) - 1

 # The index in `class_lines` that starts the definition of `all_model_classes`, `all_generative_model_classes` or
 # `pipeline_model_mapping`. This assumes they are defined in such order, and we take the start index of the last
 # block that appears in a `test_class`.
 UpperCAmelCase__							=    None
 # The indent level of the line at `class_lines[start_idx]` (if defined)
 UpperCAmelCase__							=    0
 # To record if `pipeline_model_mapping` is found in `test_class`.
 UpperCAmelCase__							=    None
 for idx, line in enumerate(_lowerCAmelCase						):
  if line.strip().startswith("all_model_classes = "						):
   UpperCAmelCase__							=    len(_lowerCAmelCase						) - len(line.lstrip()						)
   UpperCAmelCase__							=    idx
  elif line.strip().startswith("all_generative_model_classes = "						):
   UpperCAmelCase__							=    len(_lowerCAmelCase						) - len(line.lstrip()						)
   UpperCAmelCase__							=    idx
  elif line.strip().startswith("pipeline_model_mapping = "						):
   UpperCAmelCase__							=    len(_lowerCAmelCase						) - len(line.lstrip()						)
   UpperCAmelCase__							=    idx
   UpperCAmelCase__							=    line
   break

 if start_idx is None:
  return "", -1
 # Find the ending index (inclusive) of the above found block.
 UpperCAmelCase__							=    find_block_ending(_lowerCAmelCase    ,				_lowerCAmelCase    ,				_lowerCAmelCase						)

 # Extract `is_xxx_available()` from existing blocks: some models require specific libraries like `timm` and use
 # `is_timm_available()` instead of `is_torch_available()`.
 # Keep leading and trailing whitespaces
 UpperCAmelCase__							=    re.compile(R"\s(is_\S+?_available\(\))\s"						)
 for line in class_lines[start_idx : end_idx + 1]:
  UpperCAmelCase__							=    r.search(_lowerCAmelCase						)
  if backend_condition is not None:
   # replace the leading and trailing whitespaces to the space character " ".
   UpperCAmelCase__							=    " " + backend_condition[0][1:-1] + " "
   UpperCAmelCase__							=    r.sub(_lowerCAmelCase    ,				_lowerCAmelCase						)
   break

 if def_line is None:
  # `pipeline_model_mapping` is not defined. The target index is set to the ending index (inclusive) of
  # `all_model_classes` or `all_generative_model_classes`.
  UpperCAmelCase__							=    end_idx
 else:
  # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
  UpperCAmelCase__							=    start_idx - 1
  # mark the lines of the currently existing `pipeline_model_mapping` to be removed.
  for idx in range(_lowerCAmelCase    ,				end_idx + 1						):
   # These lines are going to be removed before writing to the test file.
   UpperCAmelCase__							=    None  # noqa

    # Make sure the test class is a subclass of `PipelineTesterMixin`.
 UpperCAmelCase__							=    [x.__name__ for x in test_class.__bases__]
 if "PipelineTesterMixin" not in parent_classes:
  # Put `PipelineTesterMixin` just before `unittest.TestCase`
  UpperCAmelCase__							=    [x for x in parent_classes if x != "TestCase"] + ["PipelineTesterMixin"]
  if "TestCase" in parent_classes:
   # Here we **assume** the original string is always with `unittest.TestCase`.
   _parent_classes.append("unittest.TestCase"						)
  UpperCAmelCase__							=    ", ".join(_parent_classes						)
  for idx, line in enumerate(_lowerCAmelCase						):
   # Find the ending of the declaration of `test_class`
   if line.strip().endswith("):"						):
    # mark the lines of the declaration of `test_class` to be removed
    for _idx in range(idx + 1						):
     UpperCAmelCase__							=    None  # noqa
    break
        # Add the new, one-line, class declaration for `test_class`
  UpperCAmelCase__							=    F'''class {test_class.__name__}({parent_classes}):\n'''

 # Add indentation
 UpperCAmelCase__							=    " " * indent_level + line_to_add
 # Insert `pipeline_model_mapping` to `class_lines`.
 # (The line at `target_idx` should be kept by definition!)
 UpperCAmelCase__							=    class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
 # Remove the lines that are marked to be removed
 UpperCAmelCase__							=    [x for x in class_lines if x is not None]

 # Move from test class to module (in order to write to the test file)
 UpperCAmelCase__							=    inspect.getsourcelines(inspect.getmodule(_lowerCAmelCase						)						)[0]
 # Be careful with the 1-off between line numbers and array indices
 UpperCAmelCase__							=    module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
 UpperCAmelCase__							=    "".join(_lowerCAmelCase						)

 UpperCAmelCase__							=    inspect.getsourcefile(_lowerCAmelCase						)
 with open(_lowerCAmelCase    ,				"w"    ,				encoding="UTF-8"    ,				newline="\n"						) as fp:
  fp.write(_lowerCAmelCase						)

 return line_to_add



def     lowerCAmelCase		(		_lowerCAmelCase				: Tuple    ,				_lowerCAmelCase				: List[Any]=False						):




 """simple docstring"""







 UpperCAmelCase__							=    find_test_class(_lowerCAmelCase						)
 if test_class:
  add_pipeline_model_mapping(_lowerCAmelCase    ,				overwrite=_lowerCAmelCase						)


if __name__ == "__main__":
  _lowerCAmelCase    : Tuple		= argparse.ArgumentParser()
  parser.add_argument(
      "--test_file", type=str, help="A path to the test file, starting with the repository's `tests` directory."
  )
  parser.add_argument(
      "--all",
      action="store_true",
      help="If to check and modify all test files.",
  )
  parser.add_argument(
      "--overwrite",
      action="store_true",
      help="If to overwrite a test class if it has already defined `pipeline_model_mapping`.",
  )
  _lowerCAmelCase    : Optional[Any]		= parser.parse_args()

  if not args.all and not args.test_file:
    raise ValueError("Please specify either `test_file` or pass `--all` to check/modify all test files.")
  elif args.all and args.test_file:
    raise ValueError("Only one of `--test_file` and `--all` could be specified.")

  _lowerCAmelCase    : List[str]		= []
  if args.test_file:
    _lowerCAmelCase    : Any		= [args.test_file]
  else:
    _lowerCAmelCase    : Any		= os.path.join("tests", "models", "**", "test_modeling_*.py")
    for test_file in glob.glob(pattern):
      # `Flax` is not concerned at this moment
      if not test_file.startswith("test_modeling_flax_"):
        test_files.append(test_file)

  for test_file in test_files:
    if test_file in TEST_FILE_TO_IGNORE:
      print(F'''[SKIPPED] {test_file} is skipped as it is in `TEST_FILE_TO_IGNORE` in the file {__file__}.''')
      continue
    add_pipeline_model_mapping_to_test_file(test_file, overwrite=args.overwrite)





