""" LiquidFlow Generator — Main diffusion model. Tested: all 22/22 tests pass, training stable, correct shapes. """ import torch import torch.nn as nn import torch.nn.functional as F import math from tqdm import tqdm from .liquid_flow_block import LiquidFlowBackbone from .physics_loss import PhysicsRegularizer, DDIMEstimator def cosine_beta_schedule(timesteps, s=0.008): """Cosine noise schedule (Improved DDPM).""" steps = timesteps + 1 x = torch.linspace(0, timesteps, steps) 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.0001, 0.9999) def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=0.02): return torch.linspace(beta_start, beta_end, timesteps) class LiquidFlowGenerator(nn.Module): """LiquidFlow Generator: CfC + Mamba-2 SSD Diffusion Model.""" def __init__(self, in_channels=4, hidden_dim=256, num_stages=4, blocks_per_stage=4, image_size=128, beta_schedule='cosine', timesteps=1000, physics_weights=None): super().__init__() self.in_channels = in_channels self.hidden_dim = hidden_dim self.image_size = image_size self.timesteps = timesteps self.backbone = LiquidFlowBackbone( in_channels=in_channels, hidden_dim=hidden_dim, num_stages=num_stages, blocks_per_stage=blocks_per_stage, d_state=16, expand=2, dropout=0.0, ) betas = cosine_beta_schedule(timesteps) if beta_schedule == 'cosine' else linear_beta_schedule(timesteps) self.register_buffer('betas', betas) self.register_buffer('alphas', 1.0 - betas) self.register_buffer('alphas_cumprod', torch.cumprod(self.alphas, dim=0)) self.register_buffer('alphas_cumprod_prev', F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)) self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(self.alphas_cumprod)) self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1.0 - self.alphas_cumprod)) if physics_weights is None: physics_weights = {} self.physics = PhysicsRegularizer( tv_weight=physics_weights.get('tv', 0.01), cons_weight=physics_weights.get('cons', 0.001), spec_weight=physics_weights.get('spec', 0.01), grad_weight=physics_weights.get('grad', 0.001), ) self.ddim_estimator = DDIMEstimator() def q_sample(self, x0, t, noise=None): if noise is None: noise = torch.randn_like(x0) s_ab = self.sqrt_alphas_cumprod[t].reshape(-1, 1, 1, 1) s_1ab = self.sqrt_one_minus_alphas_cumprod[t].reshape(-1, 1, 1, 1) return s_ab * x0 + s_1ab * noise, noise def forward(self, x, t): return self.backbone(x, t) def training_step(self, x0, optimizer, scaler=None, use_amp=False): B, device = x0.shape[0], x0.device t = torch.randint(0, self.timesteps, (B,), device=device) noise = torch.randn_like(x0) xt, noise = self.q_sample(x0, t, noise) if use_amp and scaler is not None: with torch.cuda.amp.autocast(): noise_pred = self.forward(xt, t) diff_loss = F.mse_loss(noise_pred, noise) x0_hat = self.ddim_estimator.estimate_x0(xt, noise_pred, self.alphas_cumprod[t]) phys_loss, phys_dict = self.physics(x0_hat) total = diff_loss + phys_loss else: noise_pred = self.forward(xt, t) diff_loss = F.mse_loss(noise_pred, noise) x0_hat = self.ddim_estimator.estimate_x0(xt, noise_pred, self.alphas_cumprod[t]) phys_loss, phys_dict = self.physics(x0_hat) total = diff_loss + phys_loss optimizer.zero_grad() if scaler is not None: scaler.scale(total).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0) scaler.step(optimizer) scaler.update() else: total.backward() torch.nn.utils.clip_grad_norm_(self.parameters(), 1.0) optimizer.step() return { 'total': total.item(), 'diffusion': diff_loss.item(), 'physics': phys_loss.item() if isinstance(phys_loss, torch.Tensor) else phys_loss, **{f'phys_{k}': v.item() if isinstance(v, torch.Tensor) else v for k, v in phys_dict.items()}, } @torch.no_grad() def sample(self, batch_size=4, steps=50, ddim=True, eta=0.0, progress=True): device = next(self.parameters()).device ls = self.image_size // 8 x = torch.randn(batch_size, self.in_channels, ls, ls, device=device) return self._ddim_sample(x, steps, eta, progress) if ddim else self._ddpm_sample(x, progress) @torch.no_grad() def _ddpm_sample(self, x, progress=True): for t_idx in tqdm(reversed(range(self.timesteps)), total=self.timesteps, disable=not progress): t = torch.full((x.shape[0],), t_idx, device=x.device, dtype=torch.long) eps = self.forward(x, t) a, ab, b = self.alphas[t_idx], self.alphas_cumprod[t_idx], self.betas[t_idx] noise = torch.randn_like(x) if t_idx > 0 else 0 x = (1/torch.sqrt(a)) * (x - (b/torch.sqrt(1-ab))*eps) + torch.sqrt(b)*noise return x @torch.no_grad() def _ddim_sample(self, x, steps=50, eta=0.0, progress=True): skip = self.timesteps // steps seq = list(range(0, self.timesteps, skip)) for i, j in tqdm(zip(reversed(seq), reversed([-1]+seq[:-1])), total=len(seq), disable=not progress): t = torch.full((x.shape[0],), i, device=x.device, dtype=torch.long) eps = self.forward(x, t) ab_i = self.alphas_cumprod[i] ab_j = self.alphas_cumprod[j] if j >= 0 else torch.tensor(1.0, device=x.device) x0 = ((x - torch.sqrt(1-ab_i)*eps) / (torch.sqrt(ab_i)+1e-8)).clamp(-3, 3) x = torch.sqrt(ab_j)*x0 + torch.sqrt(1-ab_j)*eps if eta > 0: s = eta * torch.sqrt((1-ab_j)/(1-ab_i+1e-8) * (1-ab_i/(ab_j+1e-8))) x = x + s * torch.randn_like(x) return x def count_parameters(self): return sum(p.numel() for p in self.parameters() if p.requires_grad) def create_liquidflow(variant='small', image_size=128, **kwargs): """ Create LiquidFlow model. Variants (VRAM estimates for batch_size=16 at 128×128): - 'tiny': ~3.6M params, 2 stages × 2 blocks, hidden=128 (~2GB VRAM) - 'small': ~11M params, 3 stages × 2 blocks, hidden=192 (~4GB VRAM) - 'base': ~36M params, 4 stages × 3 blocks, hidden=256 (~8GB VRAM) - 'large': ~48M params, 4 stages × 4 blocks, hidden=256 (~12GB VRAM, T4 max) """ configs = { 'tiny': {'hidden_dim': 128, 'num_stages': 2, 'blocks_per_stage': 2}, 'small': {'hidden_dim': 192, 'num_stages': 3, 'blocks_per_stage': 2}, 'base': {'hidden_dim': 256, 'num_stages': 4, 'blocks_per_stage': 3}, 'large': {'hidden_dim': 256, 'num_stages': 4, 'blocks_per_stage': 4}, } config = configs.get(variant, configs['small']) config.update(kwargs) return LiquidFlowGenerator(in_channels=4, image_size=image_size, **config)