models/Tiffusion/tiffusion_backup.py

import math
import torch
import torch.nn.functional as F

from torch import nn
from einops import reduce
from tqdm.auto import tqdm
from functools import partial
from .transformer import Transformer
from ..model_utils import default, identity, extract
from .control import *

import mlflow.pyfunc
import mlflow
from mlflow.models import infer_signature

# import matplotlib.pyplot as plt
# images_cache = []

def linear_beta_schedule(timesteps):
    scale = 1000 / timesteps
    beta_start = scale * 0.0001
    beta_end = scale * 0.02
    return torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)


def cosine_beta_schedule(timesteps, s=0.008):
    """
    cosine schedule
    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
    """
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps, dtype=torch.float64)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0, 0.999)


class Tiffusion(nn.Module):
    def __init__(
        self,
        seq_length,
        feature_size,
        n_layer_enc=3,
        n_layer_dec=6,
        d_model=None,
        timesteps=1000,
        sampling_timesteps=None,
        loss_type="l1",
        beta_schedule="cosine",
        n_heads=4,
        mlp_hidden_times=4,
        eta=0.0,
        attn_pd=0.0,
        resid_pd=0.0,
        kernel_size=None,
        padding_size=None,
        use_ff=True,
        reg_weight=None,
        control_signal={},
        moving_average=False,
        **kwargs,
    ):
        super(Tiffusion, self).__init__()

        self.eta, self.use_ff = eta, use_ff
        self.seq_length = seq_length
        self.feature_size = feature_size
        self.ff_weight = default(reg_weight, math.sqrt(self.seq_length) / 5)
        self.sum_weight = default(reg_weight, math.sqrt(self.seq_length // 10) / 50)
        self.training_control_signal = control_signal # training control signal
        self.moving_average = moving_average

        self.model: Transformer = Transformer(
            n_feat=feature_size,
            n_channel=seq_length,
            n_layer_enc=n_layer_enc,
            n_layer_dec=n_layer_dec,
            n_heads=n_heads,
            attn_pdrop=attn_pd,
            resid_pdrop=resid_pd,
            mlp_hidden_times=mlp_hidden_times,
            max_len=seq_length,
            n_embd=d_model,
            conv_params=[kernel_size, padding_size],
            **kwargs,
        )

        if beta_schedule == "linear":
            betas = linear_beta_schedule(timesteps)
        elif beta_schedule == "cosine":
            betas = cosine_beta_schedule(timesteps)
        else:
            raise ValueError(f"unknown beta schedule {beta_schedule}")

        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)

        (timesteps,) = betas.shape
        self.num_timesteps = int(timesteps)
        self.loss_type = loss_type

        # sampling related parameters
        self.sampling_timesteps = default(
            sampling_timesteps, timesteps
        )  # default num sampling timesteps to number of timesteps at training

        assert self.sampling_timesteps <= timesteps
        self.fast_sampling = self.sampling_timesteps < timesteps

        # helper function to register buffer from float64 to float32

        register_buffer = lambda name, val: self.register_buffer(
            name, val.to(torch.float32)
        )

        register_buffer("betas", betas)
        register_buffer("alphas_cumprod", alphas_cumprod)
        register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)

        # calculations for diffusion q(x_t | x_{t-1}) and others

        register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        register_buffer(
            "sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod)
        )
        register_buffer("log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod))
        register_buffer("sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod))
        register_buffer(
            "sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1)
        )

        # calculations for posterior q(x_{t-1} | x_t, x_0)

        posterior_variance = (
            betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        )

        # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)

        register_buffer("posterior_variance", posterior_variance)

        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain

        register_buffer(
            "posterior_log_variance_clipped",
            torch.log(posterior_variance.clamp(min=1e-20)),
        )
        register_buffer(
            "posterior_mean_coef1",
            betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
        )
        register_buffer(
            "posterior_mean_coef2",
            (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod),
        )

        # calculate reweighting

        register_buffer(
            "loss_weight",
            torch.sqrt(alphas) * torch.sqrt(1.0 - alphas_cumprod) / betas / 100,
        )

    def predict_noise_from_start(self, x_t, t, x0):
        return (
            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0
        ) / extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)

    def predict_start_from_noise(self, x_t, t, noise):
        return (
            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
        )

    def q_posterior(self, x_start, x_t, t):
        posterior_mean = (
            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
            + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = extract(
            self.posterior_log_variance_clipped, t, x_t.shape
        )
        return posterior_mean, posterior_variance, posterior_log_variance_clipped

    def output(self, x, t, padding_masks=None, control_signal=None):
        # if ss:=control_signal.get("sum") is not None and len(ss.shape) == 1:
        #     bs = x.shape[0]
        #     control_signal["sum"] = ss.unsqueeze(0).repeat(bs, 1)
        # print("control_signal", control_signal)
        trend, season = self.model(
            x, t, padding_masks=padding_masks, control_signal=control_signal
        )
        model_output = trend + season
        return model_output

    def model_predictions(
        self, x, t, clip_x_start=False, padding_masks=None, control_signal=None
    ):
        if padding_masks is None:
            padding_masks = torch.ones(
                x.shape[0], self.seq_length, dtype=bool, device=x.device
            )

        maybe_clip = (
            partial(torch.clamp, min=-1.0, max=1.0) if clip_x_start else identity
        )
        x_start = self.output(x, t, padding_masks, control_signal=control_signal)
        x_start = maybe_clip(x_start)
        pred_noise = self.predict_noise_from_start(x, t, x_start)
        return pred_noise, x_start

    def p_mean_variance(self, x, t, clip_denoised=True, control_signal=None):
        _, x_start = self.model_predictions(x, t, control_signal=control_signal)
        if clip_denoised:
            x_start.clamp_(-1.0, 1.0)
        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
            x_start=x_start, x_t=x, t=t
        )
        return model_mean, posterior_variance, posterior_log_variance, x_start

    def p_sample(self, x, t: int, clip_denoised=True, control_signal=None):
        batched_times = torch.full((x.shape[0],), t, device=x.device, dtype=torch.long)
        model_mean, _, model_log_variance, x_start = self.p_mean_variance(
            x=x, t=batched_times, clip_denoised=clip_denoised, control_signal=control_signal
        )
        noise = torch.randn_like(x) if t > 0 else 0.0  # no noise if t == 0
        pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
        return pred_img, x_start

    @torch.no_grad()
    def sample(self, shape, control_signal=None):
        device = self.betas.device
        img = torch.randn(shape, device=device)
        for t in tqdm(
            reversed(range(0, self.num_timesteps)),
            desc="sampling loop time step",
            total=self.num_timesteps,
        ):
            img, _ = self.p_sample(img, t, control_signal=control_signal)
        return img

    @torch.no_grad()
    def fast_sample(self, shape, clip_denoised=True, model_kwargs=None,
):
        batch, device, total_timesteps, sampling_timesteps, eta = (
            shape[0],
            self.betas.device,
            self.num_timesteps,
            self.sampling_timesteps,
            self.eta,
        )

        # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)

        times = list(reversed(times.int().tolist()))
        time_pairs = list(
            zip(times[:-1], times[1:])
        )  # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]
        img = torch.randn(shape, device=device)

        for time, time_next in tqdm(time_pairs, desc="sampling loop time step"):
            time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
            pred_noise, x_start, *_ = self.model_predictions(
                img, time_cond, clip_x_start=clip_denoised,
                control_signal=model_kwargs.get("model_control_signal", {}) if model_kwargs else {}
            )

            if time_next < 0:
                img = x_start
                continue

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]
            sigma = (
                eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            )
            c = (1 - alpha_next - sigma**2).sqrt()
            noise = torch.randn_like(img)
            img = x_start * alpha_next.sqrt() + c * pred_noise + sigma * noise

        return img

    def generate_mts(self, batch_size=16):
        feature_size, seq_length = self.feature_size, self.seq_length
        sample_fn = self.fast_sample if self.fast_sampling else self.sample
        return sample_fn((batch_size, seq_length, feature_size))

    def generate_mts_infill(self, target, partial_mask=None, clip_denoised=True, model_kwargs=None):
        sample_fn = self.fast_sample_infill_float_mask # if self.fast_sampling else self.sample_infill
        print("model_kwargs", model_kwargs)
        print("partial_mask", partial_mask.shape)
        print("target", target.shape)
        return sample_fn(
            shape=target.shape,
            target=target,
            sampling_timesteps=self.sampling_timesteps,
            partial_mask=partial_mask,
            clip_denoised=clip_denoised,
            model_kwargs=model_kwargs
        )

    @property
    def loss_fn(self):
        if self.loss_type == "l1":
            return F.l1_loss
        elif self.loss_type == "l2":
            return F.mse_loss
        else:
            raise ValueError(f"invalid loss type {self.loss_type}")

    def q_sample(self, x_start, t, noise=None):
        noise = default(noise, lambda: torch.randn_like(x_start))
        return (
            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
        )

    @torch.no_grad()
    def calculate_dynamic_window(self, t: torch.Tensor) -> torch.Tensor:
        # Batch-wise time point normalization
        t_min = 0 # t.min()
        t_max = 500 # t.max()
        # t_normalized = (t - t_min) / (t_max - t_min)
        # Compute window sizes
        # windows = ((t_normalized.exp2() - 1) * 15 // 1 + 1).long()
        # plt.scatter(t, ( (5 ** ((t - 0) / 1000))) * 15 // 7  + 1)
        windows = ((5 ** ( t / 500)) * 15 // 5  - 2).long()
        return windows

    @torch.no_grad()
    def torch_moving_average(self, bs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Compute moving average for a time series tensor with dynamically calculated 
        window sizes for each sample.
        
        Parameters:
        -----------
        bs : torch.Tensor
            Input time series tensor of shape (batch_size, sequence_length, features)
        t : torch.Tensor
            Time points tensor of shape (batch_size, sequence_length)
        
        Returns:
        --------
        torch.Tensor
            Moving average tensor with the same shape as input
        """
        # Get tensor dimensions
        batch_size, total_seq_length, num_features = bs.shape
        
        # Calculate dynamic window sizes for each sample
        windows = self.calculate_dynamic_window(t)
        
        # Create output tensor initialized with the original values
        moving_avg = bs.clone()
        
        # Compute moving average for each sample and time point
        for b in range(batch_size):
            for i in range(total_seq_length):
                # Get the window size for this sample and time point
                current_window = windows[b].item()
                
                # Determine the start and end of the window
                start = max(0, i - current_window + 1)
                window = bs[b:b+1, start:i+1, :]
                
                # Compute average along the time dimension
                window_avg = window.mean(dim=1)
                
                # Replace values where we have enough previous steps
                if i >= current_window - 1:
                    moving_avg[b, i, :] = window_avg

        return moving_avg

    def _train_loss(
        self,
        x_start,
        t,
        target=None,
        noise=None,
        padding_masks=None,
        control_signal=None,
    ):
        noise = default(noise, lambda: torch.randn_like(x_start))
        if target is None:
            target = x_start

        x = self.q_sample(x_start=x_start, t=t, noise=noise)  # noise sample

        # with torch.no_grad():
        #     if control_signal is None:
        #         control_signal = {
        #             "sum": target.mean(1),
        #             "top-peak-position": target.topk(self.seq_length // 20, dim=1)[1],
        #         }  # .unsqueeze(-1)
        #     # elif self.control_sum:
        #     #     ss = control_signal.get("sum")
        #     #     if len(ss.shape) == 1:
        #     #         bs = x.shape[0]
        #     #         control_signal["sum"] = ss.unsqueeze(0).repeat(bs, 1)
        #     #     control_signal = control_signal
        #     else:
        #         control_signal = {}

        model_out = self.output(x, t, padding_masks, control_signal=control_signal)

        # moving average according to the timestamp t, t larger means more stable, less noise
        if self.moving_average:
            target = self.torch_moving_average(target.cpu(), t.cpu()).to(model_out.device)

        train_loss = self.loss_fn(model_out, target, reduction="none")

        fourier_loss = torch.tensor([0.0])
        if self.use_ff:
            fft1 = torch.fft.fft(model_out.transpose(1, 2), norm="forward")
            fft2 = torch.fft.fft(target.transpose(1, 2), norm="forward")
            fft1, fft2 = fft1.transpose(1, 2), fft2.transpose(1, 2)
            fourier_loss = self.loss_fn(
                torch.real(fft1), torch.real(fft2), reduction="none"
            ) + self.loss_fn(torch.imag(fft1), torch.imag(fft2), reduction="none")
            train_loss += self.ff_weight * fourier_loss

        # if self.control_sum:
        #     train_loss += (
        #         self.loss_fn(model_out[..., 0].sum(1), target[..., 0].sum(1))
        #         / self.seq_length
        #     )
            # * self.sum_weight

        train_loss = reduce(train_loss, "b ... -> b (...)", "mean")
        train_loss = train_loss * extract(self.loss_weight, t, train_loss.shape)
        return train_loss.mean()

    # fmt: off
    def forward(self, x, **kwargs):
        b, c, n, device, feature_size, = *x.shape, x.device, self.feature_size
        assert n == feature_size, f'number of variable must be {feature_size}'
        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
        return self._train_loss(x_start=x, t=t, **kwargs)

    def return_components(self, x, t: int):
        b, c, n, device, feature_size, = *x.shape, x.device, self.feature_size
        assert n == feature_size, f'number of variable must be {feature_size}'
        t = torch.tensor([t])
        t = t.repeat(b).to(device)
        x = self.q_sample(x, t)
        trend, season, residual = self.model(x, t, return_res=True)
        return trend, season, residual, x

    # fmt: on
    def fast_sample_infill_float_mask(
        self,
        shape,
        target: torch.Tensor,  # target time series # [B, L, C]
        sampling_timesteps,
        partial_mask: torch.Tensor = None,  # float mask between 0 and 1 # [B, L, C]
        clip_denoised=True,
        model_kwargs=None,
    ):
        batch, device, total_timesteps, eta = (
            shape[0],
            self.betas.device,
            self.num_timesteps,
            self.eta,
        )
        # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)

        times = list(reversed(times.int().tolist()))
        time_pairs = list(
            zip(times[:-1], times[1:])
        )  # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]

        # Initialize with noise
        img = torch.randn(shape, device=device) # [B, L, C]

        for time, time_next in tqdm(
            time_pairs, desc="conditional sampling loop time step"
        ):
            time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
            pred_noise, x_start, *_ = self.model_predictions(
                img,
                time_cond,
                clip_x_start=clip_denoised,
                control_signal=model_kwargs.get("model_control_signal", {}),
            )

            if time_next < 0:
                img = x_start
                continue

            # Compute the predicted mean
            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]
            sigma = (
                eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            )
            c = (1 - alpha_next - sigma**2).sqrt()

            noise = torch.randn_like(img)
            pred_mean = x_start * alpha_next.sqrt() + c * pred_noise
            img = pred_mean + sigma * noise

            # # Apply partial mask to the current sample
            # if partial_mask is not None:
            #     target_t = self.q_sample(target, t=time_cond)
            #     img = img * (1.0 - partial_mask) + target_t * partial_mask

            # Langevin Dynamics part for additional gradient updates
            img = self.langevin_fn(
                sample=img,
                mean=pred_mean,
                sigma=sigma,
                t=time_cond,
                tgt_embs=target,
                partial_mask=partial_mask,
                enable_float_mask=True,
                **model_kwargs,
            )
            img = img * (1 - partial_mask) + target * partial_mask

        img = img * (1 - partial_mask) + target * partial_mask
        return img


    def fast_sample_infill(
        self,
        shape,
        target,
        sampling_timesteps,
        partial_mask=None,
        clip_denoised=True,
        model_kwargs=None,
    ):
        batch, device, total_timesteps, eta = (
            shape[0],
            self.betas.device,
            self.num_timesteps,
            self.eta,
        )

        # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)

        times = list(reversed(times.int().tolist()))
        time_pairs = list(
            zip(times[:-1], times[1:])
        )  # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]
        img = torch.randn(shape, device=device)

        for time, time_next in tqdm(
            time_pairs, desc="conditional sampling loop time step"
        ):
            time_cond = torch.full((batch,), time, device=device, dtype=torch.long)
            pred_noise, x_start, *_ = self.model_predictions(
                img,
                time_cond,
                clip_x_start=clip_denoised,
                control_signal=model_kwargs.get("model_control_signal", {}),
            )

            if time_next < 0:
                img = x_start
                continue

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]
            sigma = (
                eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            )
            c = (1 - alpha_next - sigma**2).sqrt()
            pred_mean = x_start * alpha_next.sqrt() + c * pred_noise
            noise = torch.randn_like(img)

            img = pred_mean + sigma * noise
            img = self.langevin_fn(
                sample=img,
                mean=pred_mean,
                sigma=sigma,
                t=time_cond,
                tgt_embs=target,
                partial_mask=partial_mask,
                # gradient_control_signal=model_kwargs.get("gradient_control_signal", {}),
                # model_control_signal=model_kwargs.get("model_control_signal", {}),
                **model_kwargs,
            )
            target_t = self.q_sample(target, t=time_cond)
            img[partial_mask] = target_t[partial_mask]

        img[partial_mask] = target[partial_mask]

        return img

    def sample_infill(
        self,
        shape,
        target,
        partial_mask=None,
        clip_denoised=True,
        model_kwargs=None,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.
        """
        batch, device = shape[0], self.betas.device
        img = torch.randn(shape, device=device)
        for t in tqdm(
            reversed(range(0, self.num_timesteps)),
            desc="conditional sampling loop time step",
            total=self.num_timesteps,
        ):
            img = self.p_sample_infill(
                x=img,
                t=t,
                clip_denoised=clip_denoised,
                target=target,
                partial_mask=partial_mask,
                model_kwargs=model_kwargs,
            )

        img[partial_mask] = target[partial_mask]
        return img

    def p_sample_infill(
        self,
        x,
        target,
        t: int,
        partial_mask=None,
        clip_denoised=True,
        model_kwargs=None,
    ):
        b, *_, device = *x.shape, self.betas.device
        batched_times = torch.full((x.shape[0],), t, device=x.device, dtype=torch.long)
        model_mean, _, model_log_variance, _ = self.p_mean_variance(
            x=x, t=batched_times, clip_denoised=clip_denoised, control_signal=model_kwargs.get("model_control_signal", {})
            # don't pass parameters to control signal, for model itself
            # Otherwise pass: control_signal=model_kwargs.get("control_signal", {})
        )
        noise = torch.randn_like(x) if t > 0 else 0.0  # no noise if t == 0
        sigma = (0.5 * model_log_variance).exp()
        pred_img = model_mean + sigma * noise

        pred_img = self.langevin_fn(
            sample=pred_img,
            mean=model_mean,
            sigma=sigma,
            t=batched_times,
            tgt_embs=target,
            partial_mask=partial_mask,
            # control_signal=model_kwargs.get("gradient_control_signal", {}),
            **model_kwargs,
        )

        # fix point (must passed points)
        target_t = self.q_sample(target, t=batched_times)
        pred_img[partial_mask] = target_t[partial_mask]

        return pred_img

    @staticmethod
    def classifier_guidance(
        x: torch.Tensor,
        t: torch.Tensor,
        y: torch.Tensor,
        classifier: torch.nn.Module
    ):
        with torch.enable_grad():
            # 激活梯度计算
            x_with_grad = x.detach().requires_grad_(True)
            # 获取 log 形式的概率分布
            logits = classifier(x_with_grad, t)
            log_prob = F.log_softmax(logits, dim=-1)
            # 选取出 y 对应的项
            selected = log_prob[range(len(logits)), y.view(-1)]
            # 计算梯度
            return torch.autograd.grad(selected.sum(), x_with_grad)[0]

    @staticmethod
    def regression_guidance(
        x: torch.Tensor,
        t: torch.Tensor,
        target_sum: torch.Tensor,  # Target sum value
        sigma: float = 1.0
    ):
        """
        Compute gradient for guiding the sum of first channel to match target value
        
        Args:
            x: Input tensor [batch_size, channels, length] or [batch_size, length, channels]
            t: Time steps
            target_sum: Target sum value [batch_size]
            sigma: Standard deviation for Gaussian likelihood
        """
        # with torch.enable_grad():
        #     x_with_grad = x.detach().requires_grad_(True)
        # normalize to 0, 1
        # x_with_grad = (x + x.min()) / (x.max() - x.min())
        # x_with_grad = x / 2  + 0.5 # [-1,1 to 0,1]
        x_with_grad = x
        # Calculate sum of first channel/feature
        # Assuming x shape is [batch_size, channels, length] or [batch_size, length, channels]
        if x_with_grad.dim() == 3:
            if x_with_grad.shape[1] < x_with_grad.shape[2]:  # [B, C, L]
                current_sum = x_with_grad[:1, 0]
                current_sum = current_sum / 2 + 0.5  # [-1, 1 to 0, 1]
                print("Current Sum: ", current_sum.max().item(), current_sum.min().item())
                current_sum = current_sum.sum(dim=1)  # Sum over length
            else:  # [B, L, C]
                current_sum = x_with_grad[:1, :, 0]
                current_sum = current_sum / 2 + 0.5  # [-1, 1 to 0, 1]
                print("Current Sum: ", current_sum.max().item(), current_sum.min().item())
                current_sum = current_sum.sum(dim=1)  # Sum over length

        # Compute log probability under Gaussian distribution
        sigma = torch.log(t) / 5
        print("sigma", sigma)
        if sigma.mean() == 0:
            pred_std = torch.ones_like(current_sum)
        else:
            pred_std = torch.ones_like(current_sum) * sigma
        
        log_prob = -0.5 * torch.log(2 * torch.pi * pred_std**2) - \
                (target_sum - current_sum)**2 / (2 * pred_std**2)
        # print(target_sum, current_sum)
        # print("Current Sum: ", current_sum.mean().item())
        # print("Current Diff: ", (target_sum - current_sum).mean().item())
        return log_prob.mean()

        # return torch.autograd.grad(log_prob.sum(), x_with_grad)[0]


    def langevin_fn(
        self,
        coef,
        partial_mask,
        tgt_embs,
        learning_rate,
        sample,
        mean,
        sigma,
        t,
        coef_=0.0,
        gradient_control_signal={},
        model_control_signal={},
        **kwargs,
    ):
        # we thus run more gradient updates at large diffusion step t to guide the generation then
        # reduce the number of gradient steps in stages to accelerate sampling.
        if t[0].item() < self.num_timesteps * 0.02 :
            K = 0
        elif t[0].item() > self.num_timesteps * 0.9:
            K = 3
        elif t[0].item() > self.num_timesteps * 0.75:
            K = 2
            learning_rate = learning_rate * 0.5
        else:
            K = 1
            learning_rate = learning_rate * 0.25

        input_embs_param = torch.nn.Parameter(sample)

        # 获取时间相关的权重调整因子
        time_weight = get_time_dependent_weights(t[0], self.num_timesteps)

        with torch.enable_grad():
            for iteration in range(K):
                # x_i+1 = x_i + noise * grad(logp(x_i)) + sqrt(2*noise) * z_i
                optimizer = torch.optim.Adagrad([input_embs_param], lr=learning_rate)
                optimizer.zero_grad()

                x_start = self.output(
                    x=input_embs_param,
                    t=t,
                    control_signal=model_control_signal,
                )

                if sigma.mean() == 0:
                    logp_term = (
                        coef * ((mean - input_embs_param) ** 2 / 1.0).mean(dim=0).sum()
                    )
                    # determine the partical_mask is float
                    if kwargs.get("enable_float_mask", False):
                        infill_loss = (x_start * (partial_mask) - tgt_embs * (partial_mask)) ** 2
                    else:
                        infill_loss = (x_start[partial_mask] - tgt_embs[partial_mask]) ** 2
                    infill_loss = infill_loss.mean(dim=0).sum()
                else:
                    logp_term = (
                        coef
                        * ((mean - input_embs_param) ** 2 / sigma).mean(dim=0).sum()
                    )
                    if kwargs.get("enable_float_mask", False):
                        infill_loss = (x_start * (partial_mask) - tgt_embs * (partial_mask)) ** 2
                    else:
                        infill_loss = (x_start[partial_mask] - tgt_embs[partial_mask]) ** 2
                    infill_loss = (infill_loss / sigma.mean()).mean(dim=0).sum()
                # 第二个等号后面最后一项消失了，因为当我们要求模型生成“狗”的图像时，扩散过程始终
                # 不变，对应的梯度也是0，可以抹掉。
                # https://lichtung612.github.io/posts/3-diffusion-models/
                # 第三个等号后面两项中，第一项是扩散模型本身的梯度引导，新增的只能是第二项，即classifier guidance只需要额外添加一个classifier的梯度来引导。

                # 控制信号损失
                gradient_scale = gradient_control_signal.get("gradient_scale", 1.0)  # 全局梯度缩放因子
                
                # Add regression guidance for sum constraint
                control_loss = 0
                # target_sum = 
                # normalize the sum to -1, 1
                # seq_length = input_embs_param.shape[1]
                # target_sum = ((target_sum / seq_length ) * 2 - 1) * seq_length
                # if target_sum:=gradient_control_signal.get("sum") is not None:
                #     # print("sigma", sigma.shape, sigma, end="  ")
                #     reg_nll = self.regression_guidance(
                #         x=input_embs_param,
                #         t=t,
                #         target_sum=target_sum,
                #         sigma=sigma
                #     )
                #     control_loss += - gradient_control_signal.get("reg_weight", 1.0) * reg_nll * (5 - K) # (reg_gradient * ).sum()

                # init control signal loss
                auc_sum, peak_points, bar_regions, target_freq = \
                    gradient_control_signal.get("auc"), gradient_control_signal.get("peak_points"), gradient_control_signal.get("bar_regions"), gradient_control_signal.get("target_freq")

                # 1. 原有的sum控制
                if auc_sum is not None:
                    sum_weight = gradient_control_signal.get("auc_weight", 1.0) * time_weight
                    auc_loss = - sum_weight * sum_guidance(
                        x=input_embs_param,
                        t=t,
                        target_sum=auc_sum,
                        gradient_scale=gradient_scale,
                        segments=gradient_control_signal.get("segments", ())
                    )
                    control_loss += auc_loss

                # 峰值引导
                if peak_points is not None:
                    peak_weight = gradient_control_signal.get("peak_weight", 1.0) * time_weight
                    peak_loss = - peak_weight * peak_guidance(
                        x=input_embs_param,
                        t=t,
                        peak_points=peak_points,
                        window_size=gradient_control_signal.get("peak_window_size", 5),
                        alpha_1=gradient_control_signal.get("peak_alpha_1", 1.2),
                        gradient_scale=gradient_scale
                    )
                    control_loss += peak_loss

                # 区间引导
                if bar_regions is not None:
                    bar_weight = gradient_control_signal.get("bar_weight", 1.0) * time_weight
                    bar_loss = -bar_weight * bar_guidance(
                        x=input_embs_param,
                        t=t,
                        bar_regions=bar_regions,
                        gradient_scale=gradient_scale
                    )
                    control_loss += bar_loss

                # 频率引导
                if target_freq is not None:
                    freq_weight = gradient_control_signal.get("freq_weight", 1.0) * time_weight
                    freq_loss = -freq_weight * frequency_guidance(
                        x=input_embs_param,
                        t=t,
                        target_freq=target_freq,
                        freq_weight=freq_weight,
                        gradient_scale=gradient_scale
                    )
                    control_loss += freq_loss


                loss = logp_term + infill_loss + control_loss
                if iteration == 0: # Only print first iteration to avoid spam
                    # print(f"Losses - Diffusion: {logp_term:.4f}, Infill: {infill_loss:.4f}, Control: {control_loss:.4f}")
                    # if target_sum is not None:
                    #     # Print current sum vs target for monitoring
                    #     if x_start.shape[1] < x_start.shape[2]:  # [B, C, L]
                    #         current_sum = input_embs_param[:, 0].sum(dim=1)
                    #     else:  # [B, L, C]
                    #         current_sum = input_embs_param[:, :, 0].sum(dim=1)
                    #     print(f"Current sum: {current_sum.data}, Target sum: {target_sum}")
                    # print(f"Losses - Diffusion: {logp_term:.4f}\tInfill: {infill_loss:.4f}", end="\t")
                    # if auc_sum is not None:
                    #     print(f"Sum Control: {auc_loss.item():.4f}", end="\t")
                    # if peak_points is not None:
                    #     print(f"Peak Control: {peak_loss.item():.4f}", end="\t")
                    # if bar_regions is not None:
                    #     print(f"Bar Control: {bar_loss.item():.4f}", end="\t")
                    # if target_freq is not None:
                    #     print(f"Freq Control: {freq_loss.item():.4f}", end="\t")
                    # print()
                    pass

                # loss = logp_term + infill_loss + auc_loss
                # print(logp_term, infill_loss, auc_loss)
                loss.backward()
                optimizer.step()
                torch.nn.utils.clip_grad_norm_([input_embs_param], gradient_control_signal.get("max_grad_norm", 1.0))

                # add more noise
                epsilon = torch.randn_like(input_embs_param.data)
                noise_scale = coef_ * sigma.mean().item() # * 2
                # noise_scale = noise_scale * time_weight # (1 - time_weight)  # 随时间减少噪声
                input_embs_param = torch.nn.Parameter(
                    (
                        input_embs_param.data + noise_scale * epsilon
                    ).detach()
                )

        if kwargs.get("enable_float_mask", False):
            sample = sample * partial_mask + input_embs_param.data * (1 - partial_mask)
        else:
            sample[~partial_mask] = input_embs_param.data[~partial_mask]

        # if t[0].item() % 10 == 9:
        #     print("Sampled Image")
        #     images_cache.append(plt.plot(sample[0,:,0].detach().cpu().numpy())[0])
            
        #     if t[0].item() == 9:
        #         plt.show()
        #         images_cache.clear()
            # plt.show()
            # plt.savefig(f"sampled_{t[0].item()}.png")
        # plt.plot(sample[0,:,0].detach().cpu().numpy())
        # plt.show()
        return sample

    # def load_weights(self, model_path):
    #     data = torch.load(model_path, map_location="cuda:0", weights_only=True)
    #     self.load_state_dict(data["model"])
    #     print("Model weights loaded successfully")

    def predict_weighted_points(
        self,
        observed_points: torch.Tensor,
        observed_mask: torch.Tensor,
        coef=1e-1,
        stepsize=1e-1,
        sampling_steps=50,
        **kargs,
    ):
        model_kwargs = {}
        model_kwargs["coef"] = coef
        model_kwargs["learning_rate"] = stepsize
        model_kwargs = {**model_kwargs, **kargs}
        assert len(observed_points.shape) == 2, "observed_points should be 2D, batch size = 1"
        x = observed_points.unsqueeze(0)
        float_mask = observed_mask.unsqueeze(0) # x != 0, 1 for observed, 0 for missing, bool tensor
        binary_mask = float_mask.clone()
        binary_mask[binary_mask > 0] = 1

        x = x * 2 - 1 # normalize
        self.device = x.device
        x, float_mask, binary_mask = x.to(self.device), float_mask.to(self.device), binary_mask.to(self.device)
        if sampling_steps == self.num_timesteps:
            print("normal sampling")
            raise NotImplementedError
            sample = self.ema.ema_model.sample_infill_float_mask(
                shape=x.shape,
                target=x * binary_mask, # x * t_m, 1 for observed, 0 for missing
                partial_mask=float_mask,
                model_kwargs=model_kwargs,
            )
            # x: partially noise : (batch_size, seq_length, feature_dim)
        else:
            print("fast sampling")
            sample = self.fast_sample_infill_float_mask(
                shape=x.shape,
                target=x * binary_mask, # x * t_m, 1 for observed, 0 for missing
                partial_mask=float_mask,
                model_kwargs=model_kwargs,
                sampling_timesteps=sampling_steps,
            )

        # unnormalize
        sample = (sample + 1) / 2
        return sample.squeeze(0).detach().cpu().numpy()


    def register_model(self, registered_model_name, model_path="tiffusion_model", conda_env=None):
        """Register the model with MLflow model registry.
        
        Args:
            registered_model_name: Name to register the model under
            model_path: Local path to save model artifacts
            conda_env: Custom conda environment for the model
        """
        # Create basic conda env if not provided
        if conda_env is None:
            conda_env = {
                'channels': ['defaults', 'conda-forge'],
                'dependencies': [
                    'python>=3.8',
                    'pytorch',
                    'einops',
                    'tqdm'
                ],
                'name': 'tiffusion_env'
            }

        # Start an MLflow run
        with mlflow.start_run() as run:
            # Log model parameters
            mlflow.log_params({
                "seq_length": self.seq_length,
                "feature_size": self.feature_size,
                "n_layer_enc": self.model.n_layer_enc,
                "n_layer_dec": self.model.n_layer_dec,
                "n_heads": self.model.n_heads,
                "timesteps": self.num_timesteps,
                "loss_type": self.loss_type
            })

            # Create a custom Python model class for MLflow
            class TiffusionWrapper(mlflow.pyfunc.PythonModel):
                def __init__(self, model):
                    self.model = model
                
                def predict(self, context, model_input):
                    # Generate predictions using the model
                    with torch.no_grad():
                        result = self.model.generate_mts(batch_size=len(model_input))
                    return result.numpy()

            # Create wrapper instance
            wrapped_model = TiffusionWrapper(self)

            # Log and register the model
            mlflow.pyfunc.log_model(
                artifact_path=model_path,
                python_model=wrapped_model,
                conda_env=conda_env,
                registered_model_name=registered_model_name
            )

            print(f"Model registered as: {registered_model_name}")
            print(f"Run ID: {run.info.run_id}")

    def load_registered_model(self, registered_model_name, version=None, stage=None):
        """Load a registered model from MLflow model registry.
        
        Args:
            registered_model_name: Name of registered model
            version: Optional specific version to load
            stage: Optional stage to load (e.g. 'Production', 'Staging')
        """
        if version:
            model_uri = f"models:/{registered_model_name}/{version}"
        elif stage:
            model_uri = f"models:/{registered_model_name}/{stage}"
        else:
            model_uri = f"models:/{registered_model_name}/latest"
            
        return mlflow.pyfunc.load_model(model_uri)