trellis/pipelines/samplers/flow_euler.py

from typing import *
import torch
import numpy as np
from tqdm import tqdm
from easydict import EasyDict as edict
from .base import Sampler
from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin
from .guidance_interval_mixin import GuidanceIntervalSamplerMixin
import math
from trellis.modules.spatial import patchify, unpatchify
from trellis.utils import render_utils, postprocessing_utils
from trellis.utils import loss_utils
import trellis.modules.sparse as sp
import torch.nn.functional as F


class FlowEulerSampler(Sampler):
    """
    Generate samples from a flow-matching model using Euler sampling.

    Args:
        sigma_min: The minimum scale of noise in flow.
    """
    def __init__(
        self,
        sigma_min: float,
    ):
        self.sigma_min = sigma_min

    def _eps_to_xstart(self, x_t, t, eps):
        assert x_t.shape == eps.shape
        return (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * eps) / (1 - t)
  
    def _xstart_to_x_t(self, x_0, t, eps):
        assert x_0.shape == eps.shape
        return (1-t) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t) * eps
        # return (1-t) * x_0  + t * eps + self.sigma_min * (1-t)  * eps 

    def _xstart_to_eps(self, x_t, t, x_0):
        assert x_t.shape == x_0.shape
        return (x_t - (1 - t) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t)

    def _v_to_xstart_eps(self, x_t, t, v):
        assert x_t.shape == v.shape
        eps = (1 - t) * v + x_t
        x_0 = (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * v
        return x_0, eps
        
    def _xstart_to_v(self, x_0, x_t, t):
        assert x_0.shape == x_t.shape
        return (x_t - (1 - self.sigma_min) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t)


    def _inference_model(self, model, x_t, t, cond=None, **kwargs):
        t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
        return model(x_t, t, cond, **kwargs)

    def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs):
        param = kwargs.pop("parameterization", "v")
        if param == "v":
            pred_v = self._inference_model(model, x_t, t, cond, **kwargs)
            pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v)
        elif param == "x0":
            pred_x_0 = self._inference_model(model, x_t, t, cond, **kwargs)
            pred_v = self._xstart_to_v(x_0=pred_x_0, x_t=x_t, t=t)
        return pred_x_0, None, pred_v

    def _get_model_gt(self, x_0, t, noise):
        gt_x_t = self._xstart_to_x_t(x_0, t, noise)
        gt_v = self._xstart_to_v(x_0, gt_x_t, t)
        return gt_x_t, gt_v

    @torch.no_grad()
    def sample_once(
        self,
        model,
        x_t,
        t: float,
        t_prev: float,
        cond: Optional[Any] = None,
        **kwargs
    ):
        """
        Sample x_{t-1} from the model using Euler method.
        
        Args:
            model: The model to sample from.
            x_t: The [N x C x ...] tensor of noisy inputs at time t.
            t: The current timestep.
            t_prev: The previous timestep.
            cond: conditional information.
            **kwargs: Additional arguments for model inference.

        Returns:
            a dict containing the following
            - 'pred_x_prev': x_{t-1}.
            - 'pred_x_0': a prediction of x_0.
        """
        pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
        pred_x_prev = x_t - (t - t_prev) * pred_v
        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0, "pred_eps": pred_eps})

    def sample_once_opt(
        self,
        model,
        x_t,
        t: float,
        t_prev: float,
        cond: Optional[Any] = None,
        **kwargs
    ):
        """
        Sample x_{t-1} from the model using Euler method.
        
        Args:
            model: The model to sample from.
            x_t: The [N x C x ...] tensor of noisy inputs at time t.
            t: The current timestep.
            t_prev: The previous timestep.
            cond: conditional information.
            **kwargs: Additional arguments for model inference.

        Returns:
            a dict containing the following
            - 'pred_x_prev': x_{t-1}.
            - 'pred_x_0': a prediction of x_0.
        """
        pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
        pred_x_prev = x_t - (t - t_prev) * pred_v
        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0, "pred_eps": pred_eps})

    def sample_once_opt_delta_v(
        self,
        model,
        slat_decoder_gs,
        slat_decoder_mesh,
        dreamsim_model,
        learning_rate,
        input_images,
        extrinsics, 
        intrinsics,
        x_t,
        t: float,
        t_prev: float,
        cond: Optional[Any] = None,
        **kwargs
    ):
        """
        Sample x_{t-1} from the model using Euler method.
        
        Args:
            model: The model to sample from.
            x_t: The [N x C x ...] tensor of noisy inputs at time t.
            t: The current timestep.
            t_prev: The previous timestep.
            cond: conditional information.
            **kwargs: Additional arguments for model inference.

        Returns:
            a dict containing the following
            - 'pred_x_prev': x_{t-1}.
            - 'pred_x_0': a prediction of x_0.
        """
        torch.cuda.empty_cache()
        with torch.no_grad():
            pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
        pred_v_opt_feat = torch.nn.Parameter(pred_v.feats.detach().clone())
        optimizer = torch.optim.Adam([pred_v_opt_feat], betas=(0.5, 0.9), lr=learning_rate)
        pred_v_opt = sp.SparseTensor(feats=pred_v_opt_feat, coords=pred_v.coords)
        total_steps = 5
        input_images = F.interpolate(input_images, size=(259, 259), mode='bilinear', align_corners=False)
        with tqdm(total=total_steps, disable=True, desc='Appearance (opt): optimizing') as pbar:
            for step in range(total_steps):
                optimizer.zero_grad()
                pred_x_0, _ = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v_opt)
                pred_gs = slat_decoder_gs(pred_x_0)
                # pred_mesh = slat_decoder_mesh(pred_x_0)
                rend_gs = render_utils.render_frames(pred_gs[0], extrinsics, intrinsics, {'resolution': 259, 'bg_color': (0, 0, 0)}, need_depth=True, opt=True)['color']
                # rend_mesh = render_utils.render_frames_opt(pred_mesh[0], extrinsics, intrinsics, {'resolution': 518, 'bg_color': (0, 0, 0)}, need_depth=True)['color']
                rend_gs = torch.stack(rend_gs, dim=0)
                loss_gs = loss_utils.l1_loss(rend_gs, input_images) + (1 - loss_utils.ssim(rend_gs, input_images)) + loss_utils.lpips(rend_gs, input_images) + dreamsim_model(rend_gs, input_images).mean()
                # loss_gs = (1 - loss_utils.ssim(rend_gs, input_images)) + loss_utils.lpips(rend_gs, input_images) + dreamsim_model(rend_gs, input_images).mean()
                # loss_mesh = loss_utils.l1_loss(rend_mesh, input_images) + 0.2 * (1 - loss_utils.ssim(rend_mesh, input_images)) + 0.2 * loss_utils.lpips(rend_mesh, input_images)
                loss = loss_gs + 0.2 * loss_utils.l1_loss(pred_v_opt_feat, pred_v.feats)
                loss.backward()
                optimizer.step()
                pbar.set_postfix({'loss': loss.item()})
                pbar.update()
        
        pred_x_prev = x_t - (t - t_prev) * pred_v_opt.detach()
        torch.cuda.empty_cache()
        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0, "pred_eps": pred_eps})

    def sample_opt(
        self,
        model,
        noise,
        cond: Optional[Any] = None,
        steps: int = 50,
        rescale_t: float = 1.0,
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        sample = noise
        t_seq = np.linspace(1, 0, steps + 1)
        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
        t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
        ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
        for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose):
            out = self.sample_once_opt(model, sample, t, t_prev, cond, **kwargs)
            sample = out.pred_x_prev
            ret.pred_x_t.append(out.pred_x_prev)
            ret.pred_x_0.append(out.pred_x_0)
        ret.samples = sample
        return ret

    def sample_opt_delta_v(
        self,
        model,
        slat_decoder_gs,
        slat_decoder_mesh,
        dreamsim_model,
        apperance_learning_rate,
        start_t,
        input_images,
        extrinsics, 
        intrinsics,
        noise,
        cond: Optional[Any] = None,
        steps: int = 50,
        rescale_t: float = 1.0,
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        sample = noise
        t_seq = np.linspace(1, 0, steps + 1)
        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
        t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
        ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
        # def cosine_anealing(step, total_steps, start_lr, end_lr):
        #     return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps))
        for i, (t, t_prev) in enumerate(tqdm(t_pairs, desc="Sampling", disable=not verbose)):
            if t > start_t:
                out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
                sample = out.pred_x_prev
                ret.pred_x_t.append(out.pred_x_prev)
                ret.pred_x_0.append(out.pred_x_0)
            else:
                # learning_rate = cosine_anealing(i - int(np.where(t_seq <= start_t)[0].min()), int(steps - np.where(t_seq <= start_t)[0].min()), apperance_learning_rate, 1e-5)
                learning_rate = apperance_learning_rate
                out = self.sample_once_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, dreamsim_model, learning_rate, input_images, extrinsics, intrinsics, sample, t, t_prev, cond, **kwargs)
                sample = out.pred_x_prev
                ret.pred_x_t.append(out.pred_x_prev)
                ret.pred_x_0.append(out.pred_x_0)
        ret.samples = sample
        return ret

    @torch.no_grad()
    def sample(
        self,
        model,
        noise,
        cond: Optional[Any] = None,
        steps: int = 50,
        rescale_t: float = 1.0,
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        sample = noise
        t_seq = np.linspace(1, 0, steps + 1)
        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
        t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
        ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
        for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose):
            out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
            sample = out.pred_x_prev
            ret.pred_x_t.append(out.pred_x_prev)
            ret.pred_x_0.append(out.pred_x_0)
        ret.samples = sample
        return ret



class LatentMatchSampler(FlowEulerSampler):
    """
    Generate samples from a Bridge Matching model using Euler sampling.
    This sampler is designed for Latent Bridge Matching (LBM), where
    the target (x_1) for training is assumed to be sampled from a Gaussian distribution,
    and the source (x_0) for inference is also typically a Gaussian noise.

    Args:
        sigma_bridge: The sigma parameter for the Bridge Matching stochastic interpolant.
                      This controls the amount of stochasticity in the SDE (LBM paper Eq 1).

    """
    def __init__(
        self,
        sigma_bridge: float = 0.1,
        **kwargs
    ):
        # Call parent constructor with a dummy sigma_min.
        # sigma_min is specific to Flow Matching's interpolant, which we override.
        super().__init__(sigma_min=0.0, **kwargs)
        self.sigma_bridge = sigma_bridge

    # Override _xstart_to_x_t for Bridge Matching's stochastic interpolant
    # This method is used to generate gt_x_t for training.
    def _xstart_to_x_t(self, x_0: torch.Tensor, t: float, noise: torch.Tensor,  x_1: torch.Tensor) -> torch.Tensor:
        """
        Calculates x_t according to the Bridge Matching stochastic interpolant.
        This function is used during training to generate noisy samples x_t from
        paired x_0 and x_1 samples. The 'x_1' argument is crucial for Bridge Matching.

        Args:
            x_0: The source latent tensor (e.g., from a data distribution or Gaussian).
            t: The current timestep (float between 0 and 1).
            eps: A random noise tensor (epsilon).
            x_1: The target latent tensor. Required for Bridge Matching.

        Returns:
            The interpolated latent tensor x_t.
        """        
        # LBM interpolant formula: x_t = (1-t)x_0 + t*x_1 + sigma_bridge*sqrt(t*(1-t))*epsilon
        return (1 - t) * x_0 + t * x_1 + self.sigma_bridge * math.sqrt(t * (1 - t)) * noise

    # This method is used to calculate gt_v for training.
    def _xstart_to_v(self, x_0: torch.Tensor, x_t: torch.Tensor, t: float, x_1: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Calculates the ground truth drift (v) that the model should predict for Bridge Matching.
        This function is used in the training objective to define the target for the model.

        Args:
            x_0: The source latent tensor.
            x_t: The interpolated latent tensor at time t.
            t: The current timestep (float between 0 and 1).
            x_1: The target latent tensor. Required for Bridge Matching.

        Returns:
            The target drift tensor v.
        """
        if x_1 is None:
            # This branch should ideally not be hit during _get_model_gt for LBM.
            raise ValueError("For Bridge Matching's target drift calculation, x_1 (target latent) must be provided.")
        assert x_t.shape == x_1.shape, "x_t and x_1 shapes must match."
        # LBM drift formula: v = (x_1 - x_t) / (1 - t)
        # Add a small epsilon to (1-t) to prevent division by zero if t is exactly 1.
        epsilon_t = 1e-5 # Small epsilon for numerical stability
        return (x_t - x_0) / (t + epsilon_t)

    # Override _get_model_gt to provide ground truth for Bridge Matching training.
    # In this simplified case, x_1 is sampled from a Gaussian distribution.
    def _get_model_gt(self, x_0: torch.Tensor, t: float, x_1: torch.Tensor):
        """
        Calculates ground truth x_t and v_target for Bridge Matching training purposes.
        In this simplified case, x_1 is sampled from a Gaussian distribution.
        
        Args:
            x_0: The source latent tensor (e.g., from a data distribution, or another Gaussian).
            t: The current timestep.
            noise: A random noise tensor (epsilon).

        Returns:
            A tuple (gt_x_t, gt_v).
        """
        # Sample x_1 from a Gaussian distribution with the same shape as x_0
        # This simulates the target distribution being Gaussian.
        if isinstance(x_0, sp.SparseTensor):
            noise = sp.SparseTensor(
                feats=torch.randn_like(x_0.feats).to(x_0.feats.device),
                coords=x_0.coords,
            )
        else:
            noise = torch.randn_like(x_0).to(x_0.device)
        # For Bridge Matching, _xstart_to_x_t needs x_1
        gt_x_t = self._xstart_to_x_t(x_0, t, noise, x_1=x_1) 
        gt_v = self._xstart_to_v(x_0, gt_x_t, t, x_1=x_1) 
        return gt_x_t, gt_v

    # Override sample_once to include the stochastic term for SDE integration.
    @torch.no_grad()
    def sample_once(
        self,
        model,
        x_t: torch.Tensor,
        t: float,
        t_prev: float,
        cond: Optional[Any] = None,
        **kwargs
    ) -> edict:
        """
        Performs a single Euler step to sample x_{t_next} from x_t for Bridge Matching.
        The model is assumed to predict the drift 'v' as per LBM's formulation.
        
        Args:
            model: The model to sample from (should be trained for Bridge Matching).
            x_t: The [N x C x ...] tensor of current latent inputs at time t.
            t: The current timestep.
            t_next: The next timestep in the forward integration sequence (t+dt).
            cond: conditional information.
            **kwargs: Additional arguments for model inference.

        Returns:
            An edict containing:
            - 'pred_x_prev': The estimated latent tensor at t_next.
            - 'pred_x_0': A prediction of x_0 (may be None as direct derivation is complex in LBM).
            - 'pred_eps': A prediction of eps (may be None).
        """
        # Get model's prediction of the drift (v)
        # We use the parent's _get_model_prediction. Its _v_to_xstart_eps uses sigma_min,
        # which is a dummy value here. For LBM, pred_v is the main output.
        pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
        
        # Calculate time step difference (dt)
        dt = t - t_prev # This is the forward step size
        
        # Sample noise for the stochastic part of the SDE
        # The SDE for LBM is dx_t = v(x_t, t) dt + sigma dB_t
        # For Euler, dB_t approx sqrt(dt) * Z, where Z ~ N(0,I)
        # noise_increment = sp.SparseTensor(
        #     feats=torch.randn_like(x_t.feats).to(x_t.feats.device),
        #     coords=x_t.coords,
        # )
        # if isinstance(x_t, sp.SparseTensor):
        #     noise_increment = sp.SparseTensor(
        #         feats=torch.randn_like(x_t.feats).to(x_t.feats.device),
        #         coords=x_t.coords,
        #     )
        # else:
        #     noise_increment = torch.randn_like(x_t).to(x_t.device)
        # noise_increment = noise_increment * self.sigma_bridge * torch.sqrt(torch.tensor(max(0.0, dt), device=x_t.device))
        # pred_x_prev = x_t - (t - t_prev) * pred_v - noise_increment
        pred_x_prev = x_t - (t - t_prev) * pred_v 
        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0, "pred_eps": pred_eps})
    
class FlowMatchingSampler(FlowEulerSampler):
    """
    Implementation of Flow Matching using Euler sampling.
    Inherits from FlowEulerSampler and modifies key methods for flow matching.
    """
    def __init__(self, sigma_min: float = 0.0):
        super().__init__(sigma_min=sigma_min)

    def _compute_velocity(self, x_t: torch.Tensor, x_0: torch.Tensor, t: float) -> torch.Tensor:
        return ((1 - self.sigma_min) * x_t - x_0 ) / (self.sigma_min + (1 - self.sigma_min) * t)

    def _get_model_gt(self, x_1: torch.Tensor, t: float, x_0: torch.Tensor = None):
        # TODO: Implement this method
        pass
        # """
        # Get ground truth for training.
        # Args:
        #     x_1: Target endpoint
        #     t: Time point
        #     noise: Initial noise to use as x_0
        # """
        # x_t = (1 - t) * x_0 + t * x_1
        # v = self._compute_velocity(x_t, x_0, t)
        # eps = x_t + (1 - t) * v  # Convert velocity to noise
        # return x_t, eps, v

    def _v_to_xstart_eps(self, x_t: torch.Tensor, t: float, v: torch.Tensor):
        """Convert velocity to x_0 and noise predictions"""
        eps = x_t + (1 - t) * v
        x_0 = self._eps_to_xstart(x_t, t, eps)
        return x_0, eps

    @torch.no_grad()
    def sample(
        self,
        model,
        x_1: torch.Tensor,
        cond: Optional[Any] = None,
        steps: int = 50,
        rescale_t: float = 1.0,
        verbose: bool = True,
        **kwargs
    ) -> Dict[str, torch.Tensor]:
        """
        Generate samples by following the flow from noise to x_1.
        Args:
            model: The model to sample from
            x_1: Target endpoint
            cond: Conditional information
            steps: Number of sampling steps
            rescale_t: Time rescaling factor
            verbose: Whether to show progress bar
            **kwargs: Additional model arguments
        Returns:
            Dictionary containing sampling trajectory and predictions
        """
        # Initialize with noise as x_0
        noise = torch.randn_like(x_1)
        current_x = noise
        
        t_seq = np.linspace(1, 0, steps + 1)
        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
        t_pairs = list(zip(t_seq[:-1], t_seq[1:]))
        
        ret = edict({
            "samples": None,
            "pred_x_t": [],
            "pred_x_0": []
        })
        
        for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose):
            out = self.sample_once(model, current_x, t, t_prev, cond, **kwargs)
            current_x = out.pred_x_prev
            ret.pred_x_t.append(out.pred_x_prev)
            ret.pred_x_0.append(out.pred_x_0)
            
        ret.samples = current_x
        return ret

    def sample_once(
        self,
        model,
        x_t: torch.Tensor,
        t: float,
        t_prev: float,
        cond: Optional[Any] = None,
        **kwargs
    ) -> Dict:
        """
        Sample x_{t-1} from the model using Euler method.
        Args:
            model: The model to sample from
            x_t: Current state
            t: Current time
            t_prev: Next time step
            cond: Conditional information
            **kwargs: Additional model arguments
        Returns:
            Dictionary containing predictions
        """
        pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
        pred_x_prev = x_t + (t_prev - t) * pred_v
        return edict({
            "pred_x_prev": pred_x_prev,
            "pred_x_0": pred_x_0,
            "pred_eps": pred_eps
        })

class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
    """
    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance.
    """
    @torch.no_grad()
    def sample(
        self,
        model,
        noise,
        cond,
        neg_cond,
        steps: int = 50,
        rescale_t: float = 1.0,
        cfg_strength: float = 3.0,
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            neg_cond: negative conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            cfg_strength: The strength of classifier-free guidance.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, **kwargs)


class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSampler):
    """
    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.
    """
    @torch.no_grad()
    def sample(
        self,
        model,
        noise,
        cond,
        neg_cond,
        steps: int = 50,
        rescale_t: float = 1.0,
        cfg_strength: float = 3.0,
        cfg_interval: Tuple[float, float] = (0.0, 1.0),
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            neg_cond: negative conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            cfg_strength: The strength of classifier-free guidance.
            cfg_interval: The interval for classifier-free guidance.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)

    def sample_opt(
        self,
        model,
        noise,
        cond,
        neg_cond,
        steps: int = 50,
        rescale_t: float = 1.0,
        cfg_strength: float = 3.0,
        cfg_interval: Tuple[float, float] = (0.0, 1.0),
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            neg_cond: negative conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            cfg_strength: The strength of classifier-free guidance.
            cfg_interval: The interval for classifier-free guidance.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        return super().sample_opt(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)

    def sample_opt_delta_v(
        self,
        model,
        slat_decoder_gs,
        slat_decoder_mesh,
        dreamsim_model,
        apperance_learning_rate,
        start_t,
        input_images,
        extrinsics, 
        intrinsics,
        noise,
        cond,
        neg_cond,
        steps: int = 50,
        rescale_t: float = 1.0,
        cfg_strength: float = 3.0,
        cfg_interval: Tuple[float, float] = (0.0, 1.0),
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            neg_cond: negative conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            cfg_strength: The strength of classifier-free guidance.
            cfg_interval: The interval for classifier-free guidance.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        return super().sample_opt_delta_v(model, slat_decoder_gs, slat_decoder_mesh, dreamsim_model, apperance_learning_rate, start_t, input_images, extrinsics, intrinsics,noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)


class LatentMatchGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, LatentMatchSampler):
    """
    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.
    """
    @torch.no_grad()
    def sample(
        self,
        model,
        noise,
        cond,
        neg_cond,
        steps: int = 50,
        rescale_t: float = 1.0,
        cfg_strength: float = 3.0,
        cfg_interval: Tuple[float, float] = (0.0, 1.0),
        verbose: bool = True,
        **kwargs
    ):
        """
        Generate samples from the model using Euler method.
        
        Args:
            model: The model to sample from.
            noise: The initial noise tensor.
            cond: conditional information.
            neg_cond: negative conditional information.
            steps: The number of steps to sample.
            rescale_t: The rescale factor for t.
            cfg_strength: The strength of classifier-free guidance.
            cfg_interval: The interval for classifier-free guidance.
            verbose: If True, show a progress bar.
            **kwargs: Additional arguments for model_inference.

        Returns:
            a dict containing the following
            - 'samples': the model samples.
            - 'pred_x_t': a list of prediction of x_t.
            - 'pred_x_0': a list of prediction of x_0.
        """
        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)