Add files using upload-large-folder tool

.gitattributes

@@ -96,3 +96,4 @@ torchmcubes-0.1.0-cp310-cp310-win_amd64.whl filter=lfs diff=lfs merge=lfs -text
 torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text
 groundingdino-0.1.0-cp310-cp310-win_amd64.whl filter=lfs diff=lfs merge=lfs -text
 groundingdino-0.1.0-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text
+Levo_Song_Generation/SongGeneration-Runtime/ckpt/models--lengyue233--content-vec-best/blobs/d8dd400e054ddf4e6be75dab5a2549db748cc99e756a097c496c099f65a4854e filter=lfs diff=lfs merge=lfs -text

Levo_Song_Generation/SongGeneration-Runtime/ckpt/encode-s12k.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e250df56b035f74c1f66f15133f4c78f664d70fa0b09aa9a752b7871bb58c02f
+size 3957949089

Levo_Song_Generation/SongGeneration-Runtime/ckpt/model_1rvq/model_2_fixed.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:339a16956b859a82defc02bfd32c3744d11ff942065f6ec9306dfd4400d62110
+size 4704507596

Levo_Song_Generation/SongGeneration-Runtime/ckpt/model_septoken/model_2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:430b7c1c245722fbe3893cd621b3d4a90076404596e9fb1ce987a4a0f2a4fc6f
+size 4808167708

Levo_Song_Generation/SongGeneration-Runtime/ckpt/models--lengyue233--content-vec-best/blobs/d8dd400e054ddf4e6be75dab5a2549db748cc99e756a097c496c099f65a4854e

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8dd400e054ddf4e6be75dab5a2549db748cc99e756a097c496c099f65a4854e
+size 378342945

Levo_Song_Generation/SongGeneration-Runtime/ckpt/models--lengyue233--content-vec-best/snapshots/c0b9ba13db21beaa4053faae94c102ebe326fd68/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8dd400e054ddf4e6be75dab5a2549db748cc99e756a097c496c099f65a4854e
+size 378342945

Levo_Song_Generation/SongGeneration-Runtime/ckpt/vae/autoencoder_music_1320k.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10ccb6c83613781ad32e998a90597ba7eb9292911a224598da1fd53728eb4cd3
+size 674920616

Levo_Song_Generation/SongGeneration-Runtime/third_party/Qwen2-7B/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2024 Alibaba Cloud
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Levo_Song_Generation/SongGeneration-Runtime/third_party/Qwen2-7B/README.md

@@ -0,0 +1,97 @@
+---
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- pretrained
+license: apache-2.0
+---
+
+# Qwen2-7B
+
+## Introduction
+
+Qwen2 is the new series of Qwen large language models. For Qwen2, we release a number of base language models and instruction-tuned language models ranging from 0.5 to 72 billion parameters, including a Mixture-of-Experts model. This repo contains the 7B Qwen2 base language model.
+
+Compared with the state-of-the-art opensource language models, including the previous released Qwen1.5, Qwen2 has generally surpassed most opensource models and demonstrated competitiveness against proprietary models across a series of benchmarks targeting for language understanding, language generation, multilingual capability, coding, mathematics, reasoning, etc.
+
+For more details, please refer to our [blog](https://qwenlm.github.io/blog/qwen2/), [GitHub](https://github.com/QwenLM/Qwen2), and [Documentation](https://qwen.readthedocs.io/en/latest/).
+<br>
+
+
+## Model Details
+Qwen2 is a language model series including decoder language models of different model sizes. For each size, we release the base language model and the aligned chat model. It is based on the Transformer architecture with SwiGLU activation, attention QKV bias, group query attention, etc. Additionally, we have an improved tokenizer adaptive to multiple natural languages and codes.
+
+## Requirements
+The code of Qwen2 has been in the latest Hugging face transformers and we advise you to install `transformers>=4.37.0`, or you might encounter the following error:
+```
+KeyError: 'qwen2'
+```
+
+
+## Usage
+
+We do not advise you to use base language models for text generation. Instead, you can apply post-training, e.g., SFT, RLHF, continued pretraining, etc., on this model.
+
+
+### Performance
+
+The evaluation of base models mainly focuses on the model performance of natural language understanding, general question answering, coding, mathematics, scientific knowledge, reasoning, multilingual capability, etc. 
+
+The datasets for evaluation include: 
+ 
+**English Tasks**: MMLU (5-shot), MMLU-Pro (5-shot), GPQA (5shot), Theorem QA (5-shot), BBH (3-shot), HellaSwag (10-shot), Winogrande (5-shot), TruthfulQA (0-shot), ARC-C (25-shot)
+ 
+**Coding Tasks**: EvalPlus (0-shot) (HumanEval, MBPP, HumanEval+, MBPP+), MultiPL-E (0-shot) (Python, C++, JAVA, PHP, TypeScript, C#, Bash, JavaScript)
+  
+**Math Tasks**: GSM8K (4-shot), MATH (4-shot)
+ 
+**Chinese Tasks**: C-Eval(5-shot), CMMLU (5-shot)
+ 
+**Multilingual Tasks**: Multi-Exam (M3Exam 5-shot, IndoMMLU 3-shot, ruMMLU 5-shot, mMMLU 5-shot), Multi-Understanding (BELEBELE 5-shot, XCOPA 5-shot, XWinograd 5-shot, XStoryCloze 0-shot, PAWS-X 5-shot), Multi-Mathematics (MGSM 8-shot), Multi-Translation (Flores-101 5-shot)
+ 
+
+  
+#### Qwen2-7B performance
+|  Datasets  |  Mistral-7B  |   Gemma-7B |   Llama-3-8B  |   Qwen1.5-7B  |  Qwen2-7B  |
+| :--------| :---------: | :------------: | :------------: | :------------: | :------------: |
+|# Params | 7.2B | 8.5B | 8.0B | 7.7B | 7.6B  |
+|# Non-emb Params | 7.0B | 7.8B | 7.0B | 6.5B | 6.5B |
+|   ***English***  |    |    |   |    |	    |
+|MMLU | 64.2 | 64.6 | 66.6 | 61.0 | **70.3** |
+|MMLU-Pro | 30.9 | 33.7 | 35.4 | 29.9 | **40.0** |
+|GPQA | 24.7 | 25.7 | 25.8 | 26.7 | **31.8** |
+|Theorem QA | 19.2 | 21.5 | 22.1 | 14.2 | **31.1** |
+|BBH  | 56.1 |  55.1  | 57.7 | 40.2 | **62.6** |
+|HellaSwag  | **83.2** |  82.2  | 82.1 | 78.5 | 80.7 |
+|Winogrande  | 78.4 |  **79.0**  | 77.4 |  71.3 |  77.0 |
+|ARC-C  | 60.0 |  **61.1**  | 59.3 | 54.2 |  60.6 |
+|TruthfulQA  | 42.2 |  44.8  | 44.0 | 51.1 |  **54.2** |
+|   ***Coding***  |    |    |   |    |	    |
+|HumanEval | 29.3 | 37.2 | 33.5 | 36.0 | **51.2**  |
+|MBPP | 51.1 | 50.6 | 53.9 | 51.6 | **65.9**  |
+|EvalPlus | 36.4 | 39.6 | 40.3 | 40.0 | **54.2**  |
+|MultiPL-E | 29.4 | 29.7 | 22.6 | 28.1 | **46.3**  |
+|   ***Mathematics***  |    |    |   |    |	    |
+|GSM8K | 52.2 |  46.4  | 56.0 | 62.5 | **79.9** |
+|MATH  | 13.1 |  24.3  | 20.5 | 20.3 | **44.2** |
+|   ***Chinese***  |    |    |   |    |	    |
+|C-Eval   | 47.4 |   43.6    |  49.5 |  74.1 |  **83.2** |
+|CMMLU   | - |   -    | 50.8 | 73.1 | **83.9** |
+|   ***Multilingual***  |    |    |   |    |	    |
+|Multi-Exam   | 47.1 |   42.7    |  52.3 |  47.7 |  **59.2** |
+|Multi-Understanding | 63.3 |  58.3    |  68.6 |  67.6 |  **72.0** |
+|Multi-Mathematics | 26.3 |   39.1    |  36.3 |  37.3 |  **57.5** |
+|Multi-Translation | 23.3 |   31.2    |  **31.9** |  28.4 |  31.5 |
+
+
+## Citation
+
+If you find our work helpful, feel free to give us a cite.
+
+```
+@article{qwen2,
+  title={Qwen2 Technical Report},
+  year={2024}
+}
+```

Levo_Song_Generation/SongGeneration-Runtime/third_party/Qwen2-7B/config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 3584,
+  "initializer_range": 0.02,
+  "intermediate_size": 18944,
+  "max_position_embeddings": 131072,
+  "max_window_layers": 28,
+  "model_type": "qwen2",
+  "num_attention_heads": 28,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 4,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sliding_window": 131072,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.37.2",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 152064
+}

Levo_Song_Generation/SongGeneration-Runtime/third_party/Qwen2-7B/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": false,
+  "eos_token_id": 151643,
+  "max_new_tokens": 2048,
+  "transformers_version": "4.37.0"
+}

Levo_Song_Generation/SongGeneration-Runtime/third_party/Qwen2-7B/tokenizer_config.json

@@ -0,0 +1,40 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": ["<|im_start|>", "<|im_end|>"],
+  "bos_token": null,
+  "chat_template": "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/ckpt/htdemucs.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4378974c3df2fbcf872d2aeb32218e4de376799494579655775a375d09931c2
+size 168138881

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/ckpt/htdemucs.yaml

@@ -0,0 +1 @@
+models: ['htdemucs']

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/apply.py

@@ -0,0 +1,315 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : apply.py
+@Time    : 2023/8/8 下午4:22
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Apply
+"""
+
+from concurrent.futures import ThreadPoolExecutor
+import torch
+import os
+import random
+import typing as tp
+
+import torch as th
+from torch import nn
+from torch.nn import functional as F
+import tqdm
+
+from .htdemucs import HTDemucs
+from .audio import load_track, save_audio
+from .utils import center_trim, DummyPoolExecutor
+
+Model = tp.Union[HTDemucs]
+
+
+class BagOfModels(nn.Module):
+    def __init__(self, models: tp.List[Model],
+                 weights: tp.Optional[tp.List[tp.List[float]]] = None,
+                 segment: tp.Optional[float] = None):
+        """
+        Represents a bag of models with specific weights.
+        You should call `apply_model` rather than calling directly the forward here for
+        optimal performance.
+
+        Args:
+            models (list[nn.Module]): list of Demucs/HDemucs models.
+            weights (list[list[float]]): list of weights. If None, assumed to
+                be all ones, otherwise it should be a list of N list (N number of models),
+                each containing S floats (S number of sources).
+            segment (None or float): overrides the `segment` attribute of each model
+                (this is performed inplace, be careful is you reuse the models passed).
+        """
+        super().__init__()
+        assert len(models) > 0
+        first = models[0]
+        for other in models:
+            assert other.sources == first.sources
+            assert other.samplerate == first.samplerate
+            assert other.audio_channels == first.audio_channels
+            if segment is not None:
+                other.segment = segment
+
+        self.audio_channels = first.audio_channels
+        self.samplerate = first.samplerate
+        self.sources = first.sources
+        self.models = nn.ModuleList(models)
+
+        if weights is None:
+            weights = [[1. for _ in first.sources] for _ in models]
+        else:
+            assert len(weights) == len(models)
+            for weight in weights:
+                assert len(weight) == len(first.sources)
+        self.weights = weights
+
+    @property
+    def max_allowed_segment(self) -> float:
+        max_allowed_segment = float('inf')
+        for model in self.models:
+            if isinstance(model, HTDemucs):
+                max_allowed_segment = min(max_allowed_segment, float(model.segment))
+        return max_allowed_segment
+
+    def forward(self, x):
+        raise NotImplementedError("Call `apply_model` on this.")
+    
+    def separate(self, source_file, output_dir, stem=None, device=None):
+        wav, _ = load_track(source_file, self.audio_channels, self.samplerate)
+        ref = wav.mean(0)
+        wav -= ref.mean()
+        wav /= ref.std()
+        sources = apply_model(self, wav[None], device=device, shifts=1, split=True, overlap=0.25, 
+                              progress=True, num_workers=0, segment=None)[0]
+        sources *= ref.std()
+        sources += ref.mean()
+
+        output_paths = []
+        name, ext = os.path.splitext(os.path.split(source_file)[-1])
+        if ext != ".flac":
+            ext = ".flac"
+        kwargs = {
+            'samplerate': self.samplerate,
+            'bitrate': 320,
+            'clip': "rescale",
+            'as_float': False,
+            'bits_per_sample': 16,
+        }
+        if stem is None:
+            for source, stem in zip(sources, self.sources):
+                output_stem_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
+                save_audio(source, output_stem_path, **kwargs)
+                output_paths.append(output_stem_path)
+        else:
+            sources = list(sources)
+            output_stem_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
+            save_audio(sources.pop(self.sources.index(stem)), output_stem_path, **kwargs)
+            other_stem = torch.zeros_like(sources[0])
+            for i in sources:
+                other_stem += i
+            output_no_stem_path = os.path.join(output_dir, f"{name}_no_{stem}{ext}")
+            save_audio(other_stem, output_no_stem_path, **kwargs)
+            output_paths = [output_stem_path, output_no_stem_path]
+
+        return output_paths
+
+
+class TensorChunk:
+    def __init__(self, tensor, offset=0, length=None):
+        total_length = tensor.shape[-1]
+        assert offset >= 0
+        assert offset < total_length
+
+        if length is None:
+            length = total_length - offset
+        else:
+            length = min(total_length - offset, length)
+
+        if isinstance(tensor, TensorChunk):
+            self.tensor = tensor.tensor
+            self.offset = offset + tensor.offset
+        else:
+            self.tensor = tensor
+            self.offset = offset
+        self.length = length
+        self.device = tensor.device
+
+    @property
+    def shape(self):
+        shape = list(self.tensor.shape)
+        shape[-1] = self.length
+        return shape
+
+    def padded(self, target_length):
+        delta = target_length - self.length
+        total_length = self.tensor.shape[-1]
+        assert delta >= 0
+
+        start = self.offset - delta // 2
+        end = start + target_length
+
+        correct_start = max(0, start)
+        correct_end = min(total_length, end)
+
+        pad_left = correct_start - start
+        pad_right = end - correct_end
+
+        out = F.pad(self.tensor[..., correct_start:correct_end], (pad_left, pad_right))
+        assert out.shape[-1] == target_length
+        return out
+
+
+def tensor_chunk(tensor_or_chunk):
+    if isinstance(tensor_or_chunk, TensorChunk):
+        return tensor_or_chunk
+    else:
+        assert isinstance(tensor_or_chunk, th.Tensor)
+        return TensorChunk(tensor_or_chunk)
+
+
+def apply_model(model: tp.Union[BagOfModels, Model],
+                mix: tp.Union[th.Tensor, TensorChunk],
+                shifts: int = 1, split: bool = True,
+                overlap: float = 0.25, transition_power: float = 1.,
+                progress: bool = False, device=None,
+                num_workers: int = 0, segment: tp.Optional[float] = None,
+                pool=None) -> th.Tensor:
+    """
+    Apply model to a given mixture.
+
+    Args:
+        shifts (int): if > 0, will shift in time `mix` by a random amount between 0 and 0.5 sec
+            and apply the oppositve shift to the output. This is repeated `shifts` time and
+            all predictions are averaged. This effectively makes the model time equivariant
+            and improves SDR by up to 0.2 points.
+        split (bool): if True, the input will be broken down in 8 seconds extracts
+            and predictions will be performed individually on each and concatenated.
+            Useful for model with large memory footprint like Tasnet.
+        progress (bool): if True, show a progress bar (requires split=True)
+        device (torch.device, str, or None): if provided, device on which to
+            execute the computation, otherwise `mix.device` is assumed.
+            When `device` is different from `mix.device`, only local computations will
+            be on `device`, while the entire tracks will be stored on `mix.device`.
+        num_workers (int): if non zero, device is 'cpu', how many threads to
+            use in parallel.
+        segment (float or None): override the model segment parameter.
+    """
+    if device is None:
+        device = mix.device
+    else:
+        device = th.device(device)
+    if pool is None:
+        if num_workers > 0 and device.type == 'cpu':
+            pool = ThreadPoolExecutor(num_workers)
+        else:
+            pool = DummyPoolExecutor()
+    kwargs: tp.Dict[str, tp.Any] = {
+        'shifts': shifts,
+        'split': split,
+        'overlap': overlap,
+        'transition_power': transition_power,
+        'progress': progress,
+        'device': device,
+        'pool': pool,
+        'segment': segment,
+    }
+    out: tp.Union[float, th.Tensor]
+    if isinstance(model, BagOfModels):
+        # Special treatment for bag of model.
+        # We explicitely apply multiple times `apply_model` so that the random shifts
+        # are different for each model.
+        estimates: tp.Union[float, th.Tensor] = 0.
+        totals = [0.] * len(model.sources)
+        for sub_model, model_weights in zip(model.models, model.weights):
+            original_model_device = next(iter(sub_model.parameters())).device
+            sub_model.to(device)
+
+            out = apply_model(sub_model, mix, **kwargs)
+            sub_model.to(original_model_device)
+            for k, inst_weight in enumerate(model_weights):
+                out[:, k, :, :] *= inst_weight
+                totals[k] += inst_weight
+            estimates += out
+            del out
+
+        assert isinstance(estimates, th.Tensor)
+        for k in range(estimates.shape[1]):
+            estimates[:, k, :, :] /= totals[k]
+        return estimates
+
+    model.to(device)
+    model.eval()
+    assert transition_power >= 1, "transition_power < 1 leads to weird behavior."
+    batch, channels, length = mix.shape
+    if shifts:
+        kwargs['shifts'] = 0
+        max_shift = int(0.5 * model.samplerate)
+        mix = tensor_chunk(mix)
+        assert isinstance(mix, TensorChunk)
+        padded_mix = mix.padded(length + 2 * max_shift)
+        out = 0.
+        for _ in range(shifts):
+            offset = random.randint(0, max_shift)
+            shifted = TensorChunk(padded_mix, offset, length + max_shift - offset)
+            shifted_out = apply_model(model, shifted, **kwargs)
+            out += shifted_out[..., max_shift - offset:]
+        out /= shifts
+        assert isinstance(out, th.Tensor)
+        return out
+    elif split:
+        kwargs['split'] = False
+        out = th.zeros(batch, len(model.sources), channels, length, device=mix.device)
+        sum_weight = th.zeros(length, device=mix.device)
+        if segment is None:
+            segment = model.segment
+        assert segment is not None and segment > 0.
+        segment_length: int = int(model.samplerate * segment)
+        stride = int((1 - overlap) * segment_length)
+        offsets = range(0, length, stride)
+        scale = float(format(stride / model.samplerate, ".2f"))
+        # We start from a triangle shaped weight, with maximal weight in the middle
+        # of the segment. Then we normalize and take to the power `transition_power`.
+        # Large values of transition power will lead to sharper transitions.
+        weight = th.cat([th.arange(1, segment_length // 2 + 1, device=device),
+                         th.arange(segment_length - segment_length // 2, 0, -1, device=device)])
+        assert len(weight) == segment_length
+        # If the overlap < 50%, this will translate to linear transition when
+        # transition_power is 1.
+        weight = (weight / weight.max())**transition_power
+        futures = []
+        for offset in offsets:
+            chunk = TensorChunk(mix, offset, segment_length)
+            future = pool.submit(apply_model, model, chunk, **kwargs)
+            futures.append((future, offset))
+            offset += segment_length
+        if progress:
+            futures = tqdm.tqdm(futures, unit_scale=scale, ncols=120, unit='seconds')
+        for future, offset in futures:
+            chunk_out = future.result()
+            chunk_length = chunk_out.shape[-1]
+            out[..., offset:offset + segment_length] += (
+                weight[:chunk_length] * chunk_out).to(mix.device)
+            sum_weight[offset:offset + segment_length] += weight[:chunk_length].to(mix.device)
+        assert sum_weight.min() > 0
+        out /= sum_weight
+        assert isinstance(out, th.Tensor)
+        return out
+    else:
+        valid_length: int
+        if isinstance(model, HTDemucs) and segment is not None:
+            valid_length = int(segment * model.samplerate)
+        elif hasattr(model, 'valid_length'):
+            valid_length = model.valid_length(length)  # type: ignore
+        else:
+            valid_length = length
+        mix = tensor_chunk(mix)
+        assert isinstance(mix, TensorChunk)
+        padded_mix = mix.padded(valid_length).to(device)
+        with th.no_grad():
+            out = model(padded_mix)
+        assert isinstance(out, th.Tensor)
+        return center_trim(out, length)

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/audio.py

@@ -0,0 +1,291 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : audio.py
+@Time    : 2023/8/8 下午7:18
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Audio
+"""
+import json
+import subprocess as sp
+import typing as tp
+from pathlib import Path
+
+import lameenc
+import julius
+import torch
+import numpy as np
+import torchaudio as ta
+
+from .utils import temp_filenames
+
+
+def _read_info(path):
+    stdout_data = sp.check_output([
+        'ffprobe', "-loglevel", "panic",
+        str(path), '-print_format', 'json', '-show_format', '-show_streams'
+    ])
+    return json.loads(stdout_data.decode('utf-8'))
+
+
+class AudioFile:
+    """
+    Allows to read audio from any format supported by ffmpeg, as well as resampling or
+    converting to mono on the fly. See :method:`read` for more details.
+    """
+    def __init__(self, path: Path):
+        self.path = Path(path)
+        self._info = None
+
+    def __repr__(self):
+        features = [("path", self.path)]
+        features.append(("samplerate", self.samplerate()))
+        features.append(("channels", self.channels()))
+        features.append(("streams", len(self)))
+        features_str = ", ".join(f"{name}={value}" for name, value in features)
+        return f"AudioFile({features_str})"
+
+    @property
+    def info(self):
+        if self._info is None:
+            self._info = _read_info(self.path)
+        return self._info
+
+    @property
+    def duration(self):
+        return float(self.info['format']['duration'])
+
+    @property
+    def _audio_streams(self):
+        return [
+            index for index, stream in enumerate(self.info["streams"])
+            if stream["codec_type"] == "audio"
+        ]
+
+    def __len__(self):
+        return len(self._audio_streams)
+
+    def channels(self, stream=0):
+        return int(self.info['streams'][self._audio_streams[stream]]['channels'])
+
+    def samplerate(self, stream=0):
+        return int(self.info['streams'][self._audio_streams[stream]]['sample_rate'])
+
+    def read(self,
+             seek_time=None,
+             duration=None,
+             streams=slice(None),
+             samplerate=None,
+             channels=None):
+        """
+        Slightly more efficient implementation than stempeg,
+        in particular, this will extract all stems at once
+        rather than having to loop over one file multiple times
+        for each stream.
+
+        Args:
+            seek_time (float):  seek time in seconds or None if no seeking is needed.
+            duration (float): duration in seconds to extract or None to extract until the end.
+            streams (slice, int or list): streams to extract, can be a single int, a list or
+                a slice. If it is a slice or list, the output will be of size [S, C, T]
+                with S the number of streams, C the number of channels and T the number of samples.
+                If it is an int, the output will be [C, T].
+            samplerate (int): if provided, will resample on the fly. If None, no resampling will
+                be done. Original sampling rate can be obtained with :method:`samplerate`.
+            channels (int): if 1, will convert to mono. We do not rely on ffmpeg for that
+                as ffmpeg automatically scale by +3dB to conserve volume when playing on speakers.
+                See https://sound.stackexchange.com/a/42710.
+                Our definition of mono is simply the average of the two channels. Any other
+                value will be ignored.
+        """
+        streams = np.array(range(len(self)))[streams]
+        single = not isinstance(streams, np.ndarray)
+        if single:
+            streams = [streams]
+
+        if duration is None:
+            target_size = None
+            query_duration = None
+        else:
+            target_size = int((samplerate or self.samplerate()) * duration)
+            query_duration = float((target_size + 1) / (samplerate or self.samplerate()))
+
+        with temp_filenames(len(streams)) as filenames:
+            command = ['ffmpeg', '-y']
+            command += ['-loglevel', 'panic']
+            if seek_time:
+                command += ['-ss', str(seek_time)]
+            command += ['-i', str(self.path)]
+            for stream, filename in zip(streams, filenames):
+                command += ['-map', f'0:{self._audio_streams[stream]}']
+                if query_duration is not None:
+                    command += ['-t', str(query_duration)]
+                command += ['-threads', '1']
+                command += ['-f', 'f32le']
+                if samplerate is not None:
+                    command += ['-ar', str(samplerate)]
+                command += [filename]
+
+            sp.run(command, check=True)
+            wavs = []
+            for filename in filenames:
+                wav = np.fromfile(filename, dtype=np.float32)
+                wav = torch.from_numpy(wav)
+                wav = wav.view(-1, self.channels()).t()
+                if channels is not None:
+                    wav = convert_audio_channels(wav, channels)
+                if target_size is not None:
+                    wav = wav[..., :target_size]
+                wavs.append(wav)
+        wav = torch.stack(wavs, dim=0)
+        if single:
+            wav = wav[0]
+        return wav
+
+
+def convert_audio_channels(wav, channels=2):
+    """Convert audio to the given number of channels."""
+    *shape, src_channels, length = wav.shape
+    if src_channels == channels:
+        pass
+    elif channels == 1:
+        # Case 1:
+        # The caller asked 1-channel audio, but the stream have multiple
+        # channels, downmix all channels.
+        wav = wav.mean(dim=-2, keepdim=True)
+    elif src_channels == 1:
+        # Case 2:
+        # The caller asked for multiple channels, but the input file have
+        # one single channel, replicate the audio over all channels.
+        wav = wav.expand(*shape, channels, length)
+    elif src_channels >= channels:
+        # Case 3:
+        # The caller asked for multiple channels, and the input file have
+        # more channels than requested. In that case return the first channels.
+        wav = wav[..., :channels, :]
+    else:
+        # Case 4: What is a reasonable choice here?
+        raise ValueError('The audio file has less channels than requested but is not mono.')
+    return wav
+
+
+def convert_audio(wav, from_samplerate, to_samplerate, channels):
+    """Convert audio from a given samplerate to a target one and target number of channels."""
+    wav = convert_audio_channels(wav, channels)
+    return julius.resample_frac(wav, from_samplerate, to_samplerate)
+
+
+def i16_pcm(wav):
+    """Convert audio to 16 bits integer PCM format."""
+    if wav.dtype.is_floating_point:
+        return (wav.clamp_(-1, 1) * (2**15 - 1)).short()
+    else:
+        return wav
+
+
+def f32_pcm(wav):
+    """Convert audio to float 32 bits PCM format."""
+    if wav.dtype.is_floating_point:
+        return wav
+    else:
+        return wav.float() / (2**15 - 1)
+
+
+def as_dtype_pcm(wav):
+    """Convert audio to either f32 pcm or i16 pcm depending on the given dtype."""
+    if wav.dtype.is_floating_point:
+        return f32_pcm(wav)
+    else:
+        return i16_pcm(wav)
+
+
+def encode_mp3(wav, path, samplerate=44100, bitrate=320, verbose=False):
+    """Save given audio as mp3. This should work on all OSes."""
+    c, _ = wav.shape
+    wav = i16_pcm(wav)
+    encoder = lameenc.Encoder()
+    encoder.set_bit_rate(bitrate)
+    encoder.set_in_sample_rate(samplerate)
+    encoder.set_channels(c)
+    encoder.set_quality(2)  # 2-highest, 7-fastest
+    if not verbose:
+        encoder.silence()
+    wav = wav.data.cpu()
+    wav = wav.transpose(0, 1).numpy()
+    mp3_data = encoder.encode(wav.tobytes())
+    mp3_data += encoder.flush()
+    with open(path, "wb") as f:
+        f.write(mp3_data)
+
+
+def prevent_clip(wav, mode='rescale'):
+    """
+    different strategies for avoiding raw clipping.
+    """
+    if mode is None or mode == 'none':
+        return wav
+    assert wav.dtype.is_floating_point, "too late for clipping"
+    if mode == 'rescale':
+        wav = wav / max(1.01 * wav.abs().max(), 1)
+    elif mode == 'clamp':
+        wav = wav.clamp(-0.99, 0.99)
+    elif mode == 'tanh':
+        wav = torch.tanh(wav)
+    else:
+        raise ValueError(f"Invalid mode {mode}")
+    return wav
+
+
+def save_audio(wav: torch.Tensor,
+               path: tp.Union[str, Path],
+               samplerate: int,
+               bitrate: int = 320,
+               clip: tp.Union[str] = 'rescale',
+               bits_per_sample: tp.Union[int] = 16,
+               as_float: bool = False):
+    """Save audio file, automatically preventing clipping if necessary
+    based on the given `clip` strategy. If the path ends in `.mp3`, this
+    will save as mp3 with the given `bitrate`.
+    """
+    wav = prevent_clip(wav, mode=clip)
+    path = Path(path)
+    suffix = path.suffix.lower()
+    if suffix == ".mp3":
+        encode_mp3(wav, path, samplerate, bitrate, verbose=True)
+    elif suffix == ".wav":
+        if as_float:
+            bits_per_sample = 32
+            encoding = 'PCM_F'
+        else:
+            encoding = 'PCM_S'
+        ta.save(str(path), wav, sample_rate=samplerate,
+                encoding=encoding, bits_per_sample=bits_per_sample)
+    elif suffix == ".flac":
+        ta.save(str(path), wav, sample_rate=samplerate, bits_per_sample=bits_per_sample)
+    else:
+        raise ValueError(f"Invalid suffix for path: {suffix}")
+
+
+def load_track(track, audio_channels, samplerate):
+    errors = {}
+    wav = None
+
+    try:
+        wav = AudioFile(track).read(
+            streams=0,
+            samplerate=samplerate,
+            channels=audio_channels)
+    except sp.CalledProcessError:
+        errors['ffmpeg'] = 'FFmpeg could not read the file.'
+
+    if wav is None:
+        try:
+            wav, sr = ta.load(str(track))
+        except RuntimeError as err:
+            errors['torchaudio'] = err.args[0]
+        else:
+            wav = convert_audio(wav, sr, samplerate, audio_channels)
+
+    return wav, errors

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/demucs.py

@@ -0,0 +1,452 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : demucs.py
+@Time    : 2023/8/8 下午4:36
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Demucs
+"""
+
+import math
+import typing as tp
+
+import julius
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .states import capture_init
+from .utils import center_trim, unfold
+from .transformer import LayerScale
+
+
+class BLSTM(nn.Module):
+    """
+    BiLSTM with same hidden units as input dim.
+    If `max_steps` is not None, input will be splitting in overlapping
+    chunks and the LSTM applied separately on each chunk.
+    """
+    def __init__(self, dim, layers=1, max_steps=None, skip=False):
+        super().__init__()
+        assert max_steps is None or max_steps % 4 == 0
+        self.max_steps = max_steps
+        self.lstm = nn.LSTM(bidirectional=True, num_layers=layers, hidden_size=dim, input_size=dim)
+        self.linear = nn.Linear(2 * dim, dim)
+        self.skip = skip
+
+    def forward(self, x):
+        b, c, t = x.shape
+        y = x
+        framed = False
+        if self.max_steps is not None and t > self.max_steps:
+            width = self.max_steps
+            stride = width // 2
+            frames = unfold(x, width, stride)
+            nframes = frames.shape[2]
+            framed = True
+            x = frames.permute(0, 2, 1, 3).reshape(-1, c, width)
+
+        x = x.permute(2, 0, 1)
+
+        x = self.lstm(x)[0]
+        x = self.linear(x)
+        x = x.permute(1, 2, 0)
+        if framed:
+            out = []
+            frames = x.reshape(b, -1, c, width)
+            limit = stride // 2
+            for k in range(nframes):
+                if k == 0:
+                    out.append(frames[:, k, :, :-limit])
+                elif k == nframes - 1:
+                    out.append(frames[:, k, :, limit:])
+                else:
+                    out.append(frames[:, k, :, limit:-limit])
+            out = torch.cat(out, -1)
+            out = out[..., :t]
+            x = out
+        if self.skip:
+            x = x + y
+        return x
+
+
+def rescale_conv(conv, reference):
+    """Rescale initial weight scale. It is unclear why it helps but it certainly does.
+    """
+    std = conv.weight.std().detach()
+    scale = (std / reference)**0.5
+    conv.weight.data /= scale
+    if conv.bias is not None:
+        conv.bias.data /= scale
+
+
+def rescale_module(module, reference):
+    for sub in module.modules():
+        if isinstance(sub, (nn.Conv1d, nn.ConvTranspose1d, nn.Conv2d, nn.ConvTranspose2d)):
+            rescale_conv(sub, reference)
+
+
+class DConv(nn.Module):
+    """
+    New residual branches in each encoder layer.
+    This alternates dilated convolutions, potentially with LSTMs and attention.
+    Also before entering each residual branch, dimension is projected on a smaller subspace,
+    e.g. of dim `channels // compress`.
+    """
+    def __init__(self, channels: int, compress: float = 4, depth: int = 2, init: float = 1e-4,
+                 norm=True, attn=False, heads=4, ndecay=4, lstm=False, gelu=True,
+                 kernel=3, dilate=True):
+        """
+        Args:
+            channels: input/output channels for residual branch.
+            compress: amount of channel compression inside the branch.
+            depth: number of layers in the residual branch. Each layer has its own
+                projection, and potentially LSTM and attention.
+            init: initial scale for LayerNorm.
+            norm: use GroupNorm.
+            attn: use LocalAttention.
+            heads: number of heads for the LocalAttention.
+            ndecay: number of decay controls in the LocalAttention.
+            lstm: use LSTM.
+            gelu: Use GELU activation.
+            kernel: kernel size for the (dilated) convolutions.
+            dilate: if true, use dilation, increasing with the depth.
+        """
+
+        super().__init__()
+        assert kernel % 2 == 1
+        self.channels = channels
+        self.compress = compress
+        self.depth = abs(depth)
+        dilate = depth > 0
+
+        norm_fn: tp.Callable[[int], nn.Module]
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm:
+            norm_fn = lambda d: nn.GroupNorm(1, d)  # noqa
+
+        hidden = int(channels / compress)
+
+        act: tp.Type[nn.Module]
+        if gelu:
+            act = nn.GELU
+        else:
+            act = nn.ReLU
+
+        self.layers = nn.ModuleList([])
+        for d in range(self.depth):
+            dilation = 2 ** d if dilate else 1
+            padding = dilation * (kernel // 2)
+            mods = [
+                nn.Conv1d(channels, hidden, kernel, dilation=dilation, padding=padding),
+                norm_fn(hidden), act(),
+                nn.Conv1d(hidden, 2 * channels, 1),
+                norm_fn(2 * channels), nn.GLU(1),
+                LayerScale(channels, init),
+            ]
+            if attn:
+                mods.insert(3, LocalState(hidden, heads=heads, ndecay=ndecay))
+            if lstm:
+                mods.insert(3, BLSTM(hidden, layers=2, max_steps=200, skip=True))
+            layer = nn.Sequential(*mods)
+            self.layers.append(layer)
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = x + layer(x)
+        return x
+
+
+class LocalState(nn.Module):
+    """Local state allows to have attention based only on data (no positional embedding),
+    but while setting a constraint on the time window (e.g. decaying penalty term).
+
+    Also a failed experiments with trying to provide some frequency based attention.
+    """
+    def __init__(self, channels: int, heads: int = 4, nfreqs: int = 0, ndecay: int = 4):
+        super().__init__()
+        assert channels % heads == 0, (channels, heads)
+        self.heads = heads
+        self.nfreqs = nfreqs
+        self.ndecay = ndecay
+        self.content = nn.Conv1d(channels, channels, 1)
+        self.query = nn.Conv1d(channels, channels, 1)
+        self.key = nn.Conv1d(channels, channels, 1)
+        if nfreqs:
+            self.query_freqs = nn.Conv1d(channels, heads * nfreqs, 1)
+        if ndecay:
+            self.query_decay = nn.Conv1d(channels, heads * ndecay, 1)
+            # Initialize decay close to zero (there is a sigmoid), for maximum initial window.
+            self.query_decay.weight.data *= 0.01
+            assert self.query_decay.bias is not None  # stupid type checker
+            self.query_decay.bias.data[:] = -2
+        self.proj = nn.Conv1d(channels + heads * nfreqs, channels, 1)
+
+    def forward(self, x):
+        b, _, t = x.shape
+        heads = self.heads
+        indexes = torch.arange(t, device=x.device, dtype=x.dtype)
+        # left index are keys, right index are queries
+        delta = indexes[:, None] - indexes[None, :]
+
+        queries = self.query(x).view(b, heads, -1, t)
+        keys = self.key(x).view(b, heads, -1, t)
+        # t are keys, s are queries
+        dots = torch.einsum("bhct,bhcs->bhts", keys, queries)
+        dots /= keys.shape[2]**0.5
+        if self.nfreqs:
+            periods = torch.arange(1, self.nfreqs + 1, device=x.device, dtype=x.dtype)
+            freq_kernel = torch.cos(2 * math.pi * delta / periods.view(-1, 1, 1))
+            freq_q = self.query_freqs(x).view(b, heads, -1, t) / self.nfreqs ** 0.5
+            dots += torch.einsum("fts,bhfs->bhts", freq_kernel, freq_q)
+        if self.ndecay:
+            decays = torch.arange(1, self.ndecay + 1, device=x.device, dtype=x.dtype)
+            decay_q = self.query_decay(x).view(b, heads, -1, t)
+            decay_q = torch.sigmoid(decay_q) / 2
+            decay_kernel = - decays.view(-1, 1, 1) * delta.abs() / self.ndecay**0.5
+            dots += torch.einsum("fts,bhfs->bhts", decay_kernel, decay_q)
+
+        # Kill self reference.
+        dots.masked_fill_(torch.eye(t, device=dots.device, dtype=torch.bool), -100)
+        weights = torch.softmax(dots, dim=2)
+
+        content = self.content(x).view(b, heads, -1, t)
+        result = torch.einsum("bhts,bhct->bhcs", weights, content)
+        if self.nfreqs:
+            time_sig = torch.einsum("bhts,fts->bhfs", weights, freq_kernel)
+            result = torch.cat([result, time_sig], 2)
+        result = result.reshape(b, -1, t)
+        return x + self.proj(result)
+
+
+class Demucs(nn.Module):
+    @capture_init
+    def __init__(self,
+                 sources,
+                 # Channels
+                 audio_channels=2,
+                 channels=64,
+                 growth=2.,
+                 # Main structure
+                 depth=6,
+                 rewrite=True,
+                 lstm_layers=0,
+                 # Convolutions
+                 kernel_size=8,
+                 stride=4,
+                 context=1,
+                 # Activations
+                 gelu=True,
+                 glu=True,
+                 # Normalization
+                 norm_starts=4,
+                 norm_groups=4,
+                 # DConv residual branch
+                 dconv_mode=1,
+                 dconv_depth=2,
+                 dconv_comp=4,
+                 dconv_attn=4,
+                 dconv_lstm=4,
+                 dconv_init=1e-4,
+                 # Pre/post processing
+                 normalize=True,
+                 resample=True,
+                 # Weight init
+                 rescale=0.1,
+                 # Metadata
+                 samplerate=44100,
+                 segment=4 * 10):
+        """
+        Args:
+            sources (list[str]): list of source names
+            audio_channels (int): stereo or mono
+            channels (int): first convolution channels
+            depth (int): number of encoder/decoder layers
+            growth (float): multiply (resp divide) number of channels by that
+                for each layer of the encoder (resp decoder)
+            depth (int): number of layers in the encoder and in the decoder.
+            rewrite (bool): add 1x1 convolution to each layer.
+            lstm_layers (int): number of lstm layers, 0 = no lstm. Deactivated
+                by default, as this is now replaced by the smaller and faster small LSTMs
+                in the DConv branches.
+            kernel_size (int): kernel size for convolutions
+            stride (int): stride for convolutions
+            context (int): kernel size of the convolution in the
+                decoder before the transposed convolution. If > 1,
+                will provide some context from neighboring time steps.
+            gelu: use GELU activation function.
+            glu (bool): use glu instead of ReLU for the 1x1 rewrite conv.
+            norm_starts: layer at which group norm starts being used.
+                decoder layers are numbered in reverse order.
+            norm_groups: number of groups for group norm.
+            dconv_mode: if 1: dconv in encoder only, 2: decoder only, 3: both.
+            dconv_depth: depth of residual DConv branch.
+            dconv_comp: compression of DConv branch.
+            dconv_attn: adds attention layers in DConv branch starting at this layer.
+            dconv_lstm: adds a LSTM layer in DConv branch starting at this layer.
+            dconv_init: initial scale for the DConv branch LayerScale.
+            normalize (bool): normalizes the input audio on the fly, and scales back
+                the output by the same amount.
+            resample (bool): upsample x2 the input and downsample /2 the output.
+            rescale (float): rescale initial weights of convolutions
+                to get their standard deviation closer to `rescale`.
+            samplerate (int): stored as meta information for easing
+                future evaluations of the model.
+            segment (float): duration of the chunks of audio to ideally evaluate the model on.
+                This is used by `demucs.apply.apply_model`.
+        """
+
+        super().__init__()
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.resample = resample
+        self.channels = channels
+        self.normalize = normalize
+        self.samplerate = samplerate
+        self.segment = segment
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+        self.skip_scales = nn.ModuleList()
+
+        if glu:
+            activation = nn.GLU(dim=1)
+            ch_scale = 2
+        else:
+            activation = nn.ReLU()
+            ch_scale = 1
+        if gelu:
+            act2 = nn.GELU
+        else:
+            act2 = nn.ReLU
+
+        in_channels = audio_channels
+        padding = 0
+        for index in range(depth):
+            norm_fn = lambda d: nn.Identity()  # noqa
+            if index >= norm_starts:
+                norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+
+            encode = []
+            encode += [
+                nn.Conv1d(in_channels, channels, kernel_size, stride),
+                norm_fn(channels),
+                act2(),
+            ]
+            attn = index >= dconv_attn
+            lstm = index >= dconv_lstm
+            if dconv_mode & 1:
+                encode += [DConv(channels, depth=dconv_depth, init=dconv_init,
+                                 compress=dconv_comp, attn=attn, lstm=lstm)]
+            if rewrite:
+                encode += [
+                    nn.Conv1d(channels, ch_scale * channels, 1),
+                    norm_fn(ch_scale * channels), activation]
+            self.encoder.append(nn.Sequential(*encode))
+
+            decode = []
+            if index > 0:
+                out_channels = in_channels
+            else:
+                out_channels = len(self.sources) * audio_channels
+            if rewrite:
+                decode += [
+                    nn.Conv1d(channels, ch_scale * channels, 2 * context + 1, padding=context),
+                    norm_fn(ch_scale * channels), activation]
+            if dconv_mode & 2:
+                decode += [DConv(channels, depth=dconv_depth, init=dconv_init,
+                                 compress=dconv_comp, attn=attn, lstm=lstm)]
+            decode += [nn.ConvTranspose1d(channels, out_channels,
+                       kernel_size, stride, padding=padding)]
+            if index > 0:
+                decode += [norm_fn(out_channels), act2()]
+            self.decoder.insert(0, nn.Sequential(*decode))
+            in_channels = channels
+            channels = int(growth * channels)
+
+        channels = in_channels
+        if lstm_layers:
+            self.lstm = BLSTM(channels, lstm_layers)
+        else:
+            self.lstm = None
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+    def valid_length(self, length):
+        """
+        Return the nearest valid length to use with the model so that
+        there is no time steps left over in a convolution, e.g. for all
+        layers, size of the input - kernel_size % stride = 0.
+
+        Note that input are automatically padded if necessary to ensure that the output
+        has the same length as the input.
+        """
+        if self.resample:
+            length *= 2
+
+        for _ in range(self.depth):
+            length = math.ceil((length - self.kernel_size) / self.stride) + 1
+            length = max(1, length)
+
+        for _ in range(self.depth):
+            length = (length - 1) * self.stride + self.kernel_size
+
+        if self.resample:
+            length = math.ceil(length / 2)
+        return int(length)
+
+    def forward(self, mix):
+        x = mix
+        length = x.shape[-1]
+
+        if self.normalize:
+            mono = mix.mean(dim=1, keepdim=True)
+            mean = mono.mean(dim=-1, keepdim=True)
+            std = mono.std(dim=-1, keepdim=True)
+            x = (x - mean) / (1e-5 + std)
+        else:
+            mean = 0
+            std = 1
+
+        delta = self.valid_length(length) - length
+        x = F.pad(x, (delta // 2, delta - delta // 2))
+
+        if self.resample:
+            x = julius.resample_frac(x, 1, 2)
+
+        saved = []
+        for encode in self.encoder:
+            x = encode(x)
+            saved.append(x)
+
+        if self.lstm:
+            x = self.lstm(x)
+
+        for decode in self.decoder:
+            skip = saved.pop(-1)
+            skip = center_trim(skip, x)
+            x = decode(x + skip)
+
+        if self.resample:
+            x = julius.resample_frac(x, 2, 1)
+        x = x * std + mean
+        x = center_trim(x, length)
+        x = x.view(x.size(0), len(self.sources), self.audio_channels, x.size(-1))
+        return x
+
+    def load_state_dict(self, state, strict=True):
+        # fix a mismatch with previous generation Demucs models.
+        for idx in range(self.depth):
+            for a in ['encoder', 'decoder']:
+                for b in ['bias', 'weight']:
+                    new = f'{a}.{idx}.3.{b}'
+                    old = f'{a}.{idx}.2.{b}'
+                    if old in state and new not in state:
+                        state[new] = state.pop(old)
+        super().load_state_dict(state, strict=strict)

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/htdemucs.py

@@ -0,0 +1,955 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : htdemucs.py
+@Time    : 2023/8/8 下午4:27
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : The spectrogram and Hybrid version of Demucs
+"""
+
+import math
+import typing as tp
+from copy import deepcopy
+from fractions import Fraction
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from einops import rearrange
+from openunmix.filtering import wiener
+
+from .transformer import CrossTransformerEncoder
+from .demucs import DConv, rescale_module
+from .states import capture_init
+from .spec import spectro, ispectro
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen."""
+    x0 = x
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            extra_pad_right = min(padding_right, extra_pad)
+            extra_pad_left = extra_pad - extra_pad_right
+            paddings = (padding_left - extra_pad_left, padding_right - extra_pad_right)
+            x = F.pad(x, (extra_pad_left, extra_pad_right))
+    out = F.pad(x, paddings, mode, value)
+    assert out.shape[-1] == length + padding_left + padding_right
+    assert (out[..., padding_left: padding_left + length] == x0).all()
+    return out
+
+
+class ScaledEmbedding(nn.Module):
+    """
+    Boost learning rate for embeddings (with `scale`).
+    Also, can make embeddings continuous with `smooth`.
+    """
+    def __init__(self, num_embeddings: int, embedding_dim: int,
+                 scale: float = 10., smooth=False):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+        if smooth:
+            weight = torch.cumsum(self.embedding.weight.data, dim=0)
+            # when summing gaussian, overscale raises as sqrt(n), so we nornalize by that.
+            weight = weight / torch.arange(1, num_embeddings + 1).to(weight).sqrt()[:, None]
+            self.embedding.weight.data[:] = weight
+        self.embedding.weight.data /= scale
+        self.scale = scale
+
+    @property
+    def weight(self):
+        return self.embedding.weight * self.scale
+
+    def forward(self, x):
+        out = self.embedding(x) * self.scale
+        return out
+
+
+class HEncLayer(nn.Module):
+    def __init__(self, chin, chout, kernel_size=8, stride=4, norm_groups=1, empty=False,
+                 freq=True, dconv=True, norm=True, context=0, dconv_kw=None, pad=True,
+                 rewrite=True):
+        """Encoder layer. This used both by the time and the frequency branch.
+        """
+        super().__init__()
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm:
+            norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+        if pad:
+            pad = kernel_size // 4
+        else:
+            pad = 0
+        klass = nn.Conv1d
+        self.freq = freq
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.empty = empty
+        self.norm = norm
+        self.pad = pad
+        if freq:
+            kernel_size = [kernel_size, 1]
+            stride = [stride, 1]
+            pad = [pad, 0]
+            klass = nn.Conv2d
+        self.conv = klass(chin, chout, kernel_size, stride, pad)
+        if self.empty:
+            return
+        self.norm1 = norm_fn(chout)
+        self.rewrite = None
+        if rewrite:
+            self.rewrite = klass(chout, 2 * chout, 1 + 2 * context, 1, context)
+            self.norm2 = norm_fn(2 * chout)
+
+        self.dconv = None
+        if dconv:
+            self.dconv = DConv(chout, **dconv_kw)
+
+    def forward(self, x, inject=None):
+        """
+        `inject` is used to inject the result from the time branch into the frequency branch,
+        when both have the same stride.
+        """
+        if not self.freq and x.dim() == 4:
+            b, c, fr, t = x.shape
+            x = x.view(b, -1, t)
+
+        if not self.freq:
+            le = x.shape[-1]
+            if not le % self.stride == 0:
+                x = F.pad(x, (0, self.stride - (le % self.stride)))
+        y = self.conv(x)
+        if self.empty:
+            return y
+        if inject is not None:
+            assert inject.shape[-1] == y.shape[-1], (inject.shape, y.shape)
+            if inject.dim() == 3 and y.dim() == 4:
+                inject = inject[:, :, None]
+            y = y + inject
+        y = F.gelu(self.norm1(y))
+        if self.dconv:
+            if self.freq:
+                b, c, fr, t = y.shape
+                y = y.permute(0, 2, 1, 3).reshape(-1, c, t)
+            y = self.dconv(y)
+            if self.freq:
+                y = y.view(b, fr, c, t).permute(0, 2, 1, 3)
+        if self.rewrite:
+            z = self.norm2(self.rewrite(y))
+            z = F.glu(z, dim=1)
+        else:
+            z = y
+        return z
+
+
+class MultiWrap(nn.Module):
+    """
+    Takes one layer and replicate it N times. each replica will act
+    on a frequency band. All is done so that if the N replica have the same weights,
+    then this is exactly equivalent to applying the original module on all frequencies.
+    """
+    def __init__(self, layer, split_ratios):
+        super().__init__()
+        self.split_ratios = split_ratios
+        self.layers = nn.ModuleList()
+        self.conv = isinstance(layer, HEncLayer)
+        assert not layer.norm
+        assert layer.freq
+        assert layer.pad
+        if not self.conv:
+            assert not layer.context_freq
+        for _ in range(len(split_ratios) + 1):
+            lay = deepcopy(layer)
+            if self.conv:
+                lay.conv.padding = (0, 0)
+            else:
+                lay.pad = False
+            for m in lay.modules():
+                if hasattr(m, 'reset_parameters'):
+                    m.reset_parameters()
+            self.layers.append(lay)
+
+    def forward(self, x, skip=None, length=None):
+        _, _, fr, _ = x.shape
+
+        ratios = list(self.split_ratios) + [1]
+        start = 0
+        outs = []
+        for ratio, layer in zip(ratios, self.layers):
+            if self.conv:
+                pad = layer.kernel_size // 4
+                if ratio == 1:
+                    limit = fr
+                    frames = -1
+                else:
+                    limit = int(round(fr * ratio))
+                    le = limit - start
+                    if start == 0:
+                        le += pad
+                    frames = round((le - layer.kernel_size) / layer.stride + 1)
+                    limit = start + (frames - 1) * layer.stride + layer.kernel_size
+                    if start == 0:
+                        limit -= pad
+                assert limit - start > 0, (limit, start)
+                assert limit <= fr, (limit, fr)
+                y = x[:, :, start:limit, :]
+                if start == 0:
+                    y = F.pad(y, (0, 0, pad, 0))
+                if ratio == 1:
+                    y = F.pad(y, (0, 0, 0, pad))
+                outs.append(layer(y))
+                start = limit - layer.kernel_size + layer.stride
+            else:
+                if ratio == 1:
+                    limit = fr
+                else:
+                    limit = int(round(fr * ratio))
+                last = layer.last
+                layer.last = True
+
+                y = x[:, :, start:limit]
+                s = skip[:, :, start:limit]
+                out, _ = layer(y, s, None)
+                if outs:
+                    outs[-1][:, :, -layer.stride:] += (
+                        out[:, :, :layer.stride] - layer.conv_tr.bias.view(1, -1, 1, 1))
+                    out = out[:, :, layer.stride:]
+                if ratio == 1:
+                    out = out[:, :, :-layer.stride // 2, :]
+                if start == 0:
+                    out = out[:, :, layer.stride // 2:, :]
+                outs.append(out)
+                layer.last = last
+                start = limit
+        out = torch.cat(outs, dim=2)
+        if not self.conv and not last:
+            out = F.gelu(out)
+        if self.conv:
+            return out
+        else:
+            return out, None
+
+
+class HDecLayer(nn.Module):
+    def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, norm_groups=1, empty=False,
+                 freq=True, dconv=True, norm=True, context=1, dconv_kw=None, pad=True,
+                 context_freq=True, rewrite=True):
+        """
+        Same as HEncLayer but for decoder. See `HEncLayer` for documentation.
+        """
+        super().__init__()
+        norm_fn = lambda d: nn.Identity()  # noqa
+        if norm:
+            norm_fn = lambda d: nn.GroupNorm(norm_groups, d)  # noqa
+        if pad:
+            pad = kernel_size // 4
+        else:
+            pad = 0
+        self.pad = pad
+        self.last = last
+        self.freq = freq
+        self.chin = chin
+        self.empty = empty
+        self.stride = stride
+        self.kernel_size = kernel_size
+        self.norm = norm
+        self.context_freq = context_freq
+        klass = nn.Conv1d
+        klass_tr = nn.ConvTranspose1d
+        if freq:
+            kernel_size = [kernel_size, 1]
+            stride = [stride, 1]
+            klass = nn.Conv2d
+            klass_tr = nn.ConvTranspose2d
+        self.conv_tr = klass_tr(chin, chout, kernel_size, stride)
+        self.norm2 = norm_fn(chout)
+        if self.empty:
+            return
+        self.rewrite = None
+        if rewrite:
+            if context_freq:
+                self.rewrite = klass(chin, 2 * chin, 1 + 2 * context, 1, context)
+            else:
+                self.rewrite = klass(chin, 2 * chin, [1, 1 + 2 * context], 1,
+                                     [0, context])
+            self.norm1 = norm_fn(2 * chin)
+
+        self.dconv = None
+        if dconv:
+            self.dconv = DConv(chin, **dconv_kw)
+
+    def forward(self, x, skip, length):
+        if self.freq and x.dim() == 3:
+            b, c, t = x.shape
+            x = x.view(b, self.chin, -1, t)
+
+        if not self.empty:
+            x = x + skip
+
+            if self.rewrite:
+                y = F.glu(self.norm1(self.rewrite(x)), dim=1)
+            else:
+                y = x
+            if self.dconv:
+                if self.freq:
+                    b, c, fr, t = y.shape
+                    y = y.permute(0, 2, 1, 3).reshape(-1, c, t)
+                y = self.dconv(y)
+                if self.freq:
+                    y = y.view(b, fr, c, t).permute(0, 2, 1, 3)
+        else:
+            y = x
+            assert skip is None
+        z = self.norm2(self.conv_tr(y))
+        if self.freq:
+            if self.pad:
+                z = z[..., self.pad:-self.pad, :]
+        else:
+            z = z[..., self.pad:self.pad + length]
+            assert z.shape[-1] == length, (z.shape[-1], length)
+        if not self.last:
+            z = F.gelu(z)
+        return z, y
+
+
+class HTDemucs(nn.Module):
+    """
+    Spectrogram and hybrid Demucs model.
+    The spectrogram model has the same structure as Demucs, except the first few layers are over the
+    frequency axis, until there is only 1 frequency, and then it moves to time convolutions.
+    Frequency layers can still access information across time steps thanks to the DConv residual.
+
+    Hybrid model have a parallel time branch. At some layer, the time branch has the same stride
+    as the frequency branch and then the two are combined. The opposite happens in the decoder.
+
+    Models can either use naive iSTFT from masking, Wiener filtering ([Ulhih et al. 2017]),
+    or complex as channels (CaC) [Choi et al. 2020]. Wiener filtering is based on
+    Open Unmix implementation [Stoter et al. 2019].
+
+    The loss is always on the temporal domain, by backpropagating through the above
+    output methods and iSTFT. This allows to define hybrid models nicely. However, this breaks
+    a bit Wiener filtering, as doing more iteration at test time will change the spectrogram
+    contribution, without changing the one from the waveform, which will lead to worse performance.
+    I tried using the residual option in OpenUnmix Wiener implementation, but it didn't improve.
+    CaC on the other hand provides similar performance for hybrid, and works naturally with
+    hybrid models.
+
+    This model also uses frequency embeddings are used to improve efficiency on convolutions
+    over the freq. axis, following [Isik et al. 2020] (https://arxiv.org/pdf/2008.04470.pdf).
+
+    Unlike classic Demucs, there is no resampling here, and normalization is always applied.
+    """
+
+    @capture_init
+    def __init__(
+        self,
+        sources,
+        # Channels
+        audio_channels=2,
+        channels=48,
+        channels_time=None,
+        growth=2,
+        # STFT
+        nfft=4096,
+        wiener_iters=0,
+        end_iters=0,
+        wiener_residual=False,
+        cac=True,
+        # Main structure
+        depth=4,
+        rewrite=True,
+        # Frequency branch
+        multi_freqs=None,
+        multi_freqs_depth=3,
+        freq_emb=0.2,
+        emb_scale=10,
+        emb_smooth=True,
+        # Convolutions
+        kernel_size=8,
+        time_stride=2,
+        stride=4,
+        context=1,
+        context_enc=0,
+        # Normalization
+        norm_starts=4,
+        norm_groups=4,
+        # DConv residual branch
+        dconv_mode=1,
+        dconv_depth=2,
+        dconv_comp=8,
+        dconv_init=1e-3,
+        # Before the Transformer
+        bottom_channels=0,
+        # Transformer
+        t_layers=5,
+        t_emb="sin",
+        t_hidden_scale=4.0,
+        t_heads=8,
+        t_dropout=0.0,
+        t_max_positions=10000,
+        t_norm_in=True,
+        t_norm_in_group=False,
+        t_group_norm=False,
+        t_norm_first=True,
+        t_norm_out=True,
+        t_max_period=10000.0,
+        t_weight_decay=0.0,
+        t_lr=None,
+        t_layer_scale=True,
+        t_gelu=True,
+        t_weight_pos_embed=1.0,
+        t_sin_random_shift=0,
+        t_cape_mean_normalize=True,
+        t_cape_augment=True,
+        t_cape_glob_loc_scale=None,
+        t_sparse_self_attn=False,
+        t_sparse_cross_attn=False,
+        t_mask_type="diag",
+        t_mask_random_seed=42,
+        t_sparse_attn_window=500,
+        t_global_window=100,
+        t_sparsity=0.95,
+        t_auto_sparsity=False,
+        # ------ Particuliar parameters
+        t_cross_first=False,
+        # Weight init
+        rescale=0.1,
+        # Metadata
+        samplerate=44100,
+        segment=10,
+        use_train_segment=True,
+    ):
+        """
+        Args:
+            sources (list[str]): list of source names.
+            audio_channels (int): input/output audio channels.
+            channels (int): initial number of hidden channels.
+            channels_time: if not None, use a different `channels` value for the time branch.
+            growth: increase the number of hidden channels by this factor at each layer.
+            nfft: number of fft bins. Note that changing this require careful computation of
+                various shape parameters and will not work out of the box for hybrid models.
+            wiener_iters: when using Wiener filtering, number of iterations at test time.
+            end_iters: same but at train time. For a hybrid model, must be equal to `wiener_iters`.
+            wiener_residual: add residual source before wiener filtering.
+            cac: uses complex as channels, i.e. complex numbers are 2 channels each
+                in input and output. no further processing is done before ISTFT.
+            depth (int): number of layers in the encoder and in the decoder.
+            rewrite (bool): add 1x1 convolution to each layer.
+            multi_freqs: list of frequency ratios for splitting frequency bands with `MultiWrap`.
+            multi_freqs_depth: how many layers to wrap with `MultiWrap`. Only the outermost
+                layers will be wrapped.
+            freq_emb: add frequency embedding after the first frequency layer if > 0,
+                the actual value controls the weight of the embedding.
+            emb_scale: equivalent to scaling the embedding learning rate
+            emb_smooth: initialize the embedding with a smooth one (with respect to frequencies).
+            kernel_size: kernel_size for encoder and decoder layers.
+            stride: stride for encoder and decoder layers.
+            time_stride: stride for the final time layer, after the merge.
+            context: context for 1x1 conv in the decoder.
+            context_enc: context for 1x1 conv in the encoder.
+            norm_starts: layer at which group norm starts being used.
+                decoder layers are numbered in reverse order.
+            norm_groups: number of groups for group norm.
+            dconv_mode: if 1: dconv in encoder only, 2: decoder only, 3: both.
+            dconv_depth: depth of residual DConv branch.
+            dconv_comp: compression of DConv branch.
+            dconv_attn: adds attention layers in DConv branch starting at this layer.
+            dconv_lstm: adds a LSTM layer in DConv branch starting at this layer.
+            dconv_init: initial scale for the DConv branch LayerScale.
+            bottom_channels: if >0 it adds a linear layer (1x1 Conv) before and after the
+                transformer in order to change the number of channels
+            t_layers: number of layers in each branch (waveform and spec) of the transformer
+            t_emb: "sin", "cape" or "scaled"
+            t_hidden_scale: the hidden scale of the Feedforward parts of the transformer
+                for instance if C = 384 (the number of channels in the transformer) and
+                t_hidden_scale = 4.0 then the intermediate layer of the FFN has dimension
+                384 * 4 = 1536
+            t_heads: number of heads for the transformer
+            t_dropout: dropout in the transformer
+            t_max_positions: max_positions for the "scaled" positional embedding, only
+                useful if t_emb="scaled"
+            t_norm_in: (bool) norm before addinf positional embedding and getting into the
+                transformer layers
+            t_norm_in_group: (bool) if True while t_norm_in=True, the norm is on all the
+                timesteps (GroupNorm with group=1)
+            t_group_norm: (bool) if True, the norms of the Encoder Layers are on all the
+                timesteps (GroupNorm with group=1)
+            t_norm_first: (bool) if True the norm is before the attention and before the FFN
+            t_norm_out: (bool) if True, there is a GroupNorm (group=1) at the end of each layer
+            t_max_period: (float) denominator in the sinusoidal embedding expression
+            t_weight_decay: (float) weight decay for the transformer
+            t_lr: (float) specific learning rate for the transformer
+            t_layer_scale: (bool) Layer Scale for the transformer
+            t_gelu: (bool) activations of the transformer are GeLU if True, ReLU else
+            t_weight_pos_embed: (float) weighting of the positional embedding
+            t_cape_mean_normalize: (bool) if t_emb="cape", normalisation of positional embeddings
+                see: https://arxiv.org/abs/2106.03143
+            t_cape_augment: (bool) if t_emb="cape", must be True during training and False
+                during the inference, see: https://arxiv.org/abs/2106.03143
+            t_cape_glob_loc_scale: (list of 3 floats) if t_emb="cape", CAPE parameters
+                see: https://arxiv.org/abs/2106.03143
+            t_sparse_self_attn: (bool) if True, the self attentions are sparse
+            t_sparse_cross_attn: (bool) if True, the cross-attentions are sparse (don't use it
+                unless you designed really specific masks)
+            t_mask_type: (str) can be "diag", "jmask", "random", "global" or any combination
+                with '_' between: i.e. "diag_jmask_random" (note that this is permutation
+                invariant i.e. "diag_jmask_random" is equivalent to "jmask_random_diag")
+            t_mask_random_seed: (int) if "random" is in t_mask_type, controls the seed
+                that generated the random part of the mask
+            t_sparse_attn_window: (int) if "diag" is in t_mask_type, for a query (i), and
+                a key (j), the mask is True id |i-j|<=t_sparse_attn_window
+            t_global_window: (int) if "global" is in t_mask_type, mask[:t_global_window, :]
+                and mask[:, :t_global_window] will be True
+            t_sparsity: (float) if "random" is in t_mask_type, t_sparsity is the sparsity
+                level of the random part of the mask.
+            t_cross_first: (bool) if True cross attention is the first layer of the
+                transformer (False seems to be better)
+            rescale: weight rescaling trick
+            use_train_segment: (bool) if True, the actual size that is used during the
+                training is used during inference.
+        """
+        super().__init__()
+        self.cac = cac
+        self.wiener_residual = wiener_residual
+        self.audio_channels = audio_channels
+        self.sources = sources
+        self.kernel_size = kernel_size
+        self.context = context
+        self.stride = stride
+        self.depth = depth
+        self.bottom_channels = bottom_channels
+        self.channels = channels
+        self.samplerate = samplerate
+        self.segment = segment
+        self.use_train_segment = use_train_segment
+        self.nfft = nfft
+        self.hop_length = nfft // 4
+        self.wiener_iters = wiener_iters
+        self.end_iters = end_iters
+        self.freq_emb = None
+        assert wiener_iters == end_iters
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        self.tencoder = nn.ModuleList()
+        self.tdecoder = nn.ModuleList()
+
+        chin = audio_channels
+        chin_z = chin  # number of channels for the freq branch
+        if self.cac:
+            chin_z *= 2
+        chout = channels_time or channels
+        chout_z = channels
+        freqs = nfft // 2
+
+        for index in range(depth):
+            norm = index >= norm_starts
+            freq = freqs > 1
+            stri = stride
+            ker = kernel_size
+            if not freq:
+                assert freqs == 1
+                ker = time_stride * 2
+                stri = time_stride
+
+            pad = True
+            last_freq = False
+            if freq and freqs <= kernel_size:
+                ker = freqs
+                pad = False
+                last_freq = True
+
+            kw = {
+                "kernel_size": ker,
+                "stride": stri,
+                "freq": freq,
+                "pad": pad,
+                "norm": norm,
+                "rewrite": rewrite,
+                "norm_groups": norm_groups,
+                "dconv_kw": {
+                    "depth": dconv_depth,
+                    "compress": dconv_comp,
+                    "init": dconv_init,
+                    "gelu": True,
+                },
+            }
+            kwt = dict(kw)
+            kwt["freq"] = 0
+            kwt["kernel_size"] = kernel_size
+            kwt["stride"] = stride
+            kwt["pad"] = True
+            kw_dec = dict(kw)
+            multi = False
+            if multi_freqs and index < multi_freqs_depth:
+                multi = True
+                kw_dec["context_freq"] = False
+
+            if last_freq:
+                chout_z = max(chout, chout_z)
+                chout = chout_z
+
+            enc = HEncLayer(
+                chin_z, chout_z, dconv=dconv_mode & 1, context=context_enc, **kw
+            )
+            if freq:
+                tenc = HEncLayer(
+                    chin,
+                    chout,
+                    dconv=dconv_mode & 1,
+                    context=context_enc,
+                    empty=last_freq,
+                    **kwt
+                )
+                self.tencoder.append(tenc)
+
+            if multi:
+                enc = MultiWrap(enc, multi_freqs)
+            self.encoder.append(enc)
+            if index == 0:
+                chin = self.audio_channels * len(self.sources)
+                chin_z = chin
+                if self.cac:
+                    chin_z *= 2
+            dec = HDecLayer(
+                chout_z,
+                chin_z,
+                dconv=dconv_mode & 2,
+                last=index == 0,
+                context=context,
+                **kw_dec
+            )
+            if multi:
+                dec = MultiWrap(dec, multi_freqs)
+            if freq:
+                tdec = HDecLayer(
+                    chout,
+                    chin,
+                    dconv=dconv_mode & 2,
+                    empty=last_freq,
+                    last=index == 0,
+                    context=context,
+                    **kwt
+                )
+                self.tdecoder.insert(0, tdec)
+            self.decoder.insert(0, dec)
+
+            chin = chout
+            chin_z = chout_z
+            chout = int(growth * chout)
+            chout_z = int(growth * chout_z)
+            if freq:
+                if freqs <= kernel_size:
+                    freqs = 1
+                else:
+                    freqs //= stride
+            if index == 0 and freq_emb:
+                self.freq_emb = ScaledEmbedding(
+                    freqs, chin_z, smooth=emb_smooth, scale=emb_scale
+                )
+                self.freq_emb_scale = freq_emb
+
+        if rescale:
+            rescale_module(self, reference=rescale)
+
+        transformer_channels = channels * growth ** (depth - 1)
+        if bottom_channels:
+            self.channel_upsampler = nn.Conv1d(transformer_channels, bottom_channels, 1)
+            self.channel_downsampler = nn.Conv1d(
+                bottom_channels, transformer_channels, 1
+            )
+            self.channel_upsampler_t = nn.Conv1d(
+                transformer_channels, bottom_channels, 1
+            )
+            self.channel_downsampler_t = nn.Conv1d(
+                bottom_channels, transformer_channels, 1
+            )
+
+            transformer_channels = bottom_channels
+
+        if t_layers > 0:
+            if t_cape_glob_loc_scale is None:
+                t_cape_glob_loc_scale = [5000.0, 1.0, 1.4]
+            self.crosstransformer = CrossTransformerEncoder(
+                dim=transformer_channels,
+                emb=t_emb,
+                hidden_scale=t_hidden_scale,
+                num_heads=t_heads,
+                num_layers=t_layers,
+                cross_first=t_cross_first,
+                dropout=t_dropout,
+                max_positions=t_max_positions,
+                norm_in=t_norm_in,
+                norm_in_group=t_norm_in_group,
+                group_norm=t_group_norm,
+                norm_first=t_norm_first,
+                norm_out=t_norm_out,
+                max_period=t_max_period,
+                weight_decay=t_weight_decay,
+                lr=t_lr,
+                layer_scale=t_layer_scale,
+                gelu=t_gelu,
+                sin_random_shift=t_sin_random_shift,
+                weight_pos_embed=t_weight_pos_embed,
+                cape_mean_normalize=t_cape_mean_normalize,
+                cape_augment=t_cape_augment,
+                cape_glob_loc_scale=t_cape_glob_loc_scale,
+                sparse_self_attn=t_sparse_self_attn,
+                sparse_cross_attn=t_sparse_cross_attn,
+                mask_type=t_mask_type,
+                mask_random_seed=t_mask_random_seed,
+                sparse_attn_window=t_sparse_attn_window,
+                global_window=t_global_window,
+                sparsity=t_sparsity,
+                auto_sparsity=t_auto_sparsity,
+            )
+        else:
+            self.crosstransformer = None
+
+    def _spec(self, x):
+        hl = self.hop_length
+        nfft = self.nfft
+
+        # We re-pad the signal in order to keep the property
+        # that the size of the output is exactly the size of the input
+        # divided by the stride (here hop_length), when divisible.
+        # This is achieved by padding by 1/4th of the kernel size (here nfft).
+        # which is not supported by torch.stft.
+        # Having all convolution operations follow this convention allow to easily
+        # align the time and frequency branches later on.
+        assert hl == nfft // 4
+        le = int(math.ceil(x.shape[-1] / hl))
+        pad = hl // 2 * 3
+        x = pad1d(x, (pad, pad + le * hl - x.shape[-1]), mode="reflect")
+
+        z = spectro(x, nfft, hl)[..., :-1, :]
+        assert z.shape[-1] == le + 4, (z.shape, x.shape, le)
+        z = z[..., 2: 2 + le]
+        return z
+
+    def _ispec(self, z, length=None, scale=0):
+        hl = self.hop_length // (4**scale)
+        z = F.pad(z, (0, 0, 0, 1))
+        z = F.pad(z, (2, 2))
+        pad = hl // 2 * 3
+        le = hl * int(math.ceil(length / hl)) + 2 * pad
+        x = ispectro(z, hl, length=le)
+        x = x[..., pad: pad + length]
+        return x
+
+    def _magnitude(self, z):
+        # return the magnitude of the spectrogram, except when cac is True,
+        # in which case we just move the complex dimension to the channel one.
+        if self.cac:
+            b, c, fr, t = z.shape
+            m = torch.view_as_real(z).permute(0, 1, 4, 2, 3)
+            m = m.reshape(b, c * 2, fr, t)
+        else:
+            m = z.abs()
+        return m
+
+    def _mask(self, z, m):
+        # Apply masking given the mixture spectrogram `z` and the estimated mask `m`.
+        # If `cac` is True, `m` is actually a full spectrogram and `z` is ignored.
+        niters = self.wiener_iters
+        if self.cac:
+            b, s, _, fr, t = m.shape
+            out = m.view(b, s, -1, 2, fr, t).permute(0, 1, 2, 4, 5, 3)
+            out = torch.view_as_complex(out.contiguous())
+            return out
+        if self.training:
+            niters = self.end_iters
+        if niters < 0:
+            z = z[:, None]
+            return z / (1e-8 + z.abs()) * m
+        else:
+            return self._wiener(m, z, niters)
+
+    def _wiener(self, mag_out, mix_stft, niters):
+        # apply wiener filtering from OpenUnmix.
+        init = mix_stft.dtype
+        wiener_win_len = 300
+        residual = self.wiener_residual
+
+        b, s, c, fq, t = mag_out.shape
+        mag_out = mag_out.permute(0, 4, 3, 2, 1)
+        mix_stft = torch.view_as_real(mix_stft.permute(0, 3, 2, 1))
+
+        outs = []
+        for sample in range(b):
+            pos = 0
+            out = []
+            for pos in range(0, t, wiener_win_len):
+                frame = slice(pos, pos + wiener_win_len)
+                z_out = wiener(
+                    mag_out[sample, frame],
+                    mix_stft[sample, frame],
+                    niters,
+                    residual=residual,
+                )
+                out.append(z_out.transpose(-1, -2))
+            outs.append(torch.cat(out, dim=0))
+        out = torch.view_as_complex(torch.stack(outs, 0))
+        out = out.permute(0, 4, 3, 2, 1).contiguous()
+        if residual:
+            out = out[:, :-1]
+        assert list(out.shape) == [b, s, c, fq, t]
+        return out.to(init)
+
+    def valid_length(self, length: int):
+        """
+        Return a length that is appropriate for evaluation.
+        In our case, always return the training length, unless
+        it is smaller than the given length, in which case this
+        raises an error.
+        """
+        if not self.use_train_segment:
+            return length
+        training_length = int(self.segment * self.samplerate)
+        if training_length < length:
+            raise ValueError(
+                    f"Given length {length} is longer than "
+                    f"training length {training_length}")
+        return training_length
+
+    def forward(self, mix):
+        length = mix.shape[-1]
+        length_pre_pad = None
+        if self.use_train_segment:
+            if self.training:
+                self.segment = Fraction(mix.shape[-1], self.samplerate)
+            else:
+                training_length = int(self.segment * self.samplerate)
+                if mix.shape[-1] < training_length:
+                    length_pre_pad = mix.shape[-1]
+                    mix = F.pad(mix, (0, training_length - length_pre_pad))
+        z = self._spec(mix)
+        mag = self._magnitude(z).to(mix.device)
+        x = mag
+
+        b, _, fq, t = x.shape
+
+        # unlike previous Demucs, we always normalize because it is easier.
+        mean = x.mean(dim=(1, 2, 3), keepdim=True)
+        std = x.std(dim=(1, 2, 3), keepdim=True)
+        x = (x - mean) / (1e-5 + std)
+        # x will be the freq. branch input.
+
+        # Prepare the time branch input.
+        xt = mix
+        meant = xt.mean(dim=(1, 2), keepdim=True)
+        stdt = xt.std(dim=(1, 2), keepdim=True)
+        xt = (xt - meant) / (1e-5 + stdt)
+
+        # okay, this is a giant mess I know...
+        saved = []  # skip connections, freq.
+        saved_t = []  # skip connections, time.
+        lengths = []  # saved lengths to properly remove padding, freq branch.
+        lengths_t = []  # saved lengths for time branch.
+        for idx, encode in enumerate(self.encoder):
+            lengths.append(x.shape[-1])
+            inject = None
+            if idx < len(self.tencoder):
+                # we have not yet merged branches.
+                lengths_t.append(xt.shape[-1])
+                tenc = self.tencoder[idx]
+                xt = tenc(xt)
+                if not tenc.empty:
+                    # save for skip connection
+                    saved_t.append(xt)
+                else:
+                    # tenc contains just the first conv., so that now time and freq.
+                    # branches have the same shape and can be merged.
+                    inject = xt
+            x = encode(x, inject)
+            if idx == 0 and self.freq_emb is not None:
+                # add frequency embedding to allow for non equivariant convolutions
+                # over the frequency axis.
+                frs = torch.arange(x.shape[-2], device=x.device)
+                emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
+                x = x + self.freq_emb_scale * emb
+
+            saved.append(x)
+        if self.crosstransformer:
+            if self.bottom_channels:
+                _, _, f, _ = x.shape
+                x = rearrange(x, "b c f t-> b c (f t)")
+                x = self.channel_upsampler(x)
+                x = rearrange(x, "b c (f t)-> b c f t", f=f)
+                xt = self.channel_upsampler_t(xt)
+
+            x, xt = self.crosstransformer(x, xt)
+
+            if self.bottom_channels:
+                x = rearrange(x, "b c f t-> b c (f t)")
+                x = self.channel_downsampler(x)
+                x = rearrange(x, "b c (f t)-> b c f t", f=f)
+                xt = self.channel_downsampler_t(xt)
+
+        for idx, decode in enumerate(self.decoder):
+            skip = saved.pop(-1)
+            x, pre = decode(x, skip, lengths.pop(-1))
+            # `pre` contains the output just before final transposed convolution,
+            # which is used when the freq. and time branch separate.
+
+            offset = self.depth - len(self.tdecoder)
+            if idx >= offset:
+                tdec = self.tdecoder[idx - offset]
+                length_t = lengths_t.pop(-1)
+                if tdec.empty:
+                    assert pre.shape[2] == 1, pre.shape
+                    pre = pre[:, :, 0]
+                    xt, _ = tdec(pre, None, length_t)
+                else:
+                    skip = saved_t.pop(-1)
+                    xt, _ = tdec(xt, skip, length_t)
+
+        # Let's make sure we used all stored skip connections.
+        assert len(saved) == 0
+        assert len(lengths_t) == 0
+        assert len(saved_t) == 0
+
+        s = len(self.sources)
+        x = x.view(b, s, -1, fq, t)
+        x = x * std[:, None] + mean[:, None]
+
+        # to cpu as mps doesnt support complex numbers
+        # demucs issue #435 ##432
+        # NOTE: in this case z already is on cpu
+        # TODO: remove this when mps supports complex numbers
+        x_is_mps = x.device.type == "mps"
+        if x_is_mps:
+            x = x.cpu()
+
+        zout = self._mask(z, x)
+        if self.use_train_segment:
+            if self.training:
+                x = self._ispec(zout, length)
+            else:
+                x = self._ispec(zout, training_length)
+        else:
+            x = self._ispec(zout, length)
+
+        # back to mps device
+        if x_is_mps:
+            x = x.to("mps")
+
+        if self.use_train_segment:
+            if self.training:
+                xt = xt.view(b, s, -1, length)
+            else:
+                xt = xt.view(b, s, -1, training_length)
+        else:
+            xt = xt.view(b, s, -1, length)
+        xt = xt * stdt[:, None] + meant[:, None]
+        x = xt + x
+        if length_pre_pad:
+            x = x[..., :length_pre_pad]
+        return x

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/pretrained.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : pretrained.py
+@Time    : 2023/8/8 下午7:22
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Loading pretrained models.
+"""
+from pathlib import Path
+
+import yaml
+
+from .apply import BagOfModels
+from .htdemucs import HTDemucs
+from .states import load_state_dict
+
+
+def add_model_flags(parser):
+    group = parser.add_mutually_exclusive_group(required=False)
+    group.add_argument("-s", "--sig", help="Locally trained XP signature.")
+    group.add_argument("-n", "--name", default=None,
+                       help="Pretrained model name or signature. Default is htdemucs.")
+    parser.add_argument("--repo", type=Path,
+                        help="Folder containing all pre-trained models for use with -n.")
+
+
+def get_model_from_yaml(yaml_file, model_file):
+    bag = yaml.safe_load(open(yaml_file))
+    model = load_state_dict(HTDemucs, model_file)
+    weights = bag.get('weights')
+    segment = bag.get('segment')
+    return BagOfModels([model], weights, segment)

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/spec.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : spec.py
+@Time    : 2023/8/8 下午5:10
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Spec
+"""
+
+import torch as th
+
+
+def spectro(x, n_fft=512, hop_length=None, pad=0):
+    *other, length = x.shape
+    x = x.reshape(-1, length)
+    is_mps = x.device.type == 'mps'
+    if is_mps:
+        x = x.cpu()
+    z = th.stft(x,
+                n_fft * (1 + pad),
+                hop_length or n_fft // 4,
+                window=th.hann_window(n_fft).to(x),
+                win_length=n_fft,
+                normalized=True,
+                center=True,
+                return_complex=True,
+                pad_mode='reflect')
+    _, freqs, frame = z.shape
+    return z.view(*other, freqs, frame)
+
+
+def ispectro(z, hop_length=None, length=None, pad=0):
+    *other, freqs, frames = z.shape
+    n_fft = 2 * freqs - 2
+    z = z.view(-1, freqs, frames)
+    win_length = n_fft // (1 + pad)
+    is_mps = z.device.type == 'mps'
+    if is_mps:
+        z = z.cpu()
+    x = th.istft(z,
+                 n_fft,
+                 hop_length,
+                 window=th.hann_window(win_length).to(z.real),
+                 win_length=win_length,
+                 normalized=True,
+                 length=length,
+                 center=True)
+    _, length = x.shape
+    return x.view(*other, length)

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/states.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : states.py
+@Time    : 2023/8/8 下午7:01
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Utilities to save and load models.
+"""
+import functools
+import inspect
+import warnings
+from pathlib import Path
+from fractions import Fraction
+
+import torch
+
+
+def load_state_dict(net, pth_path):
+    kwargs = {'sources': ['drums', 'bass', 'other', 'vocal'], 'audio_channels': 2, 'samplerate': 44100,
+              'segment': Fraction(39, 5), 'channels': 48, 'channels_time': None, 'growth': 2, 'nfft': 4096,
+              'wiener_iters': 0, 'end_iters': 0, 'wiener_residual': False, 'cac': True, 'depth': 4, 'rewrite': True,
+              'multi_freqs': [], 'multi_freqs_depth': 3, 'freq_emb': 0.2, 'emb_scale': 10, 'emb_smooth': True,
+              'kernel_size': 8, 'stride': 4, 'time_stride': 2, 'context': 1, 'context_enc': 0, 'norm_starts': 4,
+              'norm_groups': 4, 'dconv_mode': 3, 'dconv_depth': 2, 'dconv_comp': 8, 'dconv_init': 0.001,
+              'bottom_channels': 512, 't_layers': 5, 't_hidden_scale': 4.0, 't_heads': 8, 't_dropout': 0.02,
+              't_layer_scale': True, 't_gelu': True, 't_emb': 'sin', 't_max_positions': 10000, 't_max_period': 10000.0,
+              't_weight_pos_embed': 1.0, 't_cape_mean_normalize': True, 't_cape_augment': True,
+              't_cape_glob_loc_scale': [5000.0, 1.0, 1.4], 't_sin_random_shift': 0, 't_norm_in': True,
+              't_norm_in_group': False, 't_group_norm': False, 't_norm_first': True, 't_norm_out': True,
+              't_weight_decay': 0.0, 't_lr': None, 't_sparse_self_attn': False, 't_sparse_cross_attn': False,
+              't_mask_type': 'diag', 't_mask_random_seed': 42, 't_sparse_attn_window': 400, 't_global_window': 100,
+              't_sparsity': 0.95, 't_auto_sparsity': False, 't_cross_first': False, 'rescale': 0.1}
+    model = net(**kwargs)
+    state_dict = torch.load(pth_path)
+    model.load_state_dict(state_dict)
+    return model
+
+
+def load_model(path_or_package, strict=False):
+    """Load a model from the given serialized model, either given as a dict (already loaded)
+    or a path to a file on disk."""
+    if isinstance(path_or_package, dict):
+        package = path_or_package
+    elif isinstance(path_or_package, (str, Path)):
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            path = path_or_package
+            package = torch.load(path, 'cpu')
+    else:
+        raise ValueError(f"Invalid type for {path_or_package}.")
+
+    klass = package["klass"]
+    args = package["args"]
+    kwargs = package["kwargs"]
+
+    if strict:
+        model = klass(*args, **kwargs)
+    else:
+        sig = inspect.signature(klass)
+        for key in list(kwargs):
+            if key not in sig.parameters:
+                warnings.warn("Dropping inexistant parameter " + key)
+                del kwargs[key]
+        model = klass(*args, **kwargs)
+
+    state = package["state"]
+
+    set_state(model, state)
+    return model
+
+
+def get_state(model, quantizer, half=False):
+    """Get the state from a model, potentially with quantization applied.
+    If `half` is True, model are stored as half precision, which shouldn't impact performance
+    but half the state size."""
+    if quantizer is None:
+        dtype = torch.half if half else None
+        state = {k: p.data.to(device='cpu', dtype=dtype) for k, p in model.state_dict().items()}
+    else:
+        state = quantizer.get_quantized_state()
+        state['__quantized'] = True
+    return state
+
+
+def set_state(model, state, quantizer=None):
+    """Set the state on a given model."""
+    if state.get('__quantized'):
+        quantizer.restore_quantized_state(model, state['quantized'])
+    else:
+        model.load_state_dict(state)
+    return state
+
+
+def capture_init(init):
+    @functools.wraps(init)
+    def __init__(self, *args, **kwargs):
+        self._init_args_kwargs = (args, kwargs)
+        init(self, *args, **kwargs)
+
+    return __init__

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/transformer.py

@@ -0,0 +1,765 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : transformer.py
+@Time    : 2023/8/8 下午5:05
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : Transformer
+"""
+import math
+import random
+import typing as tp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from einops import rearrange
+from torch.nn import TransformerEncoderLayer, MultiheadAttention, Linear, LayerNorm
+
+
+def create_sin_embedding(
+    length: int, dim: int, shift: int = 0, device="cpu", max_period=10000
+):
+    # We aim for TBC format
+    assert dim % 2 == 0
+    pos = shift + torch.arange(length, device=device).view(-1, 1, 1)
+    half_dim = dim // 2
+    adim = torch.arange(dim // 2, device=device).view(1, 1, -1)
+    phase = pos / (max_period ** (adim / (half_dim - 1)))
+    return torch.cat(
+        [
+            torch.cos(phase),
+            torch.sin(phase),
+        ],
+        dim=-1,
+    )
+
+
+def create_2d_sin_embedding(d_model, height, width, device="cpu", max_period=10000):
+    """
+    :param d_model: dimension of the model
+    :param height: height of the positions
+    :param width: width of the positions
+    :return: d_model*height*width position matrix
+    """
+    if d_model % 4 != 0:
+        raise ValueError(
+            "Cannot use sin/cos positional encoding with "
+            "odd dimension (got dim={:d})".format(d_model)
+        )
+    pe = torch.zeros(d_model, height, width)
+    # Each dimension use half of d_model
+    d_model = int(d_model / 2)
+    div_term = torch.exp(
+        torch.arange(0.0, d_model, 2) * -(math.log(max_period) / d_model)
+    )
+    pos_w = torch.arange(0.0, width).unsqueeze(1)
+    pos_h = torch.arange(0.0, height).unsqueeze(1)
+    pe[0:d_model:2, :, :] = (
+        torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+    )
+    pe[1:d_model:2, :, :] = (
+        torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+    )
+    pe[d_model::2, :, :] = (
+        torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+    )
+    pe[d_model + 1:: 2, :, :] = (
+        torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+    )
+
+    return pe[None, :].to(device)
+
+
+def create_sin_embedding_cape(
+    length: int,
+    dim: int,
+    batch_size: int,
+    mean_normalize: bool,
+    augment: bool,  # True during training
+    max_global_shift: float = 0.0,  # delta max
+    max_local_shift: float = 0.0,  # epsilon max
+    max_scale: float = 1.0,
+    device: str = "cpu",
+    max_period: float = 10000.0,
+):
+    # We aim for TBC format
+    assert dim % 2 == 0
+    pos = 1.0 * torch.arange(length).view(-1, 1, 1)  # (length, 1, 1)
+    pos = pos.repeat(1, batch_size, 1)  # (length, batch_size, 1)
+    if mean_normalize:
+        pos -= torch.nanmean(pos, dim=0, keepdim=True)
+
+    if augment:
+        delta = np.random.uniform(
+            -max_global_shift, +max_global_shift, size=[1, batch_size, 1]
+        )
+        delta_local = np.random.uniform(
+            -max_local_shift, +max_local_shift, size=[length, batch_size, 1]
+        )
+        log_lambdas = np.random.uniform(
+            -np.log(max_scale), +np.log(max_scale), size=[1, batch_size, 1]
+        )
+        pos = (pos + delta + delta_local) * np.exp(log_lambdas)
+
+    pos = pos.to(device)
+
+    half_dim = dim // 2
+    adim = torch.arange(dim // 2, device=device).view(1, 1, -1)
+    phase = pos / (max_period ** (adim / (half_dim - 1)))
+    return torch.cat(
+        [
+            torch.cos(phase),
+            torch.sin(phase),
+        ],
+        dim=-1,
+    ).float()
+
+
+def get_causal_mask(length):
+    pos = torch.arange(length)
+    return pos > pos[:, None]
+
+
+def get_elementary_mask(
+    t1,
+    t2,
+    mask_type,
+    sparse_attn_window,
+    global_window,
+    mask_random_seed,
+    sparsity,
+    device,
+):
+    """
+    When the input of the Decoder has length T1 and the output T2
+    The mask matrix has shape (T2, T1)
+    """
+    assert mask_type in ["diag", "jmask", "random", "global"]
+
+    if mask_type == "global":
+        mask = torch.zeros(t2, t1, dtype=torch.bool)
+        mask[:, :global_window] = True
+        line_window = int(global_window * t2 / t1)
+        mask[:line_window, :] = True
+
+    if mask_type == "diag":
+
+        mask = torch.zeros(t2, t1, dtype=torch.bool)
+        rows = torch.arange(t2)[:, None]
+        cols = (
+            (t1 / t2 * rows + torch.arange(-sparse_attn_window, sparse_attn_window + 1))
+            .long()
+            .clamp(0, t1 - 1)
+        )
+        mask.scatter_(1, cols, torch.ones(1, dtype=torch.bool).expand_as(cols))
+
+    elif mask_type == "jmask":
+        mask = torch.zeros(t2 + 2, t1 + 2, dtype=torch.bool)
+        rows = torch.arange(t2 + 2)[:, None]
+        t = torch.arange(0, int((2 * t1) ** 0.5 + 1))
+        t = (t * (t + 1) / 2).int()
+        t = torch.cat([-t.flip(0)[:-1], t])
+        cols = (t1 / t2 * rows + t).long().clamp(0, t1 + 1)
+        mask.scatter_(1, cols, torch.ones(1, dtype=torch.bool).expand_as(cols))
+        mask = mask[1:-1, 1:-1]
+
+    elif mask_type == "random":
+        gene = torch.Generator(device=device)
+        gene.manual_seed(mask_random_seed)
+        mask = (
+            torch.rand(t1 * t2, generator=gene, device=device).reshape(t2, t1)
+            > sparsity
+        )
+
+    mask = mask.to(device)
+    return mask
+
+
+def get_mask(
+    t1,
+    t2,
+    mask_type,
+    sparse_attn_window,
+    global_window,
+    mask_random_seed,
+    sparsity,
+    device,
+):
+    """
+    Return a SparseCSRTensor mask that is a combination of elementary masks
+    mask_type can be a combination of multiple masks: for instance "diag_jmask_random"
+    """
+    from xformers.sparse import SparseCSRTensor
+    # create a list
+    mask_types = mask_type.split("_")
+
+    all_masks = [
+        get_elementary_mask(
+            t1,
+            t2,
+            mask,
+            sparse_attn_window,
+            global_window,
+            mask_random_seed,
+            sparsity,
+            device,
+        )
+        for mask in mask_types
+    ]
+
+    final_mask = torch.stack(all_masks).sum(axis=0) > 0
+
+    return SparseCSRTensor.from_dense(final_mask[None])
+
+
+class ScaledEmbedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        scale: float = 1.0,
+        boost: float = 3.0,
+    ):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
+        self.embedding.weight.data *= scale / boost
+        self.boost = boost
+
+    @property
+    def weight(self):
+        return self.embedding.weight * self.boost
+
+    def forward(self, x):
+        return self.embedding(x) * self.boost
+
+
+class LayerScale(nn.Module):
+    """Layer scale from [Touvron et al 2021] (https://arxiv.org/pdf/2103.17239.pdf).
+    This rescales diagonaly residual outputs close to 0 initially, then learnt.
+    """
+
+    def __init__(self, channels: int, init: float = 0, channel_last=False):
+        """
+        channel_last = False corresponds to (B, C, T) tensors
+        channel_last = True corresponds to (T, B, C) tensors
+        """
+        super().__init__()
+        self.channel_last = channel_last
+        self.scale = nn.Parameter(torch.zeros(channels, requires_grad=True))
+        self.scale.data[:] = init
+
+    def forward(self, x):
+        if self.channel_last:
+            return self.scale * x
+        else:
+            return self.scale[:, None] * x
+
+
+class MyGroupNorm(nn.GroupNorm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x):
+        """
+        x: (B, T, C)
+        if num_groups=1: Normalisation on all T and C together for each B
+        """
+        x = x.transpose(1, 2)
+        return super().forward(x).transpose(1, 2)
+
+
+class MyTransformerEncoderLayer(TransformerEncoderLayer):
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation=F.relu,
+        group_norm=0,
+        norm_first=False,
+        norm_out=False,
+        layer_norm_eps=1e-5,
+        layer_scale=False,
+        init_values=1e-4,
+        device=None,
+        dtype=None,
+        sparse=False,
+        mask_type="diag",
+        mask_random_seed=42,
+        sparse_attn_window=500,
+        global_window=50,
+        auto_sparsity=False,
+        sparsity=0.95,
+        batch_first=False,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation=activation,
+            layer_norm_eps=layer_norm_eps,
+            batch_first=batch_first,
+            norm_first=norm_first,
+            device=device,
+            dtype=dtype,
+        )
+        self.sparse = sparse
+        self.auto_sparsity = auto_sparsity
+        if sparse:
+            if not auto_sparsity:
+                self.mask_type = mask_type
+                self.sparse_attn_window = sparse_attn_window
+                self.global_window = global_window
+            self.sparsity = sparsity
+        if group_norm:
+            self.norm1 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm2 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+
+        self.norm_out = None
+        if self.norm_first & norm_out:
+            self.norm_out = MyGroupNorm(num_groups=int(norm_out), num_channels=d_model)
+        self.gamma_1 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+        self.gamma_2 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+
+        if sparse:
+            self.self_attn = MultiheadAttention(
+                d_model, nhead, dropout=dropout, batch_first=batch_first,
+                auto_sparsity=sparsity if auto_sparsity else 0,
+            )
+            self.__setattr__("src_mask", torch.zeros(1, 1))
+            self.mask_random_seed = mask_random_seed
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        """
+        if batch_first = False, src shape is (t, b, c)
+        the case where batch_first=True is not covered
+        """
+        device = src.device
+        x = src
+        t, _, _ = x.shape
+        if self.sparse and not self.auto_sparsity:
+            assert src_mask is None
+            src_mask = self.src_mask
+            if src_mask.shape[-1] != t:
+                src_mask = get_mask(
+                    t,
+                    t,
+                    self.mask_type,
+                    self.sparse_attn_window,
+                    self.global_window,
+                    self.mask_random_seed,
+                    self.sparsity,
+                    device,
+                )
+                self.__setattr__("src_mask", src_mask)
+
+        if self.norm_first:
+            x = x + self.gamma_1(
+                self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)
+            )
+            x = x + self.gamma_2(self._ff_block(self.norm2(x)))
+
+            if self.norm_out:
+                x = self.norm_out(x)
+        else:
+            x = self.norm1(
+                x + self.gamma_1(self._sa_block(x, src_mask, src_key_padding_mask))
+            )
+            x = self.norm2(x + self.gamma_2(self._ff_block(x)))
+
+        return x
+
+
+class CrossTransformerEncoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation=F.relu,
+        layer_norm_eps: float = 1e-5,
+        layer_scale: bool = False,
+        init_values: float = 1e-4,
+        norm_first: bool = False,
+        group_norm: bool = False,
+        norm_out: bool = False,
+        sparse=False,
+        mask_type="diag",
+        mask_random_seed=42,
+        sparse_attn_window=500,
+        global_window=50,
+        sparsity=0.95,
+        auto_sparsity=None,
+        device=None,
+        dtype=None,
+        batch_first=False,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        self.sparse = sparse
+        self.auto_sparsity = auto_sparsity
+        if sparse:
+            if not auto_sparsity:
+                self.mask_type = mask_type
+                self.sparse_attn_window = sparse_attn_window
+                self.global_window = global_window
+            self.sparsity = sparsity
+
+        self.cross_attn: nn.Module
+        self.cross_attn = MultiheadAttention(
+            d_model, nhead, dropout=dropout, batch_first=batch_first)
+        # Implementation of Feedforward model
+        self.linear1 = Linear(d_model, dim_feedforward, **factory_kwargs)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = Linear(dim_feedforward, d_model, **factory_kwargs)
+
+        self.norm_first = norm_first
+        self.norm1: nn.Module
+        self.norm2: nn.Module
+        self.norm3: nn.Module
+        if group_norm:
+            self.norm1 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm2 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm3 = MyGroupNorm(int(group_norm), d_model, eps=layer_norm_eps, **factory_kwargs)
+        else:
+            self.norm1 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm2 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+            self.norm3 = LayerNorm(d_model, eps=layer_norm_eps, **factory_kwargs)
+
+        self.norm_out = None
+        if self.norm_first & norm_out:
+            self.norm_out = MyGroupNorm(num_groups=int(norm_out), num_channels=d_model)
+
+        self.gamma_1 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+        self.gamma_2 = (
+            LayerScale(d_model, init_values, True) if layer_scale else nn.Identity()
+        )
+
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = self._get_activation_fn(activation)
+        else:
+            self.activation = activation
+
+        if sparse:
+            self.cross_attn = MultiheadAttention(
+                d_model, nhead, dropout=dropout, batch_first=batch_first,
+                auto_sparsity=sparsity if auto_sparsity else 0)
+            if not auto_sparsity:
+                self.__setattr__("mask", torch.zeros(1, 1))
+                self.mask_random_seed = mask_random_seed
+
+    def forward(self, q, k, mask=None):
+        """
+        Args:
+            q: tensor of shape (T, B, C)
+            k: tensor of shape (S, B, C)
+            mask: tensor of shape (T, S)
+
+        """
+        device = q.device
+        t, _, _ = q.shape
+        s, _, _ = k.shape
+        if self.sparse and not self.auto_sparsity:
+            assert mask is None
+            mask = self.mask
+            if mask.shape[-1] != s or mask.shape[-2] != t:
+                mask = get_mask(
+                    s,
+                    t,
+                    self.mask_type,
+                    self.sparse_attn_window,
+                    self.global_window,
+                    self.mask_random_seed,
+                    self.sparsity,
+                    device,
+                )
+                self.__setattr__("mask", mask)
+
+        if self.norm_first:
+            x = q + self.gamma_1(self._ca_block(self.norm1(q), self.norm2(k), mask))
+            x = x + self.gamma_2(self._ff_block(self.norm3(x)))
+            if self.norm_out:
+                x = self.norm_out(x)
+        else:
+            x = self.norm1(q + self.gamma_1(self._ca_block(q, k, mask)))
+            x = self.norm2(x + self.gamma_2(self._ff_block(x)))
+
+        return x
+
+    # self-attention block
+    def _ca_block(self, q, k, attn_mask=None):
+        x = self.cross_attn(q, k, k, attn_mask=attn_mask, need_weights=False)[0]
+        return self.dropout1(x)
+
+    # feed forward block
+    def _ff_block(self, x):
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+    def _get_activation_fn(self, activation):
+        if activation == "relu":
+            return F.relu
+        elif activation == "gelu":
+            return F.gelu
+
+        raise RuntimeError("activation should be relu/gelu, not {}".format(activation))
+
+
+# ----------------- MULTI-BLOCKS MODELS: -----------------------
+
+
+class CrossTransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        emb: str = "sin",
+        hidden_scale: float = 4.0,
+        num_heads: int = 8,
+        num_layers: int = 6,
+        cross_first: bool = False,
+        dropout: float = 0.0,
+        max_positions: int = 1000,
+        norm_in: bool = True,
+        norm_in_group: bool = False,
+        group_norm: int = False,
+        norm_first: bool = False,
+        norm_out: bool = False,
+        max_period: float = 10000.0,
+        weight_decay: float = 0.0,
+        lr: tp.Optional[float] = None,
+        layer_scale: bool = False,
+        gelu: bool = True,
+        sin_random_shift: int = 0,
+        weight_pos_embed: float = 1.0,
+        cape_mean_normalize: bool = True,
+        cape_augment: bool = True,
+        cape_glob_loc_scale: list = None,
+        sparse_self_attn: bool = False,
+        sparse_cross_attn: bool = False,
+        mask_type: str = "diag",
+        mask_random_seed: int = 42,
+        sparse_attn_window: int = 500,
+        global_window: int = 50,
+        auto_sparsity: bool = False,
+        sparsity: float = 0.95,
+    ):
+        super().__init__()
+        """
+        """
+        assert dim % num_heads == 0
+
+        hidden_dim = int(dim * hidden_scale)
+
+        self.num_layers = num_layers
+        # classic parity = 1 means that if idx%2 == 1 there is a
+        # classical encoder else there is a cross encoder
+        self.classic_parity = 1 if cross_first else 0
+        self.emb = emb
+        self.max_period = max_period
+        self.weight_decay = weight_decay
+        self.weight_pos_embed = weight_pos_embed
+        self.sin_random_shift = sin_random_shift
+        if emb == "cape":
+            self.cape_mean_normalize = cape_mean_normalize
+            self.cape_augment = cape_augment
+            if cape_glob_loc_scale is None:
+                self.cape_glob_loc_scale = [5000.0, 1.0, 1.4]
+            else:
+                self.cape_glob_loc_scale = cape_glob_loc_scale
+        if emb == "scaled":
+            self.position_embeddings = ScaledEmbedding(max_positions, dim, scale=0.2)
+
+        self.lr = lr
+
+        activation: tp.Any = F.gelu if gelu else F.relu
+
+        self.norm_in: nn.Module
+        self.norm_in_t: nn.Module
+        if norm_in:
+            self.norm_in = LayerNorm(dim)
+            self.norm_in_t = LayerNorm(dim)
+        elif norm_in_group:
+            self.norm_in = MyGroupNorm(int(norm_in_group), dim)
+            self.norm_in_t = MyGroupNorm(int(norm_in_group), dim)
+        else:
+            self.norm_in = nn.Identity()
+            self.norm_in_t = nn.Identity()
+
+        # spectrogram layers
+        self.layers = nn.ModuleList()
+        # temporal layers
+        self.layers_t = nn.ModuleList()
+
+        kwargs_common = {
+            "d_model": dim,
+            "nhead": num_heads,
+            "dim_feedforward": hidden_dim,
+            "dropout": dropout,
+            "activation": activation,
+            "group_norm": group_norm,
+            "norm_first": norm_first,
+            "norm_out": norm_out,
+            "layer_scale": layer_scale,
+            "mask_type": mask_type,
+            "mask_random_seed": mask_random_seed,
+            "sparse_attn_window": sparse_attn_window,
+            "global_window": global_window,
+            "sparsity": sparsity,
+            "auto_sparsity": auto_sparsity,
+            "batch_first": True,
+        }
+
+        kwargs_classic_encoder = dict(kwargs_common)
+        kwargs_classic_encoder.update({
+            "sparse": sparse_self_attn,
+        })
+        kwargs_cross_encoder = dict(kwargs_common)
+        kwargs_cross_encoder.update({
+            "sparse": sparse_cross_attn,
+        })
+
+        for idx in range(num_layers):
+            if idx % 2 == self.classic_parity:
+
+                self.layers.append(MyTransformerEncoderLayer(**kwargs_classic_encoder))
+                self.layers_t.append(
+                    MyTransformerEncoderLayer(**kwargs_classic_encoder)
+                )
+
+            else:
+                self.layers.append(CrossTransformerEncoderLayer(**kwargs_cross_encoder))
+
+                self.layers_t.append(
+                    CrossTransformerEncoderLayer(**kwargs_cross_encoder)
+                )
+
+    def forward(self, x, xt):
+        _, c, fr, t1 = x.shape
+        pos_emb_2d = create_2d_sin_embedding(
+            c, fr, t1, x.device, self.max_period
+        )  # (1, C, Fr, T1)
+        pos_emb_2d = rearrange(pos_emb_2d, "b c fr t1 -> b (t1 fr) c")
+        x = rearrange(x, "b c fr t1 -> b (t1 fr) c")
+        x = self.norm_in(x)
+        x = x + self.weight_pos_embed * pos_emb_2d
+
+        b, c, t2 = xt.shape
+        xt = rearrange(xt, "b c t2 -> b t2 c")  # now T2, B, C
+        pos_emb = self._get_pos_embedding(t2, b, c, x.device)
+        pos_emb = rearrange(pos_emb, "t2 b c -> b t2 c")
+        xt = self.norm_in_t(xt)
+        xt = xt + self.weight_pos_embed * pos_emb
+
+        for idx in range(self.num_layers):
+            if idx % 2 == self.classic_parity:
+                x = self.layers[idx](x)
+                xt = self.layers_t[idx](xt)
+            else:
+                old_x = x
+                x = self.layers[idx](x, xt)
+                xt = self.layers_t[idx](xt, old_x)
+
+        x = rearrange(x, "b (t1 fr) c -> b c fr t1", t1=t1)
+        xt = rearrange(xt, "b t2 c -> b c t2")
+        return x, xt
+
+    def _get_pos_embedding(self, t, b, c, device):
+        if self.emb == "sin":
+            shift = random.randrange(self.sin_random_shift + 1)
+            pos_emb = create_sin_embedding(
+                t, c, shift=shift, device=device, max_period=self.max_period
+            )
+        elif self.emb == "cape":
+            if self.training:
+                pos_emb = create_sin_embedding_cape(
+                    t,
+                    c,
+                    b,
+                    device=device,
+                    max_period=self.max_period,
+                    mean_normalize=self.cape_mean_normalize,
+                    augment=self.cape_augment,
+                    max_global_shift=self.cape_glob_loc_scale[0],
+                    max_local_shift=self.cape_glob_loc_scale[1],
+                    max_scale=self.cape_glob_loc_scale[2],
+                )
+            else:
+                pos_emb = create_sin_embedding_cape(
+                    t,
+                    c,
+                    b,
+                    device=device,
+                    max_period=self.max_period,
+                    mean_normalize=self.cape_mean_normalize,
+                    augment=False,
+                )
+
+        elif self.emb == "scaled":
+            pos = torch.arange(t, device=device)
+            pos_emb = self.position_embeddings(pos)[:, None]
+
+        return pos_emb
+
+    def make_optim_group(self):
+        group = {"params": list(self.parameters()), "weight_decay": self.weight_decay}
+        if self.lr is not None:
+            group["lr"] = self.lr
+        return group
+
+
+def scaled_query_key_softmax(q, k, att_mask):
+    from xformers.ops import masked_matmul
+    q = q / (k.size(-1)) ** 0.5
+    att = masked_matmul(q, k.transpose(-2, -1), att_mask)
+    att = torch.nn.functional.softmax(att, -1)
+    return att
+
+
+def scaled_dot_product_attention(q, k, v, att_mask, dropout):
+    att = scaled_query_key_softmax(q, k, att_mask=att_mask)
+    att = dropout(att)
+    y = att @ v
+    return y
+
+
+def _compute_buckets(x, r):
+    qq = torch.einsum('btf,bfhi->bhti', x, r)
+    qq = torch.cat([qq, -qq], dim=-1)
+    buckets = qq.argmax(dim=-1)
+
+    return buckets.permute(0, 2, 1).byte().contiguous()
+
+
+def dynamic_sparse_attention(query, key, value, sparsity, infer_sparsity=True, attn_bias=None):
+    # assert False, "The code for the custom sparse kernel is not ready for release yet."
+    from xformers.ops import find_locations, sparse_memory_efficient_attention
+    n_hashes = 32
+    proj_size = 4
+    query, key, value = [x.contiguous() for x in [query, key, value]]
+    with torch.no_grad():
+        r = torch.randn(1, query.shape[-1], n_hashes, proj_size // 2, device=query.device)
+        bucket_query = _compute_buckets(query, r)
+        bucket_key = _compute_buckets(key, r)
+        row_offsets, column_indices = find_locations(
+            bucket_query, bucket_key, sparsity, infer_sparsity)
+    return sparse_memory_efficient_attention(
+        query, key, value, row_offsets, column_indices, attn_bias)

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/models/utils.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : utils.py
+@Time    : 2023/8/8 下午4:26
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2023, Tencent
+@Desc    : utils
+"""
+from contextlib import contextmanager
+import math
+import os
+import tempfile
+import typing as tp
+import json
+import subprocess
+
+import torch
+from torch.nn import functional as F
+from torch.utils.data import Subset
+
+
+def unfold(a, kernel_size, stride):
+    """Given input of size [*OT, T], output Tensor of size [*OT, F, K]
+    with K the kernel size, by extracting frames with the given stride.
+
+    This will pad the input so that `F = ceil(T / K)`.
+
+    see https://github.com/pytorch/pytorch/issues/60466
+    """
+    *shape, length = a.shape
+    n_frames = math.ceil(length / stride)
+    tgt_length = (n_frames - 1) * stride + kernel_size
+    a = F.pad(a, (0, tgt_length - length))
+    strides = list(a.stride())
+    assert strides[-1] == 1, 'data should be contiguous'
+    strides = strides[:-1] + [stride, 1]
+    return a.as_strided([*shape, n_frames, kernel_size], strides)
+
+
+def center_trim(tensor: torch.Tensor, reference: tp.Union[torch.Tensor, int]):
+    """
+    Center trim `tensor` with respect to `reference`, along the last dimension.
+    `reference` can also be a number, representing the length to trim to.
+    If the size difference != 0 mod 2, the extra sample is removed on the right side.
+    """
+    ref_size: int
+    if isinstance(reference, torch.Tensor):
+        ref_size = reference.size(-1)
+    else:
+        ref_size = reference
+    delta = tensor.size(-1) - ref_size
+    if delta < 0:
+        raise ValueError("tensor must be larger than reference. " f"Delta is {delta}.")
+    if delta:
+        tensor = tensor[..., delta // 2:-(delta - delta // 2)]
+    return tensor
+
+
+def pull_metric(history: tp.List[dict], name: str):
+    out = []
+    for metrics in history:
+        metric = metrics
+        for part in name.split("."):
+            metric = metric[part]
+        out.append(metric)
+    return out
+
+
+def sizeof_fmt(num: float, suffix: str = 'B'):
+    """
+    Given `num` bytes, return human readable size.
+    Taken from https://stackoverflow.com/a/1094933
+    """
+    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
+        if abs(num) < 1024.0:
+            return "%3.1f%s%s" % (num, unit, suffix)
+        num /= 1024.0
+    return "%.1f%s%s" % (num, 'Yi', suffix)
+
+
+@contextmanager
+def temp_filenames(count: int, delete=True):
+    names = []
+    try:
+        for _ in range(count):
+            names.append(tempfile.NamedTemporaryFile(delete=False).name)
+        yield names
+    finally:
+        if delete:
+            for name in names:
+                os.unlink(name)
+
+
+def random_subset(dataset, max_samples: int, seed: int = 42):
+    if max_samples >= len(dataset):
+        return dataset
+
+    generator = torch.Generator().manual_seed(seed)
+    perm = torch.randperm(len(dataset), generator=generator)
+    return Subset(dataset, perm[:max_samples].tolist())
+
+
+class DummyPoolExecutor:
+    class DummyResult:
+        def __init__(self, func, *args, **kwargs):
+            self.func = func
+            self.args = args
+            self.kwargs = kwargs
+
+        def result(self):
+            return self.func(*self.args, **self.kwargs)
+
+    def __init__(self, workers=0):
+        pass
+
+    def submit(self, func, *args, **kwargs):
+        return DummyPoolExecutor.DummyResult(func, *args, **kwargs)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        return

Levo_Song_Generation/SongGeneration-Runtime/third_party/demucs/run.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    : layers.py
+@Time    : 2024/4/22 下午2:40
+@Author  : waytan
+@Contact : waytan@tencent.com
+@License : (C)Copyright 2024, Tencent
+"""
+import os
+import json
+import time
+import logging
+import argparse
+from datetime import datetime
+
+
+import torch
+
+from models.apply import BagOfModels
+from models.pretrained import get_model_from_yaml
+
+
+class Separator:
+    def __init__(self, dm_model_path, dm_config_path, gpu_id=0) -> None:
+        if torch.cuda.is_available() and gpu_id < torch.cuda.device_count():
+            self.device = torch.device(f"cuda:{gpu_id}")
+        else:
+            self.device = torch.device("cpu")
+        self.demucs_model = self.init_demucs_model(dm_model_path, dm_config_path)
+
+    def init_demucs_model(self, model_path, config_path) -> BagOfModels:
+        model = get_model_from_yaml(config_path, model_path)
+        model.to(self.device)
+        model.eval()
+        return model
+    
+    def run(self, audio_path, output_dir, ext=".flac"):
+        name, _ = os.path.splitext(os.path.split(audio_path)[-1])
+        output_paths = []
+        for stem in self.demucs_model.sources:
+            output_path = os.path.join(output_dir, f"{name}_{stem}{ext}")
+            if os.path.exists(output_path):
+                output_paths.append(output_path)
+        if len(output_paths) == 4:
+            drums_path, bass_path, other_path, vocal_path = output_paths
+        else:
+            drums_path, bass_path, other_path, vocal_path = self.demucs_model.separate(audio_path, output_dir, device=self.device)
+        data_dict = {
+            "vocal_path": vocal_path,
+            "bgm_path": [drums_path, bass_path, other_path]
+        }
+        return data_dict
+
+
+def json_io(input_json, output_json, model_dir, dst_dir, gpu_id=0):
+    current_datetime = datetime.now()
+    current_datetime_str = current_datetime.strftime('%Y-%m-%d-%H:%M')
+    logging.basicConfig(filename=os.path.join(dst_dir, f'logger-separate-{os.path.split(input_json)[1]}-{current_datetime_str}.log'), level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
+    sp = Separator(os.path.join(model_dir, "htdemucs.pth"), os.path.join(model_dir, "htdemucs.yaml"), gpu_id=gpu_id)
+    with open(input_json, "r") as fp:
+        lines = fp.readlines()
+    t1 = time.time()
+    success_num =  0
+    fail_num = 0
+    total_num = len(lines)
+    sep_items = []
+    for line in lines: 
+        item = json.loads(line)
+        flac_file = item["path"]
+        try:
+            fix_data = sp.run(flac_file, dst_dir)
+        except Exception as e:
+            fail_num += 1
+            logging.error(f"process-{success_num + fail_num}/{total_num}|success-{success_num}|fail-{fail_num}|{item['idx']} process fail for {str(e)}")
+            continue
+        
+        item["vocal_path"] = fix_data["vocal_path"]
+        item["bgm_path"] = fix_data["bgm_path"]
+        sep_items.append(item)
+        success_num += 1
+        logging.debug(f"process-{success_num + fail_num}/{total_num}|success-{success_num}|fail-{fail_num}|{item['idx']} process success")
+
+    with open(output_json, "w", encoding='utf-8') as fw:
+        for item in sep_items:
+            fw.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+    t2 = time.time()
+    logging.debug(f"total cost {round(t2-t1, 3)}s")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument("-m", dest="model_dir")
+    parser.add_argument("-d", dest="dst_dir")
+    parser.add_argument("-j", dest="input_json")
+    parser.add_argument("-o", dest="output_json")
+    parser.add_argument("-gid", dest="gpu_id", default=0, type=int)
+    args = parser.parse_args()
+
+    if not args.dst_dir:
+        dst_dir = os.path.join(os.getcwd(), "separate_result")
+        os.makedirs(dst_dir, exist_ok=True)
+    else:
+        dst_dir = os.path.join(args.dst_dir, "separate_result")
+        os.makedirs(dst_dir, exist_ok=True)
+
+    json_io(args.input_json, args.output_json, args.model_dir, dst_dir, gpu_id=args.gpu_id)

Levo_Song_Generation/SongGeneration-Runtime/third_party/hub/version.txt

@@ -0,0 +1 @@
+1

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/config/model_1920.json

@@ -0,0 +1,122 @@
+{
+    "model_type": "autoencoder",
+    "sample_size": 403200,
+    "sample_rate": 48000,
+    "audio_channels": 2,
+    "model": {
+        "encoder": {
+            "type": "oobleck",
+            "config": {
+                "in_channels": 2,
+                "channels": 128,
+                "c_mults": [1, 2, 4, 8, 16],
+                "strides": [2, 4, 4, 6, 10],
+                "latent_dim": 128,
+                "use_snake": true
+            }
+        },
+        "decoder": {
+            "type": "oobleck",
+            "config": {
+                "out_channels": 2,
+                "channels": 128,
+                "c_mults": [1, 2, 4, 8, 16],
+                "strides": [2, 4, 4, 6, 10],
+                "latent_dim": 64,
+                "use_snake": true,
+                "final_tanh": false
+            }
+        },
+        "bottleneck": {
+            "type": "vae"
+        },
+        "latent_dim": 64,
+        "downsampling_ratio": 1920,
+        "io_channels": 2
+    },
+    "training": {
+        "learning_rate": 1.5e-4,
+        "warmup_steps": 0,
+        "use_ema": true,
+        "optimizer_configs": {
+            "autoencoder": {
+                "optimizer": {
+                    "type": "AdamW",
+                    "config": {
+                        "betas": [0.8, 0.99],
+                        "lr": 1.5e-4,
+                        "weight_decay": 1e-3
+                    }
+                },
+                "scheduler": {
+                    "type": "InverseLR",
+                    "config": {
+                        "inv_gamma": 200000,
+                        "power": 0.5,
+                        "warmup": 0.999
+                    }
+                }
+            },
+            "discriminator": {
+                "optimizer": {
+                    "type": "AdamW",
+                    "config": {
+                        "betas": [0.8, 0.99],
+                        "lr": 3e-4,
+                        "weight_decay": 1e-3
+                    }
+                },
+                "scheduler": {
+                    "type": "InverseLR",
+                    "config": {
+                        "inv_gamma": 200000,
+                        "power": 0.5,
+                        "warmup": 0.999
+                    }
+                }
+            }
+        },
+        "loss_configs": {
+            "discriminator": {
+                "type": "encodec",
+                "config": {
+                    "filters": 64,
+                    "n_ffts": [2048, 1024, 512, 256, 128],
+                    "hop_lengths": [512, 256, 128, 64, 32],
+                    "win_lengths": [2048, 1024, 512, 256, 128]
+                },
+                "weights": {
+                    "adversarial": 0.1,
+                    "feature_matching": 5.0
+                }
+            },
+            "spectral": {
+                "type": "mrstft",
+                "config": {
+                    "fft_sizes": [2048, 1024, 512, 256, 128, 64, 32],
+                    "hop_sizes": [512, 256, 128, 64, 32, 16, 8],
+                    "win_lengths": [2048, 1024, 512, 256, 128, 64, 32],
+                    "perceptual_weighting": true
+                },
+                "weights": {
+                    "mrstft": 1.0
+                }
+            },
+            "time": {
+                "type": "l1",
+                "weights": {
+                    "l1": 0.0
+                }
+            },
+            "bottleneck": {
+                "type": "kl",
+                "weights": {
+                    "kl": 1e-4
+                }
+            }
+        },
+        "demo": {
+            "demo_every": 2000
+        }
+    }
+}

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/config/model_config.json

@@ -0,0 +1,122 @@
+{
+    "model_type": "autoencoder",
+    "sample_size": 409600,
+    "sample_rate": 44100,
+    "audio_channels": 2,
+    "model": {
+        "encoder": {
+            "type": "oobleck",
+            "config": {
+                "in_channels": 2,
+                "channels": 128,
+                "c_mults": [1, 2, 4, 8, 16],
+                "strides": [2, 4, 4, 8, 8],
+                "latent_dim": 128,
+                "use_snake": true
+            }
+        },
+        "decoder": {
+            "type": "oobleck",
+            "config": {
+                "out_channels": 2,
+                "channels": 128,
+                "c_mults": [1, 2, 4, 8, 16],
+                "strides": [2, 4, 4, 8, 8],
+                "latent_dim": 64,
+                "use_snake": true,
+                "final_tanh": false
+            }
+        },
+        "bottleneck": {
+            "type": "vae"
+        },
+        "latent_dim": 64,
+        "downsampling_ratio": 2048,
+        "io_channels": 2
+    },
+    "training": {
+        "learning_rate": 1.5e-4,
+        "warmup_steps": 0,
+        "use_ema": true,
+        "optimizer_configs": {
+            "autoencoder": {
+                "optimizer": {
+                    "type": "AdamW",
+                    "config": {
+                        "betas": [0.8, 0.99],
+                        "lr": 1.5e-4,
+                        "weight_decay": 1e-3
+                    }
+                },
+                "scheduler": {
+                    "type": "InverseLR",
+                    "config": {
+                        "inv_gamma": 200000,
+                        "power": 0.5,
+                        "warmup": 0.999
+                    }
+                }
+            },
+            "discriminator": {
+                "optimizer": {
+                    "type": "AdamW",
+                    "config": {
+                        "betas": [0.8, 0.99],
+                        "lr": 3e-4,
+                        "weight_decay": 1e-3
+                    }
+                },
+                "scheduler": {
+                    "type": "InverseLR",
+                    "config": {
+                        "inv_gamma": 200000,
+                        "power": 0.5,
+                        "warmup": 0.999
+                    }
+                }
+            }
+        },
+        "loss_configs": {
+            "discriminator": {
+                "type": "encodec",
+                "config": {
+                    "filters": 64,
+                    "n_ffts": [2048, 1024, 512, 256, 128],
+                    "hop_lengths": [512, 256, 128, 64, 32],
+                    "win_lengths": [2048, 1024, 512, 256, 128]
+                },
+                "weights": {
+                    "adversarial": 0.1,
+                    "feature_matching": 5.0
+                }
+            },
+            "spectral": {
+                "type": "mrstft",
+                "config": {
+                    "fft_sizes": [2048, 1024, 512, 256, 128, 64, 32],
+                    "hop_sizes": [512, 256, 128, 64, 32, 16, 8],
+                    "win_lengths": [2048, 1024, 512, 256, 128, 64, 32],
+                    "perceptual_weighting": true
+                },
+                "weights": {
+                    "mrstft": 1.0
+                }
+            },
+            "time": {
+                "type": "l1",
+                "weights": {
+                    "l1": 0.0
+                }
+            },
+            "bottleneck": {
+                "type": "kl",
+                "weights": {
+                    "kl": 1e-4
+                }
+            }
+        },
+        "demo": {
+            "demo_every": 2000
+        }
+    }
+}

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/docs/autoencoders.md

@@ -0,0 +1,357 @@
+# Autoencoders
+At a high level, autoencoders are models constructed of two parts: an *encoder*, and a *decoder*. 
+
+The *encoder* takes in an sequence (such as mono or stereo audio) and outputs a compressed representation of that sequence as a d-channel "latent sequence", usually heavily downsampled by a constant factor.
+
+The *decoder* takes in a d-channel latent sequence and upsamples it back to the original input sequence length, reversing the compression of the encoder.
+
+Autoencoders are trained with a combination of reconstruction and adversarial losses in order to create a compact and invertible representation of raw audio data that allows downstream models to work in a data-compressed "latent space", with various desirable and controllable properties such as reduced sequence length, noise resistance, and discretization.
+
+The autoencoder architectures defined in `stable-audio-tools` are largely fully-convolutional, which allows autoencoders trained on small lengths to be applied to arbitrary-length sequences. For example, an autoencoder trained on 1-second samples could be used to encode 45-second inputs to a latent diffusion model.
+
+# Model configs
+The model config file for an autoencoder should set the `model_type` to `autoencoder`, and the `model` object should have the following properties:
+
+- `encoder`
+    - Configuration for the autoencoder's encoder half
+- `decoder`
+    - Configuration for the autoencoder's decoder half
+- `latent_dim`
+    - Latent dimension of the autoencoder, used by inference scripts and downstream models
+- `downsampling_ratio`
+    - Downsampling ratio between the input sequence and the latent sequence, used by inference scripts and downstream models
+- `io_channels`
+    - Number of input and output channels for the autoencoder when they're the same, used by inference scripts and downstream models
+- `bottleneck`
+    - Configuration for the autoencoder's bottleneck
+    - Optional
+- `pretransform`
+    - A pretransform definition for the autoencoder, such as wavelet decomposition or another autoencoder
+    - See [pretransforms.md](pretransforms.md) for more information
+    - Optional
+- `in_channels`
+    - Specifies the number of input channels for the autoencoder, when it's different from `io_channels`, such as in a mono-to-stereo model
+    - Optional
+- `out_channels`
+    - Specifies the number of output channels for the autoencoder, when it's different from `io_channels`
+    - Optional
+
+# Training configs
+The `training` config in the autoencoder model config file should have the following properties:
+- `learning_rate`
+    - The learning rate to use during training
+- `use_ema`
+    - If true, a copy of the model weights is maintained during training and updated as an exponential moving average of the trained model's weights. 
+    - Optional. Default: `false`
+- `warmup_steps`
+    - The number of training steps before turning on adversarial losses
+    - Optional. Default: `0`
+- `encoder_freeze_on_warmup`
+    - If true, freezes the encoder after the warmup steps have completed, so adversarial training only affects the decoder.
+    - Optional. Default: `false`
+- `loss_configs`
+    - Configurations for the loss function calculation
+    - Optional
+- `optimizer_configs`
+    - Configuration for optimizers and schedulers
+    - Optional
+
+## Loss configs
+There are few different types of losses that are used for autoencoder training, including spectral losses, time-domain losses, adversarial losses, and bottleneck-specific losses.
+
+Hyperparameters fo these losses as well as loss weighting factors can be configured in the `loss_configs` property in the `training` config.
+
+### Spectral losses
+Multi-resolution STFT losses are the main reconstruction loss used for our audio autoencoders. We use the [auraloss](https://github.com/csteinmetz1/auraloss/tree/main/auraloss) library for our spectral loss functions. 
+
+For mono autoencoders (`io_channels` == 1), we use the [MultiResolutionSTFTLoss](https://github.com/csteinmetz1/auraloss/blob/1576b0cd6e927abc002b23cf3bfc455b660f663c/auraloss/freq.py#L329) module. 
+
+For stereo autoencoders (`io_channels` == 2), we use the [SumAndDifferenceSTFTLoss](https://github.com/csteinmetz1/auraloss/blob/1576b0cd6e927abc002b23cf3bfc455b660f663c/auraloss/freq.py#L533) module. 
+
+#### Example config
+```json
+"spectral": {
+    "type": "mrstft",
+    "config": {
+        "fft_sizes": [2048, 1024, 512, 256, 128, 64, 32],
+        "hop_sizes": [512, 256, 128, 64, 32, 16, 8],
+        "win_lengths": [2048, 1024, 512, 256, 128, 64, 32],
+        "perceptual_weighting": true
+    },
+    "weights": {
+        "mrstft": 1.0
+    }
+}
+```
+
+### Time-domain loss
+We compute the L1 distance between the original audio and the decoded audio to provide a time-domain loss.
+
+#### Example config
+```json
+"time": {
+    "type": "l1",
+    "weights": {
+        "l1": 0.1
+    }
+}
+```
+
+### Adversarial losses
+Adversarial losses bring in an ensemble of discriminator models to discriminate between real and fake audio, providing a signal to the autoencoder on perceptual discrepancies to fix.
+
+We largely rely on the [multi-scale STFT discriminator](https://github.com/facebookresearch/encodec/blob/0e2d0aed29362c8e8f52494baf3e6f99056b214f/encodec/msstftd.py#L99) from the EnCodec repo
+
+#### Example config
+```json
+"discriminator": {
+    "type": "encodec",
+    "config": {
+        "filters": 32,
+        "n_ffts": [2048, 1024, 512, 256, 128],
+        "hop_lengths": [512, 256, 128, 64, 32],
+        "win_lengths": [2048, 1024, 512, 256, 128]
+    },
+    "weights": {
+        "adversarial": 0.1,
+        "feature_matching": 5.0
+    }
+}
+```
+
+## Demo config
+The only property to set for autoencoder training demos is the `demo_every` property, determining the number of steps between demos.
+
+### Example config
+```json
+"demo": {
+    "demo_every": 2000
+}
+```
+
+# Encoder and decoder types
+Encoders and decoders are defined separately in the model configuration, so encoders and decoders from different model architectures and libraries can be used interchangeably. 
+
+## Oobleck
+Oobleck is Harmonai's in-house autoencoder architecture, implementing features from a variety of other autoencoder architectures.
+
+### Example config
+```json
+"encoder": {
+    "type": "oobleck",
+    "config": {
+        "in_channels": 2,
+        "channels": 128,
+        "c_mults": [1, 2, 4, 8],
+        "strides": [2, 4, 8, 8],
+        "latent_dim": 128,
+        "use_snake": true
+    }
+},
+"decoder": {
+    "type": "oobleck",
+    "config": {
+        "out_channels": 2,
+        "channels": 128,
+        "c_mults": [1, 2, 4, 8],
+        "strides": [2, 4, 8, 8],
+        "latent_dim": 64,
+        "use_snake": true,
+        "use_nearest_upsample": false
+    }
+}
+```
+
+## DAC
+This is the Encoder and Decoder definitions from the `descript-audio-codec` repo. It's a simple fully-convolutional autoencoder with channels doubling every level. The encoder and decoder configs are passed directly into the constructors for the DAC [Encoder](https://github.com/descriptinc/descript-audio-codec/blob/c7cfc5d2647e26471dc394f95846a0830e7bec34/dac/model/dac.py#L64) and [Decoder](https://github.com/descriptinc/descript-audio-codec/blob/c7cfc5d2647e26471dc394f95846a0830e7bec34/dac/model/dac.py#L115).
+
+**Note: This does not include the DAC quantizer, and does not load pre-trained DAC models, this is just the encoder and decoder definitions.**
+
+### Example config
+```json
+"encoder": {
+    "type": "dac",
+    "config": {
+        "in_channels": 2,
+        "latent_dim": 32,
+        "d_model": 128,
+        "strides": [2, 4, 4, 4]
+    }
+},
+"decoder": {
+    "type": "dac",
+    "config": {
+        "out_channels": 2,
+        "latent_dim": 32,
+        "channels": 1536,
+        "rates": [4, 4, 4, 2]
+    }
+}
+```
+
+## SEANet
+This is the SEANetEncoder and SEANetDecoder definitions from Meta's EnCodec repo. This is the same encoder and decoder architecture used in the EnCodec models used in MusicGen, without the quantizer.
+
+The encoder and decoder configs are passed directly into the [SEANetEncoder](https://github.com/facebookresearch/encodec/blob/0e2d0aed29362c8e8f52494baf3e6f99056b214f/encodec/modules/seanet.py#L66C12-L66C12) and [SEANetDecoder](https://github.com/facebookresearch/encodec/blob/0e2d0aed29362c8e8f52494baf3e6f99056b214f/encodec/modules/seanet.py#L147) classes directly, though we reverse the input order of the strides (ratios) in the encoder to make it consistent with the order in the decoder.
+
+### Example config
+```json
+"encoder": {
+    "type": "seanet",
+    "config": {
+        "channels": 2,
+        "dimension": 128,
+        "n_filters": 64,
+        "ratios": [4, 4, 8, 8],
+        "n_residual_layers": 1,
+        "dilation_base": 2,
+        "lstm": 2,
+        "norm": "weight_norm"
+    }
+},
+"decoder": {
+    "type": "seanet",
+    "config": {
+        "channels": 2,
+        "dimension": 64,
+        "n_filters": 64,
+        "ratios": [4, 4, 8, 8],
+        "n_residual_layers": 1,
+        "dilation_base": 2,
+        "lstm": 2,
+        "norm": "weight_norm"
+    }
+},
+```
+
+# Bottlenecks
+In our terminology, the "bottleneck" of an autoencoder is a module placed between the encoder and decoder to enforce particular constraints on the latent space the encoder creates.
+
+Bottlenecks have a similar interface to the autoencoder with `encode()` and `decode()` functions defined. Some bottlenecks return extra information in addition to the output latent series, such as quantized token indices, or additional losses to be considered during training.
+
+To define a bottleneck for the autoencoder, you can provide the `bottleneck` object in the autoencoder's model configuration, with the following 
+
+## VAE
+
+The Variational Autoencoder (VAE) bottleneck splits the encoder's output in half along the channel dimension, treats the two halves as the "mean" and "scale" parameters for VAE sampling, and performs the latent sampling. At a basic level, the "scale" values determine the amount of noise to add to the "mean" latents, which creates a noise-resistant latent space where more of the latent space decodes to perceptually "valid" audio. This is particularly helpful for diffusion models where the outpus of the diffusion sampling process leave a bit of Gaussian error noise.
+
+**Note: For the VAE bottleneck to work, the output dimension of the encoder must be twice the size of the input dimension for the decoder.**
+
+### Example config
+```json
+"bottleneck": {
+    "type": "vae"
+}
+```
+
+### Extra info
+The VAE bottleneck also returns a `kl` value in the encoder info. This is the [KL divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence) between encoded/sampled latent space and a Gaussian distribution. By including this value as a loss value to optimize, we push our latent distribution closer to a normal distribution, potentially trading off reconstruction quality.
+
+### Example loss config
+```json
+"bottleneck": {
+    "type": "kl",
+    "weights": {
+        "kl": 1e-4
+    }
+}
+```
+
+## Tanh
+This bottleneck applies the tanh function to the latent series, "soft-clipping" the latent values to be between -1 and 1. This is a quick and dirty way to enforce a limit on the variance of the latent space, but training these models can be unstable as it's seemingly easy for the latent space to saturate the values to -1 or 1 and never recover.
+
+### Example config
+```json
+"bottleneck": {
+    "type": "tanh"
+}
+```
+
+## Wasserstein
+The Wasserstein bottleneck implements the WAE-MMD regularization method from the [Wasserstein Auto-Encoders](https://arxiv.org/abs/1711.01558) paper, calculating the Maximum Mean Discrepancy (MMD) between the latent space and a Gaussian distribution. Including this value as a loss value to optimize leads to a more Gaussian latent space, but does not require stochastic sampling as with a VAE, so the encoder is deterministic.
+
+The Wasserstein bottleneck also exposes the `noise_augment_dim` property, which concatenates `noise_augment_dim` channels of Gaussian noise to the latent series before passing into the decoder. This adds some stochasticity to the latents which can be helpful for adversarial training, while keeping the encoder outputs deterministic.
+
+**Note: The MMD calculation is very VRAM-intensive for longer sequence lengths, so training a Wasserstein autoencoder is best done on autoencoders with a decent downsampling factor, or on short sequence lengths. For inference, the MMD calculation is disabled.**
+
+### Example config
+```json
+"bottleneck": {
+    "type": "wasserstein"
+}
+```
+
+### Extra info
+This bottleneck adds the `mmd` value to the encoder info, representing the Maximum Mean Discrepancy.
+
+### Example loss config
+```json
+"bottleneck": {
+    "type": "mmd",
+    "weights": {
+        "mmd": 100
+    }
+}
+```
+
+## L2 normalization (Spherical autoencoder)
+The L2 normalization bottleneck normalizes the latents across the channel-dimension, projecting the latents to a d-dimensional hypersphere. This acts as a form of latent space normalization.
+
+
+### Example config
+```json
+"bottleneck": {
+    "type": "l2_norm"
+}
+```
+
+
+## RVQ
+Residual vector quantization (RVQ) is currently the leading method for learning discrete neural audio codecs (tokenizers for audio). In vector quantization, each item in the latent sequence is individually "snapped" to the nearest vector in a discrete "codebook" of learned vectors. The index of the vector in the codebook can then be used as a token index for things like autoregressive transformers. Residual vector quantization improves the precision of normal vector quantization by adding additional codebooks. For a deeper dive into RVQ, check out [this blog post by Dr. Scott Hawley](https://drscotthawley.github.io/blog/posts/2023-06-12-RVQ.html).
+
+This RVQ bottleneck uses [lucidrains' implementation](https://github.com/lucidrains/vector-quantize-pytorch/tree/master) from the `vector-quantize-pytorch` repo, which provides a lot of different quantizer options. The bottleneck config is passed through to the `ResidualVQ`  [constructor](https://github.com/lucidrains/vector-quantize-pytorch/blob/0c6cea24ce68510b607f2c9997e766d9d55c085b/vector_quantize_pytorch/residual_vq.py#L26).
+
+**Note: This RVQ implementation uses manual replacement of codebook vectors to reduce codebook collapse. This does not work with multi-GPU training as the random replacement is not synchronized across devices.** 
+
+### Example config
+```json
+"bottleneck": {
+    "type": "rvq",
+    "config": {
+        "num_quantizers": 4,
+        "codebook_size": 2048,
+        "dim": 1024,
+        "decay": 0.99,
+    }
+}
+```
+
+## DAC RVQ
+This is the residual vector quantization implementation from the `descript-audio-codec` repo. It differs from the above implementation in that it does not use manual replacements to improve codebook usage, but instead uses learnable linear layers to project the latents down to a lower-dimensional space before performing the individual quantization operations. This means it's compatible with distributed training. 
+
+The bottleneck config is passed directly into the `ResidualVectorQuantize` [constructor](https://github.com/descriptinc/descript-audio-codec/blob/c7cfc5d2647e26471dc394f95846a0830e7bec34/dac/nn/quantize.py#L97).
+
+The `quantize_on_decode` property is also exposed, which moves the quantization process to the decoder. This should not be used during training, but is helpful when training latent diffusion models that use the quantization process as a way to remove error after the diffusion sampling process.
+
+### Example config
+```json
+"bottleneck": {
+    "type": "dac_rvq",
+    "config": {
+        "input_dim": 64,
+        "n_codebooks": 9,
+        "codebook_dim": 32,
+        "codebook_size": 1024,
+        "quantizer_dropout": 0.5
+    }
+}
+```
+
+### Extra info
+The DAC RVQ bottleneck also adds the following properties to the `info` object:
+- `pre_quantizer`
+    - The pre-quantization latent series, useful in combination with `quantize_on_decode` for training latent diffusion models.
+- `vq/commitment_loss`
+    - Commitment loss for the quantizer
+- `vq/codebook_loss`
+    - Codebook loss for the quantizer

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/docs/conditioning.md

@@ -0,0 +1,158 @@
+# Conditioning
+Conditioning, in the context of `stable-audio-tools` is the use of additional signals in a model that are used to add an additional level of control over the model's behavior. For example, we can condition the outputs of a diffusion model on a text prompt, creating a text-to-audio model.
+
+# Conditioning types
+There are a few different kinds of conditioning depending on the conditioning signal being used.
+
+## Cross attention
+Cross attention is a type of conditioning that allows us to find correlations between two sequences of potentially different lengths. For example, cross attention allows us to find correlations between a sequence of features from a text encoder and a sequence of high-level audio features.
+
+Signals used for cross-attention conditioning should be of the shape `[batch, sequence, channels]`.
+
+## Global conditioning
+Global conditioning is the use of a single n-dimensional tensor to provide conditioning information that pertains to the whole sequence being conditioned. For example, this could be the single embedding output of a CLAP model, or a learned class embedding.
+
+Signals used for global conditioning should be of the shape `[batch, channels]`.
+
+## Prepend conditioning
+Prepend conditioning involves prepending the conditioning tokens to the data tokens in the model, allowing for the information to be interpreted through the model's self-attention mechanism.
+
+This kind of conditioning is currently only supported by Transformer-based models such as diffusion transformers.
+
+Signals used for prepend conditioning should be of the shape `[batch, sequence, channels]`.
+
+## Input concatenation
+Input concatenation applies a spatial conditioning signal to the model that correlates in the sequence dimension with the model's input, and is of the same length. The conditioning signal will be concatenated with the model's input data along the channel dimension. This can be used for things like inpainting information, melody conditioning, or for creating a diffusion autoencoder.
+
+Signals used for input concatenation conditioning should be of the shape `[batch, channels, sequence]` and must be the same length as the model's input.
+
+# Conditioners and conditioning configs
+`stable-audio-tools` uses Conditioner modules to translate human-readable metadata such as text prompts or a number of seconds into tensors that the model can take as input. 
+
+Each conditioner has a corresponding `id` that it expects to find in the conditioning dictionary provided during training or inference. Each conditioner takes in the relevant conditioning data and returns a tuple containing the corresponding tensor and a mask.
+
+The ConditionedDiffusionModelWrapper manages the translation between the user-provided metadata dictionary (e.g. `{"prompt": "a beautiful song", "seconds_start": 22, "seconds_total": 193}`) and the dictionary of different conditioning types that the model uses (e.g. `{"cross_attn_cond": ...}`).
+
+To apply conditioning to a model, you must provide a `conditioning` configuration in the model's config. At the moment, we only support conditioning diffusion models though the `diffusion_cond` model type.
+
+The `conditioning` configuration should contain a `configs` array, which allows you to define multiple conditioning signals. 
+
+Each item in `configs` array should define the `id` for the corresponding metadata, the type of conditioner to be used, and the config for that conditioner.
+
+The `cond_dim` property is used to enforce the same dimension on all conditioning inputs, however that can be overridden with an explicit `output_dim` property on any of the individual configs.
+
+## Example config
+```json
+"conditioning": {
+    "configs": [
+        {
+            "id": "prompt",
+            "type": "t5",
+            "config": {
+                "t5_model_name": "t5-base",
+                "max_length": 77,
+                "project_out": true
+            }
+        }
+    ],
+    "cond_dim": 768
+}
+```
+
+# Conditioners
+
+## Text encoders
+
+### `t5`
+This uses a frozen [T5](https://huggingface.co/docs/transformers/model_doc/t5) text encoder from the `transformers` library to encode text prompts into a sequence of text features.
+
+The `t5_model_name` property determines which T5 model is loaded from the `transformers` library.
+
+The `max_length` property determines the maximum number of tokens that the text encoder will take in, as well as the sequence length of the output text features.
+
+If you set `enable_grad` to `true`, the T5 model will be un-frozen and saved with the model checkpoint, allowing you to fine-tune the T5 model.
+
+T5 encodings are only compatible with cross attention conditioning.
+
+#### Example config 
+```json
+{
+    "id": "prompt",
+    "type": "t5",
+    "config": {
+        "t5_model_name": "t5-base",
+        "max_length": 77,
+        "project_out": true
+    }
+}
+```
+
+### `clap_text`
+This loads the text encoder from a [CLAP](https://github.com/LAION-AI/CLAP) model, which can provide either a sequence of text features, or a single multimodal text/audio embedding.
+
+The CLAP model must be provided with a local file path, set in the `clap_ckpt_path` property,along with the correct `audio_model_type` and `enable_fusion` properties for the provided model.
+
+If the `use_text_features` property is set to `true`, the conditioner output will be a sequence of text features, instead of a single multimodal embedding. This allows for more fine-grained text information to be used by the model, at the cost of losing the ability to prompt with CLAP audio embeddings.
+
+By default, if `use_text_features` is true, the last layer of the CLAP text encoder's features are returned. You can return the text features of earlier layers by specifying the index of the layer to return in the `feature_layer_ix` property. For example, you can return the text features of the next-to-last layer of the CLAP model by setting `feature_layer_ix` to `-2`.
+
+If you set `enable_grad` to `true`, the CLAP model will be un-frozen and saved with the model checkpoint, allowing you to fine-tune the CLAP model.
+
+CLAP text embeddings are compatible with global conditioning and cross attention conditioning. If `use_text_features` is set to `true`, the features are not compatible with global conditioning.
+
+#### Example config
+```json
+{
+    "id": "prompt",
+    "type": "clap_text",
+    "config": {
+        "clap_ckpt_path": "/path/to/clap/model.ckpt",
+        "audio_model_type": "HTSAT-base",
+        "enable_fusion": true,
+        "use_text_features": true,
+        "feature_layer_ix": -2
+    }
+}
+```
+
+## Number encoders
+
+### `int`
+The IntConditioner takes in a list of integers in a given range, and returns a discrete learned embedding for each of those integers.
+
+The `min_val` and `max_val` properties set the range of the embedding values. Input integers are clamped to this range.
+
+This can be used for things like discrete timing embeddings, or learned class embeddings.
+
+Int embeddings are compatible with global conditioning and cross attention conditioning.
+
+#### Example config
+```json
+{
+    "id": "seconds_start",
+    "type": "int",
+    "config": {
+        "min_val": 0,
+        "max_val": 512
+    }
+}
+```
+
+### `number`
+The NumberConditioner takes in a a list of floats in a given range, and returns a continuous Fourier embedding of the provided floats.
+
+The `min_val` and `max_val` properties set the range of the float values. This is the range used to normalize the input float values.
+
+Number embeddings are compatible with global conditioning and cross attention conditioning.
+
+#### Example config
+```json
+{
+    "id": "seconds_total",
+    "type": "number",
+    "config": {
+        "min_val": 0,
+        "max_val": 512
+    }
+}
+```

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/docs/datasets.md

@@ -0,0 +1,75 @@
+# Datasets
+`stable-audio-tools` supports loading data from local file storage, as well as loading audio files and JSON files in the [WebDataset](https://github.com/webdataset/webdataset/tree/main/webdataset) format from Amazon S3 buckets.
+
+# Dataset configs
+To specify the dataset used for training, you must provide a dataset config JSON file to `train.py`.
+
+The dataset config consists of a `dataset_type` property specifying the type of data loader to use, a `datasets` array to provide multiple data sources, and a `random_crop` property, which decides if the cropped audio from the training samples is from a random place in the audio file, or always from the beginning.
+
+## Local audio files
+To use a local directory of audio samples, set the `dataset_type` property in your dataset config to `"audio_dir"`, and provide a list of objects to the `datasets` property including the `path` property, which should be the path to your directory of audio samples.
+
+This will load all of the compatible audio files from the provided directory and all subdirectories.
+
+### Example config 
+```json
+{
+    "dataset_type": "audio_dir",
+    "datasets": [
+        {
+            "id": "my_audio",
+            "path": "/path/to/audio/dataset/"
+        }
+    ],
+    "random_crop": true
+}
+```
+
+## S3 WebDataset
+To load audio files and related metadata from .tar files in the WebDataset format hosted in Amazon S3 buckets, you can set the `dataset_type` property to `s3`, and provide the `datasets` parameter with a list of objects containing the AWS S3 path to the shared S3 bucket prefix of the WebDataset .tar files. The S3 bucket will be searched recursively given the path, and assumes any .tar files found contain audio files and corresponding JSON files where the related files differ only in file extension (e.g. "000001.flac", "000001.json", "00002.flac", "00002.json", etc.)
+
+### Example config
+```json
+{
+    "dataset_type": "s3",
+    "datasets": [
+        {
+            "id": "s3-test",
+            "s3_path": "s3://my-bucket/datasets/webdataset/audio/"
+        }
+    ],
+    "random_crop": true
+}
+```
+
+# Custom metadata
+To customize the metadata provided to the conditioners during model training, you can provide a separate custom metadata module to the dataset config. This metadata module should be a Python file that must contain a function called `get_custom_metadata` that takes in two parameters, `info`, and `audio`, and returns a dictionary. 
+
+For local training, the `info` parameter will contain a few pieces of information about the loaded audio file, such as the path, and information about how the audio was cropped from the original training sample. For WebDataset datasets, it will also contain the metadata from the related JSON files. 
+
+The `audio` parameter contains the audio sample that will be passed to the model at training time. This lets you analyze the audio for extra properties that you can then pass in as extra conditioning signals.
+
+The dictionary returned from the `get_custom_metadata` function will have its properties added to the `metadata` object used at training time. For more information on how conditioning works, please see the [Conditioning documentation](./conditioning.md)
+
+## Example config and custom metadata module
+```json
+{
+    "dataset_type": "audio_dir",
+    "datasets": [
+        {
+            "id": "my_audio",
+            "path": "/path/to/audio/dataset/",
+            "custom_metadata_module": "/path/to/custom_metadata.py",
+        }
+    ],
+    "random_crop": true
+}
+```
+
+`custom_metadata.py`:
+```py
+def get_custom_metadata(info, audio):
+
+    # Pass in the relative path of the audio file as the prompt
+    return {"prompt": info["relpath"]}
+```

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/docs/diffusion.md

@@ -0,0 +1,153 @@
+# Diffusion
+
+Diffusion models learn to denoise data
+
+# Model configs
+The model config file for a diffusion model should set the `model_type` to `diffusion_cond` if the model uses conditioning, or `diffusion_uncond` if it does not, and the `model` object should have the following properties:
+
+- `diffusion`
+    - The configuration for the diffusion model itself. See below for more information on the diffusion model config
+- `pretransform`
+    - The configuration of the diffusion model's [pretransform](pretransforms.md), such as an autoencoder for latent diffusion.
+    - Optional
+- `conditioning`
+    - The configuration of the various [conditioning](conditioning.md) modules for the diffusion model
+    - Only required for `diffusion_cond`
+- `io_channels`
+    - The base number of input/output channels for the diffusion model
+    - Used by inference scripts to determine the shape of the noise to generate for the diffusion model
+
+# Diffusion configs
+- `type`
+    - The underlying model type for the transformer
+    - For conditioned diffusion models, be one of `dit` ([Diffusion Transformer](#diffusion-transformers-dit)), `DAU1d` ([Dance Diffusion U-Net](#dance-diffusion-u-net)), or `adp_cfg_1d` ([audio-diffusion-pytorch U-Net](#audio-diffusion-pytorch-u-net-adp))
+    - Unconditioned diffusion models can also use `adp_1d`
+- `cross_attention_cond_ids`
+    - Conditioner ids for conditioning information to be used as cross-attention input
+    - If multiple ids are specified, the conditioning tensors will be concatenated along the sequence dimension
+- `global_cond_ids`
+    - Conditioner ids for conditioning information to be used as global conditioning input
+    - If multiple ids are specified, the conditioning tensors will be concatenated along the channel dimension
+- `prepend_cond_ids`
+    - Conditioner ids for conditioning information to be prepended to the model input
+    - If multiple ids are specified, the conditioning tensors will be concatenated along the sequence dimension
+    - Only works with diffusion transformer models
+- `input_concat_ids`
+    - Conditioner ids for conditioning information to be concatenated to the model input
+    - If multiple ids are specified, the conditioning tensors will be concatenated along the channel dimension
+    - If the conditioning tensors are not the same length as the model input, they will be interpolated along the sequence dimension to be the same length.
+        - The interpolation algorithm is model-dependent, but usually uses nearest-neighbor resampling.
+- `config`
+    - The configuration for the model backbone itself
+    - Model-dependent
+
+# Training configs
+The `training` config in the diffusion model config file should have the following properties:
+
+- `learning_rate`
+    - The learning rate to use during training
+    - Defaults to constant learning rate, can be overridden with `optimizer_configs`
+- `use_ema`
+    - If true, a copy of the model weights is maintained during training and updated as an exponential moving average of the trained model's weights. 
+    - Optional. Default: `true`
+- `log_loss_info`
+    - If true, additional diffusion loss info will be gathered across all GPUs and displayed during training
+    - Optional. Default: `false`
+- `loss_configs`
+    - Configurations for the loss function calculation
+    - Optional
+- `optimizer_configs`
+    - Configuration for optimizers and schedulers
+    - Optional, overrides `learning_rate`
+- `demo`
+    - Configuration for the demos during training, including conditioning information
+
+## Example config
+```json
+"training": {
+    "use_ema": true,
+    "log_loss_info": false,
+    "optimizer_configs": {
+        "diffusion": {
+            "optimizer": {
+                "type": "AdamW",
+                "config": {
+                    "lr": 5e-5,
+                    "betas": [0.9, 0.999],
+                    "weight_decay": 1e-3
+                }
+            },
+            "scheduler": {
+                "type": "InverseLR",
+                "config": {
+                    "inv_gamma": 1000000,
+                    "power": 0.5,
+                    "warmup": 0.99
+                }
+            }
+        }
+    },
+    "demo": { ... }
+}
+```
+
+# Demo configs
+The `demo` config in the diffusion model training config should have the following properties:
+- `demo_every`
+    - How many training steps between demos
+- `demo_steps`
+    - Number of diffusion timesteps to run for the demos
+- `num_demos`
+    - This is the number of examples to generate in each demo
+- `demo_cond`
+    - For conditioned diffusion models, this is the conditioning metadata to provide to each example, provided as a list
+    - NOTE: List must be the same length as `num_demos`
+- `demo_cfg_scales`
+    - For conditioned diffusion models, this provides a list of classifier-free guidance (CFG) scales to render during the demos. This can be helpful to get an idea of how the model responds to different conditioning strengths as training continues.
+
+## Example config
+```json
+"demo": {
+    "demo_every": 2000,
+    "demo_steps": 250,
+    "num_demos": 4,
+    "demo_cond": [
+        {"prompt": "A beautiful piano arpeggio", "seconds_start": 0, "seconds_total": 80},
+        {"prompt": "A tropical house track with upbeat melodies, a driving bassline, and cheery vibes", "seconds_start": 0, "seconds_total": 250},
+        {"prompt": "A cool 80s glam rock song with driving drums and distorted guitars", "seconds_start": 0, "seconds_total": 180},
+        {"prompt": "A grand orchestral arrangement", "seconds_start": 0, "seconds_total": 190}
+    ],
+    "demo_cfg_scales": [3, 6, 9]
+}
+```
+
+# Model types
+
+A variety of different model types can be used as the underlying backbone for a diffusion model. At the moment, this includes variants of U-Net and Transformer models.
+
+## Diffusion Transformers (DiT)
+
+Transformers tend to consistently outperform U-Nets in terms of model quality, but are much more memory- and compute-intensive and work best on shorter sequences such as latent encodings of audio.
+
+### Continuous Transformer
+
+This is our custom implementation of a transformer model, based on the `x-transformers` implementation, but with efficiency improvements such as fused QKV layers, and Flash Attention 2 support.
+
+### `x-transformers`
+
+This model type uses the `ContinuousTransformerWrapper` class from the https://github.com/lucidrains/x-transformers repository as the diffusion transformer backbone.
+
+`x-transformers` is a great baseline transformer implementation with lots of options for various experimental settings.
+It's great for testing out experimental features without implementing them yourself, but the implementations might not be fully optimized, and breaking changes may be introduced without much warning. 
+
+## Diffusion U-Net
+
+U-Nets use a hierarchical architecture to gradually downsample the input data before more heavy processing is performed, then upsample the data again, using skip connections to pass data across the downsampling "valley" (the "U" in the name) to the upsampling layer at the same resolution. 
+
+### audio-diffusion-pytorch U-Net (ADP)
+
+This model type uses a modified implementation of the `UNetCFG1D` class from version 0.0.94 of the `https://github.com/archinetai/audio-diffusion-pytorch` repo, with added Flash Attention support.
+
+### Dance Diffusion U-Net
+
+This is a reimplementation of the U-Net used in [Dance Diffusion](https://github.com/Harmonai-org/sample-generator). It has minimal conditioning support, only really supporting global conditioning. Mostly used for unconditional diffusion models.

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/docs/pretransforms.md

@@ -0,0 +1,43 @@
+# Pretransforms
+Many models require some fixed transform to be applied to the input audio before the audio is passed in to the trainable layers of the model, as well as a corresponding inverse transform to be applied to the outputs of the model. We refer to these as "pretransforms".
+
+At the moment, `stable-audio-tools` supports two pretransforms, frozen autoencoders for latent diffusion models and wavelet decompositions.
+
+Pretransforms have a similar interface to autoencoders with "encode" and "decode" functions defined for each pretransform.
+
+## Autoencoder pretransform
+To define a model with an autoencoder pretransform, you can define the "pretransform" property in the model config, with the `type` property set to `autoencoder`. The `config` property should be an autoencoder model definition.
+
+Example:
+```json
+"pretransform": {
+    "type": "autoencoder",
+    "config": {
+        "encoder": {
+            ...
+        },
+        "decoder": {
+            ...
+        }
+        ...normal autoencoder configuration
+    }
+}
+```
+
+### Latent rescaling
+The original [Latent Diffusion paper](https://arxiv.org/abs/2112.10752) found that rescaling the latent series to unit variance before performing diffusion improved quality. To this end, we expose a `scale` property on autoencoder pretransforms that will take care of this rescaling. The scale should be set to the original standard deviation of the latents, which can be determined experimentally, or by looking at the `latent_std` value during training. The pretransform code will divide by this scale factor in the `encode` function and multiply by this scale in the `decode` function.
+
+## Wavelet pretransform
+`stable-audio-tools` also exposes wavelet decomposition as a pretransform. Wavelet decomposition is a quick way to trade off sequence length for channels in autoencoders, while maintaining a multi-band implicit bias.
+
+Wavelet pretransforms take the following properties:
+
+- `channels`
+    - The number of input and output audio channels for the wavelet transform
+- `levels`
+    - The number of successive wavelet decompositions to perform. Each level doubles the channel count and halves the sequence length
+- `wavelet`
+    - The specific wavelet from [PyWavelets](https://pywavelets.readthedocs.io/en/latest/ref/wavelets.html) to use, currently limited to `"bior2.2", "bior2.4", "bior2.6", "bior2.8", "bior4.4", "bior6.8"`
+
+## Future work
+We hope to add more filters and transforms to this list, including PQMF and STFT transforms.

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/scripts/ds_zero_to_pl_ckpt.py

@@ -0,0 +1,14 @@
+import argparse
+# from lightning.pytorch.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--save_path", type=str, help="Path to the zero checkpoint")
+    parser.add_argument("--output_path", type=str, help="Path to the output checkpoint", default="lightning_model.pt")
+    args = parser.parse_args()
+
+    # lightning deepspeed has saved a directory instead of a file
+    save_path = args.save_path
+    output_path = args.output_path
+    convert_zero_checkpoint_to_fp32_state_dict(save_path, output_path)

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/data/dataset.py

@@ -0,0 +1,654 @@
+import importlib
+import numpy as np
+import io
+import os
+import posixpath
+import random
+import re
+import subprocess
+import time
+import torch
+import torchaudio
+import webdataset as wds
+
+from aeiou.core import is_silence
+from os import path
+from pedalboard.io import AudioFile
+from torchaudio import transforms as T
+from typing import Optional, Callable, List
+
+from .utils import Stereo, Mono, PhaseFlipper, PadCrop_Normalized_T
+
+AUDIO_KEYS = ("flac", "wav", "mp3", "m4a", "ogg", "opus")
+
+# fast_scandir implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py
+
+def fast_scandir(
+    dir:str,  # top-level directory at which to begin scanning
+    ext:list,  # list of allowed file extensions,
+    #max_size = 1 * 1000 * 1000 * 1000 # Only files < 1 GB
+    ):
+    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
+    subfolders, files = [], []
+    ext = ['.'+x if x[0]!='.' else x for x in ext]  # add starting period to extensions if needed
+    try: # hope to avoid 'permission denied' by this try
+        for f in os.scandir(dir):
+            try: # 'hope to avoid too many levels of symbolic links' error
+                if f.is_dir():
+                    subfolders.append(f.path)
+                elif f.is_file():
+                    file_ext = os.path.splitext(f.name)[1].lower()
+                    is_hidden = os.path.basename(f.path).startswith(".")
+
+                    if file_ext in ext and not is_hidden:
+                        files.append(f.path)
+            except:
+                pass 
+    except:
+        pass
+
+    for dir in list(subfolders):
+        sf, f = fast_scandir(dir, ext)
+        subfolders.extend(sf)
+        files.extend(f)
+    return subfolders, files
+
+def keyword_scandir(
+    dir: str,  # top-level directory at which to begin scanning
+    ext: list,  # list of allowed file extensions
+    keywords: list,  # list of keywords to search for in the file name
+):
+    "very fast `glob` alternative. from https://stackoverflow.com/a/59803793/4259243"
+    subfolders, files = [], []
+    # make keywords case insensitive
+    keywords = [keyword.lower() for keyword in keywords]
+    # add starting period to extensions if needed
+    ext = ['.'+x if x[0] != '.' else x for x in ext]
+    banned_words = ["paxheader", "__macosx"]
+    try:  # hope to avoid 'permission denied' by this try
+        for f in os.scandir(dir):
+            try:  # 'hope to avoid too many levels of symbolic links' error
+                if f.is_dir():
+                    subfolders.append(f.path)
+                elif f.is_file():
+                    is_hidden = f.name.split("/")[-1][0] == '.'
+                    has_ext = os.path.splitext(f.name)[1].lower() in ext
+                    name_lower = f.name.lower()
+                    has_keyword = any(
+                        [keyword in name_lower for keyword in keywords])
+                    has_banned = any(
+                        [banned_word in name_lower for banned_word in banned_words])
+                    if has_ext and has_keyword and not has_banned and not is_hidden and not os.path.basename(f.path).startswith("._"):
+                        files.append(f.path)
+            except:
+                pass
+    except:
+        pass
+
+    for dir in list(subfolders):
+        sf, f = keyword_scandir(dir, ext, keywords)
+        subfolders.extend(sf)
+        files.extend(f)
+    return subfolders, files
+
+def get_audio_filenames(
+    paths: list,  # directories in which to search
+    keywords=None,
+    exts=['.wav', '.mp3', '.flac', '.ogg', '.aif', '.opus']
+):
+    "recursively get a list of audio filenames"
+    filenames = []
+    if type(paths) is str:
+        paths = [paths]
+    for path in paths:               # get a list of relevant filenames
+        if keywords is not None:
+            subfolders, files = keyword_scandir(path, exts, keywords)
+        else:
+            subfolders, files = fast_scandir(path, exts)
+        filenames.extend(files)
+    return filenames
+
+class LocalDatasetConfig:
+    def __init__(
+        self,
+        id: str,
+        path: str,
+        custom_metadata_fn: Optional[Callable[[str], str]] = None
+    ):
+        self.id = id
+        self.path = path
+        self.custom_metadata_fn = custom_metadata_fn
+
+class SampleDataset(torch.utils.data.Dataset):
+    def __init__(
+        self, 
+        configs,
+        sample_size=65536, 
+        sample_rate=48000, 
+        keywords=None, 
+        random_crop=True,
+        force_channels="stereo"
+    ):
+        super().__init__()
+        self.filenames = []
+
+        self.augs = torch.nn.Sequential(
+            PhaseFlipper(),
+        )
+
+        self.root_paths = []
+
+        self.pad_crop = PadCrop_Normalized_T(sample_size, sample_rate, randomize=random_crop)
+
+        self.force_channels = force_channels
+
+        self.encoding = torch.nn.Sequential(
+            Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
+            Mono() if self.force_channels == "mono" else torch.nn.Identity(),
+        )
+
+        self.sr = sample_rate
+
+        self.custom_metadata_fns = {}
+
+        for config in configs:
+            self.root_paths.append(config.path)
+            self.filenames.extend(get_audio_filenames(config.path, keywords))
+            if config.custom_metadata_fn is not None:
+                self.custom_metadata_fns[config.path] = config.custom_metadata_fn
+
+        print(f'Found {len(self.filenames)} files')
+
+    def load_file(self, filename):
+        ext = filename.split(".")[-1]
+
+        if ext == "mp3":
+            with AudioFile(filename) as f:
+                audio = f.read(f.frames)
+                audio = torch.from_numpy(audio)
+                in_sr = f.samplerate
+        else:
+            audio, in_sr = torchaudio.load(filename, format=ext)
+
+        if in_sr != self.sr:
+            resample_tf = T.Resample(in_sr, self.sr)
+            audio = resample_tf(audio)
+
+        return audio
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, idx):
+        audio_filename = self.filenames[idx]
+        try:
+            start_time = time.time()
+            audio = self.load_file(audio_filename)
+
+            audio, t_start, t_end, seconds_start, seconds_total, padding_mask = self.pad_crop(audio)
+
+            # Run augmentations on this sample (including random crop)
+            if self.augs is not None:
+                audio = self.augs(audio)
+
+            audio = audio.clamp(-1, 1)
+
+            # Encode the file to assist in prediction
+            if self.encoding is not None:
+                audio = self.encoding(audio)
+
+            info = {}
+
+            info["path"] = audio_filename
+
+            for root_path in self.root_paths:
+                if root_path in audio_filename:
+                    info["relpath"] = path.relpath(audio_filename, root_path)
+
+            info["timestamps"] = (t_start, t_end)
+            info["seconds_start"] = seconds_start
+            info["seconds_total"] = seconds_total
+            info["padding_mask"] = padding_mask
+
+            end_time = time.time()
+
+            info["load_time"] = end_time - start_time
+
+            for custom_md_path in self.custom_metadata_fns.keys():
+                if custom_md_path in audio_filename:
+                    custom_metadata_fn = self.custom_metadata_fns[custom_md_path]
+                    custom_metadata = custom_metadata_fn(info, audio)
+                    info.update(custom_metadata)
+
+                if "__reject__" in info and info["__reject__"]:
+                    return self[random.randrange(len(self))]
+
+            return (audio, info)
+        except Exception as e:
+            print(f'Couldn\'t load file {audio_filename}: {e}')
+            return self[random.randrange(len(self))]
+
+def group_by_keys(data, keys=wds.tariterators.base_plus_ext, lcase=True, suffixes=None, handler=None):
+    """Return function over iterator that groups key, value pairs into samples.
+    :param keys: function that splits the key into key and extension (base_plus_ext)
+    :param lcase: convert suffixes to lower case (Default value = True)
+    """
+    current_sample = None
+    for filesample in data:
+        assert isinstance(filesample, dict)
+        fname, value = filesample["fname"], filesample["data"]
+        prefix, suffix = keys(fname)
+        if wds.tariterators.trace:
+            print(
+                prefix,
+                suffix,
+                current_sample.keys() if isinstance(current_sample, dict) else None,
+            )
+        if prefix is None:
+            continue
+        if lcase:
+            suffix = suffix.lower()
+        if current_sample is None or prefix != current_sample["__key__"]:
+            if wds.tariterators.valid_sample(current_sample):
+                yield current_sample
+            current_sample = dict(__key__=prefix, __url__=filesample["__url__"])
+        if suffix in current_sample:
+            print(f"{fname}: duplicate file name in tar file {suffix} {current_sample.keys()}")
+        if suffixes is None or suffix in suffixes:
+            current_sample[suffix] = value
+    if wds.tariterators.valid_sample(current_sample):
+        yield current_sample
+
+wds.tariterators.group_by_keys = group_by_keys
+
+# S3 code and WDS preprocessing code based on implementation by Scott Hawley originally in https://github.com/zqevans/audio-diffusion/blob/main/dataset/dataset.py
+
+def get_s3_contents(dataset_path, s3_url_prefix=None, filter='', recursive=True, debug=False, profile=None):
+    """
+    Returns a list of full S3 paths to files in a given S3 bucket and directory path.
+    """
+    # Ensure dataset_path ends with a trailing slash
+    if dataset_path != '' and not dataset_path.endswith('/'):
+        dataset_path += '/'
+    # Use posixpath to construct the S3 URL path
+    bucket_path = posixpath.join(s3_url_prefix or '', dataset_path)
+    # Construct the `aws s3 ls` command
+    cmd = ['aws', 's3', 'ls', bucket_path]
+
+    if profile is not None:
+        cmd.extend(['--profile', profile])
+
+    if recursive:
+        # Add the --recursive flag if requested
+        cmd.append('--recursive')
+    
+    # Run the `aws s3 ls` command and capture the output
+    run_ls = subprocess.run(cmd, capture_output=True, check=True)
+    # Split the output into lines and strip whitespace from each line
+    contents = run_ls.stdout.decode('utf-8').split('\n')
+    contents = [x.strip() for x in contents if x]
+    # Remove the timestamp from lines that begin with a timestamp
+    contents = [re.sub(r'^\S+\s+\S+\s+\d+\s+', '', x)
+                if re.match(r'^\S+\s+\S+\s+\d+\s+', x) else x for x in contents]
+    # Construct a full S3 path for each file in the contents list
+    contents = [posixpath.join(s3_url_prefix or '', x)
+                for x in contents if not x.endswith('/')]
+    # Apply the filter, if specified
+    if filter:
+        contents = [x for x in contents if filter in x]
+    # Remove redundant directory names in the S3 URL
+    if recursive:
+        # Get the main directory name from the S3 URL
+        main_dir = "/".join(bucket_path.split('/')[3:])
+        # Remove the redundant directory names from each file path
+        contents = [x.replace(f'{main_dir}', '').replace(
+            '//', '/') for x in contents]
+    # Print debugging information, if requested
+    if debug:
+        print("contents = \n", contents)
+    # Return the list of S3 paths to files
+    return contents
+
+
+def get_all_s3_urls(
+    names=[],           # list of all valid [LAION AudioDataset] dataset names
+    # list of subsets you want from those datasets, e.g. ['train','valid']
+    subsets=[''],
+    s3_url_prefix=None,  # prefix for those dataset names
+    recursive=True,     # recursively list all tar files in all subdirs
+    filter_str='tar',   # only grab files with this substring
+    # print debugging info -- note: info displayed likely to change at dev's whims
+    debug=False,
+    profiles={},        # dictionary of profiles for each item in names, e.g. {'dataset1': 'profile1', 'dataset2': 'profile2'}
+):
+    "get urls of shards (tar files) for multiple datasets in one s3 bucket"
+    urls = []
+    for name in names:
+        # If s3_url_prefix is not specified, assume the full S3 path is included in each element of the names list
+        if s3_url_prefix is None:
+            contents_str = name
+        else:
+            # Construct the S3 path using the s3_url_prefix and the current name value
+            contents_str = posixpath.join(s3_url_prefix, name)
+        if debug:
+            print(f"get_all_s3_urls: {contents_str}:")
+        for subset in subsets:
+            subset_str = posixpath.join(contents_str, subset)
+            if debug:
+                print(f"subset_str = {subset_str}")
+            # Get the list of tar files in the current subset directory
+            profile = profiles.get(name, None)
+            tar_list = get_s3_contents(
+                subset_str, s3_url_prefix=None, recursive=recursive, filter=filter_str, debug=debug, profile=profile)
+            for tar in tar_list:
+                # Escape spaces and parentheses in the tar filename for use in the shell command
+                tar = tar.replace(" ", "\ ").replace(
+                    "(", "\(").replace(")", "\)")
+                # Construct the S3 path to the current tar file
+                s3_path = posixpath.join(name, subset, tar) + " -"
+                # Construct the AWS CLI command to download the current tar file
+                if s3_url_prefix is None:
+                    request_str = f"pipe:aws s3 --cli-connect-timeout 0 cp {s3_path}"
+                else:
+                    request_str = f"pipe:aws s3 --cli-connect-timeout 0 cp {posixpath.join(s3_url_prefix, s3_path)}"
+                if profiles.get(name):
+                    request_str += f" --profile {profiles.get(name)}"
+                if debug:
+                    print("request_str = ", request_str)
+                # Add the constructed URL to the list of URLs
+                urls.append(request_str)
+    return urls
+
+
+def log_and_continue(exn):
+    """Call in an exception handler to ignore any exception, isssue a warning, and continue."""
+    print(f"Handling webdataset error ({repr(exn)}). Ignoring.")
+    return True
+
+
+def is_valid_sample(sample):
+    has_json = "json" in sample
+    has_audio = "audio" in sample
+    is_silent = is_silence(sample["audio"])
+    is_rejected = "__reject__" in sample["json"] and sample["json"]["__reject__"]
+
+    return has_json and has_audio and not is_silent and not is_rejected
+
+class S3DatasetConfig:
+    def __init__(
+        self,
+        id: str,
+        s3_path: str,
+        custom_metadata_fn: Optional[Callable[[str], str]] = None,
+        profile: Optional[str] = None,
+    ):
+        self.id = id
+        self.path = s3_path
+        self.custom_metadata_fn = custom_metadata_fn
+        self.profile = profile
+        self.urls = []
+
+    def load_data_urls(self):
+        self.urls = get_all_s3_urls(
+            names=[self.path],
+            s3_url_prefix=None,
+            recursive=True,
+            profiles={self.path: self.profile} if self.profile else {},
+        )
+
+        return self.urls
+
+class LocalWebDatasetConfig:
+    def __init__(
+        self,
+        id: str,
+        path: str,
+        custom_metadata_fn: Optional[Callable[[str], str]] = None,
+        profile: Optional[str] = None,
+    ):
+        self.id = id
+        self.path = path
+        self.custom_metadata_fn = custom_metadata_fn
+        self.urls = []
+
+    def load_data_urls(self):
+
+        self.urls = fast_scandir(self.path, ["tar"])[1]
+
+        return self.urls
+
+def audio_decoder(key, value):
+    # Get file extension from key
+    ext = key.split(".")[-1]
+
+    if ext in AUDIO_KEYS:
+        return torchaudio.load(io.BytesIO(value))
+    else:
+        return None
+
+def collation_fn(samples):
+        batched = list(zip(*samples))
+        result = []
+        for b in batched:
+            if isinstance(b[0], (int, float)):
+                b = np.array(b)
+            elif isinstance(b[0], torch.Tensor):
+                b = torch.stack(b)
+            elif isinstance(b[0], np.ndarray):
+                b = np.array(b)
+            else:
+                b = b
+            result.append(b)
+        return result
+
+class WebDatasetDataLoader():
+    def __init__(
+        self,
+        datasets: List[S3DatasetConfig],
+        batch_size,
+        sample_size,
+        sample_rate=48000,
+        num_workers=8,
+        epoch_steps=1000,
+        random_crop=True,
+        force_channels="stereo",
+        augment_phase=True,
+        **data_loader_kwargs
+    ):
+
+        self.datasets = datasets
+
+        self.sample_size = sample_size
+        self.sample_rate = sample_rate
+        self.random_crop = random_crop
+        self.force_channels = force_channels
+        self.augment_phase = augment_phase
+
+        urls = [dataset.load_data_urls() for dataset in datasets]
+
+        # Flatten the list of lists of URLs
+        urls = [url for dataset_urls in urls for url in dataset_urls]
+
+        # Shuffle the urls
+        random.shuffle(urls)
+
+        self.dataset = wds.DataPipeline(
+            wds.ResampledShards(urls),
+            wds.tarfile_to_samples(handler=log_and_continue),
+            wds.decode(audio_decoder, handler=log_and_continue),
+            wds.map(self.wds_preprocess, handler=log_and_continue),
+            wds.select(is_valid_sample),
+            wds.to_tuple("audio", "json", handler=log_and_continue),
+            #wds.shuffle(bufsize=1000, initial=5000),
+            wds.batched(batch_size, partial=False, collation_fn=collation_fn),
+        ).with_epoch(epoch_steps//num_workers if num_workers > 0 else epoch_steps)
+
+        self.data_loader = wds.WebLoader(self.dataset, num_workers=num_workers, **data_loader_kwargs)
+
+    def wds_preprocess(self, sample):
+
+        found_key, rewrite_key = '', ''
+        for k, v in sample.items():  # print the all entries in dict
+            for akey in AUDIO_KEYS:
+                if k.endswith(akey):
+                    # to rename long/weird key with its simpler counterpart
+                    found_key, rewrite_key = k, akey
+                    break
+            if '' != found_key:
+                break
+        if '' == found_key:  # got no audio!
+            return None  # try returning None to tell WebDataset to skip this one
+
+        audio, in_sr = sample[found_key]
+        if in_sr != self.sample_rate:
+            resample_tf = T.Resample(in_sr, self.sample_rate)
+            audio = resample_tf(audio)
+
+        if self.sample_size is not None:
+            # Pad/crop and get the relative timestamp
+            pad_crop = PadCrop_Normalized_T(
+                self.sample_size, randomize=self.random_crop, sample_rate=self.sample_rate)
+            audio, t_start, t_end, seconds_start, seconds_total, padding_mask = pad_crop(
+                audio)
+            sample["json"]["seconds_start"] = seconds_start
+            sample["json"]["seconds_total"] = seconds_total
+            sample["json"]["padding_mask"] = padding_mask
+        else:
+            t_start, t_end = 0, 1
+
+        # Check if audio is length zero, initialize to a single zero if so
+        if audio.shape[-1] == 0:
+            audio = torch.zeros(1, 1)
+
+        # Make the audio stereo and augment by randomly inverting phase
+        augs = torch.nn.Sequential(
+            Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
+            Mono() if self.force_channels == "mono" else torch.nn.Identity(),
+            PhaseFlipper() if self.augment_phase else torch.nn.Identity()
+        )
+
+        audio = augs(audio)
+
+        sample["json"]["timestamps"] = (t_start, t_end)
+
+        if "text" in sample["json"]:
+            sample["json"]["prompt"] = sample["json"]["text"]
+
+        # Check for custom metadata functions
+        for dataset in self.datasets:
+            if dataset.custom_metadata_fn is None:
+                continue
+        
+            if dataset.path in sample["__url__"]:
+                custom_metadata = dataset.custom_metadata_fn(sample["json"], audio)
+                sample["json"].update(custom_metadata)
+
+        if found_key != rewrite_key:   # rename long/weird key with its simpler counterpart
+            del sample[found_key]
+
+        sample["audio"] = audio
+
+        # Add audio to the metadata as well for conditioning
+        sample["json"]["audio"] = audio
+        
+        return sample
+
+def create_dataloader_from_config(dataset_config, batch_size, sample_size, sample_rate, audio_channels=2, num_workers=4):
+
+    dataset_type = dataset_config.get("dataset_type", None)
+
+    assert dataset_type is not None, "Dataset type must be specified in dataset config"
+
+    if audio_channels == 1:
+        force_channels = "mono"
+    else:
+        force_channels = "stereo"
+
+    if dataset_type == "audio_dir":
+
+        audio_dir_configs = dataset_config.get("datasets", None)
+
+        assert audio_dir_configs is not None, "Directory configuration must be specified in datasets[\"dataset\"]"
+
+        configs = []
+
+        for audio_dir_config in audio_dir_configs:
+            audio_dir_path = audio_dir_config.get("path", None)
+            assert audio_dir_path is not None, "Path must be set for local audio directory configuration"
+
+            custom_metadata_fn = None
+            custom_metadata_module_path = audio_dir_config.get("custom_metadata_module", None)
+
+            if custom_metadata_module_path is not None:
+                spec = importlib.util.spec_from_file_location("metadata_module", custom_metadata_module_path)
+                metadata_module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(metadata_module)                
+
+                custom_metadata_fn = metadata_module.get_custom_metadata
+
+            configs.append(
+                LocalDatasetConfig(
+                    id=audio_dir_config["id"],
+                    path=audio_dir_path,
+                    custom_metadata_fn=custom_metadata_fn
+                )
+            )
+
+        train_set = SampleDataset(
+            configs,
+            sample_rate=sample_rate,
+            sample_size=sample_size,
+            random_crop=dataset_config.get("random_crop", True),
+            force_channels=force_channels
+        )
+
+        return torch.utils.data.DataLoader(train_set, batch_size, shuffle=True,
+                                num_workers=num_workers, persistent_workers=True, pin_memory=True, drop_last=True, collate_fn=collation_fn)
+
+    elif dataset_type in ["s3", "wds"]: # Support "s3" type for backwards compatibility
+        wds_configs = []
+
+        for wds_config in dataset_config["datasets"]:
+
+            custom_metadata_fn = None
+            custom_metadata_module_path = wds_config.get("custom_metadata_module", None)
+
+            if custom_metadata_module_path is not None:
+                spec = importlib.util.spec_from_file_location("metadata_module", custom_metadata_module_path)
+                metadata_module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(metadata_module)                
+
+                custom_metadata_fn = metadata_module.get_custom_metadata
+
+            if "s3_path" in wds_config:
+
+                wds_configs.append(
+                    S3DatasetConfig(
+                        id=wds_config["id"],
+                        s3_path=wds_config["s3_path"],
+                        custom_metadata_fn=custom_metadata_fn,
+                        profile=wds_config.get("profile", None),
+                    )
+                )
+            
+            elif "path" in wds_config:
+                    
+                    wds_configs.append(
+                        LocalWebDatasetConfig(
+                            id=wds_config["id"],
+                            path=wds_config["path"],
+                            custom_metadata_fn=custom_metadata_fn
+                        )
+                    )
+
+        return WebDatasetDataLoader(
+            wds_configs,
+            sample_rate=sample_rate,
+            sample_size=sample_size,
+            batch_size=batch_size,
+            random_crop=dataset_config.get("random_crop", True),
+            num_workers=num_workers,
+            persistent_workers=True,
+            force_channels=force_channels,
+            epoch_steps=dataset_config.get("epoch_steps", 2000)
+        ).data_loader

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/data/utils.py

@@ -0,0 +1,96 @@
+import math
+import random
+import torch
+
+from torch import nn
+from typing import Tuple
+
+class PadCrop(nn.Module):
+    def __init__(self, n_samples, randomize=True):
+        super().__init__()
+        self.n_samples = n_samples
+        self.randomize = randomize
+
+    def __call__(self, signal):
+        n, s = signal.shape
+        start = 0 if (not self.randomize) else torch.randint(0, max(0, s - self.n_samples) + 1, []).item()
+        end = start + self.n_samples
+        output = signal.new_zeros([n, self.n_samples])
+        output[:, :min(s, self.n_samples)] = signal[:, start:end]
+        return output
+
+class PadCrop_Normalized_T(nn.Module):
+    
+    def __init__(self, n_samples: int, sample_rate: int, randomize: bool = True):
+        
+        super().__init__()
+        
+        self.n_samples = n_samples
+        self.sample_rate = sample_rate
+        self.randomize = randomize
+
+    def __call__(self, source: torch.Tensor) -> Tuple[torch.Tensor, float, float, int, int]:
+        
+        n_channels, n_samples = source.shape
+        
+        # If the audio is shorter than the desired length, pad it
+        upper_bound = max(0, n_samples - self.n_samples)
+        
+        # If randomize is False, always start at the beginning of the audio
+        offset = 0
+        if(self.randomize and n_samples > self.n_samples):
+            offset = random.randint(0, upper_bound)
+
+        # Calculate the start and end times of the chunk
+        t_start = offset / (upper_bound + self.n_samples)
+        t_end = (offset + self.n_samples) / (upper_bound + self.n_samples)
+
+        # Create the chunk
+        chunk = source.new_zeros([n_channels, self.n_samples])
+
+        # Copy the audio into the chunk
+        chunk[:, :min(n_samples, self.n_samples)] = source[:, offset:offset + self.n_samples]
+        
+        # Calculate the start and end times of the chunk in seconds
+        seconds_start = math.floor(offset / self.sample_rate)
+        seconds_total = math.ceil(n_samples / self.sample_rate)
+
+        # Create a mask the same length as the chunk with 1s where the audio is and 0s where it isn't
+        padding_mask = torch.zeros([self.n_samples])
+        padding_mask[:min(n_samples, self.n_samples)] = 1
+        
+        
+        return (
+            chunk,
+            t_start,
+            t_end,
+            seconds_start,
+            seconds_total,
+            padding_mask
+        )
+
+class PhaseFlipper(nn.Module):
+    "Randomly invert the phase of a signal"
+    def __init__(self, p=0.5):
+        super().__init__()
+        self.p = p
+    def __call__(self, signal):
+        return -signal if (random.random() < self.p) else signal
+        
+class Mono(nn.Module):
+  def __call__(self, signal):
+    return torch.mean(signal, dim=0, keepdims=True) if len(signal.shape) > 1 else signal
+
+class Stereo(nn.Module):
+  def __call__(self, signal):
+    signal_shape = signal.shape
+    # Check if it's mono
+    if len(signal_shape) == 1: # s -> 2, s
+        signal = signal.unsqueeze(0).repeat(2, 1)
+    elif len(signal_shape) == 2:
+        if signal_shape[0] == 1: #1, s -> 2, s
+            signal = signal.repeat(2, 1)
+        elif signal_shape[0] > 2: #?, s -> 2,s
+            signal = signal[:2, :]    
+
+    return signal

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/inference/generation.py

@@ -0,0 +1,274 @@
+import numpy as np
+import torch 
+import typing as tp
+import math 
+from torchaudio import transforms as T
+
+from .utils import prepare_audio
+from .sampling import sample, sample_k, sample_rf
+from ..data.utils import PadCrop
+
+def generate_diffusion_uncond(
+        model,
+        steps: int = 250,
+        batch_size: int = 1,
+        sample_size: int = 2097152,
+        seed: int = -1,
+        device: str = "cuda",
+        init_audio: tp.Optional[tp.Tuple[int, torch.Tensor]] = None,
+        init_noise_level: float = 1.0,
+        return_latents = False,
+        **sampler_kwargs
+        ) -> torch.Tensor:
+    
+    # The length of the output in audio samples 
+    audio_sample_size = sample_size
+
+    # If this is latent diffusion, change sample_size instead to the downsampled latent size
+    if model.pretransform is not None:
+        sample_size = sample_size // model.pretransform.downsampling_ratio
+        
+    # Seed
+    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
+    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1, dtype=np.uint32)
+    print(seed)
+    torch.manual_seed(seed)
+    # Define the initial noise immediately after setting the seed
+    noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
+
+    if init_audio is not None:
+        # The user supplied some initial audio (for inpainting or variation). Let us prepare the input audio.
+        in_sr, init_audio = init_audio
+
+        io_channels = model.io_channels
+
+        # For latent models, set the io_channels to the autoencoder's io_channels
+        if model.pretransform is not None:
+            io_channels = model.pretransform.io_channels
+
+        # Prepare the initial audio for use by the model
+        init_audio = prepare_audio(init_audio, in_sr=in_sr, target_sr=model.sample_rate, target_length=audio_sample_size, target_channels=io_channels, device=device)
+
+        # For latent models, encode the initial audio into latents
+        if model.pretransform is not None:
+            init_audio = model.pretransform.encode(init_audio)
+
+        init_audio = init_audio.repeat(batch_size, 1, 1)
+    else:
+        # The user did not supply any initial audio for inpainting or variation. Generate new output from scratch. 
+        init_audio = None
+        init_noise_level = None
+
+    # Inpainting mask
+    
+    if init_audio is not None:
+        # variations
+        sampler_kwargs["sigma_max"] = init_noise_level
+        mask = None 
+    else:
+        mask = None
+
+    # Now the generative AI part:
+
+    diff_objective = model.diffusion_objective
+
+    if diff_objective == "v":    
+        # k-diffusion denoising process go!
+        sampled = sample_k(model.model, noise, init_audio, mask, steps, **sampler_kwargs, device=device)
+    elif diff_objective == "rectified_flow":
+        sampled = sample_rf(model.model, noise, init_data=init_audio, steps=steps, **sampler_kwargs, device=device)
+
+    # Denoising process done. 
+    # If this is latent diffusion, decode latents back into audio
+    if model.pretransform is not None and not return_latents:
+        sampled = model.pretransform.decode(sampled)
+
+    # Return audio
+    return sampled
+
+
+def generate_diffusion_cond(
+        model,
+        steps: int = 250,
+        cfg_scale=6,
+        conditioning: dict = None,
+        conditioning_tensors: tp.Optional[dict] = None,
+        negative_conditioning: dict = None,
+        negative_conditioning_tensors: tp.Optional[dict] = None,
+        batch_size: int = 1,
+        sample_size: int = 2097152,
+        sample_rate: int = 48000,
+        seed: int = -1,
+        device: str = "cuda",
+        init_audio: tp.Optional[tp.Tuple[int, torch.Tensor]] = None,
+        init_noise_level: float = 1.0,
+        mask_args: dict = None,
+        return_latents = False,
+        **sampler_kwargs
+        ) -> torch.Tensor: 
+    """
+    Generate audio from a prompt using a diffusion model.
+    
+    Args:
+        model: The diffusion model to use for generation.
+        steps: The number of diffusion steps to use.
+        cfg_scale: Classifier-free guidance scale 
+        conditioning: A dictionary of conditioning parameters to use for generation.
+        conditioning_tensors: A dictionary of precomputed conditioning tensors to use for generation.
+        batch_size: The batch size to use for generation.
+        sample_size: The length of the audio to generate, in samples.
+        sample_rate: The sample rate of the audio to generate (Deprecated, now pulled from the model directly)
+        seed: The random seed to use for generation, or -1 to use a random seed.
+        device: The device to use for generation.
+        init_audio: A tuple of (sample_rate, audio) to use as the initial audio for generation.
+        init_noise_level: The noise level to use when generating from an initial audio sample.
+        return_latents: Whether to return the latents used for generation instead of the decoded audio.
+        **sampler_kwargs: Additional keyword arguments to pass to the sampler.    
+    """
+
+    # The length of the output in audio samples 
+    audio_sample_size = sample_size
+
+    # If this is latent diffusion, change sample_size instead to the downsampled latent size
+    if model.pretransform is not None:
+        sample_size = sample_size // model.pretransform.downsampling_ratio
+        
+    # Seed
+    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
+    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1, dtype=np.uint32)
+    print(seed)
+    torch.manual_seed(seed)
+    # Define the initial noise immediately after setting the seed
+    noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
+
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+    torch.backends.cudnn.benchmark = False
+
+    # Conditioning
+    assert conditioning is not None or conditioning_tensors is not None, "Must provide either conditioning or conditioning_tensors"
+    if conditioning_tensors is None:
+        conditioning_tensors = model.conditioner(conditioning, device)
+    conditioning_inputs = model.get_conditioning_inputs(conditioning_tensors)
+
+    if negative_conditioning is not None or negative_conditioning_tensors is not None:
+        
+        if negative_conditioning_tensors is None:
+            negative_conditioning_tensors = model.conditioner(negative_conditioning, device)
+            
+        negative_conditioning_tensors = model.get_conditioning_inputs(negative_conditioning_tensors, negative=True)
+    else:
+        negative_conditioning_tensors = {}
+
+    if init_audio is not None:
+        # The user supplied some initial audio (for inpainting or variation). Let us prepare the input audio.
+        in_sr, init_audio = init_audio
+
+        io_channels = model.io_channels
+
+        # For latent models, set the io_channels to the autoencoder's io_channels
+        if model.pretransform is not None:
+            io_channels = model.pretransform.io_channels
+
+        # Prepare the initial audio for use by the model
+        init_audio = prepare_audio(init_audio, in_sr=in_sr, target_sr=model.sample_rate, target_length=audio_sample_size, target_channels=io_channels, device=device)
+
+        # For latent models, encode the initial audio into latents
+        if model.pretransform is not None:
+            init_audio = model.pretransform.encode(init_audio)
+
+        init_audio = init_audio.repeat(batch_size, 1, 1)
+    else:
+        # The user did not supply any initial audio for inpainting or variation. Generate new output from scratch. 
+        init_audio = None
+        init_noise_level = None
+        mask_args = None
+
+    # Inpainting mask
+    if init_audio is not None and mask_args is not None:
+        # Cut and paste init_audio according to cropfrom, pastefrom, pasteto
+        # This is helpful for forward and reverse outpainting
+        cropfrom = math.floor(mask_args["cropfrom"]/100.0 * sample_size)
+        pastefrom = math.floor(mask_args["pastefrom"]/100.0 * sample_size)
+        pasteto = math.ceil(mask_args["pasteto"]/100.0 * sample_size)
+        assert pastefrom < pasteto, "Paste From should be less than Paste To"
+        croplen = pasteto - pastefrom
+        if cropfrom + croplen > sample_size:
+            croplen = sample_size - cropfrom 
+        cropto = cropfrom + croplen
+        pasteto = pastefrom + croplen
+        cutpaste = init_audio.new_zeros(init_audio.shape)
+        cutpaste[:, :, pastefrom:pasteto] = init_audio[:,:,cropfrom:cropto]
+        #print(cropfrom, cropto, pastefrom, pasteto)
+        init_audio = cutpaste
+        # Build a soft mask (list of floats 0 to 1, the size of the latent) from the given args
+        mask = build_mask(sample_size, mask_args)
+        mask = mask.to(device)
+    elif init_audio is not None and mask_args is None:
+        # variations
+        sampler_kwargs["sigma_max"] = init_noise_level
+        mask = None 
+    else:
+        mask = None
+
+    model_dtype = next(model.model.parameters()).dtype
+    noise = noise.type(model_dtype)
+    conditioning_inputs = {k: v.type(model_dtype) if v is not None else v for k, v in conditioning_inputs.items()}
+    # Now the generative AI part:
+    # k-diffusion denoising process go!
+
+    diff_objective = model.diffusion_objective
+
+    if diff_objective == "v":    
+        # k-diffusion denoising process go!
+        sampled = sample_k(model.model, noise, init_audio, mask, steps, **sampler_kwargs, **conditioning_inputs, **negative_conditioning_tensors, cfg_scale=cfg_scale, batch_cfg=True, rescale_cfg=True, device=device)
+    elif diff_objective == "rectified_flow":
+
+        if "sigma_min" in sampler_kwargs:
+            del sampler_kwargs["sigma_min"]
+
+        if "sampler_type" in sampler_kwargs:
+            del sampler_kwargs["sampler_type"]
+
+        sampled = sample_rf(model.model, noise, init_data=init_audio, steps=steps, **sampler_kwargs, **conditioning_inputs, **negative_conditioning_tensors, cfg_scale=cfg_scale, batch_cfg=True, rescale_cfg=True, device=device)
+
+    # v-diffusion: 
+    #sampled = sample(model.model, noise, steps, 0, **conditioning_tensors, embedding_scale=cfg_scale)
+    del noise
+    del conditioning_tensors
+    del conditioning_inputs
+    torch.cuda.empty_cache()
+    # Denoising process done. 
+    # If this is latent diffusion, decode latents back into audio
+    if model.pretransform is not None and not return_latents:
+        #cast sampled latents to pretransform dtype
+        sampled = sampled.to(next(model.pretransform.parameters()).dtype)
+        sampled = model.pretransform.decode(sampled)
+
+    # Return audio
+    return sampled
+
+# builds a softmask given the parameters
+# returns array of values 0 to 1, size sample_size, where 0 means noise / fresh generation, 1 means keep the input audio, 
+# and anything between is a mixture of old/new
+# ideally 0.5 is half/half mixture but i haven't figured this out yet
+def build_mask(sample_size, mask_args):
+    maskstart = math.floor(mask_args["maskstart"]/100.0 * sample_size)
+    maskend = math.ceil(mask_args["maskend"]/100.0 * sample_size)
+    softnessL = round(mask_args["softnessL"]/100.0 * sample_size)
+    softnessR = round(mask_args["softnessR"]/100.0 * sample_size)
+    marination = mask_args["marination"]
+    # use hann windows for softening the transition (i don't know if this is correct)
+    hannL = torch.hann_window(softnessL*2, periodic=False)[:softnessL]
+    hannR = torch.hann_window(softnessR*2, periodic=False)[softnessR:]
+    # build the mask. 
+    mask = torch.zeros((sample_size))
+    mask[maskstart:maskend] = 1
+    mask[maskstart:maskstart+softnessL] = hannL
+    mask[maskend-softnessR:maskend] = hannR
+    # marination finishes the inpainting early in the denoising schedule, and lets audio get changed in the final rounds
+    if marination > 0:        
+        mask = mask * (1-marination) 
+    #print(mask)
+    return mask

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/inference/sampling.py

@@ -0,0 +1,232 @@
+import torch
+import math
+from tqdm import trange, tqdm
+
+import k_diffusion as K
+
+# Define the noise schedule and sampling loop
+def get_alphas_sigmas(t):
+    """Returns the scaling factors for the clean image (alpha) and for the
+    noise (sigma), given a timestep."""
+    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)
+
+def alpha_sigma_to_t(alpha, sigma):
+    """Returns a timestep, given the scaling factors for the clean image and for
+    the noise."""
+    return torch.atan2(sigma, alpha) / math.pi * 2
+
+def t_to_alpha_sigma(t):
+    """Returns the scaling factors for the clean image and for the noise, given
+    a timestep."""
+    return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)
+
+
+@torch.no_grad()
+def sample_discrete_euler(model, x, steps, sigma_max=1, **extra_args):
+    """Draws samples from a model given starting noise. Euler method"""
+
+    # Make tensor of ones to broadcast the single t values
+    ts = x.new_ones([x.shape[0]])
+
+    # Create the noise schedule
+    t = torch.linspace(sigma_max, 0, steps + 1)
+
+    #alphas, sigmas = 1-t, t
+
+    for t_curr, t_prev in tqdm(zip(t[:-1], t[1:])):
+            # Broadcast the current timestep to the correct shape
+            t_curr_tensor = t_curr * torch.ones(
+                (x.shape[0],), dtype=x.dtype, device=x.device
+            )
+            dt = t_prev - t_curr  # we solve backwards in our formulation
+            x = x + dt * model(x, t_curr_tensor, **extra_args) #.denoise(x, denoiser, t_curr_tensor, cond, uc)
+
+    # If we are on the last timestep, output the denoised image
+    return x
+
+@torch.no_grad()
+def sample(model, x, steps, eta, **extra_args):
+    """Draws samples from a model given starting noise. v-diffusion"""
+    ts = x.new_ones([x.shape[0]])
+
+    # Create the noise schedule
+    t = torch.linspace(1, 0, steps + 1)[:-1]
+
+    alphas, sigmas = get_alphas_sigmas(t)
+
+    # The sampling loop
+    for i in trange(steps):
+
+        # Get the model output (v, the predicted velocity)
+        with torch.cuda.amp.autocast():
+            v = model(x, ts * t[i], **extra_args).float()
+
+        # Predict the noise and the denoised image
+        pred = x * alphas[i] - v * sigmas[i]
+        eps = x * sigmas[i] + v * alphas[i]
+
+        # If we are not on the last timestep, compute the noisy image for the
+        # next timestep.
+        if i < steps - 1:
+            # If eta > 0, adjust the scaling factor for the predicted noise
+            # downward according to the amount of additional noise to add
+            ddim_sigma = eta * (sigmas[i + 1]**2 / sigmas[i]**2).sqrt() * \
+                (1 - alphas[i]**2 / alphas[i + 1]**2).sqrt()
+            adjusted_sigma = (sigmas[i + 1]**2 - ddim_sigma**2).sqrt()
+
+            # Recombine the predicted noise and predicted denoised image in the
+            # correct proportions for the next step
+            x = pred * alphas[i + 1] + eps * adjusted_sigma
+
+            # Add the correct amount of fresh noise
+            if eta:
+                x += torch.randn_like(x) * ddim_sigma
+
+    # If we are on the last timestep, output the denoised image
+    return pred
+
+# Soft mask inpainting is just shrinking hard (binary) mask inpainting
+# Given a float-valued soft mask (values between 0 and 1), get the binary mask for this particular step
+def get_bmask(i, steps, mask):
+    strength = (i+1)/(steps)
+    # convert to binary mask
+    bmask = torch.where(mask<=strength,1,0)
+    return bmask
+
+def make_cond_model_fn(model, cond_fn):
+    def cond_model_fn(x, sigma, **kwargs):
+        with torch.enable_grad():
+            x = x.detach().requires_grad_()
+            denoised = model(x, sigma, **kwargs)
+            cond_grad = cond_fn(x, sigma, denoised=denoised, **kwargs).detach()
+            cond_denoised = denoised.detach() + cond_grad * K.utils.append_dims(sigma**2, x.ndim)
+        return cond_denoised
+    return cond_model_fn
+
+# Uses k-diffusion from https://github.com/crowsonkb/k-diffusion
+# init_data is init_audio as latents (if this is latent diffusion)
+# For sampling, set both init_data and mask to None
+# For variations, set init_data 
+# For inpainting, set both init_data & mask 
+def sample_k(
+        model_fn, 
+        noise, 
+        init_data=None,
+        mask=None,
+        steps=100, 
+        sampler_type="dpmpp-2m-sde", 
+        sigma_min=0.5, 
+        sigma_max=50, 
+        rho=1.0, device="cuda", 
+        callback=None, 
+        cond_fn=None,
+        **extra_args
+    ):
+
+    denoiser = K.external.VDenoiser(model_fn)
+
+    if cond_fn is not None:
+        denoiser = make_cond_model_fn(denoiser, cond_fn)
+
+    # Make the list of sigmas. Sigma values are scalars related to the amount of noise each denoising step has
+    sigmas = K.sampling.get_sigmas_polyexponential(steps, sigma_min, sigma_max, rho, device=device)
+    # Scale the initial noise by sigma 
+    noise = noise * sigmas[0]
+
+    wrapped_callback = callback
+
+    if mask is None and init_data is not None:
+        # VARIATION (no inpainting)
+        # set the initial latent to the init_data, and noise it with initial sigma
+        x = init_data + noise 
+    elif mask is not None and init_data is not None:
+        # INPAINTING
+        bmask = get_bmask(0, steps, mask)
+        # initial noising
+        input_noised = init_data + noise
+        # set the initial latent to a mix of init_data and noise, based on step 0's binary mask
+        x = input_noised * bmask + noise * (1-bmask)
+        # define the inpainting callback function (Note: side effects, it mutates x)
+        # See https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L596C13-L596C105
+        # callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        # This is called immediately after `denoised = model(x, sigmas[i] * s_in, **extra_args)`
+        def inpainting_callback(args):
+            i = args["i"]
+            x = args["x"]
+            sigma = args["sigma"]
+            #denoised = args["denoised"]
+            # noise the init_data input with this step's appropriate amount of noise
+            input_noised = init_data + torch.randn_like(init_data) * sigma
+            # shrinking hard mask
+            bmask = get_bmask(i, steps, mask)
+            # mix input_noise with x, using binary mask
+            new_x = input_noised * bmask + x * (1-bmask)
+            # mutate x
+            x[:,:,:] = new_x[:,:,:]
+        # wrap together the inpainting callback and the user-submitted callback. 
+        if callback is None: 
+            wrapped_callback = inpainting_callback
+        else:
+            wrapped_callback = lambda args: (inpainting_callback(args), callback(args))
+    else:
+        # SAMPLING
+        # set the initial latent to noise
+        x = noise
+
+
+    with torch.cuda.amp.autocast():
+        if sampler_type == "k-heun":
+            return K.sampling.sample_heun(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-lms":
+            return K.sampling.sample_lms(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpmpp-2s-ancestral":
+            return K.sampling.sample_dpmpp_2s_ancestral(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpm-2":
+            return K.sampling.sample_dpm_2(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpm-fast":
+            return K.sampling.sample_dpm_fast(denoiser, x, sigma_min, sigma_max, steps, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "k-dpm-adaptive":
+            return K.sampling.sample_dpm_adaptive(denoiser, x, sigma_min, sigma_max, rtol=0.01, atol=0.01, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "dpmpp-2m-sde":
+            return K.sampling.sample_dpmpp_2m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+        elif sampler_type == "dpmpp-3m-sde":
+            return K.sampling.sample_dpmpp_3m_sde(denoiser, x, sigmas, disable=False, callback=wrapped_callback, extra_args=extra_args)
+
+# Uses discrete Euler sampling for rectified flow models
+# init_data is init_audio as latents (if this is latent diffusion)
+# For sampling, set both init_data and mask to None
+# For variations, set init_data 
+# For inpainting, set both init_data & mask 
+def sample_rf(
+        model_fn, 
+        noise, 
+        init_data=None,
+        steps=100, 
+        sigma_max=1,
+        device="cuda", 
+        callback=None, 
+        cond_fn=None,
+        **extra_args
+    ):
+
+    if sigma_max > 1:
+        sigma_max = 1
+
+    if cond_fn is not None:
+        denoiser = make_cond_model_fn(denoiser, cond_fn)
+
+    wrapped_callback = callback
+
+    if init_data is not None:
+        # VARIATION (no inpainting)
+        # Interpolate the init data and the noise for init audio
+        x = init_data * (1 - sigma_max) + noise * sigma_max
+    else:
+        # SAMPLING
+        # set the initial latent to noise
+        x = noise
+
+    with torch.cuda.amp.autocast():
+        # TODO: Add callback support
+        #return sample_discrete_euler(model_fn, x, steps, sigma_max, callback=wrapped_callback, **extra_args)
+        return sample_discrete_euler(model_fn, x, steps, sigma_max, **extra_args)

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/inference/utils.py

@@ -0,0 +1,35 @@
+from ..data.utils import PadCrop
+
+from torchaudio import transforms as T
+
+def set_audio_channels(audio, target_channels):
+    if target_channels == 1:
+        # Convert to mono
+        audio = audio.mean(1, keepdim=True)
+    elif target_channels == 2:
+        # Convert to stereo
+        if audio.shape[1] == 1:
+            audio = audio.repeat(1, 2, 1)
+        elif audio.shape[1] > 2:
+            audio = audio[:, :2, :]
+    return audio
+
+def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device):
+    
+    audio = audio.to(device)
+
+    if in_sr != target_sr:
+        resample_tf = T.Resample(in_sr, target_sr).to(device)
+        audio = resample_tf(audio)
+
+    audio = PadCrop(target_length, randomize=False)(audio)
+
+    # Add batch dimension
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0).unsqueeze(0)
+    elif audio.dim() == 2:
+        audio = audio.unsqueeze(0)
+
+    audio = set_audio_channels(audio, target_channels)
+
+    return audio

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/interface/gradio.py

@@ -0,0 +1,700 @@
+import gc
+import platform
+
+import numpy as np
+import gradio as gr
+import json 
+import torch
+import torchaudio
+
+from aeiou.viz import audio_spectrogram_image
+from einops import rearrange
+from safetensors.torch import load_file
+from torch.nn import functional as F
+from torchaudio import transforms as T
+
+from ..inference.generation import generate_diffusion_cond, generate_diffusion_uncond
+from ..models.factory import create_model_from_config
+from ..models.pretrained import get_pretrained_model
+from ..models.utils import load_ckpt_state_dict
+from ..inference.utils import prepare_audio
+from ..training.utils import copy_state_dict
+
+model = None
+sample_rate = 32000
+sample_size = 1920000
+
+def load_model(model_config=None, model_ckpt_path=None, pretrained_name=None, pretransform_ckpt_path=None, device="cuda", model_half=False):
+    global model, sample_rate, sample_size
+    
+    if pretrained_name is not None:
+        print(f"Loading pretrained model {pretrained_name}")
+        model, model_config = get_pretrained_model(pretrained_name)
+
+    elif model_config is not None and model_ckpt_path is not None:
+        print(f"Creating model from config")
+        model = create_model_from_config(model_config)
+
+        print(f"Loading model checkpoint from {model_ckpt_path}")
+        # Load checkpoint
+        copy_state_dict(model, load_ckpt_state_dict(model_ckpt_path))
+        #model.load_state_dict(load_ckpt_state_dict(model_ckpt_path))
+
+    sample_rate = model_config["sample_rate"]
+    sample_size = model_config["sample_size"]
+
+    if pretransform_ckpt_path is not None:
+        print(f"Loading pretransform checkpoint from {pretransform_ckpt_path}")
+        model.pretransform.load_state_dict(load_ckpt_state_dict(pretransform_ckpt_path), strict=False)
+        print(f"Done loading pretransform")
+
+    model.to(device).eval().requires_grad_(False)
+
+    if model_half:
+        model.to(torch.float16)
+        
+    print(f"Done loading model")
+
+    return model, model_config
+
+def generate_cond(
+        prompt,
+        negative_prompt=None,
+        seconds_start=0,
+        seconds_total=30,
+        cfg_scale=6.0,
+        steps=250,
+        preview_every=None,
+        seed=-1,
+        sampler_type="dpmpp-3m-sde",
+        sigma_min=0.03,
+        sigma_max=1000,
+        cfg_rescale=0.0,
+        use_init=False,
+        init_audio=None,
+        init_noise_level=1.0,
+        mask_cropfrom=None,
+        mask_pastefrom=None,
+        mask_pasteto=None,
+        mask_maskstart=None,
+        mask_maskend=None,
+        mask_softnessL=None,
+        mask_softnessR=None,
+        mask_marination=None,
+        batch_size=1    
+    ):
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+    print(f"Prompt: {prompt}")
+
+    global preview_images
+    preview_images = []
+    if preview_every == 0:
+        preview_every = None
+
+    # Return fake stereo audio
+    conditioning = [{"prompt": prompt, "seconds_start": seconds_start, "seconds_total": seconds_total}] * batch_size
+
+    if negative_prompt:
+        negative_conditioning = [{"prompt": negative_prompt, "seconds_start": seconds_start, "seconds_total": seconds_total}] * batch_size
+    else:
+        negative_conditioning = None
+        
+    #Get the device from the model
+    device = next(model.parameters()).device
+
+    seed = int(seed)
+
+    if not use_init:
+        init_audio = None
+    
+    input_sample_size = sample_size
+
+    if init_audio is not None:
+        in_sr, init_audio = init_audio
+        # Turn into torch tensor, converting from int16 to float32
+        init_audio = torch.from_numpy(init_audio).float().div(32767)
+        
+        if init_audio.dim() == 1:
+            init_audio = init_audio.unsqueeze(0) # [1, n]
+        elif init_audio.dim() == 2:
+            init_audio = init_audio.transpose(0, 1) # [n, 2] -> [2, n]
+
+        if in_sr != sample_rate:
+            resample_tf = T.Resample(in_sr, sample_rate).to(init_audio.device)
+            init_audio = resample_tf(init_audio)
+
+        audio_length = init_audio.shape[-1]
+
+        if audio_length > sample_size:
+
+            input_sample_size = audio_length + (model.min_input_length - (audio_length % model.min_input_length)) % model.min_input_length
+
+        init_audio = (sample_rate, init_audio)
+
+    def progress_callback(callback_info):
+        global preview_images
+        denoised = callback_info["denoised"]
+        current_step = callback_info["i"]
+        sigma = callback_info["sigma"]
+
+        if (current_step - 1) % preview_every == 0:
+            if model.pretransform is not None:
+                denoised = model.pretransform.decode(denoised)
+            denoised = rearrange(denoised, "b d n -> d (b n)")
+            denoised = denoised.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+            audio_spectrogram = audio_spectrogram_image(denoised, sample_rate=sample_rate)
+            preview_images.append((audio_spectrogram, f"Step {current_step} sigma={sigma:.3f})"))
+
+    # If inpainting, send mask args
+    # This will definitely change in the future
+    if mask_cropfrom is not None: 
+        mask_args = {
+            "cropfrom": mask_cropfrom,
+            "pastefrom": mask_pastefrom,
+            "pasteto": mask_pasteto,
+            "maskstart": mask_maskstart,
+            "maskend": mask_maskend,
+            "softnessL": mask_softnessL,
+            "softnessR": mask_softnessR,
+            "marination": mask_marination,
+        }
+    else:
+        mask_args = None 
+
+    # Do the audio generation
+    audio = generate_diffusion_cond(
+        model, 
+        conditioning=conditioning,
+        negative_conditioning=negative_conditioning,
+        steps=steps,
+        cfg_scale=cfg_scale,
+        batch_size=batch_size,
+        sample_size=input_sample_size,
+        sample_rate=sample_rate,
+        seed=seed,
+        device=device,
+        sampler_type=sampler_type,
+        sigma_min=sigma_min,
+        sigma_max=sigma_max,
+        init_audio=init_audio,
+        init_noise_level=init_noise_level,
+        mask_args = mask_args,
+        callback = progress_callback if preview_every is not None else None,
+        scale_phi = cfg_rescale
+    )
+
+    # Convert to WAV file
+    audio = rearrange(audio, "b d n -> d (b n)")
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+    torchaudio.save("output.wav", audio, sample_rate)
+
+    # Let's look at a nice spectrogram too
+    audio_spectrogram = audio_spectrogram_image(audio, sample_rate=sample_rate)
+
+    return ("output.wav", [audio_spectrogram, *preview_images])
+
+def generate_uncond(
+        steps=250,
+        seed=-1,
+        sampler_type="dpmpp-3m-sde",
+        sigma_min=0.03,
+        sigma_max=1000,
+        use_init=False,
+        init_audio=None,
+        init_noise_level=1.0,
+        batch_size=1,
+        preview_every=None
+        ):
+
+    global preview_images
+
+    preview_images = []
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+    #Get the device from the model
+    device = next(model.parameters()).device
+
+    seed = int(seed)
+
+    if not use_init:
+        init_audio = None
+    
+    input_sample_size = sample_size
+
+    if init_audio is not None:
+        in_sr, init_audio = init_audio
+        # Turn into torch tensor, converting from int16 to float32
+        init_audio = torch.from_numpy(init_audio).float().div(32767)
+        
+        if init_audio.dim() == 1:
+            init_audio = init_audio.unsqueeze(0) # [1, n]
+        elif init_audio.dim() == 2:
+            init_audio = init_audio.transpose(0, 1) # [n, 2] -> [2, n]
+
+        if in_sr != sample_rate:
+            resample_tf = T.Resample(in_sr, sample_rate).to(init_audio.device)
+            init_audio = resample_tf(init_audio)
+
+        audio_length = init_audio.shape[-1]
+
+        if audio_length > sample_size:
+
+            input_sample_size = audio_length + (model.min_input_length - (audio_length % model.min_input_length)) % model.min_input_length
+
+        init_audio = (sample_rate, init_audio)
+
+    def progress_callback(callback_info):
+        global preview_images
+        denoised = callback_info["denoised"]
+        current_step = callback_info["i"]
+        sigma = callback_info["sigma"]
+
+        if (current_step - 1) % preview_every == 0:
+
+            if model.pretransform is not None:
+                denoised = model.pretransform.decode(denoised)
+
+            denoised = rearrange(denoised, "b d n -> d (b n)")
+
+            denoised = denoised.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+            audio_spectrogram = audio_spectrogram_image(denoised, sample_rate=sample_rate)
+
+            preview_images.append((audio_spectrogram, f"Step {current_step} sigma={sigma:.3f})"))
+
+    audio = generate_diffusion_uncond(
+        model, 
+        steps=steps,
+        batch_size=batch_size,
+        sample_size=input_sample_size,
+        seed=seed,
+        device=device,
+        sampler_type=sampler_type,
+        sigma_min=sigma_min,
+        sigma_max=sigma_max,
+        init_audio=init_audio,
+        init_noise_level=init_noise_level,
+        callback = progress_callback if preview_every is not None else None
+    )
+
+    audio = rearrange(audio, "b d n -> d (b n)")
+
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+    torchaudio.save("output.wav", audio, sample_rate)
+
+    audio_spectrogram = audio_spectrogram_image(audio, sample_rate=sample_rate)
+
+    return ("output.wav", [audio_spectrogram, *preview_images])
+
+def generate_lm(
+        temperature=1.0,
+        top_p=0.95,
+        top_k=0,    
+        batch_size=1,
+        ):
+    
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+    #Get the device from the model
+    device = next(model.parameters()).device
+
+    audio = model.generate_audio(
+        batch_size=batch_size,
+        max_gen_len = sample_size//model.pretransform.downsampling_ratio,
+        conditioning=None,
+        temp=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        use_cache=True
+    )
+
+    audio = rearrange(audio, "b d n -> d (b n)")
+
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+    torchaudio.save("output.wav", audio, sample_rate)
+
+    audio_spectrogram = audio_spectrogram_image(audio, sample_rate=sample_rate)
+
+    return ("output.wav", [audio_spectrogram])
+
+
+def create_uncond_sampling_ui(model_config):   
+    generate_button = gr.Button("Generate", variant='primary', scale=1)
+    
+    with gr.Row(equal_height=False):
+        with gr.Column():            
+            with gr.Row():
+                # Steps slider
+                steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=100, label="Steps")
+
+            with gr.Accordion("Sampler params", open=False):
+            
+                # Seed
+                seed_textbox = gr.Textbox(label="Seed (set to -1 for random seed)", value="-1")
+
+            # Sampler params
+                with gr.Row():
+                    sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde")
+                    sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma min")
+                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=500, label="Sigma max")
+
+            with gr.Accordion("Init audio", open=False):
+                init_audio_checkbox = gr.Checkbox(label="Use init audio")
+                init_audio_input = gr.Audio(label="Init audio")
+                init_noise_level_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.01, value=0.1, label="Init noise level")
+
+        with gr.Column():
+            audio_output = gr.Audio(label="Output audio", interactive=False)
+            audio_spectrogram_output = gr.Gallery(label="Output spectrogram", show_label=False)
+            send_to_init_button = gr.Button("Send to init audio", scale=1)
+            send_to_init_button.click(fn=lambda audio: audio, inputs=[audio_output], outputs=[init_audio_input])
+    
+    generate_button.click(fn=generate_uncond, 
+        inputs=[
+            steps_slider, 
+            seed_textbox, 
+            sampler_type_dropdown, 
+            sigma_min_slider, 
+            sigma_max_slider,
+            init_audio_checkbox,
+            init_audio_input,
+            init_noise_level_slider,
+        ], 
+        outputs=[
+            audio_output, 
+            audio_spectrogram_output
+        ], 
+        api_name="generate")
+
+def create_sampling_ui(model_config, inpainting=False):
+    with gr.Row():
+        with gr.Column(scale=6):
+            prompt = gr.Textbox(show_label=False, placeholder="Prompt")
+            negative_prompt = gr.Textbox(show_label=False, placeholder="Negative prompt")
+        generate_button = gr.Button("Generate", variant='primary', scale=1)
+    
+    model_conditioning_config = model_config["model"].get("conditioning", None)
+
+    has_seconds_start = False
+    has_seconds_total = False
+
+    if model_conditioning_config is not None:
+        for conditioning_config in model_conditioning_config["configs"]:
+            if conditioning_config["id"] == "seconds_start":
+                has_seconds_start = True
+            if conditioning_config["id"] == "seconds_total":
+                has_seconds_total = True
+
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            with gr.Row(visible = has_seconds_start or has_seconds_total):
+                # Timing controls
+                seconds_start_slider = gr.Slider(minimum=0, maximum=512, step=1, value=0, label="Seconds start", visible=has_seconds_start)
+                seconds_total_slider = gr.Slider(minimum=0, maximum=512, step=1, value=sample_size//sample_rate, label="Seconds total", visible=has_seconds_total)
+            
+            with gr.Row():
+                # Steps slider
+                steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=100, label="Steps")
+
+                # Preview Every slider
+                preview_every_slider = gr.Slider(minimum=0, maximum=100, step=1, value=0, label="Preview Every")
+
+                # CFG scale 
+                cfg_scale_slider = gr.Slider(minimum=0.0, maximum=25.0, step=0.1, value=7.0, label="CFG scale")
+
+            with gr.Accordion("Sampler params", open=False):
+            
+                # Seed
+                seed_textbox = gr.Textbox(label="Seed (set to -1 for random seed)", value="-1")
+
+                # Sampler params
+                with gr.Row():
+                    sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde")
+                    sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma min")
+                    sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=500, label="Sigma max")
+                    cfg_rescale_slider = gr.Slider(minimum=0.0, maximum=1, step=0.01, value=0.0, label="CFG rescale amount")
+
+            if inpainting: 
+                # Inpainting Tab
+                with gr.Accordion("Inpainting", open=False):
+                    sigma_max_slider.maximum=1000
+                    
+                    init_audio_checkbox = gr.Checkbox(label="Do inpainting")
+                    init_audio_input = gr.Audio(label="Init audio")
+                    init_noise_level_slider = gr.Slider(minimum=0.1, maximum=100.0, step=0.1, value=80, label="Init audio noise level", visible=False) # hide this
+
+                    mask_cropfrom_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=0, label="Crop From %")
+                    mask_pastefrom_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=0, label="Paste From %")
+                    mask_pasteto_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=100, label="Paste To %")
+
+                    mask_maskstart_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=50, label="Mask Start %")
+                    mask_maskend_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=100, label="Mask End %")
+                    mask_softnessL_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=0, label="Softmask Left Crossfade Length %")
+                    mask_softnessR_slider = gr.Slider(minimum=0.0, maximum=100.0, step=0.1, value=0, label="Softmask Right Crossfade Length %")
+                    mask_marination_slider = gr.Slider(minimum=0.0, maximum=1, step=0.0001, value=0, label="Marination level", visible=False) # still working on the usefulness of this 
+
+                    inputs = [prompt, 
+                        negative_prompt,
+                        seconds_start_slider, 
+                        seconds_total_slider, 
+                        cfg_scale_slider, 
+                        steps_slider, 
+                        preview_every_slider, 
+                        seed_textbox, 
+                        sampler_type_dropdown, 
+                        sigma_min_slider, 
+                        sigma_max_slider,
+                        cfg_rescale_slider,
+                        init_audio_checkbox,
+                        init_audio_input,
+                        init_noise_level_slider,
+                        mask_cropfrom_slider,
+                        mask_pastefrom_slider,
+                        mask_pasteto_slider,
+                        mask_maskstart_slider,
+                        mask_maskend_slider,
+                        mask_softnessL_slider,
+                        mask_softnessR_slider,
+                        mask_marination_slider
+                    ]
+            else:
+                # Default generation tab
+                with gr.Accordion("Init audio", open=False):
+                    init_audio_checkbox = gr.Checkbox(label="Use init audio")
+                    init_audio_input = gr.Audio(label="Init audio")
+                    init_noise_level_slider = gr.Slider(minimum=0.1, maximum=100.0, step=0.01, value=0.1, label="Init noise level")
+
+                    inputs = [prompt, 
+                        negative_prompt,
+                        seconds_start_slider, 
+                        seconds_total_slider, 
+                        cfg_scale_slider, 
+                        steps_slider, 
+                        preview_every_slider, 
+                        seed_textbox, 
+                        sampler_type_dropdown, 
+                        sigma_min_slider, 
+                        sigma_max_slider,
+                        cfg_rescale_slider,
+                        init_audio_checkbox,
+                        init_audio_input,
+                        init_noise_level_slider
+                    ]
+
+        with gr.Column():
+            audio_output = gr.Audio(label="Output audio", interactive=False)
+            audio_spectrogram_output = gr.Gallery(label="Output spectrogram", show_label=False)
+            send_to_init_button = gr.Button("Send to init audio", scale=1)
+            send_to_init_button.click(fn=lambda audio: audio, inputs=[audio_output], outputs=[init_audio_input])
+    
+    generate_button.click(fn=generate_cond, 
+        inputs=inputs,
+        outputs=[
+            audio_output, 
+            audio_spectrogram_output
+        ], 
+        api_name="generate")
+
+
+def create_txt2audio_ui(model_config):
+    with gr.Blocks() as ui:
+        with gr.Tab("Generation"):
+            create_sampling_ui(model_config) 
+        with gr.Tab("Inpainting"):
+            create_sampling_ui(model_config, inpainting=True)    
+    return ui
+
+def create_diffusion_uncond_ui(model_config):
+    with gr.Blocks() as ui:
+        create_uncond_sampling_ui(model_config)
+    
+    return ui
+
+def autoencoder_process(audio, latent_noise, n_quantizers):
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+    #Get the device from the model
+    device = next(model.parameters()).device
+
+    in_sr, audio = audio
+
+    audio = torch.from_numpy(audio).float().div(32767).to(device)
+
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0)
+    else:
+        audio = audio.transpose(0, 1)
+
+    audio = model.preprocess_audio_for_encoder(audio, in_sr)
+    # Note: If you need to do chunked encoding, to reduce VRAM, 
+    # then add these arguments to encode_audio and decode_audio: chunked=True, overlap=32, chunk_size=128
+    # To turn it off, do chunked=False
+    # Optimal overlap and chunk_size values will depend on the model. 
+    # See encode_audio & decode_audio in autoencoders.py for more info
+    # Get dtype of model
+    dtype = next(model.parameters()).dtype
+
+    audio = audio.to(dtype)
+
+    if n_quantizers > 0:
+        latents = model.encode_audio(audio, chunked=False, n_quantizers=n_quantizers)
+    else:
+        latents = model.encode_audio(audio, chunked=False)
+
+    if latent_noise > 0:
+        latents = latents + torch.randn_like(latents) * latent_noise
+
+    audio = model.decode_audio(latents, chunked=False)
+
+    audio = rearrange(audio, "b d n -> d (b n)")
+
+    audio = audio.to(torch.float32).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+    torchaudio.save("output.wav", audio, sample_rate)
+
+    return "output.wav"
+
+def create_autoencoder_ui(model_config):
+
+    is_dac_rvq = "model" in model_config and "bottleneck" in model_config["model"] and model_config["model"]["bottleneck"]["type"] in ["dac_rvq","dac_rvq_vae"]
+
+    if is_dac_rvq:
+        n_quantizers = model_config["model"]["bottleneck"]["config"]["n_codebooks"]
+    else:
+        n_quantizers = 0
+
+    with gr.Blocks() as ui:
+        input_audio = gr.Audio(label="Input audio")
+        output_audio = gr.Audio(label="Output audio", interactive=False)
+        n_quantizers_slider = gr.Slider(minimum=1, maximum=n_quantizers, step=1, value=n_quantizers, label="# quantizers", visible=is_dac_rvq)
+        latent_noise_slider = gr.Slider(minimum=0.0, maximum=10.0, step=0.001, value=0.0, label="Add latent noise")
+        process_button = gr.Button("Process", variant='primary', scale=1)
+        process_button.click(fn=autoencoder_process, inputs=[input_audio, latent_noise_slider, n_quantizers_slider], outputs=output_audio, api_name="process")
+
+    return ui
+
+def diffusion_prior_process(audio, steps, sampler_type, sigma_min, sigma_max):
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+    #Get the device from the model
+    device = next(model.parameters()).device
+
+    in_sr, audio = audio
+
+    audio = torch.from_numpy(audio).float().div(32767).to(device)
+    
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0) # [1, n]
+    elif audio.dim() == 2:
+        audio = audio.transpose(0, 1) # [n, 2] -> [2, n]
+
+    audio = audio.unsqueeze(0)
+
+    audio = model.stereoize(audio, in_sr, steps, sampler_kwargs={"sampler_type": sampler_type, "sigma_min": sigma_min, "sigma_max": sigma_max})
+
+    audio = rearrange(audio, "b d n -> d (b n)")
+
+    audio = audio.to(torch.float32).div(torch.max(torch.abs(audio))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+
+    torchaudio.save("output.wav", audio, sample_rate)
+
+    return "output.wav"
+
+def create_diffusion_prior_ui(model_config):
+    with gr.Blocks() as ui:
+        input_audio = gr.Audio(label="Input audio")
+        output_audio = gr.Audio(label="Output audio", interactive=False)
+        # Sampler params
+        with gr.Row():
+            steps_slider = gr.Slider(minimum=1, maximum=500, step=1, value=100, label="Steps")
+            sampler_type_dropdown = gr.Dropdown(["dpmpp-2m-sde", "dpmpp-3m-sde", "k-heun", "k-lms", "k-dpmpp-2s-ancestral", "k-dpm-2", "k-dpm-fast"], label="Sampler type", value="dpmpp-3m-sde")
+            sigma_min_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.03, label="Sigma min")
+            sigma_max_slider = gr.Slider(minimum=0.0, maximum=1000.0, step=0.1, value=500, label="Sigma max")
+        process_button = gr.Button("Process", variant='primary', scale=1)
+        process_button.click(fn=diffusion_prior_process, inputs=[input_audio, steps_slider, sampler_type_dropdown, sigma_min_slider, sigma_max_slider], outputs=output_audio, api_name="process")    
+
+    return ui
+
+def create_lm_ui(model_config):
+    with gr.Blocks() as ui:
+        output_audio = gr.Audio(label="Output audio", interactive=False)
+        audio_spectrogram_output = gr.Gallery(label="Output spectrogram", show_label=False)
+
+        # Sampling params
+        with gr.Row():
+            temperature_slider = gr.Slider(minimum=0, maximum=5, step=0.01, value=1.0, label="Temperature")
+            top_p_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.95, label="Top p")
+            top_k_slider = gr.Slider(minimum=0, maximum=100, step=1, value=0, label="Top k")
+
+        generate_button = gr.Button("Generate", variant='primary', scale=1)
+        generate_button.click(
+            fn=generate_lm, 
+            inputs=[
+                temperature_slider, 
+                top_p_slider, 
+                top_k_slider
+            ], 
+            outputs=[output_audio, audio_spectrogram_output],
+            api_name="generate"
+        )
+
+    return ui
+
+def create_ui(model_config_path=None, ckpt_path=None, pretrained_name=None, pretransform_ckpt_path=None, model_half=False):
+
+    assert (pretrained_name is not None) ^ (model_config_path is not None and ckpt_path is not None), "Must specify either pretrained name or provide a model config and checkpoint, but not both"
+
+    if model_config_path is not None:
+        # Load config from json file
+        with open(model_config_path) as f:
+            model_config = json.load(f)
+    else:
+        model_config = None
+
+    try:
+        has_mps = platform.system() == "Darwin" and torch.backends.mps.is_available()
+    except Exception:
+        # In case this version of Torch doesn't even have `torch.backends.mps`...
+        has_mps = False
+
+    if has_mps:
+        device = torch.device("mps")
+    elif torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    print("Using device:", device)
+
+    _, model_config = load_model(model_config, ckpt_path, pretrained_name=pretrained_name, pretransform_ckpt_path=pretransform_ckpt_path, model_half=model_half, device=device)
+    
+    model_type = model_config["model_type"]
+
+    if model_type == "diffusion_cond":
+        ui = create_txt2audio_ui(model_config)
+    elif model_type == "diffusion_uncond":
+        ui = create_diffusion_uncond_ui(model_config)
+    elif model_type == "autoencoder" or model_type == "diffusion_autoencoder":
+        ui = create_autoencoder_ui(model_config)
+    elif model_type == "diffusion_prior":
+        ui = create_diffusion_prior_ui(model_config)
+    elif model_type == "lm":
+        ui = create_lm_ui(model_config)
+        
+    return ui

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/models/autoencoders.py

@@ -0,0 +1,794 @@
+import torch
+import math
+import numpy as np
+
+from torch import nn
+from torch.nn import functional as F
+from torchaudio import transforms as T
+from alias_free_torch import Activation1d
+from dac.nn.layers import WNConv1d, WNConvTranspose1d
+from typing import Literal, Dict, Any
+
+from ..inference.sampling import sample
+from ..inference.utils import prepare_audio
+from .blocks import SnakeBeta
+from .bottleneck import Bottleneck, DiscreteBottleneck
+from .diffusion import ConditionedDiffusionModel, DAU1DCondWrapper, UNet1DCondWrapper, DiTWrapper
+from .factory import create_pretransform_from_config, create_bottleneck_from_config
+from .pretransforms import Pretransform
+
+def checkpoint(function, *args, **kwargs):
+    kwargs.setdefault("use_reentrant", False)
+    return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
+
+def get_activation(activation: Literal["elu", "snake", "none"], antialias=False, channels=None) -> nn.Module:
+    if activation == "elu":
+        act = nn.ELU()
+    elif activation == "snake":
+        act = SnakeBeta(channels)
+    elif activation == "none":
+        act = nn.Identity()
+    else:
+        raise ValueError(f"Unknown activation {activation}")
+    
+    if antialias:
+        act = Activation1d(act)
+    
+    return act
+
+class ResidualUnit(nn.Module):
+    def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False):
+        super().__init__()
+        
+        self.dilation = dilation
+
+        padding = (dilation * (7-1)) // 2
+
+        self.layers = nn.Sequential(
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
+            WNConv1d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=7, dilation=dilation, padding=padding),
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
+            WNConv1d(in_channels=out_channels, out_channels=out_channels,
+                      kernel_size=1)
+        )
+
+    def forward(self, x):
+        res = x
+        
+        #x = checkpoint(self.layers, x)
+        x = self.layers(x)
+
+        return x + res
+
+class EncoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False):
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            ResidualUnit(in_channels=in_channels,
+                         out_channels=in_channels, dilation=1, use_snake=use_snake),
+            ResidualUnit(in_channels=in_channels,
+                         out_channels=in_channels, dilation=3, use_snake=use_snake),
+            ResidualUnit(in_channels=in_channels,
+                         out_channels=in_channels, dilation=9, use_snake=use_snake),
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
+            WNConv1d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2)),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False):
+        super().__init__()
+
+        if use_nearest_upsample:
+            upsample_layer = nn.Sequential(
+                nn.Upsample(scale_factor=stride, mode="nearest"),
+                WNConv1d(in_channels=in_channels,
+                        out_channels=out_channels, 
+                        kernel_size=2*stride,
+                        stride=1,
+                        bias=False,
+                        padding='same')
+            )
+        else:
+            upsample_layer = WNConvTranspose1d(in_channels=in_channels,
+                               out_channels=out_channels,
+                               kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2))
+
+        self.layers = nn.Sequential(
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
+            upsample_layer,
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
+                         dilation=1, use_snake=use_snake),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
+                         dilation=3, use_snake=use_snake),
+            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
+                         dilation=9, use_snake=use_snake),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+class OobleckEncoder(nn.Module):
+    def __init__(self, 
+                 in_channels=2, 
+                 channels=128, 
+                 latent_dim=32, 
+                 c_mults = [1, 2, 4, 8], 
+                 strides = [2, 4, 8, 8],
+                 use_snake=False,
+                 antialias_activation=False
+        ):
+        super().__init__()
+          
+        c_mults = [1] + c_mults
+
+        self.depth = len(c_mults)
+
+        layers = [
+            WNConv1d(in_channels=in_channels, out_channels=c_mults[0] * channels, kernel_size=7, padding=3)
+        ]
+        
+        for i in range(self.depth-1):
+            layers += [EncoderBlock(in_channels=c_mults[i]*channels, out_channels=c_mults[i+1]*channels, stride=strides[i], use_snake=use_snake)]
+
+        layers += [
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels),
+            WNConv1d(in_channels=c_mults[-1]*channels, out_channels=latent_dim, kernel_size=3, padding=1)
+        ]
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class OobleckDecoder(nn.Module):
+    def __init__(self, 
+                 out_channels=2, 
+                 channels=128, 
+                 latent_dim=32, 
+                 c_mults = [1, 2, 4, 8], 
+                 strides = [2, 4, 8, 8],
+                 use_snake=False,
+                 antialias_activation=False,
+                 use_nearest_upsample=False,
+                 final_tanh=True):
+        super().__init__()
+
+        c_mults = [1] + c_mults
+        
+        self.depth = len(c_mults)
+
+        layers = [
+            WNConv1d(in_channels=latent_dim, out_channels=c_mults[-1]*channels, kernel_size=7, padding=3),
+        ]
+        
+        for i in range(self.depth-1, 0, -1):
+            layers += [DecoderBlock(
+                in_channels=c_mults[i]*channels, 
+                out_channels=c_mults[i-1]*channels, 
+                stride=strides[i-1], 
+                use_snake=use_snake, 
+                antialias_activation=antialias_activation,
+                use_nearest_upsample=use_nearest_upsample
+                )
+            ]
+
+        layers += [
+            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels),
+            WNConv1d(in_channels=c_mults[0] * channels, out_channels=out_channels, kernel_size=7, padding=3, bias=False),
+            nn.Tanh() if final_tanh else nn.Identity()
+        ]
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class DACEncoderWrapper(nn.Module):
+    def __init__(self, in_channels=1, **kwargs):
+        super().__init__()
+
+        from dac.model.dac import Encoder as DACEncoder
+
+        latent_dim = kwargs.pop("latent_dim", None)
+
+        encoder_out_dim = kwargs["d_model"] * (2 ** len(kwargs["strides"]))
+        self.encoder = DACEncoder(d_latent=encoder_out_dim, **kwargs)
+        self.latent_dim = latent_dim
+
+        # Latent-dim support was added to DAC after this was first written, and implemented differently, so this is for backwards compatibility
+        self.proj_out = nn.Conv1d(self.encoder.enc_dim, latent_dim, kernel_size=1) if latent_dim is not None else nn.Identity()
+
+        if in_channels != 1:
+            self.encoder.block[0] = WNConv1d(in_channels, kwargs.get("d_model", 64), kernel_size=7, padding=3)
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.proj_out(x)
+        return x
+
+class DACDecoderWrapper(nn.Module):
+    def __init__(self, latent_dim, out_channels=1, **kwargs):
+        super().__init__()
+
+        from dac.model.dac import Decoder as DACDecoder
+
+        self.decoder = DACDecoder(**kwargs, input_channel = latent_dim, d_out=out_channels)
+
+        self.latent_dim = latent_dim
+
+    def forward(self, x):
+        return self.decoder(x)
+
+class AudioAutoencoder(nn.Module):
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        latent_dim,
+        downsampling_ratio,
+        sample_rate,
+        io_channels=2,
+        bottleneck: Bottleneck = None,
+        pretransform: Pretransform = None,
+        in_channels = None,
+        out_channels = None,
+        soft_clip = False
+    ):
+        super().__init__()
+
+        self.downsampling_ratio = downsampling_ratio
+        self.sample_rate = sample_rate
+
+        self.latent_dim = latent_dim
+        self.io_channels = io_channels
+        self.in_channels = io_channels
+        self.out_channels = io_channels
+
+        self.min_length = self.downsampling_ratio
+
+        if in_channels is not None:
+            self.in_channels = in_channels
+
+        if out_channels is not None:
+            self.out_channels = out_channels
+
+        self.bottleneck = bottleneck
+
+        self.encoder = encoder
+
+        self.decoder = decoder
+
+        self.pretransform = pretransform
+
+        self.soft_clip = soft_clip
+ 
+        self.is_discrete = self.bottleneck is not None and self.bottleneck.is_discrete
+
+    def encode(self, audio, return_info=False, skip_pretransform=False, iterate_batch=False, **kwargs):
+
+        info = {}
+
+        if self.pretransform is not None and not skip_pretransform:
+            if self.pretransform.enable_grad:
+                if iterate_batch:
+                    audios = []
+                    for i in range(audio.shape[0]):
+                        audios.append(self.pretransform.encode(audio[i:i+1]))
+                    audio = torch.cat(audios, dim=0)
+                else:
+                    audio = self.pretransform.encode(audio)
+            else:
+                with torch.no_grad():
+                    if iterate_batch:
+                        audios = []
+                        for i in range(audio.shape[0]):
+                            audios.append(self.pretransform.encode(audio[i:i+1]))
+                        audio = torch.cat(audios, dim=0)
+                    else:
+                        audio = self.pretransform.encode(audio)
+
+        if self.encoder is not None:
+            if iterate_batch:
+                latents = []
+                for i in range(audio.shape[0]):
+                    latents.append(self.encoder(audio[i:i+1]))
+                latents = torch.cat(latents, dim=0)
+            else:
+                latents = self.encoder(audio)
+        else:
+            latents = audio
+
+        if self.bottleneck is not None:
+            # TODO: Add iterate batch logic, needs to merge the info dicts
+            latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs)
+
+            info.update(bottleneck_info)
+        
+        if return_info:
+            return latents, info
+
+        return latents
+
+    def decode(self, latents, iterate_batch=False, **kwargs):
+
+        if self.bottleneck is not None:
+            if iterate_batch:
+                decoded = []
+                for i in range(latents.shape[0]):
+                    decoded.append(self.bottleneck.decode(latents[i:i+1]))
+                latents = torch.cat(decoded, dim=0)
+            else:
+                latents = self.bottleneck.decode(latents)
+
+        if iterate_batch:
+            decoded = []
+            for i in range(latents.shape[0]):
+                decoded.append(self.decoder(latents[i:i+1]))
+            decoded = torch.cat(decoded, dim=0)
+        else:
+            decoded = self.decoder(latents, **kwargs)
+
+        if self.pretransform is not None:
+            if self.pretransform.enable_grad:
+                if iterate_batch:
+                    decodeds = []
+                    for i in range(decoded.shape[0]):
+                        decodeds.append(self.pretransform.decode(decoded[i:i+1]))
+                    decoded = torch.cat(decodeds, dim=0)
+                else:
+                    decoded = self.pretransform.decode(decoded)
+            else:
+                with torch.no_grad():
+                    if iterate_batch:
+                        decodeds = []
+                        for i in range(latents.shape[0]):
+                            decodeds.append(self.pretransform.decode(decoded[i:i+1]))
+                        decoded = torch.cat(decodeds, dim=0)
+                    else:
+                        decoded = self.pretransform.decode(decoded)
+
+        if self.soft_clip:
+            decoded = torch.tanh(decoded)
+        
+        return decoded
+          
+    def decode_tokens(self, tokens, **kwargs):
+        '''
+        Decode discrete tokens to audio
+        Only works with discrete autoencoders
+        '''
+
+        assert isinstance(self.bottleneck, DiscreteBottleneck), "decode_tokens only works with discrete autoencoders"
+
+        latents = self.bottleneck.decode_tokens(tokens, **kwargs)
+
+        return self.decode(latents, **kwargs)
+        
+    
+    def preprocess_audio_for_encoder(self, audio, in_sr):
+        '''
+        Preprocess single audio tensor (Channels x Length) to be compatible with the encoder.
+        If the model is mono, stereo audio will be converted to mono.
+        Audio will be silence-padded to be a multiple of the model's downsampling ratio.
+        Audio will be resampled to the model's sample rate. 
+        The output will have batch size 1 and be shape (1 x Channels x Length)
+        '''
+        return self.preprocess_audio_list_for_encoder([audio], [in_sr])
+
+    def preprocess_audio_list_for_encoder(self, audio_list, in_sr_list):
+        '''
+        Preprocess a [list] of audio (Channels x Length) into a batch tensor to be compatable with the encoder. 
+        The audio in that list can be of different lengths and channels. 
+        in_sr can be an integer or list. If it's an integer it will be assumed it is the input sample_rate for every audio.
+        All audio will be resampled to the model's sample rate. 
+        Audio will be silence-padded to the longest length, and further padded to be a multiple of the model's downsampling ratio. 
+        If the model is mono, all audio will be converted to mono. 
+        The output will be a tensor of shape (Batch x Channels x Length)
+        '''
+        batch_size = len(audio_list)
+        if isinstance(in_sr_list, int):
+            in_sr_list = [in_sr_list]*batch_size
+        assert len(in_sr_list) == batch_size, "list of sample rates must be the same length of audio_list"
+        new_audio = []
+        max_length = 0
+        # resample & find the max length
+        for i in range(batch_size):
+            audio = audio_list[i]
+            in_sr = in_sr_list[i]
+            if len(audio.shape) == 3 and audio.shape[0] == 1:
+                # batchsize 1 was given by accident. Just squeeze it.
+                audio = audio.squeeze(0)
+            elif len(audio.shape) == 1:
+                # Mono signal, channel dimension is missing, unsqueeze it in
+                audio = audio.unsqueeze(0)
+            assert len(audio.shape)==2, "Audio should be shape (Channels x Length) with no batch dimension" 
+            # Resample audio
+            if in_sr != self.sample_rate:
+                resample_tf = T.Resample(in_sr, self.sample_rate).to(audio.device)
+                audio = resample_tf(audio)
+            new_audio.append(audio)
+            if audio.shape[-1] > max_length:
+                max_length = audio.shape[-1]
+        # Pad every audio to the same length, multiple of model's downsampling ratio
+        padded_audio_length = max_length + (self.min_length - (max_length % self.min_length)) % self.min_length
+        for i in range(batch_size):
+            # Pad it & if necessary, mixdown/duplicate stereo/mono channels to support model
+            new_audio[i] = prepare_audio(new_audio[i], in_sr=in_sr, target_sr=in_sr, target_length=padded_audio_length, 
+                target_channels=self.in_channels, device=new_audio[i].device).squeeze(0)
+        # convert to tensor 
+        return torch.stack(new_audio) 
+
+    def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs):
+        '''
+        Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder.
+        If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap.
+        Overlap and chunk_size params are both measured in number of latents (not audio samples) 
+        # and therefore you likely could use the same values with decode_audio. 
+        A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. 
+        Every autoencoder will have a different receptive field size, and thus ideal overlap.
+        You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff.
+        The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
+        Smaller chunk_size uses less memory, but more compute.
+        The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
+        For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
+        '''
+        if not chunked:
+            # default behavior. Encode the entire audio in parallel
+            return self.encode(audio, **kwargs)
+        else:
+            # CHUNKED ENCODING
+            # samples_per_latent is just the downsampling ratio (which is also the upsampling ratio)
+            samples_per_latent = self.downsampling_ratio
+            total_size = audio.shape[2] # in samples
+            batch_size = audio.shape[0]
+            chunk_size *= samples_per_latent # converting metric in latents to samples
+            overlap *= samples_per_latent # converting metric in latents to samples
+            hop_size = chunk_size - overlap
+            chunks = []
+            for i in range(0, total_size - chunk_size + 1, hop_size):
+                chunk = audio[:,:,i:i+chunk_size]
+                chunks.append(chunk)
+            if i+chunk_size != total_size:
+                # Final chunk
+                chunk = audio[:,:,-chunk_size:]
+                chunks.append(chunk)
+            chunks = torch.stack(chunks)
+            num_chunks = chunks.shape[0]
+            # Note: y_size might be a different value from the latent length used in diffusion training
+            # because we can encode audio of varying lengths
+            # However, the audio should've been padded to a multiple of samples_per_latent by now.
+            y_size = total_size // samples_per_latent
+            # Create an empty latent, we will populate it with chunks as we encode them
+            y_final = torch.zeros((batch_size,self.latent_dim,y_size)).to(audio.device)
+            for i in range(num_chunks):
+                x_chunk = chunks[i,:]
+                # encode the chunk
+                y_chunk = self.encode(x_chunk)
+                # figure out where to put the audio along the time domain
+                if i == num_chunks-1:
+                    # final chunk always goes at the end
+                    t_end = y_size
+                    t_start = t_end - y_chunk.shape[2]
+                else:
+                    t_start = i * hop_size // samples_per_latent
+                    t_end = t_start + chunk_size // samples_per_latent
+                #  remove the edges of the overlaps
+                ol = overlap//samples_per_latent//2
+                chunk_start = 0
+                chunk_end = y_chunk.shape[2]
+                if i > 0:
+                    # no overlap for the start of the first chunk
+                    t_start += ol
+                    chunk_start += ol
+                if i < num_chunks-1:
+                    # no overlap for the end of the last chunk
+                    t_end -= ol
+                    chunk_end -= ol
+                # paste the chunked audio into our y_final output audio
+                y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
+            return y_final
+    
+    def decode_audio(self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs):
+        '''
+        Decode latents to audio. 
+        If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, both of which are measured in number of latents. 
+        A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size. 
+        Every autoencoder will have a different receptive field size, and thus ideal overlap.
+        You can determine it empirically by diffing unchunked vs chunked audio and looking at maximum diff.
+        The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
+        Smaller chunk_size uses less memory, but more compute.
+        The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
+        For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
+        '''
+        if not chunked:
+            # default behavior. Decode the entire latent in parallel
+            return self.decode(latents, **kwargs)
+        else:
+            # chunked decoding
+            hop_size = chunk_size - overlap
+            total_size = latents.shape[2]
+            batch_size = latents.shape[0]
+            chunks = []
+            for i in range(0, total_size - chunk_size + 1, hop_size):
+                chunk = latents[:,:,i:i+chunk_size]
+                chunks.append(chunk)
+            if i+chunk_size != total_size:
+                # Final chunk
+                chunk = latents[:,:,-chunk_size:]
+                chunks.append(chunk)
+            chunks = torch.stack(chunks)
+            num_chunks = chunks.shape[0]
+            # samples_per_latent is just the downsampling ratio
+            samples_per_latent = self.downsampling_ratio
+            # Create an empty waveform, we will populate it with chunks as decode them
+            y_size = total_size * samples_per_latent
+            y_final = torch.zeros((batch_size,self.out_channels,y_size)).to(latents.device)
+            for i in range(num_chunks):
+                x_chunk = chunks[i,:]
+                # decode the chunk
+                y_chunk = self.decode(x_chunk)
+                # figure out where to put the audio along the time domain
+                if i == num_chunks-1:
+                    # final chunk always goes at the end
+                    t_end = y_size
+                    t_start = t_end - y_chunk.shape[2]
+                else:
+                    t_start = i * hop_size * samples_per_latent
+                    t_end = t_start + chunk_size * samples_per_latent
+                #  remove the edges of the overlaps
+                ol = (overlap//2) * samples_per_latent
+                chunk_start = 0
+                chunk_end = y_chunk.shape[2]
+                if i > 0:
+                    # no overlap for the start of the first chunk
+                    t_start += ol
+                    chunk_start += ol
+                if i < num_chunks-1:
+                    # no overlap for the end of the last chunk
+                    t_end -= ol
+                    chunk_end -= ol
+                # paste the chunked audio into our y_final output audio
+                y_final[:,:,t_start:t_end] = y_chunk[:,:,chunk_start:chunk_end]
+            return y_final
+
+    
+class DiffusionAutoencoder(AudioAutoencoder):
+    def __init__(
+        self,
+        diffusion: ConditionedDiffusionModel,
+        diffusion_downsampling_ratio,
+        *args,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.diffusion = diffusion
+
+        self.min_length = self.downsampling_ratio * diffusion_downsampling_ratio
+
+        if self.encoder is not None:
+            # Shrink the initial encoder parameters to avoid saturated latents
+            with torch.no_grad():
+                for param in self.encoder.parameters():
+                    param *= 0.5
+
+    def decode(self, latents, steps=100):
+
+        upsampled_length = latents.shape[2] * self.downsampling_ratio
+
+        if self.bottleneck is not None:
+            latents = self.bottleneck.decode(latents)
+
+        if self.decoder is not None:
+            latents = self.decode(latents)
+    
+        # Upsample latents to match diffusion length
+        if latents.shape[2] != upsampled_length:
+            latents = F.interpolate(latents, size=upsampled_length, mode='nearest')
+
+        noise = torch.randn(latents.shape[0], self.io_channels, upsampled_length, device=latents.device)
+        decoded = sample(self.diffusion, noise, steps, 0, input_concat_cond=latents)
+
+        if self.pretransform is not None:
+            if self.pretransform.enable_grad:
+                decoded = self.pretransform.decode(decoded)
+            else:
+                with torch.no_grad():
+                    decoded = self.pretransform.decode(decoded)
+
+        return decoded
+        
+# AE factories
+
+def create_encoder_from_config(encoder_config: Dict[str, Any]):
+    encoder_type = encoder_config.get("type", None)
+    assert encoder_type is not None, "Encoder type must be specified"
+
+    if encoder_type == "oobleck":
+        encoder = OobleckEncoder(
+            **encoder_config["config"]
+        )
+    
+    elif encoder_type == "seanet":
+        from encodec.modules import SEANetEncoder
+        seanet_encoder_config = encoder_config["config"]
+
+        #SEANet encoder expects strides in reverse order
+        seanet_encoder_config["ratios"] = list(reversed(seanet_encoder_config.get("ratios", [2, 2, 2, 2, 2])))
+        encoder = SEANetEncoder(
+            **seanet_encoder_config
+        )
+    elif encoder_type == "dac":
+        dac_config = encoder_config["config"]
+
+        encoder = DACEncoderWrapper(**dac_config)
+    elif encoder_type == "local_attn":
+        from .local_attention import TransformerEncoder1D
+
+        local_attn_config = encoder_config["config"]
+
+        encoder = TransformerEncoder1D(
+            **local_attn_config
+        )
+    else:
+        raise ValueError(f"Unknown encoder type {encoder_type}")
+    
+    requires_grad = encoder_config.get("requires_grad", True)
+    if not requires_grad:
+        for param in encoder.parameters():
+            param.requires_grad = False
+
+    return encoder
+
+def create_decoder_from_config(decoder_config: Dict[str, Any]):
+    decoder_type = decoder_config.get("type", None)
+    assert decoder_type is not None, "Decoder type must be specified"
+
+    if decoder_type == "oobleck":
+        decoder = OobleckDecoder(
+            **decoder_config["config"]
+        )
+    elif decoder_type == "seanet":
+        from encodec.modules import SEANetDecoder
+
+        decoder = SEANetDecoder(
+            **decoder_config["config"]
+        )
+    elif decoder_type == "dac":
+        dac_config = decoder_config["config"]
+
+        decoder = DACDecoderWrapper(**dac_config)
+    elif decoder_type == "local_attn":
+        from .local_attention import TransformerDecoder1D
+
+        local_attn_config = decoder_config["config"]
+
+        decoder = TransformerDecoder1D(
+            **local_attn_config
+        )
+    else:
+        raise ValueError(f"Unknown decoder type {decoder_type}")
+    
+    requires_grad = decoder_config.get("requires_grad", True)
+    if not requires_grad:
+        for param in decoder.parameters():
+            param.requires_grad = False
+
+    return decoder
+
+def create_autoencoder_from_config(config: Dict[str, Any]):
+    
+    ae_config = config["model"]
+
+    encoder = create_encoder_from_config(ae_config["encoder"])
+    decoder = create_decoder_from_config(ae_config["decoder"])
+
+    bottleneck = ae_config.get("bottleneck", None)
+
+    latent_dim = ae_config.get("latent_dim", None)
+    assert latent_dim is not None, "latent_dim must be specified in model config"
+    downsampling_ratio = ae_config.get("downsampling_ratio", None)
+    assert downsampling_ratio is not None, "downsampling_ratio must be specified in model config"
+    io_channels = ae_config.get("io_channels", None)
+    assert io_channels is not None, "io_channels must be specified in model config"
+    sample_rate = config.get("sample_rate", None)
+    assert sample_rate is not None, "sample_rate must be specified in model config"
+
+    in_channels = ae_config.get("in_channels", None)
+    out_channels = ae_config.get("out_channels", None)
+
+    pretransform = ae_config.get("pretransform", None)
+
+    if pretransform is not None:
+        pretransform = create_pretransform_from_config(pretransform, sample_rate)
+
+    if bottleneck is not None:
+        bottleneck = create_bottleneck_from_config(bottleneck)
+
+    soft_clip = ae_config["decoder"].get("soft_clip", False)
+
+    return AudioAutoencoder(
+        encoder,
+        decoder,
+        io_channels=io_channels,
+        latent_dim=latent_dim,
+        downsampling_ratio=downsampling_ratio,
+        sample_rate=sample_rate,
+        bottleneck=bottleneck,
+        pretransform=pretransform,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        soft_clip=soft_clip
+    )
+
+def create_diffAE_from_config(config: Dict[str, Any]):
+    
+    diffae_config = config["model"]
+
+    if "encoder" in diffae_config:
+        encoder = create_encoder_from_config(diffae_config["encoder"])
+    else:
+        encoder = None
+
+    if "decoder" in diffae_config:
+        decoder = create_decoder_from_config(diffae_config["decoder"])
+    else:
+        decoder = None
+
+    diffusion_model_type = diffae_config["diffusion"]["type"]
+
+    if diffusion_model_type == "DAU1d":
+        diffusion = DAU1DCondWrapper(**diffae_config["diffusion"]["config"])
+    elif diffusion_model_type == "adp_1d":
+        diffusion = UNet1DCondWrapper(**diffae_config["diffusion"]["config"])
+    elif diffusion_model_type == "dit":
+        diffusion = DiTWrapper(**diffae_config["diffusion"]["config"])
+
+    latent_dim = diffae_config.get("latent_dim", None)
+    assert latent_dim is not None, "latent_dim must be specified in model config"
+    downsampling_ratio = diffae_config.get("downsampling_ratio", None)
+    assert downsampling_ratio is not None, "downsampling_ratio must be specified in model config"
+    io_channels = diffae_config.get("io_channels", None)
+    assert io_channels is not None, "io_channels must be specified in model config"
+    sample_rate = config.get("sample_rate", None)
+    assert sample_rate is not None, "sample_rate must be specified in model config"
+
+    bottleneck = diffae_config.get("bottleneck", None)
+
+    pretransform = diffae_config.get("pretransform", None)
+
+    if pretransform is not None:
+        pretransform = create_pretransform_from_config(pretransform, sample_rate)
+
+    if bottleneck is not None:
+        bottleneck = create_bottleneck_from_config(bottleneck)
+
+    diffusion_downsampling_ratio = None,
+
+    if diffusion_model_type == "DAU1d":
+        diffusion_downsampling_ratio = np.prod(diffae_config["diffusion"]["config"]["strides"])
+    elif diffusion_model_type == "adp_1d":
+        diffusion_downsampling_ratio = np.prod(diffae_config["diffusion"]["config"]["factors"])
+    elif diffusion_model_type == "dit":
+        diffusion_downsampling_ratio = 1
+
+    return DiffusionAutoencoder(
+        encoder=encoder,
+        decoder=decoder,
+        diffusion=diffusion,
+        io_channels=io_channels,
+        sample_rate=sample_rate,
+        latent_dim=latent_dim,
+        downsampling_ratio=downsampling_ratio,
+        diffusion_downsampling_ratio=diffusion_downsampling_ratio,
+        bottleneck=bottleneck,
+        pretransform=pretransform
+    )

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/models/bottleneck.py

@@ -0,0 +1,355 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from einops import rearrange
+from vector_quantize_pytorch import ResidualVQ, FSQ
+from dac.nn.quantize import ResidualVectorQuantize as DACResidualVQ
+
+class Bottleneck(nn.Module):
+    def __init__(self, is_discrete: bool = False):
+        super().__init__()
+
+        self.is_discrete = is_discrete
+
+    def encode(self, x, return_info=False, **kwargs):
+        raise NotImplementedError
+
+    def decode(self, x):
+        raise NotImplementedError
+
+class DiscreteBottleneck(Bottleneck):
+    def __init__(self, num_quantizers, codebook_size, tokens_id):
+        super().__init__(is_discrete=True)
+
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.tokens_id = tokens_id
+
+    def decode_tokens(self, codes, **kwargs):
+        raise NotImplementedError
+    
+class TanhBottleneck(Bottleneck):
+    def __init__(self):
+        super().__init__(is_discrete=False)
+        self.tanh = nn.Tanh()
+
+    def encode(self, x, return_info=False):
+        info = {}
+
+        x = torch.tanh(x)
+
+        if return_info:
+            return x, info
+        else:
+            return x
+
+    def decode(self, x):
+        return x
+
+def vae_sample(mean, scale):
+        stdev = nn.functional.softplus(scale) + 1e-4
+        var = stdev * stdev
+        logvar = torch.log(var)
+        latents = torch.randn_like(mean) * stdev + mean
+
+        kl = (mean * mean + var - logvar - 1).sum(1).mean()
+
+        return latents, kl
+
+class VAEBottleneck(Bottleneck):
+    def __init__(self):
+        super().__init__(is_discrete=False)
+
+    def encode(self, x, return_info=False, **kwargs):
+        info = {}
+
+        mean, scale = x.chunk(2, dim=1)
+
+        x, kl = vae_sample(mean, scale)
+
+        info["kl"] = kl
+
+        if return_info:
+            return x, info
+        else:
+            return x
+
+    def decode(self, x):
+        return x
+
+def compute_mean_kernel(x, y):
+        kernel_input = (x[:, None] - y[None]).pow(2).mean(2) / x.shape[-1]
+        return torch.exp(-kernel_input).mean()
+
+def compute_mmd(latents):
+    latents_reshaped = latents.permute(0, 2, 1).reshape(-1, latents.shape[1])
+    noise = torch.randn_like(latents_reshaped)
+
+    latents_kernel = compute_mean_kernel(latents_reshaped, latents_reshaped)
+    noise_kernel = compute_mean_kernel(noise, noise)
+    latents_noise_kernel = compute_mean_kernel(latents_reshaped, noise)
+    
+    mmd = latents_kernel + noise_kernel - 2 * latents_noise_kernel
+    return mmd.mean()
+
+class WassersteinBottleneck(Bottleneck):
+    def __init__(self, noise_augment_dim: int = 0, bypass_mmd: bool = False):
+        super().__init__(is_discrete=False)
+
+        self.noise_augment_dim = noise_augment_dim
+        self.bypass_mmd = bypass_mmd
+    
+    def encode(self, x, return_info=False):
+        info = {}
+
+        if self.training and return_info:
+            if self.bypass_mmd:
+                mmd = torch.tensor(0.0)
+            else:
+                mmd = compute_mmd(x)
+                
+            info["mmd"] = mmd
+        
+        if return_info:
+            return x, info
+        
+        return x
+
+    def decode(self, x):
+
+        if self.noise_augment_dim > 0:
+            noise = torch.randn(x.shape[0], self.noise_augment_dim,
+                                x.shape[-1]).type_as(x)
+            x = torch.cat([x, noise], dim=1)
+
+        return x
+
+class L2Bottleneck(Bottleneck):
+    def __init__(self):
+        super().__init__(is_discrete=False)
+    
+    def encode(self, x, return_info=False):
+        info = {}
+
+        x = F.normalize(x, dim=1)
+
+        if return_info:
+            return x, info
+        else:
+            return x
+        
+    def decode(self, x):
+        return F.normalize(x, dim=1)
+        
+class RVQBottleneck(DiscreteBottleneck):
+    def __init__(self, **quantizer_kwargs):
+        super().__init__(num_quantizers = quantizer_kwargs["num_quantizers"], codebook_size = quantizer_kwargs["codebook_size"], tokens_id = "quantizer_indices")
+        self.quantizer = ResidualVQ(**quantizer_kwargs)
+        self.num_quantizers = quantizer_kwargs["num_quantizers"]
+
+    def encode(self, x, return_info=False, **kwargs):
+        info = {}
+
+        x = rearrange(x, "b c n -> b n c")
+        x, indices, loss = self.quantizer(x)
+        x = rearrange(x, "b n c -> b c n")
+
+        info["quantizer_indices"] = indices
+        info["quantizer_loss"] = loss.mean()
+
+        if return_info:
+            return x, info
+        else:
+            return x
+        
+    def decode(self, x):
+        return x
+    
+    def decode_tokens(self, codes, **kwargs):
+        latents = self.quantizer.get_outputs_from_indices(codes)
+
+        return self.decode(latents, **kwargs)
+    
+class RVQVAEBottleneck(DiscreteBottleneck):
+    def __init__(self, **quantizer_kwargs):
+        super().__init__(num_quantizers = quantizer_kwargs["num_quantizers"], codebook_size = quantizer_kwargs["codebook_size"], tokens_id = "quantizer_indices")
+        self.quantizer = ResidualVQ(**quantizer_kwargs)
+        self.num_quantizers = quantizer_kwargs["num_quantizers"]
+
+    def encode(self, x, return_info=False):
+        info = {}
+
+        x, kl = vae_sample(*x.chunk(2, dim=1))
+
+        info["kl"] = kl
+
+        x = rearrange(x, "b c n -> b n c")
+        x, indices, loss = self.quantizer(x)
+        x = rearrange(x, "b n c -> b c n")
+
+        info["quantizer_indices"] = indices
+        info["quantizer_loss"] = loss.mean()
+
+        if return_info:
+            return x, info
+        else:
+            return x
+        
+    def decode(self, x):
+        return x
+    
+    def decode_tokens(self, codes, **kwargs):
+        latents = self.quantizer.get_outputs_from_indices(codes)
+
+        return self.decode(latents, **kwargs)
+
+class DACRVQBottleneck(DiscreteBottleneck):
+    def __init__(self, quantize_on_decode=False, noise_augment_dim=0, **quantizer_kwargs):
+        super().__init__(num_quantizers = quantizer_kwargs["n_codebooks"], codebook_size = quantizer_kwargs["codebook_size"], tokens_id = "codes")
+        self.quantizer = DACResidualVQ(**quantizer_kwargs)
+        self.num_quantizers = quantizer_kwargs["n_codebooks"]
+        self.quantize_on_decode = quantize_on_decode
+        self.noise_augment_dim = noise_augment_dim
+
+    def encode(self, x, return_info=False, **kwargs):
+        info = {}
+
+        info["pre_quantizer"] = x
+
+        if self.quantize_on_decode:
+            return x, info if return_info else x
+
+        z, codes, latents, commitment_loss, codebook_loss = self.quantizer(x, **kwargs)
+
+        output = {
+            "z": z,
+            "codes": codes,
+            "latents": latents,
+            "vq/commitment_loss": commitment_loss,
+            "vq/codebook_loss": codebook_loss,
+        }
+
+        output["vq/commitment_loss"] /= self.num_quantizers
+        output["vq/codebook_loss"] /= self.num_quantizers
+
+        info.update(output)
+
+        if return_info:
+            return output["z"], info
+        
+        return output["z"]
+    
+    def decode(self, x):
+
+        if self.quantize_on_decode:
+            x = self.quantizer(x)[0]
+
+        if self.noise_augment_dim > 0:
+            noise = torch.randn(x.shape[0], self.noise_augment_dim,
+                                x.shape[-1]).type_as(x)
+            x = torch.cat([x, noise], dim=1)
+
+        return x
+    
+    def decode_tokens(self, codes, **kwargs):
+        latents, _, _ = self.quantizer.from_codes(codes)
+
+        return self.decode(latents, **kwargs)
+
+class DACRVQVAEBottleneck(DiscreteBottleneck):
+    def __init__(self, quantize_on_decode=False, **quantizer_kwargs):
+        super().__init__(num_quantizers = quantizer_kwargs["n_codebooks"], codebook_size = quantizer_kwargs["codebook_size"], tokens_id = "codes")
+        self.quantizer = DACResidualVQ(**quantizer_kwargs)
+        self.num_quantizers = quantizer_kwargs["n_codebooks"]
+        self.quantize_on_decode = quantize_on_decode
+
+    def encode(self, x, return_info=False, n_quantizers: int = None):
+        info = {}
+
+        mean, scale = x.chunk(2, dim=1)
+
+        x, kl = vae_sample(mean, scale)
+
+        info["pre_quantizer"] = x
+        info["kl"] = kl
+
+        if self.quantize_on_decode:
+            return x, info if return_info else x
+
+        z, codes, latents, commitment_loss, codebook_loss = self.quantizer(x, n_quantizers=n_quantizers)
+
+        output = {
+            "z": z,
+            "codes": codes,
+            "latents": latents,
+            "vq/commitment_loss": commitment_loss,
+            "vq/codebook_loss": codebook_loss,
+        }
+
+        output["vq/commitment_loss"] /= self.num_quantizers
+        output["vq/codebook_loss"] /= self.num_quantizers
+
+        info.update(output)
+
+        if return_info:
+            return output["z"], info
+        
+        return output["z"]
+    
+    def decode(self, x):
+
+        if self.quantize_on_decode:
+            x = self.quantizer(x)[0]
+
+        return x
+
+    def decode_tokens(self, codes, **kwargs):
+        latents, _, _ = self.quantizer.from_codes(codes)
+
+        return self.decode(latents, **kwargs)
+    
+class FSQBottleneck(DiscreteBottleneck):
+    def __init__(self, noise_augment_dim=0, **kwargs):
+        super().__init__(num_quantizers = kwargs.get("num_codebooks", 1), codebook_size = np.prod(kwargs["levels"]), tokens_id = "quantizer_indices")
+
+        self.noise_augment_dim = noise_augment_dim
+
+        self.quantizer = FSQ(**kwargs, allowed_dtypes=[torch.float16, torch.float32, torch.float64])
+
+    def encode(self, x, return_info=False):
+        info = {}
+
+        orig_dtype = x.dtype
+        x = x.float()
+
+        x = rearrange(x, "b c n -> b n c")
+        x, indices = self.quantizer(x)
+        x = rearrange(x, "b n c -> b c n")
+
+        x = x.to(orig_dtype)
+
+        # Reorder indices to match the expected format
+        indices = rearrange(indices, "b n q -> b q n")
+
+        info["quantizer_indices"] = indices
+
+        if return_info:
+            return x, info
+        else:
+            return x
+        
+    def decode(self, x):
+
+        if self.noise_augment_dim > 0:
+            noise = torch.randn(x.shape[0], self.noise_augment_dim,
+                                x.shape[-1]).type_as(x)
+            x = torch.cat([x, noise], dim=1)
+
+        return x
+    
+    def decode_tokens(self, tokens, **kwargs):
+        latents = self.quantizer.indices_to_codes(tokens)
+
+        return self.decode(latents, **kwargs)

Levo_Song_Generation/SongGeneration-Runtime/third_party/stable_audio_tools/stable_audio_tools/models/conditioners.py

@@ -0,0 +1,561 @@
+#Heavily influenced by https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/modules/conditioners.py
+
+import torch
+import logging, warnings
+import string
+import typing as tp
+import gc
+
+from .adp import NumberEmbedder
+from ..inference.utils import set_audio_channels
+from .factory import create_pretransform_from_config
+from .pretransforms import Pretransform
+from ..training.utils import copy_state_dict
+from .utils import load_ckpt_state_dict
+
+from torch import nn
+
+class Conditioner(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            output_dim: int,
+            project_out: bool = False
+            ):
+        
+        super().__init__()
+
+        self.dim = dim
+        self.output_dim = output_dim
+        self.proj_out = nn.Linear(dim, output_dim) if (dim != output_dim or project_out) else nn.Identity()
+
+    def forward(self, x: tp.Any) -> tp.Any:
+        raise NotImplementedError()
+    
+class IntConditioner(Conditioner):
+    def __init__(self, 
+                output_dim: int,
+                min_val: int=0,
+                max_val: int=512
+                ):
+        super().__init__(output_dim, output_dim)
+
+        self.min_val = min_val
+        self.max_val = max_val
+        self.int_embedder = nn.Embedding(max_val - min_val + 1, output_dim).requires_grad_(True)
+
+    def forward(self, ints: tp.List[int], device=None) -> tp.Any:
+            
+            #self.int_embedder.to(device)
+    
+            ints = torch.tensor(ints).to(device)
+            ints = ints.clamp(self.min_val, self.max_val)
+    
+            int_embeds = self.int_embedder(ints).unsqueeze(1)
+    
+            return [int_embeds, torch.ones(int_embeds.shape[0], 1).to(device)]
+
+class NumberConditioner(Conditioner):
+    '''
+        Conditioner that takes a list of floats, normalizes them for a given range, and returns a list of embeddings
+    '''
+    def __init__(self, 
+                output_dim: int,
+                min_val: float=0,
+                max_val: float=1
+                ):
+        super().__init__(output_dim, output_dim)
+
+        self.min_val = min_val
+        self.max_val = max_val
+
+        self.embedder = NumberEmbedder(features=output_dim)
+
+    def forward(self, floats: tp.List[float], device=None) -> tp.Any:
+    
+            # Cast the inputs to floats
+            floats = [float(x) for x in floats]
+
+            floats = torch.tensor(floats).to(device)
+
+            floats = floats.clamp(self.min_val, self.max_val)
+    
+            normalized_floats = (floats - self.min_val) / (self.max_val - self.min_val)
+
+            # Cast floats to same type as embedder
+            embedder_dtype = next(self.embedder.parameters()).dtype
+            normalized_floats = normalized_floats.to(embedder_dtype)
+
+            float_embeds = self.embedder(normalized_floats).unsqueeze(1)
+    
+            return [float_embeds, torch.ones(float_embeds.shape[0], 1).to(device)]
+
+class CLAPTextConditioner(Conditioner):
+    def __init__(self, 
+                 output_dim: int, 
+                 clap_ckpt_path,
+                 use_text_features = False,
+                 feature_layer_ix: int = -1,
+                 audio_model_type="HTSAT-base", 
+                 enable_fusion=True,
+                 project_out: bool = False,
+                 finetune: bool = False):
+        super().__init__(768 if use_text_features else 512, output_dim, project_out=project_out)
+
+        self.use_text_features = use_text_features
+        self.feature_layer_ix = feature_layer_ix
+        self.finetune = finetune
+
+        # Suppress logging from transformers
+        previous_level = logging.root.manager.disable
+        logging.disable(logging.ERROR)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            try:
+                import laion_clap
+                from laion_clap.clap_module.factory import load_state_dict as clap_load_state_dict
+                
+                model = laion_clap.CLAP_Module(enable_fusion=enable_fusion, amodel=audio_model_type, device='cpu')
+
+                if self.finetune:
+                    self.model = model
+                else: 
+                    self.__dict__["model"] = model
+
+                state_dict = clap_load_state_dict(clap_ckpt_path)
+                self.model.model.load_state_dict(state_dict, strict=False)
+
+                if self.finetune:
+                    self.model.model.text_branch.requires_grad_(True)
+                    self.model.model.text_branch.train()
+                else:
+                    self.model.model.text_branch.requires_grad_(False)
+                    self.model.model.text_branch.eval()
+
+            finally:
+                logging.disable(previous_level)
+
+        del self.model.model.audio_branch
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def get_clap_features(self, prompts, layer_ix=-2, device: tp.Any = "cuda"):
+        prompt_tokens = self.model.tokenizer(prompts)
+        attention_mask = prompt_tokens["attention_mask"].to(device=device, non_blocking=True)
+        prompt_features = self.model.model.text_branch(
+            input_ids=prompt_tokens["input_ids"].to(device=device, non_blocking=True),
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )["hidden_states"][layer_ix]
+
+        return prompt_features, attention_mask
+
+    def forward(self, texts: tp.List[str], device: tp.Any = "cuda") -> tp.Any:
+        self.model.to(device)
+
+        if self.use_text_features:
+            if len(texts) == 1:
+                text_features, text_attention_mask = self.get_clap_features([texts[0], ""], layer_ix=self.feature_layer_ix, device=device)
+                text_features = text_features[:1, ...]
+                text_attention_mask = text_attention_mask[:1, ...]
+            else:
+                text_features, text_attention_mask = self.get_clap_features(texts, layer_ix=self.feature_layer_ix, device=device)
+            return [self.proj_out(text_features), text_attention_mask]
+
+        # Fix for CLAP bug when only one text is passed
+        if len(texts) == 1:
+            text_embedding = self.model.get_text_embedding([texts[0], ""], use_tensor=True)[:1, ...]
+        else:
+            text_embedding = self.model.get_text_embedding(texts, use_tensor=True)
+
+        text_embedding = text_embedding.unsqueeze(1).to(device)
+
+        return [self.proj_out(text_embedding), torch.ones(text_embedding.shape[0], 1).to(device)]
+
+class CLAPAudioConditioner(Conditioner):
+    def __init__(self, 
+                 output_dim: int, 
+                 clap_ckpt_path,
+                 audio_model_type="HTSAT-base", 
+                 enable_fusion=True,
+                 project_out: bool = False):
+        super().__init__(512, output_dim, project_out=project_out)
+
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        # Suppress logging from transformers
+        previous_level = logging.root.manager.disable
+        logging.disable(logging.ERROR)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            try:
+                import laion_clap
+                from laion_clap.clap_module.factory import load_state_dict as clap_load_state_dict
+                
+                model = laion_clap.CLAP_Module(enable_fusion=enable_fusion, amodel=audio_model_type, device='cpu')
+
+                if self.finetune:
+                    self.model = model
+                else: 
+                    self.__dict__["model"] = model
+
+                state_dict = clap_load_state_dict(clap_ckpt_path)
+                self.model.model.load_state_dict(state_dict, strict=False)
+
+                if self.finetune:
+                    self.model.model.audio_branch.requires_grad_(True)
+                    self.model.model.audio_branch.train()
+                else:
+                    self.model.model.audio_branch.requires_grad_(False)
+                    self.model.model.audio_branch.eval()
+
+            finally:
+                logging.disable(previous_level)
+
+        del self.model.model.text_branch
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def forward(self, audios: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]] , device: tp.Any = "cuda") -> tp.Any:
+
+        self.model.to(device)
+
+        if isinstance(audios, list) or isinstance(audios, tuple):
+            audios = torch.cat(audios, dim=0)
+
+        # Convert to mono
+        mono_audios = audios.mean(dim=1)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            audio_embedding = self.model.get_audio_embedding_from_data(mono_audios.float(), use_tensor=True)
+
+        audio_embedding = audio_embedding.unsqueeze(1).to(device)
+
+        return [self.proj_out(audio_embedding), torch.ones(audio_embedding.shape[0], 1).to(device)]
+
+class T5Conditioner(Conditioner):
+
+    T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b",
+              "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large",
+              "google/flan-t5-xl", "google/flan-t5-xxl"]
+    
+    T5_MODEL_DIMS = {
+        "t5-small": 512,
+        "t5-base": 768,
+        "t5-large": 1024,
+        "t5-3b": 1024,
+        "t5-11b": 1024,
+        "t5-xl": 2048,
+        "t5-xxl": 4096,
+        "google/flan-t5-small": 512,
+        "google/flan-t5-base": 768,
+        "google/flan-t5-large": 1024,
+        "google/flan-t5-3b": 1024,
+        "google/flan-t5-11b": 1024,
+        "google/flan-t5-xl": 2048,
+        "google/flan-t5-xxl": 4096,
+    }
+
+    def __init__(
+            self,
+            output_dim: int,
+            t5_model_name: str = "t5-base",
+            max_length: str = 128,
+            enable_grad: bool = False,
+            project_out: bool = False
+    ):
+        assert t5_model_name in self.T5_MODELS, f"Unknown T5 model name: {t5_model_name}"
+        super().__init__(self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=project_out)
+        
+        from transformers import T5EncoderModel, AutoTokenizer
+
+        self.max_length = max_length
+        self.enable_grad = enable_grad
+
+        # Suppress logging from transformers
+        previous_level = logging.root.manager.disable
+        logging.disable(logging.ERROR)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            try:
+                # self.tokenizer = T5Tokenizer.from_pretrained(t5_model_name, model_max_length = max_length)
+                # model = T5EncoderModel.from_pretrained(t5_model_name, max_length=max_length).train(enable_grad).requires_grad_(enable_grad)
+                self.tokenizer = AutoTokenizer.from_pretrained(t5_model_name)
+                model = T5EncoderModel.from_pretrained(t5_model_name).train(enable_grad).requires_grad_(enable_grad).to(torch.float16)
+            finally:
+                logging.disable(previous_level)
+            
+        if self.enable_grad:
+            self.model = model
+        else: 
+            self.__dict__["model"] = model
+
+
+    def forward(self, texts: tp.List[str], device: tp.Union[torch.device, str]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+        
+        self.model.to(device)
+        self.proj_out.to(device)
+
+        encoded = self.tokenizer(
+            texts,
+            truncation=True,
+            max_length=self.max_length,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        input_ids = encoded["input_ids"].to(device)
+        attention_mask = encoded["attention_mask"].to(device).to(torch.bool)
+
+        self.model.eval()
+            
+        with torch.cuda.amp.autocast(dtype=torch.float16) and torch.set_grad_enabled(self.enable_grad):
+            embeddings = self.model(
+                input_ids=input_ids, attention_mask=attention_mask
+            )["last_hidden_state"]    
+            
+        embeddings = self.proj_out(embeddings.float())
+
+        embeddings = embeddings * attention_mask.unsqueeze(-1).float()
+
+        return embeddings, attention_mask
+    
+class PhonemeConditioner(Conditioner):
+    """
+    A conditioner that turns text into phonemes and embeds them using a lookup table
+    Only works for English text
+
+    Args:
+        output_dim: the dimension of the output embeddings
+        max_length: the maximum number of phonemes to embed
+        project_out: whether to add another linear projection to the output embeddings
+    """
+
+    def __init__(
+            self,
+            output_dim: int,
+            max_length: int = 1024,
+            project_out: bool = False,
+    ):
+        super().__init__(output_dim, output_dim, project_out=project_out)
+        
+        from g2p_en import G2p
+
+        self.max_length = max_length
+
+        self.g2p = G2p()
+
+        # Reserving 0 for padding, 1 for ignored
+        self.phoneme_embedder = nn.Embedding(len(self.g2p.phonemes) + 2, output_dim)
+
+    def forward(self, texts: tp.List[str], device: tp.Union[torch.device, str]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+        
+        self.phoneme_embedder.to(device)
+        self.proj_out.to(device)
+
+        batch_phonemes = [self.g2p(text) for text in texts] # shape [batch_size, length]
+        
+        phoneme_ignore = [" ", *string.punctuation]
+
+        # Remove ignored phonemes and cut to max length
+        batch_phonemes = [[p if p not in phoneme_ignore else "_" for p in phonemes] for phonemes in batch_phonemes]
+
+        # Convert to ids
+        phoneme_ids = [[self.g2p.p2idx[p] + 2 if p in self.g2p.p2idx else 1 for p in phonemes] for phonemes in batch_phonemes]
+
+        #Pad to match longest and make a mask tensor for the padding
+        longest = max([len(ids) for ids in phoneme_ids])
+        phoneme_ids = [ids + [0] * (longest - len(ids)) for ids in phoneme_ids]
+        
+        phoneme_ids = torch.tensor(phoneme_ids).to(device)
+
+        # Convert to embeddings
+        phoneme_embeds = self.phoneme_embedder(phoneme_ids)
+        
+        phoneme_embeds = self.proj_out(phoneme_embeds)
+
+        return phoneme_embeds, torch.ones(phoneme_embeds.shape[0], phoneme_embeds.shape[1]).to(device)
+  
+class TokenizerLUTConditioner(Conditioner):
+    """
+    A conditioner that embeds text using a lookup table on a pretrained tokenizer's vocabulary
+
+    Args:
+        tokenizer_name: the name of the tokenizer from the Hugging Face transformers library
+        output_dim: the dimension of the output embeddings
+        max_length: the maximum length of the text to embed
+        project_out: whether to add another linear projection to the output embeddings
+    """
+
+    def __init__(
+            self,
+            tokenizer_name: str, # Name of a tokenizer from the Hugging Face transformers library
+            output_dim: int,
+            max_length: int = 1024,
+            project_out: bool = False,
+    ):
+        super().__init__(output_dim, output_dim, project_out=project_out)
+        
+        from transformers import AutoTokenizer
+
+         # Suppress logging from transformers
+        previous_level = logging.root.manager.disable
+        logging.disable(logging.ERROR)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+            finally:
+                logging.disable(previous_level)
+
+        self.max_length = max_length
+
+        self.token_embedder = nn.Embedding(len(self.tokenizer), output_dim)
+
+    def forward(self, texts: tp.List[str], device: tp.Union[torch.device, str]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+        self.proj_out.to(device)
+
+        encoded = self.tokenizer(
+            texts,
+            truncation=True,
+            max_length=self.max_length,
+            padding="max_length",
+            return_tensors="pt",
+        )
+
+        input_ids = encoded["input_ids"].to(device)
+        attention_mask = encoded["attention_mask"].to(device).to(torch.bool)
+    
+        embeddings = self.token_embedder(input_ids)
+            
+        embeddings = self.proj_out(embeddings)
+
+        embeddings = embeddings * attention_mask.unsqueeze(-1).float()
+
+        return embeddings, attention_mask
+
+class PretransformConditioner(Conditioner):
+    """
+    A conditioner that uses a pretransform's encoder for conditioning
+
+    Args:
+        pretransform: an instantiated pretransform to use for conditioning
+        output_dim: the dimension of the output embeddings
+    """
+    def __init__(self, pretransform: Pretransform, output_dim: int):
+        super().__init__(pretransform.encoded_channels, output_dim)
+
+        self.pretransform = pretransform
+
+    def forward(self, audio: tp.Union[torch.Tensor, tp.List[torch.Tensor], tp.Tuple[torch.Tensor]], device: tp.Union[torch.device, str]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+
+        self.pretransform.to(device)
+        self.proj_out.to(device)
+
+        if isinstance(audio, list) or isinstance(audio, tuple):
+            audio = torch.cat(audio, dim=0)
+
+        # Convert audio to pretransform input channels
+        audio = set_audio_channels(audio, self.pretransform.io_channels)
+        
+        latents = self.pretransform.encode(audio)
+
+        latents = self.proj_out(latents)
+
+        return [latents, torch.ones(latents.shape[0], latents.shape[2]).to(latents.device)]
+
+class MultiConditioner(nn.Module):
+    """
+    A module that applies multiple conditioners to an input dictionary based on the keys
+
+    Args:
+        conditioners: a dictionary of conditioners with keys corresponding to the keys of the conditioning input dictionary (e.g. "prompt")
+        default_keys: a dictionary of default keys to use if the key is not in the input dictionary (e.g. {"prompt_t5": "prompt"})
+    """
+    def __init__(self, conditioners: tp.Dict[str, Conditioner], default_keys: tp.Dict[str, str] = {}):
+        super().__init__()
+
+        self.conditioners = nn.ModuleDict(conditioners)
+        self.default_keys = default_keys
+
+    def forward(self, batch_metadata: tp.List[tp.Dict[str, tp.Any]], device: tp.Union[torch.device, str]) -> tp.Dict[str, tp.Any]:
+        output = {}
+
+        for key, conditioner in self.conditioners.items():
+            condition_key = key
+
+            conditioner_inputs = []
+
+            for x in batch_metadata:
+
+                if condition_key not in x:
+                    if condition_key in self.default_keys:
+                        condition_key = self.default_keys[condition_key]
+                    else:
+                        raise ValueError(f"Conditioner key {condition_key} not found in batch metadata")
+
+                #Unwrap the condition info if it's a single-element list or tuple, this is to support collation functions that wrap everything in a list
+                if isinstance(x[condition_key], list) or isinstance(x[condition_key], tuple) and len(x[condition_key]) == 1:
+                    conditioner_input = x[condition_key][0]
+                    
+                else:
+                    conditioner_input = x[condition_key]
+
+                conditioner_inputs.append(conditioner_input)
+            
+            output[key] = conditioner(conditioner_inputs, device)
+
+        return output
+    
+def create_multi_conditioner_from_conditioning_config(config: tp.Dict[str, tp.Any]) -> MultiConditioner:
+    """
+    Create a MultiConditioner from a conditioning config dictionary
+
+    Args:
+        config: the conditioning config dictionary
+        device: the device to put the conditioners on
+    """
+    conditioners = {}
+    cond_dim = config["cond_dim"]
+    
+    default_keys = config.get("default_keys", {})
+
+    for conditioner_info in config["configs"]:
+        id = conditioner_info["id"]
+
+        conditioner_type = conditioner_info["type"]
+
+        conditioner_config = {"output_dim": cond_dim}
+        
+        conditioner_config.update(conditioner_info["config"])
+
+        if conditioner_type == "t5":
+            conditioners[id] = T5Conditioner(**conditioner_config)
+        elif conditioner_type == "clap_text":
+            conditioners[id] = CLAPTextConditioner(**conditioner_config)
+        elif conditioner_type == "clap_audio":
+            conditioners[id] = CLAPAudioConditioner(**conditioner_config)
+        elif conditioner_type == "int":
+            conditioners[id] = IntConditioner(**conditioner_config)
+        elif conditioner_type == "number":
+            conditioners[id] = NumberConditioner(**conditioner_config)
+        elif conditioner_type == "phoneme":
+            conditioners[id] = PhonemeConditioner(**conditioner_config)
+        elif conditioner_type == "lut":
+            conditioners[id] = TokenizerLUTConditioner(**conditioner_config)
+        elif conditioner_type == "pretransform":
+            sample_rate = conditioner_config.pop("sample_rate", None)
+            assert sample_rate is not None, "Sample rate must be specified for pretransform conditioners"
+
+            pretransform = create_pretransform_from_config(conditioner_config.pop("pretransform_config"), sample_rate=sample_rate)
+
+            if conditioner_config.get("pretransform_ckpt_path", None) is not None:
+                pretransform.load_state_dict(load_ckpt_state_dict(conditioner_config.pop("pretransform_ckpt_path")))
+
+            conditioners[id] = PretransformConditioner(pretransform, **conditioner_config)
+        else:
+            raise ValueError(f"Unknown conditioner type: {conditioner_type}")
+
+    return MultiConditioner(conditioners, default_keys=default_keys)