diff --git a/music_spectrogram_diffusion/__init__.py b/music_spectrogram_diffusion/__init__.py index 715d986..08fd347 100644 --- a/music_spectrogram_diffusion/__init__.py +++ b/music_spectrogram_diffusion/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/audio_codecs.py b/music_spectrogram_diffusion/audio_codecs.py index b64cd43..b529dd9 100644 --- a/music_spectrogram_diffusion/audio_codecs.py +++ b/music_spectrogram_diffusion/audio_codecs.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/beam/evaluation.py b/music_spectrogram_diffusion/beam/evaluation.py index fb9dff0..ac87acf 100644 --- a/music_spectrogram_diffusion/beam/evaluation.py +++ b/music_spectrogram_diffusion/beam/evaluation.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -206,7 +206,7 @@ def process(self, id_examples): tick = time.time() if self._prediction_source == 'model': - pred_encoded, _ = jax.tree_map( + pred_encoded, _ = jax.tree.map( np.asarray, self._synth_model.predict(batch)) elif self._prediction_source in ('gt_raw', 'gt_encoded'): pred_encoded = batch['decoder_target_tokens'] diff --git a/music_spectrogram_diffusion/datasets.py b/music_spectrogram_diffusion/datasets.py index cebb914..20bab4f 100644 --- a/music_spectrogram_diffusion/datasets.py +++ b/music_spectrogram_diffusion/datasets.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/dump_task.py b/music_spectrogram_diffusion/dump_task.py index 17e08a4..a01ff90 100644 --- a/music_spectrogram_diffusion/dump_task.py +++ b/music_spectrogram_diffusion/dump_task.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/event_codec.py b/music_spectrogram_diffusion/event_codec.py index b8515d9..6ac54fb 100644 --- a/music_spectrogram_diffusion/event_codec.py +++ b/music_spectrogram_diffusion/event_codec.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/event_codec_test.py b/music_spectrogram_diffusion/event_codec_test.py index aeb3d07..1f04584 100644 --- a/music_spectrogram_diffusion/event_codec_test.py +++ b/music_spectrogram_diffusion/event_codec_test.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/feature_converters.py b/music_spectrogram_diffusion/feature_converters.py index dbf3652..47399d2 100644 --- a/music_spectrogram_diffusion/feature_converters.py +++ b/music_spectrogram_diffusion/feature_converters.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/inference.py b/music_spectrogram_diffusion/inference.py index 887a9f6..cedddd5 100644 --- a/music_spectrogram_diffusion/inference.py +++ b/music_spectrogram_diffusion/inference.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/layers.py b/music_spectrogram_diffusion/layers.py index fc9a035..3d4dd27 100644 --- a/music_spectrogram_diffusion/layers.py +++ b/music_spectrogram_diffusion/layers.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/layers_test.py b/music_spectrogram_diffusion/layers_test.py index f82a415..76cf6a3 100644 --- a/music_spectrogram_diffusion/layers_test.py +++ b/music_spectrogram_diffusion/layers_test.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/metrics.py b/music_spectrogram_diffusion/metrics.py index 998400b..7367dd9 100644 --- a/music_spectrogram_diffusion/metrics.py +++ b/music_spectrogram_diffusion/metrics.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/metrics_test.py b/music_spectrogram_diffusion/metrics_test.py index a66bc89..31de1cd 100644 --- a/music_spectrogram_diffusion/metrics_test.py +++ b/music_spectrogram_diffusion/metrics_test.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/models/autoregressive/models.py b/music_spectrogram_diffusion/models/autoregressive/models.py index 73586cd..6f2b33b 100644 --- a/music_spectrogram_diffusion/models/autoregressive/models.py +++ b/music_spectrogram_diffusion/models/autoregressive/models.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/models/autoregressive/network.py b/music_spectrogram_diffusion/models/autoregressive/network.py index 6a9f675..508b268 100644 --- a/music_spectrogram_diffusion/models/autoregressive/network.py +++ b/music_spectrogram_diffusion/models/autoregressive/network.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/models/autoregressive/output_functions.py b/music_spectrogram_diffusion/models/autoregressive/output_functions.py index 8c9c358..e29d777 100644 --- a/music_spectrogram_diffusion/models/autoregressive/output_functions.py +++ b/music_spectrogram_diffusion/models/autoregressive/output_functions.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/models/diffusion/diffusion_utils.py b/music_spectrogram_diffusion/models/diffusion/diffusion_utils.py index ca33303..8ca4762 100644 --- a/music_spectrogram_diffusion/models/diffusion/diffusion_utils.py +++ b/music_spectrogram_diffusion/models/diffusion/diffusion_utils.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/models/diffusion/feature_converters.py b/music_spectrogram_diffusion/models/diffusion/feature_converters.py index f2c89bb..cd166f6 100644 --- a/music_spectrogram_diffusion/models/diffusion/feature_converters.py +++ b/music_spectrogram_diffusion/models/diffusion/feature_converters.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/models/diffusion/models.py b/music_spectrogram_diffusion/models/diffusion/models.py index 841700e..84bc74c 100644 --- a/music_spectrogram_diffusion/models/diffusion/models.py +++ b/music_spectrogram_diffusion/models/diffusion/models.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -178,7 +178,7 @@ def predict_batch_with_aux( def pred_fn(z: jnp.ndarray, time: jnp.ndarray, include_conditioning: bool) -> jnp.ndarray: - step_encodings_and_masks = jax.tree_map( + step_encodings_and_masks = jax.tree.map( lambda x: x * include_conditioning, encodings_and_masks) return self.module.apply( { @@ -373,7 +373,7 @@ def predict_batch_with_aux( def pred_fn(z: jnp.ndarray, time: jnp.ndarray, include_conditioning: bool) -> jnp.ndarray: - step_encodings_and_masks = jax.tree_map( + step_encodings_and_masks = jax.tree.map( lambda x: x * include_conditioning, encodings_and_masks) return self.module.apply( { diff --git a/music_spectrogram_diffusion/models/diffusion/network.py b/music_spectrogram_diffusion/models/diffusion/network.py index 32a65a1..120e706 100644 --- a/music_spectrogram_diffusion/models/diffusion/network.py +++ b/music_spectrogram_diffusion/models/diffusion/network.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/note_sequences.py b/music_spectrogram_diffusion/note_sequences.py index d4bcd10..f36ed0c 100644 --- a/music_spectrogram_diffusion/note_sequences.py +++ b/music_spectrogram_diffusion/note_sequences.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/note_sequences_test.py b/music_spectrogram_diffusion/note_sequences_test.py index 813dcbf..a679c74 100644 --- a/music_spectrogram_diffusion/note_sequences_test.py +++ b/music_spectrogram_diffusion/note_sequences_test.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/postprocessors.py b/music_spectrogram_diffusion/postprocessors.py index 123d6b2..51e67a5 100644 --- a/music_spectrogram_diffusion/postprocessors.py +++ b/music_spectrogram_diffusion/postprocessors.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/preprocessors.py b/music_spectrogram_diffusion/preprocessors.py index ccfb0ec..384b6b8 100644 --- a/music_spectrogram_diffusion/preprocessors.py +++ b/music_spectrogram_diffusion/preprocessors.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/run_length_encoding.py b/music_spectrogram_diffusion/run_length_encoding.py index 2679455..c46d3e7 100644 --- a/music_spectrogram_diffusion/run_length_encoding.py +++ b/music_spectrogram_diffusion/run_length_encoding.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/run_length_encoding_test.py b/music_spectrogram_diffusion/run_length_encoding_test.py index aded33b..a1c2526 100644 --- a/music_spectrogram_diffusion/run_length_encoding_test.py +++ b/music_spectrogram_diffusion/run_length_encoding_test.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/tasks.py b/music_spectrogram_diffusion/tasks.py index 61cd63a..f20f3ff 100644 --- a/music_spectrogram_diffusion/tasks.py +++ b/music_spectrogram_diffusion/tasks.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/transcription_inference.py b/music_spectrogram_diffusion/transcription_inference.py index a49e72a..639ef60 100644 --- a/music_spectrogram_diffusion/transcription_inference.py +++ b/music_spectrogram_diffusion/transcription_inference.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/vocabularies.py b/music_spectrogram_diffusion/vocabularies.py index f6c0f3e..0677ec0 100644 --- a/music_spectrogram_diffusion/vocabularies.py +++ b/music_spectrogram_diffusion/vocabularies.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/music_spectrogram_diffusion/vocabularies_test.py b/music_spectrogram_diffusion/vocabularies_test.py index a4a9cec..c231952 100644 --- a/music_spectrogram_diffusion/vocabularies_test.py +++ b/music_spectrogram_diffusion/vocabularies_test.py @@ -1,4 +1,4 @@ -# Copyright 2023 The Music Spectrogram Diffusion Authors. +# Copyright 2024 The Music Spectrogram Diffusion Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License.