



'''simple docstring'''
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 							lowercase__     (							lowerCamelCase_						, unittest.TestCase       ):
	lowercase__      =      KandinskyVaaImgaImgPipeline
	lowercase__      =      ['''image_embeds''', '''negative_image_embeds''', '''image''']
	lowercase__      =      [
	    '''image_embeds''',
	    '''negative_image_embeds''',
	    '''image''',
	]
	lowercase__      =      [
	    '''generator''',
	    '''height''',
	    '''width''',
	    '''strength''',
	    '''guidance_scale''',
	    '''num_inference_steps''',
	    '''return_dict''',
	    '''guidance_scale''',
	    '''num_images_per_prompt''',
	    '''output_type''',
	    '''return_dict''',
	]
	lowercase__      =      False





	@property
	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							return 32





	@property
	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							return 32





	@property
	def      UpperCamelCase_  (			self :    List[Any]			):







							'''simple docstring'''
							return self.time_input_dim





	@property
	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							return self.time_input_dim * 4





	@property
	def      UpperCamelCase_  (			self :    Any			):







							'''simple docstring'''
							return 100





	@property
	def      UpperCamelCase_  (			self :    List[Any]			):







							'''simple docstring'''
							torch.manual_seed(0			)

							_UpperCamelCase :		List[str]			    =       {
							    'in_channels': 4,
							    # Out channels is double in channels because predicts mean and variance
							    'out_channels': 8,
							    'addition_embed_type': 'image',
							    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
							    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
							    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
							    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
							    'layers_per_block': 1,
							    'encoder_hid_dim': self.text_embedder_hidden_size,
							    'encoder_hid_dim_type': 'image_proj',
							    'cross_attention_dim': self.cross_attention_dim,
							    'attention_head_dim': 4,
							    'resnet_time_scale_shift': 'scale_shift',
							    'class_embed_type': None,
							}

							_UpperCamelCase :		Tuple			    =       UNetaDConditionModel(**__snake_case			)
							return model





	@property
	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							return {
							    "block_out_channels": [32, 64],
							    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
							    "in_channels": 3,
							    "latent_channels": 4,
							    "layers_per_block": 1,
							    "norm_num_groups": 8,
							    "norm_type": "spatial",
							    "num_vq_embeddings": 12,
							    "out_channels": 3,
							    "up_block_types": [
							        "AttnUpDecoderBlock2D",
							        "UpDecoderBlock2D",
							    ],
							    "vq_embed_dim": 4,
							}





	@property
	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							torch.manual_seed(0			)
							_UpperCamelCase :		int			    =       VQModel(**self.dummy_movq_kwargs			)
							return model





	def      UpperCamelCase_  (			self :    Optional[int]			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[int]			    =       self.dummy_unet
							_UpperCamelCase :		Union[str, Any]			    =       self.dummy_movq

							_UpperCamelCase :		Dict			    =       {
							    'num_train_timesteps': 1000,
							    'beta_schedule': 'linear',
							    'beta_start': 0.0_0_0_8_5,
							    'beta_end': 0.0_1_2,
							    'clip_sample': False,
							    'set_alpha_to_one': False,
							    'steps_offset': 0,
							    'prediction_type': 'epsilon',
							    'thresholding': False,
							}

							_UpperCamelCase :		Dict			    =       DDIMScheduler(**__snake_case			)

							_UpperCamelCase :		Optional[Any]			    =       {
							    'unet': unet,
							    'scheduler': scheduler,
							    'movq': movq,
							}

							return components





	def      UpperCamelCase_  (			self :    int			,lowerCamelCase__ :    Union[str, Any]			,lowerCamelCase__ :    Dict=0			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       floats_tensor((1, self.text_embedder_hidden_size)			,rng=random.Random(__snake_case			)			).to(__snake_case			)
							_UpperCamelCase :		Optional[Any]			    =       floats_tensor((1, self.text_embedder_hidden_size)			,rng=random.Random(seed + 1			)			).to(
							    __snake_case			)
							# create init_image
							_UpperCamelCase :		int			    =       floats_tensor((1, 3, 64, 64)			,rng=random.Random(__snake_case			)			).to(__snake_case			)
							_UpperCamelCase :		str			    =       image.cpu().permute(0			,2			,3			,1			)[0]
							_UpperCamelCase :		Dict			    =       Image.fromarray(np.uinta(__snake_case			)			).convert('RGB'			).resize((256, 256)			)

							if str(__snake_case			).startswith('mps'			):
													_UpperCamelCase :		Optional[Any]			    =       torch.manual_seed(__snake_case			)
							else:
													_UpperCamelCase :		int			    =       torch.Generator(device=__snake_case			).manual_seed(__snake_case			)
							_UpperCamelCase :		Optional[int]			    =       {
							    'image': init_image,
							    'image_embeds': image_embeds,
							    'negative_image_embeds': negative_image_embeds,
							    'generator': generator,
							    'height': 64,
							    'width': 64,
							    'num_inference_steps': 10,
							    'guidance_scale': 7.0,
							    'strength': 0.2,
							    'output_type': 'np',
							}
							return inputs





	def      UpperCamelCase_  (			self :    List[Any]			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[int]			    =       'cpu'

							_UpperCamelCase :		Optional[int]			    =       self.get_dummy_components()

							_UpperCamelCase :		str			    =       self.pipeline_class(**__snake_case			)
							_UpperCamelCase :		Tuple			    =       pipe.to(__snake_case			)

							pipe.set_progress_bar_config(disable=__snake_case			)

							_UpperCamelCase :		List[Any]			    =       pipe(**self.get_dummy_inputs(__snake_case			)			)
							_UpperCamelCase :		Any			    =       output.images

							_UpperCamelCase :		Any			    =       pipe(
							    **self.get_dummy_inputs(__snake_case			)			,return_dict=__snake_case			,)[0]

							_UpperCamelCase :		Dict			    =       image[0, -3:, -3:, -1]
							_UpperCamelCase :		Dict			    =       image_from_tuple[0, -3:, -3:, -1]

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

							_UpperCamelCase :		str			    =       np.array(
							    [0.6_1_9_9_7_7_8, 0.6_3_9_8_4_4_0_6, 0.4_6_1_4_5_7_8_5, 0.6_2_9_4_4_9_8_4, 0.5_6_2_2_2_1_5, 0.4_7_3_0_6_1_3_2, 0.4_7_4_4_1_4_5_6, 0.4_6_0_7_6_0_6, 0.4_8_7_1_9_2_6_3]			)
							assert (
							    np.abs(image_slice.flatten() - expected_slice			).max() < 1E-2
							), F' expected_slice {expected_slice}, but got {image_slice.flatten()}'
							assert (
							    np.abs(image_from_tuple_slice.flatten() - expected_slice			).max() < 1E-2
							), F' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'






@slow
@require_torch_gpu
class 							lowercase__     (							unittest.TestCase       ):





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def      UpperCamelCase_  (			self :    Optional[int]			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/kandinskyv22/kandinskyv22_img2img_frog.npy'			)

							_UpperCamelCase :		str			    =       load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/kandinsky/cat.png'			)
							_UpperCamelCase :		Union[str, Any]			    =       'A red cartoon frog, 4k'

							_UpperCamelCase :		str			    =       KandinskyVaaPriorPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-prior'			,torch_dtype=torch.floataa			)
							pipe_prior.to(__snake_case			)

							_UpperCamelCase :		int			    =       KandinskyVaaImgaImgPipeline.from_pretrained(
							    'kandinsky-community/kandinsky-2-2-decoder'			,torch_dtype=torch.floataa			)
							_UpperCamelCase :		Optional[Any]			    =       pipeline.to(__snake_case			)

							pipeline.set_progress_bar_config(disable=__snake_case			)

							_UpperCamelCase :		Tuple			    =       torch.Generator(device='cpu'			).manual_seed(0			)
							_UpperCamelCase    ,      _UpperCamelCase :		Union[str, Any]			    =       pipe_prior(
							    __snake_case			,generator=__snake_case			,num_inference_steps=5			,negative_prompt=''			,).to_tuple()

							_UpperCamelCase :		str			    =       pipeline(
							    image=__snake_case			,image_embeds=__snake_case			,negative_image_embeds=__snake_case			,generator=__snake_case			,num_inference_steps=100			,height=768			,width=768			,strength=0.2			,output_type='np'			,)

							_UpperCamelCase :		Optional[Any]			    =       output.images[0]

							assert image.shape == (768, 768, 3)

							assert_mean_pixel_difference(__snake_case			,__snake_case			)





'''simple docstring'''
import argparse

import torch

from transformers import LxmertConfig, LxmertForPreTraining, load_tf_weights_in_lxmert
from transformers.utils import logging


logging.set_verbosity_info()






def A__   (					UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_							):
						# Initialise PyTorch model
						_UpperCamelCase :		Any			    =       LxmertConfig.from_json_file(UpperCAmelCase_							)
						print(f'Building PyTorch model from configuration: {config}'							)
						_UpperCamelCase :		int			    =       LxmertForPreTraining(UpperCAmelCase_							)

						# Load weights from tf checkpoint
						load_tf_weights_in_lxmert(UpperCAmelCase_  ,  UpperCAmelCase_  ,  UpperCAmelCase_							)

						# Save pytorch-model
						print(f'Save PyTorch model to {pytorch_dump_path}'							)
						torch.save(model.state_dict()  ,  UpperCAmelCase_							)


if __name__ == "__main__":
					snake_case_			:     int  	=  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 model. \nThis specifies the model architecture.',
					)
					parser.add_argument(
					    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
					)
					snake_case_			:     List[str]  	=  parser.parse_args()
					convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)


