



from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel



class 				__snake_case						(     UpperCamelCase_						):
							_a			      = ['''vqvae''']




							def __init__(       self     :       Tuple				,				A_     :       AutoencoderKL				,				A_     :       UNetaDConditionModel				,				A_     :       Mel				,				A_     :       Union[DDIMScheduler, DDPMScheduler]				,				):
									super().__init__()
									self.register_modules(unet=A_				,				scheduler=A_				,				mel=A_				,				vqvae=A_)




							def    UpperCAmelCase__			(       self     :       Any):
									return 5_0 if isinstance(self.scheduler				,				A_) else 1_0_0_0




							@torch.no_grad()
							def __call__(       self     :       Any				,				A_     :       int = 1				,				A_     :       str = None				,				A_     :       np.ndarray = None				,				A_     :       int = 0				,				A_     :       int = 0				,				A_     :       int = None				,				A_     :       torch.Generator = None				,				A_     :       float = 0				,				A_     :       float = 0				,				A_     :       torch.Generator = None				,				A_     :       float = 0				,				A_     :       torch.Tensor = None				,				A_     :       torch.Tensor = None				,				A_     :       int=True				,				):
									lowerCAmelCase_			:     Union[str, Any]									= steps or self.get_default_steps()
									self.scheduler.set_timesteps(A_)
									lowerCAmelCase_			:     int									= step_generator or generator
									# For backwards compatibility
									if type(self.unet.config.sample_size) == int:
											lowerCAmelCase_			:     Optional[int]									= (self.unet.config.sample_size, self.unet.config.sample_size)
									if noise is None:
											lowerCAmelCase_			:     int									= randn_tensor(
											    (
											        batch_size,
											        self.unet.config.in_channels,
											        self.unet.config.sample_size[0],
											        self.unet.config.sample_size[1],
											    )				,				generator=A_				,				device=self.device				,				)
									lowerCAmelCase_			:     Tuple									= noise
									lowerCAmelCase_			:     Union[str, Any]									= None

									if audio_file is not None or raw_audio is not None:
											self.mel.load_audio(A_				,				A_)
											lowerCAmelCase_			:     Optional[int]									= self.mel.audio_slice_to_image(A_)
											lowerCAmelCase_			:     int									= np.frombuffer(input_image.tobytes()				,				dtype='''uint8''').reshape(
											    (input_image.height, input_image.width))
											lowerCAmelCase_			:     str									= (input_image / 2_5_5) * 2 - 1
											lowerCAmelCase_			:     List[Any]									= torch.tensor(input_image[np.newaxis, :, :]				,				dtype=torch.float).to(self.device)

											if self.vqvae is not None:
													lowerCAmelCase_			:     Optional[Any]									= self.vqvae.encode(torch.unsqueeze(A_				,				0)).latent_dist.sample(
													    generator=A_)[0]
													lowerCAmelCase_			:     Optional[int]									= self.vqvae.config.scaling_factor * input_images

											if start_step > 0:
													lowerCAmelCase_			:     List[str]									= self.scheduler.add_noise(A_				,				A_				,				self.scheduler.timesteps[start_step - 1])

											lowerCAmelCase_			:     Union[str, Any]									= (
											    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
											)
											lowerCAmelCase_			:     List[Any]									= int(mask_start_secs * pixels_per_second)
											lowerCAmelCase_			:     Optional[Any]									= int(mask_end_secs * pixels_per_second)
											lowerCAmelCase_			:     str									= self.scheduler.add_noise(A_				,				A_				,				torch.tensor(self.scheduler.timesteps[start_step:]))

									for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:])):
											if isinstance(self.unet				,				A_):
													lowerCAmelCase_			:     Optional[Any]									= self.unet(A_				,				A_				,				A_)['''sample''']
											else:
													lowerCAmelCase_			:     Dict									= self.unet(A_				,				A_)['''sample''']

											if isinstance(self.scheduler				,				A_):
													lowerCAmelCase_			:     List[Any]									= self.scheduler.step(
													    model_output=A_				,				timestep=A_				,				sample=A_				,				eta=A_				,				generator=A_				,				)['''prev_sample''']
											else:
													lowerCAmelCase_			:     int									= self.scheduler.step(
													    model_output=A_				,				timestep=A_				,				sample=A_				,				generator=A_				,				)['''prev_sample''']

											if mask is not None:
													if mask_start > 0:
															lowerCAmelCase_			:     Optional[Any]									= mask[:, step, :, :mask_start]
													if mask_end > 0:
															lowerCAmelCase_			:     Any									= mask[:, step, :, -mask_end:]

									if self.vqvae is not None:
											# 0.18215 was scaling factor used in training to ensure unit variance
											lowerCAmelCase_			:     Tuple									= 1 / self.vqvae.config.scaling_factor * images
											lowerCAmelCase_			:     Tuple									= self.vqvae.decode(A_)['''sample''']

									lowerCAmelCase_			:     Optional[int]									= (images / 2 + 0.5).clamp(0				,				1)
									lowerCAmelCase_			:     List[str]									= images.cpu().permute(0				,				2				,				3				,				1).numpy()
									lowerCAmelCase_			:     Union[str, Any]									= (images * 2_5_5).round().astype('''uint8''')
									lowerCAmelCase_			:     Optional[Any]									= list(
									    (Image.fromarray(_[:, :, 0]) for _ in images)
									    if images.shape[3] == 1
									    else (Image.fromarray(A_				,				mode='''RGB''').convert('''L''') for _ in images))

									lowerCAmelCase_			:     Optional[Any]									= [self.mel.image_to_audio(A_) for _ in images]
									if not return_dict:
											return images, (self.mel.get_sample_rate(), audios)

									return BaseOutput(**AudioPipelineOutput(np.array(A_)[:, np.newaxis, :])				,				**ImagePipelineOutput(A_))




							@torch.no_grad()
							def    UpperCAmelCase__			(       self     :       Optional[int]				,				A_     :       List[Image.Image]				,				A_     :       int = 5_0):
									assert isinstance(self.scheduler				,				A_)
									self.scheduler.set_timesteps(A_)
									lowerCAmelCase_			:     int									= np.array(
									    [np.frombuffer(image.tobytes()				,				dtype='''uint8''').reshape((1, image.height, image.width)) for image in images])
									lowerCAmelCase_			:     Any									= (sample / 2_5_5) * 2 - 1
									lowerCAmelCase_			:     List[str]									= torch.Tensor(A_).to(self.device)

									for t in self.progress_bar(torch.flip(self.scheduler.timesteps				,				(0,))):
											lowerCAmelCase_			:     Optional[int]									= t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
											lowerCAmelCase_			:     Tuple									= self.scheduler.alphas_cumprod[t]
											lowerCAmelCase_			:     Optional[int]									= (
											    self.scheduler.alphas_cumprod[prev_timestep]
											    if prev_timestep >= 0
											    else self.scheduler.final_alpha_cumprod
											)
											lowerCAmelCase_			:     int									= 1 - alpha_prod_t
											lowerCAmelCase_			:     Optional[Any]									= self.unet(A_				,				A_)['''sample''']
											lowerCAmelCase_			:     Optional[Any]									= (1 - alpha_prod_t_prev) ** 0.5 * model_output
											lowerCAmelCase_			:     Tuple									= (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
											lowerCAmelCase_			:     List[str]									= sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

									return sample







							@staticmethod
							def    UpperCAmelCase__			(       A_     :       torch.Tensor				,				A_     :       torch.Tensor				,				A_     :       float):
									lowerCAmelCase_			:     Dict									= acos(torch.dot(torch.flatten(A_)				,				torch.flatten(A_)) / torch.norm(A_) / torch.norm(A_))
									return sin((1 - alpha) * theta) * xa / sin(A_) + sin(alpha * theta) * xa / sin(A_)










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						):
							_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     :       List[str]):
									return self._get_dummy_components()




							def    UpperCAmelCase__			(       self     :       List[str]				,				A_     :       List[Any]				,				A_     :       Any=0):
									if str(A_).startswith('''mps'''):
											lowerCAmelCase_			:     List[Any]									= torch.manual_seed(A_)
									else:
											lowerCAmelCase_			:     List[str]									= torch.Generator(device=A_).manual_seed(A_)

									lowerCAmelCase_			:     Optional[Any]									= {
									    '''prompt''': '''A painting of a squirrel eating a burger''',
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''output_type''': '''numpy''',
									}

									return inputs




							def    UpperCAmelCase__			(       self     :       int):
									self._test_save_load_optional_components()




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




							def    UpperCAmelCase__			(       self     :       int):
									self._test_save_load_local()




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





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




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




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

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

									lowerCAmelCase_			:     Dict									= 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_			:     List[str]									= pipe_a.encode_prompt('''anime turtle'''				,				device='''cuda''')

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

									lowerCAmelCase_			:     Tuple									= None
									lowerCAmelCase_			:     str									= 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_			:     List[str]									= IFImgaImgPipeline(**pipe_a.components)
									lowerCAmelCase_			:     Any									= 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_			:     int									= IFInpaintingPipeline(**pipe_a.components)
									lowerCAmelCase_			:     Union[str, Any]									= 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     :       int				,				A_     :       Optional[int]				,				A_     :       Any				,				A_     :       str				,				A_     :       Union[str, Any]):
									# pipeline 1

									_start_torch_memory_measurement()

									lowerCAmelCase_			:     Optional[Any]									= torch.Generator(device='''cpu''').manual_seed(0)
									lowerCAmelCase_			:     Optional[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 == (6_4, 6_4, 3)

									lowerCAmelCase_			:     Any									= torch.cuda.max_memory_allocated()
									assert mem_bytes < 1_3 * 1_0**9

									lowerCAmelCase_			:     List[str]									= 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_			:     List[str]									= torch.Generator(device='''cpu''').manual_seed(0)

									lowerCAmelCase_			:     str									= floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(0)).to(A_)

									lowerCAmelCase_			:     Optional[int]									= pipe_a(
									    prompt_embeds=A_				,				negative_prompt_embeds=A_				,				image=A_				,				generator=A_				,				num_inference_steps=2				,				output_type='''np'''				,				)

									lowerCAmelCase_			:     Union[str, Any]									= output.images[0]

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

									lowerCAmelCase_			:     Any									= torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 1_0**9

									lowerCAmelCase_			:     int									= 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     :       int				,				A_     :       Optional[int]				,				A_     :       Any				,				A_     :       List[str]				,				A_     :       List[str]):
									# pipeline 1

									_start_torch_memory_measurement()

									lowerCAmelCase_			:     Any									= floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(0)).to(A_)

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

									lowerCAmelCase_			:     Any									= pipe_a(
									    prompt_embeds=A_				,				negative_prompt_embeds=A_				,				image=A_				,				num_inference_steps=2				,				generator=A_				,				output_type='''np'''				,				)

									lowerCAmelCase_			:     Union[str, Any]									= output.images[0]

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

									lowerCAmelCase_			:     str									= torch.cuda.max_memory_allocated()
									assert mem_bytes < 1_0 * 1_0**9

									lowerCAmelCase_			:     Optional[int]									= 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_			:     int									= torch.Generator(device='''cpu''').manual_seed(0)

									lowerCAmelCase_			:     List[Any]									= floats_tensor((1, 3, 2_5_6, 2_5_6)				,				rng=random.Random(0)).to(A_)
									lowerCAmelCase_			:     Any									= floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(0)).to(A_)

									lowerCAmelCase_			:     List[str]									= pipe_a(
									    prompt_embeds=A_				,				negative_prompt_embeds=A_				,				image=A_				,				original_image=A_				,				generator=A_				,				num_inference_steps=2				,				output_type='''np'''				,				)

									lowerCAmelCase_			:     Union[str, Any]									= output.images[0]

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

									lowerCAmelCase_			:     str									= torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 1_0**9

									lowerCAmelCase_			:     Tuple									= 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     :       str				,				A_     :       Optional[Any]				,				A_     :       Optional[Any]				,				A_     :       Dict				,				A_     :       List[str]):
									# pipeline 1

									_start_torch_memory_measurement()

									lowerCAmelCase_			:     Union[str, Any]									= floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(0)).to(A_)
									lowerCAmelCase_			:     int									= floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(1)).to(A_)

									lowerCAmelCase_			:     Optional[Any]									= torch.Generator(device='''cpu''').manual_seed(0)
									lowerCAmelCase_			:     Any									= pipe_a(
									    prompt_embeds=A_				,				negative_prompt_embeds=A_				,				image=A_				,				mask_image=A_				,				num_inference_steps=2				,				generator=A_				,				output_type='''np'''				,				)

									lowerCAmelCase_			:     List[Any]									= output.images[0]

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

									lowerCAmelCase_			:     str									= torch.cuda.max_memory_allocated()
									assert mem_bytes < 1_0 * 1_0**9

									lowerCAmelCase_			:     int									= 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[Any]									= torch.Generator(device='''cpu''').manual_seed(0)

									lowerCAmelCase_			:     Optional[Any]									= floats_tensor((1, 3, 6_4, 6_4)				,				rng=random.Random(0)).to(A_)
									lowerCAmelCase_			:     Union[str, Any]									= floats_tensor((1, 3, 2_5_6, 2_5_6)				,				rng=random.Random(0)).to(A_)
									lowerCAmelCase_			:     Optional[Any]									= floats_tensor((1, 3, 2_5_6, 2_5_6)				,				rng=random.Random(1)).to(A_)

									lowerCAmelCase_			:     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_			:     Tuple									= output.images[0]

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

									lowerCAmelCase_			:     int									= torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 1_0**9

									lowerCAmelCase_			:     Any									= 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()





