Delete src/test.py

src/test.py

@@ -1,97 +0,0 @@
-import random
-import argparse
-import os
-import time
-import soundfile as sf
-import numpy as np
-import pandas as pd
-from tqdm import tqdm
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from diffusers import DDIMScheduler
-from models.conditioners import MaskDiT
-from modules.autoencoder_wrapper import Autoencoder
-from transformers import T5Tokenizer, T5EncoderModel
-from inference import inference
-from utils import scale_shift, get_lr_scheduler, compute_snr, load_yaml_with_includes
-
-
-parser = argparse.ArgumentParser()
-# config settings
-parser.add_argument('--config-name', type=str, default='configs/udit_ada.yml')
-parser.add_argument('--ckpt-path', type=str, default='../ckpts/')
-parser.add_argument('--ckpt-id', type=str, default='120')
-parser.add_argument('--save_path', type=str, default='../output/')
-parser.add_argument('--test-df', type=str, default='audiocaps_test.csv')
-# parser.add_argument('--test-split', type=str, default='test')
-
-parser.add_argument('--device', type=str, default='cuda')
-parser.add_argument('--guidance-scale', type=float, default=3)
-parser.add_argument('--guidance-rescale', type=float, default=0)
-parser.add_argument('--ddim-steps', type=int, default=50)
-parser.add_argument('--eta', type=float, default=1)
-parser.add_argument('--random-seed', type=int, default=None)
-
-args = parser.parse_args()
-params = load_yaml_with_includes(args.config_name)
-
-# args.ckpt_path = f"{args.ckpt_path}/{params['model_name']}/{args.ckpt_id}.pt"
-args.save_path = f"{args.save_path}/{params['model_name']}/{args.ckpt_id}_{args.ddim_steps}_{args.guidance_scale}_{args.guidance_rescale}/"
-args.ckpt_path = f"{args.ckpt_path}/{args.ckpt_id}.pt"
-
-if __name__ == '__main__':
-    # Codec Model
-    autoencoder = Autoencoder(ckpt_path=params['autoencoder']['path'],
-                              model_type=params['autoencoder']['name'],
-                              quantization_first=params['autoencoder']['q_first'])
-    autoencoder.to(args.device)
-    autoencoder.eval()
-
-    # text encoder
-    tokenizer = T5Tokenizer.from_pretrained(params['text_encoder']['model'])
-    text_encoder = T5EncoderModel.from_pretrained(params['text_encoder']['model'],
-                                                  device_map='cpu').to(args.device)
-    text_encoder.eval()
-
-    # main U-Net
-    unet = MaskDiT(**params['model']).to(args.device)
-    unet.eval()
-    unet.load_state_dict(torch.load(args.ckpt_path)['model'])
-
-    total_params = sum([param.nelement() for param in unet.parameters()])
-    print("Number of parameter: %.2fM" % (total_params / 1e6))
-
-    noise_scheduler = DDIMScheduler(**params['diff'])
-    # these steps reset dtype of noise_scheduler params
-    latents = torch.randn((1, 128, 128), device=args.device)
-    noise = torch.randn_like(latents)
-    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (1,), device=args.device)
-    _ = noise_scheduler.add_noise(latents, noise, timesteps)
-
-    df = pd.read_csv(args.test_df)
-    # Wdf = df[df['split'] == args.test_split]
-    df = df[df['audio_length'] != 0]
-    # df = df.sample(10)
-    os.makedirs(args.save_path, exist_ok=True)
-    audio_frames = params['data']['train_frames']
-
-    for i in tqdm(range(len(df))):
-        row = df.iloc[i]
-        text = row['caption']
-        audio_id = row['audiocap_id']
-
-        pred = inference(autoencoder, unet, None, None,
-                         tokenizer, text_encoder, 
-                         params, noise_scheduler,
-                         text, None,
-                         audio_frames,
-                         args.guidance_scale, args.guidance_rescale,
-                         args.ddim_steps, args.eta, args.random_seed,
-                         args.device)
-        pred = pred.cpu().numpy().squeeze(0).squeeze(0)
-
-        sf.write(f"{args.save_path}/{audio_id}.wav",
-                 pred, samplerate=params['data']['sr'])