Upload 66 files

Dockerfile

@@ -0,0 +1,27 @@
+# 使用 Python 3.9 作为基础镜像
+FROM python:3.9
+
+# 添加用户
+RUN useradd -m -u 1000 user
+
+# 设置工作目录
+WORKDIR /app
+
+# 切换到 root 用户以安装系统依赖
+USER root
+RUN apt-get update && apt-get install -y rubberband-cli
+
+# 切回到普通用户
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+# 复制并安装 Python 依赖
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+# 复制应用文件
+COPY --chown=user . /app
+
+# 启动 Gradio 应用，假设应用入口文件为 app.py
+# CMD ["python", "app.py"]
+CMD ["python", "app_chat.py"]

app.py

@@ -0,0 +1,107 @@
+import torch
+import numpy as np
+from tqdm import tqdm
+from model.DiffSynthSampler import DiffSynthSampler
+import soundfile as sf
+# import pyrubberband as pyrb
+from tqdm import tqdm
+from model.VQGAN import get_VQGAN
+from model.diffusion import get_diffusion_model
+from transformers import AutoTokenizer, ClapModel
+from model.diffusion_components import linear_beta_schedule
+from model.timbre_encoder_pretrain import get_timbre_encoder
+from model.multimodal_model import get_multi_modal_model
+
+
+
+import gradio as gr
+from webUI.natural_language_guided.gradio_webUI import GradioWebUI
+from webUI.natural_language_guided.text2sound import get_text2sound_module
+from webUI.natural_language_guided.sound2sound_with_text import get_sound2sound_with_text_module
+from webUI.natural_language_guided.inpaint_with_text import get_inpaint_with_text_module
+from webUI.natural_language_guided.build_instrument import get_build_instrument_module
+from webUI.natural_language_guided.README import get_readme_module
+
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+use_pretrained_CLAP = False
+
+# load VQ-GAN
+VAE_model_name = "24_1_2024-52_4x_L_D"
+modelConfig = {"in_channels": 3, "hidden_channels": [80, 160], "embedding_dim": 4, "out_channels": 3, "block_depth": 2,
+               "attn_pos":  [80, 160], "attn_with_skip": True,
+            "num_embeddings": 8192, "commitment_cost": 0.25, "decay": 0.99,
+            "norm_type": "groupnorm", "act_type": "swish", "num_groups": 16}
+VAE = get_VQGAN(modelConfig, load_pretrain=True, model_name=VAE_model_name, device=device)
+
+# load U-Net
+UNet_model_name = "history/28_1_2024_CLAP_STFT_180000" if use_pretrained_CLAP else "history/28_1_2024_TE_STFT_300000"
+unetConfig = {"in_dim": 4, "down_dims": [96, 96, 192, 384], "up_dims": [384, 384, 192, 96], "attn_type": "linear_add", "condition_type": "natural_language_prompt", "label_emb_dim": 512}
+uNet = get_diffusion_model(unetConfig, load_pretrain=True, model_name=UNet_model_name, device=device)
+
+# load LM
+CLAP_temp = ClapModel.from_pretrained("laion/clap-htsat-unfused")  # 153,492,890
+CLAP_tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")
+
+timbre_encoder_name = "24_1_2024_STFT"
+timbre_encoder_Config = {"input_dim": 512, "feature_dim": 512, "hidden_dim": 1024, "num_instrument_classes": 1006, "num_instrument_family_classes": 11, "num_velocity_classes": 128, "num_qualities": 10, "num_layers": 3}
+timbre_encoder = get_timbre_encoder(timbre_encoder_Config, load_pretrain=True, model_name=timbre_encoder_name, device=device)
+
+if use_pretrained_CLAP:
+  text_encoder = CLAP_temp
+else:
+  multimodalmodel_name = "24_1_2024"
+  multimodalmodel_config = {"text_feature_dim": 512, "spectrogram_feature_dim": 1024, "multi_modal_emb_dim": 512, "num_projection_layers": 2,
+                "temperature": 1.0, "dropout": 0.1, "freeze_text_encoder": False, "freeze_spectrogram_encoder": False}
+  mmm = get_multi_modal_model(timbre_encoder, CLAP_temp, multimodalmodel_config, load_pretrain=True, model_name=multimodalmodel_name, device=device)
+
+  text_encoder = mmm.to("cpu")
+
+
+
+
+
+
+gradioWebUI = GradioWebUI(device, VAE, uNet, text_encoder, CLAP_tokenizer, freq_resolution=512, time_resolution=256, channels=4, timesteps=1000, squared=False,
+                          VAE_scale=4, flexible_duration=True, noise_strategy="repeat", GAN_generator=None)
+
+with gr.Blocks(theme=gr.themes.Soft(), mode="dark") as demo:
+# with gr.Blocks(theme='WeixuanYuan/Soft_dark', mode="dark") as demo:
+    gr.Markdown("DiffuSynth v0.2")
+
+    reconstruction_state = gr.State(value={})
+    text2sound_state = gr.State(value={})
+    sound2sound_state = gr.State(value={})
+    inpaint_state = gr.State(value={})
+    super_resolution_state = gr.State(value={})
+    virtual_instruments_state = gr.State(value={"virtual_instruments": {}})
+
+    get_text2sound_module(gradioWebUI, text2sound_state, virtual_instruments_state)
+    get_sound2sound_with_text_module(gradioWebUI, sound2sound_state, virtual_instruments_state)
+    get_inpaint_with_text_module(gradioWebUI, inpaint_state, virtual_instruments_state)
+    get_build_instrument_module(gradioWebUI, virtual_instruments_state)
+    get_readme_module()
+
+demo.launch(debug=True, share=True)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

app_chat.py

@@ -0,0 +1,7 @@
+import gradio as gr
+
+def greet(name):
+    return "Hello " + name + "!!"
+
+demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+demo.launch()

metrics/FD.py

@@ -0,0 +1,293 @@
+import json
+import os
+
+import librosa
+import numpy as np
+import torch
+from tqdm import tqdm
+from scipy.linalg import sqrtm
+
+from metrics.pipelines import sample_pipeline, sample_pipeline_GAN
+from metrics.pipelines_STFT import sample_pipeline_STFT, sample_pipeline_GAN_STFT
+from tools import rms_normalize
+
+
+def ASTaudio2feature(device, signal, processor, AST, sampling_rate):
+    # audio file is decoded on the fly
+    inputs = processor(signal, sampling_rate=sampling_rate, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = AST(**inputs)
+
+    last_hidden_states = outputs.last_hidden_state[:, 0, :].to("cpu").detach().numpy()
+    return last_hidden_states
+
+
+# 计算两个numpy数组的均值和协方差矩阵
+def calculate_statistics(features):
+    mu = np.mean(features, axis=0)
+    sigma = np.cov(features, rowvar=False)
+    return mu, sigma
+
+
+# 计算FID
+def calculate_fid(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    # 在协方差矩阵对角线上添加一个小的正值
+    sigma1 += np.eye(sigma1.shape[0]) * eps
+    sigma2 += np.eye(sigma2.shape[0]) * eps
+
+    ssdiff = np.sum((mu1 - mu2) ** 2.0)
+    covmean = sqrtm(sigma1.dot(sigma2))
+
+    # 由于数值问题，有时可能会得到复数，只取实部
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
+
+    fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)
+    return fid
+
+
+# 计算FID
+def calculate_fid_dict(dict1, dict2, eps=1e-6):
+    # 在协方差矩阵对角线上添加一个小的正值
+    mu1, sigma1 = dict1["mu"], dict1["sigma"]
+    mu2, sigma2 = dict2["mu"], dict2["sigma"]
+    sigma1 += np.eye(sigma1.shape[0]) * eps
+    sigma2 += np.eye(sigma2.shape[0]) * eps
+
+    ssdiff = np.sum((mu1 - mu2) ** 2.0)
+    covmean = sqrtm(sigma1.dot(sigma2))
+
+    # 由于数值问题，有时可能会得到复数，只取实部
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
+
+    fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)
+    return fid
+
+
+# Todo: AudioLDM
+# def generate_features_with_AudioLDM_and_AST(device, processor, AST, AudioLDM_signals_directory_path, return_feature=False):
+
+#     diffuSynth_features = []
+
+#     # Step 1: Load all wav files in AudioLDM_signals_directory_path
+#     AudioLDM_signals = []
+#     signal_lengths = set()
+
+#     for file_name in os.listdir(AudioLDM_signals_directory_path):
+#         if file_name.endswith('.wav'):
+#             file_path = os.path.join(AudioLDM_signals_directory_path, file_name)
+#             signal, sr = librosa.load(file_path, sr=16000)  # Load audio file with sampling rate 16000
+#             # Normalize
+#             AudioLDM_signals.append(rms_normalize(signal))
+#             signal_lengths.add(len(signal))
+
+#     # Step 2: Check if all signals have the same length
+#     if len(signal_lengths) != 1:
+#         raise ValueError("Not all signals have the same length. Please ensure all audio files are of the same length.")
+
+#     # Step 3: Reshape to signal_batches [number_batches, batch_size=8, signal_length]
+#     batch_size = 8
+#     signal_length = signal_lengths.pop()  # All lengths are the same, get one of them
+
+#     # Create batches
+#     signal_batches = [AudioLDM_signals[i:i + batch_size] for i in range(0, len(AudioLDM_signals), batch_size)]
+
+#     for signal_batch in tqdm(signal_batches):
+
+#         features = ASTaudio2feature(device, signal_batch, processor, AST, sampling_rate=16000)
+#         diffuSynth_features.extend(features)
+
+#     if return_feature:
+#         return diffuSynth_features
+#     else:
+#         mu, sigma = calculate_statistics(diffuSynth_features)
+#         return {"mu": mu,  "sigma": sigma}
+
+def generate_features_with_AudioLDM_and_AST(device, processor, AST, AudioLDM_signals_directory_path, return_feature=False):
+
+    diffuSynth_features = []
+
+    # Step 1: Load all wav files in AudioLDM_signals_directory_path
+    AudioLDM_signals = []
+    signal_lengths = set()
+    target_length = 4 * 16000  # 4 seconds * 16000 samples per second
+
+    for file_name in os.listdir(AudioLDM_signals_directory_path):
+        if file_name.endswith('.wav') and not file_name.startswith('._'):
+            file_path = os.path.join(AudioLDM_signals_directory_path, file_name)
+            try:
+                signal, sr = librosa.load(file_path, sr=16000)  # Load audio file with sampling rate 16000
+                if len(signal) >= target_length:
+                    signal = signal[:target_length]  # Take only the first 4 seconds
+                else:
+                    raise ValueError(f"The file {file_name} is shorter than 4 seconds.")
+                # Normalize
+                AudioLDM_signals.append(rms_normalize(signal))
+                signal_lengths.add(len(signal))
+            except Exception as e:
+                print(f"Error loading {file_name}: {e}")
+
+    # Step 2: Check if all signals have the same length
+    if len(signal_lengths) != 1:
+        raise ValueError("Not all signals have the same length. Please ensure all audio files are of the same length.")
+
+    # Step 3: Reshape to signal_batches [number_batches, batch_size=8, signal_length]
+    batch_size = 8
+    signal_length = signal_lengths.pop()  # All lengths are the same, get one of them
+
+    # Create batches
+    signal_batches = [AudioLDM_signals[i:i + batch_size] for i in range(0, len(AudioLDM_signals), batch_size)]
+
+    for signal_batch in tqdm(signal_batches):
+        features = ASTaudio2feature(device, signal_batch, processor, AST, sampling_rate=16000)
+        diffuSynth_features.extend(features)
+
+    if return_feature:
+        return diffuSynth_features
+    else:
+        mu, sigma = calculate_statistics(diffuSynth_features)
+        return {"mu": mu, "sigma": sigma}
+
+
+
+
+def generate_features_with_diffuSynth_and_AST(device, uNet, VAE, mmm, CLAP_tokenizer, processor, AST, num_batches,
+                                              positive_prompts, negative_prompts="", CFG=1, sample_steps=10, task="spectrograms", return_feature=False):
+    diffuSynth_features = []
+
+    if task == "spectrograms":
+        pipe = sample_pipeline
+    elif task == "STFT":
+        pipe = sample_pipeline_STFT
+    else:
+        raise NotImplementedError
+
+    for _ in tqdm(range(num_batches)):
+        quantized_latent_representations, reconstruction_batch, signals = pipe(device, uNet, VAE, mmm,
+                                                                                          CLAP_tokenizer,
+                                                                                          positive_prompts=positive_prompts,
+                                                                                          negative_prompts=negative_prompts,
+                                                                                          batchsize=8,
+                                                                                          sample_steps=sample_steps,
+                                                                                          CFG=CFG, seed=None,
+                                                                                          return_latent=False)
+
+        features = ASTaudio2feature(device, signals, processor, AST, sampling_rate=16000)
+        diffuSynth_features.extend(features)
+
+    if return_feature:
+        return diffuSynth_features
+    else:
+        mu, sigma = calculate_statistics(diffuSynth_features)
+        return {"mu": mu,  "sigma": sigma}
+
+
+def generate_features_with_GAN_and_AST(device, gan_generator, VAE, mmm, CLAP_tokenizer, processor, AST, num_batches,
+                                              positive_prompts, negative_prompts="", CFG=1, sample_steps=10, task="spectrograms", return_feature=False):
+    diffuSynth_features = []
+
+    if task == "spectrograms":
+        pipe = sample_pipeline_GAN
+    elif task == "STFT":
+        pipe = sample_pipeline_GAN_STFT
+    else:
+        raise NotImplementedError
+
+    for _ in tqdm(range(num_batches)):
+        quantized_latent_representations, reconstruction_batch, signals = pipe(device, gan_generator, VAE, mmm,
+                                                                                          CLAP_tokenizer,
+                                                                                          positive_prompts=positive_prompts,
+                                                                                          negative_prompts=negative_prompts,
+                                                                                          batchsize=8,
+                                                                                          sample_steps=sample_steps,
+                                                                                          CFG=CFG, seed=None,
+                                                                                          return_latent=False)
+
+        features = ASTaudio2feature(device, signals, processor, AST, sampling_rate=16000)
+        diffuSynth_features.extend(features)
+
+    if return_feature:
+        return diffuSynth_features
+    else:
+        mu, sigma = calculate_statistics(diffuSynth_features)
+        return {"mu": mu,  "sigma": sigma}
+
+
+def get_FD(train_features, device, uNet, VAE, mmm, CLAP_tokenizer, processor, AST, num_batches, positive_prompts,
+           negative_prompts="", CFG=1, sample_steps=10):
+    diffuSynth_features = generate_features_with_diffuSynth_and_AST(device, uNet, VAE, mmm, CLAP_tokenizer, processor,
+                                                                    AST, num_batches, positive_prompts,
+                                                                    negative_prompts=negative_prompts, CFG=CFG,
+                                                                    sample_steps=sample_steps)
+
+    mu_real, sigma_real = calculate_statistics(train_features)
+    mu_gen, sigma_gen = calculate_statistics(diffuSynth_features)
+
+    fid_score = calculate_fid(mu_real, sigma_real, mu_gen, sigma_gen)
+    print('FID score:', fid_score)
+
+
+def get_fid_score(feature1, features2):
+    mu_real, sigma_real = calculate_statistics(feature1)
+    mu_gen, sigma_gen = calculate_statistics(features2)
+
+    fid_score = calculate_fid(mu_real, sigma_real, mu_gen, sigma_gen)
+    # print('FID score:', fid_score)
+    return fid_score
+
+
+def calculate_fid_matrix(features_list_1, features_list_2, get_fid_score):
+    # 初始化一个矩阵来存储FID分数
+    # 矩阵的大小为 len(features_list_1) x len(features_list_2)
+    fid_scores = [[0 for _ in range(len(features_list_2))] for _ in range(len(features_list_1))]
+
+    # 遍历两个列表，并计算每一对特征集合的FID分数
+    for i, feature1 in enumerate(features_list_1):
+        for j, feature2 in enumerate(features_list_2):
+            fid_scores[i][j] = get_fid_score(feature1, feature2)
+
+    return fid_scores
+
+
+def save_AST_feature(key, mu, sigma, path='results/AST_metric/pre_calculated_features/AST_features.json'):
+    # 尝试打开并读取现有的JSON文件
+    try:
+        with open(path, 'r') as file:
+            data = json.load(file)
+    except FileNotFoundError:
+        # 如果文件不存在，创建一个新的字典
+        data = {}
+
+    if isinstance(mu, np.ndarray):
+        mu = mu.tolist()
+    if isinstance(sigma, np.ndarray):
+        sigma = sigma.tolist()
+
+    # 添加新数据
+    data[key] = {"mu": mu, "sigma": sigma}
+
+    # 将更新后的数据写回文件
+    with open(path, 'w') as file:
+        json.dump(data, file, indent=4)
+
+
+def read_AST_features(path='results/AST_metric/pre_calculated_features/AST_features.json'):
+    try:
+        # 尝试打开并读取JSON文件
+        with open(path, 'r') as file:
+            AST_features = json.load(file)
+
+            for AST_feature_name in AST_features.keys():
+                AST_features[AST_feature_name]["mu"] = np.array(AST_features[AST_feature_name]["mu"])
+                AST_features[AST_feature_name]["sigma"] = np.array(AST_features[AST_feature_name]["sigma"])
+
+            return AST_features
+    except FileNotFoundError:
+        # 如果文件不存在，返回一个空字典
+        print(f"文件 {path} 未找到.")
+        return {}
+    except json.JSONDecodeError:
+        # 如果文件不是有效的JSON，返回一个空字典
+        print(f"文件 {path} 不是有效的JSON格式.")
+        return {}

metrics/IS.py

@@ -0,0 +1,218 @@
+import os
+
+import librosa
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from metrics.pipelines import sample_pipeline, inpaint_pipeline, sample_pipeline_GAN
+from metrics.pipelines_STFT import sample_pipeline_STFT, sample_pipeline_GAN_STFT
+from tools import rms_normalize, pad_STFT, encode_stft
+from webUI.natural_language_guided.utils import InputBatch2Encode_STFT
+
+def get_inception_score_for_AudioLDM(device, timbre_encoder, VAE, AudioLDM_signals_directory_path):
+    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+
+    diffuSynth_probabilities = []
+
+    # Step 1: Load all wav files in AudioLDM_signals_directory_path
+    AudioLDM_signals = []
+    signal_lengths = set()
+    target_length = 4 * 16000  # 4 seconds * 16000 samples per second
+
+    for file_name in os.listdir(AudioLDM_signals_directory_path):
+        if file_name.endswith('.wav') and not file_name.startswith('._'):
+            file_path = os.path.join(AudioLDM_signals_directory_path, file_name)
+            signal, sr = librosa.load(file_path, sr=16000)  # Load audio file with sampling rate 16000
+            if len(signal) >= target_length:
+                signal = signal[:target_length]  # Take only the first 4 seconds
+            else:
+                raise ValueError(f"The file {file_name} is shorter than 4 seconds.")
+            # Normalize
+            AudioLDM_signals.append(rms_normalize(signal))
+            signal_lengths.add(len(signal))
+
+    # Step 2: Check if all signals have the same length
+    if len(signal_lengths) != 1:
+        raise ValueError("Not all signals have the same length. Please ensure all audio files are of the same length.")
+
+    encoded_audios = []
+    for origin_audio in AudioLDM_signals:
+        D = librosa.stft(origin_audio, n_fft=1024, hop_length=256, win_length=1024)
+        padded_D = pad_STFT(D)
+        encoded_D = encode_stft(padded_D)
+        encoded_audios.append(encoded_D)
+    encoded_audios_np = np.array(encoded_audios)
+    origin_spectrogram_batch_tensor = torch.from_numpy(encoded_audios_np).float().to(device)
+
+    # Step 3: Reshape to signal_batches [number_batches, batch_size=8, signal_length]
+    batch_size = 8
+    num_batches = int(np.ceil(origin_spectrogram_batch_tensor.shape[0] / batch_size))
+    spectrogram_batches = []
+    for i in range(num_batches):
+        batch = origin_spectrogram_batch_tensor[i * batch_size:(i + 1) * batch_size]
+        spectrogram_batches.append(batch)
+
+    for spectrogram_batch in tqdm(spectrogram_batches):
+        spectrogram_batch = spectrogram_batch.to(device)
+        _, _, _, _, quantized_latent_representations = InputBatch2Encode_STFT(VAE_encoder, spectrogram_batch, quantizer=VAE_quantizer, squared=False)
+        quantized_latent_representations = quantized_latent_representations
+        feature, instrument_logits, instrument_family_logits, velocity_logits, qualities = timbre_encoder(quantized_latent_representations)
+        probabilities = torch.nn.functional.softmax(instrument_logits, dim=1)
+
+        diffuSynth_probabilities.extend(probabilities.to("cpu").detach().numpy())
+
+    return inception_score(np.array(diffuSynth_probabilities))
+
+
+# def get_inception_score_for_AudioLDM(device, timbre_encoder, VAE, AudioLDM_signals_directory_path):
+#     VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+#
+#     diffuSynth_probabilities = []
+#
+#     # Step 1: Load all wav files in AudioLDM_signals_directory_path
+#     AudioLDM_signals = []
+#     signal_lengths = set()
+#
+#     for file_name in os.listdir(AudioLDM_signals_directory_path):
+#         if file_name.endswith('.wav'):
+#             file_path = os.path.join(AudioLDM_signals_directory_path, file_name)
+#             signal, sr = librosa.load(file_path, sr=16000)  # Load audio file with sampling rate 16000
+#             # Normalize
+#             AudioLDM_signals.append(rms_normalize(signal))
+#             signal_lengths.add(len(signal))
+#
+#     # Step 2: Check if all signals have the same length
+#     if len(signal_lengths) != 1:
+#         raise ValueError("Not all signals have the same length. Please ensure all audio files are of the same length.")
+#
+#     encoded_audios = []
+#     for origin_audio in AudioLDM_signals:
+#         D = librosa.stft(origin_audio, n_fft=1024, hop_length=256, win_length=1024)
+#         padded_D = pad_STFT(D)
+#         encoded_D = encode_stft(padded_D)
+#         encoded_audios.append(encoded_D)
+#     encoded_audios_np = np.array(encoded_audios)
+#     origin_spectrogram_batch_tensor = torch.from_numpy(encoded_audios_np).float().to(device)
+#
+#
+#     # Step 3: Reshape to signal_batches [number_batches, batch_size=8, signal_length]
+#     batch_size = 8
+#     num_batches = int(np.ceil(origin_spectrogram_batch_tensor.shape[0] / batch_size))
+#     spectrogram_batches = []
+#     for i in range(num_batches):
+#         batch = origin_spectrogram_batch_tensor[i * batch_size:(i + 1) * batch_size]
+#         spectrogram_batches.append(batch)
+#
+#
+#     for spectrogram_batch in tqdm(spectrogram_batches):
+#         spectrogram_batch = spectrogram_batch.to(device)
+#         _, _, _, _, quantized_latent_representations = InputBatch2Encode_STFT(VAE_encoder, spectrogram_batch, quantizer=VAE_quantizer,squared=False)
+#         quantized_latent_representations = quantized_latent_representations
+#         feature, instrument_logits, instrument_family_logits, velocity_logits, qualities = timbre_encoder(quantized_latent_representations)
+#         probabilities = torch.nn.functional.softmax(instrument_logits, dim=1)
+#
+#         diffuSynth_probabilities.extend(probabilities.to("cpu").detach().numpy())
+#
+#     return inception_score(np.array(diffuSynth_probabilities))
+
+
+def get_inception_score(device, uNet, VAE, MMM, CLAP_tokenizer, timbre_encoder, num_batches, positive_prompts, negative_prompts="", CFG=1, sample_steps=10, task="spectrograms"):
+    diffuSynth_probabilities = []
+
+    if task == "spectrograms":
+        pipe = sample_pipeline
+    elif task == "STFT":
+        pipe = sample_pipeline_STFT
+    else:
+        raise NotImplementedError
+
+    for _ in tqdm(range(num_batches)):
+        quantized_latent_representations = pipe(device, uNet, VAE, MMM, CLAP_tokenizer,
+                                                           positive_prompts=positive_prompts, negative_prompts=negative_prompts,
+                                                      batchsize=8, sample_steps=sample_steps, CFG=CFG, seed=None)
+
+        quantized_latent_representations = quantized_latent_representations.to(device)
+        feature, instrument_logits, instrument_family_logits, velocity_logits, qualities = timbre_encoder(quantized_latent_representations)
+        probabilities = torch.nn.functional.softmax(instrument_logits, dim=1)
+
+        diffuSynth_probabilities.extend(probabilities.to("cpu").detach().numpy())
+
+    return inception_score(np.array(diffuSynth_probabilities))
+
+
+def get_inception_score_GAN(device, gan_generator, VAE, MMM, CLAP_tokenizer, timbre_encoder, num_batches, positive_prompts, negative_prompts="", CFG=1, sample_steps=10, task="spectrograms"):
+    diffuSynth_probabilities = []
+
+    if task == "spectrograms":
+        pipe = sample_pipeline_GAN
+    elif task == "STFT":
+        pipe = sample_pipeline_GAN_STFT
+    else:
+        raise NotImplementedError
+
+    for _ in tqdm(range(num_batches)):
+        quantized_latent_representations = pipe(device, gan_generator, VAE, MMM, CLAP_tokenizer,
+                                                           positive_prompts=positive_prompts, negative_prompts=negative_prompts,
+                                                      batchsize=8, sample_steps=sample_steps, CFG=CFG, seed=None)
+
+        quantized_latent_representations = quantized_latent_representations.to(device)
+        feature, instrument_logits, instrument_family_logits, velocity_logits, qualities = timbre_encoder(quantized_latent_representations)
+        probabilities = torch.nn.functional.softmax(instrument_logits, dim=1)
+
+        diffuSynth_probabilities.extend(probabilities.to("cpu").detach().numpy())
+
+    return inception_score(np.array(diffuSynth_probabilities))
+
+
+def predict_qualities_with_diffuSynth_sample(device, uNet, VAE, MMM, CLAP_tokenizer, timbre_encoder, num_batches, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):
+    diffuSynth_qualities = []
+    for _ in tqdm(range(num_batches)):
+        quantized_latent_representations = sample_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
+                                                           positive_prompts=positive_prompts, negative_prompts=negative_prompts,
+                                                      batchsize=8, sample_steps=sample_steps, CFG=CFG, seed=None)
+
+        quantized_latent_representations = quantized_latent_representations.to(device)
+        feature, instrument_logits, instrument_family_logits, velocity_logits, qualities = timbre_encoder(quantized_latent_representations)
+        qualities = qualities.to("cpu").detach().numpy()
+        # qualities = np.where(qualities > 0.5, 1, 0)
+
+        diffuSynth_qualities.extend(qualities)
+
+    return np.mean(diffuSynth_qualities, axis=0)
+
+
+def generate_probabilities_with_diffuSynth_inpaint(device, uNet, VAE, MMM, CLAP_tokenizer, timbre_encoder, num_batches, guidance, duration, use_dynamic_mask, noising_strength, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):
+
+    inpaint_probabilities, signals = [], []
+    for _ in tqdm(range(num_batches)):
+        quantized_latent_representations, _, rec_signals = inpaint_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
+                                                                            use_dynamic_mask=use_dynamic_mask, noising_strength=noising_strength, guidance=guidance,
+                    positive_prompts=positive_prompts, negative_prompts=negative_prompts, batchsize=8, sample_steps=sample_steps, CFG=CFG, seed=None, duration=duration, mask_flexivity=0.999,
+                    return_latent=False)
+
+        quantized_latent_representations = quantized_latent_representations.to(device)
+        feature, instrument_logits, instrument_family_logits, velocity_logits, qualities = timbre_encoder(quantized_latent_representations)
+        probabilities = torch.nn.functional.softmax(instrument_logits, dim=1)
+
+        inpaint_probabilities.extend(probabilities.to("cpu").detach().numpy())
+        signals.extend(rec_signals)
+
+    return np.array(inpaint_probabilities), signals
+
+
+def inception_score(pred):
+
+    # 计算每个图像的条件概率分布 P(y|x)
+    pyx = pred / np.sum(pred, axis=1, keepdims=True)
+
+    # 计算整个数据集的边缘概率分布 P(y)
+    py = np.mean(pyx, axis=0, keepdims=True)
+
+    # 计算KL散度
+    kl_div = pyx * (np.log(pyx + 1e-11) - np.log(py + 1e-11))
+
+    # 对所有图像求和并平均
+    kl_div_sum = np.sum(kl_div, axis=1)
+    score = np.exp(np.mean(kl_div_sum))
+    return score

metrics/P_C_T.py

@@ -0,0 +1,12 @@
+import numpy as np
+from metrics.precision_recall import knn_precision_recall_features
+
+
+# 生成样本
+real_features = np.random.normal(0, 1, size=(1600, 512))
+generated_features = np.random.normal(0, 1, size=(1600, 512))
+
+state = knn_precision_recall_features(real_features, generated_features, nhood_sizes=[1, 2, 3, 4, 5, 10],
+                                  row_batch_size=16, col_batch_size=16)
+
+print(state)

metrics/get_reference_AST_features.py

@@ -0,0 +1,63 @@
+import json
+import librosa
+import numpy as np
+from tqdm import tqdm
+from metrics.FD import ASTaudio2feature, calculate_statistics, save_AST_feature
+from tools import rms_normalize
+from transformers import AutoProcessor, ASTModel
+
+device = "cpu"
+processor = AutoProcessor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
+AST = ASTModel.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593").to(device)
+
+
+data_split = "train"
+with open(f'data/NSynth/{data_split}_examples.json') as f:
+    data = json.load(f)
+
+def read_signal(note_str):
+    y, sr = librosa.load(f"data/NSynth/nsynth-{data_split}-52/audio/{note_str}.wav", sr=16000)
+    if len(y) >= 64000:
+        y = y[:64000]
+    else:
+        y_extend = [0.0] * 64000
+        y_extend[:len(y)] = y
+        y = y_extend
+
+    return rms_normalize(y)
+
+for quality in ["bright", "dark", "distortion", "fast_decay", "long_release", "multiphonic", "nonlinear_env", "percussive", "reverb", "tempo-synced"]:
+    features = []
+    for i, (note_str, attributes) in tqdm(enumerate(data.items())):
+        if not attributes["pitch"] == 52:
+            continue
+        if not (quality in attributes['qualities_str']):
+            continue
+
+        signal = read_signal(note_str)
+        feature_for_one_signal = ASTaudio2feature(device, [signal], processor, AST, sampling_rate=16000)[0]
+        features.append(feature_for_one_signal)
+
+    mu, sigma = calculate_statistics(features)
+    print(np.shape(mu))
+    print(np.shape(sigma))
+
+    save_AST_feature(f'{data_split}_{quality}', mu.tolist(), sigma.tolist())
+
+for instrument_name in ["bass", "brass", "flute", "guitar", "keyboard", "mallet", "organ", "reed", "string", "synth_lead", "vocal"]:
+    features = []
+    for i, (note_str, attributes) in tqdm(enumerate(data.items())):
+        if not attributes["pitch"] == 52:
+            continue
+        if not (attributes["instrument_family_str"] == instrument_name):
+            continue
+
+        signal = read_signal(note_str)
+        feature_for_one_signal = ASTaudio2feature(device, [signal], processor, AST, sampling_rate=16000)[0]
+        features.append(feature_for_one_signal)
+
+    mu, sigma = calculate_statistics(features)
+    print(np.shape(mu))
+    print(np.shape(sigma))
+
+    save_AST_feature(f'{data_split}_{instrument_name}', mu.tolist(), sigma.tolist())

metrics/pipelines.py

@@ -0,0 +1,144 @@
+import librosa
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from tools import VAE_out_put_to_spc, rms_normalize, nnData2Audio
+from model.DiffSynthSampler import DiffSynthSampler
+
+def sample_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
+                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, duration=3.0,
+                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):
+
+    height = int(freq_resolution/VAE_scale)
+    width = int(time_resolution/VAE_scale)
+    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+
+    text2sound_embedding = \
+        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device)
+    negative_condition = \
+        MMM.get_text_features(**CLAP_tokenizer([negative_prompts], padding=True, return_tensors="pt"))[0].to(device)
+
+    mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy, mute=True)
+    mySampler.activate_classifier_free_guidance(CFG, negative_condition)
+
+    mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32)))
+
+    condition = text2sound_embedding.repeat(batchsize, 1)
+
+    latent_representations, initial_noise = \
+    mySampler.sample(model=uNet, shape=(batchsize, channels, height, width), seed=seed,
+                      return_tensor=True, condition=condition, sampler=sampler)
+
+    latent_representations = latent_representations[-1]
+
+    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)
+
+    if return_latent:
+        return quantized_latent_representations.detach()
+    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
+    time_resolution = int(time_resolution * ((duration+1) / 4))
+
+    rec_signals = nnData2Audio(reconstruction_batch, resolution=(freq_resolution, time_resolution))
+    rec_signals = [rms_normalize(rec_signal) for rec_signal in rec_signals]
+
+    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals
+
+def sample_pipeline_GAN(device, gan_generator, VAE, MMM, CLAP_tokenizer,
+                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, duration=3.0,
+                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):
+
+    height = int(freq_resolution/VAE_scale)
+    width = int(time_resolution/VAE_scale)
+    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+
+    text2sound_embedding = \
+        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device)
+
+    condition = text2sound_embedding.repeat(batchsize, 1)
+
+    noise = torch.randn(batchsize, channels, height, width).to(device)
+    latent_representations = gan_generator(noise, condition)
+
+    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)
+
+    if return_latent:
+        return quantized_latent_representations.detach()
+    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
+    time_resolution = int(time_resolution * ((duration+1) / 4))
+
+    rec_signals = nnData2Audio(reconstruction_batch, resolution=(freq_resolution, time_resolution))
+    rec_signals = [rms_normalize(rec_signal) for rec_signal in rec_signals]
+
+    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals
+
+def inpaint_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer, use_dynamic_mask, noising_strength, guidance,
+                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, duration=3.0, mask_flexivity=0.99,
+                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):
+
+    height = int(freq_resolution/VAE_scale)
+    width = int(time_resolution * ((duration + 1) / 4) / VAE_scale)
+    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+
+
+    text2sound_embedding = \
+        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0]
+    negative_condition = \
+        MMM.get_text_features(**CLAP_tokenizer([negative_prompts], padding=True, return_tensors="pt"))[0]
+
+
+    mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy, mute=True)
+    mySampler.activate_classifier_free_guidance(CFG, negative_condition)
+    mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32)))
+
+    condition = text2sound_embedding.repeat(batchsize, 1)
+    guidance = guidance.repeat(batchsize, 1, 1, 1).to(device)
+
+    # mask = 1, freeze
+    latent_mask = torch.zeros((batchsize, 1, height, width), dtype=torch.float32).to(device)
+    latent_mask[:, :, :, -int(time_resolution * (1 / 4) / VAE_scale):] = 1.0
+
+    latent_representations, initial_noise = \
+        mySampler.inpaint_sample(model=uNet, shape=(batchsize, channels, height, width),
+                                  noising_strength=noising_strength,
+                                  guide_img=guidance, mask=latent_mask, return_tensor=True,
+                                  condition=condition, sampler=sampler,
+                                  use_dynamic_mask=use_dynamic_mask,
+                                  end_noise_level_ratio=0.0,
+                                  mask_flexivity=mask_flexivity)
+
+    latent_representations = latent_representations[-1]
+
+    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)
+
+    if return_latent:
+        return quantized_latent_representations.detach()
+    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
+    time_resolution = int(time_resolution * ((duration+1) / 4))
+
+    rec_signals = nnData2Audio(reconstruction_batch, resolution=(freq_resolution, time_resolution))
+    rec_signals = [rms_normalize(rec_signal) for rec_signal in rec_signals]
+
+    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals
+
+
+def generate_audios_with_diffuSynth_sample(device, uNet, VAE, MMM, CLAP_tokenizer, num_batches, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):
+    diffuSynth_signals = []
+    for _ in tqdm(range(num_batches)):
+        _, _, signals = sample_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
+                                        positive_prompts=positive_prompts, negative_prompts=negative_prompts,
+                                                      batchsize=16, sample_steps=sample_steps, CFG=CFG, seed=None, return_latent=False)
+        diffuSynth_signals.extend(signals)
+    return np.array(diffuSynth_signals)
+
+
+def generate_audios_with_diffuSynth_inpaint(device, uNet, VAE, MMM, CLAP_tokenizer, num_batches, guidance, duration, use_dynamic_mask, noising_strength, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):
+
+    diffuSynth_signals = []
+    for _ in tqdm(range(num_batches)):
+        _, _, signals = inpaint_pipeline(device, uNet, VAE, MMM, CLAP_tokenizer,
+                                         use_dynamic_mask=use_dynamic_mask, noising_strength=noising_strength, guidance=guidance,
+                    positive_prompts=positive_prompts, negative_prompts=negative_prompts, batchsize=16, sample_steps=sample_steps, CFG=CFG, seed=None, duration=duration, mask_flexivity=0.999,
+                    return_latent=False)
+        diffuSynth_signals.extend(signals)
+    return np.array(diffuSynth_signals)

metrics/pipelines_STFT.py

@@ -0,0 +1,100 @@
+import librosa
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from tools import rms_normalize, decode_stft, depad_STFT
+from model.DiffSynthSampler import DiffSynthSampler
+
+def sample_pipeline_STFT(device, uNet, VAE, MMM, CLAP_tokenizer,
+                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None,
+                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):
+    "Sample a fix-length audio using a diffusion model, including 'ISTFT+' post-processing."
+
+    height = int(freq_resolution/VAE_scale)
+    width = int(time_resolution/VAE_scale)
+    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+
+    text2sound_embedding = \
+        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device)
+    negative_condition = \
+        MMM.get_text_features(**CLAP_tokenizer([negative_prompts], padding=True, return_tensors="pt"))[
+            0].to(device)
+
+    mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy, mute=True)
+    mySampler.activate_classifier_free_guidance(CFG, negative_condition)
+
+    mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32)))
+
+    condition = text2sound_embedding.repeat(batchsize, 1)
+
+    latent_representations, initial_noise = \
+    mySampler.sample(model=uNet, shape=(batchsize, channels, height, width), seed=seed,
+                      return_tensor=True, condition=condition, sampler=sampler)
+
+    latent_representations = latent_representations[-1]
+
+    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)
+
+    if return_latent:
+        return quantized_latent_representations.detach()
+
+    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
+
+    rec_signals = []
+
+    for index, STFT in enumerate(reconstruction_batch):
+        padded_D_rec = decode_stft(STFT)
+        D_rec = depad_STFT(padded_D_rec)
+        # get_audio
+        rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024)
+        rec_signals.append(rms_normalize(rec_signal))
+
+    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals
+
+def sample_pipeline_GAN_STFT(device, gan_generator, VAE, MMM, CLAP_tokenizer,
+                    positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None,
+                    freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True):
+    "Sample fix-length audio using a GAN, including 'ISTFT+' post-processing."
+
+    height = int(freq_resolution/VAE_scale)
+    width = int(time_resolution/VAE_scale)
+    VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+
+    text2sound_embedding = \
+        MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device)
+
+    condition = text2sound_embedding.repeat(batchsize, 1)
+
+    noise = torch.randn(batchsize, channels, height, width).to(device)
+    latent_representations = gan_generator(noise, condition)
+
+    quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations)
+
+    if return_latent:
+        return quantized_latent_representations.detach()
+    reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy()
+
+    rec_signals = []
+
+    for index, STFT in enumerate(reconstruction_batch):
+        padded_D_rec = decode_stft(STFT)
+        D_rec = depad_STFT(padded_D_rec)
+        # get_audio
+        rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024)
+        rec_signals.append(rms_normalize(rec_signal))
+
+    return quantized_latent_representations.detach(), reconstruction_batch, rec_signals
+
+
+def generate_audios_with_diffuSynth_sample(device, uNet, VAE, MMM, CLAP_tokenizer, num_batches, positive_prompts, negative_prompts="", CFG=6, sample_steps=10):
+    "Sample audios using a diffusion model, including 'ISTFT+' post-processing."
+
+    diffuSynth_signals = []
+    for _ in tqdm(range(num_batches)):
+        _, _, signals = sample_pipeline_STFT(device, uNet, VAE, MMM, CLAP_tokenizer,
+                                        positive_prompts=positive_prompts, negative_prompts=negative_prompts,
+                                                      batchsize=8, sample_steps=sample_steps, CFG=CFG, seed=None, return_latent=False)
+        diffuSynth_signals.extend(signals)
+    return np.array(diffuSynth_signals)
+

metrics/precision_recall.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial
+# 4.0 International License. To view a copy of this license, visit
+# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
+# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
+
+"""k-NN precision and recall."""
+
+from time import time
+
+
+# ----------------------------------------------------------------------------
+
+import numpy as np
+from tqdm import tqdm
+
+
+def batch_pairwise_distances(U, V):
+    """Compute pair-wise distance in a batch of feature."""
+
+    norm_u = np.sum(np.square(U), axis=1)
+    norm_v = np.sum(np.square(V), axis=1)
+
+    norm_u = np.reshape(norm_u, [-1, 1])
+    norm_v = np.reshape(norm_v, [1, -1])
+
+    D = np.maximum(norm_u - 2 * np.dot(U, V.T) + norm_v, 0.0)
+    return D
+
+
+# ----------------------------------------------------------------------------
+
+class DistanceBlock():
+    """Compute pair-wise distance in a batch of feature."""
+
+    def __init__(self, num_features):
+        self.num_features = num_features
+
+    def pairwise_distances(self, U, V):
+        return batch_pairwise_distances(U, V)
+
+
+
+# ----------------------------------------------------------------------------
+
+class ManifoldEstimator():
+    """Estimates the manifold of given feature vectors."""
+
+    def __init__(self, distance_block, features, row_batch_size=16, col_batch_size=16,
+                 nhood_sizes=[3], clamp_to_percentile=None, eps=1e-5, mute=False):
+        """Estimate the manifold of given feature vectors.
+
+            Args:
+                distance_block: DistanceBlock object that distributes pairwise distance
+                    calculation to multiple GPUs.
+                features (np.array/tf.Tensor): Matrix of feature vectors to estimate their manifold.
+                row_batch_size (int): Row batch size to compute pairwise distances
+                    (parameter to trade-off between memory usage and performance).
+                col_batch_size (int): Column batch size to compute pairwise distances.
+                nhood_sizes (list): Number of neighbors used to estimate the manifold.
+                clamp_to_percentile (float): Prune hyperspheres that have radius larger than
+                    the given percentile.
+                eps (float): Small number for numerical stability.
+        """
+        num_images = features.shape[0]
+        self.nhood_sizes = nhood_sizes
+        self.num_nhoods = len(nhood_sizes)
+        self.eps = eps
+        self.row_batch_size = row_batch_size
+        self.col_batch_size = col_batch_size
+        self._ref_features = features
+        self._distance_block = distance_block
+        self.mute = mute
+
+        # Estimate manifold of features by calculating distances to k-NN of each sample.
+        self.D = np.zeros([num_images, self.num_nhoods], dtype=np.float32)
+        distance_batch = np.zeros([row_batch_size, num_images], dtype=np.float32)
+        seq = np.arange(max(self.nhood_sizes) + 1, dtype=np.int32)
+
+        if mute:
+            for begin1 in range(0, num_images, row_batch_size):
+                end1 = min(begin1 + row_batch_size, num_images)
+                row_batch = features[begin1:end1]
+
+                for begin2 in range(0, num_images, col_batch_size):
+                    end2 = min(begin2 + col_batch_size, num_images)
+                    col_batch = features[begin2:end2]
+
+                    # Compute distances between batches.
+                    distance_batch[0:end1 - begin1, begin2:end2] = self._distance_block.pairwise_distances(row_batch,
+                                                                                                           col_batch)
+
+                # Find the k-nearest neighbor from the current batch.
+                self.D[begin1:end1, :] = np.partition(distance_batch[0:end1 - begin1, :], seq, axis=1)[:, self.nhood_sizes]
+        else:
+            for begin1 in tqdm(range(0, num_images, row_batch_size)):
+                end1 = min(begin1 + row_batch_size, num_images)
+                row_batch = features[begin1:end1]
+
+                for begin2 in range(0, num_images, col_batch_size):
+                    end2 = min(begin2 + col_batch_size, num_images)
+                    col_batch = features[begin2:end2]
+
+                    # Compute distances between batches.
+                    distance_batch[0:end1 - begin1, begin2:end2] = self._distance_block.pairwise_distances(row_batch,
+                                                                                                           col_batch)
+
+                # Find the k-nearest neighbor from the current batch.
+                self.D[begin1:end1, :] = np.partition(distance_batch[0:end1 - begin1, :], seq, axis=1)[:, self.nhood_sizes]
+
+        if clamp_to_percentile is not None:
+            max_distances = np.percentile(self.D, clamp_to_percentile, axis=0)
+            self.D[self.D > max_distances] = 0
+
+    def evaluate(self, eval_features, return_realism=False, return_neighbors=False):
+        """Evaluate if new feature vectors are at the manifold."""
+        num_eval_images = eval_features.shape[0]
+        num_ref_images = self.D.shape[0]
+        distance_batch = np.zeros([self.row_batch_size, num_ref_images], dtype=np.float32)
+        batch_predictions = np.zeros([num_eval_images, self.num_nhoods], dtype=np.int32)
+        max_realism_score = np.zeros([num_eval_images, ], dtype=np.float32)
+        nearest_indices = np.zeros([num_eval_images, ], dtype=np.int32)
+
+        for begin1 in range(0, num_eval_images, self.row_batch_size):
+            end1 = min(begin1 + self.row_batch_size, num_eval_images)
+            feature_batch = eval_features[begin1:end1]
+
+            for begin2 in range(0, num_ref_images, self.col_batch_size):
+                end2 = min(begin2 + self.col_batch_size, num_ref_images)
+                ref_batch = self._ref_features[begin2:end2]
+
+                distance_batch[0:end1 - begin1, begin2:end2] = self._distance_block.pairwise_distances(feature_batch,
+                                                                                                       ref_batch)
+
+            # From the minibatch of new feature vectors, determine if they are in the estimated manifold.
+            # If a feature vector is inside a hypersphere of some reference sample, then
+            # the new sample lies at the estimated manifold.
+            # The radii of the hyperspheres are determined from distances of neighborhood size k.
+            samples_in_manifold = distance_batch[0:end1 - begin1, :, None] <= self.D
+            batch_predictions[begin1:end1] = np.any(samples_in_manifold, axis=1).astype(np.int32)
+
+            max_realism_score[begin1:end1] = np.max(self.D[:, 0] / (distance_batch[0:end1 - begin1, :] + self.eps),
+                                                    axis=1)
+            nearest_indices[begin1:end1] = np.argmin(distance_batch[0:end1 - begin1, :], axis=1)
+
+        if return_realism and return_neighbors:
+            return batch_predictions, max_realism_score, nearest_indices
+        elif return_realism:
+            return batch_predictions, max_realism_score
+        elif return_neighbors:
+            return batch_predictions, nearest_indices
+
+        return batch_predictions
+
+
+# ----------------------------------------------------------------------------
+
+def knn_precision_recall_features(ref_features, eval_features, nhood_sizes=[3],
+                                  row_batch_size=10000, col_batch_size=50000, mute=False):
+    """Calculates k-NN precision and recall for two sets of feature vectors.
+
+        Args:
+            ref_features (np.array/tf.Tensor): Feature vectors of reference images.
+            eval_features (np.array/tf.Tensor): Feature vectors of generated images.
+            nhood_sizes (list): Number of neighbors used to estimate the manifold.
+            row_batch_size (int): Row batch size to compute pairwise distances
+                (parameter to trade-off between memory usage and performance).
+            col_batch_size (int): Column batch size to compute pairwise distances.
+            num_gpus (int): Number of GPUs used to evaluate precision and recall.
+
+        Returns:
+            State (dict): Dict that contains precision and recall calculated from
+            ref_features and eval_features.
+    """
+    state = dict()
+    num_images = ref_features.shape[0]
+    num_features = ref_features.shape[1]
+
+    # Initialize DistanceBlock and ManifoldEstimators.
+    distance_block = DistanceBlock(num_features)
+    ref_manifold = ManifoldEstimator(distance_block, ref_features, row_batch_size, col_batch_size, nhood_sizes, mute=mute)
+    eval_manifold = ManifoldEstimator(distance_block, eval_features, row_batch_size, col_batch_size, nhood_sizes, mute=mute)
+
+    # Evaluate precision and recall using k-nearest neighbors.
+    if not mute:
+        print('Evaluating k-NN precision and recall with %i samples...' % num_images)
+    start = time()
+
+    # Precision: How many points from eval_features are in ref_features manifold.
+    precision = ref_manifold.evaluate(eval_features)
+    state['precision'] = precision.mean(axis=0)
+
+    # Recall: How many points from ref_features are in eval_features manifold.
+    recall = eval_manifold.evaluate(ref_features)
+    state['recall'] = recall.mean(axis=0)
+
+    if not mute:
+        print('Evaluated k-NN precision and recall in: %gs' % (time() - start))
+
+    return state
+
+# ----------------------------------------------------------------------------
+

metrics/visualizations.py

@@ -0,0 +1,123 @@
+import numpy as np
+from matplotlib import pyplot as plt
+from scipy.fft import fft
+from scipy.signal import savgol_filter
+from tools import rms_normalize
+
+colors = [
+    # (0, 0, 0),       # Black
+    # (86, 180, 233),  # Sky blue
+    # (240, 228, 66),  # Yellow
+    # (204, 121, 167),  # Reddish purple
+    (213, 94, 0),  # Vermilion
+    (0, 114, 178),  # Blue
+    (230, 159, 0),  # Orange
+    (0, 158, 115),  # Bluish green
+]
+
+
+def plot_psd_multiple_signals(signals_list, labels_list, sample_rate=16000, window_size=500,
+                              figsize=(10, 6), save_path=None, normalize=False):
+    """
+    在同一张图上绘制多组音频信号的功率谱密度比较图，使用对数刻度的响度轴（以2为底），并应用平滑处理。
+
+    参数:
+    signals_list: 包含多组音频信号的列表，每组信号形状为 [sample_number, sample_length] 的numpy array
+    labels_list: 每组音频信号对应的标签字符串列表
+    sample_rate: 音频的采样率
+    """
+
+    # 确保传入的signals_list和labels_list长度相同
+    assert len(signals_list) == len(labels_list), "每组信号必须有一个对应的标签。"
+
+    signals_list = [np.array([rms_normalize(signal) for signal in signals]) for signals in signals_list]
+
+    # 绘图准备
+    plt.figure(figsize=figsize)
+
+    # 遍历所有的音频信号
+    i = 0
+    for signal, label in zip(signals_list, labels_list):
+        # 计算FFT
+        fft_signal = fft(signal, axis=1)
+
+        # 计算平均功率谱密度
+        psd_signal = np.mean(np.abs(fft_signal)**2, axis=0)
+
+        # 计算频率轴
+        freqs = np.fft.fftfreq(signal.shape[1], 1/sample_rate)
+
+        # 应用Savitzky-Golay滤波器进行平滑
+        psd_smoothed = savgol_filter(np.log2(psd_signal[:signal.shape[1] // 2] + 1), window_size, 3)  # 窗口大小51, 多项式阶数3
+
+        # Normalize each curve if normalize is True
+        if normalize:
+            psd_smoothed /= np.mean(psd_smoothed)
+
+        # 绘制每组信号的功率谱密度
+        plt.plot(freqs[:signal.shape[1] // 2], psd_smoothed, label=label, color=[x/255.0 for x in colors[i % len(colors)]], linewidth=1)
+        i += 1
+
+    # 设置图表元素
+    plt.xlabel('Frequency (Hz)')
+    plt.ylabel('Mean Log-Amplitude')
+    plt.legend()
+
+    # 根据save_path参数决定保存图像还是直接显示
+    if save_path:
+        plt.savefig(save_path)
+    else:
+        plt.show()
+
+
+def plot_amplitude_over_time(signals_list, labels_list, sample_rate=16000, window_size=500,
+                             figsize=(10, 6), save_path=None, normalize=False, start_time=0):
+    """
+    Plot the loudness of multiple sets of audio signals over time on the same graph,
+    using a logarithmic scale for the loudness axis (base 2), with smoothing applied.
+
+    Parameters:
+    signals_list: List of sets of audio signals, each set is a numpy array with shape [sample_number, sample_length]
+    labels_list: List of labels corresponding to each set of audio signals
+    sample_rate: Sampling rate of the audio
+    window_size: Window size for the Savitzky-Golay filter
+    figsize: Figure size
+    save_path: Path to save the figure, if None, the figure will be displayed
+    normalize: Whether to normalize each curve so that the sum of each curve is the same
+    start_time: Time (in seconds) to start plotting, only data after this time will be retained
+    """
+    assert len(signals_list) == len(labels_list), f"len(signals_list) != len(labels_list) for " \
+                                                  f"len(signals_list) = {len(signals_list)} and len(labels_list) = {len(labels_list)}"
+
+    # Compute starting sample index
+    start_sample = int(start_time * sample_rate)
+    
+    # Normalize signals and truncate data
+    signals_list = [np.array([rms_normalize(signal)[start_sample:] for signal in signals]) for signals in signals_list]
+    time_axis = np.arange(start_sample, start_sample + signals_list[0].shape[1]) / sample_rate
+
+    plt.figure(figsize=figsize)
+
+    i = 0
+    for signal, label in zip(signals_list, labels_list):
+        amplitude_mean = np.mean(np.abs(signal), axis=0)
+
+        amplitude_smoothed = savgol_filter(np.log2(amplitude_mean + 1), window_size, 3)
+
+        # Normalize each curve if normalize is True
+        if normalize:
+            amplitude_smoothed /= np.mean(amplitude_smoothed)
+
+        plt.plot(time_axis, amplitude_smoothed, label=label, color=[x/255.0 for x in colors[i % len(colors)]], linewidth=1)
+        i += 1
+
+    plt.xlabel('Time (seconds)')
+    plt.ylabel('Mean Log-Amplitude')
+    plt.legend()
+
+    # Save or show the figure based on save_path parameter
+    if save_path:
+        plt.savefig(save_path)
+    else:
+        plt.show()
+

model/DiffSynthSampler.py

@@ -0,0 +1,425 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res.expand(broadcast_shape)
+
+
+class DiffSynthSampler:
+
+    def __init__(self, timesteps, beta_start=0.0001, beta_end=0.02, device=None, mute=False,
+                 height=128, max_batchsize=16, max_width=256, channels=4, train_width=64, noise_strategy="repeat"):
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+        self.height = height
+        self.train_width = train_width
+        self.max_batchsize = max_batchsize
+        self.max_width = max_width
+        self.channels = channels
+        self.num_timesteps = timesteps
+        self.timestep_map = list(range(self.num_timesteps))
+        self.betas = np.array(np.linspace(beta_start, beta_end, self.num_timesteps), dtype=np.float64)
+        self.respaced = False
+        self.define_beta_schedule()
+        self.CFG = 1.0
+        self.mute = mute
+        self.noise_strategy = noise_strategy
+
+    def get_deterministic_noise_tensor_non_repeat(self, batchsize, width, reference_noise=None):
+        if reference_noise is None:
+            large_noise_tensor = torch.randn((self.max_batchsize, self.channels, self.height, self.max_width), device=self.device)
+        else:
+            assert reference_noise.shape == (batchsize, self.channels, self.height, self.max_width), "reference_noise shape mismatch"
+            large_noise_tensor = reference_noise
+        return large_noise_tensor[:batchsize, :, :, :width], None
+
+    def get_deterministic_noise_tensor(self, batchsize, width, reference_noise=None):
+        if self.noise_strategy == "repeat":
+            noise, concat_points = self.get_deterministic_noise_tensor_repeat(batchsize, width, reference_noise=reference_noise)
+            return noise, concat_points
+        else:
+            noise, concat_points = self.get_deterministic_noise_tensor_non_repeat(batchsize, width, reference_noise=reference_noise)
+            return noise, concat_points
+
+
+    def get_deterministic_noise_tensor_repeat(self, batchsize, width, reference_noise=None):
+        # 生成与训练数据长度相等的噪音
+        if reference_noise is None:
+            train_noise_tensor = torch.randn((self.max_batchsize, self.channels, self.height, self.train_width), device=self.device)
+        else:
+            assert reference_noise.shape == (batchsize, self.channels, self.height, self.train_width), "reference_noise shape mismatch"
+            train_noise_tensor = reference_noise
+
+        release_width = int(self.train_width * 1.0 / 4)
+        first_part_width = self.train_width - release_width
+
+        first_part = train_noise_tensor[:batchsize, :, :, :first_part_width]
+        release_part = train_noise_tensor[:batchsize, :, :, -release_width:]
+
+        # 如果所需 length 小于等于 origin length，去掉 first_part 的中间部分
+        if width <= self.train_width:
+            _first_part_head_width = int((width - release_width) / 2)
+            _first_part_tail_width = width - release_width - _first_part_head_width
+            all_parts = [first_part[:, :, :, :_first_part_head_width], first_part[:, :, :, -_first_part_tail_width:], release_part]
+
+            # 沿第四维度拼接张量
+            noise_tensor = torch.cat(all_parts, dim=3)
+
+            # 记录拼接点的位置
+            concat_points = [0]
+            for part in all_parts[:-1]:
+                next_point = concat_points[-1] + part.size(3)
+                concat_points.append(next_point)
+
+            return noise_tensor, concat_points
+
+        # 如果所需 length 大于 origin length，不断地从中间插入 first_part 的中间部分
+        else:
+            # 计算需要重复front_width的次数
+            repeats = (width - release_width) // first_part_width
+            extra = (width - release_width) % first_part_width
+
+            _repeat_first_part_head_width = int(first_part_width / 2)
+            _repeat_first_part_tail_width = first_part_width - _repeat_first_part_head_width
+
+            repeated_first_head_parts = [first_part[:, :, :, :_repeat_first_part_head_width] for _ in range(repeats)]
+            repeated_first_tail_parts = [first_part[:, :, :, -_repeat_first_part_tail_width:] for _ in range(repeats)]
+
+            # 计算起始索引
+            _middle_part_start_index = (first_part_width - extra) // 2
+            # 切片张量以获取中间部分
+            middle_part = first_part[:, :, :, _middle_part_start_index: _middle_part_start_index + extra]
+
+            all_parts = repeated_first_head_parts + [middle_part] + repeated_first_tail_parts + [release_part]
+
+            # 沿第四维度拼接张量
+            noise_tensor = torch.cat(all_parts, dim=3)
+
+            # 记录拼接点的位置
+            concat_points = [0]
+            for part in all_parts[:-1]:
+                next_point = concat_points[-1] + part.size(3)
+                concat_points.append(next_point)
+
+            return noise_tensor, concat_points
+
+    def define_beta_schedule(self):
+        assert self.respaced == False, "This schedule has already been respaced!"
+        # define alphas
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = np.cumprod(self.alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recip_alphas = np.sqrt(1.0 / self.alphas)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod))
+
+    def activate_classifier_free_guidance(self, CFG, unconditional_condition):
+        assert (
+                   not unconditional_condition is None) or CFG == 1.0, "For CFG != 1.0, unconditional_condition must be available"
+        self.CFG = CFG
+        self.unconditional_condition = unconditional_condition
+
+    def respace(self, use_timesteps=None):
+        if not use_timesteps is None:
+            last_alpha_cumprod = 1.0
+            new_betas = []
+            self.timestep_map = []
+            for i, _alpha_cumprod in enumerate(self.alphas_cumprod):
+                if i in use_timesteps:
+                    new_betas.append(1 - _alpha_cumprod / last_alpha_cumprod)
+                    last_alpha_cumprod = _alpha_cumprod
+                    self.timestep_map.append(i)
+            self.num_timesteps = len(use_timesteps)
+            self.betas = np.array(new_betas)
+            self.define_beta_schedule()
+            self.respaced = True
+
+    def generate_linear_noise(self, shape, variance=1.0, first_endpoint=None, second_endpoint=None):
+        assert shape[1] == self.channels, "shape[1] != self.channels"
+        assert shape[2] == self.height, "shape[2] != self.height"
+        noise = torch.empty(*shape, device=self.device)
+
+        # 第三种情况：两个端点都不是None，进行线性插值
+        if first_endpoint is not None and second_endpoint is not None:
+            for i in range(shape[0]):
+                alpha = i / (shape[0] - 1)  # 插值系数
+                noise[i] = alpha * second_endpoint + (1 - alpha) * first_endpoint
+            return noise  # 返回插值后的结果，不需要进行后续的均值和方差调整
+        else:
+            # 第一个端点不是None
+            if first_endpoint is not None:
+                noise[0] = first_endpoint
+                if shape[0] > 1:
+                    noise[1], _ = self.get_deterministic_noise_tensor(1, shape[3])[0]
+            else:
+                noise[0], _ = self.get_deterministic_noise_tensor(1, shape[3])[0]
+                if shape[0] > 1:
+                    noise[1], _ = self.get_deterministic_noise_tensor(1, shape[3])[0]
+
+            # 生成其他的噪声点
+            for i in range(2, shape[0]):
+                noise[i] = 2 * noise[i - 1] - noise[i - 2]
+
+        # 当只有一个端点被指定时
+        current_var = noise.var()
+        stddev_ratio = torch.sqrt(variance / current_var)
+        noise = noise * stddev_ratio
+
+        # 如果第一个端点被指定，进行平移调整
+        if first_endpoint is not None:
+            shift = first_endpoint - noise[0]
+            noise += shift
+
+        return noise
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        assert x_start.shape[1] == self.channels, "shape[1] != self.channels"
+        assert x_start.shape[2] == self.height, "shape[2] != self.height"
+
+        if noise is None:
+            # noise = torch.randn_like(x_start)
+            noise, _ = self.get_deterministic_noise_tensor(x_start.shape[0], x_start.shape[3])
+
+        assert noise.shape == x_start.shape
+        return (
+                _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+                + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+                * noise
+        )
+
+    @torch.no_grad()
+    def ddim_sample(self, model, x, t, condition=None, ddim_eta=0.0):
+        map_tensor = torch.tensor(self.timestep_map, device=t.device, dtype=t.dtype)
+        mapped_t = map_tensor[t]
+
+        # Todo: add CFG
+
+        if self.CFG == 1.0:
+            pred_noise = model(x, mapped_t, condition)
+        else:
+            unconditional_condition = self.unconditional_condition.unsqueeze(0).repeat(
+                *([x.shape[0]] + [1] * len(self.unconditional_condition.shape)))
+            x_in = torch.cat([x] * 2)
+            t_in = torch.cat([mapped_t] * 2)
+            c_in = torch.cat([unconditional_condition, condition])
+            noise_uncond, noise = model(x_in, t_in, c_in).chunk(2)
+            pred_noise = noise_uncond + self.CFG * (noise - noise_uncond)
+
+        # Todo: END
+
+        alpha_cumprod_t = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_cumprod_t_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+
+        pred_x0 = (x - torch.sqrt((1. - alpha_cumprod_t)) * pred_noise) / torch.sqrt(alpha_cumprod_t)
+
+        sigmas_t = (
+                ddim_eta
+                * torch.sqrt((1 - alpha_cumprod_t_prev) / (1 - alpha_cumprod_t))
+                * torch.sqrt(1 - alpha_cumprod_t / alpha_cumprod_t_prev)
+        )
+
+        pred_dir_xt = torch.sqrt(1 - alpha_cumprod_t_prev - sigmas_t ** 2) * pred_noise
+
+
+        step_noise, _ = self.get_deterministic_noise_tensor(x.shape[0], x.shape[3])
+
+
+        x_prev = torch.sqrt(alpha_cumprod_t_prev) * pred_x0 + pred_dir_xt + sigmas_t * step_noise
+
+        return x_prev
+
+    def p_sample(self, model, x, t, condition=None, sampler="ddim"):
+        if sampler == "ddim":
+            return self.ddim_sample(model, x, t, condition=condition, ddim_eta=0.0)
+        elif sampler == "ddpm":
+            return self.ddim_sample(model, x, t, condition=condition, ddim_eta=1.0)
+        else:
+            raise NotImplementedError()
+
+    def get_dynamic_masks(self, n_masks, shape, concat_points, mask_flexivity=0.8):
+        release_length = int(self.train_width / 4)
+        assert shape[3] == (concat_points[-1] + release_length), "shape[3] != (concat_points[-1] + release_length)"
+
+        fraction_lengths = [concat_points[i + 1] - concat_points[i] for i in range(len(concat_points) - 1)]
+
+        # Todo: remove hard-coding
+        n_guidance_steps = int(n_masks * mask_flexivity)
+        n_free_steps = n_masks - n_guidance_steps
+
+        masks = []
+        # Todo: 在一半的 steps 内收缩 mask。也就是说，在后程对 release 以外的区域不做inpaint，而是 img2img
+        for i in range(n_guidance_steps):
+            # mask = 1, freeze
+            step_i_mask = torch.zeros((shape[0], 1, shape[2], shape[3]), dtype=torch.float32).to(self.device)
+            step_i_mask[:, :, :, -release_length:] = 1.0
+
+            for fraction_index in range(len(fraction_lengths)):
+
+                _fraction_mask_length = int((n_guidance_steps - 1 - i) / (n_guidance_steps - 1) * fraction_lengths[fraction_index])
+
+                if fraction_index == 0:
+                    step_i_mask[:, :, :, :_fraction_mask_length] = 1.0
+                elif fraction_index == len(fraction_lengths) - 1:
+                    if not _fraction_mask_length == 0:
+                        step_i_mask[:, :, :, -_fraction_mask_length - release_length:] = 1.0
+                else:
+                    fraction_mask_start_position = int((fraction_lengths[fraction_index] - _fraction_mask_length) / 2)
+
+                    step_i_mask[:, :, :,
+                    concat_points[fraction_index] + fraction_mask_start_position:concat_points[
+                                                                                     fraction_index] + fraction_mask_start_position + _fraction_mask_length] = 1.0
+            masks.append(step_i_mask)
+
+        for i in range(n_free_steps):
+            step_i_mask = torch.zeros((shape[0], 1, shape[2], shape[3]), dtype=torch.float32).to(self.device)
+            step_i_mask[:, :, :, -release_length:] = 1.0
+            masks.append(step_i_mask)
+
+        masks.reverse()
+        return masks
+
+    @torch.no_grad()
+    def p_sample_loop(self, model, shape, initial_noise=None, start_noise_level_ratio=1.0, end_noise_level_ratio=0.0,
+                      return_tensor=False, condition=None, guide_img=None,
+                      mask=None, sampler="ddim", inpaint=False, use_dynamic_mask=False, mask_flexivity=0.8):
+
+        assert shape[1] == self.channels, "shape[1] != self.channels"
+        assert shape[2] == self.height, "shape[2] != self.height"
+
+        initial_noise, _ = self.get_deterministic_noise_tensor(shape[0], shape[3], reference_noise=initial_noise)
+        assert initial_noise.shape == shape, "initial_noise.shape != shape"
+
+        start_noise_level_index = int(self.num_timesteps * start_noise_level_ratio) # not included!!!
+        end_noise_level_index = int(self.num_timesteps * end_noise_level_ratio)
+
+        timesteps = reversed(range(end_noise_level_index, start_noise_level_index))
+
+        # configure initial img
+        assert (start_noise_level_ratio == 1.0) or (
+            not guide_img is None), "A guide_img must be given to sample from a non-pure-noise."
+
+        if guide_img is None:
+            img = initial_noise
+        else:
+            guide_img, concat_points = self.get_deterministic_noise_tensor_repeat(shape[0], shape[3], reference_noise=guide_img)
+            assert guide_img.shape == shape, "guide_img.shape != shape"
+
+            if start_noise_level_index > 0:
+                t = torch.full((shape[0],), start_noise_level_index-1, device=self.device).long()   # -1 for start_noise_level_index not included
+                img = self.q_sample(guide_img, t, noise=initial_noise)
+            else:
+                print("Zero noise added to the guidance latent representation.")
+                img = guide_img
+
+        # get masks
+        n_masks = start_noise_level_index - end_noise_level_index
+        if use_dynamic_mask:
+            masks = self.get_dynamic_masks(n_masks, shape, concat_points, mask_flexivity)
+        else:
+            masks = [mask for _ in range(n_masks)]
+
+        imgs = [img]
+        current_mask = None
+
+
+        for i in tqdm(timesteps, total=start_noise_level_index - end_noise_level_index, disable=self.mute):
+
+            # if i == 3:
+            #     return [img], initial_noise  # 第1排，第1列
+
+            img = self.p_sample(model, img, torch.full((shape[0],), i, device=self.device, dtype=torch.long),
+                                condition=condition,
+                                sampler=sampler)
+            # if i == 3:
+            #     return [img], initial_noise  # 第1排，第2列
+
+            if inpaint:
+                if i > 0:
+                    t = torch.full((shape[0],), int(i-1), device=self.device).long()
+                    img_noise_t = self.q_sample(guide_img, t, noise=initial_noise)
+                    # if i == 3:
+                    #     return [img_noise_t], initial_noise  # 第2排，第2列
+                    current_mask = masks.pop()
+                    img = current_mask * img_noise_t + (1 - current_mask) * img
+                    # if i == 3:
+                    #     return [img], initial_noise  # 第1.5排，最后1列
+                else:
+                    img = current_mask * guide_img + (1 - current_mask) * img
+
+            if return_tensor:
+                imgs.append(img)
+            else:
+                imgs.append(img.cpu().numpy())
+
+        return imgs, initial_noise
+
+
+    def sample(self, model, shape, return_tensor=False, condition=None, sampler="ddim", initial_noise=None, seed=None):
+        if not seed is None:
+            torch.manual_seed(seed)
+        return self.p_sample_loop(model, shape, initial_noise=initial_noise, start_noise_level_ratio=1.0, end_noise_level_ratio=0.0,
+                                  return_tensor=return_tensor, condition=condition, sampler=sampler)
+
+    def interpolate(self, model, shape, variance, first_endpoint=None, second_endpoint=None, return_tensor=False,
+                    condition=None, sampler="ddim", seed=None):
+        if not seed is None:
+            torch.manual_seed(seed)
+        linear_noise = self.generate_linear_noise(shape, variance, first_endpoint=first_endpoint,
+                                                  second_endpoint=second_endpoint)
+        return self.p_sample_loop(model, shape, initial_noise=linear_noise, start_noise_level_ratio=1.0,
+                                  end_noise_level_ratio=0.0,
+                                  return_tensor=return_tensor, condition=condition, sampler=sampler)
+
+    def img_guided_sample(self, model, shape, noising_strength, guide_img, return_tensor=False, condition=None,
+                          sampler="ddim", initial_noise=None, seed=None):
+        if not seed is None:
+            torch.manual_seed(seed)
+        assert guide_img.shape[-1] == shape[-1], "guide_img.shape[:-1] != shape[:-1]"
+        return self.p_sample_loop(model, shape, start_noise_level_ratio=noising_strength, end_noise_level_ratio=0.0,
+                                  return_tensor=return_tensor, condition=condition, sampler=sampler,
+                                  guide_img=guide_img, initial_noise=initial_noise)
+
+    def inpaint_sample(self, model, shape, noising_strength, guide_img, mask, return_tensor=False, condition=None,
+                       sampler="ddim", initial_noise=None, use_dynamic_mask=False, end_noise_level_ratio=0.0, seed=None,
+                       mask_flexivity=0.8):
+        if not seed is None:
+            torch.manual_seed(seed)
+        return self.p_sample_loop(model, shape, start_noise_level_ratio=noising_strength, end_noise_level_ratio=end_noise_level_ratio,
+                                  return_tensor=return_tensor, condition=condition, guide_img=guide_img, mask=mask,
+                                  sampler=sampler, inpaint=True, initial_noise=initial_noise, use_dynamic_mask=use_dynamic_mask,
+                                  mask_flexivity=mask_flexivity)

model/GAN.py

@@ -0,0 +1,262 @@
+import json
+import numpy as np
+import torch
+from torch import nn
+from six.moves import xrange
+from torch.utils.tensorboard import SummaryWriter
+import random
+
+from model.diffusion import ConditionedUnet
+from tools import create_key
+
+class Discriminator(nn.Module):
+    def __init__(self, label_emb_dim):
+        super(Discriminator, self).__init__()
+        # 特征图卷积层
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(4, 64, kernel_size=4, stride=2, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.AdaptiveAvgPool2d(1), # 添加适应性池化层
+            nn.Flatten()
+        )
+
+        # 文本嵌入处理
+        self.text_embedding = nn.Sequential(
+            nn.Linear(label_emb_dim, 512),
+            nn.LeakyReLU(0.2, inplace=True)
+        )
+
+        # 判别器最后的全连接层
+        self.fc = nn.Linear(512 + 512, 1)  # 两个512分别来自特征图和文本嵌入
+
+    def forward(self, x, text_emb):
+        x = self.conv_layers(x)
+        text_emb = self.text_embedding(text_emb)
+        combined = torch.cat((x, text_emb), dim=1)
+        output = self.fc(combined)
+        return output
+
+
+
+def evaluate_GAN(device, generator, discriminator, iterator, encodes2embeddings_mapping):
+    generator.to(device)
+    discriminator.to(device)
+    generator.eval()
+    discriminator.eval()
+
+    real_accs = []
+    fake_accs = []
+
+    with torch.no_grad():
+        for i in range(100):
+            data, attributes = next(iter(iterator))
+            data = data.to(device)
+
+            conditions = [encodes2embeddings_mapping[create_key(attribute)] for attribute in attributes]
+            selected_conditions = [random.choice(conditions_of_one_sample) for conditions_of_one_sample in conditions]
+            selected_conditions = torch.stack(selected_conditions).float().to(device)
+
+            # 将数据和标签移至设备
+            real_images = data.to(device)
+            labels = selected_conditions.to(device)
+
+            # 生成噪声和假图像
+            noise = torch.randn_like(real_images).to(device)
+            fake_images = generator(noise)
+
+            # 评估鉴别器的性能
+            real_preds = discriminator(real_images, labels).reshape(-1)
+            fake_preds = discriminator(fake_images, labels).reshape(-1)
+            real_acc = (real_preds > 0.5).float().mean().item()  # 真实图像的准确率
+            fake_acc = (fake_preds < 0.5).float().mean().item()  # 生成图像的准确率
+
+            real_accs.append(real_acc)
+            fake_accs.append(fake_acc)
+
+
+    # 计算平均准确率
+    average_real_acc = sum(real_accs) / len(real_accs)
+    average_fake_acc = sum(fake_accs) / len(fake_accs)
+
+    return average_real_acc, average_fake_acc
+
+
+def get_Generator(model_Config, load_pretrain=False, model_name=None, device="cpu"):
+    generator = ConditionedUnet(**model_Config)
+    print(f"Model intialized, size: {sum(p.numel() for p in generator.parameters() if p.requires_grad)}")
+    generator.to(device)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_generator.pth")
+        checkpoint = torch.load(f'models/{model_name}_generator.pth', map_location=device)
+        generator.load_state_dict(checkpoint['model_state_dict'])
+    generator.eval()
+    return generator
+
+
+def get_Discriminator(model_Config, load_pretrain=False, model_name=None, device="cpu"):
+    discriminator = Discriminator(**model_Config)
+    print(f"Model intialized, size: {sum(p.numel() for p in discriminator.parameters() if p.requires_grad)}")
+    discriminator.to(device)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_discriminator.pth")
+        checkpoint = torch.load(f'models/{model_name}_discriminator.pth', map_location=device)
+        discriminator.load_state_dict(checkpoint['model_state_dict'])
+    discriminator.eval()
+    return discriminator
+
+
+def train_GAN(device, init_model_name, unetConfig, BATCH_SIZE, lr_G, lr_D, max_iter, iterator, load_pretrain,
+                     encodes2embeddings_mapping, save_steps, unconditional_condition, uncondition_rate, save_model_name=None):
+
+    if save_model_name is None:
+        save_model_name = init_model_name
+
+    def save_model_hyperparameter(model_name, unetConfig, BATCH_SIZE, model_size, current_iter, current_loss):
+        model_hyperparameter = unetConfig
+        model_hyperparameter["BATCH_SIZE"] = BATCH_SIZE
+        model_hyperparameter["lr_G"] = lr_G
+        model_hyperparameter["lr_D"] = lr_D
+        model_hyperparameter["model_size"] = model_size
+        model_hyperparameter["current_iter"] = current_iter
+        model_hyperparameter["current_loss"] = current_loss
+        with open(f"models/hyperparameters/{model_name}_GAN.json", "w") as json_file:
+            json.dump(model_hyperparameter, json_file, ensure_ascii=False, indent=4)
+
+    generator = ConditionedUnet(**unetConfig)
+    discriminator = Discriminator(unetConfig["label_emb_dim"])
+    generator_size = sum(p.numel() for p in generator.parameters() if p.requires_grad)
+    discriminator_size = sum(p.numel() for p in discriminator.parameters() if p.requires_grad)
+
+    print(f"Generator trainable parameters: {generator_size}, discriminator trainable parameters: {discriminator_size}")
+    generator.to(device)
+    discriminator.to(device)
+    optimizer_G = torch.optim.Adam(filter(lambda p: p.requires_grad, generator.parameters()), lr=lr_G, amsgrad=False)
+    optimizer_D = torch.optim.Adam(filter(lambda p: p.requires_grad, discriminator.parameters()), lr=lr_D, amsgrad=False)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{init_model_name}_generator.pt")
+        checkpoint = torch.load(f'models/{init_model_name}_generator.pth')
+        generator.load_state_dict(checkpoint['model_state_dict'])
+        optimizer_G.load_state_dict(checkpoint['optimizer_state_dict'])
+        print(f"Loading weights from models/{init_model_name}_discriminator.pt")
+        checkpoint = torch.load(f'models/{init_model_name}_discriminator.pth')
+        discriminator.load_state_dict(checkpoint['model_state_dict'])
+        optimizer_D.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("Model initialized.")
+    if max_iter == 0:
+        print("Return model directly.")
+        return generator, discriminator, optimizer_G, optimizer_D
+
+
+    train_loss_G, train_loss_D = [], []
+    writer = SummaryWriter(f'runs/{save_model_name}_GAN')
+
+    # average_real_acc, average_fake_acc = evaluate_GAN(device, generator, discriminator, iterator, encodes2embeddings_mapping)
+    # print(f"average_real_acc, average_fake_acc: {average_real_acc, average_fake_acc}")
+
+    criterion = nn.BCEWithLogitsLoss()
+    generator.train()
+    for i in xrange(max_iter):
+        data, attributes = next(iter(iterator))
+        data = data.to(device)
+
+        conditions = [encodes2embeddings_mapping[create_key(attribute)] for attribute in attributes]
+        unconditional_condition_copy = torch.tensor(unconditional_condition, dtype=torch.float32).to(device).detach()
+        selected_conditions = [unconditional_condition_copy if random.random() < uncondition_rate else random.choice(
+            conditions_of_one_sample) for conditions_of_one_sample in conditions]
+        batch_size = len(selected_conditions)
+        selected_conditions = torch.stack(selected_conditions).float().to(device)
+
+        # 将数据和标签移至设备
+        real_images = data.to(device)
+        labels = selected_conditions.to(device)
+
+        # 真实和假的标签
+        real_labels = torch.ones(batch_size, 1).to(device)
+        fake_labels = torch.zeros(batch_size, 1).to(device)
+
+        # ========== 训练鉴别器 ==========
+        optimizer_D.zero_grad()
+
+        # 计算鉴别器对真实图像的损失
+        outputs_real = discriminator(real_images, labels)
+        loss_D_real = criterion(outputs_real, real_labels)
+
+        # 生成假图像
+        noise = torch.randn_like(real_images).to(device)
+        fake_images = generator(noise, labels)
+
+        # 计算鉴别器对假图像的损失
+        outputs_fake = discriminator(fake_images.detach(), labels)
+        loss_D_fake = criterion(outputs_fake, fake_labels)
+
+        # 反向传播和优化
+        loss_D = loss_D_real + loss_D_fake
+        loss_D.backward()
+        optimizer_D.step()
+
+        # ========== 训练生成器 ==========
+        optimizer_G.zero_grad()
+
+        # 计算生成器的损失
+        outputs_fake = discriminator(fake_images, labels)
+        loss_G = criterion(outputs_fake, real_labels)
+
+        # 反向传播和优化
+        loss_G.backward()
+        optimizer_G.step()
+
+
+        train_loss_G.append(loss_G.item())
+        train_loss_D.append(loss_D.item())
+        step = int(optimizer_G.state_dict()['state'][list(optimizer_G.state_dict()['state'].keys())[0]]['step'].numpy())
+
+        if (i + 1) % 100 == 0:
+            print('%d step' % (step))
+
+        if (i + 1) % save_steps == 0:
+            current_loss_D = np.mean(train_loss_D[-save_steps:])
+            current_loss_G = np.mean(train_loss_G[-save_steps:])
+            print('current_loss_G: %.5f' % current_loss_G)
+            print('current_loss_D: %.5f' % current_loss_D)
+
+            writer.add_scalar(f"current_loss_G", current_loss_G, step)
+            writer.add_scalar(f"current_loss_D", current_loss_D, step)
+
+
+            torch.save({
+                'model_state_dict': generator.state_dict(),
+                'optimizer_state_dict': optimizer_G.state_dict(),
+            }, f'models/{save_model_name}_generator.pth')
+            save_model_hyperparameter(save_model_name, unetConfig, BATCH_SIZE, generator_size, step, current_loss_G)
+            torch.save({
+                'model_state_dict': discriminator.state_dict(),
+                'optimizer_state_dict': optimizer_D.state_dict(),
+            }, f'models/{save_model_name}_discriminator.pth')
+            save_model_hyperparameter(save_model_name, unetConfig, BATCH_SIZE, discriminator_size, step, current_loss_D)
+
+        if step % 10000 == 0:
+            torch.save({
+                'model_state_dict': generator.state_dict(),
+                'optimizer_state_dict': optimizer_G.state_dict(),
+            }, f'models/history/{save_model_name}_{step}_generator.pth')
+            torch.save({
+                'model_state_dict': discriminator.state_dict(),
+                'optimizer_state_dict': optimizer_D.state_dict(),
+            }, f'models/history/{save_model_name}_{step}_discriminator.pth')
+
+    return generator, discriminator, optimizer_G, optimizer_D
+
+

model/VQGAN.py

@@ -0,0 +1,684 @@
+import json
+from torch.utils.tensorboard import SummaryWriter
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from six.moves import xrange
+from einops import rearrange
+from torchvision import models
+
+
+def Normalize(in_channels, num_groups=32, norm_type="groupnorm"):
+    """Normalization layer"""
+
+    if norm_type == "batchnorm":
+        return torch.nn.BatchNorm2d(in_channels)
+    return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+def nonlinearity(x, act_type="relu"):
+    """Nonlinear activation function"""
+
+    if act_type == "relu":
+        return F.relu(x)
+    else:
+        # swish
+        return x * torch.sigmoid(x)
+
+
+class VectorQuantizer(nn.Module):
+    """Vector quantization layer"""
+
+    def __init__(self, num_embeddings, embedding_dim, commitment_cost):
+        super(VectorQuantizer, self).__init__()
+
+        self._embedding_dim = embedding_dim
+        self._num_embeddings = num_embeddings
+
+        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
+        self._embedding.weight.data.uniform_(-1 / self._num_embeddings, 1 / self._num_embeddings)
+        self._commitment_cost = commitment_cost
+
+    def forward(self, inputs):
+        # convert inputs from BCHW -> BHWC
+        inputs = inputs.permute(0, 2, 3, 1).contiguous()
+        input_shape = inputs.shape
+
+        # Flatten input BCHW -> (BHW)C
+        flat_input = inputs.view(-1, self._embedding_dim)
+
+        # Calculate distances (input-embedding)^2
+        distances = (torch.sum(flat_input ** 2, dim=1, keepdim=True)
+                     + torch.sum(self._embedding.weight ** 2, dim=1)
+                     - 2 * torch.matmul(flat_input, self._embedding.weight.t()))
+
+        # Encoding (one-hot-encoding matrix)
+        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
+        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
+        encodings.scatter_(1, encoding_indices, 1)
+
+        # Quantize and unflatten
+        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
+
+        # Loss
+        e_latent_loss = F.mse_loss(quantized.detach(), inputs)
+        q_latent_loss = F.mse_loss(quantized, inputs.detach())
+        loss = q_latent_loss + self._commitment_cost * e_latent_loss
+
+        quantized = inputs + (quantized - inputs).detach()
+        avg_probs = torch.mean(encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        # convert quantized from BHWC -> BCHW
+        min_encodings, min_encoding_indices = None, None
+        return quantized.permute(0, 3, 1, 2).contiguous(), loss, (perplexity, min_encodings, min_encoding_indices)
+
+
+class VectorQuantizerEMA(nn.Module):
+    """Vector quantization layer based on exponential moving average"""
+
+    def __init__(self, num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5):
+        super(VectorQuantizerEMA, self).__init__()
+
+        self._embedding_dim = embedding_dim
+        self._num_embeddings = num_embeddings
+
+        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
+        self._embedding.weight.data.normal_()
+        self._commitment_cost = commitment_cost
+
+        self.register_buffer('_ema_cluster_size', torch.zeros(num_embeddings))
+        self._ema_w = nn.Parameter(torch.Tensor(num_embeddings, self._embedding_dim))
+        self._ema_w.data.normal_()
+
+        self._decay = decay
+        self._epsilon = epsilon
+
+    def forward(self, inputs):
+        # convert inputs from BCHW -> BHWC
+        inputs = inputs.permute(0, 2, 3, 1).contiguous()
+        input_shape = inputs.shape
+
+        # Flatten input
+        flat_input = inputs.view(-1, self._embedding_dim)
+
+        # Calculate distances
+        distances = (torch.sum(flat_input ** 2, dim=1, keepdim=True)
+                     + torch.sum(self._embedding.weight ** 2, dim=1)
+                     - 2 * torch.matmul(flat_input, self._embedding.weight.t()))
+
+        # Encoding
+        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
+        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
+        encodings.scatter_(1, encoding_indices, 1)
+
+        # Quantize and unflatten
+        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
+
+        # Use EMA to update the embedding vectors
+        if self.training:
+            self._ema_cluster_size = self._ema_cluster_size * self._decay + \
+                                     (1 - self._decay) * torch.sum(encodings, 0)
+
+            # Laplace smoothing of the cluster size
+            n = torch.sum(self._ema_cluster_size.data)
+            self._ema_cluster_size = (
+                    (self._ema_cluster_size + self._epsilon)
+                    / (n + self._num_embeddings * self._epsilon) * n)
+
+            dw = torch.matmul(encodings.t(), flat_input)
+            self._ema_w = nn.Parameter(self._ema_w * self._decay + (1 - self._decay) * dw)
+
+            self._embedding.weight = nn.Parameter(self._ema_w / self._ema_cluster_size.unsqueeze(1))
+
+        # Loss
+        e_latent_loss = F.mse_loss(quantized.detach(), inputs)
+        loss = self._commitment_cost * e_latent_loss
+
+        # Straight Through Estimator
+        quantized = inputs + (quantized - inputs).detach()
+        avg_probs = torch.mean(encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        # convert quantized from BHWC -> BCHW
+        min_encodings, min_encoding_indices = None, None
+        return quantized.permute(0, 3, 1, 2).contiguous(), loss, (perplexity, min_encodings, min_encoding_indices)
+
+
+class DownSample(nn.Module):
+    """DownSample layer"""
+
+    def __init__(self, in_channels, out_channels):
+        super(DownSample, self).__init__()
+        self._conv2d = nn.Conv2d(in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 kernel_size=4,
+                                 stride=2, padding=1)
+
+    def forward(self, x):
+        return self._conv2d(x)
+
+
+class UpSample(nn.Module):
+    """UpSample layer"""
+
+    def __init__(self, in_channels, out_channels):
+        super(UpSample, self).__init__()
+        self._conv2d = nn.ConvTranspose2d(in_channels=in_channels,
+                                          out_channels=out_channels,
+                                          kernel_size=4,
+                                          stride=2, padding=1)
+
+    def forward(self, x):
+        return self._conv2d(x)
+
+
+class ResnetBlock(nn.Module):
+    """ResnetBlock is a combination of non-linearity, convolution, and normalization"""
+
+    def __init__(self, *, in_channels, out_channels=None, double_conv=False, conv_shortcut=False,
+                 dropout=0.0, temb_channels=512, norm_type="groupnorm", act_type="relu", num_groups=32):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.act_type = act_type
+
+        self.norm1 = Normalize(in_channels, norm_type=norm_type, num_groups=num_groups)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+
+        self.double_conv = double_conv
+        if self.double_conv:
+            self.norm2 = Normalize(out_channels, norm_type=norm_type, num_groups=num_groups)
+            self.dropout = torch.nn.Dropout(dropout)
+            self.conv2 = torch.nn.Conv2d(out_channels,
+                                         out_channels,
+                                         kernel_size=3,
+                                         stride=1,
+                                         padding=1)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x, temb=None):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h, act_type=self.act_type)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb, act_type=self.act_type))[:, :, None, None]
+
+        if self.double_conv:
+            h = self.norm2(h)
+            h = nonlinearity(h, act_type=self.act_type)
+            h = self.dropout(h)
+            h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class LinearAttention(nn.Module):
+    """Efficient attention block based on <https://proceedings.mlr.press/v119/katharopoulos20a.html>"""
+
+    def __init__(self, dim, heads=4, dim_head=32, with_skip=True):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+        self.with_skip = with_skip
+        if self.with_skip:
+            self.nin_shortcut = torch.nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads=self.heads, qkv=3)
+        k = k.softmax(dim=-1)
+        context = torch.einsum('bhdn,bhen->bhde', k, v)
+        out = torch.einsum('bhde,bhdn->bhen', context, q)
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+
+        if self.with_skip:
+            return self.to_out(out) + self.nin_shortcut(x)
+        return self.to_out(out)
+
+
+class Encoder(nn.Module):
+    """The encoder, consisting of alternating stacks of ResNet blocks, efficient attention modules, and downsampling layers."""
+
+    def __init__(self, in_channels, hidden_channels, embedding_dim, block_depth=2,
+                 attn_pos=None, attn_with_skip=True, norm_type="groupnorm", act_type="relu", num_groups=32):
+        super(Encoder, self).__init__()
+
+        if attn_pos is None:
+            attn_pos = []
+        self._layers = nn.ModuleList([DownSample(in_channels, hidden_channels[0])])
+        current_channel = hidden_channels[0]
+
+        for i in range(1, len(hidden_channels)):
+            for _ in range(block_depth - 1):
+                self._layers.append(ResnetBlock(in_channels=current_channel,
+                                                out_channels=current_channel,
+                                                double_conv=False,
+                                                conv_shortcut=False,
+                                                norm_type=norm_type,
+                                                act_type=act_type,
+                                                num_groups=num_groups))
+                if current_channel in attn_pos:
+                    self._layers.append(LinearAttention(current_channel, 1, 32, attn_with_skip))
+
+            self._layers.append(Normalize(current_channel, norm_type=norm_type, num_groups=num_groups))
+            self._layers.append(nn.ReLU())
+            self._layers.append(DownSample(current_channel, hidden_channels[i]))
+            current_channel = hidden_channels[i]
+
+        for _ in range(block_depth - 1):
+            self._layers.append(ResnetBlock(in_channels=current_channel,
+                                            out_channels=current_channel,
+                                            double_conv=False,
+                                            conv_shortcut=False,
+                                            norm_type=norm_type,
+                                            act_type=act_type,
+                                            num_groups=num_groups))
+            if current_channel in attn_pos:
+                self._layers.append(LinearAttention(current_channel, 1, 32, attn_with_skip))
+
+        # Conv1x1: hidden_channels[-1] -> embedding_dim
+        self._layers.append(Normalize(current_channel, norm_type=norm_type, num_groups=num_groups))
+        self._layers.append(nn.ReLU())
+        self._layers.append(nn.Conv2d(in_channels=current_channel,
+                                      out_channels=embedding_dim,
+                                      kernel_size=1,
+                                      stride=1))
+
+    def forward(self, x):
+        for layer in self._layers:
+            x = layer(x)
+        return x
+
+
+class Decoder(nn.Module):
+    """The decoder, consisting of alternating stacks of ResNet blocks, efficient attention modules, and upsampling layers."""
+
+    def __init__(self, embedding_dim, hidden_channels, out_channels, block_depth=2,
+                 attn_pos=None, attn_with_skip=True, norm_type="groupnorm", act_type="relu",
+                                                num_groups=32):
+        super(Decoder, self).__init__()
+
+        if attn_pos is None:
+            attn_pos = []
+        reversed_hidden_channels = list(reversed(hidden_channels))
+
+        # Conv1x1: hidden_channels[-1] -> embedding_dim
+        self._layers = nn.ModuleList([nn.Conv2d(in_channels=embedding_dim,
+                                                out_channels=reversed_hidden_channels[0],
+                                                kernel_size=1, stride=1, bias=False)])
+
+        current_channel = reversed_hidden_channels[0]
+
+        for _ in range(block_depth - 1):
+            if current_channel in attn_pos:
+                self._layers.append(LinearAttention(current_channel, 1, 32, attn_with_skip))
+            self._layers.append(ResnetBlock(in_channels=current_channel,
+                                            out_channels=current_channel,
+                                            double_conv=False,
+                                            conv_shortcut=False,
+                                            norm_type=norm_type,
+                                            act_type=act_type,
+                                            num_groups=num_groups))
+
+        for i in range(1, len(reversed_hidden_channels)):
+            self._layers.append(Normalize(current_channel, norm_type=norm_type, num_groups=num_groups))
+            self._layers.append(nn.ReLU())
+            self._layers.append(UpSample(current_channel, reversed_hidden_channels[i]))
+            current_channel = reversed_hidden_channels[i]
+
+            for _ in range(block_depth - 1):
+                if current_channel in attn_pos:
+                    self._layers.append(LinearAttention(current_channel, 1, 32, attn_with_skip))
+                self._layers.append(ResnetBlock(in_channels=current_channel,
+                                                out_channels=current_channel,
+                                                double_conv=False,
+                                                conv_shortcut=False,
+                                                norm_type=norm_type,
+                                                act_type=act_type,
+                                                num_groups=num_groups))
+
+        self._layers.append(Normalize(current_channel, norm_type=norm_type, num_groups=num_groups))
+        self._layers.append(nn.ReLU())
+        self._layers.append(UpSample(current_channel, current_channel))
+
+        # final layers
+        self._layers.append(ResnetBlock(in_channels=current_channel,
+                                        out_channels=out_channels,
+                                        double_conv=False,
+                                        conv_shortcut=False,
+                                        norm_type=norm_type,
+                                        act_type=act_type,
+                                        num_groups=num_groups))
+
+
+    def forward(self, x):
+        for layer in self._layers:
+            x = layer(x)
+
+        log_magnitude = torch.nn.functional.softplus(x[:, 0, :, :])
+
+        cos_phase = torch.tanh(x[:, 1, :, :])
+        sin_phase = torch.tanh(x[:, 2, :, :])
+        x = torch.stack([log_magnitude, cos_phase, sin_phase], dim=1)
+
+        return x
+
+
+class VQGAN_Discriminator(nn.Module):
+    """The discriminator employs an 18-layer-ResNet architecture , with the first layer replaced by a 2D convolutional
+    layer that accommodates spectral representation inputs and the last two layers replaced by a binary classifier
+    layer."""
+
+    def __init__(self, in_channels=1):
+        super(VQGAN_Discriminator, self).__init__()
+        resnet = models.resnet18(pretrained=True)
+
+        # 修改第一层以接受单通道（黑白）图像
+        resnet.conv1 = nn.Conv2d(in_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
+
+        # 使用ResNet的特征提取部分
+        self.features = nn.Sequential(*list(resnet.children())[:-2])
+
+        # 添加判别器的额外层
+        self.classifier = nn.Sequential(
+            nn.Linear(512, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = nn.functional.adaptive_avg_pool2d(x, (1, 1))
+        x = torch.flatten(x, 1)
+        x = self.classifier(x)
+        return x
+
+
+class VQGAN(nn.Module):
+    """The VQ-GAN model. <https://openaccess.thecvf.com/content/CVPR2021/html/Esser_Taming_Transformers_for_High-Resolution_Image_Synthesis_CVPR_2021_paper.html?ref=>"""
+
+    def __init__(self, in_channels, hidden_channels, embedding_dim, out_channels, block_depth=2,
+                 attn_pos=None, attn_with_skip=True, norm_type="groupnorm", act_type="relu",
+                 num_embeddings=1024, commitment_cost=0.25, decay=0.99, num_groups=32):
+        super(VQGAN, self).__init__()
+
+        self._encoder = Encoder(in_channels, hidden_channels, embedding_dim, block_depth=block_depth,
+                                attn_pos=attn_pos, attn_with_skip=attn_with_skip, norm_type=norm_type, act_type="act_type", num_groups=num_groups)
+
+        if decay > 0.0:
+            self._vq_vae = VectorQuantizerEMA(num_embeddings, embedding_dim,
+                                              commitment_cost, decay)
+        else:
+            self._vq_vae = VectorQuantizer(num_embeddings, embedding_dim,
+                                           commitment_cost)
+        self._decoder = Decoder(embedding_dim, hidden_channels, out_channels, block_depth=block_depth,
+                                attn_pos=attn_pos, attn_with_skip=attn_with_skip, norm_type=norm_type,
+                                act_type=act_type, num_groups=num_groups)
+
+    def forward(self, x):
+        z = self._encoder(x)
+        quantized, vq_loss, (perplexity, _, _) = self._vq_vae(z)
+        x_recon = self._decoder(quantized)
+
+        return vq_loss, x_recon, perplexity
+
+
+class ReconstructionLoss(nn.Module):
+    def __init__(self, w1, w2, epsilon=1e-3):
+        super(ReconstructionLoss, self).__init__()
+        self.w1 = w1
+        self.w2 = w2
+        self.epsilon = epsilon
+
+    def weighted_mae_loss(self, y_true, y_pred):
+        # avoid divide by zero
+        y_true_safe = torch.clamp(y_true, min=self.epsilon)
+
+        # compute weighted MAE
+        loss = torch.mean(torch.abs(y_pred - y_true) / y_true_safe)
+        return loss
+
+    def mae_loss(self, y_true, y_pred):
+        loss = torch.mean(torch.abs(y_pred - y_true))
+        return loss
+
+    def forward(self, y_pred, y_true):
+        # loss for magnitude channel
+        log_magnitude_loss = self.w1 * self.weighted_mae_loss(y_pred[:, 0, :, :], y_true[:, 0, :, :])
+
+        # loss for phase channels
+        phase_loss = self.w2 * self.mae_loss(y_pred[:, 1:, :, :], y_true[:, 1:, :, :])
+
+        # sum up
+        rec_loss = log_magnitude_loss + phase_loss
+        return log_magnitude_loss, phase_loss, rec_loss
+
+
+def evaluate_VQGAN(model, discriminator, iterator, reconstructionLoss, adversarial_loss, trainingConfig):
+    model.to(trainingConfig["device"])
+    model.eval()
+    train_res_error = []
+    for i in xrange(100):
+        data = next(iter(iterator))
+        data = data.to(trainingConfig["device"])
+
+        # true/fake labels
+        real_labels = torch.ones(data.size(0), 1).to(trainingConfig["device"])
+
+        vq_loss, data_recon, perplexity = model(data)
+
+
+        fake_preds = discriminator(data_recon)
+        adver_loss = adversarial_loss(fake_preds, real_labels)
+
+        log_magnitude_loss, phase_loss, rec_loss = reconstructionLoss(data_recon, data)
+        loss = rec_loss + trainingConfig["vq_weight"] * vq_loss + trainingConfig["adver_weight"] * adver_loss
+
+        train_res_error.append(loss.item())
+    initial_loss = np.mean(train_res_error)
+    return initial_loss
+
+
+def get_VQGAN(model_Config, load_pretrain=False, model_name=None, device="cpu"):
+    VQVAE = VQGAN(**model_Config)
+    print(f"Model intialized, size: {sum(p.numel() for p in VQVAE.parameters() if p.requires_grad)}")
+    VQVAE.to(device)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_imageVQVAE.pth")
+        checkpoint = torch.load(f'models/{model_name}_imageVQVAE.pth', map_location=device)
+        VQVAE.load_state_dict(checkpoint['model_state_dict'])
+    VQVAE.eval()
+    return VQVAE
+
+
+def train_VQGAN(model_Config, trainingConfig, iterator):
+
+    def save_model_hyperparameter(model_Config, trainingConfig, current_iter,
+                                  log_magnitude_loss, phase_loss, current_perplexity, current_vq_loss,
+                                  current_loss):
+        model_name = trainingConfig["model_name"]
+        model_hyperparameter = model_Config
+        model_hyperparameter.update(trainingConfig)
+        model_hyperparameter["current_iter"] = current_iter
+        model_hyperparameter["log_magnitude_loss"] = log_magnitude_loss
+        model_hyperparameter["phase_loss"] = phase_loss
+        model_hyperparameter["erplexity"] = current_perplexity
+        model_hyperparameter["vq_loss"] = current_vq_loss
+        model_hyperparameter["total_loss"] = current_loss
+
+        with open(f"models/hyperparameters/{model_name}_VQGAN_STFT.json", "w") as json_file:
+            json.dump(model_hyperparameter, json_file, ensure_ascii=False, indent=4)
+
+    # initialize VAE
+    model = VQGAN(**model_Config)
+    print(f"VQ_VAE size: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
+    model.to(trainingConfig["device"])
+
+    VAE_optimizer = torch.optim.Adam(model.parameters(), lr=trainingConfig["lr"], amsgrad=False)
+    model_name = trainingConfig["model_name"]
+
+    if trainingConfig["load_pretrain"]:
+        print(f"Loading weights from models/{model_name}_imageVQVAE.pth")
+        checkpoint = torch.load(f'models/{model_name}_imageVQVAE.pth', map_location=trainingConfig["device"])
+        model.load_state_dict(checkpoint['model_state_dict'])
+        VAE_optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("VAE initialized.")
+    if trainingConfig["max_iter"] == 0:
+        print("Return VAE directly.")
+        return model
+
+    # initialize discriminator
+    discriminator = VQGAN_Discriminator(model_Config["in_channels"])
+    print(f"Discriminator size: {sum(p.numel() for p in discriminator.parameters() if p.requires_grad)}")
+    discriminator.to(trainingConfig["device"])
+
+    discriminator_optimizer = torch.optim.Adam(discriminator.parameters(), lr=trainingConfig["d_lr"], amsgrad=False)
+
+    if trainingConfig["load_pretrain"]:
+        print(f"Loading weights from models/{model_name}_imageVQVAE_discriminator.pth")
+        checkpoint = torch.load(f'models/{model_name}_imageVQVAE_discriminator.pth', map_location=trainingConfig["device"])
+        discriminator.load_state_dict(checkpoint['model_state_dict'])
+        discriminator_optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("Discriminator initialized.")
+
+    # Training
+
+    train_res_phase_loss, train_res_perplexity, train_res_log_magnitude_loss, train_res_vq_loss, train_res_loss = [], [], [], [], []
+    train_discriminator_loss, train_adverserial_loss = [], []
+
+    reconstructionLoss = ReconstructionLoss(w1=trainingConfig["w1"], w2=trainingConfig["w2"], epsilon=trainingConfig["threshold"])
+
+    adversarial_loss = nn.BCEWithLogitsLoss()
+    writer = SummaryWriter(f'runs/{model_name}_VQVAE_lr=1e-4')
+
+    previous_lowest_loss = evaluate_VQGAN(model, discriminator, iterator,
+                                          reconstructionLoss, adversarial_loss, trainingConfig)
+    print(f"initial_loss: {previous_lowest_loss}")
+
+    model.train()
+    for i in xrange(trainingConfig["max_iter"]):
+        data = next(iter(iterator))
+        data = data.to(trainingConfig["device"])
+
+        # true/fake labels
+        real_labels = torch.ones(data.size(0), 1).to(trainingConfig["device"])
+        fake_labels = torch.zeros(data.size(0), 1).to(trainingConfig["device"])
+
+        # update discriminator
+        discriminator_optimizer.zero_grad()
+
+        vq_loss, data_recon, perplexity = model(data)
+
+        real_preds = discriminator(data)
+        fake_preds = discriminator(data_recon.detach())
+
+        loss_real = adversarial_loss(real_preds, real_labels)
+        loss_fake = adversarial_loss(fake_preds, fake_labels)
+
+        loss_D = loss_real + loss_fake
+        loss_D.backward()
+        discriminator_optimizer.step()
+
+
+        # update VQVAE
+        VAE_optimizer.zero_grad()
+
+        fake_preds = discriminator(data_recon)
+        adver_loss = adversarial_loss(fake_preds, real_labels)
+
+        log_magnitude_loss, phase_loss, rec_loss = reconstructionLoss(data_recon, data)
+
+        loss = rec_loss + trainingConfig["vq_weight"] * vq_loss + trainingConfig["adver_weight"] * adver_loss
+        loss.backward()
+        VAE_optimizer.step()
+
+        train_discriminator_loss.append(loss_D.item())
+        train_adverserial_loss.append(trainingConfig["adver_weight"] * adver_loss.item())
+        train_res_log_magnitude_loss.append(log_magnitude_loss.item())
+        train_res_phase_loss.append(phase_loss.item())
+        train_res_perplexity.append(perplexity.item())
+        train_res_vq_loss.append(trainingConfig["vq_weight"] * vq_loss.item())
+        train_res_loss.append(loss.item())
+        step = int(VAE_optimizer.state_dict()['state'][list(VAE_optimizer.state_dict()['state'].keys())[0]]['step'].cpu().numpy())
+
+        save_steps = trainingConfig["save_steps"]
+        if (i + 1) % 100 == 0:
+            print('%d step' % (step))
+
+        if (i + 1) % save_steps == 0:
+            current_discriminator_loss = np.mean(train_discriminator_loss[-save_steps:])
+            current_adverserial_loss = np.mean(train_adverserial_loss[-save_steps:])
+            current_log_magnitude_loss = np.mean(train_res_log_magnitude_loss[-save_steps:])
+            current_phase_loss = np.mean(train_res_phase_loss[-save_steps:])
+            current_perplexity = np.mean(train_res_perplexity[-save_steps:])
+            current_vq_loss = np.mean(train_res_vq_loss[-save_steps:])
+            current_loss = np.mean(train_res_loss[-save_steps:])
+
+            print('discriminator_loss: %.3f' % current_discriminator_loss)
+            print('adverserial_loss: %.3f' % current_adverserial_loss)
+            print('log_magnitude_loss: %.3f' % current_log_magnitude_loss)
+            print('phase_loss: %.3f' % current_phase_loss)
+            print('perplexity: %.3f' % current_perplexity)
+            print('vq_loss: %.3f' % current_vq_loss)
+            print('total_loss: %.3f' % current_loss)
+            writer.add_scalar(f"log_magnitude_loss", current_log_magnitude_loss, step)
+            writer.add_scalar(f"phase_loss", current_phase_loss, step)
+            writer.add_scalar(f"perplexity", current_perplexity, step)
+            writer.add_scalar(f"vq_loss", current_vq_loss, step)
+            writer.add_scalar(f"total_loss", current_loss, step)
+            if current_loss < previous_lowest_loss:
+                previous_lowest_loss = current_loss
+
+                torch.save({
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': VAE_optimizer.state_dict(),
+                }, f'models/{model_name}_imageVQVAE.pth')
+
+                torch.save({
+                    'model_state_dict': discriminator.state_dict(),
+                    'optimizer_state_dict': discriminator_optimizer.state_dict(),
+                }, f'models/{model_name}_imageVQVAE_discriminator.pth')
+
+                save_model_hyperparameter(model_Config, trainingConfig, step,
+                                          current_log_magnitude_loss, current_phase_loss, current_perplexity, current_vq_loss,
+                                          current_loss)
+
+    return model

model/diffusion.py

@@ -0,0 +1,371 @@
+import json
+from functools import partial
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from six.moves import xrange
+from torch.utils.tensorboard import SummaryWriter
+import random
+
+from metrics.IS import get_inception_score
+from tools import create_key
+
+from model.diffusion_components import default, ConvNextBlock, ResnetBlock, SinusoidalPositionEmbeddings, Residual, \
+    PreNorm, \
+    Downsample, Upsample, exists, q_sample, get_beta_schedule, pad_and_concat, ConditionalEmbedding, \
+    LinearCrossAttention, LinearCrossAttentionAdd
+
+
+class ConditionedUnet(nn.Module):
+    def __init__(
+            self,
+            in_dim,
+            out_dim=None,
+            down_dims=None,
+            up_dims=None,
+            mid_depth=3,
+            with_time_emb=True,
+            time_dim=None,
+            resnet_block_groups=8,
+            use_convnext=True,
+            convnext_mult=2,
+            attn_type="linear_cat",
+            n_label_class=11,
+            condition_type="instrument_family",
+            label_emb_dim=128,
+    ):
+        super().__init__()
+
+        self.label_embedding = ConditionalEmbedding(int(n_label_class + 1), int(label_emb_dim), condition_type)
+
+        if up_dims is None:
+            up_dims = [128, 128, 64, 32]
+        if down_dims is None:
+            down_dims = [32, 32, 64, 128]
+
+        out_dim = default(out_dim, in_dim)
+        assert len(down_dims) == len(up_dims), "len(down_dims) != len(up_dims)"
+        assert down_dims[0] == up_dims[-1], "down_dims[0] != up_dims[-1]"
+        assert up_dims[0] == down_dims[-1], "up_dims[0] != down_dims[-1]"
+        down_in_out = list(zip(down_dims[:-1], down_dims[1:]))
+        up_in_out = list(zip(up_dims[:-1], up_dims[1:]))
+        print(f"down_in_out: {down_in_out}")
+        print(f"up_in_out: {up_in_out}")
+        time_dim = default(time_dim, int(down_dims[0] * 4))
+
+        self.init_conv = nn.Conv2d(in_dim, down_dims[0], 7, padding=3)
+
+        if use_convnext:
+            block_klass = partial(ConvNextBlock, mult=convnext_mult)
+        else:
+            block_klass = partial(ResnetBlock, groups=resnet_block_groups)
+
+        if attn_type == "linear_cat":
+            attn_klass = partial(LinearCrossAttention)
+        elif attn_type == "linear_add":
+            attn_klass = partial(LinearCrossAttentionAdd)
+        else:
+            raise NotImplementedError()
+
+        # time embeddings
+        if with_time_emb:
+            self.time_mlp = nn.Sequential(
+                SinusoidalPositionEmbeddings(down_dims[0]),
+                nn.Linear(down_dims[0], time_dim),
+                nn.GELU(),
+                nn.Linear(time_dim, time_dim),
+            )
+        else:
+            time_dim = None
+            self.time_mlp = None
+
+        # left layers
+        self.downs = nn.ModuleList([])
+        self.ups = nn.ModuleList([])
+        skip_dims = []
+
+        for down_dim_in, down_dim_out in down_in_out:
+            self.downs.append(
+                nn.ModuleList(
+                    [
+                        block_klass(down_dim_in, down_dim_out, time_emb_dim=time_dim),
+
+                        Residual(PreNorm(down_dim_out, attn_klass(down_dim_out, label_emb_dim=label_emb_dim, ))),
+                        block_klass(down_dim_out, down_dim_out, time_emb_dim=time_dim),
+                        Residual(PreNorm(down_dim_out, attn_klass(down_dim_out, label_emb_dim=label_emb_dim, ))),
+                        Downsample(down_dim_out),
+                    ]
+                )
+            )
+            skip_dims.append(down_dim_out)
+
+        # bottleneck
+        mid_dim = down_dims[-1]
+        self.mid_left = nn.ModuleList([])
+        self.mid_right = nn.ModuleList([])
+        for _ in range(mid_depth - 1):
+            self.mid_left.append(block_klass(mid_dim, mid_dim, time_emb_dim=time_dim))
+            self.mid_right.append(block_klass(mid_dim * 2, mid_dim, time_emb_dim=time_dim))
+        self.mid_mid = nn.ModuleList(
+            [
+                block_klass(mid_dim, mid_dim, time_emb_dim=time_dim),
+                Residual(PreNorm(mid_dim, attn_klass(mid_dim, label_emb_dim=label_emb_dim, ))),
+                block_klass(mid_dim, mid_dim, time_emb_dim=time_dim),
+            ]
+        )
+
+        # right layers
+        for ind, (up_dim_in, up_dim_out) in enumerate(up_in_out):
+            skip_dim = skip_dims.pop()  # down_dim_out
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        # pop&cat (h/2, w/2, down_dim_out)
+                        block_klass(up_dim_in + skip_dim, up_dim_in, time_emb_dim=time_dim),
+                        Residual(PreNorm(up_dim_in, attn_klass(up_dim_in, label_emb_dim=label_emb_dim, ))),
+                        Upsample(up_dim_in),
+                        # pop&cat (h, w, down_dim_out)
+                        block_klass(up_dim_in + skip_dim, up_dim_out, time_emb_dim=time_dim),
+                        Residual(PreNorm(up_dim_out, attn_klass(up_dim_out, label_emb_dim=label_emb_dim, ))),
+                        # pop&cat (h, w, down_dim_out)
+                        block_klass(up_dim_out + skip_dim, up_dim_out, time_emb_dim=time_dim),
+                        Residual(PreNorm(up_dim_out, attn_klass(up_dim_out, label_emb_dim=label_emb_dim, ))),
+                    ]
+                )
+            )
+
+        self.final_conv = nn.Sequential(
+            block_klass(down_dims[0] + up_dims[-1], up_dims[-1]), nn.Conv2d(up_dims[-1], out_dim, 3, padding=1)
+        )
+
+    def size(self):
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        print(f"Total parameters: {total_params}")
+        print(f"Trainable parameters: {trainable_params}")
+
+
+    def forward(self, x, time, condition=None):
+
+        if condition is not None:
+            condition_emb = self.label_embedding(condition)
+        else:
+            condition_emb = None
+
+        h = []
+
+        x = self.init_conv(x)
+        h.append(x)
+
+        time_emb = self.time_mlp(time) if exists(self.time_mlp) else None
+
+        # downsample
+        for block1, attn1, block2, attn2, downsample in self.downs:
+            x = block1(x, time_emb)
+            x = attn1(x, condition_emb)
+            h.append(x)
+            x = block2(x, time_emb)
+            x = attn2(x, condition_emb)
+            h.append(x)
+            x = downsample(x)
+            h.append(x)
+
+        # bottleneck
+
+        for block in self.mid_left:
+            x = block(x, time_emb)
+            h.append(x)
+
+        (block1, attn, block2) = self.mid_mid
+        x = block1(x, time_emb)
+        x = attn(x, condition_emb)
+        x = block2(x, time_emb)
+
+        for block in self.mid_right:
+            # This is U-Net!!!
+            x = pad_and_concat(h.pop(), x)
+            x = block(x, time_emb)
+
+        # upsample
+        for block1, attn1, upsample, block2, attn2, block3, attn3 in self.ups:
+            x = pad_and_concat(h.pop(), x)
+            x = block1(x, time_emb)
+            x = attn1(x, condition_emb)
+            x = upsample(x)
+
+            x = pad_and_concat(h.pop(), x)
+            x = block2(x, time_emb)
+            x = attn2(x, condition_emb)
+
+            x = pad_and_concat(h.pop(), x)
+            x = block3(x, time_emb)
+            x = attn3(x, condition_emb)
+
+        x = pad_and_concat(h.pop(), x)
+        x = self.final_conv(x)
+        return x
+
+
+def conditional_p_losses(denoise_model, x_start, t, condition, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod,
+                         noise=None, loss_type="l1"):
+    if noise is None:
+        noise = torch.randn_like(x_start)
+
+    x_noisy = q_sample(x_start=x_start, t=t, sqrt_alphas_cumprod=sqrt_alphas_cumprod,
+                       sqrt_one_minus_alphas_cumprod=sqrt_one_minus_alphas_cumprod, noise=noise)
+    predicted_noise = denoise_model(x_noisy, t, condition)
+
+    if loss_type == 'l1':
+        loss = F.l1_loss(noise, predicted_noise)
+    elif loss_type == 'l2':
+        loss = F.mse_loss(noise, predicted_noise)
+    elif loss_type == "huber":
+        loss = F.smooth_l1_loss(noise, predicted_noise)
+    else:
+        raise NotImplementedError()
+
+    return loss
+
+
+def evaluate_diffusion_model(device, model, iterator, BATCH_SIZE, timesteps, unetConfig, encodes2embeddings_mapping,
+                             uncondition_rate, unconditional_condition):
+    model.to(device)
+    model.eval()
+    eva_loss = []
+    sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, _, _ = get_beta_schedule(timesteps)
+    for i in xrange(500):
+        data, attributes = next(iter(iterator))
+        data = data.to(device)
+
+        conditions = [encodes2embeddings_mapping[create_key(attribute)] for attribute in attributes]
+        selected_conditions = [
+            unconditional_condition if random.random() < uncondition_rate else random.choice(conditions_of_one_sample)
+            for conditions_of_one_sample in conditions]
+
+        selected_conditions = torch.stack(selected_conditions).float().to(device)
+
+        t = torch.randint(0, timesteps, (BATCH_SIZE,), device=device).long()
+        loss = conditional_p_losses(model, data, t, selected_conditions, loss_type="huber",
+                                    sqrt_alphas_cumprod=sqrt_alphas_cumprod,
+                                    sqrt_one_minus_alphas_cumprod=sqrt_one_minus_alphas_cumprod)
+
+        eva_loss.append(loss.item())
+    initial_loss = np.mean(eva_loss)
+    return initial_loss
+
+
+def get_diffusion_model(model_Config, load_pretrain=False, model_name=None, device="cpu"):
+    UNet = ConditionedUnet(**model_Config)
+    print(f"Model intialized, size: {sum(p.numel() for p in UNet.parameters() if p.requires_grad)}")
+    UNet.to(device)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_UNet.pth")
+        checkpoint = torch.load(f'models/{model_name}_UNet.pth', map_location=device)
+        UNet.load_state_dict(checkpoint['model_state_dict'])
+    UNet.eval()
+    return UNet
+
+
+def train_diffusion_model(VAE, text_encoder, CLAP_tokenizer, timbre_encoder, device, init_model_name, unetConfig, BATCH_SIZE, timesteps, lr, max_iter, iterator, load_pretrain,
+                          encodes2embeddings_mapping, uncondition_rate, unconditional_condition, save_steps=5000, init_loss=None, save_model_name=None,
+                            n_IS_batches=50):
+
+    if save_model_name is None:
+        save_model_name = init_model_name
+
+    def save_model_hyperparameter(model_name, unetConfig, BATCH_SIZE, lr, model_size, current_iter, current_loss):
+        model_hyperparameter = unetConfig
+        model_hyperparameter["BATCH_SIZE"] = BATCH_SIZE
+        model_hyperparameter["lr"] = lr
+        model_hyperparameter["model_size"] = model_size
+        model_hyperparameter["current_iter"] = current_iter
+        model_hyperparameter["current_loss"] = current_loss
+        with open(f"models/hyperparameters/{model_name}_UNet.json", "w") as json_file:
+            json.dump(model_hyperparameter, json_file, ensure_ascii=False, indent=4)
+
+    model = ConditionedUnet(**unetConfig)
+    model_size = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Trainable parameters: {model_size}")
+    model.to(device)
+    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr, amsgrad=False)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{init_model_name}_UNet.pt")
+        checkpoint = torch.load(f'models/{init_model_name}_UNet.pth')
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("Model initialized.")
+    if max_iter == 0:
+        print("Return model directly.")
+        return model, optimizer
+
+
+    train_loss = []
+    writer = SummaryWriter(f'runs/{save_model_name}_UNet')
+    if init_loss is None:
+        previous_loss = evaluate_diffusion_model(device, model, iterator, BATCH_SIZE, timesteps, unetConfig, encodes2embeddings_mapping,
+                                                 uncondition_rate, unconditional_condition)
+    else:
+        previous_loss = init_loss
+    print(f"initial_IS: {previous_loss}")
+    sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, _, _ = get_beta_schedule(timesteps)
+
+    model.train()
+    for i in xrange(max_iter):
+        data, attributes = next(iter(iterator))
+        data = data.to(device)
+
+        conditions = [encodes2embeddings_mapping[create_key(attribute)] for attribute in attributes]
+        unconditional_condition_copy = torch.tensor(unconditional_condition, dtype=torch.float32).to(device).detach()
+        selected_conditions = [unconditional_condition_copy if random.random() < uncondition_rate else random.choice(
+            conditions_of_one_sample) for conditions_of_one_sample in conditions]
+
+        selected_conditions = torch.stack(selected_conditions).float().to(device)
+
+        optimizer.zero_grad()
+
+        t = torch.randint(0, timesteps, (BATCH_SIZE,), device=device).long()
+        loss = conditional_p_losses(model, data, t, selected_conditions, loss_type="huber",
+                                    sqrt_alphas_cumprod=sqrt_alphas_cumprod,
+                                    sqrt_one_minus_alphas_cumprod=sqrt_one_minus_alphas_cumprod)
+
+        loss.backward()
+        optimizer.step()
+
+        train_loss.append(loss.item())
+        step = int(optimizer.state_dict()['state'][list(optimizer.state_dict()['state'].keys())[0]]['step'].numpy())
+
+        if step % 100 == 0:
+            print('%d step' % (step))
+
+        if step % save_steps == 0:
+            current_loss = np.mean(train_loss[-save_steps:])
+            print(f"current_loss = {current_loss}")
+            torch.save({
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+            }, f'models/{save_model_name}_UNet.pth')
+            save_model_hyperparameter(save_model_name, unetConfig, BATCH_SIZE, lr, model_size, step, current_loss)
+
+
+        if step % 20000 == 0:
+            current_IS = get_inception_score(device, model, VAE, text_encoder, CLAP_tokenizer, timbre_encoder, n_IS_batches,
+                                     positive_prompts="", negative_prompts="", CFG=1, sample_steps=20, task="STFT")
+            print('current_IS: %.5f' % current_IS)
+            current_loss = np.mean(train_loss[-save_steps:])
+
+            writer.add_scalar(f"current_IS", current_IS, step)
+
+            torch.save({
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+            }, f'models/history/{save_model_name}_{step}_UNet.pth')
+            save_model_hyperparameter(save_model_name, unetConfig, BATCH_SIZE, lr, model_size, step, current_loss)
+
+    return model, optimizer
+
+

model/diffusion_components.py

@@ -0,0 +1,351 @@
+import torch.nn.functional as F
+import torch
+from torch import nn
+from einops import rearrange
+from inspect import isfunction
+import math
+from tqdm import tqdm
+
+
+def exists(x):
+    """Return true for x is not None."""
+    return x is not None
+
+
+def default(val, d):
+    """Helper function"""
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+class Residual(nn.Module):
+    """Skip connection"""
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, *args, **kwargs):
+        return self.fn(x, *args, **kwargs) + x
+
+
+def Upsample(dim):
+    """Upsample layer, a transposed convolution layer with stride=2"""
+    return nn.ConvTranspose2d(dim, dim, 4, 2, 1)
+
+
+def Downsample(dim):
+    """Downsample layer, a convolution layer with stride=2"""
+    return nn.Conv2d(dim, dim, 4, 2, 1)
+
+
+class SinusoidalPositionEmbeddings(nn.Module):
+    """Return sinusoidal embedding for integer time step."""
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, time):
+        device = time.device
+        half_dim = self.dim // 2
+        embeddings = math.log(10000) / (half_dim - 1)
+        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
+        embeddings = time[:, None] * embeddings[None, :]
+        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
+        return embeddings
+
+
+class Block(nn.Module):
+    """Stack of convolution, normalization, and non-linear activation"""
+
+    def __init__(self, dim, dim_out, groups=8):
+        super().__init__()
+        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)
+        self.norm = nn.GroupNorm(groups, dim_out)
+        self.act = nn.SiLU()
+
+    def forward(self, x, scale_shift=None):
+        x = self.proj(x)
+        x = self.norm(x)
+
+        if exists(scale_shift):
+            scale, shift = scale_shift
+            x = x * (scale + 1) + shift
+
+        x = self.act(x)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    """Stack of [conv + norm + act (+ scale&shift)], with positional embedding inserted <https://arxiv.org/abs/1512.03385>"""
+
+    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):
+        super().__init__()
+        self.mlp = (
+            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out))
+            if exists(time_emb_dim)
+            else None
+        )
+
+        self.block1 = Block(dim, dim_out, groups=groups)
+        self.block2 = Block(dim_out, dim_out, groups=groups)
+        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
+
+    def forward(self, x, time_emb=None):
+        h = self.block1(x)
+
+        if exists(self.mlp) and exists(time_emb):
+            time_emb = self.mlp(time_emb)
+            # Adding positional embedding to intermediate layer (by broadcasting along spatial dimension)
+            h = rearrange(time_emb, "b c -> b c 1 1") + h
+
+        h = self.block2(h)
+        return h + self.res_conv(x)
+
+
+class ConvNextBlock(nn.Module):
+    """Stack of [conv7x7 (+ condition(pos)) + norm + conv3x3 + act + norm + conv3x3 + res1x1]，with positional embedding inserted"""
+
+    def __init__(self, dim, dim_out, *, time_emb_dim=None, mult=2, norm=True):
+        super().__init__()
+        self.mlp = (
+            nn.Sequential(nn.GELU(), nn.Linear(time_emb_dim, dim))
+            if exists(time_emb_dim)
+            else None
+        )
+
+        self.ds_conv = nn.Conv2d(dim, dim, 7, padding=3, groups=dim)
+
+        self.net = nn.Sequential(
+            nn.GroupNorm(1, dim) if norm else nn.Identity(),
+            nn.Conv2d(dim, dim_out * mult, 3, padding=1),
+            nn.GELU(),
+            nn.GroupNorm(1, dim_out * mult),
+            nn.Conv2d(dim_out * mult, dim_out, 3, padding=1),
+        )
+
+        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
+
+    def forward(self, x, time_emb=None):
+        h = self.ds_conv(x)
+
+        if exists(self.mlp) and exists(time_emb):
+            assert exists(time_emb), "time embedding must be passed in"
+            condition = self.mlp(time_emb)
+            h = h + rearrange(condition, "b c -> b c 1 1")
+
+        h = self.net(h)
+        return h + self.res_conv(x)
+
+
+class PreNorm(nn.Module):
+    """Apply normalization before 'fn'"""
+
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.GroupNorm(1, dim)
+
+    def forward(self, x, *args, **kwargs):
+        x = self.norm(x)
+        return self.fn(x, *args, **kwargs)
+
+
+class ConditionalEmbedding(nn.Module):
+    """Return embedding for label and projection for text embedding"""
+
+    def __init__(self, num_labels, embedding_dim, condition_type="instrument_family"):
+        super(ConditionalEmbedding, self).__init__()
+        if condition_type == "instrument_family":
+            self.embedding = nn.Embedding(num_labels, embedding_dim)
+        elif condition_type == "natural_language_prompt":
+            self.embedding = nn.Linear(embedding_dim, embedding_dim, bias=True)
+        else:
+            raise NotImplementedError()
+
+    def forward(self, labels):
+        return self.embedding(labels)
+
+
+class LinearCrossAttention(nn.Module):
+    """Combination of efficient attention and cross attention."""
+
+    def __init__(self, dim, heads=4, label_emb_dim=128, dim_head=32):
+        super().__init__()
+        self.dim_head = dim_head
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), nn.GroupNorm(1, dim))
+
+        # embedding for key and value
+        self.label_key = nn.Linear(label_emb_dim, hidden_dim)
+        self.label_value = nn.Linear(label_emb_dim, hidden_dim)
+
+    def forward(self, x, label_embedding=None):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+        q, k, v = map(
+            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
+        )
+
+        if label_embedding is not None:
+            label_k = self.label_key(label_embedding).view(b, self.heads, self.dim_head, 1)
+            label_v = self.label_value(label_embedding).view(b, self.heads, self.dim_head, 1)
+
+            k = torch.cat([k, label_k], dim=-1)
+            v = torch.cat([v, label_v], dim=-1)
+
+        q = q.softmax(dim=-2)
+        k = k.softmax(dim=-1)
+        q = q * self.scale
+        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)
+        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
+        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)
+        return self.to_out(out)
+
+
+def pad_to_match(encoder_tensor, decoder_tensor):
+    """
+    Pads the decoder_tensor to match the spatial dimensions of encoder_tensor.
+
+    :param encoder_tensor: The feature map from the encoder.
+    :param decoder_tensor: The feature map from the decoder that needs to be upsampled.
+    :return: Padded decoder_tensor with the same spatial dimensions as encoder_tensor.
+    """
+
+    enc_shape = encoder_tensor.shape[2:]  # spatial dimensions are at index 2 and 3
+    dec_shape = decoder_tensor.shape[2:]
+
+    # assume enc_shape >= dec_shape
+    delta_w = enc_shape[1] - dec_shape[1]
+    delta_h = enc_shape[0] - dec_shape[0]
+
+    # padding
+    padding_left = delta_w // 2
+    padding_right = delta_w - padding_left
+    padding_top = delta_h // 2
+    padding_bottom = delta_h - padding_top
+    decoder_tensor_padded = F.pad(decoder_tensor, (padding_left, padding_right, padding_top, padding_bottom))
+
+    return decoder_tensor_padded
+
+
+def pad_and_concat(encoder_tensor, decoder_tensor):
+    """
+    Pads the decoder_tensor and concatenates it with the encoder_tensor along the channel dimension.
+
+    :param encoder_tensor: The feature map from the encoder.
+    :param decoder_tensor: The feature map from the decoder that needs to be concatenated with encoder_tensor.
+    :return: Concatenated tensor.
+    """
+
+    # pad decoder_tensor
+    decoder_tensor_padded = pad_to_match(encoder_tensor, decoder_tensor)
+    # concat encoder_tensor and decoder_tensor_padded
+    concatenated_tensor = torch.cat((encoder_tensor, decoder_tensor_padded), dim=1)
+    return concatenated_tensor
+
+
+class LinearCrossAttentionAdd(nn.Module):
+    def __init__(self, dim, heads=4, label_emb_dim=128, dim_head=32):
+        super().__init__()
+        self.dim = dim
+        self.dim_head = dim_head
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.label_emb_dim = label_emb_dim
+        self.dim_head = dim_head
+
+        self.hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(self.dim, self.hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Sequential(nn.Conv2d(self.hidden_dim, dim, 1), nn.GroupNorm(1, dim))
+
+        # embedding for key and value
+        self.label_key = nn.Linear(label_emb_dim, self.hidden_dim)
+        self.label_query = nn.Linear(label_emb_dim, self.hidden_dim)
+
+
+    def forward(self, x, condition=None):
+        b, c, h, w = x.shape
+
+        qkv = self.to_qkv(x).chunk(3, dim=1)
+
+        q, k, v = map(
+            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
+        )
+
+        # if condition exists，concat its key and value with origin
+        if condition is not None:
+            label_k = self.label_key(condition).view(b, self.heads, self.dim_head, 1)
+            label_q = self.label_query(condition).view(b, self.heads, self.dim_head, 1)
+            k = k + label_k
+            q = q + label_q
+
+        q = q.softmax(dim=-2)
+        k = k.softmax(dim=-1)
+        q = q * self.scale
+        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)
+        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
+        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)
+        return self.to_out(out)
+
+
+
+def linear_beta_schedule(timesteps):
+    beta_start = 0.0001
+    beta_end = 0.02
+    return torch.linspace(beta_start, beta_end, timesteps)
+
+
+def get_beta_schedule(timesteps):
+    betas = linear_beta_schedule(timesteps=timesteps)
+
+    # define alphas
+    alphas = 1. - betas
+    alphas_cumprod = torch.cumprod(alphas, axis=0)
+    alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+    sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
+
+    # calculations for diffusion q(x_t | x_{t-1}) and others
+    sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
+    sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
+
+    # calculations for posterior q(x_{t-1} | x_t, x_0)
+    posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+    return sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, posterior_variance, sqrt_recip_alphas
+
+
+def extract(a, t, x_shape):
+    batch_size = t.shape[0]
+    out = a.gather(-1, t.cpu())
+    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
+
+
+# forward diffusion
+def q_sample(x_start, t, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod, noise=None):
+    if noise is None:
+        noise = torch.randn_like(x_start)
+
+    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
+    sqrt_one_minus_alphas_cumprod_t = extract(
+        sqrt_one_minus_alphas_cumprod, t, x_start.shape
+    )
+
+    return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
+
+
+
+
+
+
+
+
+
+
+
+
+
+

model/multimodal_model.py

@@ -0,0 +1,274 @@
+import itertools
+import json
+import random
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from tools import create_key
+from model.timbre_encoder_pretrain import get_timbre_encoder
+
+
+class ProjectionLayer(nn.Module):
+    """Single-layer Linear projection with dropout, layer norm, and Gelu activation"""
+
+    def __init__(self, input_dim, output_dim, dropout):
+        super(ProjectionLayer, self).__init__()
+        self.projection = nn.Linear(input_dim, output_dim)
+        self.gelu = nn.GELU()
+        self.fc = nn.Linear(output_dim, output_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(output_dim)
+
+    def forward(self, x):
+        projected = self.projection(x)
+        x = self.gelu(projected)
+        x = self.fc(x)
+        x = self.dropout(x)
+        x = x + projected
+        x = self.layer_norm(x)
+        return x
+
+
+class ProjectionHead(nn.Module):
+    """Stack of 'ProjectionLayer'"""
+
+    def __init__(self, embedding_dim, projection_dim, dropout, num_layers=2):
+        super(ProjectionHead, self).__init__()
+        self.layers = nn.ModuleList([ProjectionLayer(embedding_dim if i == 0 else projection_dim,
+                                                     projection_dim,
+                                                     dropout) for i in range(num_layers)])
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class multi_modal_model(nn.Module):
+    """The multi-modal model for contrastive learning"""
+
+    def __init__(
+            self,
+            timbre_encoder,
+            text_encoder,
+            spectrogram_feature_dim,
+            text_feature_dim,
+            multi_modal_emb_dim,
+            temperature,
+            dropout,
+            num_projection_layers=1,
+            freeze_spectrogram_encoder=True,
+            freeze_text_encoder=True,
+    ):
+        super().__init__()
+        self.timbre_encoder = timbre_encoder
+        self.text_encoder = text_encoder
+
+        self.multi_modal_emb_dim = multi_modal_emb_dim
+
+        self.text_projection = ProjectionHead(embedding_dim=text_feature_dim,
+                                              projection_dim=self.multi_modal_emb_dim, dropout=dropout,
+                                              num_layers=num_projection_layers)
+
+        self.spectrogram_projection = ProjectionHead(embedding_dim=spectrogram_feature_dim,
+                                                     projection_dim=self.multi_modal_emb_dim, dropout=dropout,
+                                                     num_layers=num_projection_layers)
+
+        self.temperature = temperature
+
+        # Make spectrogram_encoder parameters non-trainable
+        for param in self.timbre_encoder.parameters():
+            param.requires_grad = not freeze_spectrogram_encoder
+
+        # Make text_encoder parameters non-trainable
+        for param in self.text_encoder.parameters():
+            param.requires_grad = not freeze_text_encoder
+
+    def forward(self, spectrogram_batch, tokenized_text_batch):
+        # Getting Image and Text Embeddings (with same dimension)
+        spectrogram_features, _, _, _, _ = self.timbre_encoder(spectrogram_batch)
+        text_features = self.text_encoder.get_text_features(**tokenized_text_batch)
+
+        # Concat and apply projection
+        spectrogram_embeddings = self.spectrogram_projection(spectrogram_features)
+        text_embeddings = self.text_projection(text_features)
+
+        # Calculating the Loss
+        logits = (text_embeddings @ spectrogram_embeddings.T) / self.temperature
+        images_similarity = spectrogram_embeddings @ spectrogram_embeddings.T
+        texts_similarity = text_embeddings @ text_embeddings.T
+        targets = F.softmax(
+            (images_similarity + texts_similarity) / 2 * self.temperature, dim=-1
+        )
+        texts_loss = cross_entropy(logits, targets, reduction='none')
+        images_loss = cross_entropy(logits.T, targets.T, reduction='none')
+        contrastive_loss = (images_loss + texts_loss) / 2.0  # shape: (batch_size)
+        contrastive_loss = contrastive_loss.mean()
+
+        return contrastive_loss
+
+
+    def get_text_features(self, input_ids, attention_mask):
+        text_features = self.text_encoder.get_text_features(input_ids=input_ids, attention_mask=attention_mask)
+        return self.text_projection(text_features)
+
+
+    def get_timbre_features(self, spectrogram_batch):
+        spectrogram_features, _, _, _, _ = self.timbre_encoder(spectrogram_batch)
+        return self.spectrogram_projection(spectrogram_features)
+
+
+def cross_entropy(preds, targets, reduction='none'):
+    log_softmax = nn.LogSoftmax(dim=-1)
+    loss = (-targets * log_softmax(preds)).sum(1)
+    if reduction == "none":
+        return loss
+    elif reduction == "mean":
+        return loss.mean()
+
+
+def get_multi_modal_model(timbre_encoder, text_encoder, model_Config, load_pretrain=False, model_name=None, device="cpu"):
+    mmm = multi_modal_model(timbre_encoder, text_encoder, **model_Config)
+    print(f"Model intialized, size: {sum(p.numel() for p in mmm.parameters() if p.requires_grad)}")
+    mmm.to(device)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_MMM.pth")
+        checkpoint = torch.load(f'models/{model_name}_MMM.pth', map_location=device)
+        mmm.load_state_dict(checkpoint['model_state_dict'])
+    mmm.eval()
+    return mmm
+
+
+def train_epoch(text_tokenizer, model, train_loader, labels_mapping, optimizer, device):
+    (data, attributes) = next(iter(train_loader))
+    keys = [create_key(attribute) for attribute in attributes]
+
+    while(len(set(keys)) != len(keys)):
+        (data, attributes) = next(iter(train_loader))
+        keys = [create_key(attribute) for attribute in attributes]
+
+    data = data.to(device)
+
+    texts = [labels_mapping[create_key(attribute)] for attribute in attributes]
+    selected_texts = [l[random.randint(0, len(l) - 1)] for l in texts]
+
+    tokenized_text = text_tokenizer(selected_texts, padding=True, return_tensors="pt").to(device)
+
+    loss = model(data, tokenized_text)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    return loss.item()
+
+
+def valid_epoch(text_tokenizer, model, valid_loader, labels_mapping, device):
+    (data, attributes) = next(iter(valid_loader))
+    keys = [create_key(attribute) for attribute in attributes]
+
+    while(len(set(keys)) != len(keys)):
+        (data, attributes) = next(iter(valid_loader))
+        keys = [create_key(attribute) for attribute in attributes]
+
+    data = data.to(device)
+    texts = [labels_mapping[create_key(attribute)] for attribute in attributes]
+    selected_texts = [l[random.randint(0, len(l) - 1)] for l in texts]
+
+    tokenized_text = text_tokenizer(selected_texts, padding=True, return_tensors="pt").to(device)
+
+    loss = model(data, tokenized_text)
+    return loss.item()
+
+
+def train_multi_modal_model(device, training_dataloader, labels_mapping, text_tokenizer, text_encoder,
+                            timbre_encoder_Config, MMM_config, MMM_training_config,
+                            mmm_name, BATCH_SIZE, max_iter=0, load_pretrain=True,
+                            timbre_encoder_name=None, init_loss=None, save_steps=2000):
+
+    def save_model_hyperparameter(model_name, MMM_config, MMM_training_config, BATCH_SIZE, model_size, current_iter,
+                                  current_loss):
+
+        model_hyperparameter = MMM_config
+        model_hyperparameter.update(MMM_training_config)
+        model_hyperparameter["BATCH_SIZE"] = BATCH_SIZE
+        model_hyperparameter["model_size"] = model_size
+        model_hyperparameter["current_iter"] = current_iter
+        model_hyperparameter["current_loss"] = current_loss
+        with open(f"models/hyperparameters/{model_name}_MMM.json", "w") as json_file:
+            json.dump(model_hyperparameter, json_file, ensure_ascii=False, indent=4)
+
+    timbreEncoder = get_timbre_encoder(timbre_encoder_Config, load_pretrain=True, model_name=timbre_encoder_name,
+                                       device=device)
+
+    mmm = multi_modal_model(timbreEncoder, text_encoder, **MMM_config).to(device)
+
+    print(f"spectrogram_encoder parameter: {sum(p.numel() for p in mmm.timbre_encoder.parameters())}")
+    print(f"text_encoder parameter: {sum(p.numel() for p in mmm.text_encoder.parameters())}")
+    print(f"spectrogram_projection parameter: {sum(p.numel() for p in mmm.spectrogram_projection.parameters())}")
+    print(f"text_projection parameter: {sum(p.numel() for p in mmm.text_projection.parameters())}")
+    total_parameters = sum(p.numel() for p in mmm.parameters())
+    trainable_parameters = sum(p.numel() for p in mmm.parameters() if p.requires_grad)
+    print(f"Trainable/Total parameter: {trainable_parameters}/{total_parameters}")
+
+    params = [
+        {"params": itertools.chain(
+            mmm.spectrogram_projection.parameters(),
+            mmm.text_projection.parameters(),
+        ), "lr": MMM_training_config["head_lr"], "weight_decay": MMM_training_config["head_weight_decay"]},
+    ]
+    if not MMM_config["freeze_text_encoder"]:
+        params.append({"params": mmm.text_encoder.parameters(), "lr": MMM_training_config["text_encoder_lr"],
+                       "weight_decay": MMM_training_config["text_encoder_weight_decay"]})
+    if not MMM_config["freeze_spectrogram_encoder"]:
+        params.append({"params": mmm.timbre_encoder.parameters(), "lr": MMM_training_config["spectrogram_encoder_lr"],
+                       "weight_decay": MMM_training_config["timbre_encoder_weight_decay"]})
+
+    optimizer = torch.optim.AdamW(params, weight_decay=0.)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{mmm_name}_MMM.pt")
+        checkpoint = torch.load(f'models/{mmm_name}_MMM.pth')
+        mmm.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("Model initialized.")
+
+    if max_iter == 0:
+        print("Return model directly.")
+        return mmm, optimizer
+
+    if init_loss is None:
+        previous_lowest_loss = valid_epoch(text_tokenizer, mmm, training_dataloader, labels_mapping, device)
+    else:
+        previous_lowest_loss = init_loss
+    print(f"Initial total loss: {previous_lowest_loss}")
+
+    train_loss_list = []
+    for i in range(max_iter):
+
+        mmm.train()
+        train_loss = train_epoch(text_tokenizer, mmm, training_dataloader, labels_mapping, optimizer, device)
+        train_loss_list.append(train_loss)
+
+        step = int(
+            optimizer.state_dict()['state'][list(optimizer.state_dict()['state'].keys())[0]]['step'].cpu().numpy())
+        if (i + 1) % 100 == 0:
+            print('%d step' % (step))
+
+        if (i + 1) % save_steps == 0:
+            current_loss = np.mean(train_loss_list[-save_steps:])
+            print(f"train_total_loss: {current_loss}")
+            if current_loss < previous_lowest_loss:
+                previous_lowest_loss = current_loss
+                torch.save({
+                    'model_state_dict': mmm.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                }, f'models/{mmm_name}_MMM.pth')
+                save_model_hyperparameter(mmm_name, MMM_config, MMM_training_config, BATCH_SIZE, total_parameters, step,
+                                          current_loss)
+
+    return mmm, optimizer

model/timbre_encoder_pretrain.py

@@ -0,0 +1,220 @@
+import json
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.tensorboard import SummaryWriter
+from tools import create_key
+
+
+class TimbreEncoder(nn.Module):
+    def __init__(self, input_dim, feature_dim, hidden_dim, num_instrument_classes, num_instrument_family_classes, num_velocity_classes, num_qualities, num_layers=1):
+        super(TimbreEncoder, self).__init__()
+
+        # Input layer
+        self.input_layer = nn.Linear(input_dim, feature_dim)
+
+        # LSTM Layer
+        self.lstm = nn.LSTM(feature_dim, hidden_dim, num_layers=num_layers, batch_first=True)
+
+        # Fully Connected Layers for classification
+        self.instrument_classifier_layer = nn.Linear(hidden_dim, num_instrument_classes)
+        self.instrument_family_classifier_layer = nn.Linear(hidden_dim, num_instrument_family_classes)
+        self.velocity_classifier_layer = nn.Linear(hidden_dim, num_velocity_classes)
+        self.qualities_classifier_layer = nn.Linear(hidden_dim, num_qualities)
+
+        # Softmax for converting output to probabilities
+        self.softmax = nn.LogSoftmax(dim=1)
+
+    def forward(self, x):
+        # # Merge first two dimensions
+        batch_size, _, _, seq_len = x.shape
+        x = x.view(batch_size, -1, seq_len)  # [batch_size, input_dim, seq_len]
+
+        # Forward propagate LSTM
+        x = x.permute(0, 2, 1)
+        x = self.input_layer(x)
+        feature, _ = self.lstm(x)
+        feature = feature[:, -1, :]
+
+        # Apply classification layers
+        instrument_logits = self.instrument_classifier_layer(feature)
+        instrument_family_logits = self.instrument_family_classifier_layer(feature)
+        velocity_logits = self.velocity_classifier_layer(feature)
+        qualities = self.qualities_classifier_layer(feature)
+
+        # Apply Softmax
+        instrument_logits = self.softmax(instrument_logits)
+        instrument_family_logits= self.softmax(instrument_family_logits)
+        velocity_logits = self.softmax(velocity_logits)
+        qualities = torch.sigmoid(qualities)
+
+        return feature, instrument_logits, instrument_family_logits, velocity_logits, qualities
+
+
+def get_multiclass_acc(outputs, ground_truth):
+    _, predicted = torch.max(outputs.data, 1)
+    total = ground_truth.size(0)
+    correct = (predicted == ground_truth).sum().item()
+    accuracy = 100 * correct / total
+    return accuracy
+
+def get_binary_accuracy(y_pred, y_true):
+    predictions = (y_pred > 0.5).int()
+
+    correct_predictions = (predictions == y_true).float()
+
+    accuracy = correct_predictions.mean()
+
+    return accuracy.item() * 100.0
+
+
+def get_timbre_encoder(model_Config, load_pretrain=False, model_name=None, device="cpu"):
+    timbreEncoder = TimbreEncoder(**model_Config)
+    print(f"Model intialized, size: {sum(p.numel() for p in timbreEncoder.parameters() if p.requires_grad)}")
+    timbreEncoder.to(device)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_timbre_encoder.pth")
+        checkpoint = torch.load(f'models/{model_name}_timbre_encoder.pth', map_location=device)
+        timbreEncoder.load_state_dict(checkpoint['model_state_dict'])
+    timbreEncoder.eval()
+    return timbreEncoder
+
+
+def evaluate_timbre_encoder(device, model, iterator, nll_Loss, bce_Loss, n_sample=100):
+    model.to(device)
+    model.eval()
+
+    eva_loss = []
+    for i in range(n_sample):
+        representation, attributes = next(iter(iterator))
+
+        instrument = torch.tensor([s["instrument"] for s in attributes], dtype=torch.long).to(device)
+        instrument_family = torch.tensor([s["instrument_family"] for s in attributes], dtype=torch.long).to(device)
+        velocity = torch.tensor([s["velocity"] for s in attributes], dtype=torch.long).to(device)
+        qualities = torch.tensor([[int(char) for char in create_key(attribute)[-10:]] for attribute in attributes], dtype=torch.float32).to(device)
+
+        _, instrument_logits, instrument_family_logits, velocity_logits, qualities_pred = model(representation.to(device))
+
+        # compute loss
+        instrument_loss = nll_Loss(instrument_logits, instrument)
+        instrument_family_loss = nll_Loss(instrument_family_logits, instrument_family)
+        velocity_loss = nll_Loss(velocity_logits, velocity)
+        qualities_loss = bce_Loss(qualities_pred, qualities)
+
+        loss = instrument_loss + instrument_family_loss + velocity_loss + qualities_loss
+
+        eva_loss.append(loss.item())
+
+    eva_loss = np.mean(eva_loss)
+    return eva_loss
+
+
+def train_timbre_encoder(device, model_name, timbre_encoder_Config, BATCH_SIZE, lr, max_iter, training_iterator, load_pretrain):
+    def save_model_hyperparameter(model_name, timbre_encoder_Config, BATCH_SIZE, lr, model_size, current_iter,
+                                  current_loss):
+        model_hyperparameter = timbre_encoder_Config
+        model_hyperparameter["BATCH_SIZE"] = BATCH_SIZE
+        model_hyperparameter["lr"] = lr
+        model_hyperparameter["model_size"] = model_size
+        model_hyperparameter["current_iter"] = current_iter
+        model_hyperparameter["current_loss"] = current_loss
+        with open(f"models/hyperparameters/{model_name}_timbre_encoder.json", "w") as json_file:
+            json.dump(model_hyperparameter, json_file, ensure_ascii=False, indent=4)
+
+    model = TimbreEncoder(**timbre_encoder_Config)
+    model_size = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Model size: {model_size}")
+    model.to(device)
+    nll_Loss = torch.nn.NLLLoss()
+    bce_Loss = torch.nn.BCELoss()
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, amsgrad=False)
+
+    if load_pretrain:
+        print(f"Loading weights from models/{model_name}_timbre_encoder.pt")
+        checkpoint = torch.load(f'models/{model_name}_timbre_encoder.pth')
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    else:
+        print("Model initialized.")
+    if max_iter == 0:
+        print("Return model directly.")
+        return model, model
+
+    train_loss, training_instrument_acc, training_instrument_family_acc, training_velocity_acc, training_qualities_acc = [], [], [], [], []
+    writer = SummaryWriter(f'runs/{model_name}_timbre_encoder')
+    current_best_model = model
+    previous_lowest_loss = 100.0
+    print(f"initial__loss: {previous_lowest_loss}")
+
+    for i in range(max_iter):
+        model.train()
+
+        representation, attributes = next(iter(training_iterator))
+
+        instrument = torch.tensor([s["instrument"] for s in attributes], dtype=torch.long).to(device)
+        instrument_family = torch.tensor([s["instrument_family"] for s in attributes], dtype=torch.long).to(device)
+        velocity = torch.tensor([s["velocity"] for s in attributes], dtype=torch.long).to(device)
+        qualities = torch.tensor([[int(char) for char in create_key(attribute)[-10:]] for attribute in attributes], dtype=torch.float32).to(device)
+
+        optimizer.zero_grad()
+
+        _, instrument_logits, instrument_family_logits, velocity_logits, qualities_pred = model(representation.to(device))
+
+        # compute loss
+        instrument_loss = nll_Loss(instrument_logits, instrument)
+        instrument_family_loss = nll_Loss(instrument_family_logits, instrument_family)
+        velocity_loss = nll_Loss(velocity_logits, velocity)
+        qualities_loss = bce_Loss(qualities_pred, qualities)
+
+        loss = instrument_loss + instrument_family_loss + velocity_loss + qualities_loss
+
+        loss.backward()
+        optimizer.step()
+        instrument_acc = get_multiclass_acc(instrument_logits, instrument)
+        instrument_family_acc = get_multiclass_acc(instrument_family_logits, instrument_family)
+        velocity_acc = get_multiclass_acc(velocity_logits, velocity)
+        qualities_acc = get_binary_accuracy(qualities_pred, qualities)
+
+        train_loss.append(loss.item())
+        training_instrument_acc.append(instrument_acc)
+        training_instrument_family_acc.append(instrument_family_acc)
+        training_velocity_acc.append(velocity_acc)
+        training_qualities_acc.append(qualities_acc)
+        step = int(optimizer.state_dict()['state'][list(optimizer.state_dict()['state'].keys())[0]]['step'].numpy())
+
+        if (i + 1) % 100 == 0:
+            print('%d step' % (step))
+
+        save_steps = 500
+        if (i + 1) % save_steps == 0:
+            current_loss = np.mean(train_loss[-save_steps:])
+            current_instrument_acc = np.mean(training_instrument_acc[-save_steps:])
+            current_instrument_family_acc = np.mean(training_instrument_family_acc[-save_steps:])
+            current_velocity_acc = np.mean(training_velocity_acc[-save_steps:])
+            current_qualities_acc = np.mean(training_qualities_acc[-save_steps:])
+            print('train_loss: %.5f' % current_loss)
+            print('current_instrument_acc: %.5f' % current_instrument_acc)
+            print('current_instrument_family_acc: %.5f' % current_instrument_family_acc)
+            print('current_velocity_acc: %.5f' % current_velocity_acc)
+            print('current_qualities_acc: %.5f' % current_qualities_acc)
+            writer.add_scalar(f"train_loss", current_loss, step)
+            writer.add_scalar(f"current_instrument_acc", current_instrument_acc, step)
+            writer.add_scalar(f"current_instrument_family_acc", current_instrument_family_acc, step)
+            writer.add_scalar(f"current_velocity_acc", current_velocity_acc, step)
+            writer.add_scalar(f"current_qualities_acc", current_qualities_acc, step)
+
+            if current_loss < previous_lowest_loss:
+                previous_lowest_loss = current_loss
+                current_best_model = model
+                torch.save({
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                }, f'models/{model_name}_timbre_encoder.pth')
+                save_model_hyperparameter(model_name, timbre_encoder_Config, BATCH_SIZE, lr, model_size, step,
+                                          current_loss)
+
+    return model, current_best_model
+
+

models/24_1_2024-52_4x_L_D_imageVQVAE.pth

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

models/24_1_2024-52_4x_L_D_imageVQVAE_discriminator.pth

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

models/24_1_2024_MMM.pth

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

models/24_1_2024_STFT_timbre_encoder.pth

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

models/history/28_1_2024_TE_STFT_300000_UNet.pth

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

requirements.txt

@@ -0,0 +1,15 @@
+torchmetrics==0.7.0
+torchsynth==1.0.2
+torchaudio
+soundfile
+einops
+pytorch-ssim
+piqa
+torchinfo
+mido
+tensorboard
+librosa
+transformers
+matplotlib
+gradio==3.50.2
+

tools.py

@@ -0,0 +1,344 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib
+import librosa
+from scipy.io.wavfile import write
+import torch
+
+k = 1e-16
+
+def np_log10(x):
+    """Safe log function with base 10."""
+    numerator = np.log(x + 1e-16)
+    denominator = np.log(10)
+    return numerator / denominator
+
+
+def sigmoid(x):
+    """Safe log function with base 10."""
+    s = 1 / (1 + np.exp(-x))
+    return s
+
+
+def inv_sigmoid(s):
+    """Safe inverse sigmoid function."""
+    x = np.log((s / (1 - s)) + 1e-16)
+    return x
+
+
+def spc_to_VAE_input(spc):
+    """Restrict value range from [0, infinite] to [0, 1]. (deprecated )"""
+    return spc / (1 + spc)
+
+
+def VAE_out_put_to_spc(o):
+    """Inverse transform of function 'spc_to_VAE_input'. (deprecated )"""
+    return o / (1 - o + k)
+
+
+
+def np_power_to_db(S, amin=1e-16, top_db=80.0):
+    """Helper method for numpy data scaling. (deprecated )"""
+    ref = S.max()
+
+    log_spec = 10.0 * np_log10(np.maximum(amin, S))
+    log_spec -= 10.0 * np_log10(np.maximum(amin, ref))
+
+    log_spec = np.maximum(log_spec, log_spec.max() - top_db)
+
+    return log_spec
+
+
+def show_spc(spc):
+    """Show a spectrogram. (deprecated )"""
+    s = np.shape(spc)
+    spc = np.reshape(spc, (s[0], s[1]))
+    magnitude_spectrum = np.abs(spc)
+    log_spectrum = np_power_to_db(magnitude_spectrum)
+    plt.imshow(np.flipud(log_spectrum))
+    plt.show()
+
+
+def save_results(spectrogram, spectrogram_image_path, waveform_path):
+    """Save the input 'spectrogram' and its waveform (reconstructed by Griffin Lim)
+     to path provided by 'spectrogram_image_path' and 'waveform_path'."""
+    magnitude_spectrum = np.abs(spectrogram)
+    log_spc = np_power_to_db(magnitude_spectrum)
+    log_spc = np.reshape(log_spc, (512, 256))
+    matplotlib.pyplot.imsave(spectrogram_image_path, log_spc, vmin=-100, vmax=0,
+                             origin='lower')
+
+    # save waveform
+    abs_spec = np.zeros((513, 256))
+    abs_spec[:512, :] = abs_spec[:512, :] + np.sqrt(np.reshape(spectrogram, (512, 256)))
+    rec_signal = librosa.griffinlim(abs_spec, n_iter=32, hop_length=256, win_length=1024)
+    write(waveform_path, 16000, rec_signal)
+
+
+def plot_log_spectrogram(signal: np.ndarray,
+                         path: str,
+                         n_fft=2048,
+                         frame_length=1024,
+                         frame_step=256):
+    """Save spectrogram."""
+    stft = librosa.stft(signal, n_fft=n_fft, hop_length=frame_step, win_length=frame_length)
+    amp = np.square(np.real(stft)) + np.square(np.imag(stft))
+    magnitude_spectrum = np.abs(amp)
+    log_mel = np_power_to_db(magnitude_spectrum)
+    matplotlib.pyplot.imsave(path, log_mel, vmin=-100, vmax=0, origin='lower')
+
+
+def visualize_feature_maps(device, model, inputs, channel_indices=[0, 3,]):
+    """
+    Visualize feature maps before and after quantization for given input.
+
+    Parameters:
+    - model: Your VQ-VAE model.
+    - inputs: A batch of input data.
+    - channel_indices: Indices of feature map channels to visualize.
+    """
+    model.eval()
+    inputs = inputs.to(device)
+
+    with torch.no_grad():
+        z_e = model._encoder(inputs)
+        z_q, loss, (perplexity, min_encodings, min_encoding_indices) = model._vq_vae(z_e)
+
+    # Assuming inputs have shape [batch_size, channels, height, width]
+    batch_size = z_e.size(0)
+
+    for idx in range(batch_size):
+        fig, axs = plt.subplots(1, len(channel_indices)*2, figsize=(15, 5))
+
+        for i, channel_idx in enumerate(channel_indices):
+            # Plot encoder output
+            axs[2*i].imshow(z_e[idx][channel_idx].cpu().numpy(), cmap='viridis')
+            axs[2*i].set_title(f"Encoder Output - Channel {channel_idx}")
+
+            # Plot quantized output
+            axs[2*i+1].imshow(z_q[idx][channel_idx].cpu().numpy(), cmap='viridis')
+            axs[2*i+1].set_title(f"Quantized Output - Channel {channel_idx}")
+
+        plt.show()
+
+
+def adjust_audio_length(audio, desired_length, original_sample_rate, target_sample_rate):
+    """
+    Adjust the audio length to the desired length and resample to target sample rate.
+
+    Parameters:
+    - audio (np.array): The input audio signal
+    - desired_length (int): The desired length of the output audio
+    - original_sample_rate (int): The original sample rate of the audio
+    - target_sample_rate (int): The target sample rate for the output audio
+
+    Returns:
+    - np.array: The adjusted and resampled audio
+    """
+
+    if not (original_sample_rate == target_sample_rate):
+        audio = librosa.core.resample(audio, orig_sr=original_sample_rate, target_sr=target_sample_rate)
+
+    if len(audio) > desired_length:
+        return audio[:desired_length]
+
+    elif len(audio) < desired_length:
+        padded_audio = np.zeros(desired_length)
+        padded_audio[:len(audio)] = audio
+        return padded_audio
+    else:
+        return audio
+
+
+def safe_int(s, default=0):
+    try:
+        return int(s)
+    except ValueError:
+        return default
+
+
+def pad_spectrogram(D):
+    """Resize spectrogram to (512, 256). (deprecated )"""
+    D = D[1:, :]
+
+    padding_length = 256 - D.shape[1]
+    D_padded = np.pad(D, ((0, 0), (0, padding_length)), 'constant')
+    return D_padded
+
+
+def pad_STFT(D, time_resolution=256):
+    """Resize spectral matrix by padding and cropping"""
+    D = D[1:, :]
+
+    if time_resolution is None:
+        return D
+
+    padding_length = time_resolution - D.shape[1]
+    if padding_length > 0:
+        D_padded = np.pad(D, ((0, 0), (0, padding_length)), 'constant')
+        return D_padded
+    else:
+        return D
+
+
+def depad_STFT(D_padded):
+    """Inverse function of 'pad_STFT'"""
+    zero_row = np.zeros((1, D_padded.shape[1]))
+
+    D_restored = np.concatenate([zero_row, D_padded], axis=0)
+
+    return D_restored
+
+
+def nnData2Audio(spectrogram_batch, resolution=(512, 256), squared=False):
+    """Transform batch of numpy spectrogram into signals and encodings."""
+    # Todo: remove resolution hard-coding
+    frequency_resolution, time_resolution = resolution
+
+    if isinstance(spectrogram_batch, torch.Tensor):
+        spectrogram_batch = spectrogram_batch.to("cpu").detach().numpy()
+
+    origin_signals = []
+    for spectrogram in spectrogram_batch:
+        spc = VAE_out_put_to_spc(spectrogram)
+
+        # get_audio
+        abs_spec = np.zeros((frequency_resolution+1, time_resolution))
+
+        if squared:
+            abs_spec[1:, :] = abs_spec[1:, :] + np.sqrt(np.reshape(spc, (frequency_resolution, time_resolution)))
+        else:
+            abs_spec[1:, :] = abs_spec[1:, :] + np.reshape(spc, (frequency_resolution, time_resolution))
+
+        origin_signal = librosa.griffinlim(abs_spec, n_iter=32, hop_length=256, win_length=1024)
+        origin_signals.append(origin_signal)
+
+    return origin_signals
+
+
+def amp_to_audio(amp, n_iter=50):
+    """The Griffin-Lim algorithm."""
+    y_reconstructed = librosa.griffinlim(amp, n_iter=n_iter, hop_length=256, win_length=1024)
+    return y_reconstructed
+
+
+def rescale(amp, method="log1p"):
+    """Rescale function."""
+    if method == "log1p":
+        return np.log1p(amp)
+    elif method == "NormalizedLogisticCompression":
+        return amp / (1.0 + amp)
+    else:
+        raise NotImplementedError()
+
+
+def unrescale(scaled_amp, method="NormalizedLogisticCompression"):
+    """Inverse function of 'rescale'"""
+    if method == "log1p":
+        return np.expm1(scaled_amp)
+    elif method == "NormalizedLogisticCompression":
+        return scaled_amp / (1.0 - scaled_amp + 1e-10)
+    else:
+        raise NotImplementedError()
+
+
+def create_key(attributes):
+    """Create unique key for each multi-label."""
+    qualities_str = ''.join(map(str, attributes["qualities"]))
+    instrument_source_str = attributes["instrument_source_str"]
+    instrument_family = attributes["instrument_family_str"]
+    key = f"{instrument_source_str}_{instrument_family}_{qualities_str}"
+    return key
+
+
+def merge_dictionaries(dicts):
+    """Merge dictionaries."""
+    merged_dict = {}
+    for dictionary in dicts:
+        for key, value in dictionary.items():
+            if key in merged_dict:
+                merged_dict[key] += value
+            else:
+                merged_dict[key] = value
+    return merged_dict
+
+
+def adsr_envelope(signal, sample_rate, duration, attack_time, decay_time, sustain_level, release_time):
+    """
+    Apply an ADSR envelope to an audio signal.
+
+    :param signal: The original audio signal (numpy array).
+    :param sample_rate: The sample rate of the audio signal.
+    :param attack_time: Attack time in seconds.
+    :param decay_time: Decay time in seconds.
+    :param sustain_level: Sustain level as a fraction of the peak (0 to 1).
+    :param release_time: Release time in seconds.
+    :return: The audio signal with the ADSR envelope applied.
+    """
+    # Calculate the number of samples for each ADSR phase
+    duration_samples = int(duration * sample_rate)
+
+    # assert (duration_samples + int(1.0 * sample_rate)) <= len(signal), "(duration_samples + sample_rate) > len(signal)"
+    assert release_time <= 1.0, "release_time > 1.0"
+
+    attack_samples = int(attack_time * sample_rate)
+    decay_samples = int(decay_time * sample_rate)
+    release_samples = int(release_time * sample_rate)
+    sustain_samples = max(0, duration_samples - attack_samples - decay_samples)
+
+    # Create ADSR envelope
+    attack_env = np.linspace(0, 1, attack_samples)
+    decay_env = np.linspace(1, sustain_level, decay_samples)
+    sustain_env = np.full(sustain_samples, sustain_level)
+    release_env = np.linspace(sustain_level, 0, release_samples)
+    release_env_expand = np.zeros(int(1.0 * sample_rate))
+    release_env_expand[:len(release_env)] = release_env
+
+    # Concatenate all phases to create the complete envelope
+    envelope = np.concatenate([attack_env, decay_env, sustain_env, release_env_expand])
+
+    # Apply the envelope to the signal
+    if len(envelope) <= len(signal):
+        applied_signal = signal[:len(envelope)] * envelope
+    else:
+        signal_expanded = np.zeros(len(envelope))
+        signal_expanded[:len(signal)] = signal
+        applied_signal = signal_expanded * envelope
+
+    return applied_signal
+
+
+def rms_normalize(audio, target_rms=0.1):
+    """Normalize the RMS value."""
+    current_rms = np.sqrt(np.mean(audio**2))
+    scaling_factor = target_rms / current_rms
+    normalized_audio = audio * scaling_factor
+    return normalized_audio
+
+
+def encode_stft(D):
+    """'STFT+' function that transform spectral matrix into spectral representation."""
+    magnitude = np.abs(D)
+    phase = np.angle(D)
+
+    log_magnitude = np.log1p(magnitude)
+
+    cos_phase = np.cos(phase)
+    sin_phase = np.sin(phase)
+
+    encoded_D = np.stack([log_magnitude, cos_phase, sin_phase], axis=0)
+    return encoded_D
+
+
+def decode_stft(encoded_D):
+    """'ISTFT+' function that reconstructs spectral matrix from spectral representation."""
+    log_magnitude = encoded_D[0, ...]
+    cos_phase = encoded_D[1, ...]
+    sin_phase = encoded_D[2, ...]
+
+    magnitude = np.expm1(log_magnitude)
+
+    phase = np.arctan2(sin_phase, cos_phase)
+
+    D = magnitude * (np.cos(phase) + 1j * np.sin(phase))
+    return D

webUI/deprecated/interpolationWithCondition.py

@@ -0,0 +1,178 @@
+import gradio as gr
+import numpy as np
+import torch
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import safe_int
+from webUI.natural_language_guided_STFT.utils import encodeBatch2GradioOutput, latent_representation_to_Gradio_image
+
+
+def get_interpolation_with_condition_module(gradioWebUI, interpolation_with_text_state):
+    # Load configurations
+    uNet = gradioWebUI.uNet
+    freq_resolution, time_resolution = gradioWebUI.freq_resolution, gradioWebUI.time_resolution
+    VAE_scale = gradioWebUI.VAE_scale
+    height, width, channels = int(freq_resolution/VAE_scale), int(time_resolution/VAE_scale), gradioWebUI.channels
+    timesteps = gradioWebUI.timesteps
+    VAE_quantizer = gradioWebUI.VAE_quantizer
+    VAE_decoder = gradioWebUI.VAE_decoder
+    CLAP = gradioWebUI.CLAP
+    CLAP_tokenizer = gradioWebUI.CLAP_tokenizer
+    device = gradioWebUI.device
+    squared = gradioWebUI.squared
+    sample_rate = gradioWebUI.sample_rate
+    noise_strategy = gradioWebUI.noise_strategy
+
+    def diffusion_random_sample(text2sound_prompts_1, text2sound_prompts_2, text2sound_negative_prompts, text2sound_batchsize,
+                                text2sound_duration,
+                                text2sound_guidance_scale, text2sound_sampler,
+                                text2sound_sample_steps, text2sound_seed,
+                                interpolation_with_text_dict):
+        text2sound_sample_steps = int(text2sound_sample_steps)
+        text2sound_seed = safe_int(text2sound_seed, 12345678)
+        # Todo: take care of text2sound_time_resolution/width
+        width = int(time_resolution*((text2sound_duration+1)/4) / VAE_scale)
+        text2sound_batchsize = int(text2sound_batchsize)
+
+        text2sound_embedding_1 = \
+        CLAP.get_text_features(**CLAP_tokenizer([text2sound_prompts_1], padding=True, return_tensors="pt"))[0].to(device)
+        text2sound_embedding_2 = \
+            CLAP.get_text_features(**CLAP_tokenizer([text2sound_prompts_2], padding=True, return_tensors="pt"))[0].to(device)
+
+        CFG = int(text2sound_guidance_scale)
+
+        mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy)
+        unconditional_condition = \
+        CLAP.get_text_features(**CLAP_tokenizer([text2sound_negative_prompts], padding=True, return_tensors="pt"))[0]
+        mySampler.activate_classifier_free_guidance(CFG, unconditional_condition.to(device))
+
+        mySampler.respace(list(np.linspace(0, timesteps - 1, text2sound_sample_steps, dtype=np.int32)))
+
+        condition = torch.linspace(1, 0, steps=text2sound_batchsize).unsqueeze(1).to(device) * text2sound_embedding_1 + \
+                    torch.linspace(0, 1, steps=text2sound_batchsize).unsqueeze(1).to(device) * text2sound_embedding_2
+
+        # Todo: move this code
+        torch.manual_seed(text2sound_seed)
+        initial_noise = torch.randn(text2sound_batchsize, channels, height, width).to(device)
+
+        latent_representations, initial_noise = \
+        mySampler.sample(model=uNet, shape=(text2sound_batchsize, channels, height, width), seed=text2sound_seed,
+                         return_tensor=True, condition=condition, sampler=text2sound_sampler, initial_noise=initial_noise)
+
+        latent_representations = latent_representations[-1]
+
+        interpolation_with_text_dict["latent_representations"] = latent_representations
+
+        latent_representation_gradio_images = []
+        quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_rec_signals_gradio = []
+
+        quantized_latent_representations, loss, (_, _, _) = VAE_quantizer(latent_representations)
+        # Todo: remove hard-coding
+        flipped_log_spectrums, rec_signals = encodeBatch2GradioOutput(VAE_decoder, quantized_latent_representations,
+                                                                          resolution=(512, width * VAE_scale), centralized=False,
+                                                                          squared=squared)
+
+        for i in range(text2sound_batchsize):
+            latent_representation_gradio_images.append(latent_representation_to_Gradio_image(latent_representations[i]))
+            quantized_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(quantized_latent_representations[i]))
+            new_sound_spectrogram_gradio_images.append(flipped_log_spectrums[i])
+            new_sound_rec_signals_gradio.append((sample_rate, rec_signals[i]))
+
+        def concatenate_arrays(arrays_list):
+            return np.concatenate(arrays_list, axis=1)
+
+        concatenated_spectrogram_gradio_image = concatenate_arrays(new_sound_spectrogram_gradio_images)
+
+        interpolation_with_text_dict["latent_representation_gradio_images"] = latent_representation_gradio_images
+        interpolation_with_text_dict["quantized_latent_representation_gradio_images"] = quantized_latent_representation_gradio_images
+        interpolation_with_text_dict["new_sound_spectrogram_gradio_images"] = new_sound_spectrogram_gradio_images
+        interpolation_with_text_dict["new_sound_rec_signals_gradio"] = new_sound_rec_signals_gradio
+
+        return {text2sound_latent_representation_image: interpolation_with_text_dict["latent_representation_gradio_images"][0],
+                text2sound_quantized_latent_representation_image:
+                    interpolation_with_text_dict["quantized_latent_representation_gradio_images"][0],
+                text2sound_sampled_concatenated_spectrogram_image: concatenated_spectrogram_gradio_image,
+                text2sound_sampled_spectrogram_image: interpolation_with_text_dict["new_sound_spectrogram_gradio_images"][0],
+                text2sound_sampled_audio: interpolation_with_text_dict["new_sound_rec_signals_gradio"][0],
+                text2sound_seed_textbox: text2sound_seed,
+                interpolation_with_text_state: interpolation_with_text_dict,
+                text2sound_sample_index_slider: gr.update(minimum=0, maximum=text2sound_batchsize - 1, value=0, step=1,
+                                                          visible=True,
+                                                          label="Sample index.",
+                                                          info="Swipe to view other samples")}
+
+    def show_random_sample(sample_index, text2sound_dict):
+        sample_index = int(sample_index)
+        return {text2sound_latent_representation_image: text2sound_dict["latent_representation_gradio_images"][
+            sample_index],
+                text2sound_quantized_latent_representation_image:
+                    text2sound_dict["quantized_latent_representation_gradio_images"][sample_index],
+                text2sound_sampled_spectrogram_image: text2sound_dict["new_sound_spectrogram_gradio_images"][sample_index],
+                text2sound_sampled_audio: text2sound_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+    with gr.Tab("InterpolationCond."):
+        gr.Markdown("Use interpolation to generate a gradient sound sequence.")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=3):
+                text2sound_prompts_1_textbox = gr.Textbox(label="Positive prompt 1", lines=2, value="organ")
+                text2sound_prompts_2_textbox = gr.Textbox(label="Positive prompt 2", lines=2, value="string")
+                text2sound_negative_prompts_textbox = gr.Textbox(label="Negative prompt", lines=2, value="")
+
+            with gr.Column(scale=1):
+                text2sound_sampling_button = gr.Button(variant="primary",
+                                                       value="Generate a batch of samples and show "
+                                                             "the first one",
+                                                       scale=1)
+                text2sound_sample_index_slider = gr.Slider(minimum=0, maximum=3, value=0, step=1.0, visible=False,
+                                                           label="Sample index",
+                                                           info="Swipe to view other samples")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1, variant="panel"):
+                text2sound_sample_steps_slider = gradioWebUI.get_sample_steps_slider()
+                text2sound_sampler_radio = gradioWebUI.get_sampler_radio()
+                text2sound_batchsize_slider = gradioWebUI.get_batchsize_slider(cpu_batchsize=3)
+                text2sound_duration_slider = gradioWebUI.get_duration_slider()
+                text2sound_guidance_scale_slider = gradioWebUI.get_guidance_scale_slider()
+                text2sound_seed_textbox = gradioWebUI.get_seed_textbox()
+
+            with gr.Column(scale=1):
+                with gr.Row(variant="panel"):
+                    text2sound_sampled_concatenated_spectrogram_image = gr.Image(label="Interpolations", type="numpy",
+                                                                                 height=420, scale=8)
+                    text2sound_sampled_spectrogram_image = gr.Image(label="Selected spectrogram", type="numpy",
+                                                                    height=420, scale=1)
+                text2sound_sampled_audio = gr.Audio(type="numpy", label="Play")
+
+        with gr.Row(variant="panel"):
+            text2sound_latent_representation_image = gr.Image(label="Sampled latent representation", type="numpy",
+                                                              height=200, width=100)
+            text2sound_quantized_latent_representation_image = gr.Image(label="Quantized latent representation",
+                                                                        type="numpy", height=200, width=100)
+
+    text2sound_sampling_button.click(diffusion_random_sample,
+                                     inputs=[text2sound_prompts_1_textbox,
+                                             text2sound_prompts_2_textbox,
+                                             text2sound_negative_prompts_textbox,
+                                             text2sound_batchsize_slider,
+                                             text2sound_duration_slider,
+                                             text2sound_guidance_scale_slider, text2sound_sampler_radio,
+                                             text2sound_sample_steps_slider,
+                                             text2sound_seed_textbox,
+                                             interpolation_with_text_state],
+                                     outputs=[text2sound_latent_representation_image,
+                                              text2sound_quantized_latent_representation_image,
+                                              text2sound_sampled_concatenated_spectrogram_image,
+                                              text2sound_sampled_spectrogram_image,
+                                              text2sound_sampled_audio,
+                                              text2sound_seed_textbox,
+                                              interpolation_with_text_state,
+                                              text2sound_sample_index_slider])
+    text2sound_sample_index_slider.change(show_random_sample,
+                                          inputs=[text2sound_sample_index_slider, interpolation_with_text_state],
+                                          outputs=[text2sound_latent_representation_image,
+                                                   text2sound_quantized_latent_representation_image,
+                                                   text2sound_sampled_spectrogram_image,
+                                                   text2sound_sampled_audio])

webUI/deprecated/interpolationWithXT.py

@@ -0,0 +1,173 @@
+import gradio as gr
+import numpy as np
+import torch
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import safe_int
+from webUI.natural_language_guided.utils import encodeBatch2GradioOutput, latent_representation_to_Gradio_image
+
+
+def get_interpolation_with_xT_module(gradioWebUI, interpolation_with_text_state):
+    # Load configurations
+    uNet = gradioWebUI.uNet
+    freq_resolution, time_resolution = gradioWebUI.freq_resolution, gradioWebUI.time_resolution
+    VAE_scale = gradioWebUI.VAE_scale
+    height, width, channels = int(freq_resolution/VAE_scale), int(time_resolution/VAE_scale), gradioWebUI.channels
+    timesteps = gradioWebUI.timesteps
+    VAE_quantizer = gradioWebUI.VAE_quantizer
+    VAE_decoder = gradioWebUI.VAE_decoder
+    CLAP = gradioWebUI.CLAP
+    CLAP_tokenizer = gradioWebUI.CLAP_tokenizer
+    device = gradioWebUI.device
+    squared = gradioWebUI.squared
+    sample_rate = gradioWebUI.sample_rate
+    noise_strategy = gradioWebUI.noise_strategy
+
+    def diffusion_random_sample(text2sound_prompts, text2sound_negative_prompts, text2sound_batchsize,
+                                text2sound_duration,
+                                text2sound_noise_variance, text2sound_guidance_scale, text2sound_sampler,
+                                text2sound_sample_steps, text2sound_seed,
+                                interpolation_with_text_dict):
+        text2sound_sample_steps = int(text2sound_sample_steps)
+        text2sound_seed = safe_int(text2sound_seed, 12345678)
+        # Todo: take care of text2sound_time_resolution/width
+        width = int(time_resolution*((text2sound_duration+1)/4) / VAE_scale)
+        text2sound_batchsize = int(text2sound_batchsize)
+
+        text2sound_embedding = \
+        CLAP.get_text_features(**CLAP_tokenizer([text2sound_prompts], padding=True, return_tensors="pt"))[0].to(device)
+
+        CFG = int(text2sound_guidance_scale)
+
+        mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy)
+        unconditional_condition = \
+        CLAP.get_text_features(**CLAP_tokenizer([text2sound_negative_prompts], padding=True, return_tensors="pt"))[0]
+        mySampler.activate_classifier_free_guidance(CFG, unconditional_condition.to(device))
+
+        mySampler.respace(list(np.linspace(0, timesteps - 1, text2sound_sample_steps, dtype=np.int32)))
+
+        condition = text2sound_embedding.repeat(text2sound_batchsize, 1)
+        latent_representations, initial_noise = \
+        mySampler.interpolate(model=uNet, shape=(text2sound_batchsize, channels, height, width),
+                              seed=text2sound_seed,
+                              variance=text2sound_noise_variance,
+                              return_tensor=True, condition=condition, sampler=text2sound_sampler)
+
+        latent_representations = latent_representations[-1]
+
+        interpolation_with_text_dict["latent_representations"] = latent_representations
+
+        latent_representation_gradio_images = []
+        quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_rec_signals_gradio = []
+
+        quantized_latent_representations, loss, (_, _, _) = VAE_quantizer(latent_representations)
+        # Todo: remove hard-coding
+        flipped_log_spectrums, rec_signals = encodeBatch2GradioOutput(VAE_decoder, quantized_latent_representations,
+                                                                          resolution=(512, width * VAE_scale), centralized=False,
+                                                                          squared=squared)
+
+        for i in range(text2sound_batchsize):
+            latent_representation_gradio_images.append(latent_representation_to_Gradio_image(latent_representations[i]))
+            quantized_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(quantized_latent_representations[i]))
+            new_sound_spectrogram_gradio_images.append(flipped_log_spectrums[i])
+            new_sound_rec_signals_gradio.append((sample_rate, rec_signals[i]))
+
+        def concatenate_arrays(arrays_list):
+            return np.concatenate(arrays_list, axis=1)
+
+        concatenated_spectrogram_gradio_image = concatenate_arrays(new_sound_spectrogram_gradio_images)
+
+        interpolation_with_text_dict["latent_representation_gradio_images"] = latent_representation_gradio_images
+        interpolation_with_text_dict["quantized_latent_representation_gradio_images"] = quantized_latent_representation_gradio_images
+        interpolation_with_text_dict["new_sound_spectrogram_gradio_images"] = new_sound_spectrogram_gradio_images
+        interpolation_with_text_dict["new_sound_rec_signals_gradio"] = new_sound_rec_signals_gradio
+
+        return {text2sound_latent_representation_image: interpolation_with_text_dict["latent_representation_gradio_images"][0],
+                text2sound_quantized_latent_representation_image:
+                    interpolation_with_text_dict["quantized_latent_representation_gradio_images"][0],
+                text2sound_sampled_concatenated_spectrogram_image: concatenated_spectrogram_gradio_image,
+                text2sound_sampled_spectrogram_image: interpolation_with_text_dict["new_sound_spectrogram_gradio_images"][0],
+                text2sound_sampled_audio: interpolation_with_text_dict["new_sound_rec_signals_gradio"][0],
+                text2sound_seed_textbox: text2sound_seed,
+                interpolation_with_text_state: interpolation_with_text_dict,
+                text2sound_sample_index_slider: gr.update(minimum=0, maximum=text2sound_batchsize - 1, value=0, step=1,
+                                                          visible=True,
+                                                          label="Sample index.",
+                                                          info="Swipe to view other samples")}
+
+    def show_random_sample(sample_index, text2sound_dict):
+        sample_index = int(sample_index)
+        return {text2sound_latent_representation_image: text2sound_dict["latent_representation_gradio_images"][
+            sample_index],
+                text2sound_quantized_latent_representation_image:
+                    text2sound_dict["quantized_latent_representation_gradio_images"][sample_index],
+                text2sound_sampled_spectrogram_image: text2sound_dict["new_sound_spectrogram_gradio_images"][sample_index],
+                text2sound_sampled_audio: text2sound_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+    with gr.Tab("InterpolationXT"):
+        gr.Markdown("Use interpolation to generate a gradient sound sequence.")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=3):
+                text2sound_prompts_textbox = gr.Textbox(label="Positive prompt", lines=2, value="organ")
+                text2sound_negative_prompts_textbox = gr.Textbox(label="Negative prompt", lines=2, value="")
+
+            with gr.Column(scale=1):
+                text2sound_sampling_button = gr.Button(variant="primary",
+                                                       value="Generate a batch of samples and show "
+                                                             "the first one",
+                                                       scale=1)
+                text2sound_sample_index_slider = gr.Slider(minimum=0, maximum=3, value=0, step=1.0, visible=False,
+                                                           label="Sample index",
+                                                           info="Swipe to view other samples")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1, variant="panel"):
+                text2sound_sample_steps_slider = gradioWebUI.get_sample_steps_slider()
+                text2sound_sampler_radio = gradioWebUI.get_sampler_radio()
+                text2sound_batchsize_slider = gradioWebUI.get_batchsize_slider(cpu_batchsize=3)
+                text2sound_duration_slider = gradioWebUI.get_duration_slider()
+                text2sound_guidance_scale_slider = gradioWebUI.get_guidance_scale_slider()
+                text2sound_seed_textbox = gradioWebUI.get_seed_textbox()
+                text2sound_noise_variance_slider = gr.Slider(minimum=0., maximum=5., value=1., step=0.01,
+                                                             label="Noise variance",
+                                                             info="The larger this value, the more diversity the interpolation has.")
+
+            with gr.Column(scale=1):
+                with gr.Row(variant="panel"):
+                    text2sound_sampled_concatenated_spectrogram_image = gr.Image(label="Interpolations", type="numpy",
+                                                                                 height=420, scale=8)
+                    text2sound_sampled_spectrogram_image = gr.Image(label="Selected spectrogram", type="numpy",
+                                                                    height=420, scale=1)
+                text2sound_sampled_audio = gr.Audio(type="numpy", label="Play")
+
+        with gr.Row(variant="panel"):
+            text2sound_latent_representation_image = gr.Image(label="Sampled latent representation", type="numpy",
+                                                              height=200, width=100)
+            text2sound_quantized_latent_representation_image = gr.Image(label="Quantized latent representation",
+                                                                        type="numpy", height=200, width=100)
+
+    text2sound_sampling_button.click(diffusion_random_sample,
+                                     inputs=[text2sound_prompts_textbox, text2sound_negative_prompts_textbox,
+                                             text2sound_batchsize_slider,
+                                             text2sound_duration_slider,
+                                             text2sound_noise_variance_slider,
+                                             text2sound_guidance_scale_slider, text2sound_sampler_radio,
+                                             text2sound_sample_steps_slider,
+                                             text2sound_seed_textbox,
+                                             interpolation_with_text_state],
+                                     outputs=[text2sound_latent_representation_image,
+                                              text2sound_quantized_latent_representation_image,
+                                              text2sound_sampled_concatenated_spectrogram_image,
+                                              text2sound_sampled_spectrogram_image,
+                                              text2sound_sampled_audio,
+                                              text2sound_seed_textbox,
+                                              interpolation_with_text_state,
+                                              text2sound_sample_index_slider])
+    text2sound_sample_index_slider.change(show_random_sample,
+                                          inputs=[text2sound_sample_index_slider, interpolation_with_text_state],
+                                          outputs=[text2sound_latent_representation_image,
+                                                   text2sound_quantized_latent_representation_image,
+                                                   text2sound_sampled_spectrogram_image,
+                                                   text2sound_sampled_audio])

webUI/natural_language_guided/GAN.py

@@ -0,0 +1,164 @@
+import gradio as gr
+import numpy as np
+import torch
+
+from tools import safe_int
+from webUI.natural_language_guided_STFT.utils import encodeBatch2GradioOutput, latent_representation_to_Gradio_image, \
+    add_instrument
+
+
+def get_testGAN(gradioWebUI, text2sound_state, virtual_instruments_state):
+    # Load configurations
+    gan_generator = gradioWebUI.GAN_generator
+    freq_resolution, time_resolution = gradioWebUI.freq_resolution, gradioWebUI.time_resolution
+    VAE_scale = gradioWebUI.VAE_scale
+    height, width, channels = int(freq_resolution / VAE_scale), int(time_resolution / VAE_scale), gradioWebUI.channels
+
+    timesteps = gradioWebUI.timesteps
+    VAE_quantizer = gradioWebUI.VAE_quantizer
+    VAE_decoder = gradioWebUI.VAE_decoder
+    CLAP = gradioWebUI.CLAP
+    CLAP_tokenizer = gradioWebUI.CLAP_tokenizer
+    device = gradioWebUI.device
+    squared = gradioWebUI.squared
+    sample_rate = gradioWebUI.sample_rate
+    noise_strategy = gradioWebUI.noise_strategy
+
+    def gan_random_sample(text2sound_prompts, text2sound_negative_prompts, text2sound_batchsize,
+                                text2sound_duration,
+                                text2sound_guidance_scale, text2sound_sampler,
+                                text2sound_sample_steps, text2sound_seed,
+                                text2sound_dict):
+        text2sound_seed = safe_int(text2sound_seed, 12345678)
+
+        width = int(time_resolution * ((text2sound_duration + 1) / 4) / VAE_scale)
+
+        text2sound_batchsize = int(text2sound_batchsize)
+
+        text2sound_embedding = \
+            CLAP.get_text_features(**CLAP_tokenizer([text2sound_prompts], padding=True, return_tensors="pt"))[0].to(
+                device)
+
+        CFG = int(text2sound_guidance_scale)
+
+        condition = text2sound_embedding.repeat(text2sound_batchsize, 1)
+
+        noise = torch.randn(text2sound_batchsize, channels, height, width).to(device)
+        latent_representations = gan_generator(noise, condition)
+
+        print(latent_representations[0, 0, :3, :3])
+
+        latent_representation_gradio_images = []
+        quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_rec_signals_gradio = []
+
+        quantized_latent_representations, loss, (_, _, _) = VAE_quantizer(latent_representations)
+        # Todo: remove hard-coding
+        flipped_log_spectrums, rec_signals = encodeBatch2GradioOutput(VAE_decoder, quantized_latent_representations,
+                                                                      resolution=(512, width * VAE_scale),
+                                                                      centralized=False,
+                                                                      squared=squared)
+
+        for i in range(text2sound_batchsize):
+            latent_representation_gradio_images.append(latent_representation_to_Gradio_image(latent_representations[i]))
+            quantized_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(quantized_latent_representations[i]))
+            new_sound_spectrogram_gradio_images.append(flipped_log_spectrums[i])
+            new_sound_rec_signals_gradio.append((sample_rate, rec_signals[i]))
+
+        text2sound_dict["latent_representations"] = latent_representations.to("cpu").detach().numpy()
+        text2sound_dict["quantized_latent_representations"] = quantized_latent_representations.to("cpu").detach().numpy()
+        text2sound_dict["latent_representation_gradio_images"] = latent_representation_gradio_images
+        text2sound_dict["quantized_latent_representation_gradio_images"] = quantized_latent_representation_gradio_images
+        text2sound_dict["new_sound_spectrogram_gradio_images"] = new_sound_spectrogram_gradio_images
+        text2sound_dict["new_sound_rec_signals_gradio"] = new_sound_rec_signals_gradio
+
+        text2sound_dict["condition"] = condition.to("cpu").detach().numpy()
+        # text2sound_dict["negative_condition"] = negative_condition.to("cpu").detach().numpy()
+        text2sound_dict["guidance_scale"] = CFG
+        text2sound_dict["sampler"] = text2sound_sampler
+
+        return {text2sound_latent_representation_image: text2sound_dict["latent_representation_gradio_images"][0],
+                text2sound_quantized_latent_representation_image:
+                    text2sound_dict["quantized_latent_representation_gradio_images"][0],
+                text2sound_sampled_spectrogram_image: text2sound_dict["new_sound_spectrogram_gradio_images"][0],
+                text2sound_sampled_audio: text2sound_dict["new_sound_rec_signals_gradio"][0],
+                text2sound_seed_textbox: text2sound_seed,
+                text2sound_state: text2sound_dict,
+                text2sound_sample_index_slider: gr.update(minimum=0, maximum=text2sound_batchsize - 1, value=0, step=1,
+                                                          visible=True,
+                                                          label="Sample index.",
+                                                          info="Swipe to view other samples")}
+
+    def show_random_sample(sample_index, text2sound_dict):
+        sample_index = int(sample_index)
+        text2sound_dict["sample_index"] = sample_index
+        return {text2sound_latent_representation_image: text2sound_dict["latent_representation_gradio_images"][
+            sample_index],
+                text2sound_quantized_latent_representation_image:
+                    text2sound_dict["quantized_latent_representation_gradio_images"][sample_index],
+                text2sound_sampled_spectrogram_image: text2sound_dict["new_sound_spectrogram_gradio_images"][
+                    sample_index],
+                text2sound_sampled_audio: text2sound_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+
+    with gr.Tab("Text2sound_GAN"):
+        gr.Markdown("Use neural networks to select random sounds using your favorite instrument!")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=3):
+                text2sound_prompts_textbox = gr.Textbox(label="Positive prompt", lines=2, value="organ")
+                text2sound_negative_prompts_textbox = gr.Textbox(label="Negative prompt", lines=2, value="")
+
+            with gr.Column(scale=1):
+                text2sound_sampling_button = gr.Button(variant="primary",
+                                                       value="Generate a batch of samples and show "
+                                                             "the first one",
+                                                       scale=1)
+                text2sound_sample_index_slider = gr.Slider(minimum=0, maximum=3, value=0, step=1.0, visible=False,
+                                                           label="Sample index",
+                                                           info="Swipe to view other samples")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1, variant="panel"):
+                text2sound_sample_steps_slider = gradioWebUI.get_sample_steps_slider()
+                text2sound_sampler_radio = gradioWebUI.get_sampler_radio()
+                text2sound_batchsize_slider = gradioWebUI.get_batchsize_slider()
+                text2sound_duration_slider = gradioWebUI.get_duration_slider()
+                text2sound_guidance_scale_slider = gradioWebUI.get_guidance_scale_slider()
+                text2sound_seed_textbox = gradioWebUI.get_seed_textbox()
+
+            with gr.Column(scale=1):
+                text2sound_sampled_spectrogram_image = gr.Image(label="Sampled spectrogram", type="numpy", height=420)
+                text2sound_sampled_audio = gr.Audio(type="numpy", label="Play")
+
+
+        with gr.Row(variant="panel"):
+            text2sound_latent_representation_image = gr.Image(label="Sampled latent representation", type="numpy",
+                                                              height=200, width=100)
+            text2sound_quantized_latent_representation_image = gr.Image(label="Quantized latent representation",
+                                                                        type="numpy", height=200, width=100)
+
+    text2sound_sampling_button.click(gan_random_sample,
+                                     inputs=[text2sound_prompts_textbox,
+                                             text2sound_negative_prompts_textbox,
+                                             text2sound_batchsize_slider,
+                                             text2sound_duration_slider,
+                                             text2sound_guidance_scale_slider, text2sound_sampler_radio,
+                                             text2sound_sample_steps_slider,
+                                             text2sound_seed_textbox,
+                                             text2sound_state],
+                                     outputs=[text2sound_latent_representation_image,
+                                              text2sound_quantized_latent_representation_image,
+                                              text2sound_sampled_spectrogram_image,
+                                              text2sound_sampled_audio,
+                                              text2sound_seed_textbox,
+                                              text2sound_state,
+                                              text2sound_sample_index_slider])
+
+
+    text2sound_sample_index_slider.change(show_random_sample,
+                                          inputs=[text2sound_sample_index_slider, text2sound_state],
+                                          outputs=[text2sound_latent_representation_image,
+                                                   text2sound_quantized_latent_representation_image,
+                                                   text2sound_sampled_spectrogram_image,
+                                                   text2sound_sampled_audio])

webUI/natural_language_guided/README.py

@@ -0,0 +1,53 @@
+import gradio as gr
+
+readme_content = """## Stable Diffusion for Sound Generation
+
+This project applies stable diffusion[1] to sound generation. Inspired by the work of AUTOMATIC1111, 2022[2], we have implemented a preliminary version of text2sound, sound2sound, inpaint, as well as an additional interpolation feature, all accessible through a web UI.
+
+### Neural Network Training Data:
+The neural network is trained using the filtered NSynth dataset[3], which is a large-scale and high-quality collection of annotated musical notes, comprising 305,979 musical notes. However, for this project, only samples with a pitch set to E3 were used, resulting in an actual training sample size of 4,096, making it a low-resource project.
+
+The training took place on an NVIDIA Tesla T4 GPU and spanned approximately 10 hours.
+
+### Natural Language Guidance:
+Natural language guidance is derived from the multi-label annotations of the NSynth dataset. The labels included in the training are:
+
+- **Instrument Families**: bass, brass, flute, guitar, keyboard, mallet, organ, reed, string, synth lead, vocal.
+
+- **Instrument Sources**: acoustic, electronic, synthetic.
+
+- **Note Qualities**: bright, dark, distortion, fast decay, long release, multiphonic, nonlinear env, percussive, reverb, tempo-synced.
+
+### Usage Hints:
+
+1. **Prompt Format**: It's recommended to use the format “label1, label2, label3“, e.g., ”organ, dark, long release“.
+
+2. **Unique Sounds**: If you keep generating the same sound, try setting a different seed!
+
+3. **Sample Indexing**: Drag the "Sample index slider" to view other samples within the generated batch.
+
+4. **Running on CPU**: Be cautious with the settings for 'batchsize' and 'sample_steps' when running on CPU to avoid timeouts. Recommended settings are batchsize ≤ 4 and sample_steps = 15.
+
+5. **Editing Sounds**: Generated audio can be downloaded and then re-uploaded for further editing at the sound2sound/inpaint sections.
+
+6. **Guidance Scale**: A higher 'guidance_scale' intensifies the influence of natural language conditioning on the generation[4]. It's recommended to set it between 3 and 10.
+
+7. **Noising Strength**: A smaller 'noising_strength' value makes the generated sound closer to the input sound.
+
+References:
+
+[1] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., & Ommer, B. (2022). High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 10684-10695).
+
+[2] AUTOMATIC1111. (2022). Stable Diffusion Web UI [Computer software]. Retrieved from https://github.com/AUTOMATIC1111/stable-diffusion-webui
+
+[3] Engel, J., Resnick, C., Roberts, A., Dieleman, S., Eck, D., Simonyan, K., & Norouzi, M. (2017). Neural Audio Synthesis of Musical Notes with WaveNet Autoencoders.
+
+[4] Ho, J., & Salimans, T. (2022). Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598.
+"""
+
+def get_readme_module():
+
+    with gr.Tab("README"):
+        # gr.Markdown("Use interpolation to generate a gradient sound sequence.")
+        with gr.Column(scale=3):
+            readme_textbox = gr.Textbox(label="readme", lines=40, value=readme_content, interactive=False)