


import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class 						UpperCAmelCase__   (      unittest.TestCase		):



     '''simple docstring'''




     def   lowerCAmelCase_   ( self       :       Tuple		):





            """simple docstring"""



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




     @property
     def   lowerCAmelCase_   ( self       :       Optional[Any]		):





            """simple docstring"""



            _lowercase					:				List[str]           =			1
            _lowercase					:				str           =			3
            _lowercase					:				Tuple           =			(32, 32)

            _lowercase					:				Dict           =			floats_tensor((batch_size, num_channels) + sizes	,      rng=random.Random(0		)		).to(UpperCamelCase		)
            return image




     @property
     def   lowerCAmelCase_   ( self       :       Optional[int]		):





            """simple docstring"""



            torch.manual_seed(0		)
            _lowercase					:				Any           =			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	,      )
            return model




     @property
     def   lowerCAmelCase_   ( self       :       Optional[int]		):





            """simple docstring"""



            torch.manual_seed(0		)
            _lowercase					:				Any           =			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	,      )
            return model




     @property
     def   lowerCAmelCase_   ( self       :       Dict		):





            """simple docstring"""



            torch.manual_seed(0		)
            _lowercase					:				Any           =			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=10_00	,      )
            return CLIPTextModel(UpperCamelCase		)




     @property
     def   lowerCAmelCase_   ( self       :       Dict		):





            """simple docstring"""



            def extract(*UpperCamelCase       :       Dict	,      **UpperCamelCase       :       Union[str, Any]		):

                   class 						UpperCAmelCase__   :



                        '''simple docstring'''




                        def __init__( self       :       str		):





                               """simple docstring"""



                               _lowercase					:				Tuple           =			torch.ones([0]		)




                        def   lowerCAmelCase_   ( self       :       Any	,      UpperCamelCase       :       Any		):





                               """simple docstring"""



                               self.pixel_values.to(UpperCamelCase		)
                               return self




                   return Out()

            return extract




     def   lowerCAmelCase_   ( self       :       int		):





            """simple docstring"""



            _lowercase					:				List[Any]           =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
            _lowercase					:				Optional[int]           =			self.dummy_cond_unet
            _lowercase					:				Tuple           =			DDIMScheduler(
                beta_start=0.00085	,      beta_end=0.012	,      beta_schedule='''scaled_linear'''	,      clip_sample=UpperCamelCase	,      set_alpha_to_one=UpperCamelCase	,      )

            _lowercase					:				List[Any]           =			self.dummy_vae
            _lowercase					:				str           =			self.dummy_text_encoder
            _lowercase					:				List[str]           =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

            # make sure here that pndm scheduler skips prk
            _lowercase					:				str           =			StableDiffusionPipeline(
                unet=UpperCamelCase	,      scheduler=UpperCamelCase	,      vae=UpperCamelCase	,      text_encoder=UpperCamelCase	,      tokenizer=UpperCamelCase	,      safety_checker=UpperCamelCase	,      feature_extractor=self.dummy_extractor	,      )
            _lowercase					:				Union[str, Any]           =			sd_pipe.to(UpperCamelCase		)
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Dict           =			'''A painting of a squirrel eating a burger'''

            _lowercase					:				List[str]           =			torch.Generator(device=UpperCamelCase		).manual_seed(0		)
            _lowercase					:				List[Any]           =			sd_pipe([prompt]	,      generator=UpperCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np'''		)
            _lowercase					:				Tuple           =			output.images

            _lowercase					:				Optional[int]           =			torch.Generator(device=UpperCamelCase		).manual_seed(0		)
            _lowercase					:				Optional[Any]           =			sd_pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np'''	,      return_dict=UpperCamelCase	,      )[0]

            _lowercase					:				int           =			image[0, -3:, -3:, -1]
            _lowercase					:				Optional[Any]           =			image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowercase					:				List[Any]           =			np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]		)

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




     def   lowerCAmelCase_   ( self       :       str		):





            """simple docstring"""



            _lowercase					:				str           =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
            _lowercase					:				str           =			self.dummy_cond_unet
            _lowercase					:				Optional[int]           =			PNDMScheduler(skip_prk_steps=UpperCamelCase		)
            _lowercase					:				Union[str, Any]           =			self.dummy_vae
            _lowercase					:				List[str]           =			self.dummy_text_encoder
            _lowercase					:				List[str]           =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

            # make sure here that pndm scheduler skips prk
            _lowercase					:				Optional[Any]           =			StableDiffusionPipeline(
                unet=UpperCamelCase	,      scheduler=UpperCamelCase	,      vae=UpperCamelCase	,      text_encoder=UpperCamelCase	,      tokenizer=UpperCamelCase	,      safety_checker=UpperCamelCase	,      feature_extractor=self.dummy_extractor	,      )
            _lowercase					:				int           =			sd_pipe.to(UpperCamelCase		)
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Optional[Any]           =			'''A painting of a squirrel eating a burger'''
            _lowercase					:				Union[str, Any]           =			torch.Generator(device=UpperCamelCase		).manual_seed(0		)
            _lowercase					:				Tuple           =			sd_pipe([prompt]	,      generator=UpperCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np'''		)

            _lowercase					:				str           =			output.images

            _lowercase					:				Any           =			torch.Generator(device=UpperCamelCase		).manual_seed(0		)
            _lowercase					:				int           =			sd_pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=6.0	,      num_inference_steps=2	,      output_type='''np'''	,      return_dict=UpperCamelCase	,      )[0]

            _lowercase					:				Optional[int]           =			image[0, -3:, -3:, -1]
            _lowercase					:				Union[str, Any]           =			image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowercase					:				str           =			np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]		)

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




     def   lowerCAmelCase_   ( self       :       str		):





            """simple docstring"""



            _lowercase					:				str           =			StableDiffusionPipeline.from_pretrained(
                '''hf-internal-testing/tiny-stable-diffusion-lms-pipe'''	,      safety_checker=UpperCamelCase		)
            assert isinstance(UpperCamelCase	,      UpperCamelCase		)
            assert isinstance(pipe.scheduler	,      UpperCamelCase		)
            assert pipe.safety_checker is None

            _lowercase					:				List[str]           =			pipe('''example prompt'''	,      num_inference_steps=2		).images[0]
            assert image is not None

            # check that there's no error when saving a pipeline with one of the models being None
            with tempfile.TemporaryDirectory() as tmpdirname:
                   pipe.save_pretrained(UpperCamelCase		)
                   _lowercase					:				List[Any]           =			StableDiffusionPipeline.from_pretrained(UpperCamelCase		)

            # sanity check that the pipeline still works
            assert pipe.safety_checker is None
            _lowercase					:				Optional[int]           =			pipe('''example prompt'''	,      num_inference_steps=2		).images[0]
            assert image is not None




     @unittest.skipIf(torch_device != '''cuda'''	,      '''This test requires a GPU'''		)
     def   lowerCAmelCase_   ( self       :       int		):





            """simple docstring"""



            _lowercase					:				Tuple           =			self.dummy_cond_unet
            _lowercase					:				Dict           =			PNDMScheduler(skip_prk_steps=UpperCamelCase		)
            _lowercase					:				Dict           =			self.dummy_vae
            _lowercase					:				List[str]           =			self.dummy_text_encoder
            _lowercase					:				int           =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

            # put models in fp16
            _lowercase					:				Optional[Any]           =			unet.half()
            _lowercase					:				Optional[int]           =			vae.half()
            _lowercase					:				Optional[int]           =			bert.half()

            # make sure here that pndm scheduler skips prk
            _lowercase					:				Tuple           =			StableDiffusionPipeline(
                unet=UpperCamelCase	,      scheduler=UpperCamelCase	,      vae=UpperCamelCase	,      text_encoder=UpperCamelCase	,      tokenizer=UpperCamelCase	,      safety_checker=UpperCamelCase	,      feature_extractor=self.dummy_extractor	,      )
            _lowercase					:				List[Any]           =			sd_pipe.to(UpperCamelCase		)
            sd_pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Tuple           =			'''A painting of a squirrel eating a burger'''
            _lowercase					:				Tuple           =			sd_pipe([prompt]	,      num_inference_steps=2	,      output_type='''np'''		).images

            assert image.shape == (1, 64, 64, 3)




@nightly
@require_torch_gpu
class 						UpperCAmelCase__   (      unittest.TestCase		):



     '''simple docstring'''




     def   lowerCAmelCase_   ( self       :       Dict		):





            """simple docstring"""



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




     def   lowerCAmelCase_   ( self       :       Dict		):





            """simple docstring"""



            _lowercase					:				Optional[int]           =			'''cuda'''
            _lowercase					:				Optional[int]           =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''		)
            _lowercase					:				Tuple           =			pipe.to(UpperCamelCase		)
            pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				List[Any]           =			'''a photo of a cat'''
            _lowercase					:				Optional[int]           =			{
                '''editing_prompt''': ['''sunglasses'''],
                '''reverse_editing_direction''': [False],
                '''edit_warmup_steps''': 10,
                '''edit_guidance_scale''': 6,
                '''edit_threshold''': 0.95,
                '''edit_momentum_scale''': 0.5,
                '''edit_mom_beta''': 0.6,
            }

            _lowercase					:				int           =			3
            _lowercase					:				Union[str, Any]           =			7

            # no sega enabled
            _lowercase					:				str           =			torch.Generator(UpperCamelCase		)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				int           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      )

            _lowercase					:				Optional[int]           =			output.images
            _lowercase					:				List[str]           =			image[0, -3:, -3:, -1]
            _lowercase					:				Tuple           =			[
                0.34673113,
                0.38492733,
                0.37597352,
                0.34086335,
                0.35650748,
                0.35579205,
                0.3384763,
                0.34340236,
                0.3573271,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				str           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      **UpperCamelCase	,      )

            _lowercase					:				List[Any]           =			output.images
            _lowercase					:				List[Any]           =			image[0, -3:, -3:, -1]
            _lowercase					:				Any           =			[
                0.41887826,
                0.37728766,
                0.30138272,
                0.41416335,
                0.41664985,
                0.36283392,
                0.36191246,
                0.43364465,
                0.43001732,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




     def   lowerCAmelCase_   ( self       :       Dict		):





            """simple docstring"""



            _lowercase					:				str           =			'''cuda'''
            _lowercase					:				Any           =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''		)
            _lowercase					:				Optional[int]           =			pipe.to(UpperCamelCase		)
            pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				int           =			'''an image of a crowded boulevard, realistic, 4k'''
            _lowercase					:				int           =			{
                '''editing_prompt''': '''crowd, crowded, people''',
                '''reverse_editing_direction''': True,
                '''edit_warmup_steps''': 10,
                '''edit_guidance_scale''': 8.3,
                '''edit_threshold''': 0.9,
                '''edit_momentum_scale''': 0.5,
                '''edit_mom_beta''': 0.6,
            }

            _lowercase					:				List[Any]           =			9
            _lowercase					:				Any           =			7

            # no sega enabled
            _lowercase					:				List[Any]           =			torch.Generator(UpperCamelCase		)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				Dict           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      )

            _lowercase					:				Optional[Any]           =			output.images
            _lowercase					:				Any           =			image[0, -3:, -3:, -1]
            _lowercase					:				Union[str, Any]           =			[
                0.43497998,
                0.91814065,
                0.7540739,
                0.55580205,
                0.8467265,
                0.5389691,
                0.62574506,
                0.58897763,
                0.50926757,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				Any           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      **UpperCamelCase	,      )

            _lowercase					:				Optional[int]           =			output.images
            _lowercase					:				Dict           =			image[0, -3:, -3:, -1]
            _lowercase					:				Optional[int]           =			[
                0.3089719,
                0.30500144,
                0.29016042,
                0.30630964,
                0.325687,
                0.29419225,
                0.2908091,
                0.28723598,
                0.27696294,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




     def   lowerCAmelCase_   ( self       :       Tuple		):





            """simple docstring"""



            _lowercase					:				List[str]           =			'''cuda'''
            _lowercase					:				List[Any]           =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''		)
            _lowercase					:				Tuple           =			pipe.to(UpperCamelCase		)
            pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Dict           =			'''a castle next to a river'''
            _lowercase					:				Optional[Any]           =			{
                '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
                '''reverse_editing_direction''': False,
                '''edit_warmup_steps''': [15, 18],
                '''edit_guidance_scale''': 6,
                '''edit_threshold''': [0.9, 0.8],
                '''edit_momentum_scale''': 0.5,
                '''edit_mom_beta''': 0.6,
            }

            _lowercase					:				Union[str, Any]           =			48
            _lowercase					:				List[str]           =			7

            # no sega enabled
            _lowercase					:				List[str]           =			torch.Generator(UpperCamelCase		)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				Union[str, Any]           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      )

            _lowercase					:				List[str]           =			output.images
            _lowercase					:				Optional[int]           =			image[0, -3:, -3:, -1]
            _lowercase					:				Tuple           =			[
                0.75163555,
                0.76037145,
                0.61785,
                0.9189673,
                0.8627701,
                0.85189694,
                0.8512813,
                0.87012076,
                0.8312857,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				int           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      **UpperCamelCase	,      )

            _lowercase					:				List[str]           =			output.images
            _lowercase					:				str           =			image[0, -3:, -3:, -1]
            _lowercase					:				List[Any]           =			[
                0.73553365,
                0.7537271,
                0.74341905,
                0.66480356,
                0.6472925,
                0.63039416,
                0.64812905,
                0.6749717,
                0.6517102,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




     def   lowerCAmelCase_   ( self       :       str		):





            """simple docstring"""



            _lowercase					:				Dict           =			'''cuda'''
            _lowercase					:				Any           =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	,      torch_dtype=torch.floataa		)
            _lowercase					:				Tuple           =			pipe.to(UpperCamelCase		)
            pipe.set_progress_bar_config(disable=UpperCamelCase		)

            _lowercase					:				Optional[Any]           =			'''a photo of a cat'''
            _lowercase					:				Union[str, Any]           =			{
                '''editing_prompt''': ['''sunglasses'''],
                '''reverse_editing_direction''': [False],
                '''edit_warmup_steps''': 10,
                '''edit_guidance_scale''': 6,
                '''edit_threshold''': 0.95,
                '''edit_momentum_scale''': 0.5,
                '''edit_mom_beta''': 0.6,
            }

            _lowercase					:				Optional[int]           =			3
            _lowercase					:				Optional[Any]           =			7

            # no sega enabled
            _lowercase					:				Union[str, Any]           =			torch.Generator(UpperCamelCase		)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				Optional[int]           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      )

            _lowercase					:				List[Any]           =			output.images
            _lowercase					:				Dict           =			image[0, -3:, -3:, -1]
            _lowercase					:				Dict           =			[
                0.34887695,
                0.3876953,
                0.375,
                0.34423828,
                0.3581543,
                0.35717773,
                0.3383789,
                0.34570312,
                0.359375,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(UpperCamelCase		)
            _lowercase					:				List[Any]           =			pipe(
                [prompt]	,      generator=UpperCamelCase	,      guidance_scale=UpperCamelCase	,      num_inference_steps=50	,      output_type='''np'''	,      width=5_12	,      height=5_12	,      **UpperCamelCase	,      )

            _lowercase					:				Union[str, Any]           =			output.images
            _lowercase					:				Dict           =			image[0, -3:, -3:, -1]
            _lowercase					:				Union[str, Any]           =			[
                0.42285156,
                0.36914062,
                0.29077148,
                0.42041016,
                0.41918945,
                0.35498047,
                0.3618164,
                0.4423828,
                0.43115234,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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



import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


UpperCamelCase__				       =			get_logger(__name__)

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_			)						->							List[str]:




       '''simple docstring'''






       if text.startswith(UpperCAmelCase_			):
              return text[len(UpperCAmelCase_			) :]
       return text  # or whatever

def   UpperCamelCase__   (    UpperCAmelCase_			)						->							Optional[int]:




       '''simple docstring'''






       return {remove_prefix(UpperCAmelCase_    ,     '''model.'''			): v for k, v in sd.items()}

def   UpperCamelCase__   (    UpperCAmelCase_			)						->							Any:




       '''simple docstring'''






       _lowercase					:				Optional[int]           =			{}
       for k in state_dicts[0].keys():
              _lowercase					:				Tuple           =			[sd[k] for sd in state_dicts]
              _lowercase					:				Union[str, Any]           =			sum(UpperCAmelCase_			) / len(UpperCAmelCase_			)
              assert isinstance(UpperCAmelCase_    ,     torch.Tensor			)
              _lowercase					:				Union[str, Any]           =			new_t
       return new_sd

def   UpperCamelCase__   (    UpperCAmelCase_    ,     UpperCAmelCase_    ,     UpperCAmelCase_			)						->							None:




       '''simple docstring'''






       _lowercase					:				List[str]           =			AutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase_			)
       if os.path.isfile(UpperCAmelCase_			):
              _lowercase					:				Optional[Any]           =			[pl_ckpt_path]
       else:
              assert os.path.isdir(UpperCAmelCase_			)
              _lowercase					:				List[str]           =			list(Path(UpperCAmelCase_			).glob('''*.ckpt'''			)			)
              assert ckpt_files, F'could not find any ckpt files inside the {pl_ckpt_path} directory'

       if len(UpperCAmelCase_			) > 1:
              logger.info(F'averaging the weights of {ckpt_files}'			)

       _lowercase					:				List[Any]           =			[sanitize(torch.load(UpperCAmelCase_    ,     map_location='''cpu'''			)['''state_dict''']			) for x in ckpt_files]
       _lowercase					:				Optional[Any]           =			average_state_dicts(UpperCAmelCase_			)

       _lowercase       ,	_lowercase					:				str           =			hf_model.load_state_dict(UpperCAmelCase_    ,     strict=UpperCAmelCase_			)
       assert not missing, F'missing keys: {missing}'
       hf_model.save_pretrained(UpperCAmelCase_			)
       try:
              _lowercase					:				Union[str, Any]           =			AutoTokenizer.from_pretrained(UpperCAmelCase_			)
              tok.save_pretrained(UpperCAmelCase_			)
       except Exception:
              pass
              # dont copy tokenizer if cant


if __name__ == "__main__":
   fire.Fire(convert_pl_to_hf)