




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






				'''simple docstring'''
				__UpperCamelCase         =		IFPipeline
				__UpperCamelCase         =		TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
				__UpperCamelCase         =		TEXT_TO_IMAGE_BATCH_PARAMS
				__UpperCamelCase         =		PipelineTesterMixin.required_optional_params - {"latents"}





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return self._get_dummy_components()





				def   _lowerCAmelCase							(			self						,					_a						,					_a=0					):





							"""simple docstring"""
							if str(_a					).startswith("""mps"""					):
										lowerCamelCase      =   torch.manual_seed(_a					)
							else:
										lowerCamelCase      =   torch.Generator(device=_a					).manual_seed(_a					)

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

							return inputs





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							self._test_save_load_optional_components()





				@unittest.skipIf(torch_device != """cuda"""						,					reason="""float16 requires CUDA"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# 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					):





							"""simple docstring"""
							self._test_attention_slicing_forward_pass(expected_max_diff=1e-2					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							self._test_save_load_local()





				def   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3					)






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






				'''simple docstring'''





				def   _lowerCAmelCase							(			self					):





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





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# if

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

							lowerCamelCase      =   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      =   pipe_a.encode_prompt("""anime turtle"""						,					device="""cuda"""					)

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

							lowerCamelCase      =   None
							lowerCamelCase      =   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      =   IFImgaImgPipeline(**pipe_a.components					)
							lowerCamelCase      =   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      =   IFInpaintingPipeline(**pipe_a.components					)
							lowerCamelCase      =   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   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							# pipeline 1

							_start_torch_memory_measurement()

							lowerCamelCase      =   torch.Generator(device="""cpu"""					).manual_seed(0					)
							lowerCamelCase      =   pipe_a(
							    prompt_embeds=_a						,					negative_prompt_embeds=_a						,					num_inference_steps=2						,					generator=_a						,					output_type="""np"""						,					)

							lowerCamelCase      =   output.images[0]

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

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

							lowerCamelCase      =   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      =   torch.Generator(device="""cpu"""					).manual_seed(0					)

							lowerCamelCase      =   floats_tensor((1, 3, 64, 64)						,					rng=random.Random(0					)					).to(_a					)

							lowerCamelCase      =   pipe_a(
							    prompt_embeds=_a						,					negative_prompt_embeds=_a						,					image=_a						,					generator=_a						,					num_inference_steps=2						,					output_type="""np"""						,					)

							lowerCamelCase      =   output.images[0]

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

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

							lowerCamelCase      =   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   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							# pipeline 1

							_start_torch_memory_measurement()

							lowerCamelCase      =   floats_tensor((1, 3, 64, 64)						,					rng=random.Random(0					)					).to(_a					)

							lowerCamelCase      =   torch.Generator(device="""cpu"""					).manual_seed(0					)

							lowerCamelCase      =   pipe_a(
							    prompt_embeds=_a						,					negative_prompt_embeds=_a						,					image=_a						,					num_inference_steps=2						,					generator=_a						,					output_type="""np"""						,					)

							lowerCamelCase      =   output.images[0]

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

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

							lowerCamelCase      =   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      =   torch.Generator(device="""cpu"""					).manual_seed(0					)

							lowerCamelCase      =   floats_tensor((1, 3, 256, 256)						,					rng=random.Random(0					)					).to(_a					)
							lowerCamelCase      =   floats_tensor((1, 3, 64, 64)						,					rng=random.Random(0					)					).to(_a					)

							lowerCamelCase      =   pipe_a(
							    prompt_embeds=_a						,					negative_prompt_embeds=_a						,					image=_a						,					original_image=_a						,					generator=_a						,					num_inference_steps=2						,					output_type="""np"""						,					)

							lowerCamelCase      =   output.images[0]

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

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

							lowerCamelCase      =   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   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							# pipeline 1

							_start_torch_memory_measurement()

							lowerCamelCase      =   floats_tensor((1, 3, 64, 64)						,					rng=random.Random(0					)					).to(_a					)
							lowerCamelCase      =   floats_tensor((1, 3, 64, 64)						,					rng=random.Random(1					)					).to(_a					)

							lowerCamelCase      =   torch.Generator(device="""cpu"""					).manual_seed(0					)
							lowerCamelCase      =   pipe_a(
							    prompt_embeds=_a						,					negative_prompt_embeds=_a						,					image=_a						,					mask_image=_a						,					num_inference_steps=2						,					generator=_a						,					output_type="""np"""						,					)

							lowerCamelCase      =   output.images[0]

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

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

							lowerCamelCase      =   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      =   torch.Generator(device="""cpu"""					).manual_seed(0					)

							lowerCamelCase      =   floats_tensor((1, 3, 64, 64)						,					rng=random.Random(0					)					).to(_a					)
							lowerCamelCase      =   floats_tensor((1, 3, 256, 256)						,					rng=random.Random(0					)					).to(_a					)
							lowerCamelCase      =   floats_tensor((1, 3, 256, 256)						,					rng=random.Random(1					)					).to(_a					)

							lowerCamelCase      =   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      =   output.images[0]

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

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

							lowerCamelCase      =   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        a__					(							)		->   List[str]:
			torch.cuda.empty_cache()
			torch.cuda.reset_max_memory_allocated()
			torch.cuda.reset_peak_memory_stats()









"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowerCAmelCase				:     Optional[int]					=				{"""configuration_swin""": ["""SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP""", """SwinConfig""", """SwinOnnxConfig"""]}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCAmelCase				:     int					=				[
					    """SWIN_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """SwinForImageClassification""",
					    """SwinForMaskedImageModeling""",
					    """SwinModel""",
					    """SwinPreTrainedModel""",
					    """SwinBackbone""",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCAmelCase				:     str					=				[
					    """TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """TFSwinForImageClassification""",
					    """TFSwinForMaskedImageModeling""",
					    """TFSwinModel""",
					    """TFSwinPreTrainedModel""",
					]

if TYPE_CHECKING:
					from .configuration_swin import SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinConfig, SwinOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_swin import (
										    SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
										    SwinBackbone,
										    SwinForImageClassification,
										    SwinForMaskedImageModeling,
										    SwinModel,
										    SwinPreTrainedModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_swin import (
										    TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFSwinForImageClassification,
										    TFSwinForMaskedImageModeling,
										    TFSwinModel,
										    TFSwinPreTrainedModel,
										)

else:
					import sys

					lowerCAmelCase				:     int					=				_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



