code: 
  
  
  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
style_context: 
  
  
  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)