add training code and model

flowae/configs/datasets/dae.yaml

@@ -62,7 +62,7 @@ datasets:
 
 # Visualization
 visualize_ae_dir: /mnt/nvme/dito_audio
-visualize_ae_random_n_samples: 32
+visualize_ae_random_n_samples: 8
 eval_ae_max_samples: 100
 val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
 

flowae/configs/experiments/dito-B-audio.yaml

@@ -43,6 +43,6 @@ model:
       name: fm
       args: {timescale: 1000.0}
     
-    render_sampler: {name: fm_euler_sampler_audio}
+    render_sampler: {name: fm_euler_sampler}
     render_n_steps: 50
 

flowae/configs/trainers/dito.yaml

@@ -0,0 +1,19 @@
+trainer: audio_ldm_trainer
+
+autocast_bfloat16: true
+
+max_iter: 300000
+epoch_iter: 10000
+eval_iter: 50000
+save_iter: 50000
+vis_iter: 50000
+
+optimizers:
+  encoder:
+    name: adamw
+    args: {lr: 1.e-4}
+  renderer:
+    name: adamw
+    args: {lr: 1.e-4}
+
+evaluate_ae: true

flowae/configs/trainers/glpto.yaml

@@ -0,0 +1,24 @@
+trainer: ldm_trainer
+
+autocast_bfloat16: true
+
+max_iter: 300000
+epoch_iter: 10000
+eval_iter: 50000
+save_iter: 50000
+vis_iter: 50000
+
+optimizers:
+  encoder:
+    name: adam
+    args: {lr: 1.e-4, betas: [0.5, 0.9]}
+  renderer:
+    name: adam
+    args: {lr: 1.e-4, betas: [0.5, 0.9]}
+  disc:
+    name: adam
+    args: {lr: 1.e-4, betas: [0.5, 0.9]}
+gan_start_after_iters: 50000
+find_unused_parameters: true
+
+evaluate_ae: true

flowae/configs/trainers/zdm.yaml

@@ -0,0 +1,20 @@
+trainer: ldm_trainer
+
+autocast_bfloat16: true
+
+max_iter: 400000
+epoch_iter: 10000
+eval_iter: 100000
+save_iter: 100000
+vis_iter: 100000
+ckpt_select_metric:
+  name: zdm_ema_loss
+  type: min
+
+optimizers:
+  zdm:
+    name: adamw
+    args: {lr: 1.e-4, weight_decay: 0.0}
+find_unused_parameters: true
+
+evaluate_zdm: true

flowae/models/networks/consistency_audio_decoder_unet.py

@@ -135,13 +135,13 @@ class AudioUpsample(nn.Module):
 
         gn_1 = F.silu(self.gn_1(x))
         # 1D interpolation upsampling
-        upsample = F.interpolate(gn_1, scale_factor=self.upsample_factor, mode='nearest')
+        upsample = F.interpolate(gn_1, scale_factor=self.upsample_factor, mode='linear')
         f_1 = self.f_1(upsample)
         gn_2 = self.gn_2(f_1)
 
         f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2)))
 
-        return f_2 + F.interpolate(x_skip, scale_factor=self.upsample_factor, mode='nearest')
+        return f_2 + F.interpolate(x_skip, scale_factor=self.upsample_factor, mode='linear')
 
 
 @register('audio_diffusion_unet')
@@ -272,7 +272,7 @@ class AudioDiffusionUNet(nn.Module):
             z_proj = F.interpolate(
                 z_proj, 
                 size=x.shape[-1], 
-                mode='nearest'  # or 'linear' for smoother interpolation
+                mode='linear'  # or 'linear' for smoother interpolation
             )
         
         # Add latent conditioning to audio features

flowae/run.py

@@ -13,7 +13,7 @@ def make_args():
     parser.add_argument('--tag', '-t', default=None)
     parser.add_argument('--resume', '-r', action='store_true')
     parser.add_argument('--force-replace', '-f', action='store_true')
-    parser.add_argument('--wandb', '-w', action='store_true')
+    parser.add_argument('--comet', '-c', action='store_true', help='Enable Comet ML logging')
     parser.add_argument('--save-root', default='save')
     parser.add_argument('--eval-only', action='store_true')
     args = parser.parse_args()
@@ -45,7 +45,7 @@ def make_env(args):
     env['exp_name'] = exp_name
     
     env['save_dir'] = os.path.join(args.save_root, exp_name)
-    env['wandb'] = args.wandb
+    env['comet'] = args.comet
     env['resume'] = args.resume
     env['force_replace'] = args.force_replace
     return env

flowae/run.sh

@@ -1,2 +1,2 @@
 torchrun --nnodes=1 --nproc-per-node=1 run.py --config configs/experiments/dito-B-f8c4-noise-sync.yaml --save-root /mnt/nvme/dito
-torchrun --nnodes=1 --nproc-per-node=1 run.py --config configs/experiments/dito-B-audio.yaml --save-root /mnt/nvme/ditogit ad
+torchrun --nnodes=1 --nproc-per-node=1 run.py --config configs/experiments/dito-B-audio.yaml --save-root /mnt/nvme/dit2 --comet

flowae/trainers/__init__.py

@@ -0,0 +1,4 @@
+from .trainers import register, trainers_dict
+from . import base_trainer
+from . import ldm_trainer
+from . import audio_ldm_trainer

flowae/trainers/audio_ldm_trainer.py

@@ -0,0 +1,715 @@
+import os
+import random
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+import utils
+from .trainers import register
+from trainers.base_trainer import BaseTrainer
+from models.ldm.dac.audiotools import AudioSignal
+import soundfile as sf
+import numpy as np
+import torchaudio
+import time
+
+from datetime import datetime
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+
+@register('audio_ldm_trainer')
+class AudioLDMTrainer(BaseTrainer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        
+    def make_model(self):
+        super().make_model()
+        self.has_optimizer = dict()
+        total_params = 0
+        for name, m in self.model.named_children():
+            params = utils.compute_num_params(m, text=False)
+            self.log(f'  .{name} {params}')
+            total_params = total_params + params
+            # Log to Comet
+            if self.experiment:
+                self.experiment.log_metric(f"model/{name}_params", params)
+        
+        if self.experiment:
+            self.experiment.log_metric("model/total_params", total_params)
+
+    def make_optimizers(self):
+        self.optimizers = dict()
+        self.has_optimizer = dict()
+        for name, spec in self.config.optimizers.items():
+            self.optimizers[name] = utils.make_optimizer(self.model.get_parameters(name), spec)
+            self.has_optimizer[name] = True
+
+            # Log optimizer config to Comet
+            if self.experiment:
+                self.experiment.log_parameters({
+                    f"optimizer/{name}/type": spec.get("type", "adam"),
+                    f"optimizer/{name}/lr": spec.get("lr", 1e-4),
+                    f"optimizer/{name}/weight_decay": spec.get("weight_decay", 0),
+                })
+
+    def train_step(self, data, bp=True):
+        kwargs = {'has_optimizer': self.has_optimizer}
+        
+        # Start timing
+        step_start_time = time.time()
+        # Audio-specific data preparation
+        if 'signal' in data:
+            # Convert AudioSignal to tensor format expected by model
+            audio_data = data['signal'].audio_data  # [batch, channels, samples]
+            sample_rate = data['signal'].sample_rate
+            
+            # Prepare data dict for model
+            model_data = {
+                'inp': audio_data,
+                'gt': audio_data,  # For autoencoder training
+                'sample_rate': sample_rate
+            }
+        else:
+            model_data = data
+            
+        # self.log(f'Audio data shape: {model_data["inp"].shape}')
+
+        # Log batch info to Comet
+        if self.experiment and self.iter % 500 == 0:
+            self.experiment.log_metric("train/batch_size", model_data["inp"].shape[0], step=self.iter)
+            self.experiment.log_metric("train/audio_length_samples", model_data["inp"].shape[-1], step=self.iter)
+            self.experiment.log_metric("train/audio_duration_sec", 
+                                     model_data["inp"].shape[-1] / model_data.get("sample_rate", 24000), 
+                                     step=self.iter)
+        
+
+        if self.config.get('autocast_bfloat16', False):
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                ret = self.model_ddp(model_data, mode='loss', **kwargs)
+        else:
+            ret = self.model_ddp(model_data, mode='loss', **kwargs)
+
+        loss = ret.pop('loss')
+        ret['loss'] = loss.item()
+        
+        if bp:
+            self.model_ddp.zero_grad(set_to_none=True)
+            loss.backward()
+            
+            # Log gradients to Comet
+            if self.experiment and self.iter % 5 == 0:
+                self._log_gradients()
+            
+            for name, o in self.optimizers.items():
+                if name != 'disc':
+                    o.step()
+
+
+        if hasattr(self.model, 'update_ema'):
+            self.model.update_ema()
+        
+        # Log training metrics to Comet
+        if self.experiment:
+            # Log all losses
+            for k, v in ret.items():
+                if 'loss' in k.lower():
+                    self.experiment.log_metric(f"train/{k}", v, step=self.iter)
+            
+            # Log learning rates
+            for name, opt in self.optimizers.items():
+                lr = opt.param_groups[0]['lr']
+                self.experiment.log_metric(f"train/lr_{name}", lr, step=self.iter)
+            
+            # Log timing
+            step_time = time.time() - step_start_time
+            self.experiment.log_metric("train/step_time", step_time, step=self.iter)
+            
+            # Log GPU memory usage
+            if torch.cuda.is_available():
+                self.experiment.log_metric("train/gpu_memory_allocated", 
+                                         torch.cuda.memory_allocated() / 1e9, 
+                                         step=self.iter)
+                self.experiment.log_metric("train/gpu_memory_reserved", 
+                                         torch.cuda.memory_reserved() / 1e9, 
+                                         step=self.iter)
+        return ret
+
+    def _log_gradients(self):
+        """Log gradient statistics to Comet ML"""
+        if not self.experiment:
+            return
+            
+        grad_stats = {}
+        for name, param in self.model.named_parameters():
+            if param.grad is not None:
+                grad_norm = param.grad.norm().item()
+                grad_mean = param.grad.mean().item()
+                grad_std = param.grad.std().item()
+                
+                # Log aggregate stats by module
+                module_name = name.split('.')[0]
+                if module_name not in grad_stats:
+                    grad_stats[module_name] = {
+                        'norm': [],
+                        'mean': [],
+                        'std': []
+                    }
+                grad_stats[module_name]['norm'].append(grad_norm)
+                grad_stats[module_name]['mean'].append(grad_mean)
+                grad_stats[module_name]['std'].append(grad_std)
+        
+        # Log aggregated stats
+        for module, stats in grad_stats.items():
+            self.experiment.log_metric(f"gradients/{module}/norm_mean", np.mean(stats['norm']), step=self.iter)
+            self.experiment.log_metric(f"gradients/{module}/norm_max", np.max(stats['norm']), step=self.iter)
+
+
+    def run_training(self):
+        config = self.config
+        max_iter = config['max_iter']
+        epoch_iter = config['epoch_iter']
+        assert max_iter % epoch_iter == 0
+        max_epoch = max_iter // epoch_iter
+
+        save_iter = config.get('save_iter')
+        if save_iter is not None:
+            assert save_iter % epoch_iter == 0
+            save_epoch = save_iter // epoch_iter
+            print('save_epoch', save_epoch)
+        else:
+            save_epoch = max_epoch + 1
+
+        eval_iter = config.get('eval_iter')
+        if eval_iter is not None:
+            assert eval_iter % epoch_iter == 0
+            eval_epoch = eval_iter // epoch_iter
+        else:
+            eval_epoch = max_epoch + 1
+
+        vis_iter = config.get('vis_iter')
+        if vis_iter is not None:
+            assert vis_iter % epoch_iter == 0
+            vis_epoch = vis_iter // epoch_iter
+        else:
+            vis_epoch = max_epoch + 1
+
+        if config.get('ckpt_select_metric') is not None:
+            m = config.ckpt_select_metric
+            self.ckpt_select_metric = m.name
+            self.ckpt_select_type = m.type
+            if m.type == 'min':
+                self.ckpt_select_v = 1e18
+            elif m.type == 'max':
+                self.ckpt_select_v = -1e18
+        else:
+            self.ckpt_select_metric = None
+            self.ckpt_select_v = 0
+
+        self.train_loader = self.loaders['train']
+        self.train_loader_sampler = self.loader_samplers['train']
+        self.train_loader_epoch = 0
+        self.train_loader_iter = None
+
+        self.iter = 0
+
+        if self.resume_ckpt is not None:
+            for _ in range(self.resume_ckpt['iter']):
+                self.iter += 1
+                self.at_train_iter_start()
+            self.ckpt_select_v = self.resume_ckpt['ckpt_select_v']
+            self.train_loader_epoch = self.resume_ckpt['train_loader_epoch']
+            self.train_loader_iter = None
+            self.resume_ckpt = None
+            self.log(f'Resumed iter status.')
+        
+        self.visualize()
+
+        start_epoch = self.iter // epoch_iter + 1
+
+        for epoch in range(start_epoch, max_epoch + 1):
+            self.log_buffer = [f'Epoch {epoch}']
+
+            for sampler in self.loader_samplers.values():
+                if sampler is not self.train_loader_sampler:
+                    sampler.set_epoch(epoch)
+
+            self.model_ddp.train()
+
+            pbar = range(1, epoch_iter + 1)
+            if self.is_master and epoch == start_epoch:
+                pbar = tqdm(pbar, desc='train', leave=False)
+
+            t_data = 0
+            t_nondata = 0
+            t_before_data = time.time()
+            
+            for _ in pbar:
+                self.iter += 1
+                self.at_train_iter_start()
+
+                try:
+                    if self.train_loader_iter is None:
+                        raise StopIteration
+                    data = next(self.train_loader_iter)
+                except StopIteration:
+                    self.train_loader_epoch += 1
+                    self.train_loader_sampler.set_epoch(self.train_loader_epoch)
+                    self.train_loader_iter = iter(self.train_loader)
+                    data = next(self.train_loader_iter)
+                
+                t_after_data = time.time()
+                t_data += t_after_data - t_before_data
+                
+                for k, v in data.items():
+                    data[k] = v.to(self.device) if torch.is_tensor(v) else v
+                
+                ret = self.train_step(data)
+                
+                t_before_data = time.time()
+                t_nondata += t_before_data - t_after_data
+
+                if self.is_master and epoch == start_epoch:
+                    pbar.set_description(desc=f'train: loss={ret["loss"]:.4f}')
+                
+                # save the model every 1000 iterations
+                if self.iter % 2000 == 0:
+                    self.save_ckpt(f'ckpt-{self.iter}.pth')
+            
+            self.save_ckpt('ckpt-last.pth')
+
+            if epoch % save_epoch == 0 and epoch != max_epoch:
+                self.save_ckpt(f'ckpt-{self.iter}.pth')
+
+            if epoch % eval_epoch == 0:
+                with torch.no_grad():
+                    eval_ave_scalars = self.evaluate()
+                if self.ckpt_select_metric is not None:
+                    v = eval_ave_scalars[self.ckpt_select_metric].item()
+                    if ((self.ckpt_select_type == 'min' and v < self.ckpt_select_v) or
+                        (self.ckpt_select_type == 'max' and v > self.ckpt_select_v)):
+                        self.ckpt_select_v = v
+                        self.save_ckpt('ckpt-best.pth')
+
+            if epoch % vis_epoch == 0:
+                with torch.no_grad():
+                    self.visualize()
+    
+    def evaluate(self):
+        self.model_ddp.eval()
+
+        ave_scalars = dict()
+        pbar = self.loaders['val']
+
+        for data in pbar:
+            # Prepare audio data for GPU
+            if 'signal' in data:
+                data['signal'] = data['signal'].to(self.device)
+            else:
+                for k, v in data.items():
+                    data[k] = v.to(self.device) if torch.is_tensor(v) else v
+            
+            ret = self.train_step(data, bp=False)
+
+            bs = data['signal'].batch_size if 'signal' in data else len(next(iter(data.values())))
+            for k, v in ret.items():
+                if ave_scalars.get(k) is None:
+                    ave_scalars[k] = utils.Averager()
+                ave_scalars[k].add(v, n=bs)
+        
+        self.sync_ave_scalars(ave_scalars)
+        
+        # Audio-specific evaluation
+        if self.config.get('evaluate_ae', False):
+            ave_scalars.update(self.evaluate_audio_ae())
+        
+        if self.config.get('evaluate_zdm', False):
+            ema = self.config.get('evaluate_zdm_ema', True)
+            ave_scalars.update(self.evaluate_audio_zdm(ema=ema))
+
+        logtext = 'val:'
+        for k, v in ave_scalars.items():
+            logtext += f' {k}={v.item():.4f}'
+            self.log_scalar('val/' + k, v.item())
+            
+            # Log to Comet
+            if self.experiment:
+                self.experiment.log_metric(f"val/{k}", v.item(), step=self.iter)
+        
+        self.log_buffer.append(logtext)
+        
+        return ave_scalars
+
+    def visualize(self):
+        self.model_ddp.eval()
+
+        if self.config.get('evaluate_ae', False):
+            self.visualize_audio_ae_random()
+        
+        if self.config.get('evaluate_zdm', False):
+            ema = self.config.get('evaluate_zdm_ema', True)
+            self.visualize_audio_zdm_random(ema=ema)
+
+    def evaluate_audio_ae(self):
+        """Audio autoencoder evaluation with spectral metrics"""
+        max_samples = self.config.get('eval_ae_max_samples', 1000)
+        self.loader_samplers['eval_ae'].set_epoch(0)
+        
+        l1_loss_avg = utils.Averager()
+        snr_avg = utils.Averager()
+        spectral_convergence_avg = utils.Averager()
+        cnt = 0
+
+        # Create cache directories for audio samples
+        cache_gen_dir = os.path.join(self.env['save_dir'], 'cache', 'audio_gen')
+        cache_gt_dir = os.path.join(self.env['save_dir'], 'cache', 'audio_gt')
+        if self.is_master:
+            utils.ensure_path(cache_gen_dir, force_replace=True)
+            utils.ensure_path(cache_gt_dir, force_replace=True)
+        dist.barrier()
+
+        for data in self.loaders['eval_ae']:
+            if 'signal' in data:
+                data['signal'] = data['signal'].to(self.device)
+                signal = data['signal']
+            else:
+                for k, v in data.items():
+                    data[k] = v.to(self.device) if torch.is_tensor(v) else v
+                signal = AudioSignal(data['inp'], data.get('sample_rate', 22050))
+            
+            # Get reconstruction
+            pred_audio = self.model(data, mode='pred')
+            if isinstance(pred_audio, dict):
+                pred_audio = pred_audio.get('audio', pred_audio.get('recons', pred_audio))
+            
+            recons = AudioSignal(pred_audio, signal.sample_rate)
+
+            # SNR calculation
+            signal_power = (signal.audio_data ** 2).mean()
+            noise_power = ((recons.audio_data - signal.audio_data) ** 2).mean()
+            snr = 10 * torch.log10(signal_power / (noise_power + 1e-8))
+            snr_avg.add(snr.item())
+            
+            # Spectral convergence
+            stft_transform = torchaudio.transforms.Spectrogram(
+                n_fft=1024,
+                hop_length=256,
+                power=2
+            ).to(self.device)
+            
+            orig_spec = stft_transform(signal.audio_data)
+            recon_spec = stft_transform(recons.audio_data)
+            
+            spec_diff = torch.norm(orig_spec - recon_spec, p='fro')
+            spec_norm = torch.norm(orig_spec, p='fro')
+            spectral_convergence = spec_diff / (spec_norm + 1e-8)
+            spectral_convergence_avg.add(spectral_convergence.item())
+            
+            l1_loss = torch.nn.functional.l1_loss(recons.audio_data, signal.audio_data).item()
+            l1_loss_avg.add(l1_loss)
+
+            # Save audio samples for potential subjective evaluation
+            for i in range(min(signal.batch_size, 5)):  # Save up to 5 per batch
+                idx = int(os.environ['RANK']) + cnt * int(os.environ['WORLD_SIZE'])
+                if max_samples is None or idx < max_samples:
+                    # Save as wav files
+                    sf.write(
+                        os.path.join(cache_gen_dir, f'{idx}.wav'),
+                        recons[i].audio_data.cpu().numpy().T,
+                        int(recons[i].sample_rate)
+                    )
+                    sf.write(
+                        os.path.join(cache_gt_dir, f'{idx}.wav'),
+                        signal[i].audio_data.cpu().numpy().T,
+                        int(signal[i].sample_rate)
+                    )
+                cnt += 1
+    
+        dist.barrier()
+
+        # Sync metrics across processes
+        for avg_metric in [l1_loss_avg, snr_avg, spectral_convergence_avg]:
+            vt = torch.tensor(avg_metric.item(), device=self.device)
+            dist.all_reduce(vt, op=dist.ReduceOp.SUM)
+            torch.cuda.synchronize()
+            avg_metric.v = vt.item() / int(os.environ['WORLD_SIZE'])
+
+        if self.is_master:
+            prefix = 'eval_ae'
+            ret = {
+                f'{prefix}/L1_Loss': l1_loss_avg.item(),
+                f'{prefix}/SNR': snr_avg.item(),
+                f'{prefix}/Spectral_Convergence': spectral_convergence_avg.item(),
+            }
+        else:
+            ret = {}
+        dist.barrier()
+
+        ret = {k: utils.Averager(v) for k, v in ret.items()}
+        return ret
+
+    def evaluate_audio_zdm(self, ema):
+        """Audio latent diffusion model evaluation"""
+        max_samples = self.config.get('eval_zdm_max_samples', 1000)
+        self.loader_samplers['eval_zdm'].set_epoch(0)
+    
+        cnt = 0
+        l1_loss_avg = utils.Averager()
+        cache_gen_dir = os.path.join(self.env['save_dir'], 'cache', 'audio_gen')
+        cache_gt_dir = os.path.join(self.env['save_dir'], 'cache', 'audio_gt')
+        if self.is_master:
+            utils.ensure_path(cache_gen_dir, force_replace=True)
+            utils.ensure_path(cache_gt_dir, force_replace=True)
+        dist.barrier()
+
+        for data in self.loaders['eval_zdm']:
+            if 'signal' in data:
+                data['signal'] = data['signal'].to(self.device)
+                gt_signal = data['signal']
+            else:
+                for k, v in data.items():
+                    data[k] = v.to(self.device) if torch.is_tensor(v) else v
+                gt_signal = AudioSignal(data['inp'], data.get('sample_rate', 22050))
+            
+            # Generate samples from latent diffusion model
+            net_kwargs = dict()
+            uncond_net_kwargs = dict()
+            # Add conditioning if available (e.g., for conditional generation)
+            
+            pred_audio = self.model.generate_samples(
+                batch_size=gt_signal.batch_size,
+                n_steps=self.model.zdm_n_steps,
+                net_kwargs=net_kwargs,
+                uncond_net_kwargs=uncond_net_kwargs,
+                ema=ema
+            )
+            
+            pred_signal = AudioSignal(pred_audio, gt_signal.sample_rate)
+
+            l1_loss = torch.nn.functional.l1_loss(pred_signal.audio_data, gt_signal.audio_data).item()
+            l1_loss_avg.add(l1_loss)
+
+            # Save samples
+            for i in range(min(gt_signal.batch_size, 5)):
+                idx = int(os.environ['RANK']) + cnt * int(os.environ['WORLD_SIZE'])
+                if max_samples is None or idx < max_samples:
+                    sf.write(
+                        os.path.join(cache_gen_dir, f'{idx}.wav'),
+                        pred_signal[i].audio_data.cpu().numpy().T,
+                        int(pred_signal[i].sample_rate)
+                    )
+                    sf.write(
+                        os.path.join(cache_gt_dir, f'{idx}.wav'),
+                        gt_signal[i].audio_data.cpu().numpy().T,
+                        int(gt_signal[i].sample_rate)
+                    )
+                cnt += 1
+        
+        dist.barrier()
+
+        # Sync metrics
+        for avg_metric in [l1_loss_avg]:
+            vt = torch.tensor(avg_metric.item(), device=self.device)
+            dist.all_reduce(vt, op=dist.ReduceOp.SUM)
+            torch.cuda.synchronize()
+            avg_metric.v = vt.item() / int(os.environ['WORLD_SIZE'])
+
+        if self.is_master:
+            prefix = 'eval_zdm' + ('_ema' if ema else '')
+            ret = {
+                f'{prefix}/l1_loss_avg': l1_loss_avg.item(),
+            }
+        else:
+            ret = {}
+        dist.barrier()
+
+        ret = {k: utils.Averager(v) for k, v in ret.items()}
+        return ret
+    
+    def visualize_audio_ae_random(self):
+        """Save random audio reconstructions for listening"""
+        if self.is_master:
+            idx_list = list(range(len(self.datasets['eval_ae'])))
+            random.shuffle(idx_list)
+            n_samples = self.config.get('visualize_ae_random_n_samples', 8)
+            
+            audio_samples = []
+            
+            for idx in idx_list[:n_samples]:
+                data = self.datasets['eval_ae'][idx]
+                
+                # Prepare data
+                if 'signal' in data:
+                    signal = data['signal'].unsqueeze(0).to(self.device)
+                    model_data = {
+                        'inp': signal.audio_data,
+                        'gt': signal.audio_data,
+                        'sample_rate': signal.sample_rate
+                    }
+                else:
+                    for k, v in data.items():
+                        data[k] = v.unsqueeze(0).to(self.device) if torch.is_tensor(v) else v
+                    signal = AudioSignal(data['inp'], data.get('sample_rate', 24000))
+                    model_data = data
+                
+                # Get reconstruction
+                pred_audio = self.model(model_data, mode='pred')
+                if isinstance(pred_audio, dict):
+                    pred_audio = pred_audio.get('audio', pred_audio.get('recons', pred_audio))
+                
+                recons = AudioSignal(pred_audio, signal.sample_rate)
+                
+                # Save to file and log to Comet
+                self.save_audio_sample(signal, f'audio_ae_original_{idx}')
+                self.save_audio_sample(recons, f'audio_ae_recons_{idx}')
+                
+        dist.barrier()
+    
+    def visualize_audio_zdm_random(self, ema):
+        """Save random audio generations from latent diffusion model"""
+        if self.is_master:
+            n_samples = self.config.get('visualize_zdm_random_n_samples', 8)
+            
+            for i in range(n_samples):
+                # Generate random sample
+                net_kwargs = dict()
+                uncond_net_kwargs = dict()
+                
+                # Get a reference from dataset for parameters like sample_rate
+                ref_data = self.datasets['eval_ae'][0]
+                if 'signal' in ref_data:
+                    ref_signal = ref_data['signal']
+                    sample_rate = ref_signal.sample_rate
+                    batch_size = 1
+                else:
+                    sample_rate = ref_data.get('sample_rate', 24000)
+                    batch_size = 1
+                
+                pred_audio = self.model.generate_samples(
+                    batch_size=batch_size,
+                    n_steps=self.model.zdm_n_steps,
+                    net_kwargs=net_kwargs,
+                    uncond_net_kwargs=uncond_net_kwargs,
+                    ema=ema
+                )
+                
+                pred_signal = AudioSignal(pred_audio, sample_rate)
+                
+                # Save generated audio
+                self.save_audio_sample(pred_signal, f'audio_zdm_generated_{i}')
+                
+        dist.barrier()
+
+    def save_audio_sample(self, audio_signal, name):
+        """Save audio sample and log to Comet ML"""
+        try:
+            # Ensure audio is in correct format
+            audio_data = audio_signal.audio_data.cpu()
+            
+            # Handle different dimensions
+            if audio_data.dim() == 3:  # [batch, channels, samples]
+                audio_data = audio_data[0]  # Take first sample
+            if audio_data.dim() == 2:  # [channels, samples]
+                audio_data = audio_data.transpose(0, 1)  # [samples, channels]
+            elif audio_data.dim() == 1:  # [samples]
+                audio_data = audio_data.unsqueeze(1)  # [samples, 1]
+            
+            audio_data = audio_data.numpy()
+            
+            # Normalize if needed
+            if np.abs(audio_data).max() > 1.0:
+                audio_data = audio_data / np.abs(audio_data).max()
+            
+            # Save to file
+            save_path = os.path.join(self.env['save_dir'], 'audio_samples')
+            os.makedirs(save_path, exist_ok=True)
+            
+            file_path = os.path.join(save_path, f'{name}_step_{self.iter}.wav')
+            sf.write(file_path, audio_data, int(audio_signal.sample_rate))
+            
+            # Log to Comet ML
+            if self.experiment:
+                self.experiment.log_audio(
+                    file_path,
+                    metadata={
+                        'name': name,
+                        'step': self.iter,
+                        'sample_rate': int(audio_signal.sample_rate),
+                        'duration': len(audio_data) / audio_signal.sample_rate,
+                        'channels': audio_data.shape[1] if audio_data.ndim > 1 else 1
+                    },
+                    step=self.iter
+                )
+                
+                # Also log spectrograms for visualization
+                if self.iter % self.config.get('spectrogram_log_freq', 1000) == 0:
+                    self._log_spectrogram(audio_signal, name)
+            
+            self.log(f"Saved audio sample: {file_path}")
+            
+        except Exception as e:
+            self.log(f"Error saving audio sample {name}: {e}")
+            if self.experiment:
+                self.experiment.log_text(f"Error saving audio {name}: {str(e)}", step=self.iter)
+    
+    def _log_spectrogram(self, audio_signal, name):
+        """Log spectrogram visualization to Comet ML"""
+        if not self.experiment:
+            return
+            
+        try:
+            
+            # Compute spectrogram
+            stft_transform = torchaudio.transforms.Spectrogram(
+                n_fft=2048,
+                hop_length=512,
+                power=2
+            )
+            
+            audio_data = audio_signal.audio_data
+            if audio_data.dim() == 3:
+                audio_data = audio_data[0]
+            if audio_data.dim() == 2:
+                audio_data = audio_data[0]  # Take first channel
+                
+            spec = stft_transform(audio_data.cpu())
+            spec_db = 10 * torch.log10(spec + 1e-8)
+            
+            # Create figure
+            fig, ax = plt.subplots(figsize=(10, 4))
+            im = ax.imshow(
+                spec_db.numpy(),
+                aspect='auto',
+                origin='lower',
+                cmap='viridis',
+                extent=[0, len(audio_data) / audio_signal.sample_rate, 0, audio_signal.sample_rate / 2]
+            )
+            ax.set_xlabel('Time (s)')
+            ax.set_ylabel('Frequency (Hz)')
+            ax.set_title(f'{name} - Spectrogram')
+            plt.colorbar(im, ax=ax, label='dB')
+            
+            # Log to Comet
+            self.experiment.log_figure(f"spectrogram/{name}", fig, step=self.iter)
+            plt.close(fig)
+            
+        except Exception as e:
+            self.log(f"Error logging spectrogram for {name}: {e}")
+    
+    
+    
+    def save_checkpoint(self, tag="latest"):
+        """Save checkpoint and log to Comet ML"""
+        checkpoint_path = super().save_checkpoint(tag)
+        
+        if self.experiment and checkpoint_path:
+            # Log checkpoint to Comet
+            self.experiment.log_model(
+                f"checkpoint_{tag}",
+                checkpoint_path,
+                metadata={
+                    "step": self.iter,
+                    "tag": tag,
+                    "timestamp": datetime.now().isoformat()
+                }
+            )

flowae/trainers/base_trainer.py

@@ -0,0 +1,418 @@
+import os
+import time
+import copy
+from datetime import timedelta
+
+import yaml
+import torch
+import torch.distributed as dist
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from torch.utils.data import IterableDataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel
+
+import datasets
+import models
+import utils
+from .trainers import register
+from comet_ml import Experiment
+
+from datetime import datetime
+
+@register('base_trainer')
+class BaseTrainer():
+
+    def __init__(self, env, config):
+        self.env = env
+        self.config = config
+        self.config_dict = OmegaConf.to_container(config, resolve=True)
+
+        if config.get('allow_tf32', False):
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+        
+        dist.init_process_group(backend='nccl', timeout=timedelta(minutes=240))
+        self.rank = int(os.environ['RANK'])
+        self.local_rank = int(os.environ['LOCAL_RANK'])
+        self.world_size = int(os.environ['WORLD_SIZE'])
+        self.node_id = int(os.environ['GROUP_RANK'])
+        self.node_tot = self.world_size // int(os.environ['LOCAL_WORLD_SIZE'])
+        self.is_master = (self.rank == 0)
+        
+        torch.cuda.set_device(self.local_rank)
+        self.device = torch.device('cuda', torch.cuda.current_device())
+
+        if self.is_master:
+            # Setup path
+            if env['resume']:
+                replace = False
+                force_replace = False
+            else:
+                replace = True
+                force_replace = env['force_replace']
+            utils.ensure_path(env['save_dir'], replace=replace, force_replace=force_replace)
+            
+            # Save config
+            with open(os.path.join(env['save_dir'], 'config.yaml'), 'w') as f:
+                yaml.dump(self.config_dict, f, sort_keys=False)
+
+            # Setup logging
+            logger = utils.set_logger(os.path.join(env['save_dir'], 'log.txt'))
+            self.log = logger.info
+        
+            # Initialize Comet ML experiment
+            self.experiment = None
+            if self.is_master:  # Only log from master process
+                self.experiment = Experiment(
+                    project_name=self.config.get("comet_project", "audio-ldm"),
+                    workspace=os.environ.get("COMET_WORKSPACE"),
+                    experiment_name=self.config.get("exp_name", f"audio_ldm_{datetime.now().strftime('%Y%m%d_%H%M%S')}")
+                )
+                
+                # Log hyperparameters
+                self.experiment.log_parameters(self.config)
+              
+                # Add tags
+                tags = self.config.get("tags", ["audio", "ldm", "diffusion"])
+                for tag in tags:
+                    self.experiment.add_tag(tag)
+        else:
+            self.log = lambda *args, **kwargs: None
+            self.experiment = None
+        dist.barrier()
+
+        self.log(f'Environment setup done. World size: {self.world_size}.')
+
+    def run(self, eval_only=False):
+        self.make_datasets()
+
+        resume_ckpt = os.path.join(self.env['save_dir'], 'ckpt-last.pth')
+        resume = (self.env['resume'] and os.path.isfile(resume_ckpt))
+        if resume:
+            self.resume_ckpt = torch.load(resume_ckpt, map_location='cpu')
+        else:
+            self.resume_ckpt = None
+
+        self.make_model()
+        if resume:
+            self.model.load_state_dict(self.resume_ckpt['model']['sd'])
+            self.resume_ckpt['model'] = None
+            self.log(f'Resumed model from checkpoint {resume_ckpt}.')
+
+        if eval_only:
+            self.model_ddp = self.model
+            with torch.no_grad():
+                self.log_buffer = [f'Eval']
+                self.iter = 0
+                self.evaluate()
+                self.visualize()
+                self.log(', '.join(self.log_buffer))
+        
+        else:
+            self.model_ddp = DistributedDataParallel(
+                self.model,
+                device_ids=[self.local_rank],
+                find_unused_parameters=self.config.get('find_unused_parameters', False)
+            )
+            
+            self.make_optimizers()
+            if resume:
+                for name, optimizer in self.resume_ckpt['optimizers'].items():
+                    self.optimizers[name].load_state_dict(optimizer['sd'])
+                self.resume_ckpt['optimizers'] = None
+                self.log(f'Resumed optimizers.')
+
+            self.run_training()
+
+        self.on_train_end()
+
+    def on_train_end(self):
+        """Called at the end of training"""
+        if self.experiment:
+            # Log final model
+            model_path = os.path.join(self.env['save_dir'], 'final_model.pt')
+            torch.save(self.model.state_dict(), model_path)
+            self.experiment.log_model("final_model", model_path)
+            
+            # End the experiment
+            self.experiment.end()
+
+    def make_distributed_loader(self, dataset, batch_size, shuffle, drop_last, num_workers, pin_memory):
+        assert batch_size % self.world_size == 0
+        assert num_workers % self.world_size == 0
+        if isinstance(dataset, IterableDataset):
+            sampler = None
+        else:
+            sampler = DistributedSampler(dataset, shuffle=shuffle)
+        loader = DataLoader(
+            dataset,
+            batch_size=batch_size // self.world_size,
+            drop_last=drop_last,
+            sampler=sampler,
+            num_workers=num_workers // self.world_size,
+            pin_memory=pin_memory
+        )
+        return loader, sampler
+
+    def make_datasets(self):
+        self.datasets = dict()
+        self.loaders = dict()
+        self.loader_samplers = dict()
+
+        for split, spec in self.config.datasets.items():
+            loader_spec = spec.pop('loader')
+            
+            dataset = datasets.make(spec)
+            self.datasets[split] = dataset
+            if isinstance(dataset, IterableDataset):
+                self.log(f'Dataset {split}: IterableDataset')
+            else:
+                self.log(f'Dataset {split}: len={len(dataset)}')
+
+            drop_last = loader_spec.get('drop_last', True)
+            shuffle = loader_spec.get('shuffle', True)
+            self.loaders[split], self.loader_samplers[split] = self.make_distributed_loader(
+                dataset,
+                loader_spec.batch_size,
+                shuffle,
+                drop_last,
+                loader_spec.num_workers,
+                loader_spec.get('pin_memory', True)
+            )
+    
+    def make_model(self):
+        model = models.make(self.config.model)
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+        self.model = model.to(self.device)
+        self.log(f'Model: #params={utils.compute_num_params(model)}')
+
+    def make_optimizers(self):
+        self.optimizers = {'model': utils.make_optimizer(self.model.parameters(), self.config.optimizers['model'])}
+
+    def run_training(self):
+        config = self.config
+        max_iter = config['max_iter']
+        epoch_iter = config['epoch_iter']
+        assert max_iter % epoch_iter == 0
+        max_epoch = max_iter // epoch_iter
+
+        save_iter = config.get('save_iter')
+        if save_iter is not None:
+            assert save_iter % epoch_iter == 0
+            save_epoch = save_iter // epoch_iter
+            print('save_epoch', save_epoch)
+        else:
+            save_epoch = max_epoch + 1
+
+        eval_iter = config.get('eval_iter')
+        if eval_iter is not None:
+            assert eval_iter % epoch_iter == 0
+            eval_epoch = eval_iter // epoch_iter
+        else:
+            eval_epoch = max_epoch + 1
+
+        vis_iter = config.get('vis_iter')
+        if vis_iter is not None:
+            assert vis_iter % epoch_iter == 0
+            vis_epoch = vis_iter // epoch_iter
+        else:
+            vis_epoch = max_epoch + 1
+
+        if config.get('ckpt_select_metric') is not None:
+            m = config.ckpt_select_metric
+            self.ckpt_select_metric = m.name
+            self.ckpt_select_type = m.type
+            if m.type == 'min':
+                self.ckpt_select_v = 1e18
+            elif m.type == 'max':
+                self.ckpt_select_v = -1e18
+        else:
+            self.ckpt_select_metric = None
+            self.ckpt_select_v = 0
+
+        self.train_loader = self.loaders['train']
+        self.train_loader_sampler = self.loader_samplers['train']
+        self.train_loader_epoch = 0
+        self.train_loader_iter = None
+
+        self.iter = 0
+
+        if self.resume_ckpt is not None:
+            for _ in range(self.resume_ckpt['iter']):
+                self.iter += 1
+                self.at_train_iter_start()
+            self.ckpt_select_v = self.resume_ckpt['ckpt_select_v']
+            self.train_loader_epoch = self.resume_ckpt['train_loader_epoch']
+            self.train_loader_iter = None
+            self.resume_ckpt = None
+            self.log(f'Resumed iter status.')
+        
+        if config.get('vis_before_training', False):
+            self.visualize()
+
+        start_epoch = self.iter // epoch_iter + 1
+        epoch_timer = utils.EpochTimer(max_epoch - start_epoch + 1)
+
+        for epoch in range(start_epoch, max_epoch + 1):
+            self.log_buffer = [f'Epoch {epoch}']
+
+            for sampler in self.loader_samplers.values():
+                if sampler is not self.train_loader_sampler:
+                    sampler.set_epoch(epoch)
+
+            self.model_ddp.train()
+
+            ave_scalars = dict()
+            pbar = range(1, epoch_iter + 1)
+            if self.is_master and epoch == start_epoch:
+                pbar = tqdm(pbar, desc='train', leave=False)
+
+            t_data = 0
+            t_nondata = 0
+            t_before_data = time.time()
+            
+            for _ in pbar:
+                self.iter += 1
+                self.at_train_iter_start()
+
+                try:
+                    if self.train_loader_iter is None:
+                        raise StopIteration
+                    data = next(self.train_loader_iter)
+                except StopIteration:
+                    self.train_loader_epoch += 1
+                    self.train_loader_sampler.set_epoch(self.train_loader_epoch)
+                    self.train_loader_iter = iter(self.train_loader)
+                    data = next(self.train_loader_iter)
+                
+                t_after_data = time.time()
+                t_data += t_after_data - t_before_data
+                
+                for k, v in data.items():
+                    data[k] = v.to(self.device) if torch.is_tensor(v) else v
+                
+                ret = self.train_step(data)
+
+                
+                
+                t_before_data = time.time()
+                t_nondata += t_before_data - t_after_data
+
+                if self.is_master and epoch == start_epoch:
+                    pbar.set_description(desc=f'train: loss={ret["loss"]:.4f}')
+                
+                # save the model every 1000 iterations
+                if self.iter % 100 == 0:
+                    self.save_ckpt(f'ckpt-{self.iter}.pth')
+            
+            self.save_ckpt('ckpt-last.pth')
+
+            if epoch % save_epoch == 0 and epoch != max_epoch:
+                self.save_ckpt(f'ckpt-{self.iter}.pth')
+
+            if epoch % eval_epoch == 0:
+                with torch.no_grad():
+                    eval_ave_scalars = self.evaluate()
+                if self.ckpt_select_metric is not None:
+                    v = eval_ave_scalars[self.ckpt_select_metric].item()
+                    if ((self.ckpt_select_type == 'min' and v < self.ckpt_select_v) or
+                        (self.ckpt_select_type == 'max' and v > self.ckpt_select_v)):
+                        self.ckpt_select_v = v
+                        self.save_ckpt('ckpt-best.pth')
+
+            if epoch % vis_epoch == 0:
+                with torch.no_grad():
+                    self.visualize()
+                    
+    def at_train_iter_start(self):
+        pass
+
+    def train_step(self, data, bp=True):
+
+        print('data', data)
+        if self.config.get('autocast_bfloat16', False):
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                ret = self.model_ddp(data)
+        else:
+            ret = self.model_ddp(data)
+        
+        loss = ret.pop('loss')
+        ret['loss'] = loss.item()
+        if bp:
+            self.model_ddp.zero_grad()
+            loss.backward()
+            for o in self.optimizers.values():
+                o.step()
+        return ret
+
+    def evaluate(self):
+        self.model_ddp.eval()
+
+        ave_scalars = dict()
+        pbar = self.loaders['val']
+
+        for data in pbar:
+            for k, v in data.items():
+                data[k] = v.to(self.device) if torch.is_tensor(v) else v
+            
+            ret = self.train_step(data, bp=False)
+
+            bs = len(next(iter(data.values())))
+            for k, v in ret.items():
+                if ave_scalars.get(k) is None:
+                    ave_scalars[k] = utils.Averager()
+                ave_scalars[k].add(v, n=bs)
+
+        self.sync_ave_scalars(ave_scalars)
+
+        logtext = 'val:'
+        for k, v in ave_scalars.items():
+            logtext += f' {k}={v.item():.4f}'
+            self.log_scalar('val/' + k, v.item())
+        self.log_buffer.append(logtext)
+
+        return ave_scalars
+
+    def visualize(self):
+        pass
+
+    def save_ckpt(self, filename):
+        if self.is_master:
+            model_spec = copy.copy(self.config_dict['model'])
+            model_spec['sd'] = self.model.state_dict()
+            optimizers_spec = dict()
+            for name, spec in self.config_dict['optimizers'].items():
+                spec = copy.copy(spec)
+                spec['sd'] = self.optimizers[name].state_dict()
+                optimizers_spec[name] = spec
+            ckpt = {
+                'config': self.config_dict,
+                'model': model_spec,
+                'optimizers': optimizers_spec,
+                'iter': self.iter,
+                'train_loader_epoch': self.train_loader_epoch,
+                'ckpt_select_v': self.ckpt_select_v,
+            }
+            torch.save(ckpt, os.path.join(self.env['save_dir'], filename))
+        dist.barrier()
+
+    def sync_ave_scalars(self, ave_scalars):
+        keys = sorted(list(ave_scalars.keys()))
+        for k in keys:
+            if not k.startswith('_'):
+                v = ave_scalars[k]
+                vt = torch.tensor(v.item(), device=self.device)
+                dist.all_reduce(vt, op=dist.ReduceOp.SUM)
+                torch.cuda.synchronize()
+                ave_scalars[k].v = vt.item() / self.world_size
+                ave_scalars[k].n *= self.world_size
+
+    
+    def log_scalar(self, k, v):
+        if self.experiment:
+            self.experiment.log_metric(k, v, step=self.iter)
+
+    def log_image(self, k, v):
+        if self.experiment:
+            self.experiment.log_image(k, v, step=self.iter)

flowae/trainers/ldm_trainer.py

@@ -0,0 +1,443 @@
+import os
+import random
+
+import torch
+import torch.distributed as dist
+import torch_fidelity
+import torchvision
+from PIL import Image
+from torchvision import transforms
+
+import utils
+from utils.geometry import make_coord_scale_grid
+from .trainers import register
+from trainers.base_trainer import BaseTrainer
+from models.ldm.dac.audiotools import AudioSignal
+import soundfile as sf
+import numpy as np
+
+from models.ldm.dac.loss import (GANLoss, L1Loss, MelSpectrogramLoss,
+                         MultiScaleSTFTLoss, kl_loss)
+
+@register('ldm_trainer')
+class LDMTrainer(BaseTrainer):
+
+    def make_model(self):
+        super().make_model()
+        self.has_optimizer = dict()
+        for name, m in self.model.named_children():
+            self.log(f'  .{name} {utils.compute_num_params(m)}')
+
+    def make_optimizers(self):
+        self.optimizers = dict()
+        self.has_optimizer = dict()
+        for name, spec in self.config.optimizers.items():
+            self.optimizers[name] = utils.make_optimizer(self.model.get_parameters(name), spec)
+            self.has_optimizer[name] = True
+
+    def train_step(self, data, bp=True):
+        kwargs = {'has_optimizer': self.has_optimizer}
+        print('data', data.keys())
+        print('inp', data['inp'].shape)
+        print('gt', data['gt'].shape)
+       
+        if self.config.get('autocast_bfloat16', False):
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+                ret = self.model_ddp(data, mode='loss', **kwargs)
+        else:
+            ret = self.model_ddp(data, mode='loss', **kwargs)
+
+        loss = ret.pop('loss')
+        ret['loss'] = loss.item()
+        if bp:
+            self.model_ddp.zero_grad()
+            loss.backward()
+            for name, o in self.optimizers.items():
+                if name != 'disc':
+                    o.step()
+
+        self.model.update_ema()
+
+        return ret
+    
+    def evaluate(self):
+        self.model_ddp.eval()
+
+        ave_scalars = dict()
+        pbar = self.loaders['val']
+
+        for data in pbar:
+            for k, v in data.items():
+                data[k] = v.to(self.device) if torch.is_tensor(v) else v
+            
+            ret = self.train_step(data, bp=False)
+
+            bs = len(next(iter(data.values())))
+            for k, v in ret.items():
+                if ave_scalars.get(k) is None:
+                    ave_scalars[k] = utils.Averager()
+                ave_scalars[k].add(v, n=bs)
+        
+        self.sync_ave_scalars(ave_scalars)
+        
+        # Extra evaluation #
+        if self.config.get('evaluate_ae', False):
+            ave_scalars.update(self.evaluate_ae())
+        
+        if self.config.get('evaluate_zdm', False):
+            ema = self.config.get('evaluate_zdm_ema', True)
+            ave_scalars.update(self.evaluate_zdm(ema=ema))
+        # - #
+
+        logtext = 'val:'
+        for k, v in ave_scalars.items():
+            logtext += f' {k}={v.item():.4f}'
+            self.log_scalar('val/' + k, v.item())
+        self.log_buffer.append(logtext)
+        
+        return ave_scalars
+
+    def visualize(self):
+        self.model_ddp.eval()
+
+        if self.config.get('evaluate_ae', False):
+            # self.visualize_ae_fixset()
+            self.visualize_ae_random()
+        
+        if self.config.get('evaluate_zdm', False):
+            ema = self.config.get('evaluate_zdm_ema', True)
+            # self.visualize_zdm_fixset(ema=ema)
+            self.visualize_zdm_random(ema=ema)
+            # self.visualize_zdm_denoising(ema=ema)
+
+    def evaluate_ae(self):
+        max_samples = self.config.get('eval_ae_max_samples')
+        self.loader_samplers['eval_ae'].set_epoch(0)
+        
+        to_pil = transforms.ToPILImage()
+        psnr_value = utils.Averager()
+        cnt = 0
+
+        cache_gen_dir = os.path.join(self.env['save_dir'], 'cache', 'fid_gen')
+        cache_gt_dir = os.path.join(self.env['save_dir'], 'cache', 'fid_gt')
+        if self.is_master:
+            utils.ensure_path(cache_gen_dir, force_replace=True)
+            utils.ensure_path(cache_gt_dir, force_replace=True)
+        dist.barrier()
+
+        for data in self.loaders['eval_ae']:
+            for k, v in data.items():
+                data[k] = v.to(self.device) if torch.is_tensor(v) else v
+            
+            pred = self.model(data, mode='pred')
+            gt_patch = data['gt'][:, :3, ...]
+
+            pred = (pred * 0.5 + 0.5).clamp(0, 1)
+            gt_patch = (gt_patch * 0.5 + 0.5).clamp(0, 1)
+
+            # PSNR
+            mse = (pred - gt_patch).pow(2).mean(dim=[1, 2, 3])
+            psnr_value.add((-10 * torch.log10(mse)).mean().item())
+
+            # FID
+            for i in range(len(pred)):
+                idx = int(os.environ['RANK']) + cnt * int(os.environ['WORLD_SIZE'])
+                if max_samples is None or idx < max_samples: 
+                    to_pil(pred[i]).save(os.path.join(cache_gen_dir, f'{idx}.png'))
+                    to_pil(gt_patch[i]).save(os.path.join(cache_gt_dir, f'{idx}.png'))
+                cnt += 1
+        dist.barrier()
+
+        vt = torch.tensor(psnr_value.item(), device=self.device)
+        dist.all_reduce(vt, op=dist.ReduceOp.SUM)
+        torch.cuda.synchronize()
+        psnr_value = vt.item() / int(os.environ['WORLD_SIZE'])
+
+        if self.is_master:
+            metrics = torch_fidelity.calculate_metrics(
+                input1=cache_gen_dir,
+                input2=cache_gt_dir,
+                cuda=True,
+                fid=True,
+                verbose=False, 
+            )
+            prefix = 'eval_ae'
+            ret = {
+                f'{prefix}/PSNR': psnr_value,
+                f'{prefix}/FID': metrics['frechet_inception_distance'],
+            }
+        else:
+            ret = {}
+        dist.barrier()
+
+        ret = {k: utils.Averager(v) for k, v in ret.items()}
+        return ret
+
+    def evaluate_zdm(self, ema):
+        max_samples = self.config.get('eval_zdm_max_samples')
+        self.loader_samplers['eval_zdm'].set_epoch(0)
+        
+        to_pil = transforms.ToPILImage()
+        cnt = 0
+
+        cache_gen_dir = os.path.join(self.env['save_dir'], 'cache', 'fid_gen')
+        cache_gt_dir = os.path.join(self.env['save_dir'], 'cache', 'fid_gt')
+        if self.is_master:
+            utils.ensure_path(cache_gen_dir, force_replace=True)
+            utils.ensure_path(cache_gt_dir, force_replace=True)
+        dist.barrier()
+
+        for data in self.loaders['eval_zdm']:
+            for k, v in data.items():
+                data[k] = v.to(self.device) if torch.is_tensor(v) else v
+            
+            gt_patch = data['inp']
+
+            net_kwargs = dict()
+            uncond_net_kwargs = dict()
+            if self.model.zdm_class_cond is not None:
+                net_kwargs['class_labels'] = data['class_labels']
+
+                setting = self.config['visualize_zdm_setting']
+                uncond_net_kwargs['class_labels'] = setting['n_classes'] * torch.ones(
+                    len(data['class_labels']), dtype=torch.long, device=self.device)
+            
+            pred = self.model.generate_samples(
+                batch_size=gt_patch.shape[0],
+                n_steps=self.model.zdm_n_steps,
+                net_kwargs=net_kwargs,
+                uncond_net_kwargs=uncond_net_kwargs,
+                ema=ema
+            )
+
+            pred = (pred * 0.5 + 0.5).clamp(0, 1)
+            gt_patch = (gt_patch * 0.5 + 0.5).clamp(0, 1)
+
+            # FID
+            for i in range(len(pred)):
+                idx = int(os.environ['RANK']) + cnt * int(os.environ['WORLD_SIZE'])
+                if max_samples is None or idx < max_samples: 
+                    to_pil(pred[i]).save(os.path.join(cache_gen_dir, f'{idx}.png'))
+                    to_pil(gt_patch[i]).save(os.path.join(cache_gt_dir, f'{idx}.png'))
+                cnt += 1
+        dist.barrier()
+
+        if self.is_master:
+            metrics = torch_fidelity.calculate_metrics(
+                input1=cache_gen_dir,
+                input2=cache_gt_dir,
+                cuda=True,
+                fid=True,
+                verbose=False, 
+            )
+            prefix = 'eval_zdm' + ('_ema' if ema else '')
+            ret = {
+                f'{prefix}/FID': metrics['frechet_inception_distance'],
+            }
+        else:
+            ret = {}
+        dist.barrier()
+
+        ret = {k: utils.Averager(v) for k, v in ret.items()}
+        return ret
+    
+    def visualize_ae_fixset(self):
+        if self.config.get('visualize_ae_dir') is None:
+            return
+        to_tensor = transforms.ToTensor()
+        if self.is_master:
+            files = sorted(os.listdir(self.config['visualize_ae_dir']))
+            vis_images = []
+            
+            for f in files:
+                image = Image.open(os.path.join(self.config['visualize_ae_dir'], f)).convert('RGB')
+                x = to_tensor(image).unsqueeze(0).to(self.device)
+                x = (x - 0.5) / 0.5
+                gt_dummy = torch.zeros(x.shape[0], 7, x.shape[2], x.shape[3], device=self.device)
+                
+                pred1 = self.model({'inp': x, 'gt': gt_dummy}, mode='pred')
+                pred2 = self.model({'inp': x, 'gt': gt_dummy}, mode='pred')
+                vis_images.extend([x, pred1, pred2])
+            
+            vis_images = torch.cat(vis_images, dim=0)
+            vis_images = torchvision.utils.make_grid(vis_images, normalize=True, value_range=(-1, 1), nrow=6)
+            self.log_image('vis_ae_fixset', vis_images)
+        dist.barrier()
+    
+    def visualize_ae_random(self):
+        if self.is_master:
+            idx_list = list(range(len(self.datasets['eval_ae'])))
+            random.shuffle(idx_list)
+            n_samples = self.config['visualize_ae_random_n_samples']
+            vis_images = []
+            
+            for idx in idx_list[:n_samples]:
+                data = self.datasets['eval_ae'][idx]
+                for k, v in data.items():
+                    data[k] = v.unsqueeze(0).to(self.device) if torch.is_tensor(v) else v
+                
+                pred1 = self.model(data, mode='pred')
+                pred2 = self.model(data, mode='pred')
+                gt_patch = data['gt'][:, :3, ...]
+                vis_images.extend([gt_patch, pred1, pred2])
+            
+            vis_images = torch.cat(vis_images, dim=0)
+            vis_images = torchvision.utils.make_grid(vis_images, normalize=True, value_range=(-1, 1), nrow=6)
+            self.log_image('vis_ae_random', vis_images)
+        dist.barrier()
+    
+    def visualize_zdm_fixset(self, ema):
+        if self.is_master:
+            vis_file = torch.load(self.config['visualize_zdm_file'], map_location='cpu')
+            for k, v in vis_file.items():
+                vis_file[k] = v.to(self.device) if torch.is_tensor(v) else v
+            n_samples = len(vis_file['noise'])
+            
+            batch_size = self.config.get('visualize_zdm_batch_size', 1)
+            guidance_list = [1.0] + self.config.get('visualize_zdm_guidance_list', [])
+
+            vis_images = []
+            
+            for i in range(0, n_samples, batch_size):
+                cur_batch_size = min(batch_size, n_samples - i)
+
+                net_kwargs = dict()
+                uncond_net_kwargs = dict()
+                if self.config.get('visualize_zdm_setting') is not None:
+                    setting = self.config['visualize_zdm_setting']
+                    if setting['name'] == 'class':
+                        net_kwargs['class_labels'] = vis_file['class_labels'][i:i + cur_batch_size]
+                        uncond_net_kwargs['class_labels'] = setting['n_classes'] * torch.ones(
+                            cur_batch_size, dtype=torch.long, device=self.device)
+                    else:
+                        raise NotImplementedError
+
+                for guidance in guidance_list:
+                    pred = self.model.generate_samples(
+                        batch_size=cur_batch_size,
+                        n_steps=self.model.zdm_n_steps,
+                        net_kwargs=net_kwargs,
+                        uncond_net_kwargs=uncond_net_kwargs,
+                        ema=ema,
+                        guidance=guidance,
+                        noise=vis_file['noise'][i:i + cur_batch_size],
+                    )
+                    vis_images.append(pred)
+
+            vis_images = torch.cat(vis_images, dim=0)
+            vis_images = torchvision.utils.make_grid(vis_images, normalize=True, value_range=(-1, 1), nrow=batch_size)
+            name = 'vis_zdm_fixset'
+            name += '_ema' if ema else ''
+            name += '_cfg' + str(guidance_list[1:])[1:-1] if len(guidance_list) > 1 else ''
+            self.log_image(name, vis_images)
+        dist.barrier()
+
+    def visualize_zdm_random(self, ema):
+        n_samples = self.config['visualize_zdm_random_n_samples']
+        batch_size = self.config.get('visualize_zdm_batch_size', 1)
+        guidance_list = [1.0] + self.config.get('visualize_zdm_guidance_list', [])
+
+        vis_images = []
+
+        if self.is_master:
+            for i in range(0, n_samples, batch_size):
+                cur_batch_size = min(batch_size, n_samples - i)
+
+                net_kwargs = dict()
+                uncond_net_kwargs = dict()
+                if self.config.get('visualize_zdm_setting') is not None:
+                    setting = self.config['visualize_zdm_setting']
+                    if setting['name'] == 'class':
+                        net_kwargs['class_labels'] = torch.randint(
+                            setting['n_classes'], size=(cur_batch_size,), device=self.device)
+                        uncond_net_kwargs['class_labels'] = setting['n_classes'] * torch.ones(
+                            cur_batch_size, dtype=torch.long, device=self.device)
+                    else:
+                        raise NotImplementedError
+
+                for guidance in guidance_list:
+                    pred = self.model.generate_samples(
+                        batch_size=cur_batch_size,
+                        n_steps=self.model.zdm_n_steps,
+                        net_kwargs=net_kwargs,
+                        uncond_net_kwargs=uncond_net_kwargs,
+                        ema=ema,
+                        guidance=guidance,
+                    )
+                    vis_images.append(pred)
+
+            vis_images = torch.cat(vis_images, dim=0)
+            vis_images = torchvision.utils.make_grid(vis_images, normalize=True, value_range=(-1, 1), nrow=batch_size)
+            name = 'vis_zdm_random'
+            name += '_ema' if ema else ''
+            name += '_cfg' + str(guidance_list[1:])[1:-1] if len(guidance_list) > 1 else ''
+            self.log_image(name, vis_images)
+        dist.barrier()
+
+    def visualize_zdm_denoising(self, ema, n_selected_timesteps=5):
+        if self.is_master:
+            vis_file = torch.load(self.config['visualize_zdm_denoising_file'], map_location='cpu')
+
+            vis_images = []
+            
+            for i in range(len(vis_file['inp'])):
+                x = (
+                    vis_file['inp'][i]
+                    .to(self.device)
+                    .unsqueeze(0)
+                    .expand(n_selected_timesteps, -1, -1, -1)
+                )
+                
+                z = self.model.encode(x)
+                z = self.model.normalize_for_zdm(z)
+                t = torch.linspace(0, 1, n_selected_timesteps + 1, device=self.device)[1:]
+                noise = (
+                    vis_file['noise'][i]
+                    .to(self.device)
+                    .unsqueeze(0)
+                    .expand(n_selected_timesteps, -1, -1, -1)
+                )
+                z_t, _ = self.model.zdm_diffusion.add_noise(z, t, noise=noise)
+
+                # Visualize noisy latents
+                zp = self.model.denormalize_for_zdm(z_t)
+                z_dec = self.model.decode(zp)
+                coord, scale = make_coord_scale_grid(x.shape[-2:], device=self.device, batch_size=n_selected_timesteps)
+                coord = coord.permute(0, 3, 1, 2)
+                scale = scale.permute(0, 3, 1, 2)
+                x_out = self.model.render(z_dec, coord, scale)
+                vis_images.append(x_out)
+                
+                # Generate denoised latents
+                net = self.model.zdm_net_ema if ema else self.model.zdm_net
+                net_kwargs = dict()
+                if self.config.get('visualize_zdm_setting') is not None:
+                    setting = self.config['visualize_zdm_setting']
+                    if setting['name'] == 'class':
+                        net_kwargs['class_labels'] = (
+                            vis_file['class_labels'][i]
+                            .to(self.device)
+                            .unsqueeze(0)
+                            .expand(n_selected_timesteps)
+                        )
+                    else:
+                        raise NotImplementedError
+                pred = self.model.zdm_diffusion.get_prediction(net, z_t, t, net_kwargs=net_kwargs)
+                zp = []
+                for j in range(len(pred)):
+                    zp.append(self.model.zdm_diffusion.convert_sample_prediction(z_t[j], float(t[j]), pred[j]))
+                zp = torch.stack(zp, dim=0)
+
+                # Visualize denoised latents
+                zp = self.model.denormalize_for_zdm(zp)
+                z_dec = self.model.decode(zp)
+                coord, scale = make_coord_scale_grid(x.shape[-2:], device=self.device, batch_size=n_selected_timesteps)
+                coord = coord.permute(0, 3, 1, 2)
+                scale = scale.permute(0, 3, 1, 2)
+                x_out = self.model.render(z_dec, coord, scale)
+                vis_images.append(x_out)
+            
+            vis_images = torch.cat(vis_images, dim=0)
+            vis_images = torchvision.utils.make_grid(vis_images, normalize=True, value_range=(-1, 1), nrow=n_selected_timesteps)
+            self.log_image('vis_zdm' + ('_ema' if ema else '') + '_denoising', vis_images)
+        dist.barrier() 

flowae/trainers/trainers.py

@@ -0,0 +1,8 @@
+trainers_dict = dict()
+
+
+def register(name):
+    def decorator(cls):
+        trainers_dict[name] = cls
+        return cls
+    return decorator