


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 _a   (     self       )    ->			Dict:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    @property
    def _a   (     self       )    ->			Union[str, Any]:
         __UpperCamelCase       				=1
         __UpperCamelCase       				=3
         __UpperCamelCase       				=(32, 32)

         __UpperCamelCase       				=floats_tensor((batch_size, num_channels) + sizes   ,  rng=random.Random(0       )       ).to(A_       )
         return image






    @property
    def _a   (     self       )    ->			List[Any]:
         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   ,  )
         return model






    @property
    def _a   (     self       )    ->			Optional[int]:
         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   ,  )
         return model






    @property
    def _a   (     self       )    ->			Dict:
         torch.manual_seed(0       )
         __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   ,  )
         return CLIPTextModel(A_       )






    @property
    def _a   (     self       )    ->			Dict:
         def extract(*A_   ,  **A_       ):






              class        UpperCAmelCase__			:



                  """simple docstring"""






                  def __init__(     self       )    ->			List[Any]:
                       __UpperCamelCase       				=torch.ones([0]       )






                  def _a   (     self   ,  A_       )    ->			Union[str, Any]:
                       self.pixel_values.to(A_       )
                       return self
              return Out()

         return extract






    def _a   (     self       )    ->			Dict:
         __UpperCamelCase       				='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase       				=self.dummy_cond_unet
         __UpperCamelCase       				=DDIMScheduler(
             beta_start=0.0_0085   ,  beta_end=0.012   ,  beta_schedule='scaled_linear'   ,  clip_sample=A_   ,  set_alpha_to_one=A_   ,  )

         __UpperCamelCase       				=self.dummy_vae
         __UpperCamelCase       				=self.dummy_text_encoder
         __UpperCamelCase       				=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

         # make sure here that pndm scheduler skips prk
         __UpperCamelCase       				=StableDiffusionPipeline(
             unet=A_   ,  scheduler=A_   ,  vae=A_   ,  text_encoder=A_   ,  tokenizer=A_   ,  safety_checker=A_   ,  feature_extractor=self.dummy_extractor   ,  )
         __UpperCamelCase       				=sd_pipe.to(A_       )
         sd_pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='A painting of a squirrel eating a burger'

         __UpperCamelCase       				=torch.Generator(device=A_       ).manual_seed(0       )
         __UpperCamelCase       				=sd_pipe([prompt]   ,  generator=A_   ,  guidance_scale=6.0   ,  num_inference_steps=2   ,  output_type='np'       )
         __UpperCamelCase       				=output.images

         __UpperCamelCase       				=torch.Generator(device=A_       ).manual_seed(0       )
         __UpperCamelCase       				=sd_pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=6.0   ,  num_inference_steps=2   ,  output_type='np'   ,  return_dict=A_   ,  )[0]

         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         __UpperCamelCase       				=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 _a   (     self       )    ->			Any:
         __UpperCamelCase       				='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase       				=self.dummy_cond_unet
         __UpperCamelCase       				=PNDMScheduler(skip_prk_steps=A_       )
         __UpperCamelCase       				=self.dummy_vae
         __UpperCamelCase       				=self.dummy_text_encoder
         __UpperCamelCase       				=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

         # make sure here that pndm scheduler skips prk
         __UpperCamelCase       				=StableDiffusionPipeline(
             unet=A_   ,  scheduler=A_   ,  vae=A_   ,  text_encoder=A_   ,  tokenizer=A_   ,  safety_checker=A_   ,  feature_extractor=self.dummy_extractor   ,  )
         __UpperCamelCase       				=sd_pipe.to(A_       )
         sd_pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='A painting of a squirrel eating a burger'
         __UpperCamelCase       				=torch.Generator(device=A_       ).manual_seed(0       )
         __UpperCamelCase       				=sd_pipe([prompt]   ,  generator=A_   ,  guidance_scale=6.0   ,  num_inference_steps=2   ,  output_type='np'       )

         __UpperCamelCase       				=output.images

         __UpperCamelCase       				=torch.Generator(device=A_       ).manual_seed(0       )
         __UpperCamelCase       				=sd_pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=6.0   ,  num_inference_steps=2   ,  output_type='np'   ,  return_dict=A_   ,  )[0]

         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 64, 64, 3)
         __UpperCamelCase       				=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 _a   (     self       )    ->			str:
         __UpperCamelCase       				=StableDiffusionPipeline.from_pretrained(
             'hf-internal-testing/tiny-stable-diffusion-lms-pipe'   ,  safety_checker=A_       )
         assert isinstance(A_   ,  A_       )
         assert isinstance(pipe.scheduler   ,  A_       )
         assert pipe.safety_checker is None

         __UpperCamelCase       				=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(A_       )
              __UpperCamelCase       				=StableDiffusionPipeline.from_pretrained(A_       )

         # sanity check that the pipeline still works
         assert pipe.safety_checker is None
         __UpperCamelCase       				=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 _a   (     self       )    ->			Any:
         __UpperCamelCase       				=self.dummy_cond_unet
         __UpperCamelCase       				=PNDMScheduler(skip_prk_steps=A_       )
         __UpperCamelCase       				=self.dummy_vae
         __UpperCamelCase       				=self.dummy_text_encoder
         __UpperCamelCase       				=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

         # put models in fp16
         __UpperCamelCase       				=unet.half()
         __UpperCamelCase       				=vae.half()
         __UpperCamelCase       				=bert.half()

         # make sure here that pndm scheduler skips prk
         __UpperCamelCase       				=StableDiffusionPipeline(
             unet=A_   ,  scheduler=A_   ,  vae=A_   ,  text_encoder=A_   ,  tokenizer=A_   ,  safety_checker=A_   ,  feature_extractor=self.dummy_extractor   ,  )
         __UpperCamelCase       				=sd_pipe.to(A_       )
         sd_pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='A painting of a squirrel eating a burger'
         __UpperCamelCase       				=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 _a   (     self       )    ->			Dict:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def _a   (     self       )    ->			Tuple:
         __UpperCamelCase       				='cuda'
         __UpperCamelCase       				=StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'       )
         __UpperCamelCase       				=pipe.to(A_       )
         pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='a photo of a cat'
         __UpperCamelCase       				={
             '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,
         }

         __UpperCamelCase       				=3
         __UpperCamelCase       				=7

         # no sega enabled
         __UpperCamelCase       				=torch.Generator(A_       )
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.3467_3113,
             0.3849_2733,
             0.3759_7352,
             0.3408_6335,
             0.3565_0748,
             0.3557_9205,
             0.338_4763,
             0.3434_0236,
             0.357_3271,
         ]

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

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

         # with sega enabled
         # generator = torch.manual_seed(seed)
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  **A_   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.4188_7826,
             0.3772_8766,
             0.3013_8272,
             0.4141_6335,
             0.4166_4985,
             0.3628_3392,
             0.3619_1246,
             0.4336_4465,
             0.4300_1732,
         ]

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

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






    def _a   (     self       )    ->			int:
         __UpperCamelCase       				='cuda'
         __UpperCamelCase       				=StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'       )
         __UpperCamelCase       				=pipe.to(A_       )
         pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='an image of a crowded boulevard, realistic, 4k'
         __UpperCamelCase       				={
             '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,
         }

         __UpperCamelCase       				=9
         __UpperCamelCase       				=7

         # no sega enabled
         __UpperCamelCase       				=torch.Generator(A_       )
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.4349_7998,
             0.9181_4065,
             0.754_0739,
             0.5558_0205,
             0.846_7265,
             0.538_9691,
             0.6257_4506,
             0.5889_7763,
             0.5092_6757,
         ]

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

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

         # with sega enabled
         # generator = torch.manual_seed(seed)
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  **A_   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.308_9719,
             0.3050_0144,
             0.2901_6042,
             0.3063_0964,
             0.32_5687,
             0.2941_9225,
             0.290_8091,
             0.2872_3598,
             0.2769_6294,
         ]

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

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






    def _a   (     self       )    ->			List[Any]:
         __UpperCamelCase       				='cuda'
         __UpperCamelCase       				=StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'       )
         __UpperCamelCase       				=pipe.to(A_       )
         pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='a castle next to a river'
         __UpperCamelCase       				={
             '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,
         }

         __UpperCamelCase       				=48
         __UpperCamelCase       				=7

         # no sega enabled
         __UpperCamelCase       				=torch.Generator(A_       )
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.7516_3555,
             0.7603_7145,
             0.6_1785,
             0.918_9673,
             0.862_7701,
             0.8518_9694,
             0.851_2813,
             0.8701_2076,
             0.831_2857,
         ]

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

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

         # with sega enabled
         # generator = torch.manual_seed(seed)
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  **A_   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.7355_3365,
             0.753_7271,
             0.7434_1905,
             0.6648_0356,
             0.647_2925,
             0.6303_9416,
             0.6481_2905,
             0.674_9717,
             0.651_7102,
         ]

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

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






    def _a   (     self       )    ->			str:
         __UpperCamelCase       				='cuda'
         __UpperCamelCase       				=StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5'   ,  torch_dtype=torch.floataa       )
         __UpperCamelCase       				=pipe.to(A_       )
         pipe.set_progress_bar_config(disable=A_       )

         __UpperCamelCase       				='a photo of a cat'
         __UpperCamelCase       				={
             '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,
         }

         __UpperCamelCase       				=3
         __UpperCamelCase       				=7

         # no sega enabled
         __UpperCamelCase       				=torch.Generator(A_       )
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.3488_7695,
             0.387_6953,
             0.375,
             0.3442_3828,
             0.358_1543,
             0.3571_7773,
             0.338_3789,
             0.3457_0312,
             0.35_9375,
         ]

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

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

         # with sega enabled
         # generator = torch.manual_seed(seed)
         generator.manual_seed(A_       )
         __UpperCamelCase       				=pipe(
             [prompt]   ,  generator=A_   ,  guidance_scale=A_   ,  num_inference_steps=50   ,  output_type='np'   ,  width=512   ,  height=512   ,  **A_   ,  )

         __UpperCamelCase       				=output.images
         __UpperCamelCase       				=image[0, -3:, -3:, -1]
         __UpperCamelCase       				=[
             0.4228_5156,
             0.3691_4062,
             0.2907_7148,
             0.4204_1016,
             0.4191_8945,
             0.3549_8047,
             0.361_8164,
             0.442_3828,
             0.4311_5234,
         ]

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

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









import argparse
import json

from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
from transformers.utils import logging


logging.set_verbosity_info()

def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   List[Any]     ,	SCREAMING_SNAKE_CASE__     :   Union[str, Any]     ,	SCREAMING_SNAKE_CASE__     :   Any							):
     # Initialise PyTorch model
     __UpperCamelCase       				=json.load(open(SCREAMING_SNAKE_CASE__     ,	'r'							)							)
     __UpperCamelCase       				=GPTNeoConfig(
         hidden_size=config_json['n_embd']     ,	num_layers=config_json['n_layer']     ,	num_heads=config_json['n_head']     ,	attention_types=config_json['attention_types']     ,	max_position_embeddings=config_json['n_positions']     ,	resid_dropout=config_json['res_dropout']     ,	embed_dropout=config_json['embed_dropout']     ,	attention_dropout=config_json['attn_dropout']     ,	)
     print(F'Building PyTorch model from configuration: {config}'							)
     __UpperCamelCase       				=GPTNeoForCausalLM(SCREAMING_SNAKE_CASE__							)

     # Load weights from tf checkpoint
     load_tf_weights_in_gpt_neo(SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__							)

     # Save pytorch-model
     print(F'Save PyTorch model to {pytorch_dump_path}'							)
     model.save_pretrained(SCREAMING_SNAKE_CASE__							)


if __name__ == "__main__":
   _A									=		argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
   )
   parser.add_argument(
       '--config_file',
       default=None,
       type=str,
       required=True,
       help=(
           'The config json file corresponding to the pre-trained mesh-tf model. \n'
           'This specifies the model architecture.'
       ),
   )
   parser.add_argument(
       '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
   )
   _A									=		parser.parse_args()
   convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)





