




'''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  __lowercase    (	a__							,				a__							,				unittest.TestCase			):
   _a    						=   IFPipeline
   _a    						=   TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
   _a    						=   TEXT_TO_IMAGE_BATCH_PARAMS
   _a    						=   PipelineTesterMixin.required_optional_params - {"""latents"""}







   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     return self._get_dummy_components()







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase=0   )					->     Tuple:
     if str(__a   ).startswith('mps'   ):
       __a								=       torch.manual_seed(__a   )
     else:
       __a								=       torch.Generator(device=__a   ).manual_seed(__a   )

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

     return inputs







   def 		UpperCamelCase__      (						self   )					->     Any:
     self._test_save_load_optional_components()







   @unittest.skipIf(torch_device != 'cuda'   ,				reason='float16 requires CUDA'   )
   def 		UpperCamelCase__      (						self   )					->     Tuple:
     # 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 		UpperCamelCase__      (						self   )					->     Dict:
     self._test_attention_slicing_forward_pass(expected_max_diff=1e-2   )







   def 		UpperCamelCase__      (						self   )					->     Union[str, Any]:
     self._test_save_load_local()







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





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







   def 		UpperCamelCase__      (						self   )					->     Dict:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







   def 		UpperCamelCase__      (						self   )					->     Union[str, Any]:
     # if

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

     __a								=       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'   )

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

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

     __a								=       None
     __a								=       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

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

     __a								=       IFInpaintingPipeline(**pipe_a.components   )
     __a								=       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   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     str:
     # pipeline 1

     _start_torch_memory_measurement()

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

     __a								=       output.images[0]

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

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

     __a								=       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()

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

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

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

     __a								=       output.images[0]

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

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

     __a								=       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   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     int:
     # pipeline 1

     _start_torch_memory_measurement()

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

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

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

     __a								=       output.images[0]

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

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

     __a								=       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()

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

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

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

     __a								=       output.images[0]

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

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

     __a								=       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   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     List[Any]:
     # pipeline 1

     _start_torch_memory_measurement()

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

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

     __a								=       output.images[0]

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

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

     __a								=       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()

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

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

     __a								=       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'   ,				)

     __a								=       output.images[0]

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

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

     __a								=       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     SCREAMING_SNAKE_CASE							(   ):
  torch.cuda.empty_cache()
  torch.cuda.reset_max_memory_allocated()
  torch.cuda.reset_peak_memory_stats()









'''simple docstring'''







import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 					_a (a__,		unittest.TestCase	):





     '''simple docstring'''



     lowerCAmelCase_  :     Dict          =					DanceDiffusionPipeline
     lowerCAmelCase_  :     int          =					UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     lowerCAmelCase_  :     Optional[int]          =					PipelineTesterMixin.required_optional_params - {
         """callback""",
         """latents""",
         """callback_steps""",
         """output_type""",
         """num_images_per_prompt""",
     }
     lowerCAmelCase_  :     Tuple          =					UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     lowerCAmelCase_  :     Dict          =					False
     lowerCAmelCase_  :     Tuple          =					False







     def 			snake_case_      (				self							)      ->				List[Any]:
          torch.manual_seed(0							)
          snake_case			:      List[str]        =			UNetaDModel(
              block_out_channels=(32, 32, 64)							,extra_in_channels=16							,sample_size=512							,sample_rate=16_000							,in_channels=2							,out_channels=2							,flip_sin_to_cos=__a							,use_timestep_embedding=__a							,time_embedding_type="""fourier"""							,mid_block_type="""UNetMidBlock1D"""							,down_block_types=("""DownBlock1DNoSkip""", """DownBlock1D""", """AttnDownBlock1D""")							,up_block_types=("""AttnUpBlock1D""", """UpBlock1D""", """UpBlock1DNoSkip""")							,)
          snake_case			:      Dict        =			IPNDMScheduler()

          snake_case			:      int        =			{
              """unet""": unet,
              """scheduler""": scheduler,
          }
          return components







     def 			snake_case_      (				self							,__a							,__a=0							)      ->				Tuple:
          if str(__a							).startswith("""mps"""							):
               snake_case			:      Optional[Any]        =			torch.manual_seed(__a							)
          else:
               snake_case			:      Union[str, Any]        =			torch.Generator(device=__a							).manual_seed(__a							)
          snake_case			:      str        =			{
              """batch_size""": 1,
              """generator""": generator,
              """num_inference_steps""": 4,
          }
          return inputs







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      Dict        =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
          snake_case			:      str        =			self.get_dummy_components()
          snake_case			:      List[str]        =			DanceDiffusionPipeline(**__a							)
          snake_case			:      Dict        =			pipe.to(__a							)
          pipe.set_progress_bar_config(disable=__a							)

          snake_case			:      Union[str, Any]        =			self.get_dummy_inputs(__a							)
          snake_case			:      str        =			pipe(**__a							)
          snake_case			:      Union[str, Any]        =			output.audios

          snake_case			:      List[str]        =			audio[0, -3:, -3:]

          assert audio.shape == (1, 2, components["unet"].sample_size)
          snake_case			:      Optional[Any]        =			np.array([-0.7_265, 1.0_000, -0.8_388, 0.1_175, 0.9_498, -1.0_000]							)
          assert np.abs(audio_slice.flatten() - expected_slice							).max() < 1E-2







     @skip_mps
     def 			snake_case_      (				self							)      ->				List[str]:
          return super().test_save_load_local()







     @skip_mps
     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3							)







     @skip_mps
     def 			snake_case_      (				self							)      ->				Dict:
          return super().test_save_load_optional_components()







     @skip_mps
     def 			snake_case_      (				self							)      ->				Any:
          return super().test_attention_slicing_forward_pass()







     def 			snake_case_      (				self							)      ->				str:
          super().test_inference_batch_single_identical(expected_max_diff=3E-3							)







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





     '''simple docstring'''



     def 			snake_case_      (				self							)      ->				Optional[int]:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      str        =			torch_device

          snake_case			:      List[str]        =			DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""							)
          snake_case			:      Any        =			pipe.to(__a							)
          pipe.set_progress_bar_config(disable=__a							)

          snake_case			:      int        =			torch.manual_seed(0							)
          snake_case			:      int        =			pipe(generator=__a							,num_inference_steps=100							,audio_length_in_s=4.096							)
          snake_case			:      Dict        =			output.audios

          snake_case			:      Union[str, Any]        =			audio[0, -3:, -3:]

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          snake_case			:      Tuple        =			np.array([-0.0_192, -0.0_231, -0.0_318, -0.0_059, 0.0_002, -0.0_020]							)

          assert np.abs(audio_slice.flatten() - expected_slice							).max() < 1E-2







     def 			snake_case_      (				self							)      ->				str:
          snake_case			:      str        =			torch_device

          snake_case			:      List[Any]        =			DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""							,torch_dtype=torch.floataa							)
          snake_case			:      str        =			pipe.to(__a							)
          pipe.set_progress_bar_config(disable=__a							)

          snake_case			:      str        =			torch.manual_seed(0							)
          snake_case			:      List[str]        =			pipe(generator=__a							,num_inference_steps=100							,audio_length_in_s=4.096							)
          snake_case			:      Optional[int]        =			output.audios

          snake_case			:      Any        =			audio[0, -3:, -3:]

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          snake_case			:      Optional[int]        =			np.array([-0.0_367, -0.0_488, -0.0_771, -0.0_525, -0.0_444, -0.0_341]							)

          assert np.abs(audio_slice.flatten() - expected_slice							).max() < 1E-2



