

"""simple docstring"""




import gc
import random
import unittest

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

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

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class      _UpperCAmelCase   (							_lowerCAmelCase		,    _lowerCAmelCase		,    unittest.TestCase       ):
	a__		: Optional[Any]					=					StableDiffusionXLImgaImgPipeline
	a__		: Any					=					TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
	a__		: str					=					PipelineTesterMixin.required_optional_params - {"latents"}
	a__		: Optional[int]					=					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	a__		: str					=					IMAGE_TO_IMAGE_IMAGE_PARAMS
	a__		: str					=					IMAGE_TO_IMAGE_IMAGE_PARAMS



	def        a  (  self     :       Union[str, 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''')			,    attention_head_dim=(2, 4)			,    use_linear_projection=_lowercase			,    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			,    )
		__UpperCAmelCase      		= EulerDiscreteScheduler(
		    beta_start=0.00_085			,    beta_end=0.012			,    steps_offset=1			,    beta_schedule='''scaled_linear'''			,    timestep_spacing='''leading'''			,    )
		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			,    sample_size=1_28			,    )
		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=10_00			,    hidden_act='''gelu'''			,    projection_dim=32			,    )
		__UpperCAmelCase      		= CLIPTextModel(_lowercase						)
		__UpperCAmelCase      		= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			,    local_files_only=_lowercase						)

		__UpperCAmelCase      		= CLIPTextModelWithProjection(_lowercase						)
		__UpperCAmelCase      		= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			,    local_files_only=_lowercase						)

		__UpperCAmelCase      		= {
		    '''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        a  (  self     :       List[Any]			,    _lowercase     :       List[str]			,    _lowercase     :       Any=0						):
		__UpperCAmelCase      		= floats_tensor((1, 3, 32, 32)			,    rng=random.Random(_lowercase						)						).to(_lowercase						)
		__UpperCAmelCase      		= image / 2 + 0.5
		if str(_lowercase						).startswith('''mps'''						):
			__UpperCAmelCase      		= torch.manual_seed(_lowercase						)
		else:
			__UpperCAmelCase      		= torch.Generator(device=_lowercase						).manual_seed(_lowercase						)
		__UpperCAmelCase      		= {
		    '''prompt''': '''A painting of a squirrel eating a burger''',
		    '''image''': image,
		    '''generator''': generator,
		    '''num_inference_steps''': 2,
		    '''guidance_scale''': 5.0,
		    '''output_type''': '''numpy''',
		    '''strength''': 0.75,
		}
		return inputs



	def        a  (  self     :       str						):
		__UpperCAmelCase      		= '''cpu'''  # ensure determinism for the device-dependent torch.Generator
		__UpperCAmelCase      		= self.get_dummy_components()
		__UpperCAmelCase      		= StableDiffusionXLImgaImgPipeline(**_lowercase						)
		__UpperCAmelCase      		= sd_pipe.to(_lowercase						)
		sd_pipe.set_progress_bar_config(disable=_lowercase						)

		__UpperCAmelCase      		= self.get_dummy_inputs(_lowercase						)
		__UpperCAmelCase      		= sd_pipe(**_lowercase						).images
		__UpperCAmelCase      		= image[0, -3:, -3:, -1]

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

		__UpperCAmelCase      		= np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]						)

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



	def        a  (  self     :       List[Any]						):
		super().test_attention_slicing_forward_pass(expected_max_diff=3E-3						)



	def        a  (  self     :       int						):
		super().test_inference_batch_single_identical(expected_max_diff=3E-3						)



	def        a  (  self     :       Optional[int]						):
		pass



	def        a  (  self     :       int						):
		__UpperCAmelCase      		= self.get_dummy_components()
		__UpperCAmelCase      		= StableDiffusionXLImgaImgPipeline(**_lowercase						)
		__UpperCAmelCase      		= sd_pipe.to(_lowercase						)
		__UpperCAmelCase      		= sd_pipe.to(_lowercase						)
		sd_pipe.set_progress_bar_config(disable=_lowercase						)

		# forward without prompt embeds
		__UpperCAmelCase      		= self.get_dummy_inputs(_lowercase						)
		__UpperCAmelCase      		= 3 * ['''this is a negative prompt''']
		__UpperCAmelCase      		= negative_prompt
		__UpperCAmelCase      		= 3 * [inputs['''prompt''']]

		__UpperCAmelCase      		= sd_pipe(**_lowercase						)
		__UpperCAmelCase      		= output.images[0, -3:, -3:, -1]

		# forward with prompt embeds
		__UpperCAmelCase      		= self.get_dummy_inputs(_lowercase						)
		__UpperCAmelCase      		= 3 * ['''this is a negative prompt''']
		__UpperCAmelCase      		= 3 * [inputs.pop('''prompt'''						)]

		(
		    (
		    __UpperCAmelCase
		)		,							(
		    __UpperCAmelCase
		)		,							(
		    __UpperCAmelCase
		)		,							(
		    __UpperCAmelCase
		)		,							
		)      		= sd_pipe.encode_prompt(_lowercase			,    negative_prompt=_lowercase						)

		__UpperCAmelCase      		= sd_pipe(
		    **_lowercase			,    prompt_embeds=_lowercase			,    negative_prompt_embeds=_lowercase			,    pooled_prompt_embeds=_lowercase			,    negative_pooled_prompt_embeds=_lowercase			,    )
		__UpperCAmelCase      		= 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






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



	def        a  (  self     :       Tuple						):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()



	def        a  (  self     :       Optional[int]			,    _lowercase     :       Optional[Any]			,    _lowercase     :       int="cpu"			,    _lowercase     :       List[str]=torch.floataa			,    _lowercase     :       Any=0						):
		__UpperCAmelCase      		= torch.Generator(device=_lowercase						).manual_seed(_lowercase						)
		__UpperCAmelCase      		= np.random.RandomState(_lowercase						).standard_normal((1, 4, 64, 64)						)
		__UpperCAmelCase      		= torch.from_numpy(_lowercase						).to(device=_lowercase			,    dtype=_lowercase						)
		__UpperCAmelCase      		= {
		    '''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        a  (  self     :       Tuple						):
		__UpperCAmelCase      		= DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''						)
		pipe.to(_lowercase						)
		pipe.set_progress_bar_config(disable=_lowercase						)

		__UpperCAmelCase      		= self.get_inputs(_lowercase						)
		__UpperCAmelCase      		= pipe(**_lowercase						).images
		__UpperCAmelCase      		= image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 5_12, 5_12, 3)
		__UpperCAmelCase      		= np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]						)
		assert np.abs(image_slice - expected_slice						).max() < 7E-3




'''simple docstring'''







import argparse

from transformers import TaConfig, TaForConditionalGeneration, load_tf_weights_in_ta
from transformers.utils import logging


logging.set_verbosity_info()






def       lowercase					(_A ,    _A ,    _A       ):







       """simple docstring"""







       _lowerCAmelCase     :				Optional[Any]       =							TaConfig.from_json_file(_A       )
       print(f'Building PyTorch model from configuration: {config}'       )
       _lowerCAmelCase     :				Union[str, Any]       =							TaForConditionalGeneration(_A       )

       # Load weights from tf checkpoint
       load_tf_weights_in_ta(_A ,    _A ,    _A       )

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


if __name__ == "__main__":
       lowerCAmelCase     :  str											=    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 T5 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."""
       )
       lowerCAmelCase     :  int											=    parser.parse_args()
       convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)

