import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 	lowerCamelCase_     (   UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	unittest.TestCase ):






    '''simple docstring'''






    a__	:						Tuple						=      StableDiffusionInstructPixaPixPipeline
    a__	:						Any						=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width""", """cross_attention_kwargs"""}
    a__	:						Any						=      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
    a__	:						Any						=      IMAGE_TO_IMAGE_IMAGE_PARAMS
    a__	:						Any						=      IMAGE_TO_IMAGE_IMAGE_PARAMS







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        torch.manual_seed(0)
        __UpperCamelCase     :Optional[Any]			  =		UNetaDConditionModel(
            block_out_channels=(32, 64) , layers_per_block=2 , sample_size=32 , in_channels=8 , out_channels=4 , down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''') , up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''') , cross_attention_dim=32 , )
        __UpperCamelCase     :Dict			  =		PNDMScheduler(skip_prk_steps=__lowercase)
        torch.manual_seed(0)
        __UpperCamelCase     :Tuple			  =		AutoencoderKL(
            block_out_channels=[32, 64] , in_channels=3 , out_channels=3 , down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''] , up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''] , latent_channels=4 , )
        torch.manual_seed(0)
        __UpperCamelCase     :Union[str, Any]			  =		CLIPTextConfig(
            bos_token_id=0 , eos_token_id=2 , hidden_size=32 , intermediate_size=37 , layer_norm_eps=1E-0_5 , num_attention_heads=4 , num_hidden_layers=5 , pad_token_id=1 , vocab_size=1_000 , )
        __UpperCamelCase     :List[Any]			  =		CLIPTextModel(__lowercase)
        __UpperCamelCase     :str			  =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

        __UpperCamelCase     :List[str]			  =		{
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase=0)       ->	Union[str, Any]:
        __UpperCamelCase     :Optional[Any]			  =		floats_tensor((1, 3, 32, 32) , rng=random.Random(__lowercase)).to(__lowercase)
        __UpperCamelCase     :Union[str, Any]			  =		image.cpu().permute(0 , 2 , 3 , 1)[0]
        __UpperCamelCase     :Tuple			  =		Image.fromarray(np.uinta(__lowercase)).convert('''RGB''')
        if str(__lowercase).startswith('''mps'''):
            __UpperCamelCase     :List[Any]			  =		torch.manual_seed(__lowercase)
        else:
            __UpperCamelCase     :Any			  =		torch.Generator(device=__lowercase).manual_seed(__lowercase)
        __UpperCamelCase     :Dict			  =		{
            '''prompt''': '''A painting of a squirrel eating a burger''',
            '''image''': image,
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''guidance_scale''': 6.0,
            '''image_guidance_scale''': 1,
            '''output_type''': '''numpy''',
        }
        return inputs







    def 			UpperCamelCase__   (	self)       ->	Dict:
        __UpperCamelCase     :List[str]			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Dict			  =		self.get_dummy_components()
        __UpperCamelCase     :List[Any]			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :List[str]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :List[str]			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :int			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :int			  =		image[0, -3:, -3:, -1]
        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :List[Any]			  =		np.array([0.75_26, 0.37_50, 0.45_47, 0.61_17, 0.58_66, 0.50_16, 0.43_27, 0.56_42, 0.48_15])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        __UpperCamelCase     :int			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Dict			  =		self.get_dummy_components()
        __UpperCamelCase     :int			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :List[Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :Any			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :int			  =		'''french fries'''
        __UpperCamelCase     :Tuple			  =		sd_pipe(**__lowercase , negative_prompt=__lowercase)
        __UpperCamelCase     :Dict			  =		output.images
        __UpperCamelCase     :Any			  =		image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :List[Any]			  =		np.array([0.75_11, 0.36_42, 0.45_53, 0.62_36, 0.57_97, 0.50_13, 0.43_43, 0.56_11, 0.48_31])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :int			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :str			  =		self.get_dummy_components()
        __UpperCamelCase     :str			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :List[Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :List[Any]			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :Tuple			  =		[inputs['''prompt''']] * 2

        __UpperCamelCase     :Dict			  =		np.array(inputs['''image''']).astype(np.floataa) / 2_55.0
        __UpperCamelCase     :Tuple			  =		torch.from_numpy(__lowercase).unsqueeze(0).to(__lowercase)
        __UpperCamelCase     :int			  =		image / 2 + 0.5
        __UpperCamelCase     :List[str]			  =		image.permute(0 , 3 , 1 , 2)
        __UpperCamelCase     :Dict			  =		image.repeat(2 , 1 , 1 , 1)

        __UpperCamelCase     :Tuple			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :Optional[int]			  =		image[-1, -3:, -3:, -1]

        assert image.shape == (2, 32, 32, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.58_12, 0.57_48, 0.52_22, 0.59_08, 0.56_95, 0.71_74, 0.68_04, 0.55_23, 0.55_79])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :List[Any]			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Union[str, Any]			  =		self.get_dummy_components()
        __UpperCamelCase     :Union[str, Any]			  =		EulerAncestralDiscreteScheduler(
            beta_start=0.0_00_85 , beta_end=0.0_12 , beta_schedule='''scaled_linear''')
        __UpperCamelCase     :Any			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :Union[str, Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :Optional[Any]			  =		self.get_dummy_inputs(__lowercase)
        __UpperCamelCase     :int			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :List[Any]			  =		image[0, -3:, -3:, -1]

        __UpperCamelCase     :str			  =		[round(__lowercase , 4) for x in image_slice.flatten().tolist()]
        print(''','''.join([str(__lowercase) for x in slice]))

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :str			  =		np.array([0.74_17, 0.38_42, 0.47_32, 0.57_76, 0.58_91, 0.51_39, 0.40_52, 0.56_73, 0.49_86])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :int			  =		self.get_dummy_components()
        __UpperCamelCase     :Any			  =		StableDiffusionInstructPixaPixPipeline(**__lowercase)
        __UpperCamelCase     :Union[str, Any]			  =		VaeImageProcessor(do_resize=__lowercase , do_normalize=__lowercase)
        __UpperCamelCase     :Dict			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :Any			  =		pipe(**self.get_dummy_inputs_by_type(__lowercase , input_image_type='''pt'''))[0]

        __UpperCamelCase     :int			  =		components['''vae''']
        __UpperCamelCase     :Dict			  =		self.get_dummy_inputs_by_type(__lowercase , input_image_type='''pt''')

        for image_param in self.image_latents_params:
            if image_param in inputs.keys():
                __UpperCamelCase     :str			  =		vae.encode(inputs[image_param]).latent_dist.mode()

        __UpperCamelCase     :Tuple			  =		pipe(**__lowercase)[0]

        __UpperCamelCase     :Optional[Any]			  =		np.abs(out - out_latents_inputs).max()
        self.assertLess(__lowercase , 1E-4 , '''passing latents as image input generate different result from passing image''')


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






    '''simple docstring'''






    def 			UpperCamelCase__   (	self)       ->	Any:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







    def 			UpperCamelCase__   (	self , __lowercase=0)       ->	List[Any]:
        __UpperCamelCase     :Optional[int]			  =		torch.manual_seed(__lowercase)
        __UpperCamelCase     :Optional[Any]			  =		load_image(
            '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg''')
        __UpperCamelCase     :Any			  =		{
            '''prompt''': '''turn him into a cyborg''',
            '''image''': image,
            '''generator''': generator,
            '''num_inference_steps''': 3,
            '''guidance_scale''': 7.5,
            '''image_guidance_scale''': 1.0,
            '''output_type''': '''numpy''',
        }
        return inputs







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :Any			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Union[str, Any]			  =		self.get_inputs()
        __UpperCamelCase     :Dict			  =		pipe(**__lowercase).images
        __UpperCamelCase     :Union[str, Any]			  =		image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase     :Any			  =		np.array([0.59_02, 0.60_15, 0.60_27, 0.59_83, 0.60_92, 0.60_61, 0.57_65, 0.57_85, 0.55_55])

        assert np.abs(expected_slice - image_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :int			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase)
        __UpperCamelCase     :Dict			  =		LMSDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Any			  =		self.get_inputs()
        __UpperCamelCase     :Union[str, Any]			  =		pipe(**__lowercase).images
        __UpperCamelCase     :int			  =		image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.65_78, 0.68_17, 0.69_72, 0.67_61, 0.68_56, 0.69_16, 0.64_28, 0.65_16, 0.63_01])

        assert np.abs(expected_slice - image_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Dict			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase)
        __UpperCamelCase     :List[Any]			  =		DDIMScheduler.from_config(pipe.scheduler.config)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Any			  =		self.get_inputs()
        __UpperCamelCase     :Dict			  =		pipe(**__lowercase).images
        __UpperCamelCase     :List[Any]			  =		image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.38_28, 0.38_34, 0.38_18, 0.37_92, 0.38_65, 0.37_52, 0.37_92, 0.38_47, 0.37_53])

        assert np.abs(expected_slice - image_slice).max() < 1E-3







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :Dict			  =		0

        def callback_fn(__lowercase , __lowercase , __lowercase) -> None:
            __UpperCamelCase     :Any			  =		True
            nonlocal number_of_steps
            number_of_steps += 1
            if step == 1:
                __UpperCamelCase     :str			  =		latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                __UpperCamelCase     :Optional[int]			  =		latents[0, -3:, -3:, -1]
                __UpperCamelCase     :Tuple			  =		np.array([-0.24_63, -0.46_44, -0.97_56, 1.51_76, 1.44_14, 0.78_66, 0.98_97, 0.85_21, 0.79_83])

                assert np.abs(latents_slice.flatten() - expected_slice).max() < 5E-2
            elif step == 2:
                __UpperCamelCase     :Optional[Any]			  =		latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                __UpperCamelCase     :str			  =		latents[0, -3:, -3:, -1]
                __UpperCamelCase     :int			  =		np.array([-0.26_44, -0.46_26, -0.96_53, 1.51_76, 1.45_51, 0.76_86, 0.98_05, 0.84_52, 0.81_15])

                assert np.abs(latents_slice.flatten() - expected_slice).max() < 5E-2

        __UpperCamelCase     :Tuple			  =		False

        __UpperCamelCase     :Optional[Any]			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase , torch_dtype=torch.floataa)
        __UpperCamelCase     :List[str]			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :Dict			  =		self.get_inputs()
        pipe(**__lowercase , callback=__lowercase , callback_steps=1)
        assert callback_fn.has_been_called
        assert number_of_steps == 3







    def 			UpperCamelCase__   (	self)       ->	str:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        __UpperCamelCase     :Dict			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            '''timbrooks/instruct-pix2pix''' , safety_checker=__lowercase , torch_dtype=torch.floataa)
        __UpperCamelCase     :Tuple			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing(1)
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :Any			  =		self.get_inputs()
        __UpperCamelCase     :Tuple			  =		pipe(**__lowercase)

        __UpperCamelCase     :Union[str, Any]			  =		torch.cuda.max_memory_allocated()
        # make sure that less than 2.2 GB is allocated
        assert mem_bytes < 2.2 * 10**9







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Any			  =		self.get_inputs()
        # resize to resolution that is divisible by 8 but not 16 or 32
        __UpperCamelCase     :str			  =		inputs['''image'''].resize((504, 504))

        __UpperCamelCase     :str			  =		'''timbrooks/instruct-pix2pix'''
        __UpperCamelCase     :str			  =		StableDiffusionInstructPixaPixPipeline.from_pretrained(
            __lowercase , safety_checker=__lowercase , )
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()

        __UpperCamelCase     :int			  =		pipe(**__lowercase)
        __UpperCamelCase     :Tuple			  =		output.images[0]

        __UpperCamelCase     :Dict			  =		image[255:258, 383:386, -1]

        assert image.shape == (504, 504, 3)
        __UpperCamelCase     :List[str]			  =		np.array([0.27_26, 0.25_29, 0.26_64, 0.26_55, 0.26_41, 0.26_42, 0.25_91, 0.26_49, 0.25_90])

        assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-3






import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
    from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


__lowercase							 =					logging.get_logger(__name__)  # pylint: disable=invalid-name

__lowercase							 =					256


class 	lowerCamelCase_     (   UpperCAmelCase_ ):






    '''simple docstring'''






    a__	:						Tuple						=      ["""melgan"""]







    def __init__(	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , )       ->	None:
        super().__init__()

        # From MELGAN
        __UpperCamelCase     :int			  =		math.log(1E-5)  # Matches MelGAN training.
        __UpperCamelCase     :int			  =		4.0  # Largest value for most examples
        __UpperCamelCase     :str			  =		128

        self.register_modules(
            notes_encoder=__lowercase , continuous_encoder=__lowercase , decoder=__lowercase , scheduler=__lowercase , melgan=__lowercase , )







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase=(-1.0, 1.0) , __lowercase=False)       ->	Dict:
        __UpperCamelCase  ,       __UpperCamelCase     :str			  =		output_range
        if clip:
            __UpperCamelCase     :Union[str, Any]			  =		torch.clip(__lowercase , self.min_value , self.max_value)
        # Scale to [0, 1].
        __UpperCamelCase     :Union[str, Any]			  =		(features - self.min_value) / (self.max_value - self.min_value)
        # Scale to [min_out, max_out].
        return zero_one * (max_out - min_out) + min_out







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase=(-1.0, 1.0) , __lowercase=False)       ->	Optional[int]:
        __UpperCamelCase  ,       __UpperCamelCase     :int			  =		input_range
        __UpperCamelCase     :Optional[int]			  =		torch.clip(__lowercase , __lowercase , __lowercase) if clip else outputs
        # Scale to [0, 1].
        __UpperCamelCase     :List[str]			  =		(outputs - min_out) / (max_out - min_out)
        # Scale to [self.min_value, self.max_value].
        return zero_one * (self.max_value - self.min_value) + self.min_value







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase)       ->	List[Any]:
        __UpperCamelCase     :List[str]			  =		input_tokens > 0
        __UpperCamelCase  ,       __UpperCamelCase     :Union[str, Any]			  =		self.notes_encoder(
            encoder_input_tokens=__lowercase , encoder_inputs_mask=__lowercase)

        __UpperCamelCase  ,       __UpperCamelCase     :Union[str, Any]			  =		self.continuous_encoder(
            encoder_inputs=__lowercase , encoder_inputs_mask=__lowercase)

        return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase)       ->	str:
        __UpperCamelCase     :Optional[int]			  =		noise_time
        if not torch.is_tensor(__lowercase):
            __UpperCamelCase     :str			  =		torch.tensor([timesteps] , dtype=torch.long , device=input_tokens.device)
        elif torch.is_tensor(__lowercase) and len(timesteps.shape) == 0:
            __UpperCamelCase     :Dict			  =		timesteps[None].to(input_tokens.device)

        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        __UpperCamelCase     :List[str]			  =		timesteps * torch.ones(input_tokens.shape[0] , dtype=timesteps.dtype , device=timesteps.device)

        __UpperCamelCase     :Tuple			  =		self.decoder(
            encodings_and_masks=__lowercase , decoder_input_tokens=__lowercase , decoder_noise_time=__lowercase)
        return logits







    @torch.no_grad()
    def __call__(	self , __lowercase , __lowercase = None , __lowercase = 100 , __lowercase = True , __lowercase = "numpy" , __lowercase = None , __lowercase = 1 , )       ->	Union[AudioPipelineOutput, Tuple]:
        if (callback_steps is None) or (
            callback_steps is not None and (not isinstance(__lowercase , __lowercase) or callback_steps <= 0)
        ):
            raise ValueError(
                f"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
                f""" {type(__lowercase)}.""")

        __UpperCamelCase     :Union[str, Any]			  =		np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims] , dtype=np.floataa)
        __UpperCamelCase     :Union[str, Any]			  =		np.zeros([1, 0, self.n_dims] , np.floataa)
        __UpperCamelCase     :Union[str, Any]			  =		torch.ones((1, TARGET_FEATURE_LENGTH) , dtype=__lowercase , device=self.device)

        for i, encoder_input_tokens in enumerate(__lowercase):
            if i == 0:
                __UpperCamelCase     :int			  =		torch.from_numpy(pred_mel[:1].copy()).to(
                    device=self.device , dtype=self.decoder.dtype)
                # The first chunk has no previous context.
                __UpperCamelCase     :int			  =		torch.zeros((1, TARGET_FEATURE_LENGTH) , dtype=__lowercase , device=self.device)
            else:
                # The full song pipeline does not feed in a context feature, so the mask
                # will be all 0s after the feature converter. Because we know we're
                # feeding in a full context chunk from the previous prediction, set it
                # to all 1s.
                __UpperCamelCase     :Tuple			  =		ones

            __UpperCamelCase     :Optional[Any]			  =		self.scale_features(
                __lowercase , output_range=[-1.0, 1.0] , clip=__lowercase)

            __UpperCamelCase     :int			  =		self.encode(
                input_tokens=torch.IntTensor([encoder_input_tokens]).to(device=self.device) , continuous_inputs=__lowercase , continuous_mask=__lowercase , )

            # Sample encoder_continuous_inputs shaped gaussian noise to begin loop
            __UpperCamelCase     :int			  =		randn_tensor(
                shape=encoder_continuous_inputs.shape , generator=__lowercase , device=self.device , dtype=self.decoder.dtype , )

            # set step values
            self.scheduler.set_timesteps(__lowercase)

            # Denoising diffusion loop
            for j, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
                __UpperCamelCase     :Optional[int]			  =		self.decode(
                    encodings_and_masks=__lowercase , input_tokens=__lowercase , noise_time=t / self.scheduler.config.num_train_timesteps , )

                # Compute previous output: x_t -> x_t-1
                __UpperCamelCase     :int			  =		self.scheduler.step(__lowercase , __lowercase , __lowercase , generator=__lowercase).prev_sample

            __UpperCamelCase     :Tuple			  =		self.scale_to_features(__lowercase , input_range=[-1.0, 1.0])
            __UpperCamelCase     :List[Any]			  =		mel[:1]
            __UpperCamelCase     :Optional[Any]			  =		mel.cpu().float().numpy()

            __UpperCamelCase     :Any			  =		np.concatenate([full_pred_mel, pred_mel[:1]] , axis=1)

            # call the callback, if provided
            if callback is not None and i % callback_steps == 0:
                callback(__lowercase , __lowercase)

            logger.info('''Generated segment''' , __lowercase)

        if output_type == "numpy" and not is_onnx_available():
            raise ValueError(
                '''Cannot return output in \'np\' format if ONNX is not available. Make sure to have ONNX installed or set \'output_type\' to \'mel\'.''')
        elif output_type == "numpy" and self.melgan is None:
            raise ValueError(
                '''Cannot return output in \'np\' format if melgan component is not defined. Make sure to define `self.melgan` or set \'output_type\' to \'mel\'.''')

        if output_type == "numpy":
            __UpperCamelCase     :Optional[Any]			  =		self.melgan(input_features=full_pred_mel.astype(np.floataa))
        else:
            __UpperCamelCase     :List[str]			  =		full_pred_mel

        if not return_dict:
            return (output,)

        return AudioPipelineOutput(audios=__lowercase)





