






import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
					import jax
					import jax.numpy as jnp

					from transformers.generation import (
					    FlaxForcedBOSTokenLogitsProcessor,
					    FlaxForcedEOSTokenLogitsProcessor,
					    FlaxLogitsProcessorList,
					    FlaxMinLengthLogitsProcessor,
					    FlaxTemperatureLogitsWarper,
					    FlaxTopKLogitsWarper,
					    FlaxTopPLogitsWarper,
					)
@require_flax
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





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

									'''simple docstring'''







									UpperCamelCase__         = jnp.ones((batch_size, length)			) / length
									return scores






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

									'''simple docstring'''







									UpperCamelCase__         = None
									UpperCamelCase__         = 2_0

									UpperCamelCase__         = self._get_uniform_logits(batch_size=2 ,      length=snake_case_			)

									# tweak scores to not be uniform anymore
									UpperCamelCase__         = scores.at[1, 5].set((1 / length) + 0.1			)  # peak, 1st batch
									UpperCamelCase__         = scores.at[1, 1_0].set((1 / length) - 0.4			)  # valley, 1st batch

									# compute softmax
									UpperCamelCase__         = jax.nn.softmax(snake_case_ ,      axis=-1			)

									UpperCamelCase__         = FlaxTemperatureLogitsWarper(temperature=0.5			)
									UpperCamelCase__         = FlaxTemperatureLogitsWarper(temperature=1.3			)

									UpperCamelCase__         = jax.nn.softmax(temp_dist_warper_sharper(snake_case_ ,      scores.copy() ,      cur_len=snake_case_			) ,      axis=-1			)
									UpperCamelCase__         = jax.nn.softmax(temp_dist_warper_smoother(snake_case_ ,      scores.copy() ,      cur_len=snake_case_			) ,      axis=-1			)

									# uniform distribution stays uniform
									self.assertTrue(jnp.allclose(probs[0, :] ,      warped_prob_sharp[0, :] ,      atol=1e-3			)			)
									self.assertTrue(jnp.allclose(probs[0, :] ,      warped_prob_smooth[0, :] ,      atol=1e-3			)			)

									# sharp peaks get higher, valleys get lower
									self.assertLess(probs[1, :].max() ,      warped_prob_sharp[1, :].max()			)
									self.assertGreater(probs[1, :].min() ,      warped_prob_sharp[1, :].min()			)

									# smooth peaks get lower, valleys get higher
									self.assertGreater(probs[1, :].max() ,      warped_prob_smooth[1, :].max()			)
									self.assertLess(probs[1, :].min() ,      warped_prob_smooth[1, :].min()			)






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

									'''simple docstring'''







									UpperCamelCase__         = None
									UpperCamelCase__         = 1_0
									UpperCamelCase__         = 2

									# create ramp distribution
									UpperCamelCase__         = np.broadcast_to(np.arange(snake_case_			)[None, :] ,      (batch_size, vocab_size)			).copy()
									UpperCamelCase__         = ramp_logits[1:, : vocab_size // 2] + vocab_size

									UpperCamelCase__         = FlaxTopKLogitsWarper(3			)

									UpperCamelCase__         = top_k_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)

									# check that correct tokens are filtered
									self.assertListEqual(jnp.isinf(scores[0]			).tolist() ,      7 * [True] + 3 * [False]			)
									self.assertListEqual(jnp.isinf(scores[1]			).tolist() ,      2 * [True] + 3 * [False] + 5 * [True]			)

									# check special case
									UpperCamelCase__         = 5
									UpperCamelCase__         = FlaxTopKLogitsWarper(top_k=1 ,      filter_value=0.0 ,      min_tokens_to_keep=3			)

									UpperCamelCase__         = np.broadcast_to(np.arange(snake_case_			)[None, :] ,      (batch_size, length)			).copy()
									UpperCamelCase__         = top_k_warp_safety_check(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)

									# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
									self.assertListEqual((scores == 0.0).sum(axis=-1			).tolist() ,      [2, 2]			)






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

									'''simple docstring'''







									UpperCamelCase__         = None
									UpperCamelCase__         = 1_0
									UpperCamelCase__         = 2

									# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
									UpperCamelCase__         = np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]			)			)

									UpperCamelCase__         = FlaxTopPLogitsWarper(0.8			)
									UpperCamelCase__         = np.exp(top_p_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)			)

									# dist should be filtered to keep min num values so that sum is >= top_p
									# exp (-inf) => 0
									UpperCamelCase__         = np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]			)
									self.assertTrue(np.allclose(snake_case_ ,      snake_case_ ,      atol=1e-3			)			)

									# check edge cases with negative and extreme logits
									UpperCamelCase__         = np.broadcast_to(np.arange(snake_case_			)[None, :] ,      (batch_size, vocab_size)			).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									UpperCamelCase__         = ramp_logits[1] * 1_0_0.0

									# make sure at least 2 tokens are kept
									UpperCamelCase__         = FlaxTopPLogitsWarper(0.9 ,      min_tokens_to_keep=2 ,      filter_value=0.0			)
									UpperCamelCase__         = top_p_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)

									# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
									self.assertListEqual((filtered_dist != 0.0).sum(axis=-1			).tolist() ,      [3, 2]			)






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

									'''simple docstring'''







									UpperCamelCase__         = 2_0
									UpperCamelCase__         = 4
									UpperCamelCase__         = 0

									UpperCamelCase__         = FlaxMinLengthLogitsProcessor(min_length=1_0 ,      eos_token_id=snake_case_			)

									# check that min length is applied at length 5
									UpperCamelCase__         = ids_tensor((batch_size, 2_0) ,      vocab_size=2_0			)
									UpperCamelCase__         = 5
									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = min_dist_processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist() ,      4 * [-float("""inf"""			)]			)

									# check that min length is not applied anymore at length 15
									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = 1_5
									UpperCamelCase__         = min_dist_processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									self.assertFalse(jnp.isinf(snake_case_			).any()			)






			def 				A      (   self							:						List[str]			)							->    Optional[Any]:

									'''simple docstring'''







									UpperCamelCase__         = 2_0
									UpperCamelCase__         = 4
									UpperCamelCase__         = 0

									UpperCamelCase__         = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case_			)

									# check that all scores are -inf except the bos_token_id score
									UpperCamelCase__         = ids_tensor((batch_size, 1) ,      vocab_size=2_0			)
									UpperCamelCase__         = 1
									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = logits_processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]			).all()			)
									self.assertListEqual(scores[:, bos_token_id].tolist() ,      4 * [0]			)  # score for bos_token_id shold be zero

									# check that bos_token_id is not forced if current length is greater than 1
									UpperCamelCase__         = 3
									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = logits_processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									self.assertFalse(jnp.isinf(snake_case_			).any()			)






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

									'''simple docstring'''







									UpperCamelCase__         = 2_0
									UpperCamelCase__         = 4
									UpperCamelCase__         = 0
									UpperCamelCase__         = 5

									UpperCamelCase__         = FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case_ ,      eos_token_id=snake_case_			)

									# check that all scores are -inf except the eos_token_id when max_length is reached
									UpperCamelCase__         = ids_tensor((batch_size, 4) ,      vocab_size=2_0			)
									UpperCamelCase__         = 4
									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = logits_processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]			).all()			)
									self.assertListEqual(scores[:, eos_token_id].tolist() ,      4 * [0]			)  # score for eos_token_id should be zero

									# check that eos_token_id is not forced if max_length is not reached
									UpperCamelCase__         = 3
									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = logits_processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									self.assertFalse(jnp.isinf(snake_case_			).any()			)






			def 				A      (   self							:						Dict			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = 4
									UpperCamelCase__         = 1_0
									UpperCamelCase__         = 1_5
									UpperCamelCase__         = 2
									UpperCamelCase__         = 1
									UpperCamelCase__         = 1_5

									# dummy input_ids and scores
									UpperCamelCase__         = ids_tensor((batch_size, sequence_length) ,      snake_case_			)
									UpperCamelCase__         = input_ids.copy()

									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = scores.copy()

									# instantiate all dist processors
									UpperCamelCase__         = FlaxTemperatureLogitsWarper(temperature=0.5			)
									UpperCamelCase__         = FlaxTopKLogitsWarper(3			)
									UpperCamelCase__         = FlaxTopPLogitsWarper(0.8			)

									# instantiate all logits processors
									UpperCamelCase__         = FlaxMinLengthLogitsProcessor(min_length=1_0 ,      eos_token_id=snake_case_			)
									UpperCamelCase__         = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case_			)
									UpperCamelCase__         = FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case_ ,      eos_token_id=snake_case_			)

									UpperCamelCase__         = 1_0

									# no processor list
									UpperCamelCase__         = temp_dist_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									UpperCamelCase__         = top_k_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									UpperCamelCase__         = top_p_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									UpperCamelCase__         = min_dist_proc(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									UpperCamelCase__         = bos_dist_proc(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
									UpperCamelCase__         = eos_dist_proc(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)

									# with processor list
									UpperCamelCase__         = FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
									UpperCamelCase__         = processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)

									# scores should be equal
									self.assertTrue(jnp.allclose(snake_case_ ,      snake_case_ ,      atol=1e-3			)			)

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist() ,      input_ids_comp.tolist()			)







			def 				A      (   self							:						Dict			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = 4
									UpperCamelCase__         = 1_0
									UpperCamelCase__         = 1_5
									UpperCamelCase__         = 2
									UpperCamelCase__         = 1
									UpperCamelCase__         = 1_5

									# dummy input_ids and scores
									UpperCamelCase__         = ids_tensor((batch_size, sequence_length) ,      snake_case_			)
									UpperCamelCase__         = input_ids.copy()

									UpperCamelCase__         = self._get_uniform_logits(snake_case_ ,      snake_case_			)
									UpperCamelCase__         = scores.copy()

									# instantiate all dist processors
									UpperCamelCase__         = FlaxTemperatureLogitsWarper(temperature=0.5			)
									UpperCamelCase__         = FlaxTopKLogitsWarper(3			)
									UpperCamelCase__         = FlaxTopPLogitsWarper(0.8			)

									# instantiate all logits processors
									UpperCamelCase__         = FlaxMinLengthLogitsProcessor(min_length=1_0 ,      eos_token_id=snake_case_			)
									UpperCamelCase__         = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case_			)
									UpperCamelCase__         = FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case_ ,      eos_token_id=snake_case_			)

									UpperCamelCase__         = 1_0

									# no processor list
									def run_no_processor_list(lowercase							:						str ,      lowercase							:						int ,      lowercase							:						List[Any]			):
															UpperCamelCase__         = temp_dist_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															UpperCamelCase__         = top_k_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															UpperCamelCase__         = top_p_warp(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															UpperCamelCase__         = min_dist_proc(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															UpperCamelCase__         = bos_dist_proc(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															UpperCamelCase__         = eos_dist_proc(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															return scores

									# with processor list
									def run_processor_list(lowercase							:						List[str] ,      lowercase							:						int ,      lowercase							:						List[Any]			):
															UpperCamelCase__         = FlaxLogitsProcessorList(
															    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
															UpperCamelCase__         = processor(snake_case_ ,      snake_case_ ,      cur_len=snake_case_			)
															return scores

									UpperCamelCase__         = jax.jit(snake_case_			)
									UpperCamelCase__         = jax.jit(snake_case_			)

									UpperCamelCase__         = jitted_run_no_processor_list(snake_case_ ,      snake_case_ ,      snake_case_			)
									UpperCamelCase__         = jitted_run_processor_list(snake_case_ ,      snake_case_ ,      snake_case_			)

									# scores should be equal
									self.assertTrue(jnp.allclose(snake_case_ ,      snake_case_ ,      atol=1e-3			)			)

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist() ,      input_ids_comp.tolist()			)








'''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	):


			'''simple docstring'''





			def 				A      (   self							:						Any			)							->    Any:

									'''simple docstring'''







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






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

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDModel(
									    sample_size=(3_2, 6_4) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""AttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """AttnUpBlock2D""") ,      )
									return model






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

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDConditionModel(
									    sample_size=(6_4, 3_2) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""") ,      cross_attention_dim=1_0 ,      )
									return model






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

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = AutoencoderKL(
									    sample_size=(1_2_8, 6_4) ,      in_channels=1 ,      out_channels=1 ,      latent_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""") ,      up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""") ,      )
									UpperCamelCase__         = UNetaDModel(
									    sample_size=(6_4, 3_2) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""AttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """AttnUpBlock2D""") ,      )
									return vqvae, unet







			@slow
			def 				A      (   self							:						Tuple			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = """cpu"""  # ensure determinism for the device-dependent torch.Generator
									UpperCamelCase__         = Mel(
									    x_res=self.dummy_unet.config.sample_size[1] ,      y_res=self.dummy_unet.config.sample_size[0] ,      )

									UpperCamelCase__         = DDPMScheduler()
									UpperCamelCase__         = AudioDiffusionPipeline(vqvae=lowercase ,      unet=self.dummy_unet ,      mel=lowercase ,      scheduler=lowercase			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase ,      steps=4			)
									UpperCamelCase__         = output.audios[0]
									UpperCamelCase__         = output.images[0]

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase ,      steps=4 ,      return_dict=lowercase			)
									UpperCamelCase__         = 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]
									)
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.frombuffer(image_from_tuple.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]			)

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

									UpperCamelCase__         = 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] ,      )

									UpperCamelCase__         = DDIMScheduler()
									UpperCamelCase__         = self.dummy_vqvae_and_unet
									UpperCamelCase__         = AudioDiffusionPipeline(
									    vqvae=self.dummy_vqvae_and_unet[0] ,      unet=dummy_vqvae_and_unet[1] ,      mel=lowercase ,      scheduler=lowercase			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									np.random.seed(0			)
									UpperCamelCase__         = np.random.uniform(-1 ,      1 ,      ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)			)
									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(raw_audio=lowercase ,      generator=lowercase ,      start_step=5 ,      steps=1_0			)
									UpperCamelCase__         = 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]
									)
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]			)

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

									UpperCamelCase__         = self.dummy_unet_condition
									UpperCamelCase__         = AudioDiffusionPipeline(
									    vqvae=self.dummy_vqvae_and_unet[0] ,      unet=lowercase ,      mel=lowercase ,      scheduler=lowercase			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									np.random.seed(0			)
									UpperCamelCase__         = torch.rand((1, 1, 1_0)			)
									UpperCamelCase__         = pipe(generator=lowercase ,      encoding=lowercase			)
									UpperCamelCase__         = output.images[0]
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]			)

									assert np.abs(image_slice.flatten() - expected_slice			).max() == 0
@slow
@require_torch_gpu
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						int			)							->    str:

									'''simple docstring'''







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







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

									'''simple docstring'''







									UpperCamelCase__         = torch_device

									UpperCamelCase__         = DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase			)
									UpperCamelCase__         = output.audios[0]
									UpperCamelCase__         = 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]
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]			)

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