Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,12 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Utils/JDC/bst.t7 filter=lfs diff=lfs merge=lfs -text
+Utils/PLBERT/step_1050000.t7 filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/01008270.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/kaede_san.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/riamu_zeroshot_02.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/sample_ref01.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/sample_ref02.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/shiki_fine05.wav filter=lfs diff=lfs merge=lfs -text
+reference_sample_wavs/syuukovoice_200918_3_01.wav filter=lfs diff=lfs merge=lfs -text

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

Configs/config.yml

@@ -0,0 +1,116 @@
+log_dir: "Models/LJSpeech"
+first_stage_path: "first_stage.pth"
+save_freq: 2
+log_interval: 10
+device: "cuda"
+epochs_1st: 200 # number of epochs for first stage training (pre-training)
+epochs_2nd: 100 # number of peochs for second stage training (joint training)
+batch_size: 16
+max_len: 400 # maximum number of frames
+pretrained_model: ""
+second_stage_load_pretrained: true # set to true if the pre-trained model is for 2nd stage
+load_only_params: false # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "Utils/JDC/bst.t7"
+ASR_config: "Utils/ASR/config.yml"
+ASR_path: "Utils/ASR/epoch_00080.pth"
+PLBERT_dir: 'Utils/PLBERT/'
+
+data_params:
+  train_data: "Data/train_list.txt"
+  val_data: "Data/val_list.txt"
+  root_path: "/local/LJSpeech-1.1/wavs"
+  OOD_data: "Data/OOD_texts.txt"
+  min_length: 50 # sample until texts with this size are obtained for OOD texts
+
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+
+model_params:
+  multispeaker: false
+
+  dim_in: 64 
+  hidden_dim: 512
+  max_conv_dim: 512
+  n_layer: 3
+  n_mels: 80
+
+  n_token: 178 # number of phoneme tokens
+  max_dur: 50 # maximum duration of a single phoneme
+  style_dim: 128 # style vector size
+  
+  dropout: 0.2
+
+  # config for decoder
+  decoder: 
+      type: 'istftnet' # either hifigan or istftnet
+      resblock_kernel_sizes: [3,7,11]
+      upsample_rates :  [10, 6]
+      upsample_initial_channel: 512
+      resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]]
+      upsample_kernel_sizes: [20, 12]
+      gen_istft_n_fft: 20
+      gen_istft_hop_size: 5
+      
+  # speech language model config
+  slm:
+      model: 'microsoft/wavlm-base-plus'
+      sr: 16000 # sampling rate of SLM
+      hidden: 768 # hidden size of SLM
+      nlayers: 13 # number of layers of SLM
+      initial_channel: 64 # initial channels of SLM discriminator head
+  
+  # style diffusion model config
+  diffusion:
+    embedding_mask_proba: 0.1
+    # transformer config
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+
+    # diffusion distribution config
+    dist:
+      sigma_data: 0.2 # placeholder for estimate_sigma_data set to false
+      estimate_sigma_data: true # estimate sigma_data from the current batch if set to true
+      mean: -3.0
+      std: 1.0
+  
+loss_params:
+    lambda_mel: 5. # mel reconstruction loss
+    lambda_gen: 1. # generator loss
+    lambda_slm: 1. # slm feature matching loss
+    
+    lambda_mono: 1. # monotonic alignment loss (1st stage, TMA)
+    lambda_s2s: 1. # sequence-to-sequence loss (1st stage, TMA)
+    TMA_epoch: 50 # TMA starting epoch (1st stage)
+
+    lambda_F0: 1. # F0 reconstruction loss (2nd stage)
+    lambda_norm: 1. # norm reconstruction loss (2nd stage)
+    lambda_dur: 1. # duration loss (2nd stage)
+    lambda_ce: 20. # duration predictor probability output CE loss (2nd stage)
+    lambda_sty: 1. # style reconstruction loss (2nd stage)
+    lambda_diff: 1. # score matching loss (2nd stage)
+    
+    diff_epoch: 20 # style diffusion starting epoch (2nd stage)
+    joint_epoch: 50 # joint training starting epoch (2nd stage)
+
+optimizer_params:
+  lr: 0.0001 # general learning rate
+  bert_lr: 0.00001 # learning rate for PLBERT
+  ft_lr: 0.00001 # learning rate for acoustic modules
+  
+slmadv_params:
+  min_len: 400 # minimum length of samples
+  max_len: 500 # maximum length of samples
+  batch_percentage: 0.5 # to prevent out of memory, only use half of the original batch size
+  iter: 10 # update the discriminator every this iterations of generator update
+  thresh: 5 # gradient norm above which the gradient is scaled
+  scale: 0.01 # gradient scaling factor for predictors from SLM discriminators
+  sig: 1.5 # sigma for differentiable duration modeling
+  

Configs/config_ft.yml

@@ -0,0 +1,116 @@
+log_dir: "Models/IMAS_FineTuned"
+save_freq: 1
+log_interval: 10
+device: "cuda"
+epochs: 50 # number of finetuning epoch (1 hour of data)
+batch_size: 3
+max_len: 2500 # maximum number of frames
+pretrained_model: "/home/austin/disk2/llmvcs/tt/stylekan/Models/Style_Kanade/NO_SLM_3_epoch_2nd_00002.pth"
+second_stage_load_pretrained: true # set to true if the pre-trained model is for 2nd stage
+load_only_params: true # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "/home/austin/disk2/llmvcs/tt/stylekan/Utils/JDC/bst.t7"
+ASR_config: "/home/austin/disk2/llmvcs/tt/stylekan/Utils/ASR/config.yml"
+ASR_path: "/home/austin/disk2/llmvcs/tt/stylekan/Utils/ASR/bst_00080.pth"
+
+PLBERT_dir: 'Utils/PLBERT/'
+
+data_params:
+  train_data: "/home/austin/disk2/llmvcs/tt/stylekan/Data/metadata_cleanest/FT_imas.csv"
+  val_data: "/home/austin/disk2/llmvcs/tt/stylekan/Data/metadata_cleanest/FT_imas_valid.csv"
+  root_path: ""
+  OOD_data: "/home/austin/disk2/llmvcs/tt/stylekan/Data/OOD_LargeScale_.csv"
+  min_length: 50 # sample until texts with this size are obtained for OOD texts
+
+
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+
+model_params:
+  multispeaker: true
+
+  dim_in: 64 
+  hidden_dim: 512
+  max_conv_dim: 512
+  n_layer: 3
+  n_mels: 80
+
+  n_token: 178 # number of phoneme tokens
+  max_dur: 50 # maximum duration of a single phoneme
+  style_dim: 128 # style vector size
+  
+  dropout: 0.2
+
+  decoder: 
+    type: 'istftnet' # either hifigan or istftnet
+    resblock_kernel_sizes: [3,7,11]
+    upsample_rates :  [10, 6]
+    upsample_initial_channel: 512
+    resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]]
+    upsample_kernel_sizes: [20, 12]
+    gen_istft_n_fft: 20
+    gen_istft_hop_size: 5
+
+
+      
+  # speech language model config
+  slm:
+      model: 'Respair/Whisper_Large_v2_Encoder_Block' # The model itself is hardcoded, change it through -> losses.py
+      sr: 16000 # sampling rate of SLM
+      hidden: 1280 # hidden size of SLM
+      nlayers: 33 # number of layers of SLM
+      initial_channel: 64 # initial channels of SLM discriminator head
+  
+  # style diffusion model config
+  diffusion:
+    embedding_mask_proba: 0.1
+    # transformer config
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+
+    # diffusion distribution config
+    dist:
+      sigma_data: 0.2 # placeholder for estimate_sigma_data set to false
+      estimate_sigma_data: true # estimate sigma_data from the current batch if set to true
+      mean: -3.0
+      std: 1.0
+  
+loss_params:
+    lambda_mel: 10. # mel reconstruction loss
+    lambda_gen: 1. # generator loss
+    lambda_slm: 1. # slm feature matching loss
+    
+    lambda_mono: 1. # monotonic alignment loss (1st stage, TMA)
+    lambda_s2s: 1. # sequence-to-sequence loss (1st stage, TMA)
+    TMA_epoch: 9 # TMA starting epoch (1st stage)
+
+    lambda_F0: 1. # F0 reconstruction loss (2nd stage)
+    lambda_norm: 1. # norm reconstruction loss (2nd stage)
+    lambda_dur: 1. # duration loss (2nd stage)
+    lambda_ce: 20. # duration predictor probability output CE loss (2nd stage)
+    lambda_sty: 1. # style reconstruction loss (2nd stage)
+    lambda_diff: 1. # score matching loss (2nd stage)
+    
+    diff_epoch: 0 # style diffusion starting epoch (2nd stage)
+    joint_epoch: 30 # joint training starting epoch (2nd stage)
+
+optimizer_params:
+  lr: 0.0001 # general learning rate
+  bert_lr: 0.00001 # learning rate for PLBERT
+  ft_lr: 0.00001 # learning rate for acoustic modules
+  
+slmadv_params:
+  min_len: 400 # minimum length of samples
+  max_len: 500 # maximum length of samples
+  batch_percentage: 0.5 # to prevent out of memory, only use half of the original batch size
+  iter: 20 # update the discriminator every this iterations of generator update
+  thresh: 5 # gradient norm above which the gradient is scaled
+  scale: 0.01 # gradient scaling factor for predictors from SLM discriminators
+  sig: 1.5 # sigma for differentiable duration modeling

Configs/config_kanade.yml

@@ -0,0 +1,118 @@
+log_dir: "Models/Style_Kanade_v02"
+first_stage_path: ""
+save_freq: 1
+log_interval: 10
+device: "cuda"
+epochs_1st: 30 # number of epochs for first stage training (pre-training)
+epochs_2nd: 20 # number of peochs for second stage training (joint training)
+batch_size: 64
+max_len: 560 # maximum number of frames
+pretrained_model: "Models/Style_Kanade_v02/epoch_2nd_00007.pth"
+second_stage_load_pretrained: true # set to true if the pre-trained model is for 2nd stage
+load_only_params: false # set to true if do not want to load epoch numbers and optimizer parameters
+
+F0_path: "Utils/JDC/bst.t7"
+ASR_config: "Utils/ASR/config.yml"
+ASR_path: "Utils/ASR/bst_00080.pth"
+
+PLBERT_dir: 'Utils/PLBERT/'
+
+data_params:
+  train_data: "Data/metadata_cleanest/DATA.csv"
+  val_data: "Data/VALID.txt"
+  root_path: ""
+  OOD_data: "Data/OOD_LargeScale_.csv"
+  min_length: 50 # sample until texts with this size are obtained for OOD texts
+
+
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+
+model_params:
+  multispeaker: true
+
+  dim_in: 64 
+  hidden_dim: 512
+  max_conv_dim: 512
+  n_layer: 3
+  n_mels: 80
+
+  n_token: 178 # number of phoneme tokens
+  max_dur: 50 # maximum duration of a single phoneme
+  style_dim: 128 # style vector size
+  
+  dropout: 0.2
+
+  decoder: 
+    type: 'istftnet' # either hifigan or istftnet
+    resblock_kernel_sizes: [3,7,11]
+    upsample_rates :  [10, 6]
+    upsample_initial_channel: 512
+    resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]]
+    upsample_kernel_sizes: [20, 12]
+    gen_istft_n_fft: 20
+    gen_istft_hop_size: 5
+
+
+      
+  # speech language model config
+  slm:
+      model: 'Respair/Whisper_Large_v2_Encoder_Block' # The model itself is hardcoded, change it through -> losses.py
+      sr: 16000 # sampling rate of SLM
+      hidden: 1280 # hidden size of SLM
+      nlayers: 33 # number of layers of SLM
+      initial_channel: 64 # initial channels of SLM discriminator head
+  
+  # style diffusion model config
+  diffusion:
+    embedding_mask_proba: 0.1
+    # transformer config
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+
+    # diffusion distribution config
+    dist:
+      sigma_data: 0.2 # placeholder for estimate_sigma_data set to false
+      estimate_sigma_data: true # estimate sigma_data from the current batch if set to true
+      mean: -3.0
+      std: 1.0
+  
+loss_params:
+    lambda_mel: 10. # mel reconstruction loss
+    lambda_gen: 1. # generator loss
+    lambda_slm: 1. # slm feature matching loss
+    
+    lambda_mono: 1. # monotonic alignment loss (1st stage, TMA)
+    lambda_s2s: 1. # sequence-to-sequence loss (1st stage, TMA)
+    TMA_epoch: 5 # TMA starting epoch (1st stage)
+
+    lambda_F0: 1. # F0 reconstruction loss (2nd stage)
+    lambda_norm: 1. # norm reconstruction loss (2nd stage)
+    lambda_dur: 1. # duration loss (2nd stage)
+    lambda_ce: 20. # duration predictor probability output CE loss (2nd stage)
+    lambda_sty: 1. # style reconstruction loss (2nd stage)
+    lambda_diff: 1. # score matching loss (2nd stage)
+    
+    diff_epoch: 4 # style diffusion starting epoch (2nd stage)
+    joint_epoch: 999 # joint training starting epoch (2nd stage)
+
+optimizer_params:
+  lr: 0.0001 # general learning rate
+  bert_lr: 0.00001 # learning rate for PLBERT
+  ft_lr: 0.00001 # learning rate for acoustic modules
+  
+slmadv_params:
+  min_len: 400 # minimum length of samples
+  max_len: 500 # maximum length of samples
+  batch_percentage: 0.5 # to prevent out of memory, only use half of the original batch size
+  iter: 20 # update the discriminator every this iterations of generator update
+  thresh: 5 # gradient norm above which the gradient is scaled
+  scale: 0.01 # gradient scaling factor for predictors from SLM discriminators
+  sig: 1.5 # sigma for differentiable duration modeling

Inference/input_for_prompt.txt

@@ -0,0 +1,4 @@
+この俺に何度も同じことを説明させるな！! お前たちは俺の忠告を完全に無視して、とんでもない結果を招いてしまった。これが最後の警告だ。次は絶対に許さないぞ！
+時には、静けさの中にこそ、本当の答えが見つかるものですね。慌てる必要はないのです。
+人生には、表現しきれないほどの驚きがあると思うよ。それは、目には見えない力で、人々を繋ぐ不思議な絆だ。私は、その驚きを胸に秘め、日々を楽しく過ごしているんだ。言葉を伝えるたびに、未来への期待を込めて、元気に話す。それは、夢を叶えるための魔法のようなものだ。
+かなたの次元より迫り来る混沌の使者たちよ、貴様らの野望を我が焔の業火で焼き尽くす！ 運命の歯車は、我が意思と共にすでに動き出したのだ。我が宿命の敵に立ち向かうため、禁断の呪文を紡ぐ時は今ここに訪れる。さあ、見るがよい。我が力を！！

Inference/prompt.txt

@@ -0,0 +1,4 @@
+A male voice that resonates with deep, thunderous intensity. His rapid-fire words slam like aggressive drumbeats, each syllable charged with intense rage. The expressive tone fluctuates between restrained fury and explosive outbursts.
+A female voice with a distinctively deep, low pitch that commands attention. Her slightly monotone delivery creates an air of composure and gravitas, while maintaining a calm, measured pace. Her voice carries a soothing weight, like gentle thunder in the distance, making her words feel grounded and reassuring.
+a female voice that is gentle and soft, with a slightly high pitch that adds a comforting warmth to her words. Her tone is moderate, neither too expressive nor too flat, creating a balanced and soothing atmosphere. The slow speed of her speech gives her words a deliberate and thoughtful cadence, allowing each phrase to resonate fully. There's a sense of wonder and optimism in her voice, as if she is constantly marveling at the mysteries of life. Her gentle demeanor and soft delivery make her sound approachable and kind, inviting listeners to share in her sense of wonder.
+A female voice that resonates with deep intensity and a distinctly low pitch. Her words flow with the force and rhythm of a relentless tide, each syllable weighted with profound determination. The expressive tone navigates between measured intensity and powerful surges.

Inference/random_texts.txt

@@ -0,0 +1,14 @@
+Akashi: 不思議な人ですね、レザさんは。たまには子供扱いしてくれてちょっとむきになりますけど、とても頼りがいのある人だと思いますよ？
+Kimiji: 人生は、果てしない探求の旅のようなもの。私たちは、自分自身や周囲の世界について、常に新しい発見をしていく。それは、時として喜びをもたらすこともあれば、困難に直面することもある。しかしそれら全てが、自分を形作る貴重な経験である。
+Reira: 私に何度も同じことを説明させないでよ! お前たちは私の忠告を完全に無視して、とんでもない結果を招いてしまった!! これが最後の警告だ。次は絶対に許さないぞ！ ------------------------------------------- (NOTE: enable the diffusion, then set the Intensity to 3, also remove this line!)
+Yoriko: この世には、言葉にできない悲しみがある。 それは、胸の奥に沈んでいくような重さで、時間が経つにつれて、じわじわと広がっていく。私は、その悲しみを抱えながら、日々を過ごしていいた。 言葉を発するたびに、心の中で何度も繰り返し、慎重に選び抜いている。それは、痛みを和らげるための儀式のようなものだ.
+Kimiji: 人生には、表現しきれないほどの驚きがあると思います。それは、目には見えない力で、人々を繋ぐ不思議な絆です。私は、その驚きを胸に秘め、日々を楽しく過ごしています。言葉を伝えるたびに、未来への期待を込めて、元気に話す。それは、夢を叶えるための魔法のようなものです。
+Teppei: そうだな、この新しいシステムの仕組みについて説明しておこう。基本的には三層構造になっていて、各層が独立して機能している。一番下のレイヤーでデータの処理と保存を行い、真ん中の層でビジネスロジックを実装している。ユーザーが直接触れるのは最上位層だけで、そこでインターフェースの制御をしているんだ。パフォーマンスを考えると、データベースへのアクセスを最小限に抑える必要があるから、キャッシュの実装も検討している。ただ、システムの規模を考えると、当面は現状の構成で十分だと思われる。将来的にスケールする必要が出てきた場合は、その時点で見直しを検討すればいいだろう。
+Kirara: べ、別にそんなことじゃないってば！あんたのことなんて全然気にしてないんだからね！
+Maiko: ねえ、ちょっと今日の空を見上げて。朝から少しずつ変わっていく雲の形が、まるで漫画の中の風景みたい。東の方からゆっくりと暖かい風が吹いてきて、桜の花びらが舞い散るように、優しく大地を撫でていくの。春の陽気が徐々に夏の暑さに変わっていくこの季節は、なんだかわくわくするよね。でも、週間予報によると、明後日あたりから天気が崩れるみたいで、しばらくは傘の出番かもしれないんだ。梅雨の時期が近づいているから、空気も少しずつ湿っぽくなってきているのを感じない？でも、雨上がりの空気って、なんだか特別な匂いがして、私、結構好きなんだよね。
+Amane: そ、そうかな？あなたと過ごした今までの時間は、私に、とても、とーっても大事だよ？だって、あなたの言葉には、どこか特別な温もりがあるような気がする。まるで、心の中に直接響いてくるようで、ほんの少しの間だけ、世界が優しくなった気がするよ。
+Shioji: つい昨日のことみたいなのに、もう随分と昔のお話なんだね ! 今でも古い本の匂いがすると、あの静かで穏やかな時間がまるで透明で綺麗な海みたいに心の中でふわりと溶けていくの。
+Riemi: かなたの次元より迫り来る混沌の使者たちよ、貴様らの野望を我が焔の業火で焼き尽くす！ 運命の歯車は、我が意思と共にすでに動き出したのだ。我が宿命の敵に立ち向かうため、禁断の呪文を紡ぐ時は今ここに訪れる。さあ、見るがよい。我が力を！！
+だめですよ？そんなことは。もっと慎重に行動してくださいね。
+Kazunari: ojousama wa, katachi dake no tsukibito ga ireba sore de ii, to omotte orareru you desuga, katachi bakari de mo, sannen ka o tomo ni seikatsu o suru koto ni narimasɯ. --------------------------------------- (NOTE!: your spacing will impact the intonations. also remove this line!)
+Kimiu: お前 wa muzukashiku 考えると kekka to shite 空回り suru taipu dakara na. 大体, sonna kanji de oboete ikeba iin da. -------------------------------------------------- (NOTE!: your spacing will impact the intonations. also remove this line!)

LICENSE

@@ -0,0 +1,25 @@
+LICENSE FOR STYLETTS2 DEMO PAGE: GPL (I KNOW I DON'T LIKE GPL BUT I HAVE TO BECAUSE OF PHONEMIZER REQUIREMENT)
+
+
+styletts 2 license:
+
+Copyright (c) 2023 Aaron (Yinghao) Li
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

Models/Style_Tsukasa_v02/Top_ckpt_24khz.pth

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

Models/Style_Tsukasa_v02/config_kanade.yml

@@ -0,0 +1,121 @@
+log_dir: "/home/austin/disk2/llmvcs/tt/stylekan/Models/Style_Kanade_v02"
+first_stage_path: "/home/austin/disk2/llmvcs/tt/stylekan/Models/Style_Kanade_v02/epoch_1st_00026.pth"
+save_freq: 1
+log_interval: 10
+device: "cuda"
+epochs_1st: 30 # number of epochs for first stage training (pre-training)
+epochs_2nd: 20 # number of peochs for second stage training (joint training)
+batch_size: 64
+max_len: 4000 # maximum number of frames
+pretrained_model: "/home/austin/disk2/llmvcs/tt/stylekan/Models/Style_Kanade_v02/epoch_2nd_00007.pth"
+second_stage_load_pretrained: true # set to true if the pre-trained model is for 2nd stage
+load_only_params: false # set to true if do not want to load epoch numbers and optimizer parameters
+
+# CUDA_VISIBLE_DEVICES=1,2,3 accelerate launch train_first.py --config_path ./Configs/config_kanade.yml
+# CUDA_VISIBLE_DEVICES=6,7 accelerate launch accelerate_train_second.py --config_path ./Configs/config_kanade_test.yml
+
+F0_path: "/home/austin/disk2/llmvcs/tt/stylekan/Utils/JDC/bst.t7"
+ASR_config: "/home/austin/disk2/llmvcs/tt/stylekan/Utils/ASR/config.yml"
+ASR_path: "/home/austin/disk2/llmvcs/tt/stylekan/Utils/ASR/bst_00080.pth"
+
+PLBERT_dir: 'Utils/PLBERT/'
+
+data_params:
+  train_data: "/home/austin/disk2/llmvcs/tt/stylekan/Data/metadata_cleanest/filtered_train_list_no_nsp_plus.csv"
+  val_data: "/home/austin/disk2/llmvcs/tt/stylekan/Data/mg_valid.txt"
+  root_path: ""
+  OOD_data: "/home/austin/disk2/llmvcs/tt/stylekan/Data/OOD_LargeScale_.csv"
+  min_length: 50 # sample until texts with this size are obtained for OOD texts
+
+
+preprocess_params:
+  sr: 24000
+  spect_params:
+    n_fft: 2048
+    win_length: 1200
+    hop_length: 300
+
+model_params:
+  multispeaker: true
+
+  dim_in: 64 
+  hidden_dim: 512
+  max_conv_dim: 512
+  n_layer: 3
+  n_mels: 80
+
+  n_token: 178 # number of phoneme tokens
+  max_dur: 50 # maximum duration of a single phoneme
+  style_dim: 128 # style vector size
+  
+  dropout: 0.2
+
+  decoder: 
+    type: 'istftnet' # either hifigan or istftnet
+    resblock_kernel_sizes: [3,7,11]
+    upsample_rates :  [10, 6]
+    upsample_initial_channel: 512
+    resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]]
+    upsample_kernel_sizes: [20, 12]
+    gen_istft_n_fft: 20
+    gen_istft_hop_size: 5
+
+
+      
+  # speech language model config
+  slm:
+      model: 'Respair/Whisper_Large_v2_Encoder_Block' # The model itself is hardcoded, change it through -> losses.py
+      sr: 16000 # sampling rate of SLM
+      hidden: 1280 # hidden size of SLM
+      nlayers: 33 # number of layers of SLM
+      initial_channel: 64 # initial channels of SLM discriminator head
+  
+  # style diffusion model config
+  diffusion:
+    embedding_mask_proba: 0.1
+    # transformer config
+    transformer:
+      num_layers: 3
+      num_heads: 8
+      head_features: 64
+      multiplier: 2
+
+    # diffusion distribution config
+    dist:
+      sigma_data: 0.2 # placeholder for estimate_sigma_data set to false
+      estimate_sigma_data: true # estimate sigma_data from the current batch if set to true
+      mean: -3.0
+      std: 1.0
+  
+loss_params:
+    lambda_mel: 10. # mel reconstruction loss
+    lambda_gen: 1. # generator loss
+    lambda_slm: 1. # slm feature matching loss
+    
+    lambda_mono: 1. # monotonic alignment loss (1st stage, TMA)
+    lambda_s2s: 1. # sequence-to-sequence loss (1st stage, TMA)
+    TMA_epoch: 5 # TMA starting epoch (1st stage)
+
+    lambda_F0: 1. # F0 reconstruction loss (2nd stage)
+    lambda_norm: 1. # norm reconstruction loss (2nd stage)
+    lambda_dur: 1. # duration loss (2nd stage)
+    lambda_ce: 20. # duration predictor probability output CE loss (2nd stage)
+    lambda_sty: 1. # style reconstruction loss (2nd stage)
+    lambda_diff: 1. # score matching loss (2nd stage)
+    
+    diff_epoch: 4 # style diffusion starting epoch (2nd stage)
+    joint_epoch: 999 # joint training starting epoch (2nd stage)
+
+optimizer_params:
+  lr: 0.0001 # general learning rate
+  bert_lr: 0.00001 # learning rate for PLBERT
+  ft_lr: 0.00001 # learning rate for acoustic modules
+  
+slmadv_params:
+  min_len: 400 # minimum length of samples
+  max_len: 500 # maximum length of samples
+  batch_percentage: 0.5 # to prevent out of memory, only use half of the original batch size
+  iter: 20 # update the discriminator every this iterations of generator update
+  thresh: 5 # gradient norm above which the gradient is scaled
+  scale: 0.01 # gradient scaling factor for predictors from SLM discriminators
+  sig: 1.5 # sigma for differentiable duration modeling

Modules/KotoDama_sampler.py

@@ -0,0 +1,269 @@
+from transformers import AutoModelForSequenceClassification, PreTrainedModel, AutoConfig, AutoModel, AutoTokenizer
+import torch
+import torch.nn as nn
+from text_utils import TextCleaner
+textclenaer = TextCleaner()
+
+
+def length_to_mask(lengths):
+    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
+    mask = torch.gt(mask+1, lengths.unsqueeze(1))
+    return mask
+
+
+
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+# tokenizer_koto_prompt = AutoTokenizer.from_pretrained("google/mt5-small", trust_remote_code=True)
+tokenizer_koto_prompt = AutoTokenizer.from_pretrained("ku-nlp/deberta-v3-base-japanese", trust_remote_code=True)
+tokenizer_koto_text = AutoTokenizer.from_pretrained("line-corporation/line-distilbert-base-japanese", trust_remote_code=True)
+
+class KotoDama_Prompt(PreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        self.backbone = AutoModel.from_config(config)
+
+        self.output = nn.Sequential(nn.Linear(config.hidden_size, 512),
+                                    nn.LeakyReLU(0.2), 
+                                    nn.Linear(512, config.num_labels))
+
+                                    
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        labels=None,
+    ):
+        outputs = self.backbone(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+        )
+
+
+        sequence_output = outputs.last_hidden_state[:, 0, :]
+        outputs = self.output(sequence_output)
+
+        # if labels, then we are training
+        loss = None
+        if labels is not None:
+
+            loss_fn = nn.MSELoss()
+            # labels = labels.unsqueeze(1)
+            loss = loss_fn(outputs, labels)
+            
+        return {
+            "loss": loss,
+            "logits": outputs
+        }
+
+
+class KotoDama_Text(PreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        self.backbone = AutoModel.from_config(config)
+
+        self.output = nn.Sequential(nn.Linear(config.hidden_size, 512),
+                                    nn.LeakyReLU(0.2), 
+                                    nn.Linear(512, config.num_labels))
+
+                                    
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        # token_type_ids=None,
+        # position_ids=None,
+        labels=None,
+    ):
+        outputs = self.backbone(
+            input_ids,
+            attention_mask=attention_mask,
+            # token_type_ids=token_type_ids,
+            # position_ids=position_ids,
+        )
+
+
+        sequence_output = outputs.last_hidden_state[:, 0, :]
+        outputs = self.output(sequence_output)
+
+        # if labels, then we are training
+        loss = None
+        if labels is not None:
+
+            loss_fn = nn.MSELoss()
+            # labels = labels.unsqueeze(1)
+            loss = loss_fn(outputs, labels)
+            
+        return {
+            "loss": loss,
+            "logits": outputs
+        }
+
+
+def inference(model, diffusion_sampler, text=None, ref_s=None, alpha = 0.3, beta = 0.7, diffusion_steps=5, embedding_scale=1, rate_of_speech=1.):
+
+    tokens = textclenaer(text)
+    tokens.insert(0, 0)
+    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+    
+    with torch.no_grad():
+        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
+
+        text_mask = length_to_mask(input_lengths).to(device)
+
+        t_en = model.text_encoder(tokens, input_lengths, text_mask)
+        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
+        d_en = model.bert_encoder(bert_dur).transpose(-1, -2) 
+        
+
+
+        s_pred = diffusion_sampler(noise = torch.randn((1, 256)).unsqueeze(1).to(device), 
+                                          embedding=bert_dur,
+                                          embedding_scale=embedding_scale,
+                                            features=ref_s, # reference from the same speaker as the embedding
+                                             num_steps=diffusion_steps).squeeze(1)
+
+
+        s = s_pred[:, 128:]
+        ref = s_pred[:, :128]
+
+        ref = alpha * ref + (1 - alpha)  * ref_s[:, :128]
+        s = beta * s + (1 - beta)  * ref_s[:, 128:]
+
+        d = model.predictor.text_encoder(d_en, 
+                                         s, input_lengths, text_mask)
+        
+  
+
+        x = model.predictor.lstm(d)
+        x_mod =  model.predictor.prepare_projection(x) 
+        duration = model.predictor.duration_proj(x_mod) 
+
+
+        duration = torch.sigmoid(duration).sum(axis=-1) / rate_of_speech
+        
+        pred_dur = torch.round(duration.squeeze()).clamp(min=1) 
+
+
+
+        pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
+        
+        c_frame = 0
+        for i in range(pred_aln_trg.size(0)):
+            pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1
+            c_frame += int(pred_dur[i].data)
+
+        # encode prosody
+        en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))
+
+
+
+        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+
+        asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
+
+
+        out = model.decoder(asr, 
+                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))
+    
+        
+    return out.squeeze().cpu().numpy()[..., :-50] 
+
+
+def Longform(model, diffusion_sampler, text, s_prev, ref_s, alpha = 0.3, beta = 0.7, t = 0.7, diffusion_steps=5, embedding_scale=1, rate_of_speech=1.0):
+
+    tokens = textclenaer(text)
+    tokens.insert(0, 0)
+    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
+    
+    with torch.no_grad():
+        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
+        text_mask = length_to_mask(input_lengths).to(device)
+
+        t_en = model.text_encoder(tokens, input_lengths, text_mask)
+        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
+        d_en = model.bert_encoder(bert_dur).transpose(-1, -2) 
+
+        s_pred = diffusion_sampler(noise = torch.randn((1, 256)).unsqueeze(1).to(device), 
+                                          embedding=bert_dur,
+                                          embedding_scale=embedding_scale,
+                                            features=ref_s, 
+                                             num_steps=diffusion_steps).squeeze(1)
+        
+        if s_prev is not None:
+            # convex combination of previous and current style
+            s_pred = t * s_prev + (1 - t) * s_pred
+        
+        s = s_pred[:, 128:]
+        ref = s_pred[:, :128]
+        
+        ref = alpha * ref + (1 - alpha)  * ref_s[:, :128]
+        s = beta * s + (1 - beta)  * ref_s[:, 128:]
+
+        s_pred = torch.cat([ref, s], dim=-1)
+
+        d = model.predictor.text_encoder(d_en, 
+                                         s, input_lengths, text_mask)
+
+        x = model.predictor.lstm(d)
+        x_mod =  model.predictor.prepare_projection(x) # 640 -> 512
+        duration = model.predictor.duration_proj(x_mod)
+
+        duration = torch.sigmoid(duration).sum(axis=-1) / rate_of_speech
+        pred_dur = torch.round(duration.squeeze()).clamp(min=1) 
+
+
+        pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
+        c_frame = 0
+        for i in range(pred_aln_trg.size(0)):
+            pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1
+            c_frame += int(pred_dur[i].data)
+
+        # encode prosody
+        en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))
+
+        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+
+        asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))
+
+        out = model.decoder(asr, 
+                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))
+    
+        
+    return out.squeeze().cpu().numpy()[..., :-100], s_pred 
+
+
+def merge_short_elements(lst):
+    i = 0
+    while i < len(lst):
+        if i > 0 and len(lst[i]) < 10:
+            lst[i-1] += ' ' + lst[i]
+            lst.pop(i)
+        else:
+            i += 1
+    return lst
+
+
+def merge_three(text_list, maxim=2):
+
+    merged_list = []
+    for i in range(0, len(text_list), maxim):
+        merged_text = ' '.join(text_list[i:i+maxim])
+        merged_list.append(merged_text)
+    return merged_list
+
+
+def merging_sentences(lst):
+    return merge_three(merge_short_elements(lst))

Modules/__init__.py

@@ -0,0 +1 @@
+

Modules/diffusion/__init__.py

@@ -0,0 +1 @@
+

Modules/diffusion/audio_diffusion_pytorch/__init__.py

@@ -0,0 +1,20 @@
+from .components import LTPlugin, MelSpectrogram, UNetV0, XUNet
+from .diffusion import (
+    Diffusion,
+    Distribution,
+    LinearSchedule,
+    Sampler,
+    Schedule,
+    UniformDistribution,
+    VDiffusion,
+    VInpainter,
+    VSampler,
+)
+from .models import (
+    DiffusionAE,
+    DiffusionAR,
+    DiffusionModel,
+    DiffusionUpsampler,
+    DiffusionVocoder,
+    EncoderBase,
+)

Modules/diffusion/audio_diffusion_pytorch/components.py

@@ -0,0 +1,236 @@
+from typing import Callable, Optional, Sequence
+
+import torch
+import torch.nn.functional as F
+from a_unet import (
+    ClassifierFreeGuidancePlugin,
+    Conv,
+    Module,
+    TextConditioningPlugin,
+    TimeConditioningPlugin,
+    default,
+    exists,
+)
+from a_unet.apex import (
+    AttentionItem,
+    CrossAttentionItem,
+    InjectChannelsItem,
+    ModulationItem,
+    ResnetItem,
+    SkipCat,
+    SkipModulate,
+    XBlock,
+    XUNet,
+)
+from einops import pack, unpack
+from torch import Tensor, nn
+from torchaudio import transforms
+
+"""
+UNets (built with a-unet: https://github.com/archinetai/a-unet)
+"""
+
+
+def UNetV0(
+    dim: int,
+    in_channels: int,
+    channels: Sequence[int],
+    factors: Sequence[int],
+    items: Sequence[int],
+    attentions: Optional[Sequence[int]] = None,
+    cross_attentions: Optional[Sequence[int]] = None,
+    context_channels: Optional[Sequence[int]] = None,
+    attention_features: Optional[int] = None,
+    attention_heads: Optional[int] = None,
+    embedding_features: Optional[int] = None,
+    resnet_groups: int = 8,
+    use_modulation: bool = True,
+    modulation_features: int = 1024,
+    embedding_max_length: Optional[int] = None,
+    use_time_conditioning: bool = True,
+    use_embedding_cfg: bool = False,
+    use_text_conditioning: bool = False,
+    out_channels: Optional[int] = None,
+):
+    # Set defaults and check lengths
+    num_layers = len(channels)
+    attentions = default(attentions, [0] * num_layers)
+    cross_attentions = default(cross_attentions, [0] * num_layers)
+    context_channels = default(context_channels, [0] * num_layers)
+    xs = (channels, factors, items, attentions, cross_attentions, context_channels)
+    assert all(len(x) == num_layers for x in xs)  # type: ignore
+
+    # Define UNet type
+    UNetV0 = XUNet
+
+    if use_embedding_cfg:
+        msg = "use_embedding_cfg requires embedding_max_length"
+        assert exists(embedding_max_length), msg
+        UNetV0 = ClassifierFreeGuidancePlugin(UNetV0, embedding_max_length)
+
+    if use_text_conditioning:
+        UNetV0 = TextConditioningPlugin(UNetV0)
+
+    if use_time_conditioning:
+        assert use_modulation, "use_time_conditioning requires use_modulation=True"
+        UNetV0 = TimeConditioningPlugin(UNetV0)
+
+    # Build
+    return UNetV0(
+        dim=dim,
+        in_channels=in_channels,
+        out_channels=out_channels,
+        blocks=[
+            XBlock(
+                channels=channels,
+                factor=factor,
+                context_channels=ctx_channels,
+                items=(
+                    [ResnetItem]
+                    + [ModulationItem] * use_modulation
+                    + [InjectChannelsItem] * (ctx_channels > 0)
+                    + [AttentionItem] * att
+                    + [CrossAttentionItem] * cross
+                )
+                * items,
+            )
+            for channels, factor, items, att, cross, ctx_channels in zip(*xs)  # type: ignore # noqa
+        ],
+        skip_t=SkipModulate if use_modulation else SkipCat,
+        attention_features=attention_features,
+        attention_heads=attention_heads,
+        embedding_features=embedding_features,
+        modulation_features=modulation_features,
+        resnet_groups=resnet_groups,
+    )
+
+
+"""
+Plugins
+"""
+
+
+def LTPlugin(
+    net_t: Callable, num_filters: int, window_length: int, stride: int
+) -> Callable[..., nn.Module]:
+    """Learned Transform Plugin"""
+
+    def Net(
+        dim: int, in_channels: int, out_channels: Optional[int] = None, **kwargs
+    ) -> nn.Module:
+        out_channels = default(out_channels, in_channels)
+        in_channel_transform = in_channels * num_filters
+        out_channel_transform = out_channels * num_filters  # type: ignore
+
+        padding = window_length // 2 - stride // 2
+        encode = Conv(
+            dim=dim,
+            in_channels=in_channels,
+            out_channels=in_channel_transform,
+            kernel_size=window_length,
+            stride=stride,
+            padding=padding,
+            padding_mode="reflect",
+            bias=False,
+        )
+        decode = nn.ConvTranspose1d(
+            in_channels=out_channel_transform,
+            out_channels=out_channels,  # type: ignore
+            kernel_size=window_length,
+            stride=stride,
+            padding=padding,
+            bias=False,
+        )
+        net = net_t(  # type: ignore
+            dim=dim,
+            in_channels=in_channel_transform,
+            out_channels=out_channel_transform,
+            **kwargs
+        )
+
+        def forward(x: Tensor, *args, **kwargs):
+            x = encode(x)
+            x = net(x, *args, **kwargs)
+            x = decode(x)
+            return x
+
+        return Module([encode, decode, net], forward)
+
+    return Net
+
+
+def AppendChannelsPlugin(
+    net_t: Callable,
+    channels: int,
+):
+    def Net(
+        in_channels: int, out_channels: Optional[int] = None, **kwargs
+    ) -> nn.Module:
+        out_channels = default(out_channels, in_channels)
+        net = net_t(  # type: ignore
+            in_channels=in_channels + channels, out_channels=out_channels, **kwargs
+        )
+
+        def forward(x: Tensor, *args, append_channels: Tensor, **kwargs):
+            x = torch.cat([x, append_channels], dim=1)
+            return net(x, *args, **kwargs)
+
+        return Module([net], forward)
+
+    return Net
+
+
+"""
+Other
+"""
+
+
+class MelSpectrogram(nn.Module):
+    def __init__(
+        self,
+        n_fft: int,
+        hop_length: int,
+        win_length: int,
+        sample_rate: int,
+        n_mel_channels: int,
+        center: bool = False,
+        normalize: bool = False,
+        normalize_log: bool = False,
+    ):
+        super().__init__()
+        self.padding = (n_fft - hop_length) // 2
+        self.normalize = normalize
+        self.normalize_log = normalize_log
+        self.hop_length = hop_length
+
+        self.to_spectrogram = transforms.Spectrogram(
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            center=center,
+            power=None,
+        )
+
+        self.to_mel_scale = transforms.MelScale(
+            n_mels=n_mel_channels, n_stft=n_fft // 2 + 1, sample_rate=sample_rate
+        )
+
+    def forward(self, waveform: Tensor) -> Tensor:
+        # Pack non-time dimension
+        waveform, ps = pack([waveform], "* t")
+        # Pad waveform
+        waveform = F.pad(waveform, [self.padding] * 2, mode="reflect")
+        # Compute STFT
+        spectrogram = self.to_spectrogram(waveform)
+        # Compute magnitude
+        spectrogram = torch.abs(spectrogram)
+        # Convert to mel scale
+        mel_spectrogram = self.to_mel_scale(spectrogram)
+        # Normalize
+        if self.normalize:
+            mel_spectrogram = mel_spectrogram / torch.max(mel_spectrogram)
+            mel_spectrogram = 2 * torch.pow(mel_spectrogram, 0.25) - 1
+        if self.normalize_log:
+            mel_spectrogram = torch.log(torch.clamp(mel_spectrogram, min=1e-5))
+        # Unpack non-spectrogram dimension
+        return unpack(mel_spectrogram, ps, "* f l")[0]

Modules/diffusion/audio_diffusion_pytorch/diffusion.py

@@ -0,0 +1,354 @@
+from math import pi
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import Tensor
+from tqdm import tqdm
+
+from .utils import default
+
+""" Distributions """
+
+
+class Distribution:
+    """Interface used by different distributions"""
+
+    def __call__(self, num_samples: int, device: torch.device):
+        raise NotImplementedError()
+
+
+class UniformDistribution(Distribution):
+    def __init__(self, vmin: float = 0.0, vmax: float = 1.0):
+        super().__init__()
+        self.vmin, self.vmax = vmin, vmax
+
+    def __call__(self, num_samples: int, device: torch.device = torch.device("cpu")):
+        vmax, vmin = self.vmax, self.vmin
+        return (vmax - vmin) * torch.rand(num_samples, device=device) + vmin
+
+
+""" Diffusion Methods """
+
+
+def pad_dims(x: Tensor, ndim: int) -> Tensor:
+    # Pads additional ndims to the right of the tensor
+    return x.view(*x.shape, *((1,) * ndim))
+
+
+def clip(x: Tensor, dynamic_threshold: float = 0.0):
+    if dynamic_threshold == 0.0:
+        return x.clamp(-1.0, 1.0)
+    else:
+        # Dynamic thresholding
+        # Find dynamic threshold quantile for each batch
+        x_flat = rearrange(x, "b ... -> b (...)")
+        scale = torch.quantile(x_flat.abs(), dynamic_threshold, dim=-1)
+        # Clamp to a min of 1.0
+        scale.clamp_(min=1.0)
+        # Clamp all values and scale
+        scale = pad_dims(scale, ndim=x.ndim - scale.ndim)
+        x = x.clamp(-scale, scale) / scale
+        return x
+
+
+def extend_dim(x: Tensor, dim: int):
+    # e.g. if dim = 4: shape [b] => [b, 1, 1, 1],
+    return x.view(*x.shape + (1,) * (dim - x.ndim))
+
+
+class Diffusion(nn.Module):
+    """Interface used by different diffusion methods"""
+
+    pass
+
+
+class VDiffusion(Diffusion):
+    def __init__(
+        self, net: nn.Module, sigma_distribution: Distribution = UniformDistribution(), loss_fn: Any = F.mse_loss
+    ):
+        super().__init__()
+        self.net = net
+        self.sigma_distribution = sigma_distribution
+        self.loss_fn = loss_fn
+
+    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
+        angle = sigmas * pi / 2
+        alpha, beta = torch.cos(angle), torch.sin(angle)
+        return alpha, beta
+
+    def forward(self, x: Tensor, **kwargs) -> Tensor:  # type: ignore
+        batch_size, device = x.shape[0], x.device
+        # Sample amount of noise to add for each batch element
+        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
+        sigmas_batch = extend_dim(sigmas, dim=x.ndim)
+        # Get noise
+        noise = torch.randn_like(x)
+        # Combine input and noise weighted by half-circle
+        alphas, betas = self.get_alpha_beta(sigmas_batch)
+        x_noisy = alphas * x + betas * noise
+        v_target = alphas * noise - betas * x
+        # Predict velocity and return loss
+        v_pred = self.net(x_noisy, sigmas, **kwargs)
+        return self.loss_fn(v_pred, v_target)
+
+
+class ARVDiffusion(Diffusion):
+    def __init__(self, net: nn.Module, length: int, num_splits: int, loss_fn: Any = F.mse_loss):
+        super().__init__()
+        assert length % num_splits == 0, "length must be divisible by num_splits"
+        self.net = net
+        self.length = length
+        self.num_splits = num_splits
+        self.split_length = length // num_splits
+        self.loss_fn = loss_fn
+
+    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
+        angle = sigmas * pi / 2
+        alpha, beta = torch.cos(angle), torch.sin(angle)
+        return alpha, beta
+
+    def forward(self, x: Tensor, **kwargs) -> Tensor:
+        """Returns diffusion loss of v-objective with different noises per split"""
+        b, _, t, device, dtype = *x.shape, x.device, x.dtype
+        assert t == self.length, "input length must match length"
+        # Sample amount of noise to add for each split
+        sigmas = torch.rand((b, 1, self.num_splits), device=device, dtype=dtype)
+        sigmas = repeat(sigmas, "b 1 n -> b 1 (n l)", l=self.split_length)
+        # Get noise
+        noise = torch.randn_like(x)
+        # Combine input and noise weighted by half-circle
+        alphas, betas = self.get_alpha_beta(sigmas)
+        x_noisy = alphas * x + betas * noise
+        v_target = alphas * noise - betas * x
+        # Sigmas will be provided as additional channel
+        channels = torch.cat([x_noisy, sigmas], dim=1)
+        # Predict velocity and return loss
+        v_pred = self.net(channels, **kwargs)
+        return self.loss_fn(v_pred, v_target)
+
+""" Schedules """
+
+
+class Schedule(nn.Module):
+    """Interface used by different sampling schedules"""
+
+    def forward(self, num_steps: int, device: torch.device) -> Tensor:
+        raise NotImplementedError()
+
+
+class LinearSchedule(Schedule):
+    def __init__(self, start: float = 1.0, end: float = 0.0):
+        super().__init__()
+        self.start, self.end = start, end
+
+    def forward(self, num_steps: int, device: Any) -> Tensor:
+        return torch.linspace(self.start, self.end, num_steps, device=device)
+
+
+""" Samplers """
+
+
+class Sampler(nn.Module):
+    pass
+
+
+class VSampler(Sampler):
+
+    diffusion_types = [VDiffusion]
+
+    def __init__(self, net: nn.Module, schedule: Schedule = LinearSchedule()):
+        super().__init__()
+        self.net = net
+        self.schedule = schedule
+
+    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
+        angle = sigmas * pi / 2
+        alpha, beta = torch.cos(angle), torch.sin(angle)
+        return alpha, beta
+
+    @torch.no_grad()
+    def forward(  # type: ignore
+        self, x_noisy: Tensor, num_steps: int, show_progress: bool = False, **kwargs
+    ) -> Tensor:
+        b = x_noisy.shape[0]
+        sigmas = self.schedule(num_steps + 1, device=x_noisy.device)
+        sigmas = repeat(sigmas, "i -> i b", b=b)
+        sigmas_batch = extend_dim(sigmas, dim=x_noisy.ndim + 1)
+        alphas, betas = self.get_alpha_beta(sigmas_batch)
+        progress_bar = tqdm(range(num_steps), disable=not show_progress)
+
+        for i in progress_bar:
+            v_pred = self.net(x_noisy, sigmas[i], **kwargs)
+            x_pred = alphas[i] * x_noisy - betas[i] * v_pred
+            noise_pred = betas[i] * x_noisy + alphas[i] * v_pred
+            x_noisy = alphas[i + 1] * x_pred + betas[i + 1] * noise_pred
+            progress_bar.set_description(f"Sampling (noise={sigmas[i+1,0]:.2f})")
+
+        return x_noisy
+
+
+class ARVSampler(Sampler):
+    def __init__(self, net: nn.Module, in_channels: int, length: int, num_splits: int):
+        super().__init__()
+        assert length % num_splits == 0, "length must be divisible by num_splits"
+        self.length = length
+        self.in_channels = in_channels
+        self.num_splits = num_splits
+        self.split_length = length // num_splits
+        self.net = net
+
+    @property
+    def device(self):
+        return next(self.net.parameters()).device
+
+    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
+        angle = sigmas * pi / 2
+        alpha = torch.cos(angle)
+        beta = torch.sin(angle)
+        return alpha, beta
+
+    def get_sigmas_ladder(self, num_items: int, num_steps_per_split: int) -> Tensor:
+        b, n, l, i = num_items, self.num_splits, self.split_length, num_steps_per_split
+        n_half = n // 2  # Only half ladder, rest is zero, to leave some context
+        sigmas = torch.linspace(1, 0, i * n_half, device=self.device)
+        sigmas = repeat(sigmas, "(n i) -> i b 1 (n l)", b=b, l=l, n=n_half)
+        sigmas = torch.flip(sigmas, dims=[-1])  # Lowest noise level first
+        sigmas = F.pad(sigmas, pad=[0, 0, 0, 0, 0, 0, 0, 1])  # Add index i+1
+        sigmas[-1, :, :, l:] = sigmas[0, :, :, :-l]  # Loop back at index i+1
+        return torch.cat([torch.zeros_like(sigmas), sigmas], dim=-1)
+
+    def sample_loop(
+        self, current: Tensor, sigmas: Tensor, show_progress: bool = False, **kwargs
+    ) -> Tensor:
+        num_steps = sigmas.shape[0] - 1
+        alphas, betas = self.get_alpha_beta(sigmas)
+        progress_bar = tqdm(range(num_steps), disable=not show_progress)
+
+        for i in progress_bar:
+            channels = torch.cat([current, sigmas[i]], dim=1)
+            v_pred = self.net(channels, **kwargs)
+            x_pred = alphas[i] * current - betas[i] * v_pred
+            noise_pred = betas[i] * current + alphas[i] * v_pred
+            current = alphas[i + 1] * x_pred + betas[i + 1] * noise_pred
+            progress_bar.set_description(f"Sampling (noise={sigmas[i+1,0,0,0]:.2f})")
+
+        return current
+
+    def sample_start(self, num_items: int, num_steps: int, **kwargs) -> Tensor:
+        b, c, t = num_items, self.in_channels, self.length
+        # Same sigma schedule over all chunks
+        sigmas = torch.linspace(1, 0, num_steps + 1, device=self.device)
+        sigmas = repeat(sigmas, "i -> i b 1 t", b=b, t=t)
+        noise = torch.randn((b, c, t), device=self.device) * sigmas[0]
+        # Sample start
+        return self.sample_loop(current=noise, sigmas=sigmas, **kwargs)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        num_items: int,
+        num_chunks: int,
+        num_steps: int,
+        start: Optional[Tensor] = None,
+        show_progress: bool = False,
+        **kwargs,
+    ) -> Tensor:
+        assert_message = f"required at least {self.num_splits} chunks"
+        assert num_chunks >= self.num_splits, assert_message
+
+        # Sample initial chunks
+        start = self.sample_start(num_items=num_items, num_steps=num_steps, **kwargs)
+        # Return start if only num_splits chunks
+        if num_chunks == self.num_splits:
+            return start
+
+        # Get sigmas for autoregressive ladder
+        b, n = num_items, self.num_splits
+        assert num_steps >= n, "num_steps must be greater than num_splits"
+        sigmas = self.get_sigmas_ladder(
+            num_items=b,
+            num_steps_per_split=num_steps // self.num_splits,
+        )
+        alphas, betas = self.get_alpha_beta(sigmas)
+
+        # Noise start to match ladder and set starting chunks
+        start_noise = alphas[0] * start + betas[0] * torch.randn_like(start)
+        chunks = list(start_noise.chunk(chunks=n, dim=-1))
+
+        # Loop over ladder shifts
+        num_shifts = num_chunks  # - self.num_splits
+        progress_bar = tqdm(range(num_shifts), disable=not show_progress)
+
+        for j in progress_bar:
+            # Decrease ladder noise of last n chunks
+            updated = self.sample_loop(
+                current=torch.cat(chunks[-n:], dim=-1), sigmas=sigmas, **kwargs
+            )
+            # Update chunks
+            chunks[-n:] = list(updated.chunk(chunks=n, dim=-1))
+            # Add fresh noise chunk
+            shape = (b, self.in_channels, self.split_length)
+            chunks += [torch.randn(shape, device=self.device)]
+
+        return torch.cat(chunks[:num_chunks], dim=-1)
+
+
+"""  Inpainters """
+
+
+class Inpainter(nn.Module):
+    pass
+
+
+class VInpainter(Inpainter):
+
+    diffusion_types = [VDiffusion]
+
+    def __init__(self, net: nn.Module, schedule: Schedule = LinearSchedule()):
+        super().__init__()
+        self.net = net
+        self.schedule = schedule
+
+    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
+        angle = sigmas * pi / 2
+        alpha, beta = torch.cos(angle), torch.sin(angle)
+        return alpha, beta
+
+    @torch.no_grad()
+    def forward(  # type: ignore
+        self,
+        source: Tensor,
+        mask: Tensor,
+        num_steps: int,
+        num_resamples: int,
+        show_progress: bool = False,
+        x_noisy: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Tensor:
+        x_noisy = default(x_noisy, lambda: torch.randn_like(source))
+        b = x_noisy.shape[0]
+        sigmas = self.schedule(num_steps + 1, device=x_noisy.device)
+        sigmas = repeat(sigmas, "i -> i b", b=b)
+        sigmas_batch = extend_dim(sigmas, dim=x_noisy.ndim + 1)
+        alphas, betas = self.get_alpha_beta(sigmas_batch)
+        progress_bar = tqdm(range(num_steps), disable=not show_progress)
+
+        for i in progress_bar:
+            for r in range(num_resamples):
+                v_pred = self.net(x_noisy, sigmas[i], **kwargs)
+                x_pred = alphas[i] * x_noisy - betas[i] * v_pred
+                noise_pred = betas[i] * x_noisy + alphas[i] * v_pred
+                # Renoise to current noise level if resampling
+                j = r == num_resamples - 1
+                x_noisy = alphas[i + j] * x_pred + betas[i + j] * noise_pred
+                s_noisy = alphas[i + j] * source + betas[i + j] * torch.randn_like(
+                    source
+                )
+                x_noisy = s_noisy * mask + x_noisy * ~mask
+
+            progress_bar.set_description(f"Inpainting (noise={sigmas[i+1,0]:.2f})")
+
+        return x_noisy

Modules/diffusion/audio_diffusion_pytorch/models.py

@@ -0,0 +1,250 @@
+from abc import ABC, abstractmethod
+from math import floor
+from typing import Any, Callable, Optional, Sequence, Tuple, Union
+
+import torch
+from einops import pack, rearrange, unpack
+from torch import Generator, Tensor, nn
+
+from .components import AppendChannelsPlugin, MelSpectrogram
+from .diffusion import ARVDiffusion, ARVSampler, VDiffusion, VSampler
+from .utils import (
+    closest_power_2,
+    default,
+    downsample,
+    exists,
+    groupby,
+    randn_like,
+    upsample,
+)
+
+
+class DiffusionModel(nn.Module):
+    def __init__(
+        self,
+        net_t: Callable,
+        diffusion_t: Callable = VDiffusion,
+        sampler_t: Callable = VSampler,
+        loss_fn: Callable = torch.nn.functional.mse_loss,
+        dim: int = 1,
+        **kwargs,
+    ):
+        super().__init__()
+        diffusion_kwargs, kwargs = groupby("diffusion_", kwargs)
+        sampler_kwargs, kwargs = groupby("sampler_", kwargs)
+
+        self.net = net_t(dim=dim, **kwargs)
+        self.diffusion = diffusion_t(net=self.net, loss_fn=loss_fn, **diffusion_kwargs)
+        self.sampler = sampler_t(net=self.net, **sampler_kwargs)
+
+    def forward(self, *args, **kwargs) -> Tensor:
+        return self.diffusion(*args, **kwargs)
+
+    @torch.no_grad()
+    def sample(self, *args, **kwargs) -> Tensor:
+        return self.sampler(*args, **kwargs)
+
+
+class EncoderBase(nn.Module, ABC):
+    """Abstract class for DiffusionAE encoder"""
+
+    @abstractmethod
+    def __init__(self):
+        super().__init__()
+        self.out_channels = None
+        self.downsample_factor = None
+
+
+class AdapterBase(nn.Module, ABC):
+    """Abstract class for DiffusionAE encoder"""
+
+    @abstractmethod
+    def encode(self, x: Tensor) -> Tensor:
+        pass
+
+    @abstractmethod
+    def decode(self, x: Tensor) -> Tensor:
+        pass
+
+
+class DiffusionAE(DiffusionModel):
+    """Diffusion Auto Encoder"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        channels: Sequence[int],
+        encoder: EncoderBase,
+        inject_depth: int,
+        latent_factor: Optional[int] = None,
+        adapter: Optional[AdapterBase] = None,
+        **kwargs,
+    ):
+        context_channels = [0] * len(channels)
+        context_channels[inject_depth] = encoder.out_channels
+        super().__init__(
+            in_channels=in_channels,
+            channels=channels,
+            context_channels=context_channels,
+            **kwargs,
+        )
+        self.in_channels = in_channels
+        self.encoder = encoder
+        self.inject_depth = inject_depth
+        # Optional custom latent factor and adapter
+        self.latent_factor = default(latent_factor, self.encoder.downsample_factor)
+        self.adapter = adapter.requires_grad_(False) if exists(adapter) else None
+
+    def forward(  # type: ignore
+        self, x: Tensor, with_info: bool = False, **kwargs
+    ) -> Union[Tensor, Tuple[Tensor, Any]]:
+        # Encode input to latent channels
+        latent, info = self.encode(x, with_info=True)
+        channels = [None] * self.inject_depth + [latent]
+        # Adapt input to diffusion if adapter provided
+        x = self.adapter.encode(x) if exists(self.adapter) else x
+        # Compute diffusion loss
+        loss = super().forward(x, channels=channels, **kwargs)
+        return (loss, info) if with_info else loss
+
+    def encode(self, *args, **kwargs):
+        return self.encoder(*args, **kwargs)
+
+    @torch.no_grad()
+    def decode(
+        self, latent: Tensor, generator: Optional[Generator] = None, **kwargs
+    ) -> Tensor:
+        b = latent.shape[0]
+        noise_length = closest_power_2(latent.shape[2] * self.latent_factor)
+        # Compute noise by inferring shape from latent length
+        noise = torch.randn(
+            (b, self.in_channels, noise_length),
+            device=latent.device,
+            dtype=latent.dtype,
+            generator=generator,
+        )
+        # Compute context from latent
+        channels = [None] * self.inject_depth + [latent]  # type: ignore
+        # Decode by sampling while conditioning on latent channels
+        out = super().sample(noise, channels=channels, **kwargs)
+        # Decode output with adapter if provided
+        return self.adapter.decode(out) if exists(self.adapter) else out
+
+
+class DiffusionUpsampler(DiffusionModel):
+    def __init__(
+        self,
+        in_channels: int,
+        upsample_factor: int,
+        net_t: Callable,
+        **kwargs,
+    ):
+        self.upsample_factor = upsample_factor
+        super().__init__(
+            net_t=AppendChannelsPlugin(net_t, channels=in_channels),
+            in_channels=in_channels,
+            **kwargs,
+        )
+
+    def reupsample(self, x: Tensor) -> Tensor:
+        x = x.clone()
+        x = downsample(x, factor=self.upsample_factor)
+        x = upsample(x, factor=self.upsample_factor)
+        return x
+
+    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:  # type: ignore
+        reupsampled = self.reupsample(x)
+        return super().forward(x, *args, append_channels=reupsampled, **kwargs)
+
+    @torch.no_grad()
+    def sample(  # type: ignore
+        self, downsampled: Tensor, generator: Optional[Generator] = None, **kwargs
+    ) -> Tensor:
+        reupsampled = upsample(downsampled, factor=self.upsample_factor)
+        noise = randn_like(reupsampled, generator=generator)
+        return super().sample(noise, append_channels=reupsampled, **kwargs)
+
+
+class DiffusionVocoder(DiffusionModel):
+    def __init__(
+        self,
+        net_t: Callable,
+        mel_channels: int,
+        mel_n_fft: int,
+        mel_hop_length: Optional[int] = None,
+        mel_win_length: Optional[int] = None,
+        in_channels: int = 1,  # Ignored: channels are automatically batched.
+        **kwargs,
+    ):
+        mel_hop_length = default(mel_hop_length, floor(mel_n_fft) // 4)
+        mel_win_length = default(mel_win_length, mel_n_fft)
+        mel_kwargs, kwargs = groupby("mel_", kwargs)
+        super().__init__(
+            net_t=AppendChannelsPlugin(net_t, channels=1),
+            in_channels=1,
+            **kwargs,
+        )
+        self.to_spectrogram = MelSpectrogram(
+            n_fft=mel_n_fft,
+            hop_length=mel_hop_length,
+            win_length=mel_win_length,
+            n_mel_channels=mel_channels,
+            **mel_kwargs,
+        )
+        self.to_flat = nn.ConvTranspose1d(
+            in_channels=mel_channels,
+            out_channels=1,
+            kernel_size=mel_win_length,
+            stride=mel_hop_length,
+            padding=(mel_win_length - mel_hop_length) // 2,
+            bias=False,
+        )
+
+    def forward(self, x: Tensor, *args, **kwargs) -> Tensor:  # type: ignore
+        # Get spectrogram, pack channels and flatten
+        spectrogram = rearrange(self.to_spectrogram(x), "b c f l -> (b c) f l")
+        spectrogram_flat = self.to_flat(spectrogram)
+        # Pack wave channels
+        x = rearrange(x, "b c t -> (b c) 1 t")
+        return super().forward(x, *args, append_channels=spectrogram_flat, **kwargs)
+
+    @torch.no_grad()
+    def sample(  # type: ignore
+        self, spectrogram: Tensor, generator: Optional[Generator] = None, **kwargs
+    ) -> Tensor:  # type: ignore
+        # Pack channels and flatten spectrogram
+        spectrogram, ps = pack([spectrogram], "* f l")
+        spectrogram_flat = self.to_flat(spectrogram)
+        # Get start noise and sample
+        noise = randn_like(spectrogram_flat, generator=generator)
+        waveform = super().sample(noise, append_channels=spectrogram_flat, **kwargs)
+        # Unpack wave channels
+        waveform = rearrange(waveform, "... 1 t -> ... t")
+        waveform = unpack(waveform, ps, "* t")[0]
+        return waveform
+
+
+class DiffusionAR(DiffusionModel):
+    def __init__(
+        self,
+        in_channels: int,
+        length: int,
+        num_splits: int,
+        diffusion_t: Callable = ARVDiffusion,
+        sampler_t: Callable = ARVSampler,
+        **kwargs,
+    ):
+        super().__init__(
+            in_channels=in_channels + 1,
+            out_channels=in_channels,
+            diffusion_t=diffusion_t,
+            diffusion_length=length,
+            diffusion_num_splits=num_splits,
+            sampler_t=sampler_t,
+            sampler_in_channels=in_channels,
+            sampler_length=length,
+            sampler_num_splits=num_splits,
+            use_time_conditioning=False,
+            use_modulation=False,
+            **kwargs,
+        )

Modules/diffusion/audio_diffusion_pytorch/utils.py

@@ -0,0 +1,125 @@
+from functools import reduce
+from inspect import isfunction
+from math import ceil, floor, log2, pi
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Generator, Tensor
+from typing_extensions import TypeGuard
+
+T = TypeVar("T")
+
+
+def exists(val: Optional[T]) -> TypeGuard[T]:
+    return val is not None
+
+
+def iff(condition: bool, value: T) -> Optional[T]:
+    return value if condition else None
+
+
+def is_sequence(obj: T) -> TypeGuard[Union[list, tuple]]:
+    return isinstance(obj, list) or isinstance(obj, tuple)
+
+
+def default(val: Optional[T], d: Union[Callable[..., T], T]) -> T:
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def to_list(val: Union[T, Sequence[T]]) -> List[T]:
+    if isinstance(val, tuple):
+        return list(val)
+    if isinstance(val, list):
+        return val
+    return [val]  # type: ignore
+
+
+def prod(vals: Sequence[int]) -> int:
+    return reduce(lambda x, y: x * y, vals)
+
+
+def closest_power_2(x: float) -> int:
+    exponent = log2(x)
+    distance_fn = lambda z: abs(x - 2 ** z)  # noqa
+    exponent_closest = min((floor(exponent), ceil(exponent)), key=distance_fn)
+    return 2 ** int(exponent_closest)
+
+
+"""
+Kwargs Utils
+"""
+
+
+def group_dict_by_prefix(prefix: str, d: Dict) -> Tuple[Dict, Dict]:
+    return_dicts: Tuple[Dict, Dict] = ({}, {})
+    for key in d.keys():
+        no_prefix = int(not key.startswith(prefix))
+        return_dicts[no_prefix][key] = d[key]
+    return return_dicts
+
+
+def groupby(prefix: str, d: Dict, keep_prefix: bool = False) -> Tuple[Dict, Dict]:
+    kwargs_with_prefix, kwargs = group_dict_by_prefix(prefix, d)
+    if keep_prefix:
+        return kwargs_with_prefix, kwargs
+    kwargs_no_prefix = {k[len(prefix) :]: v for k, v in kwargs_with_prefix.items()}
+    return kwargs_no_prefix, kwargs
+
+
+def prefix_dict(prefix: str, d: Dict) -> Dict:
+    return {prefix + str(k): v for k, v in d.items()}
+
+
+"""
+DSP Utils
+"""
+
+
+def resample(
+    waveforms: Tensor,
+    factor_in: int,
+    factor_out: int,
+    rolloff: float = 0.99,
+    lowpass_filter_width: int = 6,
+) -> Tensor:
+    """Resamples a waveform using sinc interpolation, adapted from torchaudio"""
+    b, _, length = waveforms.shape
+    length_target = int(factor_out * length / factor_in)
+    d = dict(device=waveforms.device, dtype=waveforms.dtype)
+
+    base_factor = min(factor_in, factor_out) * rolloff
+    width = ceil(lowpass_filter_width * factor_in / base_factor)
+    idx = torch.arange(-width, width + factor_in, **d)[None, None] / factor_in  # type: ignore # noqa
+    t = torch.arange(0, -factor_out, step=-1, **d)[:, None, None] / factor_out + idx  # type: ignore # noqa
+    t = (t * base_factor).clamp(-lowpass_filter_width, lowpass_filter_width) * pi
+
+    window = torch.cos(t / lowpass_filter_width / 2) ** 2
+    scale = base_factor / factor_in
+    kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t)
+    kernels *= window * scale
+
+    waveforms = rearrange(waveforms, "b c t -> (b c) t")
+    waveforms = F.pad(waveforms, (width, width + factor_in))
+    resampled = F.conv1d(waveforms[:, None], kernels, stride=factor_in)
+    resampled = rearrange(resampled, "(b c) k l -> b c (l k)", b=b)
+    return resampled[..., :length_target]
+
+
+def downsample(waveforms: Tensor, factor: int, **kwargs) -> Tensor:
+    return resample(waveforms, factor_in=factor, factor_out=1, **kwargs)
+
+
+def upsample(waveforms: Tensor, factor: int, **kwargs) -> Tensor:
+    return resample(waveforms, factor_in=1, factor_out=factor, **kwargs)
+
+
+""" Torch Utils """
+
+
+def randn_like(tensor: Tensor, *args, generator: Optional[Generator] = None, **kwargs):
+    """randn_like that supports generator"""
+    return torch.randn(tensor.shape, *args, generator=generator, **kwargs).to(tensor)

Modules/diffusion/diffusion.py

@@ -0,0 +1,94 @@
+from math import pi
+from random import randint
+from typing import Any, Optional, Sequence, Tuple, Union
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+from tqdm import tqdm
+
+from .utils import *
+from .sampler import *
+
+"""
+Diffusion Classes (generic for 1d data)
+"""
+
+
+class Model1d(nn.Module):
+    def __init__(self, unet_type: str = "base", **kwargs):
+        super().__init__()
+        diffusion_kwargs, kwargs = groupby("diffusion_", kwargs)
+        self.unet = None
+        self.diffusion = None
+
+    def forward(self, x: Tensor, **kwargs) -> Tensor:
+        return self.diffusion(x, **kwargs)
+
+    def sample(self, *args, **kwargs) -> Tensor:
+        return self.diffusion.sample(*args, **kwargs)
+
+
+"""
+Audio Diffusion Classes (specific for 1d audio data)
+"""
+
+
+def get_default_model_kwargs():
+    return dict(
+        channels=128,
+        patch_size=16,
+        multipliers=[1, 2, 4, 4, 4, 4, 4],
+        factors=[4, 4, 4, 2, 2, 2],
+        num_blocks=[2, 2, 2, 2, 2, 2],
+        attentions=[0, 0, 0, 1, 1, 1, 1],
+        attention_heads=8,
+        attention_features=64,
+        attention_multiplier=2,
+        attention_use_rel_pos=False,
+        diffusion_type="v",
+        diffusion_sigma_distribution=UniformDistribution(),
+    )
+
+
+def get_default_sampling_kwargs():
+    return dict(sigma_schedule=LinearSchedule(), sampler=VSampler(), clamp=True)
+
+
+class AudioDiffusionModel(Model1d):
+    def __init__(self, **kwargs):
+        super().__init__(**{**get_default_model_kwargs(), **kwargs})
+
+    def sample(self, *args, **kwargs):
+        return super().sample(*args, **{**get_default_sampling_kwargs(), **kwargs})
+
+
+class AudioDiffusionConditional(Model1d):
+    def __init__(
+        self,
+        embedding_features: int,
+        embedding_max_length: int,
+        embedding_mask_proba: float = 0.1,
+        **kwargs,
+    ):
+        self.embedding_mask_proba = embedding_mask_proba
+        default_kwargs = dict(
+            **get_default_model_kwargs(),
+            unet_type="cfg",
+            context_embedding_features=embedding_features,
+            context_embedding_max_length=embedding_max_length,
+        )
+        super().__init__(**{**default_kwargs, **kwargs})
+
+    def forward(self, *args, **kwargs):
+        default_kwargs = dict(embedding_mask_proba=self.embedding_mask_proba)
+        return super().forward(*args, **{**default_kwargs, **kwargs})
+
+    def sample(self, *args, **kwargs):
+        default_kwargs = dict(
+            **get_default_sampling_kwargs(),
+            embedding_scale=5.0,
+        )
+        return super().sample(*args, **{**default_kwargs, **kwargs})
+
+

Modules/diffusion/modules.py

@@ -0,0 +1,693 @@
+from math import floor, log, pi
+from typing import Any, List, Optional, Sequence, Tuple, Union
+
+from .utils import *
+
+import torch
+import torch.nn as nn
+from einops import rearrange, reduce, repeat
+from einops.layers.torch import Rearrange
+from einops_exts import rearrange_many
+from torch import Tensor, einsum
+
+
+"""
+Utils
+"""
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, style_dim, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.fc = nn.Linear(style_dim, channels*2)
+
+    def forward(self, x, s):
+        x = x.transpose(-1, -2)
+        x = x.transpose(1, -1)
+                
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
+        
+        
+        x = F.layer_norm(x, (self.channels,), eps=self.eps)
+        x = (1 + gamma) * x + beta
+        return x.transpose(1, -1).transpose(-1, -2)
+
+class StyleTransformer1d(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        channels: int,
+        num_heads: int,
+        head_features: int,
+        multiplier: int,
+        use_context_time: bool = True,
+        use_rel_pos: bool = False,
+        context_features_multiplier: int = 1,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+        context_features: Optional[int] = None,
+        context_embedding_features: Optional[int] = None,
+        embedding_max_length: int = 512,
+    ):
+        super().__init__()
+
+        self.blocks = nn.ModuleList(
+            [
+                StyleTransformerBlock(
+                    features=channels + context_embedding_features,
+                    head_features=head_features,
+                    num_heads=num_heads,
+                    multiplier=multiplier,
+                    style_dim=context_features,
+                    use_rel_pos=use_rel_pos,
+                    rel_pos_num_buckets=rel_pos_num_buckets,
+                    rel_pos_max_distance=rel_pos_max_distance,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.to_out = nn.Sequential(
+            Rearrange("b t c -> b c t"),
+            nn.Conv1d(
+                in_channels=channels + context_embedding_features,
+                out_channels=channels,
+                kernel_size=1,
+            ),
+        )
+        
+        use_context_features = exists(context_features)
+        self.use_context_features = use_context_features
+        self.use_context_time = use_context_time
+
+        if use_context_time or use_context_features:
+            context_mapping_features = channels + context_embedding_features
+
+            self.to_mapping = nn.Sequential(
+                nn.Linear(context_mapping_features, context_mapping_features),
+                nn.GELU(),
+                nn.Linear(context_mapping_features, context_mapping_features),
+                nn.GELU(),
+            )
+        
+        if use_context_time:
+            assert exists(context_mapping_features)
+            self.to_time = nn.Sequential(
+                TimePositionalEmbedding(
+                    dim=channels, out_features=context_mapping_features
+                ),
+                nn.GELU(),
+            )
+
+        if use_context_features:
+            assert exists(context_features) and exists(context_mapping_features)
+            self.to_features = nn.Sequential(
+                nn.Linear(
+                    in_features=context_features, out_features=context_mapping_features
+                ),
+                nn.GELU(),
+            )
+            
+        self.fixed_embedding = FixedEmbedding(
+            max_length=embedding_max_length, features=context_embedding_features
+        )
+        
+
+    def get_mapping(
+        self, time: Optional[Tensor] = None, features: Optional[Tensor] = None
+    ) -> Optional[Tensor]:
+        """Combines context time features and features into mapping"""
+        items, mapping = [], None
+        # Compute time features
+        if self.use_context_time:
+            assert_message = "use_context_time=True but no time features provided"
+            assert exists(time), assert_message
+            items += [self.to_time(time)]
+        # Compute features
+        if self.use_context_features:
+            assert_message = "context_features exists but no features provided"
+            assert exists(features), assert_message
+            items += [self.to_features(features)]
+
+        # Compute joint mapping
+        if self.use_context_time or self.use_context_features:
+            mapping = reduce(torch.stack(items), "n b m -> b m", "sum")
+            mapping = self.to_mapping(mapping)
+
+        return mapping
+            
+    def run(self, x, time, embedding, features):
+        
+        mapping = self.get_mapping(time, features)
+        x = torch.cat([x.expand(-1, embedding.size(1), -1), embedding], axis=-1)
+        mapping = mapping.unsqueeze(1).expand(-1, embedding.size(1), -1)
+        
+        for block in self.blocks:
+            x = x + mapping
+            x = block(x, features)
+        
+        x = x.mean(axis=1).unsqueeze(1)
+        x = self.to_out(x)
+        x = x.transpose(-1, -2)
+        
+        return x
+        
+    def forward(self, x: Tensor, 
+                time: Tensor, 
+                embedding_mask_proba: float = 0.0,
+                embedding: Optional[Tensor] = None, 
+                features: Optional[Tensor] = None,
+               embedding_scale: float = 1.0) -> Tensor:
+        
+        b, device = embedding.shape[0], embedding.device
+        fixed_embedding = self.fixed_embedding(embedding)
+        if embedding_mask_proba > 0.0:
+            # Randomly mask embedding
+            batch_mask = rand_bool(
+                shape=(b, 1, 1), proba=embedding_mask_proba, device=device
+            )
+            embedding = torch.where(batch_mask, fixed_embedding, embedding)
+
+        if embedding_scale != 1.0:
+            # Compute both normal and fixed embedding outputs
+            out = self.run(x, time, embedding=embedding, features=features)
+            out_masked = self.run(x, time, embedding=fixed_embedding, features=features)
+            # Scale conditional output using classifier-free guidance
+            return out_masked + (out - out_masked) * embedding_scale
+        else:
+            return self.run(x, time, embedding=embedding, features=features)
+        
+        return x
+
+
+class StyleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        num_heads: int,
+        head_features: int,
+        style_dim: int,
+        multiplier: int,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+        context_features: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.use_cross_attention = exists(context_features) and context_features > 0
+
+        self.attention = StyleAttention(
+            features=features,
+            style_dim=style_dim,
+            num_heads=num_heads,
+            head_features=head_features,
+            use_rel_pos=use_rel_pos,
+            rel_pos_num_buckets=rel_pos_num_buckets,
+            rel_pos_max_distance=rel_pos_max_distance,
+        )
+
+        if self.use_cross_attention:
+            self.cross_attention = StyleAttention(
+                features=features,
+                style_dim=style_dim,
+                num_heads=num_heads,
+                head_features=head_features,
+                context_features=context_features,
+                use_rel_pos=use_rel_pos,
+                rel_pos_num_buckets=rel_pos_num_buckets,
+                rel_pos_max_distance=rel_pos_max_distance,
+            )
+
+        self.feed_forward = FeedForward(features=features, multiplier=multiplier)
+
+    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
+        x = self.attention(x, s) + x
+        if self.use_cross_attention:
+            x = self.cross_attention(x, s, context=context) + x
+        x = self.feed_forward(x) + x
+        return x
+
+class StyleAttention(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        *,
+        style_dim: int,
+        head_features: int,
+        num_heads: int,
+        context_features: Optional[int] = None,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+    ):
+        super().__init__()
+        self.context_features = context_features
+        mid_features = head_features * num_heads
+        context_features = default(context_features, features)
+
+        self.norm = AdaLayerNorm(style_dim, features)
+        self.norm_context = AdaLayerNorm(style_dim, context_features)
+        self.to_q = nn.Linear(
+            in_features=features, out_features=mid_features, bias=False
+        )
+        self.to_kv = nn.Linear(
+            in_features=context_features, out_features=mid_features * 2, bias=False
+        )
+        self.attention = AttentionBase(
+            features,
+            num_heads=num_heads,
+            head_features=head_features,
+            use_rel_pos=use_rel_pos,
+            rel_pos_num_buckets=rel_pos_num_buckets,
+            rel_pos_max_distance=rel_pos_max_distance,
+        )
+
+    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
+        assert_message = "You must provide a context when using context_features"
+        assert not self.context_features or exists(context), assert_message
+        # Use context if provided
+        context = default(context, x)
+        # Normalize then compute q from input and k,v from context
+        x, context = self.norm(x, s), self.norm_context(context, s)
+        
+        q, k, v = (self.to_q(x), *torch.chunk(self.to_kv(context), chunks=2, dim=-1))
+        # Compute and return attention
+        return self.attention(q, k, v)
+        
+class Transformer1d(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        channels: int,
+        num_heads: int,
+        head_features: int,
+        multiplier: int,
+        use_context_time: bool = True,
+        use_rel_pos: bool = False,
+        context_features_multiplier: int = 1,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+        context_features: Optional[int] = None,
+        context_embedding_features: Optional[int] = None,
+        embedding_max_length: int = 512,
+    ):
+        super().__init__()
+
+        self.blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    features=channels + context_embedding_features,
+                    head_features=head_features,
+                    num_heads=num_heads,
+                    multiplier=multiplier,
+                    use_rel_pos=use_rel_pos,
+                    rel_pos_num_buckets=rel_pos_num_buckets,
+                    rel_pos_max_distance=rel_pos_max_distance,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.to_out = nn.Sequential(
+            Rearrange("b t c -> b c t"),
+            nn.Conv1d(
+                in_channels=channels + context_embedding_features,
+                out_channels=channels,
+                kernel_size=1,
+            ),
+        )
+        
+        use_context_features = exists(context_features)
+        self.use_context_features = use_context_features
+        self.use_context_time = use_context_time
+
+        if use_context_time or use_context_features:
+            context_mapping_features = channels + context_embedding_features
+
+            self.to_mapping = nn.Sequential(
+                nn.Linear(context_mapping_features, context_mapping_features),
+                nn.GELU(),
+                nn.Linear(context_mapping_features, context_mapping_features),
+                nn.GELU(),
+            )
+        
+        if use_context_time:
+            assert exists(context_mapping_features)
+            self.to_time = nn.Sequential(
+                TimePositionalEmbedding(
+                    dim=channels, out_features=context_mapping_features
+                ),
+                nn.GELU(),
+            )
+
+        if use_context_features:
+            assert exists(context_features) and exists(context_mapping_features)
+            self.to_features = nn.Sequential(
+                nn.Linear(
+                    in_features=context_features, out_features=context_mapping_features
+                ),
+                nn.GELU(),
+            )
+            
+        self.fixed_embedding = FixedEmbedding(
+            max_length=embedding_max_length, features=context_embedding_features
+        )
+        
+
+    def get_mapping(
+        self, time: Optional[Tensor] = None, features: Optional[Tensor] = None
+    ) -> Optional[Tensor]:
+        """Combines context time features and features into mapping"""
+        items, mapping = [], None
+        # Compute time features
+        if self.use_context_time:
+            assert_message = "use_context_time=True but no time features provided"
+            assert exists(time), assert_message
+            items += [self.to_time(time)]
+        # Compute features
+        if self.use_context_features:
+            assert_message = "context_features exists but no features provided"
+            assert exists(features), assert_message
+            items += [self.to_features(features)]
+
+        # Compute joint mapping
+        if self.use_context_time or self.use_context_features:
+            mapping = reduce(torch.stack(items), "n b m -> b m", "sum")
+            mapping = self.to_mapping(mapping)
+
+        return mapping
+            
+    def run(self, x, time, embedding, features):
+        
+        mapping = self.get_mapping(time, features)
+        x = torch.cat([x.expand(-1, embedding.size(1), -1), embedding], axis=-1)
+        mapping = mapping.unsqueeze(1).expand(-1, embedding.size(1), -1)
+        
+        for block in self.blocks:
+            x = x + mapping
+            x = block(x)
+        
+        x = x.mean(axis=1).unsqueeze(1)
+        x = self.to_out(x)
+        x = x.transpose(-1, -2)
+        
+        return x
+        
+    def forward(self, x: Tensor, 
+                time: Tensor, 
+                embedding_mask_proba: float = 0.0,
+                embedding: Optional[Tensor] = None, 
+                features: Optional[Tensor] = None,
+               embedding_scale: float = 1.0) -> Tensor:
+        
+        b, device = embedding.shape[0], embedding.device
+        fixed_embedding = self.fixed_embedding(embedding)
+        if embedding_mask_proba > 0.0:
+            # Randomly mask embedding
+            batch_mask = rand_bool(
+                shape=(b, 1, 1), proba=embedding_mask_proba, device=device
+            )
+            embedding = torch.where(batch_mask, fixed_embedding, embedding)
+
+        if embedding_scale != 1.0:
+            # Compute both normal and fixed embedding outputs
+            out = self.run(x, time, embedding=embedding, features=features)
+            out_masked = self.run(x, time, embedding=fixed_embedding, features=features)
+            # Scale conditional output using classifier-free guidance
+            return out_masked + (out - out_masked) * embedding_scale
+        else:
+            return self.run(x, time, embedding=embedding, features=features)
+        
+        return x
+
+
+"""
+Attention Components
+"""
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, num_buckets: int, max_distance: int, num_heads: int):
+        super().__init__()
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.num_heads = num_heads
+        self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position: Tensor, num_buckets: int, max_distance: int
+    ):
+        num_buckets //= 2
+        ret = (relative_position >= 0).to(torch.long) * num_buckets
+        n = torch.abs(relative_position)
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = (
+            max_exact
+            + (
+                torch.log(n.float() / max_exact)
+                / log(max_distance / max_exact)
+                * (num_buckets - max_exact)
+            ).long()
+        )
+        val_if_large = torch.min(
+            val_if_large, torch.full_like(val_if_large, num_buckets - 1)
+        )
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, num_queries: int, num_keys: int) -> Tensor:
+        i, j, device = num_queries, num_keys, self.relative_attention_bias.weight.device
+        q_pos = torch.arange(j - i, j, dtype=torch.long, device=device)
+        k_pos = torch.arange(j, dtype=torch.long, device=device)
+        rel_pos = rearrange(k_pos, "j -> 1 j") - rearrange(q_pos, "i -> i 1")
+
+        relative_position_bucket = self._relative_position_bucket(
+            rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance
+        )
+
+        bias = self.relative_attention_bias(relative_position_bucket)
+        bias = rearrange(bias, "m n h -> 1 h m n")
+        return bias
+
+
+def FeedForward(features: int, multiplier: int) -> nn.Module:
+    mid_features = features * multiplier
+    return nn.Sequential(
+        nn.Linear(in_features=features, out_features=mid_features),
+        nn.GELU(),
+        nn.Linear(in_features=mid_features, out_features=features),
+    )
+
+
+class AttentionBase(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        *,
+        head_features: int,
+        num_heads: int,
+        use_rel_pos: bool,
+        out_features: Optional[int] = None,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+    ):
+        super().__init__()
+        self.scale = head_features ** -0.5
+        self.num_heads = num_heads
+        self.use_rel_pos = use_rel_pos
+        mid_features = head_features * num_heads
+
+        if use_rel_pos:
+            assert exists(rel_pos_num_buckets) and exists(rel_pos_max_distance)
+            self.rel_pos = RelativePositionBias(
+                num_buckets=rel_pos_num_buckets,
+                max_distance=rel_pos_max_distance,
+                num_heads=num_heads,
+            )
+        if out_features is None:
+            out_features = features
+            
+        self.to_out = nn.Linear(in_features=mid_features, out_features=out_features)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        # Split heads
+        q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)
+        # Compute similarity matrix
+        sim = einsum("... n d, ... m d -> ... n m", q, k)
+        sim = (sim + self.rel_pos(*sim.shape[-2:])) if self.use_rel_pos else sim
+        sim = sim * self.scale
+        # Get attention matrix with softmax
+        attn = sim.softmax(dim=-1)
+        # Compute values
+        out = einsum("... n m, ... m d -> ... n d", attn, v)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        return self.to_out(out)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        *,
+        head_features: int,
+        num_heads: int,
+        out_features: Optional[int] = None,
+        context_features: Optional[int] = None,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+    ):
+        super().__init__()
+        self.context_features = context_features
+        mid_features = head_features * num_heads
+        context_features = default(context_features, features)
+
+        self.norm = nn.LayerNorm(features)
+        self.norm_context = nn.LayerNorm(context_features)
+        self.to_q = nn.Linear(
+            in_features=features, out_features=mid_features, bias=False
+        )
+        self.to_kv = nn.Linear(
+            in_features=context_features, out_features=mid_features * 2, bias=False
+        )
+
+        self.attention = AttentionBase(
+            features,
+            out_features=out_features,
+            num_heads=num_heads,
+            head_features=head_features,
+            use_rel_pos=use_rel_pos,
+            rel_pos_num_buckets=rel_pos_num_buckets,
+            rel_pos_max_distance=rel_pos_max_distance,
+        )
+
+    def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
+        assert_message = "You must provide a context when using context_features"
+        assert not self.context_features or exists(context), assert_message
+        # Use context if provided
+        context = default(context, x)
+        # Normalize then compute q from input and k,v from context
+        x, context = self.norm(x), self.norm_context(context)
+        q, k, v = (self.to_q(x), *torch.chunk(self.to_kv(context), chunks=2, dim=-1))
+        # Compute and return attention
+        return self.attention(q, k, v)
+
+
+"""
+Transformer Blocks
+"""
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        num_heads: int,
+        head_features: int,
+        multiplier: int,
+        use_rel_pos: bool,
+        rel_pos_num_buckets: Optional[int] = None,
+        rel_pos_max_distance: Optional[int] = None,
+        context_features: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.use_cross_attention = exists(context_features) and context_features > 0
+
+        self.attention = Attention(
+            features=features,
+            num_heads=num_heads,
+            head_features=head_features,
+            use_rel_pos=use_rel_pos,
+            rel_pos_num_buckets=rel_pos_num_buckets,
+            rel_pos_max_distance=rel_pos_max_distance,
+        )
+
+        if self.use_cross_attention:
+            self.cross_attention = Attention(
+                features=features,
+                num_heads=num_heads,
+                head_features=head_features,
+                context_features=context_features,
+                use_rel_pos=use_rel_pos,
+                rel_pos_num_buckets=rel_pos_num_buckets,
+                rel_pos_max_distance=rel_pos_max_distance,
+            )
+
+        self.feed_forward = FeedForward(features=features, multiplier=multiplier)
+
+    def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
+        x = self.attention(x) + x
+        if self.use_cross_attention:
+            x = self.cross_attention(x, context=context) + x
+        x = self.feed_forward(x) + x
+        return x
+
+
+
+"""
+Time Embeddings
+"""
+
+
+class SinusoidalEmbedding(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: Tensor) -> Tensor:
+        device, half_dim = x.device, self.dim // 2
+        emb = torch.tensor(log(10000) / (half_dim - 1), device=device)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = rearrange(x, "i -> i 1") * rearrange(emb, "j -> 1 j")
+        return torch.cat((emb.sin(), emb.cos()), dim=-1)
+
+
+class LearnedPositionalEmbedding(nn.Module):
+    """Used for continuous time"""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        assert (dim % 2) == 0
+        half_dim = dim // 2
+        self.weights = nn.Parameter(torch.randn(half_dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = rearrange(x, "b -> b 1")
+        freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * pi
+        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
+        fouriered = torch.cat((x, fouriered), dim=-1)
+        return fouriered
+
+
+def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
+    return nn.Sequential(
+        LearnedPositionalEmbedding(dim),
+        nn.Linear(in_features=dim + 1, out_features=out_features),
+    )
+
+class FixedEmbedding(nn.Module):
+    def __init__(self, max_length: int, features: int):
+        super().__init__()
+        self.max_length = max_length
+        self.embedding = nn.Embedding(max_length, features)
+
+    def forward(self, x: Tensor) -> Tensor:
+        batch_size, length, device = *x.shape[0:2], x.device
+        assert_message = "Input sequence length must be <= max_length"
+        assert length <= self.max_length, assert_message
+        position = torch.arange(length, device=device)
+        fixed_embedding = self.embedding(position)
+        fixed_embedding = repeat(fixed_embedding, "n d -> b n d", b=batch_size)
+        return fixed_embedding

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/.github/workflows/python-publish.yml

@@ -0,0 +1,39 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+.mypy_cache

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/.pre-commit-config.yaml

@@ -0,0 +1,41 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+
+# Formats code correctly
+-   repo: https://github.com/psf/black
+    rev: 21.12b0
+    hooks:
+    -   id: black
+        args: [
+            '--experimental-string-processing'
+        ]
+
+# Sorts imports
+-   repo: https://github.com/pycqa/isort
+    rev: 5.10.1
+    hooks:
+    -   id: isort
+        name: isort (python)
+        args: ["--profile", "black"]
+
+# Checks unused imports, like lengths, etc
+-   repo: https://gitlab.com/pycqa/flake8
+    rev: 4.0.0
+    hooks:
+    -   id: flake8
+        args: [
+            '--per-file-ignores=__init__.py:F401',
+            '--max-line-length=88',
+            '--ignore=E203,W503'
+        ]
+
+# Checks types
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.971'
+    hooks:
+    -   id: mypy
+        additional_dependencies: [data-science-types>=0.2, torch>=1.6]

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 archinet.ai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/README.md

@@ -0,0 +1,251 @@
+<img src="./LOGO.png"></img>
+
+A fully featured audio diffusion library, for PyTorch. Includes models for unconditional audio generation, text-conditional audio generation, diffusion autoencoding, upsampling, and vocoding. The provided models are waveform-based, however, the U-Net (built using [`a-unet`](https://github.com/archinetai/a-unet)), `DiffusionModel`, diffusion method, and diffusion samplers are both generic to any dimension and highly customizable to work on other formats. **Notes: (1) no pre-trained models are provided here, (2) the configs shown are indicative and untested, see [Moûsai](https://arxiv.org/abs/2301.11757) for the configs used in the paper.**
+
+
+## Install
+
+```bash
+pip install audio-diffusion-pytorch
+```
+
+[![PyPI - Python Version](https://img.shields.io/pypi/v/audio-diffusion-pytorch?style=flat&colorA=black&colorB=black)](https://pypi.org/project/audio-diffusion-pytorch/)
+[![Downloads](https://static.pepy.tech/personalized-badge/audio-diffusion-pytorch?period=total&units=international_system&left_color=black&right_color=black&left_text=Downloads)](https://pepy.tech/project/audio-diffusion-pytorch)
+
+
+## Usage
+
+### Unconditional Generator
+
+```py
+from audio_diffusion_pytorch import DiffusionModel, UNetV0, VDiffusion, VSampler
+
+model = DiffusionModel(
+    net_t=UNetV0, # The model type used for diffusion (U-Net V0 in this case)
+    in_channels=2, # U-Net: number of input/output (audio) channels
+    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
+    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer
+    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
+    attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer
+    attention_heads=8, # U-Net: number of attention heads per attention item
+    attention_features=64, # U-Net: number of attention features per attention item
+    diffusion_t=VDiffusion, # The diffusion method used
+    sampler_t=VSampler, # The diffusion sampler used
+)
+
+# Train model with audio waveforms
+audio = torch.randn(1, 2, 2**18) # [batch_size, in_channels, length]
+loss = model(audio)
+loss.backward()
+
+# Turn noise into new audio sample with diffusion
+noise = torch.randn(1, 2, 2**18) # [batch_size, in_channels, length]
+sample = model.sample(noise, num_steps=10) # Suggested num_steps 10-100
+```
+
+### Text-Conditional Generator
+A text-to-audio diffusion model that conditions the generation with `t5-base` text embeddings, requires `pip install transformers`.
+```py
+from audio_diffusion_pytorch import DiffusionModel, UNetV0, VDiffusion, VSampler
+
+model = DiffusionModel(
+    # ... same as unconditional model
+    use_text_conditioning=True, # U-Net: enables text conditioning (default T5-base)
+    use_embedding_cfg=True, # U-Net: enables classifier free guidance
+    embedding_max_length=64, # U-Net: text embedding maximum length (default for T5-base)
+    embedding_features=768, # U-Net: text mbedding features (default for T5-base)
+    cross_attentions=[0, 0, 0, 1, 1, 1, 1, 1, 1], # U-Net: cross-attention enabled/disabled at each layer
+)
+
+# Train model with audio waveforms
+audio_wave = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+loss = model(
+    audio_wave,
+    text=['The audio description'], # Text conditioning, one element per batch
+    embedding_mask_proba=0.1 # Probability of masking text with learned embedding (Classifier-Free Guidance Mask)
+)
+loss.backward()
+
+# Turn noise into new audio sample with diffusion
+noise = torch.randn(1, 2, 2**18)
+sample = model.sample(
+    noise,
+    text=['The audio description'],
+    embedding_scale=5.0, # Higher for more text importance, suggested range: 1-15 (Classifier-Free Guidance Scale)
+    num_steps=2 # Higher for better quality, suggested num_steps: 10-100
+)
+```
+
+### Diffusion Upsampler
+Upsample audio from a lower sample rate to higher sample rate using diffusion, e.g. 3kHz to 48kHz.
+```py
+from audio_diffusion_pytorch import DiffusionUpsampler, UNetV0, VDiffusion, VSampler
+
+upsampler = DiffusionUpsampler(
+    net_t=UNetV0, # The model type used for diffusion
+    upsample_factor=16, # The upsample factor (e.g. 16 can be used for 3kHz to 48kHz)
+    in_channels=2, # U-Net: number of input/output (audio) channels
+    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
+    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer
+    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
+    diffusion_t=VDiffusion, # The diffusion method used
+    sampler_t=VSampler, # The diffusion sampler used
+)
+
+# Train model with high sample rate audio waveforms
+audio = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+loss = upsampler(audio)
+loss.backward()
+
+# Turn low sample rate audio into high sample rate
+downsampled_audio = torch.randn(1, 2, 2**14) # [batch, in_channels, length]
+sample = upsampler.sample(downsampled_audio, num_steps=10) # Output has shape: [1, 2, 2**18]
+```
+
+### Diffusion Vocoder
+Convert a mel-spectrogram to wavefrom using diffusion.
+```py
+from audio_diffusion_pytorch import DiffusionVocoder, UNetV0, VDiffusion, VSampler
+
+vocoder = DiffusionVocoder(
+    mel_n_fft=1024, # Mel-spectrogram n_fft
+    mel_channels=80, # Mel-spectrogram channels
+    mel_sample_rate=48000, # Mel-spectrogram sample rate
+    mel_normalize_log=True, # Mel-spectrogram log normalization (alternative is mel_normalize=True for [-1,1] power normalization)
+    net_t=UNetV0, # The model type used for diffusion vocoding
+    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
+    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer
+    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
+    diffusion_t=VDiffusion, # The diffusion method used
+    sampler_t=VSampler, # The diffusion sampler used
+)
+
+# Train model on waveforms (automatically converted to mel internally)
+audio = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+loss = vocoder(audio)
+loss.backward()
+
+# Turn mel spectrogram into waveform
+mel_spectrogram = torch.randn(1, 2, 80, 1024) # [batch, in_channels, mel_channels, mel_length]
+sample = vocoder.sample(mel_spectrogram, num_steps=10) # Output has shape: [1, 2, 2**18]
+```
+
+### Diffusion Autoencoder
+Autoencode audio into a compressed latent using diffusion. Any encoder can be provided as long as it subclasses the `EncoderBase` class or contains an `out_channels` and `downsample_factor` field.
+```py
+from audio_diffusion_pytorch import DiffusionAE, UNetV0, VDiffusion, VSampler
+from audio_encoders_pytorch import MelE1d, TanhBottleneck
+
+autoencoder = DiffusionAE(
+    encoder=MelE1d( # The encoder used, in this case a mel-spectrogram encoder
+        in_channels=2,
+        channels=512,
+        multipliers=[1, 1],
+        factors=[2],
+        num_blocks=[12],
+        out_channels=32,
+        mel_channels=80,
+        mel_sample_rate=48000,
+        mel_normalize_log=True,
+        bottleneck=TanhBottleneck(),
+    ),
+    inject_depth=6,
+    net_t=UNetV0, # The model type used for diffusion upsampling
+    in_channels=2, # U-Net: number of input/output (audio) channels
+    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
+    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer
+    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
+    diffusion_t=VDiffusion, # The diffusion method used
+    sampler_t=VSampler, # The diffusion sampler used
+)
+
+# Train autoencoder with audio samples
+audio = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+loss = autoencoder(audio)
+loss.backward()
+
+# Encode/decode audio
+audio = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+latent = autoencoder.encode(audio) # Encode
+sample = autoencoder.decode(latent, num_steps=10) # Decode by sampling diffusion model conditioning on latent
+```
+
+## Other
+
+### Inpainting
+```py
+from audio_diffusion_pytorch import UNetV0, VInpainter
+
+# The diffusion UNetV0 (this is an example, the net must be trained to work)
+net = UNetV0(
+    dim=1,
+    in_channels=2, # U-Net: number of input/output (audio) channels
+    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
+    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer
+    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
+    attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer
+    attention_heads=8, # U-Net: number of attention heads per attention block
+    attention_features=64, # U-Net: number of attention features per attention block,
+)
+
+# Instantiate inpainter with trained net
+inpainter = VInpainter(net=net)
+
+# Inpaint source
+y = inpainter(
+    source=torch.randn(1, 2, 2**18), # Start source
+    mask=torch.randint(0, 2, (1, 2, 2 ** 18), dtype=torch.bool),  # Set to `True` the parts you want to keep
+    num_steps=10, # Number of inpainting steps
+    num_resamples=2, # Number of resampling steps
+    show_progress=True,
+) # [1, 2, 2 ** 18]
+```
+
+## Appreciation
+
+* [StabilityAI](https://stability.ai/) for the compute, [Zach Evans](https://github.com/zqevans) and everyone else from [HarmonAI](https://www.harmonai.org/) for the interesting research discussions.
+* [ETH Zurich](https://inf.ethz.ch/) for the resources, [Zhijing Jin](https://zhijing-jin.com/), [Bernhard Schoelkopf](https://is.mpg.de/~bs), and [Mrinmaya Sachan](http://www.mrinmaya.io/) for supervising this Thesis.
+* [Phil Wang](https://github.com/lucidrains) for the beautiful open source contributions on [diffusion](https://github.com/lucidrains/denoising-diffusion-pytorch) and [Imagen](https://github.com/lucidrains/imagen-pytorch).
+* [Katherine Crowson](https://github.com/crowsonkb) for the experiments with [k-diffusion](https://github.com/crowsonkb/k-diffusion) and the insane collection of samplers.
+
+## Citations
+
+DDPM Diffusion
+```bibtex
+@misc{2006.11239,
+Author = {Jonathan Ho and Ajay Jain and Pieter Abbeel},
+Title = {Denoising Diffusion Probabilistic Models},
+Year = {2020},
+Eprint = {arXiv:2006.11239},
+}
+```
+
+DDIM (V-Sampler)
+```bibtex
+@misc{2010.02502,
+Author = {Jiaming Song and Chenlin Meng and Stefano Ermon},
+Title = {Denoising Diffusion Implicit Models},
+Year = {2020},
+Eprint = {arXiv:2010.02502},
+}
+```
+
+V-Diffusion
+```bibtex
+@misc{2202.00512,
+Author = {Tim Salimans and Jonathan Ho},
+Title = {Progressive Distillation for Fast Sampling of Diffusion Models},
+Year = {2022},
+Eprint = {arXiv:2202.00512},
+}
+```
+
+Imagen (T5 Text Conditioning)
+```bibtex
+@misc{2205.11487,
+Author = {Chitwan Saharia and William Chan and Saurabh Saxena and Lala Li and Jay Whang and Emily Denton and Seyed Kamyar Seyed Ghasemipour and Burcu Karagol Ayan and S. Sara Mahdavi and Rapha Gontijo Lopes and Tim Salimans and Jonathan Ho and David J Fleet and Mohammad Norouzi},
+Title = {Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding},
+Year = {2022},
+Eprint = {arXiv:2205.11487},
+}
+```

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/setup.py

@@ -0,0 +1,29 @@
+from setuptools import find_packages, setup
+
+setup(
+    name="audio-diffusion-pytorch",
+    packages=find_packages(exclude=[]),
+    version="0.1.3",
+    license="MIT",
+    description="Audio Diffusion - PyTorch",
+    long_description_content_type="text/markdown",
+    author="Flavio Schneider",
+    author_email="archinetai@protonmail.com",
+    url="https://github.com/archinetai/audio-diffusion-pytorch",
+    keywords=["artificial intelligence", "deep learning", "audio generation"],
+    install_requires=[
+        "tqdm",
+        "torch>=1.6",
+        "torchaudio",
+        "data-science-types>=0.2",
+        "einops>=0.6",
+        "a-unet",
+    ],
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3.6",
+    ],
+)

Modules/diffusion/reconstruction_head/audio-diffusion-pytorch/tests/testcustomloss.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn.functional as F
+from audio_diffusion_pytorch import DiffusionAE, UNetV0, VDiffusion, VSampler
+from audio_encoders_pytorch import MelE1d, TanhBottleneck
+from auraloss.freq import MultiResolutionSTFTLoss
+
+autoencoder = DiffusionAE(
+    encoder=MelE1d( # The encoder used, in this case a mel-spectrogram encoder
+        in_channels=2,
+        channels=512,
+        multipliers=[1, 1],
+        factors=[2],
+        num_blocks=[12],
+        out_channels=32,
+        mel_channels=80,
+        mel_sample_rate=48000,
+        mel_normalize_log=True,
+        bottleneck=TanhBottleneck(),
+    ),
+    inject_depth=6,
+    net_t=UNetV0, # The model type used for diffusion upsampling
+    in_channels=2, # U-Net: number of input/output (audio) channels
+    channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer
+    factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer
+    items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer
+    diffusion_t=VDiffusion, # The diffusion method used
+    sampler_t=VSampler, # The diffusion sampler used
+    loss_fn=MultiResolutionSTFTLoss(),  # The loss function used
+)
+
+# Train autoencoder with audio samples
+audio = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+loss = autoencoder(audio)
+loss.backward()
+
+# Encode/decode audio
+audio = torch.randn(1, 2, 2**18) # [batch, in_channels, length]
+latent = autoencoder.encode(audio) # Encode
+sample = autoencoder.decode(latent, num_steps=10) # Decode by sampling diffusion model conditioning on latent

Modules/diffusion/sampler.py

@@ -0,0 +1,691 @@
+from math import atan, cos, pi, sin, sqrt
+from typing import Any, Callable, List, Optional, Tuple, Type
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, reduce
+from torch import Tensor
+
+from .utils import *
+
+"""
+Diffusion Training
+"""
+
+""" Distributions """
+
+
+class Distribution:
+    def __call__(self, num_samples: int, device: torch.device):
+        raise NotImplementedError()
+
+
+class LogNormalDistribution(Distribution):
+    def __init__(self, mean: float, std: float):
+        self.mean = mean
+        self.std = std
+
+    def __call__(
+        self, num_samples: int, device: torch.device = torch.device("cpu")
+    ) -> Tensor:
+        normal = self.mean + self.std * torch.randn((num_samples,), device=device)
+        return normal.exp()
+
+
+class UniformDistribution(Distribution):
+    def __call__(self, num_samples: int, device: torch.device = torch.device("cpu")):
+        return torch.rand(num_samples, device=device)
+
+
+class VKDistribution(Distribution):
+    def __init__(
+        self,
+        min_value: float = 0.0,
+        max_value: float = float("inf"),
+        sigma_data: float = 1.0,
+    ):
+        self.min_value = min_value
+        self.max_value = max_value
+        self.sigma_data = sigma_data
+
+    def __call__(
+        self, num_samples: int, device: torch.device = torch.device("cpu")
+    ) -> Tensor:
+        sigma_data = self.sigma_data
+        min_cdf = atan(self.min_value / sigma_data) * 2 / pi
+        max_cdf = atan(self.max_value / sigma_data) * 2 / pi
+        u = (max_cdf - min_cdf) * torch.randn((num_samples,), device=device) + min_cdf
+        return torch.tan(u * pi / 2) * sigma_data
+
+
+""" Diffusion Classes """
+
+
+def pad_dims(x: Tensor, ndim: int) -> Tensor:
+    # Pads additional ndims to the right of the tensor
+    return x.view(*x.shape, *((1,) * ndim))
+
+
+def clip(x: Tensor, dynamic_threshold: float = 0.0):
+    if dynamic_threshold == 0.0:
+        return x.clamp(-1.0, 1.0)
+    else:
+        # Dynamic thresholding
+        # Find dynamic threshold quantile for each batch
+        x_flat = rearrange(x, "b ... -> b (...)")
+        scale = torch.quantile(x_flat.abs(), dynamic_threshold, dim=-1)
+        # Clamp to a min of 1.0
+        scale.clamp_(min=1.0)
+        # Clamp all values and scale
+        scale = pad_dims(scale, ndim=x.ndim - scale.ndim)
+        x = x.clamp(-scale, scale) / scale
+        return x
+
+
+def to_batch(
+    batch_size: int,
+    device: torch.device,
+    x: Optional[float] = None,
+    xs: Optional[Tensor] = None,
+) -> Tensor:
+    assert exists(x) ^ exists(xs), "Either x or xs must be provided"
+    # If x provided use the same for all batch items
+    if exists(x):
+        xs = torch.full(size=(batch_size,), fill_value=x).to(device)
+    assert exists(xs)
+    return xs
+
+
+class Diffusion(nn.Module):
+
+    alias: str = ""
+
+    """Base diffusion class"""
+
+    def denoise_fn(
+        self,
+        x_noisy: Tensor,
+        sigmas: Optional[Tensor] = None,
+        sigma: Optional[float] = None,
+        **kwargs,
+    ) -> Tensor:
+        raise NotImplementedError("Diffusion class missing denoise_fn")
+
+    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
+        raise NotImplementedError("Diffusion class missing forward function")
+
+
+class VDiffusion(Diffusion):
+
+    alias = "v"
+
+    def __init__(self, net: nn.Module, *, sigma_distribution: Distribution):
+        super().__init__()
+        self.net = net
+        self.sigma_distribution = sigma_distribution
+
+    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
+        angle = sigmas * pi / 2
+        alpha = torch.cos(angle)
+        beta = torch.sin(angle)
+        return alpha, beta
+
+    def denoise_fn(
+        self,
+        x_noisy: Tensor,
+        sigmas: Optional[Tensor] = None,
+        sigma: Optional[float] = None,
+        **kwargs,
+    ) -> Tensor:
+        batch_size, device = x_noisy.shape[0], x_noisy.device
+        sigmas = to_batch(x=sigma, xs=sigmas, batch_size=batch_size, device=device)
+        return self.net(x_noisy, sigmas, **kwargs)
+
+    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
+        batch_size, device = x.shape[0], x.device
+
+        # Sample amount of noise to add for each batch element
+        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
+        sigmas_padded = rearrange(sigmas, "b -> b 1 1")
+
+        # Get noise
+        noise = default(noise, lambda: torch.randn_like(x))
+
+        # Combine input and noise weighted by half-circle
+        alpha, beta = self.get_alpha_beta(sigmas_padded)
+        x_noisy = x * alpha + noise * beta
+        x_target = noise * alpha - x * beta
+
+        # Denoise and return loss
+        x_denoised = self.denoise_fn(x_noisy, sigmas, **kwargs)
+        return F.mse_loss(x_denoised, x_target)
+
+
+class KDiffusion(Diffusion):
+    """Elucidated Diffusion (Karras et al. 2022): https://arxiv.org/abs/2206.00364"""
+
+    alias = "k"
+
+    def __init__(
+        self,
+        net: nn.Module,
+        *,
+        sigma_distribution: Distribution,
+        sigma_data: float,  # data distribution standard deviation
+        dynamic_threshold: float = 0.0,
+    ):
+        super().__init__()
+        self.net = net
+        self.sigma_data = sigma_data
+        self.sigma_distribution = sigma_distribution
+        self.dynamic_threshold = dynamic_threshold
+
+    def get_scale_weights(self, sigmas: Tensor) -> Tuple[Tensor, ...]:
+        sigma_data = self.sigma_data
+        c_noise = torch.log(sigmas) * 0.25
+        sigmas = rearrange(sigmas, "b -> b 1 1")
+        c_skip = (sigma_data ** 2) / (sigmas ** 2 + sigma_data ** 2)
+        c_out = sigmas * sigma_data * (sigma_data ** 2 + sigmas ** 2) ** -0.5
+        c_in = (sigmas ** 2 + sigma_data ** 2) ** -0.5
+        return c_skip, c_out, c_in, c_noise
+
+    def denoise_fn(
+        self,
+        x_noisy: Tensor,
+        sigmas: Optional[Tensor] = None,
+        sigma: Optional[float] = None,
+        **kwargs,
+    ) -> Tensor:
+        batch_size, device = x_noisy.shape[0], x_noisy.device
+        sigmas = to_batch(x=sigma, xs=sigmas, batch_size=batch_size, device=device)
+
+        # Predict network output and add skip connection
+        c_skip, c_out, c_in, c_noise = self.get_scale_weights(sigmas)
+        x_pred = self.net(c_in * x_noisy, c_noise, **kwargs)
+        x_denoised = c_skip * x_noisy + c_out * x_pred
+
+        return x_denoised
+
+    def loss_weight(self, sigmas: Tensor) -> Tensor:
+        # Computes weight depending on data distribution
+        return (sigmas ** 2 + self.sigma_data ** 2) * (sigmas * self.sigma_data) ** -2
+
+    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
+        batch_size, device = x.shape[0], x.device
+        from einops import rearrange, reduce
+
+        # Sample amount of noise to add for each batch element
+        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
+        sigmas_padded = rearrange(sigmas, "b -> b 1 1")
+
+        # Add noise to input
+        noise = default(noise, lambda: torch.randn_like(x))
+        x_noisy = x + sigmas_padded * noise
+        
+        # Compute denoised values
+        x_denoised = self.denoise_fn(x_noisy, sigmas=sigmas, **kwargs)
+
+        # Compute weighted loss
+        losses = F.mse_loss(x_denoised, x, reduction="none")
+        losses = reduce(losses, "b ... -> b", "mean")
+        losses = losses * self.loss_weight(sigmas)
+        loss = losses.mean()
+        return loss
+
+
+class VKDiffusion(Diffusion):
+
+    alias = "vk"
+
+    def __init__(self, net: nn.Module, *, sigma_distribution: Distribution):
+        super().__init__()
+        self.net = net
+        self.sigma_distribution = sigma_distribution
+
+    def get_scale_weights(self, sigmas: Tensor) -> Tuple[Tensor, ...]:
+        sigma_data = 1.0
+        sigmas = rearrange(sigmas, "b -> b 1 1")
+        c_skip = (sigma_data ** 2) / (sigmas ** 2 + sigma_data ** 2)
+        c_out = -sigmas * sigma_data * (sigma_data ** 2 + sigmas ** 2) ** -0.5
+        c_in = (sigmas ** 2 + sigma_data ** 2) ** -0.5
+        return c_skip, c_out, c_in
+
+    def sigma_to_t(self, sigmas: Tensor) -> Tensor:
+        return sigmas.atan() / pi * 2
+
+    def t_to_sigma(self, t: Tensor) -> Tensor:
+        return (t * pi / 2).tan()
+
+    def denoise_fn(
+        self,
+        x_noisy: Tensor,
+        sigmas: Optional[Tensor] = None,
+        sigma: Optional[float] = None,
+        **kwargs,
+    ) -> Tensor:
+        batch_size, device = x_noisy.shape[0], x_noisy.device
+        sigmas = to_batch(x=sigma, xs=sigmas, batch_size=batch_size, device=device)
+
+        # Predict network output and add skip connection
+        c_skip, c_out, c_in = self.get_scale_weights(sigmas)
+        x_pred = self.net(c_in * x_noisy, self.sigma_to_t(sigmas), **kwargs)
+        x_denoised = c_skip * x_noisy + c_out * x_pred
+        return x_denoised
+
+    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
+        batch_size, device = x.shape[0], x.device
+
+        # Sample amount of noise to add for each batch element
+        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
+        sigmas_padded = rearrange(sigmas, "b -> b 1 1")
+
+        # Add noise to input
+        noise = default(noise, lambda: torch.randn_like(x))
+        x_noisy = x + sigmas_padded * noise
+
+        # Compute model output
+        c_skip, c_out, c_in = self.get_scale_weights(sigmas)
+        x_pred = self.net(c_in * x_noisy, self.sigma_to_t(sigmas), **kwargs)
+
+        # Compute v-objective target
+        v_target = (x - c_skip * x_noisy) / (c_out + 1e-7)
+
+        # Compute loss
+        loss = F.mse_loss(x_pred, v_target)
+        return loss
+
+
+"""
+Diffusion Sampling
+"""
+
+""" Schedules """
+
+
+class Schedule(nn.Module):
+    """Interface used by different sampling schedules"""
+
+    def forward(self, num_steps: int, device: torch.device) -> Tensor:
+        raise NotImplementedError()
+
+
+class LinearSchedule(Schedule):
+    def forward(self, num_steps: int, device: Any) -> Tensor:
+        sigmas = torch.linspace(1, 0, num_steps + 1)[:-1]
+        return sigmas
+
+
+class KarrasSchedule(Schedule):
+    """https://arxiv.org/abs/2206.00364 equation 5"""
+
+    def __init__(self, sigma_min: float, sigma_max: float, rho: float = 7.0):
+        super().__init__()
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.rho = rho
+
+    def forward(self, num_steps: int, device: Any) -> Tensor:
+        rho_inv = 1.0 / self.rho
+        steps = torch.arange(num_steps, device=device, dtype=torch.float32)
+        sigmas = (
+            self.sigma_max ** rho_inv
+            + (steps / (num_steps - 1))
+            * (self.sigma_min ** rho_inv - self.sigma_max ** rho_inv)
+        ) ** self.rho
+        sigmas = F.pad(sigmas, pad=(0, 1), value=0.0)
+        return sigmas
+
+
+""" Samplers """
+
+
+class Sampler(nn.Module):
+
+    diffusion_types: List[Type[Diffusion]] = []
+
+    def forward(
+        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
+    ) -> Tensor:
+        raise NotImplementedError()
+
+    def inpaint(
+        self,
+        source: Tensor,
+        mask: Tensor,
+        fn: Callable,
+        sigmas: Tensor,
+        num_steps: int,
+        num_resamples: int,
+    ) -> Tensor:
+        raise NotImplementedError("Inpainting not available with current sampler")
+
+
+class VSampler(Sampler):
+
+    diffusion_types = [VDiffusion]
+
+    def get_alpha_beta(self, sigma: float) -> Tuple[float, float]:
+        angle = sigma * pi / 2
+        alpha = cos(angle)
+        beta = sin(angle)
+        return alpha, beta
+
+    def forward(
+        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
+    ) -> Tensor:
+        x = sigmas[0] * noise
+        alpha, beta = self.get_alpha_beta(sigmas[0].item())
+
+        for i in range(num_steps - 1):
+            is_last = i == num_steps - 1
+
+            x_denoised = fn(x, sigma=sigmas[i])
+            x_pred = x * alpha - x_denoised * beta
+            x_eps = x * beta + x_denoised * alpha
+
+            if not is_last:
+                alpha, beta = self.get_alpha_beta(sigmas[i + 1].item())
+                x = x_pred * alpha + x_eps * beta
+
+        return x_pred
+
+
+class KarrasSampler(Sampler):
+    """https://arxiv.org/abs/2206.00364 algorithm 1"""
+
+    diffusion_types = [KDiffusion, VKDiffusion]
+
+    def __init__(
+        self,
+        s_tmin: float = 0,
+        s_tmax: float = float("inf"),
+        s_churn: float = 0.0,
+        s_noise: float = 1.0,
+    ):
+        super().__init__()
+        self.s_tmin = s_tmin
+        self.s_tmax = s_tmax
+        self.s_noise = s_noise
+        self.s_churn = s_churn
+
+    def step(
+        self, x: Tensor, fn: Callable, sigma: float, sigma_next: float, gamma: float
+    ) -> Tensor:
+        """Algorithm 2 (step)"""
+        # Select temporarily increased noise level
+        sigma_hat = sigma + gamma * sigma
+        # Add noise to move from sigma to sigma_hat
+        epsilon = self.s_noise * torch.randn_like(x)
+        x_hat = x + sqrt(sigma_hat ** 2 - sigma ** 2) * epsilon
+        # Evaluate ∂x/∂sigma at sigma_hat
+        d = (x_hat - fn(x_hat, sigma=sigma_hat)) / sigma_hat
+        # Take euler step from sigma_hat to sigma_next
+        x_next = x_hat + (sigma_next - sigma_hat) * d
+        # Second order correction
+        if sigma_next != 0:
+            model_out_next = fn(x_next, sigma=sigma_next)
+            d_prime = (x_next - model_out_next) / sigma_next
+            x_next = x_hat + 0.5 * (sigma - sigma_hat) * (d + d_prime)
+        return x_next
+
+    def forward(
+        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
+    ) -> Tensor:
+        x = sigmas[0] * noise
+        # Compute gammas
+        gammas = torch.where(
+            (sigmas >= self.s_tmin) & (sigmas <= self.s_tmax),
+            min(self.s_churn / num_steps, sqrt(2) - 1),
+            0.0,
+        )
+        # Denoise to sample
+        for i in range(num_steps - 1):
+            x = self.step(
+                x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1], gamma=gammas[i]  # type: ignore # noqa
+            )
+
+        return x
+
+
+class AEulerSampler(Sampler):
+
+    diffusion_types = [KDiffusion, VKDiffusion]
+
+    def get_sigmas(self, sigma: float, sigma_next: float) -> Tuple[float, float]:
+        sigma_up = sqrt(sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2)
+        sigma_down = sqrt(sigma_next ** 2 - sigma_up ** 2)
+        return sigma_up, sigma_down
+
+    def step(self, x: Tensor, fn: Callable, sigma: float, sigma_next: float) -> Tensor:
+        # Sigma steps
+        sigma_up, sigma_down = self.get_sigmas(sigma, sigma_next)
+        # Derivative at sigma (∂x/∂sigma)
+        d = (x - fn(x, sigma=sigma)) / sigma
+        # Euler method
+        x_next = x + d * (sigma_down - sigma)
+        # Add randomness
+        x_next = x_next + torch.randn_like(x) * sigma_up
+        return x_next
+
+    def forward(
+        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
+    ) -> Tensor:
+        x = sigmas[0] * noise
+        # Denoise to sample
+        for i in range(num_steps - 1):
+            x = self.step(x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1])  # type: ignore # noqa
+        return x
+
+
+class ADPM2Sampler(Sampler):
+    """https://www.desmos.com/calculator/jbxjlqd9mb"""
+
+    diffusion_types = [KDiffusion, VKDiffusion]
+
+    def __init__(self, rho: float = 1.0):
+        super().__init__()
+        self.rho = rho
+
+    def get_sigmas(self, sigma: float, sigma_next: float) -> Tuple[float, float, float]:
+        r = self.rho
+        sigma_up = sqrt(sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2)
+        sigma_down = sqrt(sigma_next ** 2 - sigma_up ** 2)
+        sigma_mid = ((sigma ** (1 / r) + sigma_down ** (1 / r)) / 2) ** r
+        return sigma_up, sigma_down, sigma_mid
+
+    def step(self, x: Tensor, fn: Callable, sigma: float, sigma_next: float) -> Tensor:
+        # Sigma steps
+        sigma_up, sigma_down, sigma_mid = self.get_sigmas(sigma, sigma_next)
+        # Derivative at sigma (∂x/∂sigma)
+        d = (x - fn(x, sigma=sigma)) / sigma
+        # Denoise to midpoint
+        x_mid = x + d * (sigma_mid - sigma)
+        # Derivative at sigma_mid (∂x_mid/∂sigma_mid)
+        d_mid = (x_mid - fn(x_mid, sigma=sigma_mid)) / sigma_mid
+        # Denoise to next
+        x = x + d_mid * (sigma_down - sigma)
+        # Add randomness
+        x_next = x + torch.randn_like(x) * sigma_up
+        return x_next
+
+    def forward(
+        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
+    ) -> Tensor:
+        x = sigmas[0] * noise
+        # Denoise to sample
+        for i in range(num_steps - 1):
+            x = self.step(x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1])  # type: ignore # noqa
+        return x
+
+    def inpaint(
+        self,
+        source: Tensor,
+        mask: Tensor,
+        fn: Callable,
+        sigmas: Tensor,
+        num_steps: int,
+        num_resamples: int,
+    ) -> Tensor:
+        x = sigmas[0] * torch.randn_like(source)
+
+        for i in range(num_steps - 1):
+            # Noise source to current noise level
+            source_noisy = source + sigmas[i] * torch.randn_like(source)
+            for r in range(num_resamples):
+                # Merge noisy source and current then denoise
+                x = source_noisy * mask + x * ~mask
+                x = self.step(x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1])  # type: ignore # noqa
+                # Renoise if not last resample step
+                if r < num_resamples - 1:
+                    sigma = sqrt(sigmas[i] ** 2 - sigmas[i + 1] ** 2)
+                    x = x + sigma * torch.randn_like(x)
+
+        return source * mask + x * ~mask
+
+
+""" Main Classes """
+
+
+class DiffusionSampler(nn.Module):
+    def __init__(
+        self,
+        diffusion: Diffusion,
+        *,
+        sampler: Sampler,
+        sigma_schedule: Schedule,
+        num_steps: Optional[int] = None,
+        clamp: bool = True,
+    ):
+        super().__init__()
+        self.denoise_fn = diffusion.denoise_fn
+        self.sampler = sampler
+        self.sigma_schedule = sigma_schedule
+        self.num_steps = num_steps
+        self.clamp = clamp
+
+        # Check sampler is compatible with diffusion type
+        sampler_class = sampler.__class__.__name__
+        diffusion_class = diffusion.__class__.__name__
+        message = f"{sampler_class} incompatible with {diffusion_class}"
+        assert diffusion.alias in [t.alias for t in sampler.diffusion_types], message
+
+    def forward(
+        self, noise: Tensor, num_steps: Optional[int] = None, **kwargs
+    ) -> Tensor:
+        device = noise.device
+        num_steps = default(num_steps, self.num_steps)  # type: ignore
+        assert exists(num_steps), "Parameter `num_steps` must be provided"
+        # Compute sigmas using schedule
+        sigmas = self.sigma_schedule(num_steps, device)
+        # Append additional kwargs to denoise function (used e.g. for conditional unet)
+        fn = lambda *a, **ka: self.denoise_fn(*a, **{**ka, **kwargs})  # noqa
+        # Sample using sampler
+        x = self.sampler(noise, fn=fn, sigmas=sigmas, num_steps=num_steps)
+        x = x.clamp(-1.0, 1.0) if self.clamp else x
+        return x
+
+
+class DiffusionInpainter(nn.Module):
+    def __init__(
+        self,
+        diffusion: Diffusion,
+        *,
+        num_steps: int,
+        num_resamples: int,
+        sampler: Sampler,
+        sigma_schedule: Schedule,
+    ):
+        super().__init__()
+        self.denoise_fn = diffusion.denoise_fn
+        self.num_steps = num_steps
+        self.num_resamples = num_resamples
+        self.inpaint_fn = sampler.inpaint
+        self.sigma_schedule = sigma_schedule
+
+    @torch.no_grad()
+    def forward(self, inpaint: Tensor, inpaint_mask: Tensor) -> Tensor:
+        x = self.inpaint_fn(
+            source=inpaint,
+            mask=inpaint_mask,
+            fn=self.denoise_fn,
+            sigmas=self.sigma_schedule(self.num_steps, inpaint.device),
+            num_steps=self.num_steps,
+            num_resamples=self.num_resamples,
+        )
+        return x
+
+
+def sequential_mask(like: Tensor, start: int) -> Tensor:
+    length, device = like.shape[2], like.device
+    mask = torch.ones_like(like, dtype=torch.bool)
+    mask[:, :, start:] = torch.zeros((length - start,), device=device)
+    return mask
+
+
+class SpanBySpanComposer(nn.Module):
+    def __init__(
+        self,
+        inpainter: DiffusionInpainter,
+        *,
+        num_spans: int,
+    ):
+        super().__init__()
+        self.inpainter = inpainter
+        self.num_spans = num_spans
+
+    def forward(self, start: Tensor, keep_start: bool = False) -> Tensor:
+        half_length = start.shape[2] // 2
+
+        spans = list(start.chunk(chunks=2, dim=-1)) if keep_start else []
+        # Inpaint second half from first half
+        inpaint = torch.zeros_like(start)
+        inpaint[:, :, :half_length] = start[:, :, half_length:]
+        inpaint_mask = sequential_mask(like=start, start=half_length)
+
+        for i in range(self.num_spans):
+            # Inpaint second half
+            span = self.inpainter(inpaint=inpaint, inpaint_mask=inpaint_mask)
+            # Replace first half with generated second half
+            second_half = span[:, :, half_length:]
+            inpaint[:, :, :half_length] = second_half
+            # Save generated span
+            spans.append(second_half)
+
+        return torch.cat(spans, dim=2)
+
+
+class XDiffusion(nn.Module):
+    def __init__(self, type: str, net: nn.Module, **kwargs):
+        super().__init__()
+
+        diffusion_classes = [VDiffusion, KDiffusion, VKDiffusion]
+        aliases = [t.alias for t in diffusion_classes]  # type: ignore
+        message = f"type='{type}' must be one of {*aliases,}"
+        assert type in aliases, message
+        self.net = net
+
+        for XDiffusion in diffusion_classes:
+            if XDiffusion.alias == type:  # type: ignore
+                self.diffusion = XDiffusion(net=net, **kwargs)
+
+    def forward(self, *args, **kwargs) -> Tensor:
+        return self.diffusion(*args, **kwargs)
+
+    def sample(
+        self,
+        noise: Tensor,
+        num_steps: int,
+        sigma_schedule: Schedule,
+        sampler: Sampler,
+        clamp: bool,
+        **kwargs,
+    ) -> Tensor:
+        diffusion_sampler = DiffusionSampler(
+            diffusion=self.diffusion,
+            sampler=sampler,
+            sigma_schedule=sigma_schedule,
+            num_steps=num_steps,
+            clamp=clamp,
+        )
+        return diffusion_sampler(noise, **kwargs)

Modules/diffusion/utils.py

@@ -0,0 +1,82 @@
+from functools import reduce
+from inspect import isfunction
+from math import ceil, floor, log2, pi
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Generator, Tensor
+from typing_extensions import TypeGuard
+
+T = TypeVar("T")
+
+
+def exists(val: Optional[T]) -> TypeGuard[T]:
+    return val is not None
+
+
+def iff(condition: bool, value: T) -> Optional[T]:
+    return value if condition else None
+
+
+def is_sequence(obj: T) -> TypeGuard[Union[list, tuple]]:
+    return isinstance(obj, list) or isinstance(obj, tuple)
+
+
+def default(val: Optional[T], d: Union[Callable[..., T], T]) -> T:
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def to_list(val: Union[T, Sequence[T]]) -> List[T]:
+    if isinstance(val, tuple):
+        return list(val)
+    if isinstance(val, list):
+        return val
+    return [val]  # type: ignore
+
+
+def prod(vals: Sequence[int]) -> int:
+    return reduce(lambda x, y: x * y, vals)
+
+
+def closest_power_2(x: float) -> int:
+    exponent = log2(x)
+    distance_fn = lambda z: abs(x - 2 ** z)  # noqa
+    exponent_closest = min((floor(exponent), ceil(exponent)), key=distance_fn)
+    return 2 ** int(exponent_closest)
+
+def rand_bool(shape, proba, device = None):
+    if proba == 1:
+        return torch.ones(shape, device=device, dtype=torch.bool)
+    elif proba == 0:
+        return torch.zeros(shape, device=device, dtype=torch.bool)
+    else:
+        return torch.bernoulli(torch.full(shape, proba, device=device)).to(torch.bool)
+
+
+"""
+Kwargs Utils
+"""
+
+
+def group_dict_by_prefix(prefix: str, d: Dict) -> Tuple[Dict, Dict]:
+    return_dicts: Tuple[Dict, Dict] = ({}, {})
+    for key in d.keys():
+        no_prefix = int(not key.startswith(prefix))
+        return_dicts[no_prefix][key] = d[key]
+    return return_dicts
+
+
+def groupby(prefix: str, d: Dict, keep_prefix: bool = False) -> Tuple[Dict, Dict]:
+    kwargs_with_prefix, kwargs = group_dict_by_prefix(prefix, d)
+    if keep_prefix:
+        return kwargs_with_prefix, kwargs
+    kwargs_no_prefix = {k[len(prefix) :]: v for k, v in kwargs_with_prefix.items()}
+    return kwargs_no_prefix, kwargs
+
+
+def prefix_dict(prefix: str, d: Dict) -> Dict:
+    return {prefix + str(k): v for k, v in d.items()}

Modules/discriminators.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.nn import Conv1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, spectral_norm
+
+from .utils import get_padding
+
+LRELU_SLOPE = 0.1
+
+def stft(x, fft_size, hop_size, win_length, window):
+    """Perform STFT and convert to magnitude spectrogram.
+    Args:
+        x (Tensor): Input signal tensor (B, T).
+        fft_size (int): FFT size.
+        hop_size (int): Hop size.
+        win_length (int): Window length.
+        window (str): Window function type.
+    Returns:
+        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
+    """
+    x_stft = torch.stft(x, fft_size, hop_size, win_length, window,
+            return_complex=True)
+    real = x_stft[..., 0]
+    imag = x_stft[..., 1]
+
+    return torch.abs(x_stft).transpose(2, 1)
+
+class SpecDiscriminator(nn.Module):
+    """docstring for Discriminator."""
+
+    def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window", use_spectral_norm=False):
+        super(SpecDiscriminator, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.fft_size = fft_size
+        self.shift_size = shift_size
+        self.win_length = win_length
+        self.window = getattr(torch, window)(win_length)
+        self.discriminators = nn.ModuleList([
+            norm_f(nn.Conv2d(1, 32, kernel_size=(3, 9), padding=(1, 4))),
+            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1,2), padding=(1, 4))),
+            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1,2), padding=(1, 4))),
+            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1,2), padding=(1, 4))),
+            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 3), stride=(1,1), padding=(1, 1))),
+        ])
+
+        self.out = norm_f(nn.Conv2d(32, 1, 3, 1, 1))
+
+    def forward(self, y):
+
+        fmap = []
+        y = y.squeeze(1)
+        y = stft(y, self.fft_size, self.shift_size, self.win_length, self.window.to(y.get_device()))
+        y = y.unsqueeze(1)
+        for i, d in enumerate(self.discriminators):
+            y = d(y)
+            y = F.leaky_relu(y, LRELU_SLOPE)
+            fmap.append(y)
+
+        y = self.out(y)
+        fmap.append(y)
+
+        return torch.flatten(y, 1, -1), fmap
+
+class MultiResSpecDiscriminator(torch.nn.Module):
+
+    def __init__(self,
+                 fft_sizes=[1024, 2048, 512],
+                 hop_sizes=[120, 240, 50],
+                 win_lengths=[600, 1200, 240],
+                 window="hann_window"):
+
+        super(MultiResSpecDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            SpecDiscriminator(fft_sizes[0], hop_sizes[0], win_lengths[0], window),
+            SpecDiscriminator(fft_sizes[1], hop_sizes[1], win_lengths[1], window),
+            SpecDiscriminator(fft_sizes[2], hop_sizes[2], win_lengths[2], window)
+            ])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorP(torch.nn.Module):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        super(DiscriminatorP, self).__init__()
+        self.period = period
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
+            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
+        ])
+        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0: # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super(MultiPeriodDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            DiscriminatorP(2),
+            DiscriminatorP(3),
+            DiscriminatorP(5),
+            DiscriminatorP(7),
+            DiscriminatorP(11),
+        ])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+    
+class WavLMDiscriminator(nn.Module):
+    """docstring for Discriminator."""
+
+    def __init__(self, slm_hidden=768, 
+                 slm_layers=13, 
+                 initial_channel=64, 
+                 use_spectral_norm=False):
+        super(WavLMDiscriminator, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.pre = norm_f(Conv1d(slm_hidden * slm_layers, initial_channel, 1, 1, padding=0))
+        
+        self.convs = nn.ModuleList([
+            norm_f(nn.Conv1d(initial_channel, initial_channel * 2, kernel_size=5, padding=2)),
+            norm_f(nn.Conv1d(initial_channel * 2, initial_channel * 4, kernel_size=5, padding=2)),
+            norm_f(nn.Conv1d(initial_channel * 4, initial_channel * 4, 5, 1, padding=2)),
+        ])
+
+        self.conv_post = norm_f(Conv1d(initial_channel * 4, 1, 3, 1, padding=1))
+        
+    def forward(self, x):
+        x = self.pre(x)
+        
+        fmap = []
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            fmap.append(x)
+        x = self.conv_post(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x