misc

codes/data/audio/unsupervised_audio_dataset.py

@@ -208,6 +208,8 @@ if __name__ == '__main__':
     for b in tqdm(dl):
         for b_ in range(b['clip'].shape[0]):
             #pass
-            #torchaudio.save(f'{i}_clip_{b_}.wav', b['clip'][b_], ds.sampling_rate)
-            #torchaudio.save(f'{i}_alt_clip_{b_}.wav', b['alt_clips'][b_], ds.sampling_rate)
+            torchaudio.save(f'{i}_clip_{b_}.wav', b['clip'][b_], ds.sampling_rate)
+            torchaudio.save(f'{i}_alt_clip_{b_}.wav', b['alt_clips'][b_], ds.sampling_rate)
             i += 1
+        if i > 200:
+            break

codes/models/audio/music/transformer_diffusion9.py

@@ -1,334 +0,0 @@
-import itertools
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from models.audio.music.music_quantizer2 import MusicQuantizer2
-from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear
-from models.diffusion.unet_diffusion import TimestepBlock, AttentionBlock, TimestepEmbedSequential
-from models.lucidrains.x_transformers import Encoder
-from trainer.networks import register_model
-from utils.util import checkpoint, print_network
-
-
-def is_latent(t):
-    return t.dtype == torch.float
-
-def is_sequence(t):
-    return t.dtype == torch.long
-
-
-class MultiGroupEmbedding(nn.Module):
-    def __init__(self, tokens, groups, dim):
-        super().__init__()
-        self.m = nn.ModuleList([nn.Embedding(tokens, dim // groups) for _ in range(groups)])
-
-    def forward(self, x):
-        h = [embedding(x[:, :, i]) for i, embedding in enumerate(self.m)]
-        return torch.cat(h, dim=-1)
-
-
-class ResBlock(TimestepBlock):
-    def __init__(
-        self,
-        channels,
-        emb_channels,
-        dropout,
-        out_channels=None,
-        dims=2,
-        kernel_size=3,
-        efficient_config=False,
-        use_scale_shift_norm=False,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.emb_channels = emb_channels
-        self.dropout = dropout
-        self.out_channels = out_channels or channels
-        self.use_scale_shift_norm = use_scale_shift_norm
-        padding = {1: 0, 3: 1, 5: 2}[kernel_size]
-        eff_kernel = 1 if efficient_config else 3
-        eff_padding = 0 if efficient_config else 1
-
-        self.in_layers = nn.Sequential(
-            normalization(channels),
-            nn.SiLU(),
-            conv_nd(dims, channels, self.out_channels, eff_kernel, padding=eff_padding),
-        )
-
-        self.emb_layers = nn.Sequential(
-            nn.SiLU(),
-            linear(
-                emb_channels,
-                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
-            ),
-        )
-        self.out_layers = nn.Sequential(
-            normalization(self.out_channels),
-            nn.SiLU(),
-            nn.Dropout(p=dropout),
-            zero_module(
-                conv_nd(dims, self.out_channels, self.out_channels, kernel_size, padding=padding)
-            ),
-        )
-
-        if self.out_channels == channels:
-            self.skip_connection = nn.Identity()
-        else:
-            self.skip_connection = conv_nd(dims, channels, self.out_channels, eff_kernel, padding=eff_padding)
-
-    def forward(self, x, emb):
-        """
-        Apply the block to a Tensor, conditioned on a timestep embedding.
-
-        :param x: an [N x C x ...] Tensor of features.
-        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
-        :return: an [N x C x ...] Tensor of outputs.
-        """
-        return checkpoint(
-            self._forward, x, emb
-        )
-
-    def _forward(self, x, emb):
-        h = self.in_layers(x)
-        emb_out = self.emb_layers(emb).type(h.dtype)
-        while len(emb_out.shape) < len(h.shape):
-            emb_out = emb_out[..., None]
-        if self.use_scale_shift_norm:
-            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
-            scale, shift = torch.chunk(emb_out, 2, dim=1)
-            h = out_norm(h) * (1 + scale) + shift
-            h = out_rest(h)
-        else:
-            h = h + emb_out
-            h = self.out_layers(h)
-        return self.skip_connection(x) + h
-
-
-class DiffusionLayer(TimestepBlock):
-    def __init__(self, model_channels, dropout, num_heads):
-        super().__init__()
-        self.resblk = ResBlock(model_channels, model_channels, dropout, model_channels, dims=1, use_scale_shift_norm=True)
-        self.attn = AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True)
-
-    def forward(self, x, time_emb):
-        y = self.resblk(x, time_emb)
-        return self.attn(y)
-    
-
-class TransformerDiffusion(nn.Module):
-    """
-    A diffusion model composed entirely of stacks of transformer layers. Why would you do it any other way?
-    """
-    def __init__(
-            self,
-            model_channels=512,
-            prenet_layers=3,
-            num_layers=8,
-            in_channels=256,
-            input_vec_dim=512,
-            out_channels=512,  # mean and variance
-            dropout=0,
-            use_fp16=False,
-            # Parameters for regularization.
-            unconditioned_percentage=.1,  # This implements a mechanism similar to what is used in classifier-free training.
-    ):
-        super().__init__()
-
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-        self.out_channels = out_channels
-        self.dropout = dropout
-        self.unconditioned_percentage = unconditioned_percentage
-        self.enable_fp16 = use_fp16
-
-        self.inp_block = conv_nd(1, in_channels, model_channels, 3, 1, 1)
-
-        self.time_embed = nn.Sequential(
-            linear(model_channels, model_channels),
-            nn.SiLU(),
-            linear(model_channels, model_channels),
-        )
-
-        self.input_converter = nn.Linear(input_vec_dim, model_channels)
-        self.code_converter = Encoder(
-                    dim=model_channels,
-                    depth=prenet_layers,
-                    heads=model_channels//64,
-                    ff_dropout=dropout,
-                    attn_dropout=dropout,
-                    use_rmsnorm=True,
-                    ff_glu=True,
-                    rotary_pos_emb=True,
-                    zero_init_branch_output=True,
-                    ff_mult=1,
-                )
-
-        self.unconditioned_embedding = nn.Parameter(torch.randn(1,1,model_channels))
-        self.intg = nn.Conv1d(model_channels*2, model_channels, kernel_size=1)
-        self.layers = TimestepEmbedSequential(*[DiffusionLayer(model_channels, dropout, model_channels // 64) for _ in range(num_layers)])
-
-        self.out = nn.Sequential(
-            normalization(model_channels),
-            nn.SiLU(),
-            zero_module(conv_nd(1, model_channels, out_channels, 3, padding=1)),
-        )
-
-        self.debug_codes = {}
-
-    def get_grad_norm_parameter_groups(self):
-        groups = {
-            'layers': list(self.layers.parameters()) + list(self.inp_block.parameters()),
-            'code_converters': list(self.input_converter.parameters()) + list(self.code_converter.parameters()),
-            'time_embed': list(self.time_embed.parameters()),
-        }
-        return groups
-
-    def timestep_independent(self, codes, expected_seq_len):
-        code_emb = self.input_converter(codes)
-
-        # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance.
-        if self.training and self.unconditioned_percentage > 0:
-            unconditioned_batches = torch.rand((code_emb.shape[0], 1, 1),
-                                               device=code_emb.device) < self.unconditioned_percentage
-            code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(codes.shape[0], 1, 1),
-                                   code_emb)
-        code_emb = self.code_converter(code_emb)
-
-        expanded_code_emb = F.interpolate(code_emb.permute(0,2,1), size=expected_seq_len, mode='nearest').permute(0,2,1)
-        return expanded_code_emb
-
-    def forward(self, x, timesteps, codes=None, conditioning_input=None, precomputed_code_embeddings=None, conditioning_free=False):
-        if precomputed_code_embeddings is not None:
-            assert codes is None and conditioning_input is None, "Do not provide precomputed embeddings and the other parameters. It is unclear what you want me to do here."
-
-        unused_params = []
-        if conditioning_free:
-            code_emb = self.unconditioned_embedding.repeat(x.shape[0], x.shape[-1], 1)
-            unused_params.extend(list(self.code_converter.parameters()))
-        else:
-            if precomputed_code_embeddings is not None:
-                code_emb = precomputed_code_embeddings
-            else:
-                code_emb = self.timestep_independent(codes, x.shape[-1])
-            unused_params.append(self.unconditioned_embedding)
-        code_emb = code_emb.permute(0,2,1)
-
-        blk_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
-        x = self.inp_block(x)
-        x = self.intg(torch.cat([x, code_emb], dim=1))
-        for layer in self.layers:
-            x = checkpoint(layer, x, blk_emb)
-
-        x = x.float()
-        out = self.out(x)
-
-        # Involve probabilistic or possibly unused parameters in loss so we don't get DDP errors.
-        extraneous_addition = 0
-        for p in unused_params:
-            extraneous_addition = extraneous_addition + p.mean()
-        out = out + extraneous_addition * 0
-
-        return out
-
-
-class TransformerDiffusionWithQuantizer(nn.Module):
-    def __init__(self, freeze_quantizer_until=20000, **kwargs):
-        super().__init__()
-
-        self.internal_step = 0
-        self.freeze_quantizer_until = freeze_quantizer_until
-        self.diff = TransformerDiffusion(**kwargs)
-        self.quantizer = MusicQuantizer2(inp_channels=256, inner_dim=[1024], codevector_dim=1024, codebook_size=256,
-                                        codebook_groups=2, max_gumbel_temperature=4, min_gumbel_temperature=.5)
-        self.quantizer.quantizer.temperature = self.quantizer.min_gumbel_temperature
-        del self.quantizer.up
-
-    def update_for_step(self, step, *args):
-        self.internal_step = step
-        qstep = max(0, self.internal_step - self.freeze_quantizer_until)
-        self.quantizer.quantizer.temperature = max(
-            self.quantizer.max_gumbel_temperature * self.quantizer.gumbel_temperature_decay ** qstep,
-                    self.quantizer.min_gumbel_temperature,
-                )
-
-    def forward(self, x, timesteps, truth_mel, conditioning_input=None, disable_diversity=False, conditioning_free=False):
-        quant_grad_enabled = self.internal_step > self.freeze_quantizer_until
-        with torch.set_grad_enabled(quant_grad_enabled):
-            proj, diversity_loss = self.quantizer(truth_mel, return_decoder_latent=True)
-            proj = proj.permute(0,2,1)
-
-        # Make sure this does not cause issues in DDP by explicitly using the parameters for nothing.
-        if not quant_grad_enabled:
-            unused = 0
-            for p in self.quantizer.parameters():
-                unused = unused + p.mean() * 0
-            proj = proj + unused
-            diversity_loss = diversity_loss * 0
-
-        diff = self.diff(x, timesteps, codes=proj, conditioning_input=conditioning_input, conditioning_free=conditioning_free)
-        if disable_diversity:
-            return diff
-        return diff, diversity_loss
-
-    def get_debug_values(self, step, __):
-        if self.quantizer.total_codes > 0:
-            return {'histogram_codes': self.quantizer.codes[:self.quantizer.total_codes]}
-        else:
-            return {}
-
-    def get_grad_norm_parameter_groups(self):
-        groups = {
-            'attention_layers': list(itertools.chain.from_iterable([lyr.attn.parameters() for lyr in self.diff.layers])),
-            'res_layers': list(itertools.chain.from_iterable([lyr.resblk.parameters() for lyr in self.diff.layers])),
-            'quantizer_encoder': list(self.quantizer.encoder.parameters()),
-            'quant_codebook': [self.quantizer.quantizer.codevectors],
-            'out': list(self.diff.out.parameters()),
-            'x_proj': list(self.diff.inp_block.parameters()),
-            'layers': list(self.diff.layers.parameters()),
-            'code_converters': list(self.diff.input_converter.parameters()) + list(self.diff.code_converter.parameters()),
-            'time_embed': list(self.diff.time_embed.parameters()),
-        }
-        return groups
-
-
-@register_model
-def register_transformer_diffusion9(opt_net, opt):
-    return TransformerDiffusion(**opt_net['kwargs'])
-
-
-@register_model
-def register_transformer_diffusion9_with_quantizer(opt_net, opt):
-    return TransformerDiffusionWithQuantizer(**opt_net['kwargs'])
-
-
-"""
-# For TFD5
-if __name__ == '__main__':
-    clip = torch.randn(2, 256, 400)
-    aligned_sequence = torch.randn(2,100,512)
-    cond = torch.randn(2, 256, 400)
-    ts = torch.LongTensor([600, 600])
-    model = TransformerDiffusion(model_channels=3072, model_channels=1536, model_channels=1536)
-    torch.save(model, 'sample.pth')
-    print_network(model)
-    o = model(clip, ts, aligned_sequence, cond)
-"""
-
-if __name__ == '__main__':
-    clip = torch.randn(2, 256, 400)
-    cond = torch.randn(2, 256, 400)
-    ts = torch.LongTensor([600, 600])
-    model = TransformerDiffusionWithQuantizer(model_channels=1024, input_vec_dim=1024, num_layers=16, prenet_layers=6)
-    model.get_grad_norm_parameter_groups()
-
-    quant_weights = torch.load('D:\\dlas\\experiments\\train_music_quant_r4\\models\\5000_generator.pth')
-    #diff_weights = torch.load('X:\\dlas\\experiments\\train_music_diffusion_tfd5\\models\\48000_generator_ema.pth')
-    model.quantizer.load_state_dict(quant_weights, strict=False)
-    #model.diff.load_state_dict(diff_weights)
-
-    torch.save(model.state_dict(), 'sample.pth')
-    print_network(model)
-    o = model(clip, ts, clip, cond)
-

codes/scripts/audio/gen/ctc_codes.py

@@ -1,9 +1,15 @@
 from itertools import groupby
 
 import torch
+import torchaudio
 from transformers import Wav2Vec2CTCTokenizer
 
+from data.audio.voice_tokenizer import VoiceBpeTokenizer
 from models.audio.tts.ctc_code_generator import CtcCodeGenerator
+from models.audio.tts.transformer_diffusion_tts import TransformerDiffusionTTS
+from scripts.audio.gen.speech_synthesis_utils import load_discrete_vocoder_diffuser, load_univnet_vocoder, load_clvp
+from trainer.injectors.audio_injectors import TorchMelSpectrogramInjector, denormalize_mel
+from utils.util import load_audio
 
 
 def get_ctc_metadata(codes):
@@ -28,26 +34,61 @@ def get_ctc_metadata(codes):
     return rcodes, pads, repeats
 
 
+def decode_ctc_metadata(rcodes, pads, repeats):
+    outp = []
+    for s in range(rcodes.shape[-1]):
+        outp = outp + [0 for _ in range(pads[s])]
+        outp = outp + [rcodes[s].item() for _ in range(repeats[s])]
+    return torch.tensor(outp, device=rcodes.device)
+
+
+def diffuse(text, codes, cond):
+    RATIO = 263/140
+
+    codes = codes.cuda();
+    cond = cond.cuda()
+
+    bpe_tokenizer = VoiceBpeTokenizer('../experiments/bpe_lowercase_asr_256.json')
+    clvp = load_clvp().cuda()
+    diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=200, schedule='linear',
+                                             enable_conditioning_free_guidance=False,
+                                             conditioning_free_k=1)
+    diffusion_model = TransformerDiffusionTTS(model_channels=896, num_layers=16, in_channels=100, in_latent_channels=1024,
+                                              token_count=256, out_channels=200, dropout=0, unconditioned_percentage=0)
+    diffusion_model.load_state_dict(torch.load('X:\\dlas\\experiments\\train_speech_diffusion_from_ctc_tfd5\\models\\26500_generator_ema.pth'))
+    diffusion_model = diffusion_model.cuda().eval()
+    with torch.no_grad():
+        text_codes = torch.LongTensor(bpe_tokenizer.encode(text)).unsqueeze(0).to(codes.device)
+        clvp_latent = clvp.embed_text(text_codes)
+        cond_mel = TorchMelSpectrogramInjector({'n_mel_channels': 100, 'mel_fmax': 11000, 'filter_length': 8000, 'normalize': True,
+                                                    'true_normalization': True, 'in': 'in', 'out': 'out'}, {})({'in': cond})['out']
+        gen = diffuser.p_sample_loop(diffusion_model, (1,100,int(codes.shape[-1]*RATIO)), model_kwargs={'codes': codes,
+                                                                                                   'conditioning_input': cond_mel,
+                                                                                                   'type': torch.tensor([0], device=codes.device),
+                                                                                                   'clvp_input': clvp_latent})
+        gen_denorm = denormalize_mel(gen)
+        vocoder = load_univnet_vocoder().cuda()
+        gen_wav = vocoder.inference(gen_denorm)
+    return gen_wav
+
+
 if __name__ == '__main__':
     model = CtcCodeGenerator(model_dim=512, layers=16, dropout=0).eval().cuda()
     model.load_state_dict(torch.load('../experiments/train_encoder_build_ctc_alignments_toy/models/76000_generator_ema.pth'))
 
     tokenizer = Wav2Vec2CTCTokenizer.from_pretrained('jbetker/tacotron-symbols')
-    text = "and now, what do you want."
-    seq = [0, 0, 0, 38, 51, 51, 41, 11, 11, 51, 51, 0, 0, 0, 0, 52, 0, 60, 0, 0, 0, 0, 0, 0, 6, 11, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 60, 45, 0, 38, 57, 57, 11, 0, 41, 52, 52, 11, 11, 62, 52, 52, 58, 0, 11, 11, 60, 0, 0, 0, 0, 38, 0, 0, 51, 51, 0, 0, 57, 0, 0, 7, 7, 0, 0, 0]
-    codes, pads, repeats = get_ctc_metadata(seq)
+    text = "Can I have tea and a pot of butter, please?"
+    #seq = [0, 0, 0, 38, 51, 51, 41, 11, 11, 51, 51, 0, 0, 0, 0, 52, 0, 60, 0, 0, 0, 0, 0, 0, 6, 11, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 60, 45, 0, 38, 57, 57, 11, 0, 41, 52, 52, 11, 11, 62, 52, 52, 58, 0, 11, 11, 60, 0, 0, 0, 0, 38, 0, 0, 51, 51, 0, 0, 57, 0, 0, 7, 7, 0, 0, 0]
+    #codes, pads, repeats = get_ctc_metadata(seq)
+    codes = tokenizer.encode(text)
 
     with torch.no_grad():
-        codes = codes.cuda().unsqueeze(0)
-        pads = pads.cuda().unsqueeze(0)
-        repeats = repeats.cuda().unsqueeze(0)
-
-        ppads = pads.clone()
-        prepeats = repeats.clone()
-        mask = torch.zeros_like(pads)
-        conf_str = tokenizer.decode(codes[0].tolist())
+        codes = torch.tensor(codes).cuda().unsqueeze(0)
+        ppads = torch.zeros_like(codes)
+        prepeats = torch.zeros_like(codes)
+        mask = torch.zeros_like(codes)
         for s in range(codes.shape[-1]):
-            logits, confidences = model.inference(codes, pads * mask, repeats * mask)
+            logits, confidences = model.inference(codes, ppads * mask, prepeats * mask)
 
             confidences = confidences * mask.logical_not()  # prevent prediction of tokens that have already been predicted.
             i = confidences.argmax(dim=-1)
@@ -59,5 +100,9 @@ if __name__ == '__main__':
             prepeats[0,i] = pred_repeats
             mask[0,i] = 1
 
-            conf_str = conf_str[:i] + conf_str[i].upper() + conf_str[i+1:]
-            print(f"conf: {conf_str} pads={pred_pads}:{pads[0,i].item()} repeats={pred_repeats}:{repeats[0,i].item()}")
+            #print(f"conf: {conf_str} pads={pred_pads}:{pads[0,i].item()} repeats={pred_repeats}:{repeats[0,i].item()}")
+
+        decoded_codes = decode_ctc_metadata(codes[0], ppads[0], prepeats[0]).unsqueeze(0)
+        cond = load_audio('D:\\tortoise-tts\\tortoise\\voices\\train_dotrice\\1.wav', 22050).unsqueeze(0).cuda()
+        decoded_wav = diffuse(text, decoded_codes, cond)
+        torchaudio.save('output.wav', decoded_wav.cpu()[0], 24000)

codes/trainer/eval/music_diffusion_fid.py

@@ -62,6 +62,8 @@ class MusicDiffusionFid(evaluator.Evaluator):
             self.local_modules['codegen'] = get_music_codegen()
         elif 'from_codes_quant' == mode:
             self.diffusion_fn = self.perform_diffusion_from_codes_quant
+        elif 'from_codes_quant_gradual_decode' == mode:
+            self.diffusion_fn = self.perform_diffusion_from_codes_quant_gradual_decode
         self.spec_fn = TorchMelSpectrogramInjector({'n_mel_channels': 256, 'mel_fmax': 11000, 'filter_length': 16000,
                                                     'normalize': True, 'in': 'in', 'out': 'out'}, {})
 
@@ -133,6 +135,38 @@ class MusicDiffusionFid(evaluator.Evaluator):
 
         return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate
 
+    def perform_diffusion_from_codes_quant_gradual_decode(self, audio, sample_rate=22050):
+        if sample_rate != sample_rate:
+            real_resampled = torchaudio.functional.resample(audio, 22050, sample_rate).unsqueeze(0)
+        else:
+            real_resampled = audio
+        audio = audio.unsqueeze(0)
+
+        mel = self.spec_fn({'in': audio})['out']
+        mel_norm = normalize_mel(mel)
+        guidance = torch.zeros_like(mel_norm)
+        mask = torch.zeros_like(mel_norm)
+        GRADS = 4
+        for k in range(GRADS):
+            gen_mel = self.diffuser.p_sample_loop_with_guidance(self.model,
+                                                                guidance_input=guidance, mask=mask,
+                                                                model_kwargs={'truth_mel': mel,
+                                                                              'conditioning_input': torch.zeros_like(mel_norm[:,:,:390]),
+                                                                              'disable_diversity': True})
+            pk = int(k*(mel_norm.shape[1]/GRADS))
+            ek = int((k+1)*(mel_norm.shape[1]/GRADS))
+            guidance[:, pk:ek] = gen_mel[:, pk:ek]
+            mask[:, :ek] = 1
+
+        gen_mel_denorm = denormalize_mel(gen_mel)
+        output_shape = (1,16,audio.shape[-1]//16)
+        self.spec_decoder = self.spec_decoder.to(audio.device)
+        gen_wav = self.diffuser.p_sample_loop(self.spec_decoder, output_shape,
+                                              model_kwargs={'aligned_conditioning': gen_mel_denorm})
+        gen_wav = pixel_shuffle_1d(gen_wav, 16)
+
+        return gen_wav, real_resampled, gen_mel, mel_norm, sample_rate
+
 
     def project(self, sample, sample_rate):
         sample = torchaudio.functional.resample(sample, sample_rate, 22050)
@@ -201,13 +235,15 @@ class MusicDiffusionFid(evaluator.Evaluator):
 
 
 if __name__ == '__main__':
-    diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_diffusion_tfd_quant7.yml', 'generator',
+    diffusion = load_model_from_config('X:\\dlas\\experiments\\train_music_diffusion_ar_prior.yml', 'generator',
                                        also_load_savepoint=False,
-                                       load_path='X:\\dlas\\experiments\\train_music_diffusion_unet_music\\models\\46500_generator_ema.pth'
+                                       load_path='X:\\dlas\\experiments\\train_music_diffusion_ar_prior\\models\\22000_generator_ema.pth'
                                        ).cuda()
-    opt_eval = {'path': 'Y:\\split\\yt-music-eval', 'diffusion_steps': 100,
-                'conditioning_free': True, 'conditioning_free_k': 1,
+    opt_eval = {#'path': 'Y:\\split\\yt-music-eval',
+                'path': 'E:\\music_eval',
+                'diffusion_steps': 100,
+                'conditioning_free': False, 'conditioning_free_k': 1,
                 'diffusion_schedule': 'linear', 'diffusion_type': 'from_codes_quant'}
-    env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 561, 'device': 'cuda', 'opt': {}}
+    env = {'rank': 0, 'base_path': 'D:\\tmp\\test_eval_music', 'step': 502, 'device': 'cuda', 'opt': {}}
     eval = MusicDiffusionFid(diffusion, opt_eval, env)
     print(eval.perform_eval())