# Copyright (c) 2025 Hansheng Chen import sys import inspect import torch import diffusers import mmcv from typing import Optional from copy import deepcopy from mmgen.models.architectures.common import get_module_device from mmgen.models.builder import MODULES, build_module from . import GaussianFlow, schedulers from lakonlab.ops.gmflow_ops.gmflow_ops import ( gm_to_sample, gm_to_mean, gaussian_samples_to_gm_samples, gm_samples_to_gaussian_samples, iso_gaussian_mul_iso_gaussian, gm_mul_iso_gaussian, gm_to_iso_gaussian, gm_transpose_t_first) @torch.jit.script def probabilistic_guidance_jit( cond_mean, total_var, uncond_mean, guidance_scale: float, orthogonal: float = 1.0, orthogonal_axis: Optional[torch.Tensor] = None): dim = list(range(1, cond_mean.dim())) bias = cond_mean - uncond_mean if orthogonal > 0.0: if orthogonal_axis is None: orthogonal_axis = cond_mean bias = bias - ((bias * orthogonal_axis).mean( dim=dim, keepdim=True ) / (orthogonal_axis * orthogonal_axis).mean( dim=dim, keepdim=True ).clamp(min=1e-6) * orthogonal_axis).mul(orthogonal) bias_power = (bias * bias).mean(dim=dim, keepdim=True) avg_var = total_var.mean(dim=dim, keepdim=True) bias = bias * ((avg_var / bias_power.clamp(min=1e-6)).sqrt() * guidance_scale) gaussian_output = dict( mean=cond_mean + bias, var=total_var * (1 - (guidance_scale * guidance_scale))) return gaussian_output, bias, avg_var @torch.jit.script def gmflow_posterior_jit( sigma_t_src, sigma_t, x_t_src, x_t, gm_means, gm_vars, gm_logstds, gm_logweights, eps: float, gm_dim: int = -4, channel_dim: int = -3): alpha_t_src = 1 - sigma_t_src alpha_t = 1 - sigma_t sigma_t_src_sq = sigma_t_src.square() sigma_t_sq = sigma_t.square() # compute gaussian params denom = (alpha_t.square() * sigma_t_src_sq - alpha_t_src.square() * sigma_t_sq).clamp(min=eps) # ζ g_mean = (alpha_t * sigma_t_src_sq * x_t - alpha_t_src * sigma_t_sq * x_t_src) / denom # ν / ζ g_var = sigma_t_sq * sigma_t_src_sq / denom # gm_mul_iso_gaussian g_mean = g_mean.unsqueeze(gm_dim) # (bs, *, 1, out_channels, h, w) g_var = g_var.unsqueeze(gm_dim) # (bs, *, 1, 1, 1, 1) g_logstd = g_var.clamp(min=eps).log() / 2 gm_diffs = gm_means - g_mean # (bs, *, num_gaussians, out_channels, h, w) norm_factor = (g_var + gm_vars).clamp(min=eps) out_means = (g_var * gm_means + gm_vars * g_mean) / norm_factor # (bs, *, num_gaussians, 1, h, w) logweights_delta = gm_diffs.square().sum(dim=channel_dim, keepdim=True) * (-0.5 / norm_factor) out_logweights = (gm_logweights + logweights_delta).log_softmax(dim=gm_dim) out_logstds = gm_logstds + g_logstd - 0.5 * torch.log(norm_factor) return out_means, out_logstds, out_logweights @torch.jit.script def gmflow_posterior_mean_jit( sigma_t_src, sigma_t, x_t_src, x_t, gm_means, gm_vars, gm_logweights, eps: float, gm_dim: int = -4, channel_dim: int = -3): alpha_t_src = 1 - sigma_t_src alpha_t = 1 - sigma_t sigma_t_src_sq = sigma_t_src.square() sigma_t_sq = sigma_t.square() # compute gaussian params denom = (alpha_t.square() * sigma_t_src_sq - alpha_t_src.square() * sigma_t_sq).clamp(min=eps) # ζ g_mean = (alpha_t * sigma_t_src_sq * x_t - alpha_t_src * sigma_t_sq * x_t_src) / denom # ν / ζ g_var = sigma_t_sq * sigma_t_src_sq / denom # gm_mul_iso_gaussian g_mean = g_mean.unsqueeze(gm_dim) # (bs, *, 1, out_channels, h, w) g_var = g_var.unsqueeze(gm_dim) # (bs, *, 1, 1, 1, 1) gm_diffs = gm_means - g_mean # (bs, *, num_gaussians, out_channels, h, w) norm_factor = (g_var + gm_vars).clamp(min=eps) out_means = (g_var * gm_means + gm_vars * g_mean) / norm_factor # (bs, *, num_gaussians, 1, h, w) logweights_delta = gm_diffs.square().sum(dim=channel_dim, keepdim=True) * (-0.5 / norm_factor) out_weights = (gm_logweights + logweights_delta).softmax(dim=gm_dim) out_mean = (out_means * out_weights).sum(dim=gm_dim) return out_mean class GMFlowMixin: @property def time_scaling(self): if hasattr(self, 'scheduler'): # for diffusers pipelines return self.scheduler.config.num_train_timesteps elif hasattr(self, 'num_timesteps'): return self.num_timesteps else: raise ValueError('num_timesteps or scheduler is not defined.') def u_to_x_0(self, denoising_output, x_t, t=None, sigma=None, eps=1e-6): if isinstance(denoising_output, dict) and 'logweights' in denoising_output: x_t = x_t.unsqueeze(-4) if sigma is None: if not isinstance(t, torch.Tensor): t = torch.tensor(t, device=x_t.device) t = t.reshape(*t.size(), *((x_t.dim() - t.dim()) * [1])) sigma = t / self.time_scaling else: assert sigma.dim() == x_t.dim() - 1 sigma = sigma.unsqueeze(-4) if isinstance(denoising_output, dict): if 'logweights' in denoising_output: means_x_0 = x_t - sigma * denoising_output['means'] logstds_x_0 = denoising_output['logstds'] + torch.log(sigma.clamp(min=eps)) return dict( means=means_x_0, logstds=logstds_x_0, logweights=denoising_output['logweights']) elif 'var' in denoising_output: mean = x_t - sigma * denoising_output['mean'] var = denoising_output['var'] * sigma.square() return dict(mean=mean, var=var) else: raise ValueError('Invalid denoising_output.') else: # sample mode x_0 = x_t - sigma * denoising_output return x_0 def gmflow_posterior_mean( self, gm, x_t, x_t_src, t=None, t_src=None, sigma_t_src=None, sigma_t=None, eps=1e-6, prediction_type='x0', checkpointing=False): """ Fuse gmflow_posterior and gm_to_mean to avoid redundant computation. """ assert isinstance(gm, dict) if sigma_t_src is None: if not isinstance(t_src, torch.Tensor): t_src = torch.tensor(t_src, device=x_t_src.device) t_src = t_src.reshape(*t_src.size(), *((x_t_src.dim() - t_src.dim()) * [1])) sigma_t_src = t_src / self.time_scaling if sigma_t is None: if not isinstance(t, torch.Tensor): t = torch.tensor(t, device=x_t_src.device) t = t.reshape(*t.size(), *((x_t_src.dim() - t.dim()) * [1])) sigma_t = t / self.time_scaling if prediction_type == 'u': gm = self.u_to_x_0(gm, x_t_src, sigma=sigma_t_src) else: assert prediction_type == 'x0' gm_means = gm['means'] # (bs, *, num_gaussians, out_channels, h, w) gm_logweights = gm['logweights'] # (bs, *, num_gaussians, 1, h, w) if 'gm_vars' in gm: gm_vars = gm['gm_vars'] else: gm_vars = (gm['logstds'] * 2).exp() # (bs, *, 1, 1, 1, 1) gm['gm_vars'] = gm_vars if checkpointing and torch.is_grad_enabled(): return torch.utils.checkpoint.checkpoint( gmflow_posterior_mean_jit, sigma_t_src, sigma_t, x_t_src, x_t, gm_means, gm_vars, gm_logweights, eps, use_reentrant=True) # use_reentrant=False does not work with jit else: return gmflow_posterior_mean_jit( sigma_t_src, sigma_t, x_t_src, x_t, gm_means, gm_vars, gm_logweights, eps) def reverse_transition(self, denoising_output, x_t_high, t_low, t_high, eps=1e-6, prediction_type='u'): if isinstance(denoising_output, dict): x_t_high = x_t_high.unsqueeze(-4) bs = x_t_high.size(0) if not isinstance(t_low, torch.Tensor): t_low = torch.tensor(t_low, device=x_t_high.device) if not isinstance(t_high, torch.Tensor): t_high = torch.tensor(t_high, device=x_t_high.device) if t_low.dim() == 0: t_low = t_low.expand(bs) if t_high.dim() == 0: t_high = t_high.expand(bs) t_low = t_low.reshape(*t_low.size(), *((x_t_high.dim() - t_low.dim()) * [1])) t_high = t_high.reshape(*t_high.size(), *((x_t_high.dim() - t_high.dim()) * [1])) sigma = t_high / self.time_scaling sigma_to = t_low / self.time_scaling alpha = 1 - sigma alpha_to = 1 - sigma_to sigma_to_over_sigma = sigma_to / sigma.clamp(min=eps) alpha_over_alpha_to = alpha / alpha_to.clamp(min=eps) beta_over_sigma_sq = 1 - (sigma_to_over_sigma * alpha_over_alpha_to) ** 2 c1 = sigma_to_over_sigma ** 2 * alpha_over_alpha_to c2 = beta_over_sigma_sq * alpha_to if isinstance(denoising_output, dict): c3 = beta_over_sigma_sq * sigma_to ** 2 if prediction_type == 'u': means_x_0 = x_t_high - sigma * denoising_output['means'] logstds_x_t_low = torch.logaddexp( (denoising_output['logstds'] + torch.log((sigma * c2).clamp(min=eps))) * 2, torch.log(c3.clamp(min=eps)) ) / 2 elif prediction_type == 'x0': means_x_0 = denoising_output['means'] logstds_x_t_low = torch.logaddexp( (denoising_output['logstds'] + torch.log(c2.clamp(min=eps))) * 2, torch.log(c3.clamp(min=eps)) ) / 2 else: raise ValueError('Invalid prediction_type.') means_x_t_low = c1 * x_t_high + c2 * means_x_0 return dict( means=means_x_t_low, logstds=logstds_x_t_low, logweights=denoising_output['logweights']) else: # sample mode c3_sqrt = beta_over_sigma_sq ** 0.5 * sigma_to noise = torch.randn_like(denoising_output) if prediction_type == 'u': x_0 = x_t_high - sigma * denoising_output elif prediction_type == 'x0': x_0 = denoising_output else: raise ValueError('Invalid prediction_type.') x_t_low = c1 * x_t_high + c2 * x_0 + c3_sqrt * noise return x_t_low @staticmethod def gm_sample(gm, power_spectrum=None, n_samples=1, generator=None): device = gm['means'].device if power_spectrum is not None: power_spectrum = power_spectrum.to(dtype=torch.float32).unsqueeze(-4) shape = list(gm['means'].size()) shape[-4] = n_samples half_size = shape[-1] // 2 + 1 spectral_samples = torch.randn( shape, dtype=torch.float32, device=device, generator=generator) * (power_spectrum / 2).exp() z_1 = spectral_samples.roll((-1, -1), dims=(-2, -1)).flip((-2, -1))[..., :half_size] z_0 = spectral_samples[..., :half_size] z_kr = torch.complex(z_0 + z_1, z_0 - z_1) / 2 gaussian_samples = torch.fft.irfft2(z_kr, norm='ortho') samples = gaussian_samples_to_gm_samples(gm, gaussian_samples, axis_aligned=True) else: samples = gm_to_sample(gm, n_samples=n_samples) spectral_samples = None return samples, spectral_samples def gm_to_model_output(self, gm, output_mode, power_spectrum=None, generator=None): assert output_mode in ['mean', 'sample'] if output_mode == 'mean': output = gm_to_mean(gm) else: # sample output = self.gm_sample(gm, power_spectrum, generator=generator)[0].squeeze(-4) return output def gm_2nd_order( self, gm_output, gaussian_output, x_t, t, h, guidance_scale=0.0, gm_cond=None, gaussian_cond=None, avg_var=None, cfg_bias=None, ca=0.005, cb=1.0, gm2_correction_steps=0): if self.prev_gm is not None: dim = list(range(1, x_t.dim())) if cfg_bias is not None: gm_mean = gm_to_mean(gm_output) base_gaussian = gaussian_cond base_gm = gm_cond else: gm_mean = gaussian_output['mean'] base_gaussian = gaussian_output base_gaussian['var'] = base_gaussian['var'].mean(dim=dim[:-3] + dim[-2:], keepdim=True) # exclude channel dim avg_var = base_gaussian['var'].mean(dim=dim, keepdim=True) base_gm = gm_output mean_from_prev = self.gmflow_posterior_mean( self.prev_gm, x_t, self.prev_x_t, t, self.prev_t, prediction_type='x0') self.prev_gm = gm_output # Compute rescaled 2nd-order mean difference k = 0.5 * h / self.prev_h prev_h_norm = self.prev_h / self.time_scaling _guidance_scale = guidance_scale * cb err_power = avg_var * (_guidance_scale * _guidance_scale + ca) mean_diff = (gm_mean - mean_from_prev) * ( (1 - err_power / (prev_h_norm * prev_h_norm)).clamp(min=0).sqrt() * k) bias = mean_diff # Here we fuse probabilistic guidance bias and 2nd-order mean difference to perform one single # update to the base GM, which avoids cumulative errors. if cfg_bias is not None: bias = mean_diff + cfg_bias bias_power = bias.square().mean(dim=dim, keepdim=True) bias = bias * (avg_var / bias_power.clamp(min=1e-6)).clamp(max=1).sqrt() gaussian_output = dict( mean=base_gaussian['mean'] + bias, var=base_gaussian['var'] * (1 - bias_power / avg_var.clamp(min=1e-6)).clamp(min=1e-6)) gm_output = gm_mul_iso_gaussian( base_gm, iso_gaussian_mul_iso_gaussian(gaussian_output, base_gaussian, 1, -1), 1, 1)[0] # Additional correction steps for strictly matching the 2nd order mean difference if gm2_correction_steps > 0: adjusted_bias = bias tgt_bias = mean_diff + gm_mean - base_gaussian['mean'] for _ in range(gm2_correction_steps): out_bias = gm_to_mean(gm_output) - base_gaussian['mean'] err = out_bias - tgt_bias adjusted_bias = adjusted_bias - err * ( adjusted_bias.norm(dim=-3, keepdim=True) / out_bias.norm(dim=-3, keepdim=True).clamp(min=1e-6) ).clamp(max=1) adjusted_bias_power = adjusted_bias.square().mean(dim=dim, keepdim=True) adjusted_bias = adjusted_bias * (avg_var / adjusted_bias_power.clamp(min=1e-6)).clamp(max=1).sqrt() adjusted_gaussian_output = dict( mean=base_gaussian['mean'] + adjusted_bias, var=base_gaussian['var'] * (1 - adjusted_bias_power / avg_var.clamp(min=1e-6)).clamp(min=1e-6)) gm_output = gm_mul_iso_gaussian( base_gm, iso_gaussian_mul_iso_gaussian(adjusted_gaussian_output, base_gaussian, 1, -1), 1, 1)[0] else: self.prev_gm = gm_output self.prev_x_t = x_t self.prev_t = t self.prev_h = h return gm_output, gaussian_output def init_gm_cache(self): self.prev_gm = None self.prev_x_t = None self.prev_t = None self.prev_h = None @MODULES.register_module() class GMFlow(GaussianFlow, GMFlowMixin): def __init__( self, *args, spectrum_net=None, spectral_loss_weight=1.0, **kwargs): super().__init__(*args, **kwargs) self.spectrum_net = build_module(spectrum_net) if spectrum_net is not None else None self.spectral_loss_weight = spectral_loss_weight self.intermediate_x_t = [] self.intermediate_x_0 = [] def loss(self, denoising_output, x_t_low, x_t_high, t_low, t_high): """ GMFlow transition loss. """ x_t_low = x_t_low.float() x_t_high = x_t_high.float() t_low = t_low.float() t_high = t_high.float() x_t_low_gm = self.reverse_transition(denoising_output, x_t_high, t_low, t_high) loss_kwargs = {k: v for k, v in x_t_low_gm.items()} loss_kwargs.update(x_t_low=x_t_low, timesteps=t_high) return self.flow_loss(loss_kwargs) def spectral_loss(self, denoising_output, x_0, x_t, t, eps=1e-6): x_0 = x_0.float() x_t = x_t.float() t = t.float() t = t.reshape(*t.size(), *((x_t.dim() - t.dim()) * [1])) inv_sigma = self.num_timesteps / t.clamp(min=eps) with torch.no_grad(): output_g = self.u_to_x_0(gm_to_iso_gaussian(denoising_output)[0], x_t, t) u = (x_t - x_0) * inv_sigma z_kr = gm_samples_to_gaussian_samples( denoising_output, u.unsqueeze(-4), axis_aligned=True).squeeze(-4) z_kr_fft = torch.fft.fft2(z_kr, norm='ortho') z_kr_fft = z_kr_fft.real + z_kr_fft.imag log_var = self.spectrum_net(output_g) loss = z_kr_fft.square() * (torch.exp(-log_var) - 1) + log_var loss = loss.mean() * (0.5 * self.spectral_loss_weight) return loss def pred(self, x_t=None, t=None, **kwargs): ndim = x_t.dim() assert ndim in [4, 5], f'Invalid x_t shape: {x_t.shape}. Expected 4D or 5D tensor.' if ndim == 5: # (bs, t, c, h, w) x_t = x_t.permute(0, 2, 1, 3, 4) # (bs, c, t, h, w) output = super().pred(x_t=x_t, t=t, **kwargs) if ndim == 5: output = gm_transpose_t_first(output) # (bs, t, c, h, w) return output def forward_train(self, x_0, **kwargs): device = get_module_device(self) num_batches = x_0.size(0) seq_len = x_0.shape[2:].numel() # h * w or t * h * w ndim = x_0.dim() assert ndim in [4, 5], f'Invalid x_0 shape: {x_0.shape}. Expected 4D or 5D tensor.' if ndim == 5: # (bs, c, t, h, w) x_0 = x_0.permute(0, 2, 1, 3, 4) # (bs, t, c, h, w) trans_ratio = self.train_cfg.get('trans_ratio', 1.0) eps = self.train_cfg.get('eps', 1e-4) t_high = self.timestep_sampler( num_batches, seq_len=seq_len).to(device).clamp(min=eps, max=self.num_timesteps) t_low = t_high * (1 - trans_ratio) t_low = torch.minimum(t_low, t_high - eps).clamp(min=0) noise = torch.randn((num_batches * 2, *x_0.shape[1:]), device=device, dtype=x_0.dtype) noise_0, noise_1 = torch.chunk(noise, 2, dim=0) x_t_low, _, _ = self.sample_forward_diffusion(x_0, t_low, noise_0) x_t_high = self.sample_forward_transition(x_t_low, noise_1, t_src=t_low, t_tgt=t_high) denoising_output = self.pred(x_t_high, t_high, **kwargs) loss = self.loss(denoising_output, x_t_low, x_t_high, t_low, t_high) log_vars = self.flow_loss.log_vars log_vars.update(loss_transition=float(loss)) if self.spectrum_net is not None: # Note: only support 2D power spectrum for now. loss_spectral = self.spectral_loss(denoising_output, x_0, x_t_high, t_high) log_vars.update(loss_spectral=float(loss_spectral)) loss = loss + loss_spectral return loss, log_vars def forward_test( self, x_0=None, noise=None, guidance_scale=0.0, test_cfg_override=dict(), show_pbar=False, **kwargs): x_t = torch.randn_like(x_0) if noise is None else noise num_batches = x_t.size(0) seq_len = x_t.shape[2:].numel() # h * w or t * h * w ori_dtype = x_t.dtype x_t = x_t.float() ndim = x_t.dim() assert ndim in [4, 5], f'Invalid x_t shape: {x_t.shape}. Expected 4D or 5D tensor.' if ndim == 5: # (bs, c, t, h, w) x_t = x_t.permute(0, 2, 1, 3, 4) # (bs, t, c, h, w) cfg = deepcopy(self.test_cfg) cfg.update(test_cfg_override) output_mode = cfg.get('output_mode', 'mean') assert output_mode in ['mean', 'sample'] sampler = cfg['sampler'] sampler_class = getattr(diffusers.schedulers, sampler + 'Scheduler', None) if sampler_class is None: sampler_class = getattr(schedulers, sampler + 'Scheduler', None) if sampler_class is None: raise AttributeError(f'Cannot find sampler [{sampler}].') sampler_kwargs = cfg.get('sampler_kwargs', {}) signatures = inspect.signature(sampler_class).parameters.keys() for key in ['shift', 'use_dynamic_shifting', 'base_seq_len', 'max_seq_len', 'base_logshift', 'max_logshift']: if key in signatures and key not in sampler_kwargs: sampler_kwargs[key] = cfg.get(key, getattr(self.timestep_sampler, key)) sampler = sampler_class(self.num_timesteps, **sampler_kwargs) num_timesteps = cfg.get('num_timesteps', self.num_timesteps) num_substeps = cfg.get('num_substeps', 1) guidance_interval = cfg.get('guidance_interval', [0, self.num_timesteps]) orthogonal_guidance = cfg.get('orthogonal_guidance', 1.0) save_intermediate = cfg.get('save_intermediate', False) order = cfg.get('order', 1) gm2_coefs = cfg.get('gm2_coefs', [0.005, 1.0]) gm2_correction_steps = cfg.get('gm2_correction_steps', 0) set_timesteps_signatures = inspect.signature(sampler.set_timesteps).parameters.keys() if 'seq_len' in set_timesteps_signatures: sampler.set_timesteps(num_timesteps * num_substeps, seq_len=seq_len, device=x_t.device) else: sampler.set_timesteps(num_timesteps * num_substeps, device=x_t.device) timesteps = sampler.timesteps self.intermediate_x_t = [] self.intermediate_x_0 = [] self.intermediate_gm_x_0 = [] self.intermediate_gm_trans = [] self.intermediate_t = [] use_guidance = 0.0 < guidance_scale < 1.0 assert order in [1, 2] if show_pbar: pbar = mmcv.ProgressBar(num_timesteps) # ========== Main sampling loop ========== self.init_gm_cache() for timestep_id in range(num_timesteps): t = timesteps[timestep_id * num_substeps] if save_intermediate: self.intermediate_x_t.append(x_t) self.intermediate_t.append(t) x_t_input = x_t _kwargs = kwargs if use_guidance: guidance_active = guidance_interval[0] <= t <= guidance_interval[1] if guidance_active: x_t_input = torch.cat([x_t_input, x_t_input], dim=0) else: _kwargs = { k: v[num_batches:] if isinstance(v, torch.Tensor) and v.size(0) == 2 * num_batches else v for k, v in kwargs.items()} gm_output = self.pred(x_t_input, t, **_kwargs) assert isinstance(gm_output, dict) gm_output = self.u_to_x_0(gm_output, x_t_input, t) # ========== Probabilistic CFG ========== if use_guidance and guidance_active: gm_cond = {k: v[num_batches:] for k, v in gm_output.items()} gm_uncond = {k: v[:num_batches] for k, v in gm_output.items()} uncond_mean = gm_to_mean(gm_uncond) gaussian_cond = gm_to_iso_gaussian(gm_cond)[0] if ndim == 5: gaussian_cond['var'] = gaussian_cond['var'].mean(dim=(-4, -2, -1), keepdim=True) # exclude channel dim else: gaussian_cond['var'] = gaussian_cond['var'].mean(dim=(-2, -1), keepdim=True) # exclude channel dim gaussian_output, cfg_bias, avg_var = probabilistic_guidance_jit( gaussian_cond['mean'], gaussian_cond['var'], uncond_mean, guidance_scale, orthogonal=orthogonal_guidance) gm_output = gm_mul_iso_gaussian( gm_cond, iso_gaussian_mul_iso_gaussian(gaussian_output, gaussian_cond, 1, -1), 1, 1)[0] else: gaussian_output = gm_to_iso_gaussian(gm_output)[0] gm_cond = gaussian_cond = avg_var = cfg_bias = None # ========== 2nd order GM ========== if order == 2: if timestep_id < num_timesteps - 1: h = t - timesteps[(timestep_id + 1) * num_substeps] else: h = t gm_output, gaussian_output = self.gm_2nd_order( gm_output, gaussian_output, x_t, t, h, guidance_scale if guidance_active else 0.0, gm_cond, gaussian_cond, avg_var, cfg_bias, ca=gm2_coefs[0], cb=gm2_coefs[1], gm2_correction_steps=gm2_correction_steps) if save_intermediate: self.intermediate_gm_x_0.append(gm_output) if timestep_id < num_timesteps - 1: t_next = timesteps[(timestep_id + 1) * num_substeps] else: t_next = 0 gm_trans = self.reverse_transition(gm_output, x_t, t_next, t, prediction_type='x0') self.intermediate_gm_trans.append(gm_trans) # ========== GM SDE step or GM ODE substeps ========== x_t_base = x_t t_base = t for substep_id in range(num_substeps): if substep_id == 0: if self.spectrum_net is not None and output_mode == 'sample': # Note: only support 2D power spectrum for now. power_spectrum = self.spectrum_net(gaussian_output) else: power_spectrum = None model_output = self.gm_to_model_output(gm_output, output_mode, power_spectrum=power_spectrum) else: assert output_mode == 'mean' t = timesteps[timestep_id * num_substeps + substep_id] model_output = self.gmflow_posterior_mean( gm_output, x_t, x_t_base, t, t_base, prediction_type='x0') x_t = sampler.step(model_output, t, x_t, return_dict=False, prediction_type='x0')[0] if save_intermediate: self.intermediate_x_0.append(model_output) if show_pbar: pbar.update() if show_pbar: sys.stdout.write('\n') if ndim == 5: # (bs, t, c, h, w) x_t = x_t.permute(0, 2, 1, 3, 4) # (bs, c, t, h, w) return x_t.to(ori_dtype) def forward_u(self, x_t, t, guidance_scale=0.0, test_cfg_override=dict(), **kwargs): ori_dtype = x_t.dtype x_t = x_t.float() ndim = x_t.dim() assert ndim in [4, 5], f'Invalid x_t shape: {x_t.shape}. Expected 4D or 5D tensor.' if ndim == 5: # (bs, c, t, h, w) x_t = x_t.permute(0, 2, 1, 3, 4) # (bs, t, c, h, w) cfg = deepcopy(self.test_cfg) cfg.update(test_cfg_override) orthogonal_guidance = cfg.get('orthogonal_guidance', 1.0) guidance_interval = cfg.get('guidance_interval', [0, self.num_timesteps]) use_guidance = 0.0 < guidance_scale < 1.0 x_t_input = x_t t_input = t if use_guidance: x_t_input = torch.cat([x_t_input, x_t_input], dim=0) t_input = torch.cat([t_input, t_input], dim=0) gm_output = self.pred(x_t_input, t_input, **kwargs) assert isinstance(gm_output, dict) # ========== Probabilistic CFG ========== if use_guidance: num_batches = x_t.size(0) gm_cond = {k: v[num_batches:] for k, v in gm_output.items()} gm_uncond = {k: v[:num_batches] for k, v in gm_output.items()} uncond_mean = gm_to_mean(gm_uncond) gaussian_cond = gm_to_iso_gaussian(gm_cond)[0] if ndim == 5: gaussian_cond['var'] = gaussian_cond['var'].mean(dim=(-4, -2, -1), keepdim=True) # exclude channel dim else: gaussian_cond['var'] = gaussian_cond['var'].mean(dim=(-2, -1), keepdim=True) gaussian_output = probabilistic_guidance_jit( gaussian_cond['mean'], gaussian_cond['var'], uncond_mean, guidance_scale, orthogonal=orthogonal_guidance, orthogonal_axis=self.u_to_x_0(gaussian_cond['mean'], x_t, t))[0] gm_output = gm_mul_iso_gaussian( gm_cond, iso_gaussian_mul_iso_gaussian(gaussian_output, gaussian_cond, 1, -1), 1, 1)[0] if guidance_interval[0] > 0 or guidance_interval[1] < self.num_timesteps: guidance_active = ((t >= guidance_interval[0]) & (t <= guidance_interval[1])).reshape( [num_batches] + [1] * ndim) gm_output = {k: torch.where(guidance_active, v, gm_cond[k]) for k, v in gm_output.items()} u = gm_to_mean(gm_output) if ndim == 5: # (bs, t, c, h, w) u = u.permute(0, 2, 1, 3, 4) # (bs, c, t, h, w) return u.to(ori_dtype)