Upload 49 files

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
+
+

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)

webUI/natural_language_guided/build_instrument.py

@@ -0,0 +1,274 @@
+import librosa
+import numpy as np
+import torch
+import gradio as gr
+import mido
+from io import BytesIO
+import pyrubberband as pyrb
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import adsr_envelope, adjust_audio_length
+from webUI.natural_language_guided.track_maker import DiffSynth
+from webUI.natural_language_guided.utils import encodeBatch2GradioOutput_STFT, phase_to_Gradio_image, \
+    spectrogram_to_Gradio_image
+
+
+def get_build_instrument_module(gradioWebUI, virtual_instruments_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 select_sound(virtual_instrument_name, virtual_instruments_dict):
+        virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+        virtual_instrument = virtual_instruments[virtual_instrument_name]
+
+        return {source_sound_spectrogram_image: virtual_instrument["spectrogram_gradio_image"],
+                source_sound_phase_image: virtual_instrument["phase_gradio_image"],
+                source_sound_audio: virtual_instrument["signal"]}
+
+    def make_track(inpaint_steps, midi, noising_strength, attack, before_release, instrument_names, virtual_instruments_dict):
+
+        if noising_strength < 1:
+          print(f"Warning: making track with noising_strength = {noising_strength} < 1")
+        virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+        sample_steps = int(inpaint_steps)
+
+        instrument_names = instrument_names.split("@")
+        instruments_configs = {}
+        for virtual_instrument_name in instrument_names:
+            virtual_instrument = virtual_instruments[virtual_instrument_name]
+
+            latent_representation = torch.tensor(virtual_instrument["latent_representation"], dtype=torch.float32).to(device)
+            sampler = virtual_instrument["sampler"]
+
+            batchsize = 1
+
+            latent_representation = latent_representation.repeat(batchsize, 1, 1, 1)
+
+            mid = mido.MidiFile(file=BytesIO(midi))
+            instruments_configs[virtual_instrument_name] = {
+                    'sample_steps': sample_steps,
+                    'sampler': sampler,
+                    'noising_strength': noising_strength,
+                    'latent_representation': latent_representation,
+                    'attack': attack,
+                    'before_release': before_release}
+
+        diffSynth = DiffSynth(instruments_configs, uNet, VAE_quantizer, VAE_decoder, CLAP, CLAP_tokenizer, device)
+
+        full_audio = diffSynth.get_music(mid, instrument_names)
+
+        return {track_audio: (sample_rate, full_audio)}
+
+    def test_duration_inpaint(virtual_instrument_name, inpaint_steps, duration, noising_strength, end_noise_level_ratio, attack, before_release, mask_flexivity, virtual_instruments_dict, use_dynamic_mask):
+        width = int(time_resolution * ((duration + 1) / 4) / VAE_scale)
+
+        virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+        virtual_instrument = virtual_instruments[virtual_instrument_name]
+
+        latent_representation = torch.tensor(virtual_instrument["latent_representation"], dtype=torch.float32).to(device)
+        sample_steps = int(inpaint_steps)
+        sampler = virtual_instrument["sampler"]
+        batchsize = 1
+
+        mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy)
+        mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32)))
+
+        latent_representation = latent_representation.repeat(batchsize, 1, 1, 1)
+
+        # mask = 1, freeze
+        latent_mask = torch.zeros((batchsize, 1, height, width), dtype=torch.float32).to(device)
+
+        latent_mask[:, :, :, :int(time_resolution * (attack / 4) / VAE_scale)] = 1.0
+        latent_mask[:, :, :, -int(time_resolution * ((before_release+1) / 4) / VAE_scale):] = 1.0
+
+
+        text2sound_embedding = \
+            CLAP.get_text_features(**CLAP_tokenizer([""], padding=True, return_tensors="pt"))[0].to(
+                device)
+        condition = text2sound_embedding.repeat(1, 1)
+
+
+        latent_representations, initial_noise = \
+            mySampler.inpaint_sample(model=uNet, shape=(batchsize, channels, height, width),
+                                     noising_strength=noising_strength,
+                                     guide_img=latent_representation, mask=latent_mask, return_tensor=True,
+                                     condition=condition, sampler=sampler,
+                                     use_dynamic_mask=use_dynamic_mask,
+                                     end_noise_level_ratio=end_noise_level_ratio,
+                                     mask_flexivity=mask_flexivity)
+
+        latent_representations = latent_representations[-1]
+
+        quantized_latent_representations, loss, (_, _, _) = VAE_quantizer(latent_representations)
+        # Todo: remove hard-coding
+        flipped_log_spectrums, flipped_phases, rec_signals, _, _, _ = encodeBatch2GradioOutput_STFT(VAE_decoder,
+                                                                                                            quantized_latent_representations,
+                                                                                                            resolution=(
+                                                                                                                512,
+                                                                                                                width * VAE_scale),
+                                                                                                            original_STFT_batch=None
+                                                                                                     )
+
+
+        return {test_duration_spectrogram_image: flipped_log_spectrums[0],
+                test_duration_phase_image: flipped_phases[0],
+                test_duration_audio: (sample_rate, rec_signals[0])}
+
+    def test_duration_envelope(virtual_instrument_name, duration, noising_strength, attack, before_release, release, virtual_instruments_dict):
+
+        virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+        virtual_instrument = virtual_instruments[virtual_instrument_name]
+        sample_rate, signal = virtual_instrument["signal"]
+
+        applied_signal = adsr_envelope(signal=signal, sample_rate=sample_rate, duration=duration,
+                                       attack_time=0.0, decay_time=0.0, sustain_level=1.0, release_time=release)
+
+        D = librosa.stft(applied_signal, n_fft=1024, hop_length=256, win_length=1024)[1:, :]
+        spc = np.abs(D)
+        phase = np.angle(D)
+
+        flipped_log_spectrum = spectrogram_to_Gradio_image(spc)
+        flipped_phase = phase_to_Gradio_image(phase)
+
+        return {test_duration_spectrogram_image: flipped_log_spectrum,
+                test_duration_phase_image: flipped_phase,
+                test_duration_audio: (sample_rate, applied_signal)}
+
+    def test_duration_stretch(virtual_instrument_name, duration, noising_strength, attack, before_release, release, virtual_instruments_dict):
+
+        virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+        virtual_instrument = virtual_instruments[virtual_instrument_name]
+        sample_rate, signal = virtual_instrument["signal"]
+
+        s = 3 / duration
+        applied_signal = pyrb.time_stretch(signal, sample_rate, s)
+        applied_signal = adjust_audio_length(applied_signal, int((duration+1) * sample_rate), sample_rate, sample_rate)
+
+        D = librosa.stft(applied_signal, n_fft=1024, hop_length=256, win_length=1024)[1:, :]
+        spc = np.abs(D)
+        phase = np.angle(D)
+
+        flipped_log_spectrum = spectrogram_to_Gradio_image(spc)
+        flipped_phase = phase_to_Gradio_image(phase)
+
+        return {test_duration_spectrogram_image: flipped_log_spectrum,
+                test_duration_phase_image: flipped_phase,
+                test_duration_audio: (sample_rate, applied_signal)}
+
+
+    with gr.Tab("TestInTrack"):
+        gr.Markdown("Make music with generated sounds!")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=3):
+                instrument_name_textbox = gr.Textbox(label="Instrument name", lines=1,
+                                                     placeholder="Name of your instrument", scale=1)
+                select_instrument_button = gr.Button(variant="primary", value="Select", scale=1)
+            with gr.Column(scale=3):
+                inpaint_steps_slider = gr.Slider(minimum=5.0, maximum=999.0, value=20.0, step=1.0, label="inpaint_steps")
+                noising_strength_slider = gradioWebUI.get_noising_strength_slider(default_noising_strength=1.)
+                end_noise_level_ratio_slider = gr.Slider(minimum=0.0, maximum=1., value=0.0, step=0.01, label="end_noise_level_ratio")
+                attack_slider = gr.Slider(minimum=0.0, maximum=1.5, value=0.5, step=0.01, label="attack in sec")
+                before_release_slider = gr.Slider(minimum=0.0, maximum=1.5, value=0.5, step=0.01, label="before_release in sec")
+                release_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, step=0.01, label="release in sec")
+                mask_flexivity_slider = gr.Slider(minimum=0.01, maximum=1.00, value=1., step=0.01, label="mask_flexivity")
+            with gr.Column(scale=3):
+                use_dynamic_mask_checkbox = gr.Checkbox(label="Use dynamic mask", value=True)
+                test_duration_envelope_button = gr.Button(variant="primary", value="Apply envelope", scale=1)
+                test_duration_stretch_button = gr.Button(variant="primary", value="Apply stretch", scale=1)
+                test_duration_inpaint_button = gr.Button(variant="primary", value="Inpaint different duration", scale=1)
+                duration_slider = gradioWebUI.get_duration_slider()
+
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=2):
+                with gr.Row(variant="panel"):
+                    source_sound_spectrogram_image = gr.Image(label="New sound spectrogram", type="numpy",
+                                                              height=600, scale=1)
+                    source_sound_phase_image = gr.Image(label="New sound phase", type="numpy",
+                                                              height=600, scale=1)
+                source_sound_audio = gr.Audio(type="numpy", label="Play new sound", interactive=False)
+
+            with gr.Column(scale=3):
+                with gr.Row(variant="panel"):
+                    test_duration_spectrogram_image = gr.Image(label="New sound spectrogram", type="numpy",
+                                                              height=600, scale=1)
+                    test_duration_phase_image = gr.Image(label="New sound phase", type="numpy",
+                                                              height=600, scale=1)
+                test_duration_audio = gr.Audio(type="numpy", label="Play new sound", interactive=False)
+
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                # track_spectrogram_image = gr.Image(label="New sound spectrogram", type="numpy",
+                #                                    height=420, scale=1)
+                midi_file = gr.File(label="Upload midi file", type="binary")
+                instrument_names_textbox = gr.Textbox(label="Instrument names", lines=2,
+                                                     placeholder="Names of your instrument used to play the midi", scale=1)
+                track_audio = gr.Audio(type="numpy", label="Play new sound", interactive=False)
+            make_track_button = gr.Button(variant="primary", value="Make track", scale=1)
+
+    select_instrument_button.click(select_sound,
+                                   inputs=[instrument_name_textbox, virtual_instruments_state],
+                                   outputs=[source_sound_spectrogram_image,
+                                            source_sound_phase_image,
+                                            source_sound_audio])
+
+    test_duration_envelope_button.click(test_duration_envelope,
+                                      inputs=[instrument_name_textbox, duration_slider,
+                                              noising_strength_slider,
+                                              attack_slider,
+                                              before_release_slider,
+                                              release_slider,
+                                              virtual_instruments_state,
+                                              ],
+                                      outputs=[test_duration_spectrogram_image,
+                                               test_duration_phase_image,
+                                               test_duration_audio])
+
+    test_duration_stretch_button.click(test_duration_stretch,
+                                      inputs=[instrument_name_textbox, duration_slider,
+                                              noising_strength_slider,
+                                              attack_slider,
+                                              before_release_slider,
+                                              release_slider,
+                                              virtual_instruments_state,
+                                              ],
+                                      outputs=[test_duration_spectrogram_image,
+                                               test_duration_phase_image,
+                                               test_duration_audio])
+
+    test_duration_inpaint_button.click(test_duration_inpaint,
+                                      inputs=[instrument_name_textbox,
+                                              inpaint_steps_slider,
+                                              duration_slider,
+                                              noising_strength_slider,
+                                              end_noise_level_ratio_slider,
+                                              attack_slider,
+                                              before_release_slider,
+                                              mask_flexivity_slider,
+                                              virtual_instruments_state,
+                                              use_dynamic_mask_checkbox],
+                                      outputs=[test_duration_spectrogram_image,
+                                               test_duration_phase_image,
+                                               test_duration_audio])
+
+    make_track_button.click(make_track,
+                                      inputs=[inpaint_steps_slider, midi_file,
+                                              noising_strength_slider,
+                                              attack_slider,
+                                              before_release_slider,
+                                              instrument_names_textbox,
+                                              virtual_instruments_state],
+                                      outputs=[track_audio])
+

webUI/natural_language_guided/gradio_webUI.py

@@ -0,0 +1,68 @@
+import gradio as gr
+
+
+class GradioWebUI():
+
+    def __init__(self, device, VAE, uNet, CLAP, CLAP_tokenizer,
+                 freq_resolution=512, time_resolution=256, channels=4, timesteps=1000,
+                 sample_rate=16000, squared=False, VAE_scale=4,
+                 flexible_duration=False, noise_strategy="repeat",
+                 GAN_generator = None):
+        self.device = device
+        self.VAE_encoder, self.VAE_quantizer, self.VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder
+        self.uNet = uNet
+        self.CLAP, self.CLAP_tokenizer = CLAP, CLAP_tokenizer
+        self.freq_resolution, self.time_resolution = freq_resolution, time_resolution
+        self.channels = channels
+        self.GAN_generator = GAN_generator
+
+        self.timesteps = timesteps
+        self.sample_rate = sample_rate
+        self.squared = squared
+        self.VAE_scale = VAE_scale
+        self.flexible_duration = flexible_duration
+        self.noise_strategy = noise_strategy
+
+        self.text2sound_state = gr.State(value={})
+        self.interpolation_state = gr.State(value={})
+        self.sound2sound_state = gr.State(value={})
+        self.inpaint_state = gr.State(value={})
+
+    def get_sample_steps_slider(self):
+        default_steps = 10 if (self.device == "cpu") else 20
+        return gr.Slider(minimum=10, maximum=100, value=default_steps, step=1,
+                         label="Sample steps",
+                         info="Sampling steps. The more sampling steps, the better the "
+                              "theoretical result, but the time it consumes.")
+
+    def get_sampler_radio(self):
+        # return gr.Radio(choices=["ddpm", "ddim", "dpmsolver++", "dpmsolver"], value="ddim", label="Sampler")
+        return gr.Radio(choices=["ddpm", "ddim"], value="ddim", label="Sampler")
+
+    def get_batchsize_slider(self, cpu_batchsize=1):
+        return gr.Slider(minimum=1., maximum=16, value=cpu_batchsize if (self.device == "cpu") else 8, step=1, label="Batchsize")
+
+    def get_time_resolution_slider(self):
+        return gr.Slider(minimum=16., maximum=int(1024/self.VAE_scale), value=int(256/self.VAE_scale), step=1, label="Time resolution", interactive=True)
+
+    def get_duration_slider(self):
+        if self.flexible_duration:
+            return gr.Slider(minimum=0.25, maximum=8., value=3., step=0.01, label="duration in sec")
+        else:
+            return gr.Slider(minimum=1., maximum=8., value=3., step=1., label="duration in sec")
+
+    def get_guidance_scale_slider(self):
+        return gr.Slider(minimum=0., maximum=20., value=6., step=1.,
+                         label="Guidance scale",
+                         info="The larger this value, the more the generated sound is "
+                              "influenced by the condition. Setting it to 0 is equivalent to "
+                              "the negative case.")
+
+    def get_noising_strength_slider(self, default_noising_strength=0.7):
+        return gr.Slider(minimum=0.0, maximum=1.00, value=default_noising_strength, step=0.01,
+                         label="noising strength",
+                         info="The smaller this value, the more the generated sound is "
+                              "closed to the origin.")
+
+    def get_seed_textbox(self):
+        return gr.Textbox(label="Seed", lines=1, placeholder="seed", value=0)

webUI/natural_language_guided/inpaint_with_text.py

@@ -0,0 +1,441 @@
+import librosa
+import numpy as np
+import torch
+import gradio as gr
+from scipy.ndimage import zoom
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import adjust_audio_length, safe_int, pad_STFT, encode_stft
+from webUI.natural_language_guided.utils import latent_representation_to_Gradio_image, InputBatch2Encode_STFT, encodeBatch2GradioOutput_STFT, add_instrument
+
+
+def get_triangle_mask(height, width):
+    mask = np.zeros((height, width))
+    slope = 8 / 3
+    for i in range(height):
+        for j in range(width):
+            if i < slope * j:
+                mask[i, j] = 1
+    return mask
+
+
+def get_inpaint_with_text_module(gradioWebUI, inpaintWithText_state, virtual_instruments_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_encoder = gradioWebUI.VAE_encoder
+    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 receive_uopoad_origin_audio(sound2sound_duration, sound2sound_origin_source, sound2sound_origin_upload, sound2sound_origin_microphone,
+                                    inpaintWithText_dict):
+
+        if sound2sound_origin_source == "upload":
+            origin_sr, origin_audio = sound2sound_origin_upload
+        else:
+            origin_sr, origin_audio = sound2sound_origin_microphone
+
+        origin_audio = origin_audio / np.max(np.abs(origin_audio))
+
+        width = int(time_resolution*((sound2sound_duration+1)/4) / VAE_scale)
+        audio_length = 256 * (VAE_scale * width - 1)
+        origin_audio = adjust_audio_length(origin_audio, audio_length, origin_sr, sample_rate)
+
+        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)
+
+        # Todo: justify batchsize to 1
+        origin_spectrogram_batch_tensor = torch.from_numpy(
+            np.repeat(encoded_D[np.newaxis, :, :, :], 1, axis=0)).float().to(device)
+
+        # Todo: remove hard-coding
+        origin_flipped_log_spectrums, origin_flipped_phases, origin_signals, origin_latent_representations, quantized_origin_latent_representations = InputBatch2Encode_STFT(
+            VAE_encoder, origin_spectrogram_batch_tensor, resolution=(512, width * VAE_scale), quantizer=VAE_quantizer, squared=squared)
+
+        if sound2sound_origin_source == "upload":
+            inpaintWithText_dict["origin_upload_latent_representations"] = origin_latent_representations.tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_upload_latent_representation_image"] = latent_representation_to_Gradio_image(
+                origin_latent_representations[0]).tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_upload_quantized_latent_representation_image"] = latent_representation_to_Gradio_image(
+                quantized_origin_latent_representations[0]).tolist()
+            return {sound2sound_origin_spectrogram_upload_image: origin_flipped_log_spectrums[0],
+                    sound2sound_origin_phase_upload_image: origin_flipped_phases[0],
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(),
+                    sound2sound_origin_phase_microphone_image: gr.update(),
+                    sound2sound_origin_upload_latent_representation_image: latent_representation_to_Gradio_image(
+                        origin_latent_representations[0]),
+                    sound2sound_origin_upload_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                        quantized_origin_latent_representations[0]),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(),
+                    inpaintWithText_state: inpaintWithText_dict}
+        else:
+            inpaintWithText_dict["origin_microphone_latent_representations"] = origin_latent_representations.tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_microphone_latent_representation_image"] = latent_representation_to_Gradio_image(
+                origin_latent_representations[0]).tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_microphone_quantized_latent_representation_image"] = latent_representation_to_Gradio_image(
+                quantized_origin_latent_representations[0]).tolist()
+            return {sound2sound_origin_spectrogram_upload_image: origin_flipped_log_spectrums[0],
+                    sound2sound_origin_phase_upload_image: origin_flipped_phases[0],
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(),
+                    sound2sound_origin_phase_microphone_image: gr.update(),
+                    sound2sound_origin_upload_latent_representation_image: latent_representation_to_Gradio_image(
+                        origin_latent_representations[0]),
+                    sound2sound_origin_upload_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                        quantized_origin_latent_representations[0]),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(),
+                    inpaintWithText_state: inpaintWithText_dict}
+
+    def sound2sound_sample(sound2sound_origin_spectrogram_upload, sound2sound_origin_spectrogram_microphone,
+                           text2sound_prompts, text2sound_negative_prompts, sound2sound_batchsize,
+                           sound2sound_guidance_scale, sound2sound_sampler,
+                           sound2sound_sample_steps, sound2sound_origin_source,
+                           sound2sound_noising_strength, sound2sound_seed, sound2sound_inpaint_area,
+                           mask_time_begin, mask_time_end, mask_frequency_begin, mask_frequency_end, inpaintWithText_dict
+                           ):
+
+        # input preprocessing
+        sound2sound_seed = safe_int(sound2sound_seed, 12345678)
+        sound2sound_batchsize = int(sound2sound_batchsize)
+        noising_strength = sound2sound_noising_strength
+        sound2sound_sample_steps = int(sound2sound_sample_steps)
+        CFG = int(sound2sound_guidance_scale)
+
+        text2sound_embedding = \
+        CLAP.get_text_features(**CLAP_tokenizer([text2sound_prompts], padding=True, return_tensors="pt"))[0].to(device)
+
+        if sound2sound_origin_source == "upload":
+            origin_latent_representations = torch.tensor(
+                inpaintWithText_dict["origin_upload_latent_representations"]).repeat(sound2sound_batchsize, 1, 1, 1).to(
+                device)
+            mask = np.array(sound2sound_origin_spectrogram_upload["mask"])
+        elif sound2sound_origin_source == "microphone":
+            origin_latent_representations = torch.tensor(
+                inpaintWithText_dict["origin_microphone_latent_representations"]).repeat(sound2sound_batchsize, 1, 1, 1).to(
+                device)
+            mask = np.array(sound2sound_origin_spectrogram_microphone["mask"])
+        else:
+            print("Input source not in ['upload', 'microphone']!")
+            raise NotImplementedError()
+
+        merged_mask = np.all(mask == 255, axis=2).astype(np.uint8)
+        latent_mask = zoom(merged_mask, (1 / VAE_scale, 1 / VAE_scale))
+        latent_mask = np.clip(latent_mask, 0, 1)
+        print(f"latent_mask.avg = {np.mean(latent_mask)}")
+        latent_mask[int(mask_frequency_begin):int(mask_frequency_end), int(mask_time_begin*time_resolution/(VAE_scale*4)):int(mask_time_end*time_resolution/(VAE_scale*4))] = 1
+
+
+        # latent_mask = get_triangle_mask(128, 64)
+
+        print(f"latent_mask.avg = {np.mean(latent_mask)}")
+        if sound2sound_inpaint_area == "inpaint masked":
+            latent_mask = 1 - latent_mask
+        latent_mask = torch.from_numpy(latent_mask).unsqueeze(0).unsqueeze(1).repeat(sound2sound_batchsize, channels, 1,
+                                                                                     1).float().to(device)
+        latent_mask = torch.flip(latent_mask, [2])
+
+        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))
+
+        normalized_sample_steps = int(sound2sound_sample_steps / noising_strength)
+
+        mySampler.respace(list(np.linspace(0, timesteps - 1, normalized_sample_steps, dtype=np.int32)))
+
+        # Todo: remove hard-coding
+        width = origin_latent_representations.shape[-1]
+        condition = text2sound_embedding.repeat(sound2sound_batchsize, 1)
+
+        new_sound_latent_representations, initial_noise = \
+            mySampler.inpaint_sample(model=uNet, shape=(sound2sound_batchsize, channels, height, width),
+                                     seed=sound2sound_seed,
+                                     noising_strength=noising_strength,
+                                     guide_img=origin_latent_representations, mask=latent_mask, return_tensor=True,
+                                     condition=condition, sampler=sound2sound_sampler)
+
+        new_sound_latent_representations = new_sound_latent_representations[-1]
+
+        # Quantize new sound latent representations
+        quantized_new_sound_latent_representations, loss, (_, _, _) = VAE_quantizer(new_sound_latent_representations)
+        new_sound_flipped_log_spectrums, new_sound_flipped_phases, new_sound_signals, _, _, _ = encodeBatch2GradioOutput_STFT(VAE_decoder,
+                                                                                                            quantized_new_sound_latent_representations,
+                                                                                                            resolution=(
+                                                                                                                512,
+                                                                                                                width * VAE_scale),
+                                                                                                            original_STFT_batch=None
+                                                                                                     )
+
+        new_sound_latent_representation_gradio_images = []
+        new_sound_quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_phase_gradio_images = []
+        new_sound_rec_signals_gradio = []
+        for i in range(sound2sound_batchsize):
+            new_sound_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(new_sound_latent_representations[i]))
+            new_sound_quantized_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(quantized_new_sound_latent_representations[i]))
+            new_sound_spectrogram_gradio_images.append(new_sound_flipped_log_spectrums[i])
+            new_sound_phase_gradio_images.append(new_sound_flipped_phases[i])
+            new_sound_rec_signals_gradio.append((sample_rate, new_sound_signals[i]))
+
+        inpaintWithText_dict[
+            "new_sound_latent_representation_gradio_images"] = new_sound_latent_representation_gradio_images
+        inpaintWithText_dict[
+            "new_sound_quantized_latent_representation_gradio_images"] = new_sound_quantized_latent_representation_gradio_images
+        inpaintWithText_dict["new_sound_spectrogram_gradio_images"] = new_sound_spectrogram_gradio_images
+        inpaintWithText_dict["new_sound_phase_gradio_images"] = new_sound_phase_gradio_images
+        inpaintWithText_dict["new_sound_rec_signals_gradio"] = new_sound_rec_signals_gradio
+
+        inpaintWithText_dict["latent_representations"] = new_sound_latent_representations.to("cpu").detach().numpy()
+        inpaintWithText_dict["quantized_latent_representations"] = quantized_new_sound_latent_representations.to("cpu").detach().numpy()
+        inpaintWithText_dict["sampler"] = sound2sound_sampler
+
+        return {sound2sound_new_sound_latent_representation_image: latent_representation_to_Gradio_image(
+            new_sound_latent_representations[0]),
+            sound2sound_new_sound_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                quantized_new_sound_latent_representations[0]),
+            sound2sound_new_sound_spectrogram_image: new_sound_flipped_log_spectrums[0],
+            sound2sound_new_sound_phase_image: new_sound_flipped_phases[0],
+            sound2sound_new_sound_audio: (sample_rate, new_sound_signals[0]),
+            sound2sound_sample_index_slider: gr.update(minimum=0, maximum=sound2sound_batchsize - 1, value=0,
+                                                       step=1.0,
+                                                       visible=True,
+                                                       label="Sample index",
+                                                       info="Swipe to view other samples"),
+            sound2sound_seed_textbox: sound2sound_seed,
+            inpaintWithText_state: inpaintWithText_dict}
+
+    def show_sound2sound_sample(sound2sound_sample_index, inpaintWithText_dict):
+        sample_index = int(sound2sound_sample_index)
+        return {sound2sound_new_sound_latent_representation_image:
+                    inpaintWithText_dict["new_sound_latent_representation_gradio_images"][sample_index],
+                sound2sound_new_sound_quantized_latent_representation_image:
+                    inpaintWithText_dict["new_sound_quantized_latent_representation_gradio_images"][sample_index],
+                sound2sound_new_sound_spectrogram_image: inpaintWithText_dict["new_sound_spectrogram_gradio_images"][
+                    sample_index],
+                sound2sound_new_sound_phase_image: inpaintWithText_dict["new_sound_phase_gradio_images"][
+                    sample_index],
+                sound2sound_new_sound_audio: inpaintWithText_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+    def sound2sound_switch_origin_source(sound2sound_origin_source):
+
+        if sound2sound_origin_source == "upload":
+            return {sound2sound_origin_upload_audio: gr.update(visible=True),
+                    sound2sound_origin_microphone_audio: gr.update(visible=False),
+                    sound2sound_origin_spectrogram_upload_image: gr.update(visible=True),
+                    sound2sound_origin_phase_upload_image: gr.update(visible=True),
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(visible=False),
+                    sound2sound_origin_phase_microphone_image: gr.update(visible=False),
+                    sound2sound_origin_upload_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(visible=False)}
+        elif sound2sound_origin_source == "microphone":
+            return {sound2sound_origin_upload_audio: gr.update(visible=False),
+                    sound2sound_origin_microphone_audio: gr.update(visible=True),
+                    sound2sound_origin_spectrogram_upload_image: gr.update(visible=False),
+                    sound2sound_origin_phase_upload_image: gr.update(visible=False),
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(visible=True),
+                    sound2sound_origin_phase_microphone_image: gr.update(visible=True),
+                    sound2sound_origin_upload_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(visible=True)}
+        else:
+            print("Input source not in ['upload', 'microphone']!")
+
+    def save_virtual_instrument(sample_index, virtual_instrument_name, sound2sound_dict, virtual_instruments_dict):
+
+      virtual_instruments_dict = add_instrument(sound2sound_dict, virtual_instruments_dict, virtual_instrument_name, sample_index)
+      return {virtual_instruments_state: virtual_instruments_dict,
+              sound2sound_instrument_name_textbox: gr.Textbox(label="Instrument name", lines=1,
+                                                          placeholder=f"Saved as {virtual_instrument_name}!")}
+
+    with gr.Tab("Inpaint"):
+        gr.Markdown("Select the area to inpaint and use the prompt to guide the synthesis of a new sound!")
+        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):
+                sound2sound_sample_button = gr.Button(variant="primary", value="Generate", scale=1)
+
+                sound2sound_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):
+                with gr.Tab("Origin sound"):
+                    sound2sound_duration_slider = gradioWebUI.get_duration_slider()
+                    sound2sound_origin_source_radio = gr.Radio(choices=["upload", "microphone"], value="upload",
+                                                               label="Input source")
+
+                    sound2sound_origin_upload_audio = gr.Audio(type="numpy", label="Upload", source="upload",
+                                                               interactive=True, visible=True)
+                    sound2sound_origin_microphone_audio = gr.Audio(type="numpy", label="Record", source="microphone",
+                                                                   interactive=True, visible=False)
+                    with gr.Row(variant="panel"):
+                        sound2sound_origin_spectrogram_upload_image = gr.Image(label="Original upload spectrogram",
+                                                                               type="numpy", height=600,
+                                                                            visible=True, tool="sketch")
+                        sound2sound_origin_phase_upload_image = gr.Image(label="Original upload phase",
+                                                                               type="numpy", height=600,
+                                                                               visible=True)
+                        sound2sound_origin_spectrogram_microphone_image = gr.Image(label="Original microphone spectrogram",
+                                                                                   type="numpy", height=600,
+                                                                                   visible=False, tool="sketch")
+                        sound2sound_origin_phase_microphone_image = gr.Image(label="Original microphone phase",
+                                                                                   type="numpy", height=600,
+                                                                                   visible=False)
+                    sound2sound_inpaint_area_radio = gr.Radio(choices=["inpaint masked", "inpaint not masked"],
+                                                              value="inpaint masked")
+
+                with gr.Tab("Sound2sound settings"):
+                    sound2sound_sample_steps_slider = gradioWebUI.get_sample_steps_slider()
+                    sound2sound_sampler_radio = gradioWebUI.get_sampler_radio()
+                    sound2sound_batchsize_slider = gradioWebUI.get_batchsize_slider()
+                    sound2sound_noising_strength_slider = gradioWebUI.get_noising_strength_slider()
+                    sound2sound_guidance_scale_slider = gradioWebUI.get_guidance_scale_slider()
+                    sound2sound_seed_textbox = gradioWebUI.get_seed_textbox()
+
+                with gr.Tab("Mask prototypes"):
+                    with gr.Tab("Mask along time axis"):
+                        mask_time_begin_slider = gr.Slider(minimum=0.0, maximum=4.00, value=0.0, step=0.01, label="Begin time")
+                        mask_time_end_slider = gr.Slider(minimum=0.0, maximum=4.00, value=0.0, step=0.01, label="End time")
+                    with gr.Tab("Mask along frequency axis"):
+                        mask_frequency_begin_slider = gr.Slider(minimum=0, maximum=127, value=0, step=1, label="Begin freq pixel")
+                        mask_frequency_end_slider = gr.Slider(minimum=0, maximum=127, value=0, step=1, label="End freq pixel")
+
+            with gr.Column(scale=1):
+                sound2sound_new_sound_audio = gr.Audio(type="numpy", label="Play new sound", interactive=False)
+                with gr.Row(variant="panel"):
+                    sound2sound_new_sound_spectrogram_image = gr.Image(label="New sound spectrogram", type="numpy",
+                                                                       height=600, scale=1)
+                    sound2sound_new_sound_phase_image = gr.Image(label="New sound phase", type="numpy",
+                                                                       height=600, scale=1)
+
+
+                with gr.Row(variant="panel"):
+                    sound2sound_instrument_name_textbox = gr.Textbox(label="Instrument name", lines=1,
+                                                                    placeholder="Name of your instrument")
+                    sound2sound_save_instrument_button = gr.Button(variant="primary",
+                                                                  value="Save instrument",
+                                                                  scale=1)
+
+        with gr.Row(variant="panel"):
+            sound2sound_origin_upload_latent_representation_image = gr.Image(label="Original latent representation",
+                                                                             type="numpy", height=800,
+                                                                             visible=True)
+            sound2sound_origin_upload_quantized_latent_representation_image = gr.Image(
+                label="Original quantized latent representation", type="numpy", height=800, visible=True)
+
+            sound2sound_origin_microphone_latent_representation_image = gr.Image(label="Original latent representation",
+                                                                                 type="numpy", height=800,
+                                                                                 visible=False)
+            sound2sound_origin_microphone_quantized_latent_representation_image = gr.Image(
+                label="Original quantized latent representation", type="numpy", height=800, visible=False)
+
+            sound2sound_new_sound_latent_representation_image = gr.Image(label="New latent representation",
+                                                                         type="numpy", height=800)
+            sound2sound_new_sound_quantized_latent_representation_image = gr.Image(
+                label="New sound quantized latent representation", type="numpy", height=800)
+
+    sound2sound_origin_upload_audio.change(receive_uopoad_origin_audio,
+                                           inputs=[sound2sound_duration_slider, sound2sound_origin_source_radio, sound2sound_origin_upload_audio,
+                                                   sound2sound_origin_microphone_audio, inpaintWithText_state],
+                                           outputs=[sound2sound_origin_spectrogram_upload_image,
+                                                    sound2sound_origin_phase_upload_image,
+                                                    sound2sound_origin_spectrogram_microphone_image,
+                                                    sound2sound_origin_phase_microphone_image,
+                                                    sound2sound_origin_upload_latent_representation_image,
+                                                    sound2sound_origin_upload_quantized_latent_representation_image,
+                                                    sound2sound_origin_microphone_latent_representation_image,
+                                                    sound2sound_origin_microphone_quantized_latent_representation_image,
+                                                    inpaintWithText_state])
+    sound2sound_origin_microphone_audio.change(receive_uopoad_origin_audio,
+                                               inputs=[sound2sound_duration_slider, sound2sound_origin_source_radio, sound2sound_origin_upload_audio,
+                                                       sound2sound_origin_microphone_audio, inpaintWithText_state],
+                                               outputs=[sound2sound_origin_spectrogram_upload_image,
+                                                        sound2sound_origin_phase_upload_image,
+                                                        sound2sound_origin_spectrogram_microphone_image,
+                                                        sound2sound_origin_phase_microphone_image,
+                                                        sound2sound_origin_upload_latent_representation_image,
+                                                        sound2sound_origin_upload_quantized_latent_representation_image,
+                                                        sound2sound_origin_microphone_latent_representation_image,
+                                                        sound2sound_origin_microphone_quantized_latent_representation_image,
+                                                        inpaintWithText_state])
+
+    sound2sound_sample_button.click(sound2sound_sample,
+                                    inputs=[sound2sound_origin_spectrogram_upload_image,
+                                            sound2sound_origin_spectrogram_microphone_image,
+                                            text2sound_prompts_textbox,
+                                            text2sound_negative_prompts_textbox,
+                                            sound2sound_batchsize_slider,
+                                            sound2sound_guidance_scale_slider,
+                                            sound2sound_sampler_radio,
+                                            sound2sound_sample_steps_slider,
+                                            sound2sound_origin_source_radio,
+                                            sound2sound_noising_strength_slider,
+                                            sound2sound_seed_textbox,
+                                            sound2sound_inpaint_area_radio,
+                                            mask_time_begin_slider,
+                                            mask_time_end_slider,
+                                            mask_frequency_begin_slider,
+                                            mask_frequency_end_slider,
+                                            inpaintWithText_state],
+                                    outputs=[sound2sound_new_sound_latent_representation_image,
+                                             sound2sound_new_sound_quantized_latent_representation_image,
+                                             sound2sound_new_sound_spectrogram_image,
+                                             sound2sound_new_sound_phase_image,
+                                             sound2sound_new_sound_audio,
+                                             sound2sound_sample_index_slider,
+                                             sound2sound_seed_textbox,
+                                             inpaintWithText_state])
+
+    sound2sound_sample_index_slider.change(show_sound2sound_sample,
+                                           inputs=[sound2sound_sample_index_slider, inpaintWithText_state],
+                                           outputs=[sound2sound_new_sound_latent_representation_image,
+                                                    sound2sound_new_sound_quantized_latent_representation_image,
+                                                    sound2sound_new_sound_spectrogram_image,
+                                                    sound2sound_new_sound_phase_image,
+                                                    sound2sound_new_sound_audio])
+
+    sound2sound_origin_source_radio.change(sound2sound_switch_origin_source,
+                                           inputs=[sound2sound_origin_source_radio],
+                                           outputs=[sound2sound_origin_upload_audio,
+                                                    sound2sound_origin_microphone_audio,
+                                                    sound2sound_origin_spectrogram_upload_image,
+                                                    sound2sound_origin_phase_upload_image,
+                                                    sound2sound_origin_spectrogram_microphone_image,
+                                                    sound2sound_origin_phase_microphone_image,
+                                                    sound2sound_origin_upload_latent_representation_image,
+                                                    sound2sound_origin_upload_quantized_latent_representation_image,
+                                                    sound2sound_origin_microphone_latent_representation_image,
+                                                    sound2sound_origin_microphone_quantized_latent_representation_image])
+
+    sound2sound_save_instrument_button.click(save_virtual_instrument,
+                                        inputs=[sound2sound_sample_index_slider,
+                                                sound2sound_instrument_name_textbox,
+                                                inpaintWithText_state,
+                                                virtual_instruments_state],
+                                        outputs=[virtual_instruments_state,
+                                                  sound2sound_instrument_name_textbox])

webUI/natural_language_guided/rec.py

@@ -0,0 +1,190 @@
+import gradio as gr
+
+from data_generation.nsynth import get_nsynth_dataloader
+from webUI.natural_language_guided_STFT.utils import encodeBatch2GradioOutput_STFT, InputBatch2Encode_STFT, \
+    latent_representation_to_Gradio_image
+
+
+def get_recSTFT_module(gradioWebUI, reconstruction_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_encoder = gradioWebUI.VAE_encoder
+    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 generate_reconstruction_samples(sample_source, batchsize_slider, encodeCache,
+                                        reconstruction_samples):
+
+        vae_batchsize = int(batchsize_slider)
+
+        if sample_source == "text2sound_trainSTFT":
+            training_dataset_path = f'data/NSynth/nsynth-STFT-train-52.hdf5'  # Make sure to use your actual path
+            iterator = get_nsynth_dataloader(training_dataset_path, batch_size=vae_batchsize, shuffle=True,
+                                                        get_latent_representation=False, with_meta_data=False,
+                                                        task="STFT")
+        elif sample_source == "text2sound_validSTFT":
+            training_dataset_path = f'data/NSynth/nsynth-STFT-valid-52.hdf5'  # Make sure to use your actual path
+            iterator = get_nsynth_dataloader(training_dataset_path, batch_size=vae_batchsize, shuffle=True,
+                                                        get_latent_representation=False, with_meta_data=False,
+                                                        task="STFT")
+        elif sample_source == "text2sound_testSTFT":
+            training_dataset_path = f'data/NSynth/nsynth-STFT-test-52.hdf5'  # Make sure to use your actual path
+            iterator = get_nsynth_dataloader(training_dataset_path, batch_size=vae_batchsize, shuffle=True,
+                                                        get_latent_representation=False, with_meta_data=False,
+                                                        task="STFT")
+        else:
+            raise NotImplementedError()
+
+        spectrogram_batch = next(iter(iterator))
+
+        origin_flipped_log_spectrums, origin_flipped_phases, origin_signals, latent_representations, quantized_latent_representations = InputBatch2Encode_STFT(
+            VAE_encoder, spectrogram_batch, resolution=(512, width * VAE_scale), quantizer=VAE_quantizer, squared=squared)
+
+        latent_representation_gradio_images, quantized_latent_representation_gradio_images = [], []
+        for i in range(vae_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]))
+
+        if quantized_latent_representations is None:
+            quantized_latent_representations = latent_representations
+        reconstruction_flipped_log_spectrums, reconstruction_flipped_phases, reconstruction_signals, reconstruction_flipped_log_spectrums_WOA, reconstruction_flipped_phases_WOA, reconstruction_signals_WOA = encodeBatch2GradioOutput_STFT(VAE_decoder,
+                                                                                                            quantized_latent_representations,
+                                                                                                            resolution=(
+                                                                                                                512,
+                                                                                                                width * VAE_scale),
+                                                                                                            original_STFT_batch=spectrogram_batch
+                                                                                                     )
+
+        reconstruction_samples["origin_flipped_log_spectrums"] = origin_flipped_log_spectrums
+        reconstruction_samples["origin_flipped_phases"] = origin_flipped_phases
+        reconstruction_samples["origin_signals"] = origin_signals
+        reconstruction_samples["latent_representation_gradio_images"] = latent_representation_gradio_images
+        reconstruction_samples[
+            "quantized_latent_representation_gradio_images"] = quantized_latent_representation_gradio_images
+        reconstruction_samples[
+            "reconstruction_flipped_log_spectrums"] = reconstruction_flipped_log_spectrums
+        reconstruction_samples[
+            "reconstruction_flipped_phases"] = reconstruction_flipped_phases
+        reconstruction_samples["reconstruction_signals"] = reconstruction_signals
+        reconstruction_samples[
+            "reconstruction_flipped_log_spectrums_WOA"] = reconstruction_flipped_log_spectrums_WOA
+        reconstruction_samples[
+            "reconstruction_flipped_phases_WOA"] = reconstruction_flipped_phases_WOA
+        reconstruction_samples["reconstruction_signals_WOA"] = reconstruction_signals_WOA
+        reconstruction_samples["sampleRate"] = sample_rate
+
+        latent_representation_gradio_image = reconstruction_samples["latent_representation_gradio_images"][0]
+        quantized_latent_representation_gradio_image = \
+            reconstruction_samples["quantized_latent_representation_gradio_images"][0]
+        origin_flipped_log_spectrum = reconstruction_samples["origin_flipped_log_spectrums"][0]
+        origin_flipped_phase = reconstruction_samples["origin_flipped_phases"][0]
+        origin_signal = reconstruction_samples["origin_signals"][0]
+        reconstruction_flipped_log_spectrum = reconstruction_samples["reconstruction_flipped_log_spectrums"][0]
+        reconstruction_flipped_phase = reconstruction_samples["reconstruction_flipped_phases"][0]
+        reconstruction_signal = reconstruction_samples["reconstruction_signals"][0]
+        reconstruction_flipped_log_spectrum_WOA = reconstruction_samples["reconstruction_flipped_log_spectrums_WOA"][0]
+        reconstruction_flipped_phase_WOA = reconstruction_samples["reconstruction_flipped_phases_WOA"][0]
+        reconstruction_signal_WOA = reconstruction_samples["reconstruction_signals_WOA"][0]
+
+        return {origin_amplitude_image_output: origin_flipped_log_spectrum,
+                origin_phase_image_output: origin_flipped_phase,
+                origin_audio_output: (sample_rate, origin_signal),
+                latent_representation_image_output: latent_representation_gradio_image,
+                quantized_latent_representation_image_output: quantized_latent_representation_gradio_image,
+                reconstruction_amplitude_image_output: reconstruction_flipped_log_spectrum,
+                reconstruction_phase_image_output: reconstruction_flipped_phase,
+                reconstruction_audio_output: (sample_rate, reconstruction_signal),
+                reconstruction_amplitude_image_output_WOA: reconstruction_flipped_log_spectrum_WOA,
+                reconstruction_phase_image_output_WOA: reconstruction_flipped_phase_WOA,
+                reconstruction_audio_output_WOA: (sample_rate, reconstruction_signal_WOA),
+                sample_index_slider: gr.update(minimum=0, maximum=vae_batchsize - 1, value=0, step=1.0,
+                                               label="Sample index.",
+                                               info="Slide to view other samples", scale=1, visible=True),
+                reconstruction_state: encodeCache,
+                reconstruction_samples_state: reconstruction_samples}
+
+    def show_reconstruction_sample(sample_index, encodeCache_state, reconstruction_samples_state):
+        sample_index = int(sample_index)
+        sampleRate = reconstruction_samples_state["sampleRate"]
+        latent_representation_gradio_image = reconstruction_samples_state["latent_representation_gradio_images"][
+            sample_index]
+        quantized_latent_representation_gradio_image = \
+            reconstruction_samples_state["quantized_latent_representation_gradio_images"][sample_index]
+        origin_flipped_log_spectrum = reconstruction_samples_state["origin_flipped_log_spectrums"][sample_index]
+        origin_flipped_phase = reconstruction_samples_state["origin_flipped_phases"][sample_index]
+        origin_signal = reconstruction_samples_state["origin_signals"][sample_index]
+        reconstruction_flipped_log_spectrum = reconstruction_samples_state["reconstruction_flipped_log_spectrums"][
+            sample_index]
+        reconstruction_flipped_phase = reconstruction_samples_state["reconstruction_flipped_phases"][
+            sample_index]
+        reconstruction_signal = reconstruction_samples_state["reconstruction_signals"][sample_index]
+        reconstruction_flipped_log_spectrum_WOA = reconstruction_samples_state["reconstruction_flipped_log_spectrums_WOA"][
+            sample_index]
+        reconstruction_flipped_phase_WOA = reconstruction_samples_state["reconstruction_flipped_phases_WOA"][
+            sample_index]
+        reconstruction_signal_WOA = reconstruction_samples_state["reconstruction_signals_WOA"][sample_index]
+        return origin_flipped_log_spectrum, origin_flipped_phase, (sampleRate, origin_signal), \
+            latent_representation_gradio_image, quantized_latent_representation_gradio_image, \
+            reconstruction_flipped_log_spectrum, reconstruction_flipped_phase, (sampleRate, reconstruction_signal), \
+            reconstruction_flipped_log_spectrum_WOA, reconstruction_flipped_phase_WOA, (sampleRate, reconstruction_signal_WOA), \
+            encodeCache_state, reconstruction_samples_state
+
+    with gr.Tab("Reconstruction"):
+        reconstruction_samples_state = gr.State(value={})
+        gr.Markdown("Test reconstruction.")
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                sample_source_radio = gr.Radio(
+                    choices=["synthetic", "external", "text2sound_trainSTFT", "text2sound_testSTFT", "text2sound_validSTFT"],
+                    value="text2sound_trainf", info="Info placeholder", scale=2)
+                batchsize_slider = gr.Slider(minimum=1., maximum=16., value=4., step=1.,
+                                             label="batchsize")
+            with gr.Column():
+                generate_button = gr.Button(variant="primary", value="Generate reconstruction samples", scale=1)
+                sample_index_slider = gr.Slider(minimum=0, maximum=3, value=0, step=1.0, label="Sample index.",
+                                                info="Slide to view other samples", scale=1, visible=False)
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                origin_amplitude_image_output = gr.Image(label="Spectrogram", type="numpy", height=300, width=100, scale=1)
+                origin_phase_image_output = gr.Image(label="Phase", type="numpy", height=300, width=100, scale=1)
+                origin_audio_output = gr.Audio(type="numpy", label="Play the example!")
+            with gr.Column():
+                reconstruction_amplitude_image_output = gr.Image(label="Spectrogram", type="numpy", height=300, width=100, scale=1)
+                reconstruction_phase_image_output = gr.Image(label="Phase", type="numpy", height=300, width=100, scale=1)
+                reconstruction_audio_output = gr.Audio(type="numpy", label="Play the example!")
+            with gr.Column():
+                reconstruction_amplitude_image_output_WOA = gr.Image(label="Spectrogram", type="numpy", height=300, width=100, scale=1)
+                reconstruction_phase_image_output_WOA = gr.Image(label="Phase", type="numpy", height=300, width=100, scale=1)
+                reconstruction_audio_output_WOA = gr.Audio(type="numpy", label="Play the example!")
+        with gr.Row(variant="panel", equal_height=True):
+            latent_representation_image_output = gr.Image(label="latent_representation", type="numpy", height=300, width=100)
+            quantized_latent_representation_image_output = gr.Image(label="quantized", type="numpy", height=300, width=100)
+
+    generate_button.click(generate_reconstruction_samples,
+                          inputs=[sample_source_radio, batchsize_slider, reconstruction_state,
+                                  reconstruction_samples_state],
+                          outputs=[origin_amplitude_image_output, origin_phase_image_output, origin_audio_output,
+                                   latent_representation_image_output, quantized_latent_representation_image_output,
+                                   reconstruction_amplitude_image_output, reconstruction_phase_image_output, reconstruction_audio_output,
+                                   reconstruction_amplitude_image_output_WOA, reconstruction_phase_image_output_WOA, reconstruction_audio_output_WOA,
+                                   sample_index_slider, reconstruction_state, reconstruction_samples_state])
+
+    sample_index_slider.change(show_reconstruction_sample,
+                               inputs=[sample_index_slider, reconstruction_state, reconstruction_samples_state],
+                               outputs=[origin_amplitude_image_output, origin_phase_image_output, origin_audio_output,
+                                        latent_representation_image_output, quantized_latent_representation_image_output,
+                                        reconstruction_amplitude_image_output, reconstruction_phase_image_output, reconstruction_audio_output,
+                                        reconstruction_amplitude_image_output_WOA, reconstruction_phase_image_output_WOA, reconstruction_audio_output_WOA,
+                                        reconstruction_state, reconstruction_samples_state])

webUI/natural_language_guided/sound2sound_with_text.py

@@ -0,0 +1,416 @@
+import gradio as gr
+import librosa
+import numpy as np
+import torch
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import pad_STFT, encode_stft
+from tools import safe_int, adjust_audio_length
+from webUI.natural_language_guided.utils import InputBatch2Encode_STFT, encodeBatch2GradioOutput_STFT, \
+    latent_representation_to_Gradio_image
+
+
+def get_sound2sound_with_text_module(gradioWebUI, sound2sound_with_text_state, virtual_instruments_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_encoder = gradioWebUI.VAE_encoder
+    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 receive_upload_origin_audio(sound2sound_duration, sound2sound_origin_source,
+                                    sound2sound_origin_upload, sound2sound_origin_microphone,
+                                    sound2sound_with_text_dict, virtual_instruments_dict):
+
+        if sound2sound_origin_source == "upload":
+            origin_sr, origin_audio = sound2sound_origin_upload
+        else:
+            origin_sr, origin_audio = sound2sound_origin_microphone
+
+        origin_audio = origin_audio / np.max(np.abs(origin_audio))
+
+        width = int(time_resolution*((sound2sound_duration+1)/4) / VAE_scale)
+        audio_length = 256 * (VAE_scale * width - 1)
+        origin_audio = adjust_audio_length(origin_audio, audio_length, origin_sr, sample_rate)
+
+        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)
+
+        # Todo: justify batchsize to 1
+        origin_spectrogram_batch_tensor = torch.from_numpy(
+            np.repeat(encoded_D[np.newaxis, :, :, :], 1, axis=0)).float().to(device)
+
+        # Todo: remove hard-coding
+        origin_flipped_log_spectrums, origin_flipped_phases, origin_signals, origin_latent_representations, quantized_origin_latent_representations = InputBatch2Encode_STFT(
+            VAE_encoder, origin_spectrogram_batch_tensor, resolution=(512, width * VAE_scale), quantizer=VAE_quantizer, squared=squared)
+
+        default_condition = CLAP.get_text_features(**CLAP_tokenizer([""], padding=True, return_tensors="pt"))[0].to("cpu").detach().numpy()
+
+        if sound2sound_origin_source == "upload":
+            sound2sound_with_text_dict["origin_upload_latent_representations"] = origin_latent_representations.tolist()
+            sound2sound_with_text_dict[
+                "sound2sound_origin_upload_latent_representation_image"] = latent_representation_to_Gradio_image(
+                origin_latent_representations[0]).tolist()
+            sound2sound_with_text_dict[
+                "sound2sound_origin_upload_quantized_latent_representation_image"] = latent_representation_to_Gradio_image(
+                quantized_origin_latent_representations[0]).tolist()
+
+            virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+            virtual_instrument = {"condition": default_condition,
+                                  "negative_condition": default_condition,  # care!!!
+                                  "CFG": 1,
+                                  "latent_representation": origin_latent_representations[0].to("cpu").detach().numpy(),
+                                  "quantized_latent_representation": quantized_origin_latent_representations[0].to("cpu").detach().numpy(),
+                                  "sampler": "ddim",
+                                  "signal": (sample_rate, origin_audio),
+                                  "spectrogram_gradio_image": origin_flipped_log_spectrums[0],
+                                  "phase_gradio_image": origin_flipped_phases[0]}
+            virtual_instruments["s2sup"] = virtual_instrument
+            virtual_instruments_dict["virtual_instruments"] = virtual_instruments
+
+            return {sound2sound_origin_spectrogram_upload_image: origin_flipped_log_spectrums[0],
+                    sound2sound_origin_phase_upload_image: origin_flipped_phases[0],
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(),
+                    sound2sound_origin_phase_microphone_image: gr.update(),
+                    sound2sound_origin_upload_latent_representation_image: latent_representation_to_Gradio_image(
+                        origin_latent_representations[0]),
+                    sound2sound_origin_upload_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                        quantized_origin_latent_representations[0]),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(),
+                    sound2sound_with_text_state: sound2sound_with_text_dict,
+                    virtual_instruments_state: virtual_instruments_dict}
+        else:
+            sound2sound_with_text_dict["origin_microphone_latent_representations"] = origin_latent_representations.tolist()
+            sound2sound_with_text_dict[
+                "sound2sound_origin_microphone_latent_representation_image"] = latent_representation_to_Gradio_image(
+                origin_latent_representations[0]).tolist()
+            sound2sound_with_text_dict[
+                "sound2sound_origin_microphone_quantized_latent_representation_image"] = latent_representation_to_Gradio_image(
+                quantized_origin_latent_representations[0]).tolist()
+
+            virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+            virtual_instrument = {"condition": default_condition,
+                                  "negative_condition": default_condition,  # care!!!
+                                  "CFG": 1,
+                                  "latent_representation": origin_latent_representations[0],
+                                  "quantized_latent_representation": quantized_origin_latent_representations[0],
+                                  "sampler": "ddim",
+                                  "signal": origin_audio,
+                                  "spectrogram_gradio_image": origin_flipped_log_spectrums[0]}
+            virtual_instruments["s2sre"] = virtual_instrument
+            virtual_instruments_dict["virtual_instruments"] = virtual_instruments
+
+            return {sound2sound_origin_spectrogram_upload_image: gr.update(),
+                    sound2sound_origin_phase_upload_image: gr.update(),
+                    sound2sound_origin_spectrogram_microphone_image: origin_flipped_log_spectrums[0],
+                    sound2sound_origin_phase_microphone_image: origin_flipped_phases[0],
+                    sound2sound_origin_upload_latent_representation_image: gr.update(),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(),
+                    sound2sound_origin_microphone_latent_representation_image: latent_representation_to_Gradio_image(
+                        origin_latent_representations[0]),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                        quantized_origin_latent_representations[0]),
+                    sound2sound_with_text_state: sound2sound_with_text_dict,
+                    virtual_instruments_state: virtual_instruments_dict}
+
+    def sound2sound_sample(sound2sound_prompts, sound2sound_negative_prompts, sound2sound_batchsize,
+                           sound2sound_guidance_scale, sound2sound_sampler,
+                           sound2sound_sample_steps,
+                           sound2sound_origin_source,
+                           sound2sound_noising_strength, sound2sound_seed, sound2sound_dict, virtual_instruments_dict):
+
+        # input processing
+        sound2sound_seed = safe_int(sound2sound_seed, 12345678)
+        sound2sound_batchsize = int(sound2sound_batchsize)
+        noising_strength = sound2sound_noising_strength
+        sound2sound_sample_steps = int(sound2sound_sample_steps)
+        CFG = int(sound2sound_guidance_scale)
+
+        if sound2sound_origin_source == "upload":
+            origin_latent_representations = torch.tensor(
+                sound2sound_dict["origin_upload_latent_representations"]).repeat(sound2sound_batchsize, 1, 1, 1).to(
+                device)
+        elif sound2sound_origin_source == "microphone":
+            origin_latent_representations = torch.tensor(
+                sound2sound_dict["origin_microphone_latent_representations"]).repeat(sound2sound_batchsize, 1, 1, 1).to(
+                device)
+        else:
+            print("Input source not in ['upload', 'microphone']!")
+            raise NotImplementedError()
+
+        # sound2sound
+        text2sound_embedding = \
+            CLAP.get_text_features(**CLAP_tokenizer([sound2sound_prompts], padding=True, return_tensors="pt"))[0].to(
+                device)
+
+        mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy)
+        unconditional_condition = \
+            CLAP.get_text_features(**CLAP_tokenizer([sound2sound_negative_prompts], padding=True, return_tensors="pt"))[
+                0]
+        mySampler.activate_classifier_free_guidance(CFG, unconditional_condition.to(device))
+
+        normalized_sample_steps = int(sound2sound_sample_steps / noising_strength)
+        mySampler.respace(list(np.linspace(0, timesteps - 1, normalized_sample_steps, dtype=np.int32)))
+
+        condition = text2sound_embedding.repeat(sound2sound_batchsize, 1)
+
+        # Todo: remove-hard coding
+        width = origin_latent_representations.shape[-1]
+        new_sound_latent_representations, initial_noise = \
+            mySampler.img_guided_sample(model=uNet, shape=(sound2sound_batchsize, channels, height, width),
+                                        seed=sound2sound_seed,
+                                        noising_strength=noising_strength,
+                                        guide_img=origin_latent_representations, return_tensor=True,
+                                        condition=condition,
+                                        sampler=sound2sound_sampler)
+
+        new_sound_latent_representations = new_sound_latent_representations[-1]
+
+        # Quantize new sound latent representations
+        quantized_new_sound_latent_representations, loss, (_, _, _) = VAE_quantizer(new_sound_latent_representations)
+
+        new_sound_flipped_log_spectrums, new_sound_flipped_phases, new_sound_signals, _, _, _ = encodeBatch2GradioOutput_STFT(VAE_decoder,
+                                                                                                            quantized_new_sound_latent_representations,
+                                                                                                            resolution=(
+                                                                                                                512,
+                                                                                                                width * VAE_scale),
+                                                                                                            original_STFT_batch=None
+                                                                                                     )
+
+
+
+        new_sound_latent_representation_gradio_images = []
+        new_sound_quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_phase_gradio_images = []
+        new_sound_rec_signals_gradio = []
+        for i in range(sound2sound_batchsize):
+            new_sound_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(new_sound_latent_representations[i]))
+            new_sound_quantized_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(quantized_new_sound_latent_representations[i]))
+            new_sound_spectrogram_gradio_images.append(new_sound_flipped_log_spectrums[i])
+            new_sound_phase_gradio_images.append(new_sound_flipped_phases[i])
+            new_sound_rec_signals_gradio.append((sample_rate, new_sound_signals[i]))
+
+        sound2sound_dict[
+            "new_sound_latent_representation_gradio_images"] = new_sound_latent_representation_gradio_images
+        sound2sound_dict[
+            "new_sound_quantized_latent_representation_gradio_images"] = new_sound_quantized_latent_representation_gradio_images
+        sound2sound_dict["new_sound_spectrogram_gradio_images"] = new_sound_spectrogram_gradio_images
+        sound2sound_dict["new_sound_phase_gradio_images"] = new_sound_phase_gradio_images
+        sound2sound_dict["new_sound_rec_signals_gradio"] = new_sound_rec_signals_gradio
+
+        return {sound2sound_new_sound_latent_representation_image: latent_representation_to_Gradio_image(
+            new_sound_latent_representations[0]),
+            sound2sound_new_sound_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                quantized_new_sound_latent_representations[0]),
+            sound2sound_new_sound_spectrogram_image: new_sound_flipped_log_spectrums[0],
+            sound2sound_new_sound_phase_image: new_sound_flipped_phases[0],
+            sound2sound_new_sound_audio: (sample_rate, new_sound_signals[0]),
+            sound2sound_sample_index_slider: gr.update(minimum=0, maximum=sound2sound_batchsize - 1, value=0,
+                                                       step=1.0,
+                                                       visible=True,
+                                                       label="Sample index",
+                                                       info="Swipe to view other samples"),
+            sound2sound_seed_textbox: sound2sound_seed,
+            sound2sound_with_text_state: sound2sound_dict,
+            virtual_instruments_state: virtual_instruments_dict}
+
+    def show_sound2sound_sample(sound2sound_sample_index, sound2sound_with_text_dict):
+        sample_index = int(sound2sound_sample_index)
+        return {sound2sound_new_sound_latent_representation_image:
+                    sound2sound_with_text_dict["new_sound_latent_representation_gradio_images"][sample_index],
+                sound2sound_new_sound_quantized_latent_representation_image:
+                    sound2sound_with_text_dict["new_sound_quantized_latent_representation_gradio_images"][sample_index],
+                sound2sound_new_sound_spectrogram_image: sound2sound_with_text_dict["new_sound_spectrogram_gradio_images"][
+                    sample_index],
+                sound2sound_new_sound_phase_image: sound2sound_with_text_dict["new_sound_phase_gradio_images"][
+                    sample_index],
+                sound2sound_new_sound_audio: sound2sound_with_text_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+    def sound2sound_switch_origin_source(sound2sound_origin_source):
+
+        if sound2sound_origin_source == "upload":
+            return {sound2sound_origin_upload_audio: gr.update(visible=True),
+                    sound2sound_origin_microphone_audio: gr.update(visible=False),
+                    sound2sound_origin_spectrogram_upload_image: gr.update(visible=True),
+                    sound2sound_origin_phase_upload_image: gr.update(visible=True),
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(visible=False),
+                    sound2sound_origin_phase_microphone_image: gr.update(visible=False),
+                    sound2sound_origin_upload_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(visible=False)}
+        elif sound2sound_origin_source == "microphone":
+            return {sound2sound_origin_upload_audio: gr.update(visible=False),
+                    sound2sound_origin_microphone_audio: gr.update(visible=True),
+                    sound2sound_origin_spectrogram_upload_image: gr.update(visible=False),
+                    sound2sound_origin_phase_upload_image: gr.update(visible=False),
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(visible=True),
+                    sound2sound_origin_phase_microphone_image: gr.update(visible=True),
+                    sound2sound_origin_upload_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(visible=True)}
+        else:
+            print("Input source not in ['upload', 'microphone']!")
+
+    with gr.Tab("Sound2Sound"):
+        gr.Markdown("Generate new sound based on a given sound!")
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=3):
+                sound2sound_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):
+                sound2sound_sample_button = gr.Button(variant="primary", value="Generate", scale=1)
+
+                sound2sound_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):
+                with gr.Tab("Origin sound"):
+                    sound2sound_duration_slider = gradioWebUI.get_duration_slider()
+                    sound2sound_origin_source_radio = gr.Radio(choices=["upload", "microphone"], value="upload",
+                                                               label="Input source")
+
+                    sound2sound_origin_upload_audio = gr.Audio(type="numpy", label="Upload", source="upload",
+                                                               interactive=True, visible=True)
+                    sound2sound_origin_microphone_audio = gr.Audio(type="numpy", label="Record", source="microphone",
+                                                                   interactive=True, visible=False)
+                    with gr.Row(variant="panel"):
+                        sound2sound_origin_spectrogram_upload_image = gr.Image(label="Original upload spectrogram",
+                                                                               type="numpy", height=600,
+                                                                            visible=True)
+                        sound2sound_origin_phase_upload_image = gr.Image(label="Original upload phase",
+                                                                               type="numpy", height=600,
+                                                                               visible=True)
+                        sound2sound_origin_spectrogram_microphone_image = gr.Image(label="Original microphone spectrogram",
+                                                                                   type="numpy", height=600,
+                                                                                   visible=False)
+                        sound2sound_origin_phase_microphone_image = gr.Image(label="Original microphone phase",
+                                                                                   type="numpy", height=600,
+                                                                                   visible=False)
+
+                with gr.Tab("Sound2sound settings"):
+                    sound2sound_sample_steps_slider = gradioWebUI.get_sample_steps_slider()
+                    sound2sound_sampler_radio = gradioWebUI.get_sampler_radio()
+                    sound2sound_batchsize_slider = gradioWebUI.get_batchsize_slider()
+                    sound2sound_noising_strength_slider = gradioWebUI.get_noising_strength_slider()
+                    sound2sound_guidance_scale_slider = gradioWebUI.get_guidance_scale_slider()
+                    sound2sound_seed_textbox = gradioWebUI.get_seed_textbox()
+
+            with gr.Column(scale=1):
+                sound2sound_new_sound_audio = gr.Audio(type="numpy", label="Play new sound", interactive=False)
+                with gr.Row(variant="panel"):
+                    sound2sound_new_sound_spectrogram_image = gr.Image(label="New sound spectrogram", type="numpy",
+                                                                       height=600, scale=1)
+                    sound2sound_new_sound_phase_image = gr.Image(label="New sound phase", type="numpy",
+                                                                       height=600, scale=1)
+
+        with gr.Row(variant="panel"):
+            sound2sound_origin_upload_latent_representation_image = gr.Image(label="Original latent representation",
+                                                                             type="numpy", height=800,
+                                                                             visible=True)
+            sound2sound_origin_upload_quantized_latent_representation_image = gr.Image(
+                label="Original quantized latent representation", type="numpy", height=800, visible=True)
+
+            sound2sound_origin_microphone_latent_representation_image = gr.Image(label="Original latent representation",
+                                                                                 type="numpy", height=800,
+                                                                                 visible=False)
+            sound2sound_origin_microphone_quantized_latent_representation_image = gr.Image(
+                label="Original quantized latent representation", type="numpy", height=800, visible=False)
+
+            sound2sound_new_sound_latent_representation_image = gr.Image(label="New latent representation",
+                                                                         type="numpy", height=800)
+            sound2sound_new_sound_quantized_latent_representation_image = gr.Image(
+                label="New sound quantized latent representation", type="numpy", height=800)
+
+    sound2sound_origin_upload_audio.change(receive_upload_origin_audio,
+                                           inputs=[sound2sound_duration_slider, sound2sound_origin_source_radio,
+                                                   sound2sound_origin_upload_audio,
+                                                   sound2sound_origin_microphone_audio, sound2sound_with_text_state,
+                                                   virtual_instruments_state],
+                                           outputs=[sound2sound_origin_spectrogram_upload_image,
+                                                    sound2sound_origin_phase_upload_image,
+                                                    sound2sound_origin_spectrogram_microphone_image,
+                                                    sound2sound_origin_phase_microphone_image,
+                                                    sound2sound_origin_upload_latent_representation_image,
+                                                    sound2sound_origin_upload_quantized_latent_representation_image,
+                                                    sound2sound_origin_microphone_latent_representation_image,
+                                                    sound2sound_origin_microphone_quantized_latent_representation_image,
+                                                    sound2sound_with_text_state,
+                                                    virtual_instruments_state])
+
+    sound2sound_origin_microphone_audio.change(receive_upload_origin_audio,
+                                               inputs=[sound2sound_duration_slider,
+                                                       sound2sound_origin_source_radio, sound2sound_origin_upload_audio,
+                                                       sound2sound_origin_microphone_audio, sound2sound_with_text_state,
+                                                       virtual_instruments_state],
+                                               outputs=[sound2sound_origin_spectrogram_upload_image,
+                                                        sound2sound_origin_phase_upload_image,
+                                                        sound2sound_origin_spectrogram_microphone_image,
+                                                        sound2sound_origin_phase_microphone_image,
+                                                        sound2sound_origin_upload_latent_representation_image,
+                                                        sound2sound_origin_upload_quantized_latent_representation_image,
+                                                        sound2sound_origin_microphone_latent_representation_image,
+                                                        sound2sound_origin_microphone_quantized_latent_representation_image,
+                                                        sound2sound_with_text_state,
+                                                        virtual_instruments_state])
+
+    sound2sound_sample_button.click(sound2sound_sample,
+                                    inputs=[sound2sound_prompts_textbox,
+                                            text2sound_negative_prompts_textbox,
+                                            sound2sound_batchsize_slider,
+                                            sound2sound_guidance_scale_slider,
+                                            sound2sound_sampler_radio,
+                                            sound2sound_sample_steps_slider,
+                                            sound2sound_origin_source_radio,
+                                            sound2sound_noising_strength_slider,
+                                            sound2sound_seed_textbox,
+                                            sound2sound_with_text_state,
+                                            virtual_instruments_state],
+                                    outputs=[sound2sound_new_sound_latent_representation_image,
+                                             sound2sound_new_sound_quantized_latent_representation_image,
+                                             sound2sound_new_sound_spectrogram_image,
+                                             sound2sound_new_sound_phase_image,
+                                             sound2sound_new_sound_audio,
+                                             sound2sound_sample_index_slider,
+                                             sound2sound_seed_textbox,
+                                             sound2sound_with_text_state,
+                                             virtual_instruments_state])
+
+    sound2sound_sample_index_slider.change(show_sound2sound_sample,
+                                           inputs=[sound2sound_sample_index_slider, sound2sound_with_text_state],
+                                           outputs=[sound2sound_new_sound_latent_representation_image,
+                                                    sound2sound_new_sound_quantized_latent_representation_image,
+                                                    sound2sound_new_sound_spectrogram_image,
+                                                    sound2sound_new_sound_phase_image,
+                                                    sound2sound_new_sound_audio])
+
+    sound2sound_origin_source_radio.change(sound2sound_switch_origin_source,
+                                           inputs=[sound2sound_origin_source_radio],
+                                           outputs=[sound2sound_origin_upload_audio,
+                                                    sound2sound_origin_microphone_audio,
+                                                    sound2sound_origin_spectrogram_upload_image,
+                                                    sound2sound_origin_phase_upload_image,
+                                                    sound2sound_origin_spectrogram_microphone_image,
+                                                    sound2sound_origin_phase_microphone_image,
+                                                    sound2sound_origin_upload_latent_representation_image,
+                                                    sound2sound_origin_upload_quantized_latent_representation_image,
+                                                    sound2sound_origin_microphone_latent_representation_image,
+                                                    sound2sound_origin_microphone_quantized_latent_representation_image])

webUI/natural_language_guided/super_resolution_with_text.py

@@ -0,0 +1,387 @@
+import librosa
+import numpy as np
+import torch
+import gradio as gr
+from scipy.ndimage import zoom
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import adjust_audio_length, rescale, safe_int, pad_STFT, encode_stft
+from webUI.natural_language_guided_STFT.utils import latent_representation_to_Gradio_image
+from webUI.natural_language_guided_STFT.utils import InputBatch2Encode_STFT, encodeBatch2GradioOutput_STFT
+
+
+def get_super_resolution_with_text_module(gradioWebUI, inpaintWithText_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_encoder = gradioWebUI.VAE_encoder
+    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 receive_uopoad_origin_audio(sound2sound_duration, sound2sound_origin_source, sound2sound_origin_upload, sound2sound_origin_microphone,
+                                    inpaintWithText_dict):
+
+        if sound2sound_origin_source == "upload":
+            origin_sr, origin_audio = sound2sound_origin_upload
+        else:
+            origin_sr, origin_audio = sound2sound_origin_microphone
+
+        origin_audio = origin_audio / np.max(np.abs(origin_audio))
+
+        width = int(time_resolution*((sound2sound_duration+1)/4) / VAE_scale)
+        audio_length = 256 * (VAE_scale * width - 1)
+        origin_audio = adjust_audio_length(origin_audio, audio_length, origin_sr, sample_rate)
+
+        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)
+
+        # Todo: justify batchsize to 1
+        origin_spectrogram_batch_tensor = torch.from_numpy(
+            np.repeat(encoded_D[np.newaxis, :, :, :], 1, axis=0)).float().to(device)
+
+        # Todo: remove hard-coding
+        origin_flipped_log_spectrums, origin_flipped_phases, origin_signals, origin_latent_representations, quantized_origin_latent_representations = InputBatch2Encode_STFT(
+            VAE_encoder, origin_spectrogram_batch_tensor, resolution=(512, width * VAE_scale), quantizer=VAE_quantizer, squared=squared)
+
+        if sound2sound_origin_source == "upload":
+            inpaintWithText_dict["origin_upload_latent_representations"] = origin_latent_representations.tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_upload_latent_representation_image"] = latent_representation_to_Gradio_image(
+                origin_latent_representations[0]).tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_upload_quantized_latent_representation_image"] = latent_representation_to_Gradio_image(
+                quantized_origin_latent_representations[0]).tolist()
+            return {sound2sound_origin_spectrogram_upload_image: origin_flipped_log_spectrums[0],
+                    sound2sound_origin_phase_upload_image: origin_flipped_phases[0],
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(),
+                    sound2sound_origin_phase_microphone_image: gr.update(),
+                    sound2sound_origin_upload_latent_representation_image: latent_representation_to_Gradio_image(
+                        origin_latent_representations[0]),
+                    sound2sound_origin_upload_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                        quantized_origin_latent_representations[0]),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(),
+                    inpaintWithText_state: inpaintWithText_dict}
+        else:
+            inpaintWithText_dict["origin_microphone_latent_representations"] = origin_latent_representations.tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_microphone_latent_representation_image"] = latent_representation_to_Gradio_image(
+                origin_latent_representations[0]).tolist()
+            inpaintWithText_dict[
+                "sound2sound_origin_microphone_quantized_latent_representation_image"] = latent_representation_to_Gradio_image(
+                quantized_origin_latent_representations[0]).tolist()
+            return {sound2sound_origin_spectrogram_upload_image: origin_flipped_log_spectrums[0],
+                    sound2sound_origin_phase_upload_image: origin_flipped_phases[0],
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(),
+                    sound2sound_origin_phase_microphone_image: gr.update(),
+                    sound2sound_origin_upload_latent_representation_image: latent_representation_to_Gradio_image(
+                        origin_latent_representations[0]),
+                    sound2sound_origin_upload_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                        quantized_origin_latent_representations[0]),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(),
+                    inpaintWithText_state: inpaintWithText_dict}
+
+    def sound2sound_sample(sound2sound_origin_spectrogram_upload, sound2sound_origin_spectrogram_microphone,
+                           text2sound_prompts, text2sound_negative_prompts, sound2sound_batchsize,
+                           sound2sound_guidance_scale, sound2sound_sampler,
+                           sound2sound_sample_steps, sound2sound_origin_source,
+                           sound2sound_noising_strength, sound2sound_seed, sound2sound_inpaint_area, inpaintWithText_dict
+                           ):
+
+        # input preprocessing
+        sound2sound_seed = safe_int(sound2sound_seed, 12345678)
+        sound2sound_batchsize = int(sound2sound_batchsize)
+        noising_strength = sound2sound_noising_strength
+        sound2sound_sample_steps = int(sound2sound_sample_steps)
+        CFG = int(sound2sound_guidance_scale)
+
+        text2sound_embedding = \
+        CLAP.get_text_features(**CLAP_tokenizer([text2sound_prompts], padding=True, return_tensors="pt"))[0].to(device)
+
+        if sound2sound_origin_source == "upload":
+            origin_latent_representations = torch.tensor(
+                inpaintWithText_dict["origin_upload_latent_representations"]).repeat(sound2sound_batchsize, 1, 1, 1).to(
+                device)
+        elif sound2sound_origin_source == "microphone":
+            origin_latent_representations = torch.tensor(
+                inpaintWithText_dict["origin_microphone_latent_representations"]).repeat(sound2sound_batchsize, 1, 1, 1).to(
+                device)
+        else:
+            print("Input source not in ['upload', 'microphone']!")
+            raise NotImplementedError()
+
+        high_resolution_latent_representations = torch.zeros((sound2sound_batchsize, channels, 256, 64)).to(device)
+        high_resolution_latent_representations[:, :, :128, :] = origin_latent_representations
+        latent_mask = np.ones((256, 64))
+        latent_mask[192:, :] = 0.0
+        print(f"latent_mask mean: {np.mean(latent_mask)}")
+
+        if sound2sound_inpaint_area == "inpaint masked":
+            latent_mask = 1 - latent_mask
+        latent_mask = torch.from_numpy(latent_mask).unsqueeze(0).unsqueeze(1).repeat(sound2sound_batchsize, channels, 1,
+                                                                                     1).float().to(device)
+        latent_mask = torch.flip(latent_mask, [2])
+
+        mySampler = DiffSynthSampler(timesteps, height=height*2, 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))
+
+        normalized_sample_steps = int(sound2sound_sample_steps / noising_strength)
+
+        mySampler.respace(list(np.linspace(0, timesteps - 1, normalized_sample_steps, dtype=np.int32)))
+
+        # Todo: remove hard-coding
+        width = high_resolution_latent_representations.shape[-1]
+        condition = text2sound_embedding.repeat(sound2sound_batchsize, 1)
+
+        new_sound_latent_representations, initial_noise = \
+            mySampler.inpaint_sample(model=uNet, shape=(sound2sound_batchsize, channels, height*2, width),
+                                     seed=sound2sound_seed,
+                                     noising_strength=noising_strength,
+                                     guide_img=high_resolution_latent_representations, mask=latent_mask, return_tensor=True,
+                                     condition=condition, sampler=sound2sound_sampler)
+
+        new_sound_latent_representations = new_sound_latent_representations[-1]
+
+        # Quantize new sound latent representations
+        quantized_new_sound_latent_representations, loss, (_, _, _) = VAE_quantizer(new_sound_latent_representations)
+        new_sound_flipped_log_spectrums, new_sound_flipped_phases, new_sound_signals, _, _, _ = encodeBatch2GradioOutput_STFT(VAE_decoder,
+                                                                                                            quantized_new_sound_latent_representations,
+                                                                                                            resolution=(
+                                                                                                                1024,
+                                                                                                                width * VAE_scale),
+                                                                                                            original_STFT_batch=None
+                                                                                                     )
+
+        new_sound_latent_representation_gradio_images = []
+        new_sound_quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_phase_gradio_images = []
+        new_sound_rec_signals_gradio = []
+        for i in range(sound2sound_batchsize):
+            new_sound_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(new_sound_latent_representations[i]))
+            new_sound_quantized_latent_representation_gradio_images.append(
+                latent_representation_to_Gradio_image(quantized_new_sound_latent_representations[i]))
+            new_sound_spectrogram_gradio_images.append(new_sound_flipped_log_spectrums[i])
+            new_sound_phase_gradio_images.append(new_sound_flipped_phases[i])
+            new_sound_rec_signals_gradio.append((sample_rate, new_sound_signals[i]))
+
+        inpaintWithText_dict[
+            "new_sound_latent_representation_gradio_images"] = new_sound_latent_representation_gradio_images
+        inpaintWithText_dict[
+            "new_sound_quantized_latent_representation_gradio_images"] = new_sound_quantized_latent_representation_gradio_images
+        inpaintWithText_dict["new_sound_spectrogram_gradio_images"] = new_sound_spectrogram_gradio_images
+        inpaintWithText_dict["new_sound_phase_gradio_images"] = new_sound_phase_gradio_images
+        inpaintWithText_dict["new_sound_rec_signals_gradio"] = new_sound_rec_signals_gradio
+
+        return {sound2sound_new_sound_latent_representation_image: latent_representation_to_Gradio_image(
+            new_sound_latent_representations[0]),
+            sound2sound_new_sound_quantized_latent_representation_image: latent_representation_to_Gradio_image(
+                quantized_new_sound_latent_representations[0]),
+            sound2sound_new_sound_spectrogram_image: new_sound_flipped_log_spectrums[0],
+            sound2sound_new_sound_phase_image: new_sound_flipped_phases[0],
+            sound2sound_new_sound_audio: (sample_rate, new_sound_signals[0]),
+            sound2sound_sample_index_slider: gr.update(minimum=0, maximum=sound2sound_batchsize - 1, value=0,
+                                                       step=1.0,
+                                                       visible=True,
+                                                       label="Sample index",
+                                                       info="Swipe to view other samples"),
+            sound2sound_seed_textbox: sound2sound_seed,
+            inpaintWithText_state: inpaintWithText_dict}
+
+    def show_sound2sound_sample(sound2sound_sample_index, inpaintWithText_dict):
+        sample_index = int(sound2sound_sample_index)
+        return {sound2sound_new_sound_latent_representation_image:
+                    inpaintWithText_dict["new_sound_latent_representation_gradio_images"][sample_index],
+                sound2sound_new_sound_quantized_latent_representation_image:
+                    inpaintWithText_dict["new_sound_quantized_latent_representation_gradio_images"][sample_index],
+                sound2sound_new_sound_spectrogram_image: inpaintWithText_dict["new_sound_spectrogram_gradio_images"][
+                    sample_index],
+                sound2sound_new_sound_phase_image: inpaintWithText_dict["new_sound_phase_gradio_images"][
+                    sample_index],
+                sound2sound_new_sound_audio: inpaintWithText_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+    def sound2sound_switch_origin_source(sound2sound_origin_source):
+
+        if sound2sound_origin_source == "upload":
+            return {sound2sound_origin_upload_audio: gr.update(visible=True),
+                    sound2sound_origin_microphone_audio: gr.update(visible=False),
+                    sound2sound_origin_spectrogram_upload_image: gr.update(visible=True),
+                    sound2sound_origin_phase_upload_image: gr.update(visible=True),
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(visible=False),
+                    sound2sound_origin_phase_microphone_image: gr.update(visible=False),
+                    sound2sound_origin_upload_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(visible=False)}
+        elif sound2sound_origin_source == "microphone":
+            return {sound2sound_origin_upload_audio: gr.update(visible=False),
+                    sound2sound_origin_microphone_audio: gr.update(visible=True),
+                    sound2sound_origin_spectrogram_upload_image: gr.update(visible=False),
+                    sound2sound_origin_phase_upload_image: gr.update(visible=False),
+                    sound2sound_origin_spectrogram_microphone_image: gr.update(visible=True),
+                    sound2sound_origin_phase_microphone_image: gr.update(visible=True),
+                    sound2sound_origin_upload_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_upload_quantized_latent_representation_image: gr.update(visible=False),
+                    sound2sound_origin_microphone_latent_representation_image: gr.update(visible=True),
+                    sound2sound_origin_microphone_quantized_latent_representation_image: gr.update(visible=True)}
+        else:
+            print("Input source not in ['upload', 'microphone']!")
+
+    with gr.Tab("Super Resolution"):
+        gr.Markdown("Select the area to inpaint and use the prompt to guide the synthesis of a new sound!")
+        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):
+                sound2sound_sample_button = gr.Button(variant="primary", value="Generate", scale=1)
+
+                sound2sound_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):
+                with gr.Tab("Origin sound"):
+                    sound2sound_duration_slider = gradioWebUI.get_duration_slider()
+                    sound2sound_origin_source_radio = gr.Radio(choices=["upload", "microphone"], value="upload",
+                                                               label="Input source")
+
+                    sound2sound_origin_upload_audio = gr.Audio(type="numpy", label="Upload", source="upload",
+                                                               interactive=True, visible=True)
+                    sound2sound_origin_microphone_audio = gr.Audio(type="numpy", label="Record", source="microphone",
+                                                                   interactive=True, visible=False)
+                    with gr.Row(variant="panel"):
+                        sound2sound_origin_spectrogram_upload_image = gr.Image(label="Original upload spectrogram",
+                                                                               type="numpy", height=600,
+                                                                            visible=True, tool="sketch")
+                        sound2sound_origin_phase_upload_image = gr.Image(label="Original upload phase",
+                                                                               type="numpy", height=600,
+                                                                               visible=True)
+                        sound2sound_origin_spectrogram_microphone_image = gr.Image(label="Original microphone spectrogram",
+                                                                                   type="numpy", height=600,
+                                                                                   visible=False, tool="sketch")
+                        sound2sound_origin_phase_microphone_image = gr.Image(label="Original microphone phase",
+                                                                                   type="numpy", height=600,
+                                                                                   visible=False)
+                    sound2sound_inpaint_area_radio = gr.Radio(choices=["inpaint masked", "inpaint not masked"],
+                                                              value="inpaint masked")
+
+                with gr.Tab("Sound2sound settings"):
+                    sound2sound_sample_steps_slider = gradioWebUI.get_sample_steps_slider()
+                    sound2sound_sampler_radio = gradioWebUI.get_sampler_radio()
+                    sound2sound_batchsize_slider = gradioWebUI.get_batchsize_slider()
+                    sound2sound_noising_strength_slider = gradioWebUI.get_noising_strength_slider(default_noising_strength=1.0)
+                    sound2sound_guidance_scale_slider = gradioWebUI.get_guidance_scale_slider()
+                    sound2sound_seed_textbox = gradioWebUI.get_seed_textbox()
+
+
+            with gr.Column(scale=1):
+                sound2sound_new_sound_audio = gr.Audio(type="numpy", label="Play new sound", interactive=False)
+                with gr.Row(variant="panel"):
+                    sound2sound_new_sound_spectrogram_image = gr.Image(label="New sound spectrogram", type="numpy",
+                                                                       height=1200, scale=1)
+                    sound2sound_new_sound_phase_image = gr.Image(label="New sound phase", type="numpy",
+                                                                       height=1200, scale=1)
+
+        with gr.Row(variant="panel"):
+            sound2sound_origin_upload_latent_representation_image = gr.Image(label="Original latent representation",
+                                                                             type="numpy", height=1200,
+                                                                             visible=True)
+            sound2sound_origin_upload_quantized_latent_representation_image = gr.Image(
+                label="Original quantized latent representation", type="numpy", height=1200, visible=True)
+
+            sound2sound_origin_microphone_latent_representation_image = gr.Image(label="Original latent representation",
+                                                                                 type="numpy", height=1200,
+                                                                                 visible=False)
+            sound2sound_origin_microphone_quantized_latent_representation_image = gr.Image(
+                label="Original quantized latent representation", type="numpy", height=1200, visible=False)
+
+            sound2sound_new_sound_latent_representation_image = gr.Image(label="New latent representation",
+                                                                         type="numpy", height=1200)
+            sound2sound_new_sound_quantized_latent_representation_image = gr.Image(
+                label="New sound quantized latent representation", type="numpy", height=1200)
+
+    sound2sound_origin_upload_audio.change(receive_uopoad_origin_audio,
+                                           inputs=[sound2sound_duration_slider, sound2sound_origin_source_radio, sound2sound_origin_upload_audio,
+                                                   sound2sound_origin_microphone_audio, inpaintWithText_state],
+                                           outputs=[sound2sound_origin_spectrogram_upload_image,
+                                                    sound2sound_origin_phase_upload_image,
+                                                    sound2sound_origin_spectrogram_microphone_image,
+                                                    sound2sound_origin_phase_microphone_image,
+                                                    sound2sound_origin_upload_latent_representation_image,
+                                                    sound2sound_origin_upload_quantized_latent_representation_image,
+                                                    sound2sound_origin_microphone_latent_representation_image,
+                                                    sound2sound_origin_microphone_quantized_latent_representation_image,
+                                                    inpaintWithText_state])
+    sound2sound_origin_microphone_audio.change(receive_uopoad_origin_audio,
+                                               inputs=[sound2sound_duration_slider, sound2sound_origin_source_radio, sound2sound_origin_upload_audio,
+                                                       sound2sound_origin_microphone_audio, inpaintWithText_state],
+                                               outputs=[sound2sound_origin_spectrogram_upload_image,
+                                                        sound2sound_origin_phase_upload_image,
+                                                        sound2sound_origin_spectrogram_microphone_image,
+                                                        sound2sound_origin_phase_microphone_image,
+                                                        sound2sound_origin_upload_latent_representation_image,
+                                                        sound2sound_origin_upload_quantized_latent_representation_image,
+                                                        sound2sound_origin_microphone_latent_representation_image,
+                                                        sound2sound_origin_microphone_quantized_latent_representation_image,
+                                                        inpaintWithText_state])
+
+    sound2sound_sample_button.click(sound2sound_sample,
+                                    inputs=[sound2sound_origin_spectrogram_upload_image,
+                                            sound2sound_origin_spectrogram_microphone_image,
+                                            text2sound_prompts_textbox,
+                                            text2sound_negative_prompts_textbox,
+                                            sound2sound_batchsize_slider,
+                                            sound2sound_guidance_scale_slider,
+                                            sound2sound_sampler_radio,
+                                            sound2sound_sample_steps_slider,
+                                            sound2sound_origin_source_radio,
+                                            sound2sound_noising_strength_slider,
+                                            sound2sound_seed_textbox,
+                                            sound2sound_inpaint_area_radio,
+                                            inpaintWithText_state],
+                                    outputs=[sound2sound_new_sound_latent_representation_image,
+                                             sound2sound_new_sound_quantized_latent_representation_image,
+                                             sound2sound_new_sound_spectrogram_image,
+                                             sound2sound_new_sound_phase_image,
+                                             sound2sound_new_sound_audio,
+                                             sound2sound_sample_index_slider,
+                                             sound2sound_seed_textbox,
+                                             inpaintWithText_state])
+
+    sound2sound_sample_index_slider.change(show_sound2sound_sample,
+                                           inputs=[sound2sound_sample_index_slider, inpaintWithText_state],
+                                           outputs=[sound2sound_new_sound_latent_representation_image,
+                                                    sound2sound_new_sound_quantized_latent_representation_image,
+                                                    sound2sound_new_sound_spectrogram_image,
+                                                    sound2sound_new_sound_phase_image,
+                                                    sound2sound_new_sound_audio])
+
+    sound2sound_origin_source_radio.change(sound2sound_switch_origin_source,
+                                           inputs=[sound2sound_origin_source_radio],
+                                           outputs=[sound2sound_origin_upload_audio,
+                                                    sound2sound_origin_microphone_audio,
+                                                    sound2sound_origin_spectrogram_upload_image,
+                                                    sound2sound_origin_phase_upload_image,
+                                                    sound2sound_origin_spectrogram_microphone_image,
+                                                    sound2sound_origin_phase_microphone_image,
+                                                    sound2sound_origin_upload_latent_representation_image,
+                                                    sound2sound_origin_upload_quantized_latent_representation_image,
+                                                    sound2sound_origin_microphone_latent_representation_image,
+                                                    sound2sound_origin_microphone_quantized_latent_representation_image])

webUI/natural_language_guided/text2sound.py

@@ -0,0 +1,212 @@
+import gradio as gr
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+
+from model.DiffSynthSampler import DiffSynthSampler
+from tools import safe_int
+from webUI.natural_language_guided.utils import latent_representation_to_Gradio_image, \
+    encodeBatch2GradioOutput_STFT, add_instrument
+
+
+def get_text2sound_module(gradioWebUI, text2sound_state, virtual_instruments_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_guidance_scale, text2sound_sampler,
+                                text2sound_sample_steps, text2sound_seed,
+                                text2sound_dict):
+        text2sound_sample_steps = int(text2sound_sample_steps)
+        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)
+
+        mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy)
+        negative_condition = \
+            CLAP.get_text_features(**CLAP_tokenizer([text2sound_negative_prompts], padding=True, return_tensors="pt"))[
+                0]
+        mySampler.activate_classifier_free_guidance(CFG, negative_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.sample(model=uNet, shape=(text2sound_batchsize, channels, height, width), seed=text2sound_seed,
+                         return_tensor=True, condition=condition, sampler=text2sound_sampler)
+
+        latent_representations = latent_representations[-1]
+        print(latent_representations[0, 0, :3, :3])
+
+        latent_representation_gradio_images = []
+        quantized_latent_representation_gradio_images = []
+        new_sound_spectrogram_gradio_images = []
+        new_sound_phase_gradio_images = []
+        new_sound_rec_signals_gradio = []
+
+        quantized_latent_representations, loss, (_, _, _) = VAE_quantizer(latent_representations)
+        # Todo: remove hard-coding
+        flipped_log_spectrums, flipped_phases, rec_signals, _, _, _ = encodeBatch2GradioOutput_STFT(VAE_decoder,
+                                                                                                            quantized_latent_representations,
+                                                                                                            resolution=(
+                                                                                                                512,
+                                                                                                                width * VAE_scale),
+                                                                                                            original_STFT_batch=None
+                                                                                                     )
+
+        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_phase_gradio_images.append(flipped_phases[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_phase_gradio_images"] = new_sound_phase_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_phase_image: text2sound_dict["new_sound_phase_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_phase_image: text2sound_dict["new_sound_phase_gradio_images"][
+                    sample_index],
+                text2sound_sampled_audio: text2sound_dict["new_sound_rec_signals_gradio"][sample_index]}
+
+
+
+    def save_virtual_instrument(sample_index, virtual_instrument_name, text2sound_dict, virtual_instruments_dict):
+
+        virtual_instruments_dict = add_instrument(text2sound_dict, virtual_instruments_dict, virtual_instrument_name, sample_index)
+
+        return {virtual_instruments_state: virtual_instruments_dict,
+                text2sound_instrument_name_textbox: gr.Textbox(label="Instrument name", lines=1,
+                                                            placeholder=f"Saved as {virtual_instrument_name}!")}
+
+    with gr.Tab("Text2sound"):
+        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(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(variant="panel"):
+                with gr.Row(variant="panel"):
+                    text2sound_sampled_spectrogram_image = gr.Image(label="Sampled spectrogram", type="numpy", height=600)
+                    text2sound_sampled_phase_image = gr.Image(label="Sampled phase", type="numpy", height=600)
+                text2sound_sampled_audio = gr.Audio(type="numpy", label="Play")
+
+        with gr.Row(variant="panel"):
+            text2sound_instrument_name_textbox = gr.Textbox(label="Instrument name", lines=1,
+                                                            placeholder="Name of your instrument")
+            text2sound_save_instrument_button = gr.Button(variant="primary",
+                                                          value="Save instrument",
+                                                          scale=1)
+
+        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_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_phase_image,
+                                              text2sound_sampled_audio,
+                                              text2sound_seed_textbox,
+                                              text2sound_state,
+                                              text2sound_sample_index_slider])
+
+    text2sound_save_instrument_button.click(save_virtual_instrument,
+                                            inputs=[text2sound_sample_index_slider,
+                                                    text2sound_instrument_name_textbox,
+                                                    text2sound_state,
+                                                    virtual_instruments_state],
+                                            outputs=[virtual_instruments_state,
+                                                     text2sound_instrument_name_textbox])
+
+    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_phase_image,
+                                                   text2sound_sampled_audio])

webUI/natural_language_guided/track_maker.py

@@ -0,0 +1,192 @@
+import numpy as np
+import torch
+
+from model.DiffSynthSampler import DiffSynthSampler
+from webUI.natural_language_guided.utils import encodeBatch2GradioOutput_STFT
+import mido
+import pyrubberband as pyrb
+from tqdm import tqdm
+
+class NoteEvent:
+    def __init__(self, note, velocity, start_time, duration):
+        self.note = note
+        self.velocity = velocity
+        self.start_time = start_time  # In ticks
+        self.duration = duration      # In ticks
+
+    def __str__(self):
+        return f"Note {self.note}, velocity {self.velocity}, start_time {self.start_time}, duration {self.duration}"
+
+
+class Track:
+    def __init__(self, track, ticks_per_beat):
+        self.tempo_events = self._parse_tempo_events(track)
+        self.events = self._parse_note_events(track)
+        self.ticks_per_beat = ticks_per_beat
+
+    def _parse_tempo_events(self, track):
+        tempo_events = []
+        current_tempo = 500000  # Default MIDI tempo is 120 BPM which is 500000 microseconds per beat
+        for msg in track:
+            if msg.type == 'set_tempo':
+                tempo_events.append((msg.time, msg.tempo))
+            elif not msg.is_meta:
+                tempo_events.append((msg.time, current_tempo))
+        return tempo_events
+
+    def _parse_note_events(self, track):
+        events = []
+        start_time = 0
+        for msg in track:
+            if not msg.is_meta:
+                start_time += msg.time
+                if msg.type == 'note_on' and msg.velocity > 0:
+                    note_on_time = start_time
+                elif msg.type == 'note_on' and msg.velocity == 0:
+                    duration = start_time - note_on_time
+                    events.append(NoteEvent(msg.note, msg.velocity, note_on_time, duration))
+        return events
+
+    def synthesize_track(self, diffSynthSampler, sample_rate=16000):
+        track_audio = np.zeros(int(self._get_total_time() * sample_rate), dtype=np.float32)
+        current_tempo = 500000  # Start with default MIDI tempo 120 BPM
+        duration_note_mapping = {}
+
+        for event in tqdm(self.events[:50]):
+            current_tempo = self._get_tempo_at(event.start_time)
+            seconds_per_tick = mido.tick2second(1, self.ticks_per_beat, current_tempo)
+            start_time_sec = event.start_time * seconds_per_tick
+            # Todo: set a minimum duration
+            duration_sec = event.duration * seconds_per_tick
+            duration_sec = max(duration_sec, 0.75)
+            start_sample = int(start_time_sec * sample_rate)
+            if not (str(duration_sec) in duration_note_mapping):
+                note_sample = diffSynthSampler(event.velocity, duration_sec)
+                duration_note_mapping[str(duration_sec)] = note_sample / np.max(np.abs(note_sample))
+
+            note_audio = pyrb.pitch_shift(duration_note_mapping[str(duration_sec)], sample_rate, event.note - 52)
+            end_sample = start_sample + len(note_audio)
+            track_audio[start_sample:end_sample] += note_audio
+
+        return track_audio
+
+    def _get_tempo_at(self, time_tick):
+        current_tempo = 500000  # Start with default MIDI tempo 120 BPM
+        elapsed_ticks = 0
+
+        for tempo_change in self.tempo_events:
+            if elapsed_ticks + tempo_change[0] > time_tick:
+                return current_tempo
+            elapsed_ticks += tempo_change[0]
+            current_tempo = tempo_change[1]
+
+        return current_tempo
+
+    def _get_total_time(self):
+        total_time = 0
+        current_tempo = 500000  # Start with default MIDI tempo 120 BPM
+
+        for event in self.events:
+            current_tempo = self._get_tempo_at(event.start_time)
+            seconds_per_tick = mido.tick2second(1, self.ticks_per_beat, current_tempo)
+            total_time += event.duration * seconds_per_tick
+
+        return total_time
+
+
+class DiffSynth:
+    def __init__(self, instruments_configs, noise_prediction_model, VAE_quantizer, VAE_decoder, text_encoder, CLAP_tokenizer, device,
+                               model_sample_rate=16000, timesteps=1000, channels=4, freq_resolution=512, time_resolution=256, VAE_scale=4, squared=False):
+
+        self.noise_prediction_model = noise_prediction_model
+        self.VAE_quantizer = VAE_quantizer
+        self.VAE_decoder = VAE_decoder
+        self.device = device
+        self.model_sample_rate = model_sample_rate
+        self.timesteps = timesteps
+        self.channels = channels
+        self.freq_resolution = freq_resolution
+        self.time_resolution = time_resolution
+        self.height = int(freq_resolution/VAE_scale)
+        self.VAE_scale = VAE_scale
+        self.squared = squared
+        self.text_encoder = text_encoder
+        self.CLAP_tokenizer = CLAP_tokenizer
+
+        # instruments_configs 是字典 string -> (condition, negative_condition, guidance_scale, sample_steps, seed, initial_noise, sampler)
+        self.instruments_configs = instruments_configs
+        self.diffSynthSamplers = {}
+        self._update_instruments()
+
+
+    def _update_instruments(self):
+
+        def diffSynthSamplerWrapper(instruments_config):
+
+            def diffSynthSampler(velocity, duration_sec, sample_rate=16000):
+
+                condition = self.text_encoder.get_text_features(**self.CLAP_tokenizer([""], padding=True, return_tensors="pt")).to(self.device)
+                sample_steps = instruments_config['sample_steps']
+                sampler = instruments_config['sampler']
+                noising_strength = instruments_config['noising_strength']
+                latent_representation = instruments_config['latent_representation']
+                attack = instruments_config['attack']
+                before_release = instruments_config['before_release']
+
+                assert sample_rate == self.model_sample_rate, "sample_rate != model_sample_rate"
+
+                width = int(self.time_resolution * ((duration_sec + 1) / 4) / self.VAE_scale)
+
+                mySampler = DiffSynthSampler(self.timesteps, height=128, channels=4, noise_strategy="repeat", mute=True)
+                mySampler.respace(list(np.linspace(0, self.timesteps - 1, sample_steps, dtype=np.int32)))
+
+                # mask = 1, freeze
+                latent_mask = torch.zeros((1, 1, self.height, width), dtype=torch.float32).to(self.device)
+                latent_mask[:, :, :, :int(self.time_resolution * (attack / 4) / self.VAE_scale)] = 1.0
+                latent_mask[:, :, :, -int(self.time_resolution * ((before_release+1) / 4) / self.VAE_scale):] = 1.0
+
+                latent_representations, _ = \
+                    mySampler.inpaint_sample(model=self.noise_prediction_model, shape=(1, self.channels, self.height, width),
+                                            noising_strength=noising_strength, condition=condition,
+                                            guide_img=latent_representation, mask=latent_mask, return_tensor=True,
+                                            sampler=sampler,
+                                            use_dynamic_mask=True, end_noise_level_ratio=0.0,
+                                            mask_flexivity=1.0)
+
+
+                latent_representations = latent_representations[-1]
+
+                quantized_latent_representations, _, (_, _, _) = self.VAE_quantizer(latent_representations)
+                # Todo: remove hard-coding
+
+                flipped_log_spectrums, flipped_phases, rec_signals, _, _, _ = encodeBatch2GradioOutput_STFT(self.VAE_decoder,
+                                                                  quantized_latent_representations,
+                                                                  resolution=(
+                                                                      512,
+                                                                      width * self.VAE_scale),
+                                                                  original_STFT_batch=None,
+                                                            )
+
+
+                return rec_signals[0]
+
+            return diffSynthSampler
+
+        for key in self.instruments_configs.keys():
+            self.diffSynthSamplers[key] = diffSynthSamplerWrapper(self.instruments_configs[key])
+
+    def get_music(self, mid, instrument_names, sample_rate=16000):
+        tracks = [Track(t, mid.ticks_per_beat) for t in mid.tracks]
+        assert len(tracks) == len(instrument_names), f"len(tracks) = {len(tracks)} != {len(instrument_names)} = len(instrument_names)"
+
+        track_audios = [track.synthesize_track(self.diffSynthSamplers[instrument_names[i]], sample_rate=sample_rate) for i, track in enumerate(tracks)]
+
+        # 将所有音轨填充至最长音轨的长度，以便它们可以被叠加
+        max_length = max(len(audio) for audio in track_audios)
+        full_audio = np.zeros(max_length, dtype=np.float32)  # 初始化全音频数组为零
+        for audio in track_audios:
+            # 音轨可能不够长，需要填充零
+            padded_audio = np.pad(audio, (0, max_length - len(audio)), 'constant')
+            full_audio += padded_audio  # 叠加音轨
+
+        return full_audio

webUI/natural_language_guided/utils.py

@@ -0,0 +1,174 @@
+import librosa
+import numpy as np
+import torch
+
+from tools import np_power_to_db, decode_stft, depad_STFT
+
+
+def spectrogram_to_Gradio_image(spc):
+    ### input: spc [np.ndarray]
+    frequency_resolution, time_resolution = spc.shape[-2], spc.shape[-1]
+    spc = np.reshape(spc, (frequency_resolution, time_resolution))
+
+    # Todo:
+    magnitude_spectrum = np.abs(spc)
+    log_spectrum = np_power_to_db(magnitude_spectrum)
+    flipped_log_spectrum = np.flipud(log_spectrum)
+
+    colorful_spc = np.ones((frequency_resolution, time_resolution, 3)) * -80.0
+    colorful_spc[:, :, 0] = flipped_log_spectrum
+    colorful_spc[:, :, 1] = flipped_log_spectrum
+    colorful_spc[:, :, 2] = np.ones((frequency_resolution, time_resolution)) * -60.0
+    # Rescale to 0-255 and convert to uint8
+    rescaled = (colorful_spc + 80.0) / 80.0
+    rescaled = (255.0 * rescaled).astype(np.uint8)
+    return rescaled
+
+
+def phase_to_Gradio_image(phase):
+    ### input: spc [np.ndarray]
+    frequency_resolution, time_resolution = phase.shape[-2], phase.shape[-1]
+    phase = np.reshape(phase, (frequency_resolution, time_resolution))
+
+    # Todo:
+    flipped_phase = np.flipud(phase)
+    flipped_phase = (flipped_phase + 1.0) / 2.0
+
+    colorful_spc = np.zeros((frequency_resolution, time_resolution, 3))
+    colorful_spc[:, :, 0] = flipped_phase
+    colorful_spc[:, :, 1] = flipped_phase
+    colorful_spc[:, :, 2] = 0.2
+    # Rescale to 0-255 and convert to uint8
+    rescaled = (255.0 * colorful_spc).astype(np.uint8)
+    return rescaled
+
+
+def latent_representation_to_Gradio_image(latent_representation):
+    # input: latent_representation [torch.tensor]
+    if not isinstance(latent_representation, np.ndarray):
+        latent_representation = latent_representation.to("cpu").detach().numpy()
+    image = latent_representation
+
+    def normalize_image(img):
+        min_val = img.min()
+        max_val = img.max()
+        normalized_img = ((img - min_val) / (max_val - min_val) * 255)
+        return normalized_img
+
+    image[0, :, :] = normalize_image(image[0, :, :])
+    image[1, :, :] = normalize_image(image[1, :, :])
+    image[2, :, :] = normalize_image(image[2, :, :])
+    image[3, :, :] = normalize_image(image[3, :, :])
+    image_transposed = np.transpose(image, (1, 2, 0))
+    enlarged_image = np.repeat(image_transposed, 8, axis=0)
+    enlarged_image = np.repeat(enlarged_image, 8, axis=1)
+    return np.flipud(enlarged_image).astype(np.uint8)
+
+
+def InputBatch2Encode_STFT(encoder, STFT_batch, resolution=(512, 256), quantizer=None, squared=True):
+    """Transform batch of numpy spectrogram's into signals and encodings."""
+    # Todo: remove resolution hard-coding
+    frequency_resolution, time_resolution = resolution
+
+    device = next(encoder.parameters()).device
+    if not (quantizer is None):
+        latent_representation_batch = encoder(STFT_batch.to(device))
+        quantized_latent_representation_batch, loss, (_, _, _) = quantizer(latent_representation_batch)
+    else:
+        mu, logvar, latent_representation_batch = encoder(STFT_batch.to(device))
+        quantized_latent_representation_batch = None
+
+    STFT_batch = STFT_batch.to("cpu").detach().numpy()
+
+    origin_flipped_log_spectrums, origin_flipped_phases, origin_signals = [], [], []
+    for STFT in STFT_batch:
+
+        padded_D_rec = decode_stft(STFT)
+        D_rec = depad_STFT(padded_D_rec)
+        spc = np.abs(D_rec)
+        phase = np.angle(D_rec)
+
+        flipped_log_spectrum = spectrogram_to_Gradio_image(spc)
+        flipped_phase = phase_to_Gradio_image(phase)
+
+        # get_audio
+        rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024)
+
+        origin_flipped_log_spectrums.append(flipped_log_spectrum)
+        origin_flipped_phases.append(flipped_phase)
+        origin_signals.append(rec_signal)
+
+    return origin_flipped_log_spectrums, origin_flipped_phases, origin_signals, \
+        latent_representation_batch, quantized_latent_representation_batch
+
+
+def encodeBatch2GradioOutput_STFT(decoder, latent_vector_batch, resolution=(512, 256), original_STFT_batch=None):
+    """Show a spectrogram."""
+    # Todo: remove resolution hard-coding
+    frequency_resolution, time_resolution = resolution
+
+    if isinstance(latent_vector_batch, np.ndarray):
+        latent_vector_batch = torch.from_numpy(latent_vector_batch).to(next(decoder.parameters()).device)
+
+    reconstruction_batch = decoder(latent_vector_batch).to("cpu").detach().numpy()
+
+    flipped_log_spectrums, flipped_phases, rec_signals = [], [], []
+    flipped_log_spectrums_with_original_amp, flipped_phases_with_original_amp, rec_signals_with_original_amp = [], [], []
+
+    for index, STFT in enumerate(reconstruction_batch):
+        padded_D_rec = decode_stft(STFT)
+        D_rec = depad_STFT(padded_D_rec)
+        spc = np.abs(D_rec)
+        phase = np.angle(D_rec)
+
+        flipped_log_spectrum = spectrogram_to_Gradio_image(spc)
+        flipped_phase = phase_to_Gradio_image(phase)
+
+        # get_audio
+        rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024)
+
+        flipped_log_spectrums.append(flipped_log_spectrum)
+        flipped_phases.append(flipped_phase)
+        rec_signals.append(rec_signal)
+
+        ##########################################
+
+        if original_STFT_batch is not None:
+            STFT[0, :, :] = original_STFT_batch[index, 0, :, :]
+
+            padded_D_rec = decode_stft(STFT)
+            D_rec = depad_STFT(padded_D_rec)
+            spc = np.abs(D_rec)
+            phase = np.angle(D_rec)
+
+            flipped_log_spectrum = spectrogram_to_Gradio_image(spc)
+            flipped_phase = phase_to_Gradio_image(phase)
+
+            # get_audio
+            rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024)
+
+            flipped_log_spectrums_with_original_amp.append(flipped_log_spectrum)
+            flipped_phases_with_original_amp.append(flipped_phase)
+            rec_signals_with_original_amp.append(rec_signal)
+
+
+    return flipped_log_spectrums, flipped_phases, rec_signals, \
+        flipped_log_spectrums_with_original_amp, flipped_phases_with_original_amp, rec_signals_with_original_amp
+
+
+
+def add_instrument(source_dict, virtual_instruments_dict, virtual_instrument_name, sample_index):
+
+    virtual_instruments = virtual_instruments_dict["virtual_instruments"]
+    virtual_instrument = {
+                          "latent_representation": source_dict["latent_representations"][sample_index],
+                          "quantized_latent_representation": source_dict["quantized_latent_representations"][sample_index],
+                          "sampler": source_dict["sampler"],
+                          "signal": source_dict["new_sound_rec_signals_gradio"][sample_index],
+                          "spectrogram_gradio_image": source_dict["new_sound_spectrogram_gradio_images"][
+                              sample_index],
+                          "phase_gradio_image": source_dict["new_sound_phase_gradio_images"][
+                              sample_index]}
+    virtual_instruments[virtual_instrument_name] = virtual_instrument
+    virtual_instruments_dict["virtual_instruments"] = virtual_instruments
+    return virtual_instruments_dict