




"""simple docstring"""






import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin




@skip_mps
class   a__			(       UpperCamelCase_					,	UpperCamelCase_					,	unittest.TestCase  ):
						snake_case__             = IFPipeline
						snake_case__             = TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
						snake_case__             = TEXT_TO_IMAGE_BATCH_PARAMS
						snake_case__             = PipelineTesterMixin.required_optional_params - {'''latents'''}





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




											"""simple docstring"""







											return self._get_dummy_components()





						def       __UpperCamelCase     (							self      :	Optional[Any]     ,a__      :	Union[str, Any]     ,a__      :	Optional[int]=0)						->       Tuple:




											"""simple docstring"""







											if str(a__).startswith('''mps'''):
																_lowerCAmelCase:Union[str, Any]									=  torch.manual_seed(a__)
											else:
																_lowerCAmelCase:Optional[Any]									=  torch.Generator(device=a__).manual_seed(a__)

											_lowerCAmelCase:int									=  {
											    '''prompt''': '''A painting of a squirrel eating a burger''',
											    '''generator''': generator,
											    '''num_inference_steps''': 2,
											    '''output_type''': '''numpy''',
											}

											return inputs





						def       __UpperCamelCase     (							self      :	Dict)						->       Optional[int]:




											"""simple docstring"""







											self._test_save_load_optional_components()





						@unittest.skipIf(torch_device != '''cuda'''     ,reason='''float16 requires CUDA''')
						def       __UpperCamelCase     (							self      :	Union[str, Any])						->       List[str]:




											"""simple docstring"""







											super().test_save_load_floataa(expected_max_diff=1E-1)





						def       __UpperCamelCase     (							self      :	List[Any])						->       Optional[int]:




											"""simple docstring"""







											self._test_attention_slicing_forward_pass(expected_max_diff=1E-2)





						def       __UpperCamelCase     (							self      :	Any)						->       Union[str, Any]:




											"""simple docstring"""







											self._test_save_load_local()





						def       __UpperCamelCase     (							self      :	str)						->       Optional[Any]:




											"""simple docstring"""







											self._test_inference_batch_single_identical(
											    expected_max_diff=1E-2     ,)







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




											"""simple docstring"""







											self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)




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





						def       __UpperCamelCase     (							self      :	str)						->       str:




											"""simple docstring"""







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





						def       __UpperCamelCase     (							self      :	Optional[int])						->       Optional[int]:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]									=  IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''     ,variant='''fp16'''     ,torch_dtype=torch.floataa)

											_lowerCAmelCase:List[Any]									=  IFSuperResolutionPipeline.from_pretrained(
											    '''DeepFloyd/IF-II-L-v1.0'''     ,variant='''fp16'''     ,torch_dtype=torch.floataa     ,text_encoder=a__     ,tokenizer=a__)

											# pre compute text embeddings and remove T5 to save memory

											pipe_a.text_encoder.to('''cuda''')

											_lowerCAmelCase							,   _lowerCAmelCase:str									=  pipe_a.encode_prompt('''anime turtle'''     ,device='''cuda''')

											del pipe_a.tokenizer
											del pipe_a.text_encoder
											gc.collect()

											_lowerCAmelCase:List[Any]									=  None
											_lowerCAmelCase:Optional[int]									=  None

											pipe_a.enable_model_cpu_offload()
											pipe_a.enable_model_cpu_offload()

											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())
											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())

											self._test_if(a__     ,a__     ,a__     ,a__)

											pipe_a.remove_all_hooks()
											pipe_a.remove_all_hooks()

											# img2img

											_lowerCAmelCase:Union[str, Any]									=  IFImgaImgPipeline(**pipe_a.components)
											_lowerCAmelCase:Tuple									=  IFImgaImgSuperResolutionPipeline(**pipe_a.components)

											pipe_a.enable_model_cpu_offload()
											pipe_a.enable_model_cpu_offload()

											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())
											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())

											self._test_if_imgaimg(a__     ,a__     ,a__     ,a__)

											pipe_a.remove_all_hooks()
											pipe_a.remove_all_hooks()

											# inpainting

											_lowerCAmelCase:List[str]									=  IFInpaintingPipeline(**pipe_a.components)
											_lowerCAmelCase:Tuple									=  IFInpaintingSuperResolutionPipeline(**pipe_a.components)

											pipe_a.enable_model_cpu_offload()
											pipe_a.enable_model_cpu_offload()

											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())
											pipe_a.unet.set_attn_processor(AttnAddedKVProcessor())

											self._test_if_inpainting(a__     ,a__     ,a__     ,a__)





						def       __UpperCamelCase     (							self      :	List[str]     ,a__      :	Optional[int]     ,a__      :	Any     ,a__      :	Any     ,a__      :	Dict)						->       List[str]:




											"""simple docstring"""







											_start_torch_memory_measurement()

											_lowerCAmelCase:int									=  torch.Generator(device='''cpu''').manual_seed(0)
											_lowerCAmelCase:Union[str, Any]									=  pipe_a(
											    prompt_embeds=a__     ,negative_prompt_embeds=a__     ,num_inference_steps=2     ,generator=a__     ,output_type='''np'''     ,)

											_lowerCAmelCase:Dict									=  output.images[0]

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

											_lowerCAmelCase:int									=  torch.cuda.max_memory_allocated()
											assert mem_bytes < 13 * 10**9

											_lowerCAmelCase:Dict									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy''')
											assert_mean_pixel_difference(a__     ,a__)

											# pipeline 2

											_start_torch_memory_measurement()

											_lowerCAmelCase:Tuple									=  torch.Generator(device='''cpu''').manual_seed(0)

											_lowerCAmelCase:Any									=  floats_tensor((1, 3, 64, 64)     ,rng=random.Random(0)).to(a__)

											_lowerCAmelCase:Dict									=  pipe_a(
											    prompt_embeds=a__     ,negative_prompt_embeds=a__     ,image=a__     ,generator=a__     ,num_inference_steps=2     ,output_type='''np'''     ,)

											_lowerCAmelCase:Tuple									=  output.images[0]

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

											_lowerCAmelCase:List[str]									=  torch.cuda.max_memory_allocated()
											assert mem_bytes < 4 * 10**9

											_lowerCAmelCase:Union[str, Any]									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy''')
											assert_mean_pixel_difference(a__     ,a__)





						def       __UpperCamelCase     (							self      :	str     ,a__      :	List[str]     ,a__      :	Optional[int]     ,a__      :	Dict     ,a__      :	List[Any])						->       List[str]:




											"""simple docstring"""







											_start_torch_memory_measurement()

											_lowerCAmelCase:Optional[int]									=  floats_tensor((1, 3, 64, 64)     ,rng=random.Random(0)).to(a__)

											_lowerCAmelCase:Optional[int]									=  torch.Generator(device='''cpu''').manual_seed(0)

											_lowerCAmelCase:Union[str, Any]									=  pipe_a(
											    prompt_embeds=a__     ,negative_prompt_embeds=a__     ,image=a__     ,num_inference_steps=2     ,generator=a__     ,output_type='''np'''     ,)

											_lowerCAmelCase:Optional[int]									=  output.images[0]

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

											_lowerCAmelCase:Tuple									=  torch.cuda.max_memory_allocated()
											assert mem_bytes < 10 * 10**9

											_lowerCAmelCase:List[Any]									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy''')
											assert_mean_pixel_difference(a__     ,a__)

											# pipeline 2

											_start_torch_memory_measurement()

											_lowerCAmelCase:Optional[Any]									=  torch.Generator(device='''cpu''').manual_seed(0)

											_lowerCAmelCase:Optional[int]									=  floats_tensor((1, 3, 256, 256)     ,rng=random.Random(0)).to(a__)
											_lowerCAmelCase:str									=  floats_tensor((1, 3, 64, 64)     ,rng=random.Random(0)).to(a__)

											_lowerCAmelCase:int									=  pipe_a(
											    prompt_embeds=a__     ,negative_prompt_embeds=a__     ,image=a__     ,original_image=a__     ,generator=a__     ,num_inference_steps=2     ,output_type='''np'''     ,)

											_lowerCAmelCase:str									=  output.images[0]

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

											_lowerCAmelCase:Tuple									=  torch.cuda.max_memory_allocated()
											assert mem_bytes < 4 * 10**9

											_lowerCAmelCase:Dict									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy''')
											assert_mean_pixel_difference(a__     ,a__)







						def       __UpperCamelCase     (							self      :	Tuple     ,a__      :	Dict     ,a__      :	Union[str, Any]     ,a__      :	Union[str, Any]     ,a__      :	Any)						->       List[Any]:




											"""simple docstring"""







											_start_torch_memory_measurement()

											_lowerCAmelCase:Tuple									=  floats_tensor((1, 3, 64, 64)     ,rng=random.Random(0)).to(a__)
											_lowerCAmelCase:str									=  floats_tensor((1, 3, 64, 64)     ,rng=random.Random(1)).to(a__)

											_lowerCAmelCase:Optional[int]									=  torch.Generator(device='''cpu''').manual_seed(0)
											_lowerCAmelCase:Dict									=  pipe_a(
											    prompt_embeds=a__     ,negative_prompt_embeds=a__     ,image=a__     ,mask_image=a__     ,num_inference_steps=2     ,generator=a__     ,output_type='''np'''     ,)

											_lowerCAmelCase:Dict									=  output.images[0]

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

											_lowerCAmelCase:Optional[int]									=  torch.cuda.max_memory_allocated()
											assert mem_bytes < 10 * 10**9

											_lowerCAmelCase:Union[str, Any]									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy''')
											assert_mean_pixel_difference(a__     ,a__)

											# pipeline 2

											_start_torch_memory_measurement()

											_lowerCAmelCase:Optional[int]									=  torch.Generator(device='''cpu''').manual_seed(0)

											_lowerCAmelCase:Tuple									=  floats_tensor((1, 3, 64, 64)     ,rng=random.Random(0)).to(a__)
											_lowerCAmelCase:List[Any]									=  floats_tensor((1, 3, 256, 256)     ,rng=random.Random(0)).to(a__)
											_lowerCAmelCase:Optional[Any]									=  floats_tensor((1, 3, 256, 256)     ,rng=random.Random(1)).to(a__)

											_lowerCAmelCase:Optional[int]									=  pipe_a(
											    prompt_embeds=a__     ,negative_prompt_embeds=a__     ,image=a__     ,mask_image=a__     ,original_image=a__     ,generator=a__     ,num_inference_steps=2     ,output_type='''np'''     ,)

											_lowerCAmelCase:List[str]									=  output.images[0]

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

											_lowerCAmelCase:str									=  torch.cuda.max_memory_allocated()
											assert mem_bytes < 4 * 10**9

											_lowerCAmelCase:str									=  load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy''')
											assert_mean_pixel_difference(a__     ,a__)





def 							UpperCAmelCase		(		):
					torch.cuda.empty_cache()
					torch.cuda.reset_max_memory_allocated()
					torch.cuda.reset_peak_memory_stats()












"""simple docstring"""






import argparse
import json
import re
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileNetVaConfig,
    MobileNetVaForImageClassification,
    MobileNetVaImageProcessor,
    load_tf_weights_in_mobilenet_va,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase__						=			logging.get_logger(__name__)





def 							UpperCAmelCase		(		snake_case			:					Optional[int]    ):
					_lowerCAmelCase:List[Any]									=  MobileNetVaConfig(layer_norm_eps=0.0_01    )

					if "_quant" in model_name:
										raise ValueError('''Quantized models are not supported.'''    )

					_lowerCAmelCase:List[Any]									=  re.match(R'''^mobilenet_v1_([^_]*)_([^_]*)$'''      ,       snake_case    )
					if matches:
										_lowerCAmelCase:str									=  float(matches[1]    )
										_lowerCAmelCase:str									=  int(matches[2]    )

					# The TensorFlow version of MobileNetV1 predicts 1001 classes instead of
					# the usual 1000. The first class (index 0) is "background".
					_lowerCAmelCase:int									=  1001
					_lowerCAmelCase:Union[str, Any]									=  '''imagenet-1k-id2label.json'''
					_lowerCAmelCase:Optional[int]									=  '''huggingface/label-files'''
					_lowerCAmelCase:List[str]									=  json.load(open(hf_hub_download(snake_case      ,       snake_case      ,       repo_type='''dataset'''    )      ,       '''r'''    )    )
					_lowerCAmelCase:Union[str, Any]									=  {int(snake_case    ) + 1: v for k, v in idalabel.items()}
					_lowerCAmelCase:Optional[Any]									=  '''background'''
					_lowerCAmelCase:Optional[Any]									=  idalabel
					_lowerCAmelCase:List[Any]									=  {v: k for k, v in idalabel.items()}

					return config





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:Tuple									=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
					_lowerCAmelCase:Optional[Any]									=  Image.open(requests.get(snake_case      ,       stream=snake_case    ).raw    )
					return im





@torch.no_grad()
def 							UpperCAmelCase		(		snake_case			:					Any      ,       snake_case			:					Union[str, Any]      ,       snake_case			:					List[Any]      ,       snake_case			:					Any=False    ):
					_lowerCAmelCase:List[str]									=  get_mobilenet_va_config(snake_case    )

					# Load 🤗 model
					_lowerCAmelCase:Any									=  MobileNetVaForImageClassification(snake_case    ).eval()

					# Load weights from TensorFlow checkpoint
					load_tf_weights_in_mobilenet_va(snake_case      ,       snake_case      ,       snake_case    )

					# Check outputs on an image, prepared by MobileNetV1ImageProcessor
					_lowerCAmelCase:Dict									=  MobileNetVaImageProcessor(
					    crop_size={'''width''': config.image_size, '''height''': config.image_size}      ,       size={'''shortest_edge''': config.image_size + 32}      ,       )
					_lowerCAmelCase:List[str]									=  image_processor(images=prepare_img()      ,       return_tensors='''pt'''    )
					_lowerCAmelCase:List[str]									=  model(**snake_case    )
					_lowerCAmelCase:int									=  outputs.logits

					assert logits.shape == (1, 1001)

					if model_name == "mobilenet_v1_1.0_224":
										_lowerCAmelCase:Union[str, Any]									=  torch.tensor([-4.17_39, -1.12_33, 3.12_05]    )
					elif model_name == "mobilenet_v1_0.75_192":
										_lowerCAmelCase:Any									=  torch.tensor([-3.94_40, -2.31_41, -0.33_33]    )
					else:
										_lowerCAmelCase:Dict									=  None

					if expected_logits is not None:
										assert torch.allclose(logits[0, :3]      ,       snake_case      ,       atol=1e-4    )

					Path(snake_case    ).mkdir(exist_ok=snake_case    )
					print(F'Saving model {model_name} to {pytorch_dump_folder_path}'    )
					model.save_pretrained(snake_case    )
					print(F'Saving image processor to {pytorch_dump_folder_path}'    )
					image_processor.save_pretrained(snake_case    )

					if push_to_hub:
										print('''Pushing to the hub...'''    )
										_lowerCAmelCase:List[str]									=  '''google/''' + model_name
										image_processor.push_to_hub(snake_case    )
										model.push_to_hub(snake_case    )


if __name__ == "__main__":
				UpperCamelCase__						=			argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '''--model_name''',
				    default='''mobilenet_v1_1.0_224''',
				    type=str,
				    help='''Name of the MobileNetV1 model you\'d like to convert. Should in the form \'mobilenet_v1_<depth>_<size>\'.''',
				)
				parser.add_argument(
				    '''--checkpoint_path''', required=True, type=str, help='''Path to the original TensorFlow checkpoint (.ckpt file).'''
				)
				parser.add_argument(
				    '''--pytorch_dump_folder_path''', required=True, type=str, help='''Path to the output PyTorch model directory.'''
				)
				parser.add_argument(
				    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
				)

				UpperCamelCase__						=			parser.parse_args()
				convert_movilevit_checkpoint(
				    args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
				)






