trellis/trainers/flow_matching/discrete_diffusion.py

from typing import *
import copy
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np
from easydict import EasyDict as edict
import os
import torch.distributed as dist
import matplotlib.pyplot as plt 
import io
from PIL import Image
import wandb
import re
from pytorch3d.ops import cubify
from trimesh import Trimesh
from einops import rearrange
from ..basic import BasicTrainer
from contextlib import nullcontext
from ...pipelines import samplers 
from ...utils.general_utils import dict_reduce
from .mixins.classifier_free_guidance import DiscreteClassifierFreeGuidanceMixin
from .mixins.text_conditioned import TextConditionedMixin
from .mixins.image_conditioned import ImageConditionedMixin
from ...utils.data_utils import recursive_to_device
from tqdm import tqdm
import json
from ...utils.mask_utils import generate_block_masks_ragged


# This script is build upon TRELLIS https://github.com/microsoft/TRELLIS and DUO https://github.com/s-sahoo/duo

"""
This is the trainer of uniform state discrete diffusion model (USDM).

"""

class LogLinear(torch.nn.Module):
  def __init__(self, eps=1e-3):
    super().__init__()
    self.eps = eps  # To be consistent with SEDD: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion/blob/0605786da5ccb5747545e26d66fdf477187598b6/noise_lib.py#L56

  def forward(self, t):
    t = (1 - self.eps) * t
    alpha_t = 1 - t 
    dalpha_t = - (1 - self.eps)
    return dalpha_t, alpha_t
  

class DiscreteDiffusionTrainer(BasicTrainer):
    """
    Trainer for diffusion model with USDM objective.
    
    Args:
        models (dict[str, nn.Module]): Models to train.
        dataset (torch.utils.data.Dataset): Dataset.
        output_dir (str): Output directory.
        load_dir (str): Load directory.
        step (int): Step to load.
        batch_size (int): Batch size.
        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
        batch_split (int): Split batch with gradient accumulation.
        max_steps (int): Max steps.
        optimizer (dict): Optimizer config.
        lr_scheduler (dict): Learning rate scheduler config.
        elastic (dict): Elastic memory management config.
        grad_clip (float or dict): Gradient clip config.
        ema_rate (float or list): Exponential moving average rates.
        fp16_mode (str): FP16 mode.
            - None: No FP16.
            - 'inflat_all': Hold a inflated fp32 master param for all params.
            - 'amp': Automatic mixed precision.
        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
        finetune_ckpt (dict): Finetune checkpoint.
        log_param_stats (bool): Log parameter stats.
        i_print (int): Print interval.
        i_log (int): Log interval.
        i_sample (int): Sample interval.
        i_save (int): Save interval.
        i_ddpcheck (int): DDP check interval.

        t_schedule (dict): Time schedule for flow matching.
        sigma_min (float): Minimum noise level.
    """
    def __init__(
        self,
        *args,
        t_schedule: dict = {
            'name': 'logitNormal',
            'args': {
                'mean': 0.0,
                'std': 1.0,
            }
        },
        sigma_min: float = 1e-5,
        vocab_size: int = 2,
        multi_mask_ratio: float = 0.0,
        use_cross_entropy: bool = False,
        reweight_elbo: bool=False,
        **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.t_schedule = t_schedule
        self.noise = LogLinear(sigma_min)
        self.vocab_size = vocab_size
        self.sigma_min = sigma_min
        self.multi_mask_ratio = multi_mask_ratio
        self.use_cross_entropy = use_cross_entropy  
        self.reweight_elbo = reweight_elbo

    def diffuse(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None, multi_mask: Optional[bool]=False) -> torch.Tensor:
        """
        Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). 
        If multi_mask is True, the perturbation pattern will be blocked-structured (BSP). Actual_t is returned which calculates the proportion of noised tokens.

        Args:
            x_0: The [N x C x ...] tensor of noiseless inputs.
            t: The [N] tensor of diffusion steps [0-1].
            noise: If specified, use this noise instead of generating new noise.

        Returns:
            x_t: the noisy version of x_0 under timestep t, same shape as x_0.
            actual_t: Optional. Only returned when multi_mask is True. The actual t used for each sample, same shape as t.
        """
        if noise is None:
            noise = torch.randint(0, self.vocab_size, x_0.shape, device=x_0.device)
        assert noise.shape == x_0.shape, "noise must have same shape as x_0"

        if not multi_mask:
            _, alpha_t = self.noise(t)
            alpha_t = alpha_t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
            move_indices = torch.rand(x_0.shape, device=x_0.device) < 1 - alpha_t

            x_t = torch.where(move_indices, noise, x_0)
            return x_t
        else:
            # only works for 64
            lengths = torch.tensor([1,2,4,8,16,32,48,54,60], device = x_0.device)
            move_indices, actual_t = generate_block_masks_ragged(
                B=x_0.shape[0], N=64, dim=3, lengths=lengths, t=1-self.noise(t)[1], device=x_0.device, round_from=9
            )
            move_indices = move_indices.unsqueeze(1)


            # ------------------------- perturb t_new portion inside blocks ---------------------------
            t_new = self.sample_t(x_0.shape[0]).to(x_0.device).float()
            _, alpha_t = self.noise(t_new)
            alpha_t = alpha_t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
            move_indices_new = torch.rand(x_0.shape, device=x_0.device) < 1 - alpha_t
            move_indices_use = move_indices & move_indices_new 
            actual_t = t_new*actual_t

            x_t = torch.where(move_indices_use, noise, x_0)
            return x_t, actual_t



    def reverse_diffuse(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
        """
        Get original image from noisy version under timestep t.
        Here assume noise is the sample of x_0 prediction.
        """
        # raise NotImplementedError
        assert noise.dtype in (torch.int32, torch.int64), "noise must be integer type"
        return noise

    def get_v(self, x_0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Compute the velocity of the diffusion process at time t.
        """
        raise NotImplementedError # to be consistent with flow matching trainer

    def get_cond(self, cond, **kwargs):
        """
        Get the conditioning data.
        """
        return cond
    
    def get_inference_cond(self, cond, **kwargs):
        """
        Get the conditioning data for inference.
        """
        return {'cond': cond, **kwargs}

    def get_sampler(self, **kwargs) -> samplers.USDMAncestralSampler:
        """
        Get the sampler for the diffusion process.
        """
        return samplers.USDMAncestralSampler(self.sigma_min, self.vocab_size)
    
    def vis_cond(self, **kwargs):
        """
        Visualize the conditioning data.
        """
        return {}

    def sample_beta_uniform_mix(self, n, p_beta=0.98):
        # vectorized mixture sampling: with probability p_beta draw from beta_dist, else uniform(0,1)
        # more t close to gt is sampled
        beta_dist =  torch.distributions.Beta(3,1)
        mask = torch.rand(n) < p_beta
        n_beta = int(mask.sum().item())
        out = torch.empty(n)
        if n_beta > 0:
            out[mask] = beta_dist.sample((n_beta,))
        if n - n_beta > 0:
            out[~mask] = torch.rand(n - n_beta)
        return out
    
    def sample_logitnormal_uniform_mix(self, n, p_logitnormal=0.90, mean=1.0, std=1.0):
        # vectorized mixture sampling: with probability p_logitnormal draw from logitnormal, else uniform(0,1)
        
        mask = torch.rand(n) < p_logitnormal
        n_lognormal = int(mask.sum().item())
        out = torch.empty(n)
        if n_lognormal > 0:
            out[mask] = torch.sigmoid(torch.randn(n_lognormal) * std + mean)
        if n - n_lognormal > 0:
            out[~mask] = torch.rand(n - n_lognormal)
        return out


    def sample_t(self, batch_size: int) -> torch.Tensor:
        """
        Sample timesteps.
        """

        if self.t_schedule['name'] == 'uniform':
            # do antithetic sampling for t
            _eps_t = torch.rand(batch_size)
            offset = torch.arange(batch_size) / batch_size
            _eps_t = (_eps_t + offset) % 1
            t = _eps_t
        elif self.t_schedule['name'] == 'logitNormal':
            mean = self.t_schedule['args']['mean']
            std = self.t_schedule['args']['std']
            t = torch.sigmoid(torch.randn(batch_size) * std + mean)
        elif self.t_schedule['name'] == 'beta_uniform':
            t = self.sample_beta_uniform_mix(batch_size, self.t_schedule['args']['p_beta'])
        elif self.t_schedule['name'] == 'logitNormal_uniform':
            t = self.sample_logitnormal_uniform_mix(batch_size, self.t_schedule['args']['p_logitnormal'], self.t_schedule['args']['mean'], self.t_schedule['args']['std'])
        else:
            raise ValueError(f"Unknown t_schedule: {self.t_schedule['name']}")
        return t

    def nll_per_token(self, log_x_theta, xt, x0, alpha_t,
                    dalpha_t,):

        """
        Calculate the (variant of) negative log likelihood per token, used as diffusion loss. Adapted from https://github.com/s-sahoo/duo.
        
        Args:
            log_x_theta: The [N, L, vocab_size] tensor of log probabilities predicted by the model.
            xt: [N, L] int tensor of noised inputs at time t.
            x0:  [N, L] int tensor of original inputs.
            alpha_t: [N,1] tensor of alpha_t.
            d_alpha_t: scalar or [N,1] tensor of the time derivative of alpha_t.

        """

        assert alpha_t.ndim == 2
        assert x0.ndim == 2
        assert xt.ndim == 2
        assert not torch.is_tensor(dalpha_t) or dalpha_t.ndim == 2
        if not self.use_cross_entropy:
            if self.reweight_elbo:
                # reweight elbo to avoid large variance components in loss
                alpha_t = torch.clamp(alpha_t, max=0.6)

            x_reconst = log_x_theta.exp()
            x_bar_theta = self.vocab_size * alpha_t[
                :, :, None] * x_reconst + 1 - alpha_t[:, :, None]
            coeff = dalpha_t / (self.vocab_size * alpha_t)
            x_eq_xt = (x0 == xt).float()
            x_neq_xt = 1 - x_eq_xt
            xbar_xt = (1 - alpha_t) + self.vocab_size * alpha_t * x_eq_xt
            xbar_theta_xt = torch.gather(
            x_bar_theta, -1, xt.unsqueeze(-1)).squeeze(-1)
            xbar_theta_x = torch.gather(
            x_bar_theta, -1, x0.unsqueeze(-1)).squeeze(-1)
            term1 = self.vocab_size * (1 / xbar_xt
                                        - 1 / xbar_theta_xt)
            
            const = (1 - alpha_t) / (self.vocab_size * alpha_t
                                    + 1 - alpha_t)
            term2_coefs = x_eq_xt * const + x_neq_xt
            term2_offset = ((self.vocab_size - 1) * const * x_eq_xt
                            - (1 / const) * x_neq_xt) * const.log()
            term2_theta = - term2_coefs * (
            x_bar_theta.log().sum(-1)
            - self.vocab_size * xbar_theta_xt.log())
            term2_theta = (
            term2_theta
            - self.vocab_size * alpha_t / (1 - alpha_t) * (
                xbar_theta_x.log() - xbar_theta_xt.log()) * x_neq_xt)
            term2 = term2_theta + term2_offset
            diffusion_loss = coeff * (term1 - term2)
            assert diffusion_loss.ndim == 2
        else:
            # use cross entropy loss
            diffusion_loss = F.cross_entropy(log_x_theta.permute(0,2,1), x0, reduction='none')
        return diffusion_loss   

    def training_losses(
        self,
        x_0: torch.Tensor,
        cond=None,
        **kwargs
    ) -> Tuple[Dict, Dict]:
        """
        Compute training losses for a single timestep.

        Args:
            x_0: The [N x C x ...] tensor of noiseless inputs.
            cond: The [N x ...] tensor of additional conditions.
            kwargs: Additional arguments to pass to the backbone.

        Returns:
            a dict with the key "loss" containing a tensor of shape [N].
            may also contain other keys for different terms.
        """
        bsz = x_0.shape[0]
        noise = torch.randint(0, self.vocab_size, x_0.shape, device=x_0.device)
        if self.multi_mask_ratio > 0.0:
            # for multi_mask, we use t as uniform by default
            num_multi_mask = (torch.rand(bsz, device=x_0.device) < self.multi_mask_ratio).sum()
            if num_multi_mask > 0:
                t_multi = torch.rand(num_multi_mask, device=x_0.device)
                x_t_multi_mask, t_multi_actual = self.diffuse(x_0[:num_multi_mask], t_multi, noise[:num_multi_mask], multi_mask = True)
                t_origin_mask = self.sample_t(bsz - num_multi_mask).to(x_0.device).float()
                x_t_origin = self.diffuse(x_0[num_multi_mask:], t_origin_mask, noise[num_multi_mask:], multi_mask=False)

                x_t = torch.cat([x_t_multi_mask, x_t_origin], dim=0)
                t = torch.cat([t_multi_actual, t_origin_mask], dim=0)
            else:
                t = self.sample_t(bsz).to(x_0.device).float()
                x_t = self.diffuse(x_0, t, noise=noise)
        else:
            t = self.sample_t(bsz).to(x_0.device).float()
            x_t = self.diffuse(x_0, t, noise=noise)
        cond = self.get_cond(cond, **kwargs)
        
        pred_logits = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
        pred_logits_normalized = F.log_softmax(pred_logits, dim=-1) # shape [B,...,vocab_size]

        
        dalpha_t, alpha_t = self.noise(t)
        dalpha_t = torch.ones_like(alpha_t) * dalpha_t  
        # reshape to (B, L) while computing loss

        loss_nll = self.nll_per_token(
            rearrange(pred_logits_normalized, 'b d h w v -> b (d h w) v'),
            rearrange(x_t, 'b 1 d h w -> b (d h w)'),
            rearrange(x_0, 'b 1 d h w -> b (d h w)'),
            alpha_t.unsqueeze(-1),
            dalpha_t.unsqueeze(-1),
        )   
        terms = edict()
        terms["nll"] = loss_nll.mean()
        terms["loss"] = terms["nll"]

        # log loss with time bins, be consistent with continous
        with torch.no_grad():
            nll_per_instance = loss_nll.detach().mean(dim=-1)
        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
        for i in range(10):
            if (time_bin == i).sum() != 0:
                terms[f"bin_{i}"] = {"nll": nll_per_instance[time_bin == i].mean()}

        return terms, {}
    
    @torch.no_grad()
    def run_snapshot(
        self,
        num_samples: int,
        batch_size: int,
        verbose: bool = False,
    ) -> Dict:
        dataloader = DataLoader(
            copy.deepcopy(self.dataset),
            batch_size=batch_size,
            shuffle=True,
            num_workers=0,
            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
        )

        # inference
        sampler = self.get_sampler() # overwrite by classifier_free_guidance mixin
        sample_gt = []
        sample = []
        uids: List[str] = []
        cond_vis = []
        # persistent iterator to avoid repeating the first batch every time
        data_iter = iter(dataloader)
        # choose autocast context matching training precision
        if self.fp16_mode == 'amp':
            autocast_ctx = torch.autocast(device_type='cuda', dtype=torch.float16)
        elif self.fp16_mode == 'bf16':
            autocast_ctx = torch.autocast(device_type='cuda', dtype=torch.bfloat16)
        else:
            autocast_ctx = nullcontext()

        for i in range(0, num_samples, batch_size):
            batch = min(batch_size, num_samples - i)
            data = next(data_iter)
            data = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
            # collect uids if present
            if 'uid' in data:
                # ensure this is a list of strings
                batch_uids = data['uid']
                if isinstance(batch_uids, (list, tuple)):
                    uids.extend([str(x) for x in batch_uids])
                else:
                    # fallback to string, one per item expected
                    uids.extend([str(batch_uids)] * batch)

            noise = torch.randint(0, self.vocab_size, data['x_0'].shape, device=data['x_0'].device)
            sample_gt.append(data['x_0'])
            cond_vis.append(self.vis_cond(**data))
            del data['x_0']
            args = self.get_inference_cond(**data)
            # ensure inference runs under autocast so FlashAttention sees half/bfloat tensors
            with autocast_ctx:
                res = sampler.sample(
                    self.models['denoiser'],
                    noise=noise,
                    **args,
                    steps=256, cfg_strength=2.0, verbose=verbose, rescale_t='cosine',
                ) # cfg here not used
            sample.append(res.samples)

        sample_gt = torch.cat(sample_gt, dim=0)
        sample = torch.cat(sample, dim=0)
        assert sample.ndim == 5, "data should be [B, 1, D, H, W]"
        sample_dict = {
            'sample_gt': {'value': sample_gt, 'type': 'sample_gt'},
            'sample': {'value': sample, 'type': 'sample'},
        }
        # check in cond_vis is text or image

        if cond_vis[0]['cond_vis']['type'] == 'image':
            sample_dict.update(dict_reduce(cond_vis, None, {
                'value': lambda x: torch.cat(x, dim=0),
                'type': lambda x: x[0],
            }))
        elif cond_vis[0]['cond_vis']['type'] == 'text':
            sample_dict.update(dict_reduce(cond_vis, None, {
                'value': lambda x: x,
                'type': lambda x: x[0],
            }))
        if len(uids) > 0:
            sample_dict['uids'] = {'value': uids, 'type': 'uid_list'}
        
        return sample_dict


    @torch.no_grad()
    def slice_voxel(self, voxels, index = -1):
        """
        Slice voxels for visualization.
        Args:
            voxels: (B, D, H, W) tensor
            index: index to slice
        """
        pils = []
        assert voxels.ndim == 5, "voxels should be [B, 1, D, H, W]"
        voxels = voxels.squeeze(1)
        if index == -1:
            index = voxels.shape[-1] // 2 - 1 
        if voxels.shape[-1] == 1:
            num_samples = voxels.shape[0]
            fig, axes = plt.subplots(1, min(20,num_samples), figsize=(2*num_samples, 3))
            for i in range(min(20,num_samples)):
                img = voxels[i].reshape(28, 28).cpu().numpy()
                axes[i].imshow(img, cmap='gray', vmin=0, vmax=self.vocab_size-1)
                axes[i].axis('off')
            buf = io.BytesIO()
            fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0.1)
            plt.close(fig)              # important to free memory
            buf.seek(0)
            pils.append(Image.open(buf).convert("RGB"))
        else:
            for i in range(voxels.shape[0]):
                slice_x = voxels[i, :, index, :].cpu().numpy()
                slice_y = voxels[i, :, :, index].cpu().numpy()
                slice_z = voxels[i, index, :, :].cpu().numpy()

                fig, axs = plt.subplots(1, 3, figsize=(15, 5))
                axs[0].imshow(slice_x, cmap='gray')
                axs[0].set_title('Slice along X-axis')
                axs[0].axis('off')
                axs[1].imshow(slice_y, cmap='gray')
                axs[1].set_title('Slice along Y-axis')
                axs[1].axis('off')
                axs[2].imshow(slice_z, cmap='gray')
                axs[2].set_title('Slice along Z-axis')
                axs[2].axis('off')

                buf = io.BytesIO()
                fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0.1)
                plt.close(fig)              # important to free memory
                buf.seek(0)
                pils.append(Image.open(buf).convert("RGB"))
        return pils


    @torch.no_grad()
    def snapshot(self, suffix=None, num_samples=4, batch_size=4, verbose=False):
        """
        Sample images from the model.
        """
        if self.is_master:
            print(f'\nSampling {num_samples} images...', end='')

        if suffix is None:
            suffix = f'step{self.step:07d}'

        # Assign tasks
        num_samples_per_process = int(np.ceil(num_samples / self.world_size))
        samples = self.run_snapshot(num_samples_per_process, batch_size=batch_size, verbose=verbose) 


        # Gather results
        if self.world_size > 1:
            # iterate over a copy of keys as we'll mutate values
            for key in list(samples.keys()):
                # Special handling for uid lists (non-tensor)
                if samples[key]['type'] == 'uid_list':
                    if self.is_master:
                        all_uids: Optional[List[List[str]]] = [None for _ in range(self.world_size)]
                    else:
                        all_uids = None
                    dist.gather_object(samples[key]['value'], all_uids, dst=0)
                    if self.is_master:
                        # flatten and trim
                        flat = []
                        for part in all_uids:
                            if part is not None:
                                flat.extend(part)
                        samples[key]['value'] = flat[:num_samples]
                    continue

                if samples[key]['type'] == 'text':
                    if self.is_master:
                        all_uids: Optional[List[List[str]]] = [None for _ in range(self.world_size)]
                    else:
                        all_uids = None
                    dist.gather_object(samples[key]['value'], all_uids, dst=0)
                    if self.is_master:
                        # flatten and trim
                        flat = []
                        for part in all_uids:
                            if part is not None:
                                flat.extend(part)
                        samples[key]['value'] = flat[:num_samples]
                    continue

                samples[key]['value'] = samples[key]['value'].contiguous()
                if self.is_master:
                    all_images = [torch.empty_like(samples[key]['value']) for _ in range(self.world_size)]
                else:
                    all_images = None
                dist.gather(samples[key]['value'], all_images, dst=0)
                if self.is_master:
                    samples[key]['value'] = torch.cat(all_images, dim=0)[:num_samples]

        # Save images
        if self.is_master:
            # helper to sanitize uid for filenames
            def _sanitize_uid(s: str) -> str:
                s = str(s)
                s = s.strip()
                # replace all non safe filename chars with '_'
                return re.sub(r"[^A-Za-z0-9._-]+", "_", s)[:128]

            # extracted uids if available
            uid_list: Optional[List[str]] = None
            if 'uids' in samples and isinstance(samples['uids'].get('value', None), list):
                uid_list = samples['uids']['value']

            text_list: Optional[List[str]] = None

            if samples['cond_vis']['type'] == 'text':
                text_list = samples['cond_vis']['value']

            os.makedirs(os.path.join(self.output_dir, 'samples', suffix), exist_ok=True)
            os.makedirs(os.path.join(self.output_dir, 'gt', suffix), exist_ok=True)
            # dump all uids to json under samples/<suffix>/uids.json
            if uid_list is not None:
                uid_json_path = os.path.join(self.output_dir, 'samples', suffix, 'uids.json')
                try:
                    with open(uid_json_path, 'w', encoding='utf-8') as f:
                        import json as _json
                        _json.dump(uid_list, f, ensure_ascii=False, indent=2)
                except Exception as e:
                    print(f"[snapshot] Failed to write uids.json: {e}")

            if text_list is not None:
                text_json_path = os.path.join(self.output_dir, 'samples', suffix, 'text.json')
                try:
                    with open(text_json_path, 'w', encoding='utf-8') as f:
                        import json as _json
                        _json.dump(text_list, f, ensure_ascii=False, indent=2)
                except Exception as e:
                    print(f"[snapshot] Failed to write uids.json: {e}")
                    
            for key in samples.keys():

                if samples[key]['type'] == 'image':
                    wandb.log({'samples/condition': [wandb.Image(f) for f in samples[key]['value']]}, step=self.step)
                elif samples[key]['type'] == 'sample_gt':
                    # assume samples
                    # save voxel slices
                    pils = self.slice_voxel(samples[key]['value'])

                    wandb.log({'samples/voxel_slices_gt': [wandb.Image(f) for f in pils]}, step=self.step)
                    # log voxels
                    if samples[key]['value'].shape[-1] == 1:
                        continue
                    meshes_gt = []
                    cubified_mesh = cubify(samples[key]['value'].squeeze(1).float(), thresh=0.5, align="center")
                    for i, (verts, faces) in enumerate(zip(cubified_mesh.verts_list(), cubified_mesh.faces_list())):
                        uid_tag = _sanitize_uid(uid_list[i]) if uid_list is not None and i < len(uid_list) else str(i)
                        verts_np = verts.cpu().numpy()
                        faces_np = faces.cpu().numpy()
                        # export to obj
                        mesh = Trimesh(vertices=verts_np, faces=faces_np, process=False)
                        path = os.path.join(self.output_dir, 'gt', suffix, f'{key}_{suffix}_{uid_tag}.obj')
                        mesh.export(path)
                        meshes_gt.append(path)
                elif samples[key]['type'] == 'sample':
                    # assume samples
                    # save voxel slices
                    pils = self.slice_voxel(samples[key]['value'])
                    for i, pil in enumerate(pils):
                        uid_tag = _sanitize_uid(uid_list[i]) if uid_list is not None and i < len(uid_list) else str(i)
                        pil.save(os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}_{uid_tag}.png'))
                    
                    wandb.log({'samples/voxel_slices_generated': [wandb.Image(f) for f in pils]}, step=self.step)
                    # log voxels
                    if samples[key]['value'].shape[-1] == 1:
                        continue
                    meshes_sample = []
                    cubified_mesh = cubify(samples[key]['value'].squeeze(1).float(), thresh=0.5, align="center")
                    for i, (verts, faces) in enumerate(zip(cubified_mesh.verts_list(), cubified_mesh.faces_list())):
                        uid_tag = _sanitize_uid(uid_list[i]) if uid_list is not None and i < len(uid_list) else str(i)
                        verts_np = verts.cpu().numpy()
                        faces_np = faces.cpu().numpy()
                        # export to obj
                        mesh = Trimesh(vertices=verts_np, faces=faces_np, process=False)
                        path = os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}_{uid_tag}.obj')
                        mesh.export(path)
                        meshes_sample.append(path)
                        # save the corresponding voxels as npz
                        voxel_path = os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}_{uid_tag}.npz')
                        np.savez_compressed(voxel_path, voxels=samples[key]['value'][i].cpu().numpy())

                    
 

        if self.is_master:
            print(' Done.')
    

    def validation_step(self, data_list):
        """
        Run a training step.
        """
        
        # Choose autocast context: fp16 (amp), bf16, or disabled
        if self.fp16_mode == 'amp':
            autocast_ctx = torch.autocast(device_type='cuda', dtype=torch.float16)
        elif self.fp16_mode == 'bf16':
            autocast_ctx = torch.autocast(device_type='cuda', dtype=torch.bfloat16)
        else:
            autocast_ctx = nullcontext()


        for i, mb_data in enumerate(data_list):
            
            # data_list should have length 1, need to set batch_split=1 in trainer
            with autocast_ctx:
                loss, status = self.eval_nll(**mb_data)
                l = loss / len(data_list)

        dicts = {'nll_sum': l}
        dicts.update(status)
        return dicts

    def eval_nll(
            self,
            x_0: torch.Tensor,
            cond=None,
            **kwargs
        ) -> Tuple[Dict, Dict]:
            """
            Compute training losses for a single timestep.

            Args:
                x_0: The [N x C x ...] tensor of noiseless inputs.
                cond: The [N x ...] tensor of additional conditions.
                kwargs: Additional arguments to pass to the backbone.

            Returns:
                a dict with the key "loss" containing a tensor of shape [N].
                may also contain other keys for different terms.
            """
            bsz = x_0.shape[0]
            noise = torch.randint(0, self.vocab_size, x_0.shape, device=x_0.device)
            t = self.sample_t(bsz).to(x_0.device).float()
            x_t = self.diffuse(x_0, t, noise=noise)
            cond = self.get_cond(cond, **kwargs)
            
            pred_logits = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
            pred_logits_normalized = F.log_softmax(pred_logits, dim=-1) # shape [B,...,vocab_size]
            
            dalpha_t, alpha_t = self.noise(t)
            dalpha_t = torch.ones_like(alpha_t) * dalpha_t  
            # reshape to (B, L) while computing loss
            log_x_theta = rearrange(pred_logits_normalized, 'b d h w v -> b (d h w) v')
            x0 = rearrange(x_0, 'b 1 d h w -> b (d h w)')
            log_p_theta = torch.gather(
                input=log_x_theta,
                dim=-1,
                index = x0[:,:, None]
            ).squeeze(-1)

            loss_nll =  -log_p_theta.mean(-1).sum()
            terms = edict()

            with torch.no_grad():
                nll_per_instance = -log_p_theta.detach().mean(dim=-1)
            time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
            for i in range(10):
                if (time_bin == i).sum() != 0:
                    terms[f"bin_{i}"] = {"nll": nll_per_instance[time_bin == i].cpu()}

            return loss_nll, terms


    @torch.no_grad()
    def validate_nll(self, runs: int = 1):
        if not self.is_master:
            return
        self.training_models['denoiser'].eval()
        avg_nll = torch.tensor(0.0, dtype = torch.float64)
        avg_nlls = []
        bin_collection = {f'bin_{i}': torch.tensor([]) for i in range(10)}

        # prepare dataloader
        dataloader = DataLoader(
            self.dataset,
            batch_size=self.batch_size_per_gpu,
            num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
            pin_memory=True,
            drop_last=False,
            persistent_workers=True,
            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
            sampler=self.data_sampler,
        )

        for run in range(runs):

            total_nll = torch.tensor(0.0, dtype = torch.float64)

            for this_data in tqdm(dataloader):
                if self.prefetch_data:
                    if self._data_prefetched is None:
                        self._data_prefetched = recursive_to_device(this_data, self.device, non_blocking=True)
                    data = self._data_prefetched
                    self._data_prefetched = recursive_to_device(this_data, self.device, non_blocking=True)
                else:
                    data = recursive_to_device(this_data, self.device, non_blocking=True)
                
                # if the data is a dict, we need to split it into multiple dicts with batch_size_per_gpu
                assert self.batch_split == 1, "NLL validation currently only supports batch_split=1"
                if isinstance(data, dict):
                    if self.batch_split == 1:
                        data_list = [data]
                    else:
                        batch_size = list(data.values())[0].shape[0]
                        data_list = [
                            {k: v[i * batch_size // self.batch_split:(i + 1) * batch_size // self.batch_split] for k, v in data.items()}
                            for i in range(self.batch_split)
                        ]
                elif isinstance(data, list):
                    data_list = data
                else:
                    raise ValueError('Data must be a dict or a list of dicts.')

                # now we have data_list for data
                step_log = self.validation_step(data_list)
                total_nll += step_log['nll_sum'].cpu().double()
                print('Current NLL:', step_log['nll_sum'].item()/self.batch_size_per_gpu)
                # collect nlls in bins
                for i in range(10):
                    bin_key = f'bin_{i}'
                    if bin_key in step_log:
                        bin_collection[bin_key] = torch.concat(
                            [bin_collection[bin_key], step_log[bin_key]['nll'].double()]
                        )
            avg_nlls.append((total_nll / len(self.dataset)).item())
            avg_nll += total_nll / len(self.dataset) / runs
        print(f'Validation NLL: {avg_nll.item()}')
        path = os.path.join(self.output_dir, 'nlls.json')

        # calculate stats for bin collection
        bin_log = {}
        for i in range(10):
            bin_key = f'bin_{i}'
            if bin_collection[bin_key].numel() > 0:
                bin_log[bin_key] = {
                    'mean': bin_collection[bin_key].mean().item(),
                    'std': bin_collection[bin_key].std().item(),
                    'count': bin_collection[bin_key].numel(),
                }
            else:
                bin_log[bin_key] = {
                    'mean': None,
                    'std': None,
                    'count': 0,
                }
        print('NLL per time bin:', bin_log)
        logs = {}
        logs.update({
            'nlls': avg_nlls,
            'mean': avg_nll.item(),
            'std': np.std(avg_nlls),
            'bin_stats': bin_log
        })
        with open(path, 'w') as f:
            json.dump(logs, f, indent=4)

        # plot figrue and save
        import matplotlib.pyplot as plt
        bins = sorted(bin_log.keys(), key=lambda k: int(k.split('_')[-1]))
        means = [bin_log[k]['mean'] for k in bins]
        stds  = [bin_log[k]['std']  for k in bins]
        counts = [bin_log[k]['count'] for k in bins]

        fig, ax = plt.subplots(figsize=(8, 4.5))
        x = range(len(bins))
        bars = ax.bar(x, means, yerr=stds, capsize=4)
        ax.set_xticks(x, bins, rotation=45, ha='right')
        ax.set_ylabel('Mean')
        ax.set_title('Mean ± Std per Bin')
        fig.tight_layout()
        out_path = 'bin_mean_std.png'
        fig.savefig(os.path.join(self.output_dir, out_path), dpi=300)
        plt.close(fig)

class DiscreteDiffusionCFGTrainer(DiscreteClassifierFreeGuidanceMixin, DiscreteDiffusionTrainer):
    """
    Trainer for diffusion model with flow matching objective and classifier-free guidance.
    
    Args:
        models (dict[str, nn.Module]): Models to train.
        dataset (torch.utils.data.Dataset): Dataset.
        output_dir (str): Output directory.
        load_dir (str): Load directory.
        step (int): Step to load.
        batch_size (int): Batch size.
        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
        batch_split (int): Split batch with gradient accumulation.
        max_steps (int): Max steps.
        optimizer (dict): Optimizer config.
        lr_scheduler (dict): Learning rate scheduler config.
        elastic (dict): Elastic memory management config.
        grad_clip (float or dict): Gradient clip config.
        ema_rate (float or list): Exponential moving average rates.
        fp16_mode (str): FP16 mode.
            - None: No FP16.
            - 'inflat_all': Hold a inflated fp32 master param for all params.
            - 'amp': Automatic mixed precision.
        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
        finetune_ckpt (dict): Finetune checkpoint.
        log_param_stats (bool): Log parameter stats.
        i_print (int): Print interval.
        i_log (int): Log interval.
        i_sample (int): Sample interval.
        i_save (int): Save interval.
        i_ddpcheck (int): DDP check interval.

        t_schedule (dict): Time schedule for flow matching.
        sigma_min (float): Minimum noise level.
        p_uncond (float): Probability of dropping conditions.
    """
    pass


class TextConditionedDiscreteDiffusionCFGTrainer(TextConditionedMixin, DiscreteDiffusionCFGTrainer):
    """
    Trainer for text-conditioned diffusion model with flow matching objective and classifier-free guidance.
    
    Args:
        models (dict[str, nn.Module]): Models to train.
        dataset (torch.utils.data.Dataset): Dataset.
        output_dir (str): Output directory.
        load_dir (str): Load directory.
        step (int): Step to load.
        batch_size (int): Batch size.
        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
        batch_split (int): Split batch with gradient accumulation.
        max_steps (int): Max steps.
        optimizer (dict): Optimizer config.
        lr_scheduler (dict): Learning rate scheduler config.
        elastic (dict): Elastic memory management config.
        grad_clip (float or dict): Gradient clip config.
        ema_rate (float or list): Exponential moving average rates.
        fp16_mode (str): FP16 mode.
            - None: No FP16.
            - 'inflat_all': Hold a inflated fp32 master param for all params.
            - 'amp': Automatic mixed precision.
        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
        finetune_ckpt (dict): Finetune checkpoint.
        log_param_stats (bool): Log parameter stats.
        i_print (int): Print interval.
        i_log (int): Log interval.
        i_sample (int): Sample interval.
        i_save (int): Save interval.
        i_ddpcheck (int): DDP check interval.

        t_schedule (dict): Time schedule for flow matching.
        sigma_min (float): Minimum noise level.
        p_uncond (float): Probability of dropping conditions.
        text_cond_model(str): Text conditioning model.
    """
    pass


class ImageConditionedDiscreteDiffusionCFGTrainer(ImageConditionedMixin, DiscreteDiffusionCFGTrainer):
    """
    Trainer for image-conditioned diffusion model with flow matching objective and classifier-free guidance.
    
    Args:
        models (dict[str, nn.Module]): Models to train.
        dataset (torch.utils.data.Dataset): Dataset.
        output_dir (str): Output directory.
        load_dir (str): Load directory.
        step (int): Step to load.
        batch_size (int): Batch size.
        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
        batch_split (int): Split batch with gradient accumulation.
        max_steps (int): Max steps.
        optimizer (dict): Optimizer config.
        lr_scheduler (dict): Learning rate scheduler config.
        elastic (dict): Elastic memory management config.
        grad_clip (float or dict): Gradient clip config.
        ema_rate (float or list): Exponential moving average rates.
        fp16_mode (str): FP16 mode.
            - None: No FP16.
            - 'inflat_all': Hold a inflated fp32 master param for all params.
            - 'amp': Automatic mixed precision.
        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
        finetune_ckpt (dict): Finetune checkpoint.
        log_param_stats (bool): Log parameter stats.
        i_print (int): Print interval.
        i_log (int): Log interval.
        i_sample (int): Sample interval.
        i_save (int): Save interval.
        i_ddpcheck (int): DDP check interval.

        t_schedule (dict): Time schedule for flow matching.
        sigma_min (float): Minimum noise level.
        p_uncond (float): Probability of dropping conditions.
        image_cond_model (str): Image conditioning model.
    """
    pass