"""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        snake_case (				__UpperCAmelCase   ,					__UpperCAmelCase   ,					unittest.TestCase       ):




					"""simple docstring"""

					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 	__lowerCAmelCase		(       self     :			Any    ):
									return self._get_dummy_components()







					def 	__lowerCAmelCase		(       self     :			Any					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			Dict=0    ):
									if str(lowerCamelCase__    ).startswith('mps'    ):
													UpperCAmelCase__    =							torch.manual_seed(lowerCamelCase__    )
									else:
													UpperCAmelCase__    =							torch.Generator(device=lowerCamelCase__    ).manual_seed(lowerCamelCase__    )

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

									return inputs







					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									self._test_save_load_optional_components()







					@unittest.skipIf(torch_device != 'cuda'					,reason='float16 requires CUDA'    )
					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
									super().test_save_load_floataa(expected_max_diff=1e-1    )







					def 	__lowerCAmelCase		(       self     :			Tuple    ):
									self._test_attention_slicing_forward_pass(expected_max_diff=1e-2    )







					def 	__lowerCAmelCase		(       self     :			Tuple    ):
									self._test_save_load_local()







					def 	__lowerCAmelCase		(       self     :			Dict    ):
									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 	__lowerCAmelCase		(       self     :			Dict    ):
									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3    )






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




					"""simple docstring"""

					def 	__lowerCAmelCase		(       self     :			Union[str, Any]    ):
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()







					def 	__lowerCAmelCase		(       self     :			List[str]    ):
									# if

									UpperCAmelCase__    =							IFPipeline.from_pretrained('DeepFloyd/IF-I-XL-v1.0'					,variant='fp16'					,torch_dtype=torch.floataa    )

									UpperCAmelCase__    =							IFSuperResolutionPipeline.from_pretrained(
									    'DeepFloyd/IF-II-L-v1.0'					,variant='fp16'					,torch_dtype=torch.floataa					,text_encoder=lowerCamelCase__					,tokenizer=lowerCamelCase__    )

									# pre compute text embeddings and remove T5 to save memory

									pipe_a.text_encoder.to('cuda'    )

									UpperCAmelCase__   ,			UpperCAmelCase__    =							pipe_a.encode_prompt('anime turtle'					,device='cuda'    )

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

									UpperCAmelCase__    =							None
									UpperCAmelCase__    =							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(lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__    )

									pipe_a.remove_all_hooks()
									pipe_a.remove_all_hooks()

									# img2img

									UpperCAmelCase__    =							IFImgaImgPipeline(**pipe_a.components    )
									UpperCAmelCase__    =							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(lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__    )

									pipe_a.remove_all_hooks()
									pipe_a.remove_all_hooks()

									# inpainting

									UpperCAmelCase__    =							IFInpaintingPipeline(**pipe_a.components    )
									UpperCAmelCase__    =							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(lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Optional[int]					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			Tuple					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			List[str]    ):
									# pipeline 1

									_start_torch_memory_measurement()

									UpperCAmelCase__    =							torch.Generator(device='cpu'    ).manual_seed(0    )
									UpperCAmelCase__    =							pipe_a(
									    prompt_embeds=lowerCamelCase__					,negative_prompt_embeds=lowerCamelCase__					,num_inference_steps=2					,generator=lowerCamelCase__					,output_type='np'					,)

									UpperCAmelCase__    =							output.images[0]

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

									UpperCAmelCase__    =							torch.cuda.max_memory_allocated()
									assert mem_bytes < 13 * 10**9

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

									# pipeline 2

									_start_torch_memory_measurement()

									UpperCAmelCase__    =							torch.Generator(device='cpu'    ).manual_seed(0    )

									UpperCAmelCase__    =							floats_tensor((1, 3, 64, 64)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )

									UpperCAmelCase__    =							pipe_a(
									    prompt_embeds=lowerCamelCase__					,negative_prompt_embeds=lowerCamelCase__					,image=lowerCamelCase__					,generator=lowerCamelCase__					,num_inference_steps=2					,output_type='np'					,)

									UpperCAmelCase__    =							output.images[0]

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

									UpperCAmelCase__    =							torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**9

									UpperCAmelCase__    =							load_numpy(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'    )
									assert_mean_pixel_difference(lowerCamelCase__					,lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			int					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			Optional[int]					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			Dict    ):
									# pipeline 1

									_start_torch_memory_measurement()

									UpperCAmelCase__    =							floats_tensor((1, 3, 64, 64)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )

									UpperCAmelCase__    =							torch.Generator(device='cpu'    ).manual_seed(0    )

									UpperCAmelCase__    =							pipe_a(
									    prompt_embeds=lowerCamelCase__					,negative_prompt_embeds=lowerCamelCase__					,image=lowerCamelCase__					,num_inference_steps=2					,generator=lowerCamelCase__					,output_type='np'					,)

									UpperCAmelCase__    =							output.images[0]

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

									UpperCAmelCase__    =							torch.cuda.max_memory_allocated()
									assert mem_bytes < 10 * 10**9

									UpperCAmelCase__    =							load_numpy(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'    )
									assert_mean_pixel_difference(lowerCamelCase__					,lowerCamelCase__    )

									# pipeline 2

									_start_torch_memory_measurement()

									UpperCAmelCase__    =							torch.Generator(device='cpu'    ).manual_seed(0    )

									UpperCAmelCase__    =							floats_tensor((1, 3, 256, 256)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )
									UpperCAmelCase__    =							floats_tensor((1, 3, 64, 64)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )

									UpperCAmelCase__    =							pipe_a(
									    prompt_embeds=lowerCamelCase__					,negative_prompt_embeds=lowerCamelCase__					,image=lowerCamelCase__					,original_image=lowerCamelCase__					,generator=lowerCamelCase__					,num_inference_steps=2					,output_type='np'					,)

									UpperCAmelCase__    =							output.images[0]

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

									UpperCAmelCase__    =							torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**9

									UpperCAmelCase__    =							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(lowerCamelCase__					,lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			int					,lowerCamelCase__     :			Dict					,lowerCamelCase__     :			Union[str, Any]					,lowerCamelCase__     :			Optional[Any]					,lowerCamelCase__     :			Union[str, Any]    ):
									# pipeline 1

									_start_torch_memory_measurement()

									UpperCAmelCase__    =							floats_tensor((1, 3, 64, 64)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )
									UpperCAmelCase__    =							floats_tensor((1, 3, 64, 64)					,rng=random.Random(1    )    ).to(lowerCamelCase__    )

									UpperCAmelCase__    =							torch.Generator(device='cpu'    ).manual_seed(0    )
									UpperCAmelCase__    =							pipe_a(
									    prompt_embeds=lowerCamelCase__					,negative_prompt_embeds=lowerCamelCase__					,image=lowerCamelCase__					,mask_image=lowerCamelCase__					,num_inference_steps=2					,generator=lowerCamelCase__					,output_type='np'					,)

									UpperCAmelCase__    =							output.images[0]

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

									UpperCAmelCase__    =							torch.cuda.max_memory_allocated()
									assert mem_bytes < 10 * 10**9

									UpperCAmelCase__    =							load_numpy(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'    )
									assert_mean_pixel_difference(lowerCamelCase__					,lowerCamelCase__    )

									# pipeline 2

									_start_torch_memory_measurement()

									UpperCAmelCase__    =							torch.Generator(device='cpu'    ).manual_seed(0    )

									UpperCAmelCase__    =							floats_tensor((1, 3, 64, 64)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )
									UpperCAmelCase__    =							floats_tensor((1, 3, 256, 256)					,rng=random.Random(0    )    ).to(lowerCamelCase__    )
									UpperCAmelCase__    =							floats_tensor((1, 3, 256, 256)					,rng=random.Random(1    )    ).to(lowerCamelCase__    )

									UpperCAmelCase__    =							pipe_a(
									    prompt_embeds=lowerCamelCase__					,negative_prompt_embeds=lowerCamelCase__					,image=lowerCamelCase__					,mask_image=lowerCamelCase__					,original_image=lowerCamelCase__					,generator=lowerCamelCase__					,num_inference_steps=2					,output_type='np'					,)

									UpperCAmelCase__    =							output.images[0]

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

									UpperCAmelCase__    =							torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**9

									UpperCAmelCase__    =							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(lowerCamelCase__					,lowerCamelCase__    )






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




"""simple docstring"""






# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os

from accelerate.utils import ComputeEnvironment

from .cluster import get_cluster_input
from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
from .sagemaker import get_sagemaker_input


lowerCAmelCase__						:						Dict		  =			'Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine'



def 						a_		(      ):
				UpperCAmelCase__    =							_ask_options(
				    'In which compute environment are you running?'      ,    ['This machine', 'AWS (Amazon SageMaker)']      ,    _convert_compute_environment      ,    )
				if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
								UpperCAmelCase__    =							get_sagemaker_input()
				else:
								UpperCAmelCase__    =							get_cluster_input()
				return config



def 						a_		(      lowerCamelCase=None ):
				if subparsers is not None:
								UpperCAmelCase__    =							subparsers.add_parser('config'      ,    description=lowerCamelCase )
				else:
								UpperCAmelCase__    =							argparse.ArgumentParser('Accelerate config command'      ,    description=lowerCamelCase )

				parser.add_argument(
				    '--config_file'      ,    default=lowerCamelCase      ,    help=(
				        'The path to use to store the config file. Will default to a file named default_config.yaml in the cache '
				        'location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '
				        'such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '
				        'with \'huggingface\'.'
				    )      ,    )

				if subparsers is not None:
								parser.set_defaults(func=lowerCamelCase )
				return parser



def 						a_		(      lowerCamelCase ):
				UpperCAmelCase__    =							get_user_input()
				if args.config_file is not None:
								UpperCAmelCase__    =							args.config_file
				else:
								if not os.path.isdir(lowerCamelCase ):
												os.makedirs(lowerCamelCase )
								UpperCAmelCase__    =							default_yaml_config_file

				if config_file.endswith('.json' ):
								config.to_json_file(lowerCamelCase )
				else:
								config.to_yaml_file(lowerCamelCase )
				print(f'''accelerate configuration saved at {config_file}''' )



def 						a_		(      ):
				UpperCAmelCase__    =							config_command_parser()
				UpperCAmelCase__    =							parser.parse_args()
				config_command(lowerCamelCase )


if __name__ == "__main__":
						main()



