


'''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     lowercase__  (						unittest.TestCase					):





							'''simple docstring'''


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






							@property
							def UpperCAmelCase_	(    self  ):
											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      UNetaDModel(
											    sample_size=(32, 64)    ,					in_channels=1    ,					out_channels=1    ,					layers_per_block=2    ,					block_out_channels=(128, 128)    ,					down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")    ,					up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")    ,					)
											return model






							@property
							def UpperCAmelCase_	(    self  ):
											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      UNetaDConditionModel(
											    sample_size=(64, 32)    ,					in_channels=1    ,					out_channels=1    ,					layers_per_block=2    ,					block_out_channels=(128, 128)    ,					down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")    ,					up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")    ,					cross_attention_dim=10    ,					)
											return model






							@property
							def UpperCAmelCase_	(    self  ):
											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      AutoencoderKL(
											    sample_size=(128, 64)    ,					in_channels=1    ,					out_channels=1    ,					latent_channels=1    ,					layers_per_block=2    ,					block_out_channels=(128, 128)    ,					down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")    ,					up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")    ,					)
											_SCREAMING_SNAKE_CASE :       List[str]     		=      UNetaDModel(
											    sample_size=(64, 32)    ,					in_channels=1    ,					out_channels=1    ,					layers_per_block=2    ,					block_out_channels=(128, 128)    ,					down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")    ,					up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")    ,					)
											return vqvae, unet






							@slow
							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       str     		=      """cpu"""  # ensure determinism for the device-dependent torch.Generator
											_SCREAMING_SNAKE_CASE :       str     		=      Mel(
											    x_res=self.dummy_unet.config.sample_size[1]    ,					y_res=self.dummy_unet.config.sample_size[0]    ,					)

											_SCREAMING_SNAKE_CASE :       Any     		=      DDPMScheduler()
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      AudioDiffusionPipeline(vqvae=__snake_case    ,					unet=self.dummy_unet    ,					mel=__snake_case    ,					scheduler=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      torch.Generator(device=__snake_case  ).manual_seed(42  )
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      pipe(generator=__snake_case    ,					steps=4  )
											_SCREAMING_SNAKE_CASE :       int     		=      output.audios[0]
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      output.images[0]

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      torch.Generator(device=__snake_case  ).manual_seed(42  )
											_SCREAMING_SNAKE_CASE :       int     		=      pipe(generator=__snake_case    ,					steps=4    ,					return_dict=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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 :       List[str]     		=      np.frombuffer(image.tobytes()    ,					dtype="""uint8"""  )[:10]
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      np.frombuffer(image_from_tuple.tobytes()    ,					dtype="""uint8"""  )[:10]
											_SCREAMING_SNAKE_CASE :       List[str]     		=      np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]  )

											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 :       str     		=      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[Any]     		=      DDIMScheduler()
											_SCREAMING_SNAKE_CASE :       str     		=      self.dummy_vqvae_and_unet
											_SCREAMING_SNAKE_CASE :       Dict     		=      AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]    ,					unet=dummy_vqvae_and_unet[1]    ,					mel=__snake_case    ,					scheduler=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )

											np.random.seed(0  )
											_SCREAMING_SNAKE_CASE :       str     		=      np.random.uniform(-1    ,					1    ,					((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)  )
											_SCREAMING_SNAKE_CASE :       int     		=      torch.Generator(device=__snake_case  ).manual_seed(42  )
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      pipe(raw_audio=__snake_case    ,					generator=__snake_case    ,					start_step=5    ,					steps=10  )
											_SCREAMING_SNAKE_CASE :       List[str]     		=      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 :       List[Any]     		=      np.frombuffer(image.tobytes()    ,					dtype="""uint8"""  )[:10]
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]  )

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

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

											np.random.seed(0  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      torch.rand((1, 1, 10)  )
											_SCREAMING_SNAKE_CASE :       int     		=      pipe(generator=__snake_case    ,					encoding=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      output.images[0]
											_SCREAMING_SNAKE_CASE :       int     		=      np.frombuffer(image.tobytes()    ,					dtype="""uint8"""  )[:10]
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]  )

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





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





							'''simple docstring'''


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






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      torch_device

											_SCREAMING_SNAKE_CASE :       str     		=      DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      torch.Generator(device=__snake_case  ).manual_seed(42  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      pipe(generator=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      output.audios[0]
											_SCREAMING_SNAKE_CASE :       int     		=      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 :       Tuple     		=      np.frombuffer(image.tobytes()    ,					dtype="""uint8"""  )[:10]
											_SCREAMING_SNAKE_CASE :       int     		=      np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]  )

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







import os

import numpy
import onnx





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    a.name
							snake_case_							=    b.name

							snake_case_							=    ''
							snake_case_							=    ''

							snake_case_							=    a == b

							snake_case_							=    name_a
							snake_case_							=    name_b

							return res





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							for i, input_name in enumerate(node_proto.input      ):
														if input_name == name:
																					node_proto.input.insert(UpperCamelCase__   ,							UpperCamelCase__      )
																					node_proto.input.pop(i + 1      )

							if node_proto.op_type == "If":
														_graph_replace_input_with(node_proto.attribute[0].g   ,							UpperCamelCase__   ,							UpperCamelCase__      )
														_graph_replace_input_with(node_proto.attribute[1].g   ,							UpperCamelCase__   ,							UpperCamelCase__      )
							if node_proto.op_type == "Loop":
														_graph_replace_input_with(node_proto.attribute[0].g   ,							UpperCamelCase__   ,							UpperCamelCase__      )





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							for n in graph_proto.node:
														_node_replace_input_with(UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      )





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    list(model.graph.initializer      )
							snake_case_							=    list(model_without_ext.graph.initializer      )
							for i, ref_i in ind_to_replace:
														assert inits_with_data[i].name == inits[i].name
														assert inits_with_data[ref_i].name == inits[ref_i].name
														assert i > ref_i

														snake_case_							=    inits[i].name
														snake_case_							=    inits[ref_i].name

														model_without_ext.graph.initializer.remove(inits[i]      )

														# for n in model.graph.node:
														_graph_replace_input_with(model_without_ext.graph   ,							UpperCamelCase__   ,							UpperCamelCase__      )





def 					__lowerCamelCase	( UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    os.path.dirname(UpperCamelCase__      )
							snake_case_							=    os.path.basename(UpperCamelCase__      )

							snake_case_							=    onnx.load(os.path.join(UpperCamelCase__   ,							UpperCamelCase__      )      )

							snake_case_							=    list(model.graph.initializer      )

							snake_case_							=    set()
							snake_case_							=    {}
							snake_case_							=    []

							snake_case_							=    0

							for i in range(len(UpperCamelCase__      )      ):
														if i in dup_set:
																					continue

														for j in range(i + 1   ,							len(UpperCamelCase__      )      ):
																					if j in dup_set:
																												continue
																					if _is_equal_tensor_proto(inits[i]   ,							inits[j]      ):
																												dup_set.add(UpperCamelCase__      )
																												dup_set.add(UpperCamelCase__      )

																												snake_case_							=    inits[j].data_type
																												snake_case_							=    numpy.prod(inits[j].dims      )
																												if dtype == 1:
																																			mem_size *= 4
																												elif dtype == 6:
																																			mem_size *= 4
																												elif dtype == 7 or dtype == 11:
																																			mem_size *= 8
																												else:
																																			print('unexpected data type: '   ,							UpperCamelCase__      )
																												total_reduced_size += mem_size

																												snake_case_							=    inits[i].name
																												snake_case_							=    inits[j].name

																												if name_i in dup_map:
																																			dup_map[name_i].append(UpperCamelCase__      )
																												else:
																																			snake_case_							=    [name_j]
																												ind_to_replace.append((j, i)      )

							print('total reduced size: '   ,							total_reduced_size / 1024 / 1024 / 1024   ,							'GB'      )

							snake_case_							=    sorted(UpperCamelCase__      )
							_remove_dup_initializers_from_model(UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      )

							snake_case_							=    'optimized_' + model_file_name
							snake_case_							=    os.path.join(UpperCamelCase__   ,							UpperCamelCase__      )
							onnx.save(UpperCamelCase__   ,							UpperCamelCase__      )

							return new_model




