




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

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


lowerCamelCase	:	int				=   False




@skip_mps
class 						__lowercase		(UpperCamelCase__					,		UpperCamelCase__					,		UpperCamelCase__					,		unittest.TestCase							):




						"""simple docstring"""





						_snake_case												=						StableDiffusionAttendAndExcitePipeline
						_snake_case												=						False
						_snake_case												=						TEXT_TO_IMAGE_PARAMS
						_snake_case												=						TEXT_TO_IMAGE_BATCH_PARAMS.union({"""token_indices"""}							)
						_snake_case												=						TEXT_TO_IMAGE_IMAGE_PARAMS
						_snake_case												=						TEXT_TO_IMAGE_IMAGE_PARAMS

						@classmethod
						def      UpperCAmelCase				(				cls     )		->    str:
									super().setUpClass()
									torch.use_deterministic_algorithms(_snake_case     )

						@classmethod
						def      UpperCAmelCase				(				cls     )		->    Optional[Any]:
									super().tearDownClass()
									torch.use_deterministic_algorithms(_snake_case     )

						def      UpperCAmelCase				(				self     )		->    List[str]:
									torch.manual_seed(0     )
									snake_case				:    Optional[Any]           =	UNetaDConditionModel(
									    block_out_channels=(3_2, 6_4)    ,				layers_per_block=1    ,				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    ,				attention_head_dim=(2, 4)    ,				use_linear_projection=_snake_case    ,				)
									snake_case				:    int           =	DDIMScheduler(
									    beta_start=0.0_00_85    ,				beta_end=0.0_12    ,				beta_schedule="""scaled_linear"""    ,				clip_sample=_snake_case    ,				set_alpha_to_one=_snake_case    ,				)
									torch.manual_seed(0     )
									snake_case				:    Tuple           =	AutoencoderKL(
									    block_out_channels=[3_2, 6_4]    ,				in_channels=3    ,				out_channels=3    ,				down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]    ,				up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]    ,				latent_channels=4    ,				sample_size=1_2_8    ,				)
									torch.manual_seed(0     )
									snake_case				:    Union[str, Any]           =	CLIPTextConfig(
									    bos_token_id=0    ,				eos_token_id=2    ,				hidden_size=3_2    ,				intermediate_size=3_7    ,				layer_norm_eps=1e-0_5    ,				num_attention_heads=4    ,				num_hidden_layers=5    ,				pad_token_id=1    ,				vocab_size=1_0_0_0    ,				hidden_act="""gelu"""    ,				projection_dim=5_1_2    ,				)
									snake_case				:    Dict           =	CLIPTextModel(_snake_case     )
									snake_case				:    Tuple           =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

									snake_case				:    str           =	{
									    """unet""": unet,
									    """scheduler""": scheduler,
									    """vae""": vae,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									    """safety_checker""": None,
									    """feature_extractor""": None,
									}

									return components

						def      UpperCAmelCase				(				self    ,				A    ,				A=0     )		->    Optional[Any]:


									if str(_snake_case     ).startswith("""mps"""     ):
												snake_case				:    Optional[int]           =	torch.manual_seed(_snake_case     )
									else:
												snake_case				:    int           =	torch.Generator(device=_snake_case     ).manual_seed(_snake_case     )
									snake_case				:    Optional[Any]           =	{
									    """prompt""": """a cat and a frog""",
									    """token_indices""": [2, 5],
									    """generator""": generator,
									    """num_inference_steps""": 1,
									    """guidance_scale""": 6.0,
									    """output_type""": """numpy""",
									    """max_iter_to_alter""": 2,
									    """thresholds""": {0: 0.7},
									}
									return inputs

						def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
									snake_case				:    str           =	"""cpu"""

									snake_case				:    Optional[int]           =	self.get_dummy_components()
									snake_case				:    List[str]           =	self.pipeline_class(**_snake_case     )
									pipe.to(_snake_case     )
									pipe.set_progress_bar_config(disable=_snake_case     )

									snake_case				:    List[Any]           =	self.get_dummy_inputs(_snake_case     )
									snake_case				:    Tuple           =	pipe(**_snake_case     ).images
									snake_case				:    Tuple           =	image[0, -3:, -3:, -1]

									self.assertEqual(image.shape    ,				(1, 6_4, 6_4, 3)     )
									snake_case				:    List[str]           =	np.array(
									    [0.63_90_53_64, 0.62_89_73_07, 0.48_59_90_17, 0.5_13_36_24, 0.5_55_00_48, 0.45_76_95_16, 0.50_32_69_73, 0.5_02_31_39, 0.45_38_44_96]     )
									snake_case				:    Tuple           =	np.abs(image_slice.flatten() - expected_slice     ).max()
									self.assertLessEqual(_snake_case    ,				1e-3     )

						def      UpperCAmelCase				(				self     )		->    List[Any]:
									super().test_cpu_offload_forward_pass(expected_max_diff=5e-4     )

						def      UpperCAmelCase				(				self     )		->    Dict:
									self._test_inference_batch_consistent(batch_sizes=[1, 2]     )

						def      UpperCAmelCase				(				self     )		->    str:
									self._test_inference_batch_single_identical(batch_size=2    ,				expected_max_diff=7e-4     )

						def      UpperCAmelCase				(				self     )		->    Any:
									super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3     )

						def      UpperCAmelCase				(				self     )		->    List[Any]:
									super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4     )

						def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
									super().test_save_load_local(expected_max_difference=5e-4     )

						def      UpperCAmelCase				(				self     )		->    Any:
									super().test_save_load_optional_components(expected_max_difference=4e-4     )




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




						"""simple docstring"""





						@classmethod
						def      UpperCAmelCase				(				cls     )		->    Tuple:
									super().setUpClass()
									torch.use_deterministic_algorithms(_snake_case     )

						@classmethod
						def      UpperCAmelCase				(				cls     )		->    List[Any]:
									super().tearDownClass()
									torch.use_deterministic_algorithms(_snake_case     )

						def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()

						def      UpperCAmelCase				(				self     )		->    List[str]:
									snake_case				:    List[Any]           =	torch.manual_seed(5_1     )

									snake_case				:    int           =	StableDiffusionAttendAndExcitePipeline.from_pretrained(
									    """CompVis/stable-diffusion-v1-4"""    ,				safety_checker=_snake_case    ,				torch_dtype=torch.floataa     )
									pipe.to("""cuda"""     )

									snake_case				:    str           =	"""a painting of an elephant with glasses"""
									snake_case				:    int           =	[5, 7]

									snake_case				:    str           =	pipe(
									    prompt=_snake_case    ,				token_indices=_snake_case    ,				guidance_scale=7.5    ,				generator=_snake_case    ,				num_inference_steps=5    ,				max_iter_to_alter=5    ,				output_type="""numpy"""    ,				).images[0]

									snake_case				:    Optional[Any]           =	load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy"""     )
									assert np.abs((expected_image - image).max()     ) < 5e-1











from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       List[str]:
   return ConvertCommand(
       args.model_type						,args.tf_checkpoint						,args.pytorch_dump_output						,args.config						,args.finetuning_task_name					)


lowerCamelCase	:	Tuple				=   '\ntransformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires\nTensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.\n'




class 						__lowercase		(UpperCamelCase__							):




      """simple docstring"""





      @staticmethod
      def      UpperCAmelCase				(				A     )		->    List[str]:
         snake_case				:    Union[str, Any]           =	parser.add_parser(
             """convert"""    ,				help="""CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints."""    ,				)
         train_parser.add_argument("""--model_type"""    ,				type=A    ,				required=A    ,				help="""Model's type."""     )
         train_parser.add_argument(
             """--tf_checkpoint"""    ,				type=A    ,				required=A    ,				help="""TensorFlow checkpoint path or folder."""     )
         train_parser.add_argument(
             """--pytorch_dump_output"""    ,				type=A    ,				required=A    ,				help="""Path to the PyTorch saved model output."""     )
         train_parser.add_argument("""--config"""    ,				type=A    ,				default=""""""    ,				help="""Configuration file path or folder."""     )
         train_parser.add_argument(
             """--finetuning_task_name"""    ,				type=A    ,				default=A    ,				help="""Optional fine-tuning task name if the TF model was a finetuned model."""    ,				)
         train_parser.set_defaults(func=A     )

      def __init__(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				*A    ,				)		->    List[Any]:
         snake_case				:    Any           =	logging.get_logger("""transformers-cli/converting"""     )

         self._logger.info(f"""Loading model {model_type}"""     )
         snake_case				:    int           =	model_type
         snake_case				:    Any           =	tf_checkpoint
         snake_case				:    int           =	pytorch_dump_output
         snake_case				:    List[str]           =	config
         snake_case				:    Tuple           =	finetuning_task_name

      def      UpperCAmelCase				(				self     )		->    int:
         if self._model_type == "albert":
            try:
               from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "bert":
            try:
               from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "funnel":
            try:
               from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "t5":
            try:
               from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "gpt":
            from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
                convert_openai_checkpoint_to_pytorch,
            )

            convert_openai_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "transfo_xl":
            try:
               from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
                   convert_transfo_xl_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            if "ckpt" in self._tf_checkpoint.lower():
               snake_case				:    List[Any]           =	self._tf_checkpoint
               snake_case				:    Tuple           =	""""""
            else:
               snake_case				:    Tuple           =	self._tf_checkpoint
               snake_case				:    Tuple           =	""""""
            convert_transfo_xl_checkpoint_to_pytorch(
                A    ,				self._config    ,				self._pytorch_dump_output    ,				A     )
         elif self._model_type == "gpt2":
            try:
               from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
                   convert_gpta_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "xlnet":
            try:
               from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
                   convert_xlnet_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_xlnet_checkpoint_to_pytorch(
                self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output    ,				self._finetuning_task_name     )
         elif self._model_type == "xlm":
            from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
                convert_xlm_checkpoint_to_pytorch,
            )

            convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._pytorch_dump_output     )
         elif self._model_type == "lxmert":
            from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
                convert_lxmert_checkpoint_to_pytorch,
            )

            convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._pytorch_dump_output     )
         elif self._model_type == "rembert":
            from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
                convert_rembert_tf_checkpoint_to_pytorch,
            )

            convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         else:
            raise ValueError(
                """--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]"""     )





