


"""simple docstring"""




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 				lowerCamelCase__       (	unittest.TestCase):





     def 					__A (self	)					->       str:
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





     @property
     def 					__A (self	)					->       Any:
      _lowercase											=1
      _lowercase											=3
      _lowercase											=(3_2, 3_2)

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





     @property
     def 					__A (self	)					->       Union[str, Any]:
      torch.manual_seed(0	)
      _lowercase											=UNetaDConditionModel(
          block_out_channels=(3_2, 6_4)		,					layers_per_block=2		,					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		,					)
      return model





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





     @property
     def 					__A (self	)					->       int:
      torch.manual_seed(0	)
      _lowercase											=CLIPTextConfig(
          bos_token_id=0		,					eos_token_id=2		,					hidden_size=3_2		,					intermediate_size=3_7		,					layer_norm_eps=1e-05		,					num_attention_heads=4		,					num_hidden_layers=5		,					pad_token_id=1		,					vocab_size=1_0_0_0		,					)
      return CLIPTextModel(UpperCAmelCase	)





     @property
     def 					__A (self	)					->       Tuple:
      def extract(*UpperCAmelCase		,					**UpperCAmelCase	):







       class 				lowerCamelCase__       :





            def __init__(self	)					->       Dict:
             _lowercase											=torch.ones([0]	)





            def 					__A (self		,					UpperCAmelCase	)					->       Any:
             self.pixel_values.to(UpperCAmelCase	)
             return self



       return Out()

      return extract





     def 					__A (self	)					->       List[str]:
      _lowercase											='''cpu'''  # ensure determinism for the device-dependent torch.Generator
      _lowercase											=self.dummy_cond_unet
      _lowercase											=DDIMScheduler(
          beta_start=0.0_0085		,					beta_end=0.012		,					beta_schedule='''scaled_linear'''		,					clip_sample=UpperCAmelCase		,					set_alpha_to_one=UpperCAmelCase		,					)

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

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

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

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

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

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

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

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

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

      _lowercase											=output.images

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

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

      assert image.shape == (1, 6_4, 6_4, 3)
      _lowercase											=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	)					->       Union[str, Any]:
      _lowercase											=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											=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											=StableDiffusionPipeline.from_pretrained(UpperCAmelCase	)

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

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

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

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

      assert image.shape == (1, 6_4, 6_4, 3)







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





     def 					__A (self	)					->       Optional[Any]:
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





     def 					__A (self	)					->       str:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

      _lowercase											='''a photo of a cat'''
      _lowercase											={
          '''editing_prompt''': ['''sunglasses'''],
          '''reverse_editing_direction''': [False],
          '''edit_warmup_steps''': 1_0,
          '''edit_guidance_scale''': 6,
          '''edit_threshold''': 0.95,
          '''edit_momentum_scale''': 0.5,
          '''edit_mom_beta''': 0.6,
      }

      _lowercase											=3
      _lowercase											=7

      # no sega enabled
      _lowercase											=torch.Generator(UpperCAmelCase	)
      generator.manual_seed(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 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											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 3)

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





     def 					__A (self	)					->       Union[str, Any]:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

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

      _lowercase											=9
      _lowercase											=7

      # no sega enabled
      _lowercase											=torch.Generator(UpperCAmelCase	)
      generator.manual_seed(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 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											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 3)

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





     def 					__A (self	)					->       str:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

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

      _lowercase											=4_8
      _lowercase											=7

      # no sega enabled
      _lowercase											=torch.Generator(UpperCAmelCase	)
      generator.manual_seed(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 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											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 3)

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





     def 					__A (self	)					->       Tuple:
      _lowercase											='''cuda'''
      _lowercase											=StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''		,					torch_dtype=torch.floataa	)
      _lowercase											=pipe.to(UpperCAmelCase	)
      pipe.set_progress_bar_config(disable=UpperCAmelCase	)

      _lowercase											='''a photo of a cat'''
      _lowercase											={
          '''editing_prompt''': ['''sunglasses'''],
          '''reverse_editing_direction''': [False],
          '''edit_warmup_steps''': 1_0,
          '''edit_guidance_scale''': 6,
          '''edit_threshold''': 0.95,
          '''edit_momentum_scale''': 0.5,
          '''edit_mom_beta''': 0.6,
      }

      _lowercase											=3
      _lowercase											=7

      # no sega enabled
      _lowercase											=torch.Generator(UpperCAmelCase	)
      generator.manual_seed(UpperCAmelCase	)
      _lowercase											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 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											=pipe(
          [prompt]		,					generator=UpperCAmelCase		,					guidance_scale=UpperCAmelCase		,					num_inference_steps=5_0		,					output_type='''np'''		,					width=5_1_2		,					height=5_1_2		,					**UpperCAmelCase		,					)

      _lowercase											=output.images
      _lowercase											=image[0, -3:, -3:, -1]
      _lowercase											=[
          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, 5_1_2, 5_1_2, 3)

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





import argparse

from . import (
    BertConfig,
    BertGenerationConfig,
    BertGenerationDecoder,
    BertGenerationEncoder,
    load_tf_weights_in_bert_generation,
    logging,
)


logging.set_verbosity_info()
def     UpperCAmelCase_ (							__snake_case					,						__snake_case					,						__snake_case					,						__snake_case					,						__snake_case						)					->      Tuple:







 """simple docstring"""


 _lowercase											=BertConfig.from_pretrained(
     '''bert-large-cased'''					,						vocab_size=__snake_case					,						max_position_embeddings=512					,						is_decoder=__snake_case					,						add_cross_attention=__snake_case					,						)
 _lowercase											=bert_config.to_dict()
 del bert_config_dict["type_vocab_size"]
 _lowercase											=BertGenerationConfig(**__snake_case						)
 if is_encoder:
  _lowercase											=BertGenerationEncoder(__snake_case						)
 else:
  _lowercase											=BertGenerationDecoder(__snake_case						)
 print(F"Building PyTorch model from configuration: {config}"						)

 # Load weights from tf checkpoint
 load_tf_weights_in_bert_generation(
     __snake_case					,						__snake_case					,						model_class='''bert'''					,						is_encoder_named_decoder=__snake_case					,						is_encoder=__snake_case					,						)

 # Save pytorch-model
 print(F"Save PyTorch model and config to {pytorch_dump_path}"						)
 model.save_pretrained(__snake_case						)


if __name__ == "__main__":
  UpperCAmelCase__            =							argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '''--tf_hub_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
  )
  parser.add_argument(
      '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
  )
  parser.add_argument(
      '''--is_encoder_named_decoder''',
      action='''store_true''',
      help='''If decoder has to be renamed to encoder in PyTorch model.''',
  )
  parser.add_argument('''--is_encoder''', action='''store_true''', help='''If model is an encoder.''')
  parser.add_argument('''--vocab_size''', default=5_0358, type=int, help='''Vocab size of model''')
  UpperCAmelCase__            =							parser.parse_args()
  convert_tf_checkpoint_to_pytorch(
      args.tf_hub_path,
      args.pytorch_dump_path,
      args.is_encoder_named_decoder,
      args.vocab_size,
      is_encoder=args.is_encoder,
  )

