



"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

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


enable_full_determinism()



class     A__     (						__SCREAMING_SNAKE_CASE						,	__SCREAMING_SNAKE_CASE						,	__SCREAMING_SNAKE_CASE						,	unittest.TestCase							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE					   =			StableUnCLIPImgaImgPipeline
	SCREAMING_SNAKE_CASE					   =			TEXT_GUIDED_IMAGE_VARIATION_PARAMS
	SCREAMING_SNAKE_CASE					   =			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
	SCREAMING_SNAKE_CASE					   =			frozenset(
	    []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
	SCREAMING_SNAKE_CASE					   =			frozenset([]							)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> Dict:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[Any]				= 32
				__lowerCAmelCase   :						Union[str, Any]				= embedder_hidden_size

				# image encoding components

				__lowerCAmelCase   :						Optional[Any]				= CLIPImageProcessor(crop_size=32 ,    size=32)

				torch.manual_seed(0)
				__lowerCAmelCase   :						int				= CLIPVisionModelWithProjection(
				    CLIPVisionConfig(
				        hidden_size=_SCREAMING_SNAKE_CASE ,    projection_dim=_SCREAMING_SNAKE_CASE ,    num_hidden_layers=5 ,    num_attention_heads=4 ,    image_size=32 ,    intermediate_size=37 ,    patch_size=1 ,    ))

				# regular denoising components

				torch.manual_seed(0)
				__lowerCAmelCase   :						Optional[Any]				= StableUnCLIPImageNormalizer(embedding_dim=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Optional[Any]				= DDPMScheduler(beta_schedule="squaredcos_cap_v2")

				torch.manual_seed(0)
				__lowerCAmelCase   :						Optional[Any]				= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

				torch.manual_seed(0)
				__lowerCAmelCase   :						Dict				= CLIPTextModel(
				    CLIPTextConfig(
				        bos_token_id=0 ,    eos_token_id=2 ,    hidden_size=_SCREAMING_SNAKE_CASE ,    projection_dim=32 ,    intermediate_size=37 ,    layer_norm_eps=1e-05 ,    num_attention_heads=4 ,    num_hidden_layers=5 ,    pad_token_id=1 ,    vocab_size=1000 ,    ))

				torch.manual_seed(0)
				__lowerCAmelCase   :						Any				= UNetaDConditionModel(
				    sample_size=32 ,    in_channels=4 ,    out_channels=4 ,    down_block_types=("CrossAttnDownBlock2D", "DownBlock2D") ,    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D") ,    block_out_channels=(32, 64) ,    attention_head_dim=(2, 4) ,    class_embed_type="projection" ,    projection_class_embeddings_input_dim=embedder_projection_dim * 2 ,    cross_attention_dim=_SCREAMING_SNAKE_CASE ,    layers_per_block=1 ,    upcast_attention=_SCREAMING_SNAKE_CASE ,    use_linear_projection=_SCREAMING_SNAKE_CASE ,    )

				torch.manual_seed(0)
				__lowerCAmelCase   :						str				= DDIMScheduler(
				    beta_schedule="scaled_linear" ,    beta_start=0.0_0085 ,    beta_end=0.012 ,    prediction_type="v_prediction" ,    set_alpha_to_one=_SCREAMING_SNAKE_CASE ,    steps_offset=1 ,    )

				torch.manual_seed(0)
				__lowerCAmelCase   :						Optional[Any]				= AutoencoderKL()

				__lowerCAmelCase   :						Optional[Any]				= {
				    # image encoding components
				    "feature_extractor": feature_extractor,
				    "image_encoder": image_encoder.eval(),
				    # image noising components
				    "image_normalizer": image_normalizer.eval(),
				    "image_noising_scheduler": image_noising_scheduler,
				    # regular denoising components
				    "tokenizer": tokenizer,
				    "text_encoder": text_encoder.eval(),
				    "unet": unet.eval(),
				    "scheduler": scheduler,
				    "vae": vae.eval(),
				}

				return components

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any ,    _SCREAMING_SNAKE_CASE:	Dict ,    _SCREAMING_SNAKE_CASE:	str=0 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=True)       -> Dict:

				"""simple docstring"""

				if str(_SCREAMING_SNAKE_CASE).startswith("mps"):
							__lowerCAmelCase   :						List[str]				= torch.manual_seed(_SCREAMING_SNAKE_CASE)
				else:
							__lowerCAmelCase   :						Dict				= torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						List[Any]				= floats_tensor((1, 3, 32, 32) ,    rng=random.Random(_SCREAMING_SNAKE_CASE)).to(_SCREAMING_SNAKE_CASE)

				if pil_image:
							__lowerCAmelCase   :						List[Any]				= input_image * 0.5 + 0.5
							__lowerCAmelCase   :						Dict				= input_image.clamp(0 ,    1)
							__lowerCAmelCase   :						Union[str, Any]				= input_image.cpu().permute(0 ,    2 ,    3 ,    1).float().numpy()
							__lowerCAmelCase   :						Union[str, Any]				= DiffusionPipeline.numpy_to_pil(_SCREAMING_SNAKE_CASE)[0]

				return {
				    "prompt": "An anime racoon running a marathon",
				    "image": input_image,
				    "generator": generator,
				    "num_inference_steps": 2,
				    "output_type": "np",
				}

	@skip_mps
	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> Union[str, Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= "cpu"  # ensure determinism for the device-dependent torch.Generator
				__lowerCAmelCase   :						List[Any]				= self.get_dummy_components()
				__lowerCAmelCase   :						Dict				= StableUnCLIPImgaImgPipeline(**_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= sd_pipe.to(_SCREAMING_SNAKE_CASE)
				sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						int				= self.get_dummy_inputs(_SCREAMING_SNAKE_CASE)
				inputs.update({"image_embeds": None})
				__lowerCAmelCase   :						str				= sd_pipe(**_SCREAMING_SNAKE_CASE).images
				__lowerCAmelCase   :						Any				= image[0, -3:, -3:, -1]

				assert image.shape == (1, 32, 32, 3)
				__lowerCAmelCase   :						Any				= np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078])

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

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[int])       -> str:

				"""simple docstring"""

				__lowerCAmelCase   :						Tuple				= torch_device in ["cpu", "mps"]

				self._test_attention_slicing_forward_pass(test_max_difference=_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> str:

				"""simple docstring"""

				__lowerCAmelCase   :						List[str]				= torch_device in ["cpu", "mps"]

				self._test_inference_batch_single_identical(test_max_difference=_SCREAMING_SNAKE_CASE)







	@unittest.skipIf(
	    torch_device != "cuda" or not is_xformers_available() ,    reason="XFormers attention is only available with CUDA and `xformers` installed" ,    )
	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> int:

				"""simple docstring"""

				self._test_xformers_attention_forwardGenerator_pass(test_max_difference=_SCREAMING_SNAKE_CASE)






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




	'''simple docstring'''





	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> Dict:

				"""simple docstring"""

				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Dict				= load_image(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png")

				__lowerCAmelCase   :						Union[str, Any]				= load_numpy(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy")

				__lowerCAmelCase   :						Optional[Any]				= StableUnCLIPImgaImgPipeline.from_pretrained(
				    "fusing/stable-unclip-2-1-l-img2img" ,    torch_dtype=torch.floataa)
				pipe.to(_SCREAMING_SNAKE_CASE)
				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)
				# stable unclip will oom when integration tests are run on a V100,
				# so turn on memory savings
				pipe.enable_attention_slicing()
				pipe.enable_sequential_cpu_offload()

				__lowerCAmelCase   :						List[Any]				= torch.Generator(device="cpu").manual_seed(0)
				__lowerCAmelCase   :						Optional[int]				= pipe(_SCREAMING_SNAKE_CASE ,    "anime turle" ,    generator=_SCREAMING_SNAKE_CASE ,    output_type="np")

				__lowerCAmelCase   :						List[Any]				= output.images[0]

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

				assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> Any:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[Any]				= load_image(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png")

				__lowerCAmelCase   :						Dict				= load_numpy(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy")

				__lowerCAmelCase   :						Optional[Any]				= StableUnCLIPImgaImgPipeline.from_pretrained(
				    "fusing/stable-unclip-2-1-h-img2img" ,    torch_dtype=torch.floataa)
				pipe.to(_SCREAMING_SNAKE_CASE)
				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)
				# stable unclip will oom when integration tests are run on a V100,
				# so turn on memory savings
				pipe.enable_attention_slicing()
				pipe.enable_sequential_cpu_offload()

				__lowerCAmelCase   :						Optional[int]				= torch.Generator(device="cpu").manual_seed(0)
				__lowerCAmelCase   :						List[str]				= pipe(_SCREAMING_SNAKE_CASE ,    "anime turle" ,    generator=_SCREAMING_SNAKE_CASE ,    output_type="np")

				__lowerCAmelCase   :						Tuple				= output.images[0]

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

				assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)







	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> str:

				"""simple docstring"""

				__lowerCAmelCase   :						Tuple				= load_image(
				    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png")

				torch.cuda.empty_cache()
				torch.cuda.reset_max_memory_allocated()
				torch.cuda.reset_peak_memory_stats()

				__lowerCAmelCase   :						List[Any]				= StableUnCLIPImgaImgPipeline.from_pretrained(
				    "fusing/stable-unclip-2-1-h-img2img" ,    torch_dtype=torch.floataa)
				__lowerCAmelCase   :						Dict				= pipe.to(_SCREAMING_SNAKE_CASE)
				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE)
				pipe.enable_attention_slicing()
				pipe.enable_sequential_cpu_offload()

				__lowerCAmelCase   :						Tuple				= pipe(
				    _SCREAMING_SNAKE_CASE ,    "anime turtle" ,    num_inference_steps=2 ,    output_type="np" ,    )

				__lowerCAmelCase   :						Dict				= torch.cuda.max_memory_allocated()
				# make sure that less than 7 GB is allocated
				assert mem_bytes < 7 * 10**9




"""simple docstring"""






import argparse

import torch

from ...utils import logging
from . import AlbertConfig, AlbertForPreTraining, load_tf_weights_in_albert


logging.set_verbosity_info()


def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case			)		->     Tuple:
			# Initialise PyTorch model
			__lowerCAmelCase   :						Any				= AlbertConfig.from_json_file(__snake_case			)
			print(F"""Building PyTorch model from configuration: {config}"""			)
			__lowerCAmelCase   :						str				= AlbertForPreTraining(__snake_case			)

			# Load weights from tf checkpoint
			load_tf_weights_in_albert(__snake_case      ,__snake_case      ,__snake_case			)

			# Save pytorch-model
			print(F"""Save PyTorch model to {pytorch_dump_path}"""			)
			torch.save(model.state_dict()      ,__snake_case			)


if __name__ == "__main__":
	__snake_case							:     Union[str, Any]   =					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(
	    '--albert_config_file',
	    default=None,
	    type=str,
	    required=True,
	    help=(
	        'The config json file corresponding to the pre-trained ALBERT 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.'
	)
	__snake_case							:     str   =					parser.parse_args()
	convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.albert_config_file, args.pytorch_dump_path)