


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 						lowerCamelCase_							(					lowerCAmelCase__  ,  lowerCAmelCase__  ,  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 A   (					self    ) ->      int:







								'''simple docstring'''

								return self._get_dummy_components()






					def A   (					self			,							snake_case_			,							snake_case_=0    ) ->      Optional[int]:







								'''simple docstring'''

								if str(snake_case_    ).startswith('''mps'''    ):
											__lowercase			       = torch.manual_seed(snake_case_    )
								else:
											__lowercase			       = torch.Generator(device=snake_case_    ).manual_seed(snake_case_    )

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

								return inputs






					def A   (					self    ) ->      Any:







								'''simple docstring'''

								self._test_save_load_optional_components()






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







								'''simple docstring'''

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






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







								'''simple docstring'''

								self._test_attention_slicing_forward_pass(expected_max_diff=1e-2    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								self._test_save_load_local()






					def A   (					self    ) ->      Union[str, 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 A   (					self    ) ->      Tuple:







								'''simple docstring'''

								self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3    )




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



					'''simple docstring'''







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







								'''simple docstring'''

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






					def A   (					self    ) ->      Tuple:







								'''simple docstring'''

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

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

								# pre compute text embeddings and remove T5 to save memory

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

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

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

								__lowercase			       = None
								__lowercase			       = 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(snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    )

								pipe_a.remove_all_hooks()
								pipe_a.remove_all_hooks()

								# img2img

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

								pipe_a.remove_all_hooks()
								pipe_a.remove_all_hooks()

								# inpainting

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[Any]:







								'''simple docstring'''

								_start_torch_memory_measurement()

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

								__lowercase			       = output.images[0]

								assert image.shape == (6_4, 6_4, 3)

								__lowercase			       = torch.cuda.max_memory_allocated()
								assert mem_bytes < 1_3 * 1_0**9

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

								# pipeline 2

								_start_torch_memory_measurement()

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

								__lowercase			       = floats_tensor((1, 3, 6_4, 6_4)			,							rng=random.Random(0    )    ).to(snake_case_    )

								__lowercase			       = pipe_a(
								    prompt_embeds=snake_case_			,							negative_prompt_embeds=snake_case_			,							image=snake_case_			,							generator=snake_case_			,							num_inference_steps=2			,							output_type='''np'''			,							)

								__lowercase			       = output.images[0]

								assert image.shape == (2_5_6, 2_5_6, 3)

								__lowercase			       = torch.cuda.max_memory_allocated()
								assert mem_bytes < 4 * 1_0**9

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

								_start_torch_memory_measurement()

								__lowercase			       = floats_tensor((1, 3, 6_4, 6_4)			,							rng=random.Random(0    )    ).to(snake_case_    )

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

								__lowercase			       = pipe_a(
								    prompt_embeds=snake_case_			,							negative_prompt_embeds=snake_case_			,							image=snake_case_			,							num_inference_steps=2			,							generator=snake_case_			,							output_type='''np'''			,							)

								__lowercase			       = output.images[0]

								assert image.shape == (6_4, 6_4, 3)

								__lowercase			       = torch.cuda.max_memory_allocated()
								assert mem_bytes < 1_0 * 1_0**9

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

								# pipeline 2

								_start_torch_memory_measurement()

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

								__lowercase			       = floats_tensor((1, 3, 2_5_6, 2_5_6)			,							rng=random.Random(0    )    ).to(snake_case_    )
								__lowercase			       = floats_tensor((1, 3, 6_4, 6_4)			,							rng=random.Random(0    )    ).to(snake_case_    )

								__lowercase			       = pipe_a(
								    prompt_embeds=snake_case_			,							negative_prompt_embeds=snake_case_			,							image=snake_case_			,							original_image=snake_case_			,							generator=snake_case_			,							num_inference_steps=2			,							output_type='''np'''			,							)

								__lowercase			       = output.images[0]

								assert image.shape == (2_5_6, 2_5_6, 3)

								__lowercase			       = torch.cuda.max_memory_allocated()
								assert mem_bytes < 4 * 1_0**9

								__lowercase			       = 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(snake_case_			,							snake_case_    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

								_start_torch_memory_measurement()

								__lowercase			       = floats_tensor((1, 3, 6_4, 6_4)			,							rng=random.Random(0    )    ).to(snake_case_    )
								__lowercase			       = floats_tensor((1, 3, 6_4, 6_4)			,							rng=random.Random(1    )    ).to(snake_case_    )

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

								__lowercase			       = output.images[0]

								assert image.shape == (6_4, 6_4, 3)

								__lowercase			       = torch.cuda.max_memory_allocated()
								assert mem_bytes < 1_0 * 1_0**9

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

								# pipeline 2

								_start_torch_memory_measurement()

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

								__lowercase			       = floats_tensor((1, 3, 6_4, 6_4)			,							rng=random.Random(0    )    ).to(snake_case_    )
								__lowercase			       = floats_tensor((1, 3, 2_5_6, 2_5_6)			,							rng=random.Random(0    )    ).to(snake_case_    )
								__lowercase			       = floats_tensor((1, 3, 2_5_6, 2_5_6)			,							rng=random.Random(1    )    ).to(snake_case_    )

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

								__lowercase			       = output.images[0]

								assert image.shape == (2_5_6, 2_5_6, 3)

								__lowercase			       = torch.cuda.max_memory_allocated()
								assert mem_bytes < 4 * 1_0**9

								__lowercase			       = 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(snake_case_			,							snake_case_    )




def 		lowercase_	(  ):


			'''simple docstring'''







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






from typing import Callable, Optional

from .. import Features
from ..packaged_modules.generator.generator import Generator
from .abc import AbstractDatasetInputStream







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			def __init__(				self   ,    snake_case   ,    snake_case = None   ,    snake_case = None   ,    snake_case = False   ,    snake_case = False   ,    snake_case = None   ,    snake_case = None   ,    **snake_case   ,    ):

						'''simple docstring'''


						super().__init__(
						    features=snake_case   ,    cache_dir=snake_case   ,    keep_in_memory=snake_case   ,    streaming=snake_case   ,    num_proc=snake_case   ,    **snake_case   ,    )
						UpperCamelCase__       =							Generator(
						    cache_dir=snake_case   ,    features=snake_case   ,    generator=snake_case   ,    gen_kwargs=snake_case   ,    **snake_case   ,    )







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						if self.streaming:
									UpperCamelCase__       =							self.builder.as_streaming_dataset(split="train"  )
						# Build regular (map-style) dataset
						else:
									UpperCamelCase__       =							None
									UpperCamelCase__       =							None
									UpperCamelCase__       =							None
									UpperCamelCase__       =							None

									self.builder.download_and_prepare(
									    download_config=snake_case   ,    download_mode=snake_case   ,    verification_mode=snake_case   ,    base_path=snake_case   ,    num_proc=self.num_proc   ,    )
									UpperCamelCase__       =							self.builder.as_dataset(
									    split="train"   ,    verification_mode=snake_case   ,    in_memory=self.keep_in_memory  )
						return dataset


