





"""simple docstring"""
import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class  _SCREAMING_SNAKE_CASE(       unittest.TestCase  ):




       def    _UpperCamelCase					(						self							)      ->				List[Any]:





             """simple docstring"""






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




       @property
       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






             torch.manual_seed(0							)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  UNetaDModel(
                 sample_size=(32, 64)		,in_channels=1		,out_channels=1		,layers_per_block=2		,block_out_channels=(1_28, 1_28)		,down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')		,up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')		,)
             return model




       @property
       def    _UpperCamelCase					(						self							)      ->				Union[str, Any]:





             """simple docstring"""






             torch.manual_seed(0							)
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  UNetaDConditionModel(
                 sample_size=(64, 32)		,in_channels=1		,out_channels=1		,layers_per_block=2		,block_out_channels=(1_28, 1_28)		,down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')		,up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')		,cross_attention_dim=10		,)
             return model




       @property
       def    _UpperCamelCase					(						self							)      ->				int:





             """simple docstring"""






             torch.manual_seed(0							)
             __SCREAMING_SNAKE_CASE       :List[str]		    =  AutoencoderKL(
                 sample_size=(1_28, 64)		,in_channels=1		,out_channels=1		,latent_channels=1		,layers_per_block=2		,block_out_channels=(1_28, 1_28)		,down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D''')		,up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D''')		,)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  UNetaDModel(
                 sample_size=(64, 32)		,in_channels=1		,out_channels=1		,layers_per_block=2		,block_out_channels=(1_28, 1_28)		,down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')		,up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')		,)
             return vqvae, unet




       @slow
       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
             __SCREAMING_SNAKE_CASE       :List[str]		    =  Mel(
                 x_res=self.dummy_unet.config.sample_size[1]		,y_res=self.dummy_unet.config.sample_size[0]		,)

             __SCREAMING_SNAKE_CASE       :Tuple		    =  DDPMScheduler()
             __SCREAMING_SNAKE_CASE       :Tuple		    =  AudioDiffusionPipeline(vqvae=SCREAMING_SNAKE_CASE__		,unet=self.dummy_unet		,mel=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE       :Any		    =  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(42							)
             __SCREAMING_SNAKE_CASE       :Any		    =  pipe(generator=SCREAMING_SNAKE_CASE__		,steps=4							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  output.audios[0]
             __SCREAMING_SNAKE_CASE       :int		    =  output.images[0]

             __SCREAMING_SNAKE_CASE       :Dict		    =  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(42							)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  pipe(generator=SCREAMING_SNAKE_CASE__		,steps=4		,return_dict=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  output[0][0]

             assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
             assert (
                 image.height == self.dummy_unet.config.sample_size[0]
                 and image.width == self.dummy_unet.config.sample_size[1]
             )
             __SCREAMING_SNAKE_CASE       :Any		    =  np.frombuffer(image.tobytes()		,dtype='''uint8'''							)[:10]
             __SCREAMING_SNAKE_CASE       :Any		    =  np.frombuffer(image_from_tuple.tobytes()		,dtype='''uint8'''							)[:10]
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() == 0
             assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() == 0

             __SCREAMING_SNAKE_CASE       :Dict		    =  Mel(
                 x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]		,y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]		,)

             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  DDIMScheduler()
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  self.dummy_vqvae_and_unet
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  AudioDiffusionPipeline(
                 vqvae=self.dummy_vqvae_and_unet[0]		,unet=dummy_vqvae_and_unet[1]		,mel=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Any		    =  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             np.random.seed(0							)
             __SCREAMING_SNAKE_CASE       :Any		    =  np.random.uniform(-1		,1		,((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)							)
             __SCREAMING_SNAKE_CASE       :Dict		    =  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(42							)
             __SCREAMING_SNAKE_CASE       :List[str]		    =  pipe(raw_audio=SCREAMING_SNAKE_CASE__		,generator=SCREAMING_SNAKE_CASE__		,start_step=5		,steps=10							)
             __SCREAMING_SNAKE_CASE       :Dict		    =  output.images[0]

             assert (
                 image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
                 and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
             )
             __SCREAMING_SNAKE_CASE       :int		    =  np.frombuffer(image.tobytes()		,dtype='''uint8'''							)[:10]
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() == 0

             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  self.dummy_unet_condition
             __SCREAMING_SNAKE_CASE       :List[str]		    =  AudioDiffusionPipeline(
                 vqvae=self.dummy_vqvae_and_unet[0]		,unet=SCREAMING_SNAKE_CASE__		,mel=SCREAMING_SNAKE_CASE__		,scheduler=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Dict		    =  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             np.random.seed(0							)
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  torch.rand((1, 1, 10)							)
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  pipe(generator=SCREAMING_SNAKE_CASE__		,encoding=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :List[str]		    =  output.images[0]
             __SCREAMING_SNAKE_CASE       :List[str]		    =  np.frombuffer(image.tobytes()		,dtype='''uint8'''							)[:10]
             __SCREAMING_SNAKE_CASE       :Any		    =  np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() == 0






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




       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :str		    =  torch_device

             __SCREAMING_SNAKE_CASE       :str		    =  DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''							)
             __SCREAMING_SNAKE_CASE       :List[str]		    =  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE       :int		    =  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(42							)
             __SCREAMING_SNAKE_CASE       :str		    =  pipe(generator=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Dict		    =  output.audios[0]
             __SCREAMING_SNAKE_CASE       :List[str]		    =  output.images[0]

             assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
             assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
             __SCREAMING_SNAKE_CASE       :int		    =  np.frombuffer(image.tobytes()		,dtype='''uint8'''							)[:10]
             __SCREAMING_SNAKE_CASE       :Any		    =  np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26]							)

             assert np.abs(image_slice.flatten() - expected_slice							).max() == 0






"""simple docstring"""
import numpy as np
import torch
import tqdm

from ...models.unet_ad import UNetaDModel
from ...pipelines import DiffusionPipeline
from ...utils import randn_tensor
from ...utils.dummy_pt_objects import DDPMScheduler






class  _SCREAMING_SNAKE_CASE(       A  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,)      ->				Optional[int]:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE       :Any		    =  value_function
             __SCREAMING_SNAKE_CASE       :List[str]		    =  unet
             __SCREAMING_SNAKE_CASE       :int		    =  scheduler
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  env
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  env.get_dataset()
             __SCREAMING_SNAKE_CASE       :str		    =  {}
             for key in self.data.keys():
                   try:
                         __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  self.data[key].mean()
                   except:  # noqa: E722
                         pass
             __SCREAMING_SNAKE_CASE       :Dict		    =  {}
             for key in self.data.keys():
                   try:
                         __SCREAMING_SNAKE_CASE       :str		    =  self.data[key].std()
                   except:  # noqa: E722
                         pass
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  env.observation_space.shape[0]
             __SCREAMING_SNAKE_CASE       :int		    =  env.action_space.shape[0]




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)      ->				str:





             """simple docstring"""






             return (x_in - self.means[key]) / self.stds[key]




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)      ->				Dict:





             """simple docstring"""






             return x_in * self.stds[key] + self.means[key]




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





             """simple docstring"""






             if type(SCREAMING_SNAKE_CASE__							) is dict:
                   return {k: self.to_torch(SCREAMING_SNAKE_CASE__							) for k, v in x_in.items()}
             elif torch.is_tensor(SCREAMING_SNAKE_CASE__							):
                   return x_in.to(self.unet.device							)
             return torch.tensor(SCREAMING_SNAKE_CASE__		,device=self.unet.device							)




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





             """simple docstring"""






             for key, val in cond.items():
                   __SCREAMING_SNAKE_CASE       :Dict		    =  val.clone()
             return x_in




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)      ->				int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  x.shape[0]
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  None
             for i in tqdm.tqdm(self.scheduler.timesteps							):
                   # create batch of timesteps to pass into model
                   __SCREAMING_SNAKE_CASE       :Tuple		    =  torch.full((batch_size,)		,SCREAMING_SNAKE_CASE__		,device=self.unet.device		,dtype=torch.long							)
                   for _ in range(SCREAMING_SNAKE_CASE__							):
                         with torch.enable_grad():
                               x.requires_grad_()

                               # permute to match dimension for pre-trained models
                               __SCREAMING_SNAKE_CASE       :str		    =  self.value_function(x.permute(0		,2		,1							)		,SCREAMING_SNAKE_CASE__							).sample
                               __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  torch.autograd.grad([y.sum()]		,[x]							)[0]

                               __SCREAMING_SNAKE_CASE       :int		    =  self.scheduler._get_variance(SCREAMING_SNAKE_CASE__							)
                               __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  torch.exp(0.5 * posterior_variance							)
                               __SCREAMING_SNAKE_CASE       :List[str]		    =  model_std * grad

                         __SCREAMING_SNAKE_CASE       :Dict		    =  0
                         __SCREAMING_SNAKE_CASE       :List[Any]		    =  x.detach()
                         __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  x + scale * grad
                         __SCREAMING_SNAKE_CASE       :Any		    =  self.reset_xa(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,self.action_dim							)

                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.unet(x.permute(0		,2		,1							)		,SCREAMING_SNAKE_CASE__							).sample.permute(0		,2		,1							)

                   # TODO: verify deprecation of this kwarg
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.scheduler.step(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,predict_epsilon=SCREAMING_SNAKE_CASE__							)['''prev_sample''']

                   # apply conditions to the trajectory (set the initial state)
                   __SCREAMING_SNAKE_CASE       :List[str]		    =  self.reset_xa(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,self.action_dim							)
                   __SCREAMING_SNAKE_CASE       :Dict		    =  self.to_torch(SCREAMING_SNAKE_CASE__							)
             return x, y




       def __call__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=64		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=2		,SCREAMING_SNAKE_CASE__=0.1							)      ->				Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Tuple		    =  self.normalize(SCREAMING_SNAKE_CASE__		,'''observations'''							)
             __SCREAMING_SNAKE_CASE       :List[Any]		    =  obs[None].repeat(SCREAMING_SNAKE_CASE__		,axis=0							)

             __SCREAMING_SNAKE_CASE       :str		    =  {0: self.to_torch(SCREAMING_SNAKE_CASE__							)}
             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  (batch_size, planning_horizon, self.state_dim + self.action_dim)

             # generate initial noise and apply our conditions (to make the trajectories start at current state)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  randn_tensor(SCREAMING_SNAKE_CASE__		,device=self.unet.device							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  self.reset_xa(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,self.action_dim							)
             __SCREAMING_SNAKE_CASE       :Any		    =  self.to_torch(SCREAMING_SNAKE_CASE__							)

             # run the diffusion process
             __SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE       :Tuple		    =  self.run_diffusion(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)

             # sort output trajectories by value
             __SCREAMING_SNAKE_CASE       :Any		    =  y.argsort(0		,descending=SCREAMING_SNAKE_CASE__							).squeeze()
             __SCREAMING_SNAKE_CASE       :Any		    =  x[sorted_idx]
             __SCREAMING_SNAKE_CASE       :str		    =  sorted_values[:, :, : self.action_dim]
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  actions.detach().cpu().numpy()
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.de_normalize(SCREAMING_SNAKE_CASE__		,key='''actions'''							)

             # select the action with the highest value
             if y is not None:
                   __SCREAMING_SNAKE_CASE       :Optional[int]		    =  0
             else:
                   # if we didn't run value guiding, select a random action
                   __SCREAMING_SNAKE_CASE       :Any		    =  np.random.randint(0		,SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE       :int		    =  denorm_actions[selected_index, 0]
             return denorm_actions