



'''simple docstring'''





import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class  lowerCamelCase_   (    SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       unittest.TestCase					):
  """simple docstring"""




  a_							 =StableDiffusionXLPipeline
  a_							 =TEXT_TO_IMAGE_PARAMS
  a_							 =TEXT_TO_IMAGE_BATCH_PARAMS
  a_							 =TEXT_TO_IMAGE_IMAGE_PARAMS
  a_							 =TEXT_TO_IMAGE_IMAGE_PARAMS




  def 						_lowercase					(			self				:     List[str]       ) -> List[str]:
    torch.manual_seed(0       )
    __lowerCamelCase	:	Tuple   	=			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')      ,     attention_head_dim=(2, 4)      ,     use_linear_projection=_a      ,     addition_embed_type='text_time'      ,     addition_time_embed_dim=8      ,     transformer_layers_per_block=(1, 2)      ,     projection_class_embeddings_input_dim=80      ,     cross_attention_dim=64      ,     )
    __lowerCamelCase	:	Tuple   	=			EulerDiscreteScheduler(
        beta_start=0.00085      ,     beta_end=0.012      ,     steps_offset=1      ,     beta_schedule='scaled_linear'      ,     timestep_spacing='leading'      ,     )
    torch.manual_seed(0       )
    __lowerCamelCase	:	Optional[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      ,     sample_size=128      ,     )
    torch.manual_seed(0       )
    __lowerCamelCase	:	str   	=			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      ,     hidden_act='gelu'      ,     projection_dim=32      ,     )
    __lowerCamelCase	:	Union[str, Any]   	=			CLIPTextModel(_a       )
    __lowerCamelCase	:	int   	=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      ,     local_files_only=_a       )

    __lowerCamelCase	:	List[Any]   	=			CLIPTextModelWithProjection(_a       )
    __lowerCamelCase	:	Optional[int]   	=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      ,     local_files_only=_a       )

    __lowerCamelCase	:	List[Any]   	=			{
        'unet': unet,
        'scheduler': scheduler,
        'vae': vae,
        'text_encoder': text_encoder,
        'tokenizer': tokenizer,
        'text_encoder_2': text_encoder_a,
        'tokenizer_2': tokenizer_a,
        # "safety_checker": None,
        # "feature_extractor": None,
    }
    return components




  def 						_lowercase					(			self				:     Optional[Any]      ,     _a				:     Any      ,     _a				:     int=0       ) -> Any:
    if str(_a       ).startswith('mps'       ):
      __lowerCamelCase	:	List[str]   	=			torch.manual_seed(_a       )
    else:
      __lowerCamelCase	:	int   	=			torch.Generator(device=_a       ).manual_seed(_a       )
    __lowerCamelCase	:	Union[str, Any]   	=			{
        'prompt': 'A painting of a squirrel eating a burger',
        'generator': generator,
        'num_inference_steps': 2,
        'guidance_scale': 5.0,
        'output_type': 'numpy',
    }
    return inputs




  def 						_lowercase					(			self				:     Union[str, Any]       ) -> str:
    __lowerCamelCase	:	Dict   	=			'cpu'  # ensure determinism for the device-dependent torch.Generator
    __lowerCamelCase	:	Optional[Any]   	=			self.get_dummy_components()
    __lowerCamelCase	:	Dict   	=			StableDiffusionXLPipeline(**_a       )
    __lowerCamelCase	:	Optional[int]   	=			sd_pipe.to(_a       )
    sd_pipe.set_progress_bar_config(disable=_a       )

    __lowerCamelCase	:	Optional[int]   	=			self.get_dummy_inputs(_a       )
    __lowerCamelCase	:	Optional[int]   	=			sd_pipe(**_a       ).images
    __lowerCamelCase	:	Union[str, Any]   	=			image[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    __lowerCamelCase	:	Optional[int]   	=			np.array([0.5873, 0.6128, 0.4797, 0.5122, 0.5674, 0.4639, 0.5227, 0.5149, 0.4747]       )

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




  def 						_lowercase					(			self				:     Dict       ) -> str:
    __lowerCamelCase	:	List[str]   	=			self.get_dummy_components()
    __lowerCamelCase	:	str   	=			StableDiffusionXLPipeline(**_a       )
    __lowerCamelCase	:	Union[str, Any]   	=			sd_pipe.to(_a       )
    __lowerCamelCase	:	Any   	=			sd_pipe.to(_a       )
    sd_pipe.set_progress_bar_config(disable=_a       )

    # forward without prompt embeds
    __lowerCamelCase	:	Optional[int]   	=			self.get_dummy_inputs(_a       )
    __lowerCamelCase	:	List[Any]   	=			3 * ['this is a negative prompt']
    __lowerCamelCase	:	List[str]   	=			negative_prompt
    __lowerCamelCase	:	int   	=			3 * [inputs['prompt']]

    __lowerCamelCase	:	Optional[Any]   	=			sd_pipe(**_a       )
    __lowerCamelCase	:	int   	=			output.images[0, -3:, -3:, -1]

    # forward with prompt embeds
    __lowerCamelCase	:	Tuple   	=			self.get_dummy_inputs(_a       )
    __lowerCamelCase	:	Optional[int]   	=			3 * ['this is a negative prompt']
    __lowerCamelCase	:	Dict   	=			3 * [inputs.pop('prompt'       )]

    (
        (
        __lowerCamelCase
    )	,(
        __lowerCamelCase
    )	,(
        __lowerCamelCase
    )	,(
        __lowerCamelCase
    )	,
    )	:	int   	=			sd_pipe.encode_prompt(_a      ,     negative_prompt=_a       )

    __lowerCamelCase	:	Optional[int]   	=			sd_pipe(
        **_a      ,     prompt_embeds=_a      ,     negative_prompt_embeds=_a      ,     pooled_prompt_embeds=_a      ,     negative_pooled_prompt_embeds=_a      ,     )
    __lowerCamelCase	:	Union[str, Any]   	=			output.images[0, -3:, -3:, -1]

    # make sure that it's equal
    assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()       ).max() < 1e-4




  def 						_lowercase					(			self				:     Any       ) -> Any:
    super().test_attention_slicing_forward_pass(expected_max_diff=3e-3       )




  def 						_lowercase					(			self				:     List[str]       ) -> Union[str, Any]:
    super().test_inference_batch_single_identical(expected_max_diff=3e-3       )






@slow
@require_torch_gpu
class  lowerCamelCase_   (    unittest.TestCase					):
  """simple docstring"""




  def 						_lowercase					(			self				:     Dict       ) -> str:
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




  def 						_lowercase					(			self				:     Dict      ,     _a				:     List[str]      ,     _a				:     Union[str, Any]="cpu"      ,     _a				:     int=torch.floataa      ,     _a				:     Any=0       ) -> Any:
    __lowerCamelCase	:	Dict   	=			torch.Generator(device=_a       ).manual_seed(_a       )
    __lowerCamelCase	:	int   	=			np.random.RandomState(_a       ).standard_normal((1, 4, 64, 64)       )
    __lowerCamelCase	:	int   	=			torch.from_numpy(_a       ).to(device=_a      ,     dtype=_a       )
    __lowerCamelCase	:	Optional[int]   	=			{
        'prompt': 'a photograph of an astronaut riding a horse',
        'latents': latents,
        'generator': generator,
        'num_inference_steps': 3,
        'guidance_scale': 7.5,
        'output_type': 'numpy',
    }
    return inputs




  def 						_lowercase					(			self				:     List[Any]       ) -> str:
    __lowerCamelCase	:	Optional[Any]   	=			DiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-base'       )
    pipe.to(_a       )
    pipe.set_progress_bar_config(disable=_a       )

    __lowerCamelCase	:	int   	=			self.get_inputs(_a       )
    __lowerCamelCase	:	Optional[int]   	=			pipe(**_a       ).images
    __lowerCamelCase	:	Union[str, Any]   	=			image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase	:	List[Any]   	=			np.array([0.49493, 0.47896, 0.40798, 0.54214, 0.53212, 0.48202, 0.47656, 0.46329, 0.48506]       )
    assert np.abs(image_slice - expected_slice       ).max() < 7e-3











'''simple docstring'''





import argparse

from transformers import CanineConfig, CanineModel, CanineTokenizer, load_tf_weights_in_canine
from transformers.utils import logging


logging.set_verbosity_info()

def 	a_					(							_lowerCAmelCase	,_lowerCAmelCase	)				->				List[Any]:
  # Initialize PyTorch model
  __lowerCamelCase	:	Union[str, Any]   	=			CanineConfig()
  __lowerCamelCase	:	Optional[Any]   	=			CanineModel(_lowerCAmelCase	)
  model.eval()

  print(F'Building PyTorch model from configuration: {config}'	)

  # Load weights from tf checkpoint
  load_tf_weights_in_canine(_lowerCAmelCase	,_lowerCAmelCase	,_lowerCAmelCase	)

  # Save pytorch-model (weights and configuration)
  print(F'Save PyTorch model to {pytorch_dump_path}'	)
  model.save_pretrained(_lowerCAmelCase	)

  # Save tokenizer files
  __lowerCamelCase	:	str   	=			CanineTokenizer()
  print(F'Save tokenizer files to {pytorch_dump_path}'	)
  tokenizer.save_pretrained(_lowerCAmelCase	)


if __name__ == "__main__":
   _UpperCamelCase							    =						argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--tf_checkpoint_path',
       default=None,
       type=str,
       required=True,
       help='Path to the TensorFlow checkpoint. Should end with model.ckpt',
   )
   parser.add_argument(
       '--pytorch_dump_path',
       default=None,
       type=str,
       required=True,
       help='Path to a folder where the PyTorch model will be placed.',
   )
   _UpperCamelCase							    =						parser.parse_args()
   convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.pytorch_dump_path)






