aaa / wavvae3(1).py

Upload 2 files

19be270 verified over 1 year ago

12.9 kB

import argparse
import filecmp
import multiprocessing
import os
import subprocess
import librosa
from functools import partial
from multiprocessing import Pool, Process

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.optim import AdamW

from modules.vocoder.commons.stft_loss import MultiResolutionSTFTLoss
from modules.vocoder.hifigan.hifigan import MultiPeriodDiscriminator, MultiScaleDiscriminator, \
    generator_loss, feature_loss, discriminator_loss
from modules.vocoder.hifigan.mel_utils import mel_spectrogram
from modules.vocoder.univnet.mrd import MultiResolutionDiscriminator
from modules.tts.wavvae.decoder.wavvae_v3 import WavVAE_V3
from tasks.tts.utils.audio import torch_wav2spec
from tasks.tts.utils.audio.align import mel2token_to_dur
from utils.commons.ckpt_utils import load_ckpt
from utils.commons.hparams import hparams

from attrdict import AttrDict
from tasks.tts.dataset_mixin import TTSDatasetMixin
from utils.commons.base_task import BaseTask
from utils.commons.import_utils import import_module_bystr
from utils.nn.schedulers import WarmupSchedule, CosineSchedule


class WavVAETask(TTSDatasetMixin, BaseTask):
    def __init__(self):
        super().__init__()
        self.dataset_cls = import_module_bystr(hparams['dataset_cls'])
        self.val_dataset_cls = import_module_bystr(hparams['val_dataset_cls'])
        self.processer_fn = import_module_bystr(hparams['processer_fn'])
        self.build_fast_dataloader = import_module_bystr(hparams['build_fast_dataloader'])
        self.hparams = hparams
        self.config = AttrDict(hparams)

        # Online load mel with GPU
        sample_rate = hparams["audio_sample_rate"]
        fft_size = hparams["win_size"]
        win_size = hparams["win_size"]
        hop_size = hparams["hop_size"]
        num_mels = hparams["audio_num_mel_bins"]
        fmin = hparams["fmin"]
        fmax = hparams["fmax"]
        mel_basis = librosa.filters.mel(
            sr=sample_rate, n_fft=fft_size, n_mels=num_mels, fmin=fmin, fmax=fmax
        )
        self.torch_wav2spec_ = partial(
            torch_wav2spec, mel_basis=mel_basis, fft_size=fft_size, hop_size=hop_size, win_length=win_size,
        )

    def build_model(self):
        self.model_gen = WavVAE_V3(hparams=hparams)

        self.model_disc = torch.nn.ModuleDict()
        self.model_disc['mpd'] = MultiPeriodDiscriminator(hparams['mpd'], use_cond=hparams['use_cond_disc'])
        self.model_disc['msd'] = MultiScaleDiscriminator(use_cond=hparams['use_cond_disc'])
        if hparams['use_mrd']:
            self.model_disc['mrd'] = MultiResolutionDiscriminator(hparams)
        self.stft_loss = MultiResolutionSTFTLoss()

        # load_ckpt(self.model_gen.encoder, './checkpoints/1231_megatts3_wavvae_v2_25hz', 'model.module.encoder', strict=False)
        # load_ckpt(self.model_gen.decoder, './checkpoints/1117_melgan-nsf_full_1', 'model_gen', force=True, strict=True)
        # load_ckpt(self.model_disc, './checkpoints/1117_melgan-nsf_full_1', 'model_disc', force=True, strict=True)
        return {'trainable': [self.model_gen, self.model_disc['mpd'], self.model_disc['msd'], self.model_disc['mrd']], 'others': []}

    def load_model(self):
        if hparams.get('load_ckpt', '') != '':
            load_ckpt(self.model, hparams['load_ckpt'], 'model', strict=False)

    def build_optimizer(self):
        optimizer_gen = torch.optim.AdamW(self.model_gen.parameters(), lr=hparams['lr'],
                                          betas=[hparams['adam_b1'], hparams['adam_b2']])
        optimizer_disc = torch.optim.AdamW(self.model_disc.parameters(),
                                           lr=hparams.get('disc_lr', hparams['lr']),
                                           betas=[hparams['adam_b1'], hparams['adam_b2']])
        return [optimizer_gen, optimizer_disc]

    def build_scheduler(self, optimizer):
        return None

    def _training_step(self, sample, batch_idx, optimizer_idx):
        log_outputs = {}
        loss_weights = {}
        sample['wavs'] = sample['wavs'].float()
        # return None, {}

        if self.global_step % 100 == 0:
            devices = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(",")
            for d in devices:
                os.system(f'pkill -f "voidgpu{d}"')

        y = sample['wavs']
        loss_output = {}
        if optimizer_idx == 0:
            #######################
            #      Generator      #
            #######################
            y_, posterior = self.model_gen(y)
            y = y.unsqueeze(1)
            y_mel = mel_spectrogram(y.squeeze(1), hparams).transpose(1, 2)
            y_hat_mel = mel_spectrogram(y_.squeeze(1), hparams).transpose(1, 2)
            loss_output['mel'] = F.l1_loss(y_hat_mel, y_mel) * hparams['lambda_mel']
            if self.training:
                _, y_p_hat_g, fmap_f_r, fmap_f_g = self.model_disc['mpd'](y, y_, None)
                _, y_s_hat_g, fmap_s_r, fmap_s_g = self.model_disc['msd'](y, y_, None)
                loss_output['a_p'] = generator_loss(y_p_hat_g) * hparams['lambda_adv'] * hparams.get('lambda_mpd', 1.0)
                loss_output['a_s'] = generator_loss(y_s_hat_g) * hparams['lambda_adv'] * hparams.get('lambda_msd', 1.0)
                if hparams['use_mrd']:
                    y_r_hat_g = [x[1] for x in self.model_disc['mrd'](y_)]
                    loss_output['a_r'] = generator_loss(y_r_hat_g) \
                                         * hparams['lambda_adv'] * hparams.get('lambda_mrd', 1.0)
                if hparams['use_ms_stft']:
                    loss_output['sc'], loss_output['mag'] = self.stft_loss(y.squeeze(1), y_.squeeze(1))
                loss_output['kl_loss'] = posterior.kl().mean() * hparams.get('lambda_kl', 1.0)
            self.y_ = y_.detach()
        else:
            #######################
            #    Discriminator    #
            #######################
            if not self.training:
                return None
            y = y.unsqueeze(1)
            y_ = self.y_
            # MPD
            y_p_hat_r, y_p_hat_g, _, _ = self.model_disc['mpd'](y, y_.detach(), None)
            loss_output['r_p'], loss_output['f_p'] = discriminator_loss(y_p_hat_r, y_p_hat_g)
            # MSD
            y_s_hat_r, y_s_hat_g, _, _ = self.model_disc['msd'](y, y_.detach(), None)
            loss_output['r_s'], loss_output['f_s'] = discriminator_loss(y_s_hat_r, y_s_hat_g)
            # MRD
            if hparams['use_mrd']:
                y_r_hat_r = [x[1] for x in self.model_disc['mrd'](y)]
                y_r_hat_g = [x[1] for x in self.model_disc['mrd'](y_.detach())]
                loss_output['r_r'], loss_output['f_r'] = discriminator_loss(y_r_hat_r, y_r_hat_g)
        total_loss = sum(loss_output.values())
        loss_output['bs'] = sample['wavs'].shape[0]
        return total_loss, loss_output

    def save_valid_result(self, sample, batch_idx, model_out):
        sr = hparams['audio_sample_rate']
        mel_out = model_out.get('mel_out')
        f0 = sample.get('f0')
        f0_gt = sample.get('f0')
        if f0 is not None:
            f0_gt = f0_gt.cpu()[-1]
        if mel_out is not None:
            f0_pred = self.predict_f0(sample['mels'])
            self.plot_mel(batch_idx, sample['mels'], mel_out, f0s={'f0': f0_pred, 'f0g': f0_gt})
        # gt wav
        if self.global_step <= hparams['valid_infer_interval']:
            mel_gt = sample['mels'][-1].cpu()
            f0 = self.predict_f0(sample['mels'][-1:])
            wav_gt = self.vocoder.spec2wav(mel_gt, f0=f0)
            self.logger.add_audio(f'wav_gt_{batch_idx}', wav_gt, self.global_step, sr)

        if self.global_step >= 0:
            # with gt duration
            model_out = self.run_model(sample, infer=True, infer_use_gt_dur=True)
            # dur_info = self.get_plot_dur_info(sample, model_out)
            # del dur_info['dur_pred']
            dur_info = None

            f0 = self.predict_f0(model_out['mel_out'])
            wav_pred = self.vocoder.spec2wav(model_out['mel_out'][-1].cpu(), f0=f0)
            self.logger.add_audio(f'wav_gdur_{batch_idx}', wav_pred, self.global_step, sr)
            self.plot_mel(batch_idx, sample['mels'][-1:], model_out['mel_out'][-1], f'mel_gdur_{batch_idx}',
                          dur_info=dur_info, f0s={'f0': f0, 'f0g': f0_gt})

            # with pred duration
            if not hparams['use_gt_dur'] and not hparams['use_gt_latent']:
                model_out = self.run_model(sample, infer=True, infer_use_gt_dur=False)
                # dur_info = self.get_plot_dur_info(sample, model_out)
                dur_info = None
                f0 = self.predict_f0(model_out['mel_out'])
                self.plot_mel(
                    batch_idx, sample['mels'], model_out['mel_out'][-1], f'mel_pdur_{batch_idx}',
                    dur_info=dur_info, f0s={'f0': f0, 'f0g': f0_gt})
                wav_pred = self.vocoder.spec2wav(model_out['mel_out'][-1].cpu(), f0=f0)
                self.logger.add_audio(f'wav_pdur_{batch_idx}', wav_pred, self.global_step, sr)

    def get_plot_dur_info(self, sample, model_out):
        T_txt = sample['txt_tokens'].shape[1]
        dur_gt = mel2token_to_dur(sample['mel2ph'], T_txt)[-1]
        dur_pred = model_out['dur'] if 'dur' in model_out else dur_gt
        txt = self.token_encoder.decode(sample['txt_tokens'][-1].cpu().numpy())
        txt = txt.split(" ")
        return {'dur_gt': dur_gt, 'dur_pred': dur_pred, 'txt': txt}

    def on_before_optimization(self, opt_idx):
        if opt_idx == 0:
            nn.utils.clip_grad_norm_(self.model_gen.parameters(), hparams['generator_grad_norm'])
        else:
            nn.utils.clip_grad_norm_(self.model_disc.parameters(), hparams["discriminator_grad_norm"])

    def to(self, device=None, dtype=None):
        super().to(device=device, dtype=dtype)
        # trainer doesn't move ema to device automatically, we do it mannually
        if hparams.get('use_ema', False):
            self.ema.to(device=device, dtype=dtype)

    def cuda(self,device):
        super().cuda(device)
        if hparams.get('use_ema', False):
            self.ema.to(device=device)
    
    @torch.no_grad()
    def validation_step(self, sample, batch_idx):
        infer_steps = self.hparams.get('infer_steps', 12)
        outputs = self._validation_step(sample, batch_idx, infer_steps)
        return outputs

    def _validation_step(self, sample, batch_idx, infer_steps):
        outputs = {}
        if self.trainer.proc_rank == 0:
            # self.vae.eval()
            # with torch.inference_mode():
            #     with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
            #         lat = self.vae.get_latent(sample["mels"])
            #         lat_lens = latent_lengths.clamp(max=lat.size(1))
            # mel = self.vae.decode(lat)
            pass
            # outputs['losses'], _ = self.run_model(sample)
            # _, model_out = self.run_model(sample, infer=True, infer_steps=infer_steps)
            # outputs = tensors_to_scalars(outputs)
            # output_ldm = model_out['ldm_out']
            # T = output_ldm.shape[1]
            # ldm = sample['kps'][:, :T]  # [B, T, nkp, kp_dim] [0, 1]
            # B, T, nkp, kp_dim = ldm.shape
            # output_ldm = self.denormalize_ldm(output_ldm)
            # recon_ldm = model_out['recon_ldm']
            # recon_ldm = self.denormalize_ldm(recon_ldm)

            # results_dir = f"{hparams['work_dir']}/results/{self.global_step}_infersteps{infer_steps}_cfg{hparams['cfg_w']}"
            # os.makedirs(results_dir, exist_ok=True)
            # n_ctx = model_out['ctx_mask'][0, :, 0].sum().long().item()
            # writer_kp = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_kp.sil.mp4", fps=25)
            # writer_gt = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_gt.sil.mp4", fps=25)
            # writer_pred = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_pred.sil.mp4", fps=25)
            # for i in range(T):
            #     img = self.draw_ldm(recon_ldm[0, i])
            #     writer_gt.append_data(img)
            #     img = self.draw_ldm(ldm[0, i])
            #     writer_kp.append_data(img)
            #     if i < n_ctx:
            #         writer_pred.append_data(img)
            #     else:
            #         img = self.draw_ldm(
            #             output_ldm[0, i], color=(255, 255, 0),
            #         )
            #         writer_pred.append_data(img)
            # writer_gt.close()
            # writer_kp.close()
            # writer_pred.close()
        return outputs

    @torch.no_grad()
    def test_step(self, sample, batch_idx):
        infer_steps = hparams['infer_steps']
        return self._validation_step(sample, batch_idx, infer_steps)

	import argparse
	import filecmp
	import multiprocessing
	import os
	import subprocess
	import librosa
	from functools import partial
	from multiprocessing import Pool, Process

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import torch.distributed as dist
	from torch.optim import AdamW

	from modules.vocoder.commons.stft_loss import MultiResolutionSTFTLoss
	from modules.vocoder.hifigan.hifigan import MultiPeriodDiscriminator, MultiScaleDiscriminator, \
	generator_loss, feature_loss, discriminator_loss
	from modules.vocoder.hifigan.mel_utils import mel_spectrogram
	from modules.vocoder.univnet.mrd import MultiResolutionDiscriminator
	from modules.tts.wavvae.decoder.wavvae_v3 import WavVAE_V3
	from tasks.tts.utils.audio import torch_wav2spec
	from tasks.tts.utils.audio.align import mel2token_to_dur
	from utils.commons.ckpt_utils import load_ckpt
	from utils.commons.hparams import hparams

	from attrdict import AttrDict
	from tasks.tts.dataset_mixin import TTSDatasetMixin
	from utils.commons.base_task import BaseTask
	from utils.commons.import_utils import import_module_bystr
	from utils.nn.schedulers import WarmupSchedule, CosineSchedule


	class WavVAETask(TTSDatasetMixin, BaseTask):
	def __init__(self):
	super().__init__()
	self.dataset_cls = import_module_bystr(hparams['dataset_cls'])
	self.val_dataset_cls = import_module_bystr(hparams['val_dataset_cls'])
	self.processer_fn = import_module_bystr(hparams['processer_fn'])
	self.build_fast_dataloader = import_module_bystr(hparams['build_fast_dataloader'])
	self.hparams = hparams
	self.config = AttrDict(hparams)

	# Online load mel with GPU
	sample_rate = hparams["audio_sample_rate"]
	fft_size = hparams["win_size"]
	win_size = hparams["win_size"]
	hop_size = hparams["hop_size"]
	num_mels = hparams["audio_num_mel_bins"]
	fmin = hparams["fmin"]
	fmax = hparams["fmax"]
	mel_basis = librosa.filters.mel(
	sr=sample_rate, n_fft=fft_size, n_mels=num_mels, fmin=fmin, fmax=fmax
	)
	self.torch_wav2spec_ = partial(
	torch_wav2spec, mel_basis=mel_basis, fft_size=fft_size, hop_size=hop_size, win_length=win_size,
	)

	def build_model(self):
	self.model_gen = WavVAE_V3(hparams=hparams)

	self.model_disc = torch.nn.ModuleDict()
	self.model_disc['mpd'] = MultiPeriodDiscriminator(hparams['mpd'], use_cond=hparams['use_cond_disc'])
	self.model_disc['msd'] = MultiScaleDiscriminator(use_cond=hparams['use_cond_disc'])
	if hparams['use_mrd']:
	self.model_disc['mrd'] = MultiResolutionDiscriminator(hparams)
	self.stft_loss = MultiResolutionSTFTLoss()

	# load_ckpt(self.model_gen.encoder, './checkpoints/1231_megatts3_wavvae_v2_25hz', 'model.module.encoder', strict=False)
	# load_ckpt(self.model_gen.decoder, './checkpoints/1117_melgan-nsf_full_1', 'model_gen', force=True, strict=True)
	# load_ckpt(self.model_disc, './checkpoints/1117_melgan-nsf_full_1', 'model_disc', force=True, strict=True)
	return {'trainable': [self.model_gen, self.model_disc['mpd'], self.model_disc['msd'], self.model_disc['mrd']], 'others': []}

	def load_model(self):
	if hparams.get('load_ckpt', '') != '':
	load_ckpt(self.model, hparams['load_ckpt'], 'model', strict=False)

	def build_optimizer(self):
	optimizer_gen = torch.optim.AdamW(self.model_gen.parameters(), lr=hparams['lr'],
	betas=[hparams['adam_b1'], hparams['adam_b2']])
	optimizer_disc = torch.optim.AdamW(self.model_disc.parameters(),
	lr=hparams.get('disc_lr', hparams['lr']),
	betas=[hparams['adam_b1'], hparams['adam_b2']])
	return [optimizer_gen, optimizer_disc]

	def build_scheduler(self, optimizer):
	return None

	def _training_step(self, sample, batch_idx, optimizer_idx):
	log_outputs = {}
	loss_weights = {}
	sample['wavs'] = sample['wavs'].float()
	# return None, {}

	if self.global_step % 100 == 0:
	devices = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(",")
	for d in devices:
	os.system(f'pkill -f "voidgpu{d}"')

	y = sample['wavs']
	loss_output = {}
	if optimizer_idx == 0:
	#######################
	# Generator #
	#######################
	y_, posterior = self.model_gen(y)
	y = y.unsqueeze(1)
	y_mel = mel_spectrogram(y.squeeze(1), hparams).transpose(1, 2)
	y_hat_mel = mel_spectrogram(y_.squeeze(1), hparams).transpose(1, 2)
	loss_output['mel'] = F.l1_loss(y_hat_mel, y_mel) * hparams['lambda_mel']
	if self.training:
	_, y_p_hat_g, fmap_f_r, fmap_f_g = self.model_disc['mpd'](y, y_, None)
	_, y_s_hat_g, fmap_s_r, fmap_s_g = self.model_disc['msd'](y, y_, None)
	loss_output['a_p'] = generator_loss(y_p_hat_g) * hparams['lambda_adv'] * hparams.get('lambda_mpd', 1.0)
	loss_output['a_s'] = generator_loss(y_s_hat_g) * hparams['lambda_adv'] * hparams.get('lambda_msd', 1.0)
	if hparams['use_mrd']:
	y_r_hat_g = [x[1] for x in self.model_disc['mrd'](y_)]
	loss_output['a_r'] = generator_loss(y_r_hat_g) \
	* hparams['lambda_adv'] * hparams.get('lambda_mrd', 1.0)
	if hparams['use_ms_stft']:
	loss_output['sc'], loss_output['mag'] = self.stft_loss(y.squeeze(1), y_.squeeze(1))
	loss_output['kl_loss'] = posterior.kl().mean() * hparams.get('lambda_kl', 1.0)
	self.y_ = y_.detach()
	else:
	#######################
	# Discriminator #
	#######################
	if not self.training:
	return None
	y = y.unsqueeze(1)
	y_ = self.y_
	# MPD
	y_p_hat_r, y_p_hat_g, _, _ = self.model_disc['mpd'](y, y_.detach(), None)
	loss_output['r_p'], loss_output['f_p'] = discriminator_loss(y_p_hat_r, y_p_hat_g)
	# MSD
	y_s_hat_r, y_s_hat_g, _, _ = self.model_disc['msd'](y, y_.detach(), None)
	loss_output['r_s'], loss_output['f_s'] = discriminator_loss(y_s_hat_r, y_s_hat_g)
	# MRD
	if hparams['use_mrd']:
	y_r_hat_r = [x[1] for x in self.model_disc['mrd'](y)]
	y_r_hat_g = [x[1] for x in self.model_disc['mrd'](y_.detach())]
	loss_output['r_r'], loss_output['f_r'] = discriminator_loss(y_r_hat_r, y_r_hat_g)
	total_loss = sum(loss_output.values())
	loss_output['bs'] = sample['wavs'].shape[0]
	return total_loss, loss_output

	def save_valid_result(self, sample, batch_idx, model_out):
	sr = hparams['audio_sample_rate']
	mel_out = model_out.get('mel_out')
	f0 = sample.get('f0')
	f0_gt = sample.get('f0')
	if f0 is not None:
	f0_gt = f0_gt.cpu()[-1]
	if mel_out is not None:
	f0_pred = self.predict_f0(sample['mels'])
	self.plot_mel(batch_idx, sample['mels'], mel_out, f0s={'f0': f0_pred, 'f0g': f0_gt})
	# gt wav
	if self.global_step <= hparams['valid_infer_interval']:
	mel_gt = sample['mels'][-1].cpu()
	f0 = self.predict_f0(sample['mels'][-1:])
	wav_gt = self.vocoder.spec2wav(mel_gt, f0=f0)
	self.logger.add_audio(f'wav_gt_{batch_idx}', wav_gt, self.global_step, sr)

	if self.global_step >= 0:
	# with gt duration
	model_out = self.run_model(sample, infer=True, infer_use_gt_dur=True)
	# dur_info = self.get_plot_dur_info(sample, model_out)
	# del dur_info['dur_pred']
	dur_info = None

	f0 = self.predict_f0(model_out['mel_out'])
	wav_pred = self.vocoder.spec2wav(model_out['mel_out'][-1].cpu(), f0=f0)
	self.logger.add_audio(f'wav_gdur_{batch_idx}', wav_pred, self.global_step, sr)
	self.plot_mel(batch_idx, sample['mels'][-1:], model_out['mel_out'][-1], f'mel_gdur_{batch_idx}',
	dur_info=dur_info, f0s={'f0': f0, 'f0g': f0_gt})

	# with pred duration
	if not hparams['use_gt_dur'] and not hparams['use_gt_latent']:
	model_out = self.run_model(sample, infer=True, infer_use_gt_dur=False)
	# dur_info = self.get_plot_dur_info(sample, model_out)
	dur_info = None
	f0 = self.predict_f0(model_out['mel_out'])
	self.plot_mel(
	batch_idx, sample['mels'], model_out['mel_out'][-1], f'mel_pdur_{batch_idx}',
	dur_info=dur_info, f0s={'f0': f0, 'f0g': f0_gt})
	wav_pred = self.vocoder.spec2wav(model_out['mel_out'][-1].cpu(), f0=f0)
	self.logger.add_audio(f'wav_pdur_{batch_idx}', wav_pred, self.global_step, sr)

	def get_plot_dur_info(self, sample, model_out):
	T_txt = sample['txt_tokens'].shape[1]
	dur_gt = mel2token_to_dur(sample['mel2ph'], T_txt)[-1]
	dur_pred = model_out['dur'] if 'dur' in model_out else dur_gt
	txt = self.token_encoder.decode(sample['txt_tokens'][-1].cpu().numpy())
	txt = txt.split(" ")
	return {'dur_gt': dur_gt, 'dur_pred': dur_pred, 'txt': txt}

	def on_before_optimization(self, opt_idx):
	if opt_idx == 0:
	nn.utils.clip_grad_norm_(self.model_gen.parameters(), hparams['generator_grad_norm'])
	else:
	nn.utils.clip_grad_norm_(self.model_disc.parameters(), hparams["discriminator_grad_norm"])

	def to(self, device=None, dtype=None):
	super().to(device=device, dtype=dtype)
	# trainer doesn't move ema to device automatically, we do it mannually
	if hparams.get('use_ema', False):
	self.ema.to(device=device, dtype=dtype)

	def cuda(self,device):
	super().cuda(device)
	if hparams.get('use_ema', False):
	self.ema.to(device=device)

	@torch.no_grad()
	def validation_step(self, sample, batch_idx):
	infer_steps = self.hparams.get('infer_steps', 12)
	outputs = self._validation_step(sample, batch_idx, infer_steps)
	return outputs

	def _validation_step(self, sample, batch_idx, infer_steps):
	outputs = {}
	if self.trainer.proc_rank == 0:
	# self.vae.eval()
	# with torch.inference_mode():
	# with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
	# lat = self.vae.get_latent(sample["mels"])
	# lat_lens = latent_lengths.clamp(max=lat.size(1))
	# mel = self.vae.decode(lat)
	pass
	# outputs['losses'], _ = self.run_model(sample)
	# _, model_out = self.run_model(sample, infer=True, infer_steps=infer_steps)
	# outputs = tensors_to_scalars(outputs)
	# output_ldm = model_out['ldm_out']
	# T = output_ldm.shape[1]
	# ldm = sample['kps'][:, :T] # [B, T, nkp, kp_dim] [0, 1]
	# B, T, nkp, kp_dim = ldm.shape
	# output_ldm = self.denormalize_ldm(output_ldm)
	# recon_ldm = model_out['recon_ldm']
	# recon_ldm = self.denormalize_ldm(recon_ldm)

	# results_dir = f"{hparams['work_dir']}/results/{self.global_step}_infersteps{infer_steps}_cfg{hparams['cfg_w']}"
	# os.makedirs(results_dir, exist_ok=True)
	# n_ctx = model_out['ctx_mask'][0, :, 0].sum().long().item()
	# writer_kp = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_kp.sil.mp4", fps=25)
	# writer_gt = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_gt.sil.mp4", fps=25)
	# writer_pred = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_pred.sil.mp4", fps=25)
	# for i in range(T):
	# img = self.draw_ldm(recon_ldm[0, i])
	# writer_gt.append_data(img)
	# img = self.draw_ldm(ldm[0, i])
	# writer_kp.append_data(img)
	# if i < n_ctx:
	# writer_pred.append_data(img)
	# else:
	# img = self.draw_ldm(
	# output_ldm[0, i], color=(255, 255, 0),
	# )
	# writer_pred.append_data(img)
	# writer_gt.close()
	# writer_kp.close()
	# writer_pred.close()
	return outputs

	@torch.no_grad()
	def test_step(self, sample, batch_idx):
	infer_steps = hparams['infer_steps']
	return self._validation_step(sample, batch_idx, infer_steps)