"""Jenus conditional DDPM model with Hugging Face AutoModel support.""" import math import random as py_random import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from .configuration_jenus_ddpm import JenusDDPMConfig class ResBlock(nn.Module): def __init__(self, in_ch, out_ch, time_emb_dim): super().__init__() self.norm1 = nn.GroupNorm(min(8, in_ch), in_ch) self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1) self.time_mlp = nn.Linear(time_emb_dim, out_ch) self.norm2 = nn.GroupNorm(min(8, out_ch), out_ch) self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1) self.shortcut = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity() def forward(self, x, t_emb): h = self.conv1(F.silu(self.norm1(x))) h = h + self.time_mlp(F.silu(t_emb))[:, :, None, None] h = self.conv2(F.silu(self.norm2(h))) return h + self.shortcut(x) class AttentionBlock(nn.Module): def __init__(self, channels): super().__init__() self.norm = nn.GroupNorm(min(8, channels), channels) self.qkv = nn.Conv2d(channels, channels * 3, 1) self.proj = nn.Conv2d(channels, channels, 1) self.scale = channels ** -0.5 def forward(self, x): batch, channels, height, width = x.shape h = self.norm(x) qkv = self.qkv(h).reshape(batch, 3, channels, height * width) q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] attn = torch.softmax(torch.bmm(q.transpose(1, 2), k) * self.scale, dim=-1) h = torch.bmm(v, attn.transpose(1, 2)).reshape(batch, channels, height, width) return x + self.proj(h) class DownBlock(nn.Module): def __init__(self, in_ch, out_ch, time_emb_dim, has_attn=False): super().__init__() self.res = ResBlock(in_ch, out_ch, time_emb_dim) self.attn = AttentionBlock(out_ch) if has_attn else nn.Identity() self.down = nn.Conv2d(out_ch, out_ch, 3, stride=2, padding=1) def forward(self, x, t_emb): x = self.res(x, t_emb) skip = self.attn(x) return self.down(skip), skip class UpBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch, time_emb_dim, has_attn=False): super().__init__() self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2) self.res = ResBlock(in_ch + skip_ch, out_ch, time_emb_dim) self.attn = AttentionBlock(out_ch) if has_attn else nn.Identity() def forward(self, x, skip, t_emb): x = self.up(x) x = torch.cat([x, skip], dim=1) x = self.res(x, t_emb) return self.attn(x) def get_timestep_embedding(timesteps, embedding_dim): half_dim = embedding_dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, device=timesteps.device) * -emb) emb = timesteps[:, None].float() * emb[None, :] return torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) class ConditionalUNet(nn.Module): def __init__(self, in_channels=2, out_channels=1, base_channels=64, time_emb_dim=256): super().__init__() self.time_emb_dim = time_emb_dim self.time_mlp = nn.Sequential( nn.Linear(time_emb_dim, time_emb_dim * 4), nn.SiLU(), nn.Linear(time_emb_dim * 4, time_emb_dim), ) ch = base_channels self.conv_in = nn.Conv2d(in_channels, ch, 3, padding=1) self.down1 = DownBlock(ch, ch * 2, time_emb_dim, has_attn=False) self.down2 = DownBlock(ch * 2, ch * 4, time_emb_dim, has_attn=False) self.down3 = DownBlock(ch * 4, ch * 4, time_emb_dim, has_attn=False) self.mid_res1 = ResBlock(ch * 4, ch * 4, time_emb_dim) self.mid_attn = AttentionBlock(ch * 4) self.mid_res2 = ResBlock(ch * 4, ch * 4, time_emb_dim) self.up1 = UpBlock(ch * 4, ch * 4, ch * 4, time_emb_dim, has_attn=False) self.up2 = UpBlock(ch * 4, ch * 4, ch * 2, time_emb_dim, has_attn=False) self.up3 = UpBlock(ch * 2, ch * 2, ch, time_emb_dim, has_attn=False) self.conv_out = nn.Sequential( nn.GroupNorm(8, ch), nn.SiLU(), nn.Conv2d(ch, out_channels, 3, padding=1), ) def forward(self, x, t, cond): x = torch.cat([x, cond], dim=1) t_emb = get_timestep_embedding(t, self.time_emb_dim) t_emb = self.time_mlp(t_emb) h = self.conv_in(x) h, skip1 = self.down1(h, t_emb) h, skip2 = self.down2(h, t_emb) h, skip3 = self.down3(h, t_emb) h = self.mid_res1(h, t_emb) h = self.mid_attn(h) h = self.mid_res2(h, t_emb) h = self.up1(h, skip3, t_emb) h = self.up2(h, skip2, t_emb) h = self.up3(h, skip1, t_emb) return self.conv_out(h) class DDPMScheduler: def __init__(self, num_timesteps=1000, beta_start=1e-4, beta_end=0.02): self.num_timesteps = num_timesteps self.betas = torch.linspace(beta_start, beta_end, num_timesteps) self.alphas = 1 - self.betas self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) def to(self, device): self.betas = self.betas.to(device) self.alphas = self.alphas.to(device) self.alphas_cumprod = self.alphas_cumprod.to(device) return self class DDIMScheduler: def __init__(self, ddpm_scheduler, num_inference_steps=50): self.num_timesteps = ddpm_scheduler.num_timesteps self.alphas_cumprod = ddpm_scheduler.alphas_cumprod step_ratio = self.num_timesteps // num_inference_steps self.timesteps = torch.flip(torch.arange(0, num_inference_steps) * step_ratio, [0]) def step(self, model_output, t, t_prev, x_t, eta=0.0): alpha_t = self.alphas_cumprod[t].view(-1, 1, 1, 1) alpha_t_prev = ( self.alphas_cumprod[t_prev].view(-1, 1, 1, 1) if t_prev >= 0 else torch.ones_like(alpha_t) ) pred_x0 = (x_t - torch.sqrt(1 - alpha_t) * model_output) / torch.sqrt(alpha_t) pred_x0 = torch.clamp(pred_x0, -1, 1) sigma_t = eta * torch.sqrt((1 - alpha_t_prev) / (1 - alpha_t)) * torch.sqrt( 1 - alpha_t / alpha_t_prev ) pred_dir = torch.sqrt(1 - alpha_t_prev - sigma_t ** 2) * model_output x_prev = torch.sqrt(alpha_t_prev) * pred_x0 + pred_dir if eta > 0 and t_prev >= 0: x_prev = x_prev + sigma_t * torch.randn_like(x_t) return x_prev def to(self, device): self.alphas_cumprod = self.alphas_cumprod.to(device) self.timesteps = self.timesteps.to(device) return self class JenusDDPMModel(PreTrainedModel): """Hugging Face compatible Jenus conditional DDPM model.""" config_class = JenusDDPMConfig base_model_prefix = "jenus_ddpm" main_input_name = "x" _tied_weights_keys = [] all_tied_weights_keys = {} def __init__(self, config): super().__init__(config) self.unet = ConditionalUNet( in_channels=config.in_channels, out_channels=config.out_channels, base_channels=config.base_channels, time_emb_dim=config.time_emb_dim, ) self.ddpm_scheduler = None def forward(self, x, t, cond): return self.unet(x, t, cond) def to(self, *args, **kwargs): return super().to(*args, **kwargs) @property def device(self): return next(self.parameters()).device def _prepare_condition(self, raw_density): arr = raw_density.astype(np.float32) np.nan_to_num(arr, copy=False) raw_min, raw_max = float(arr.min()), float(arr.max()) img_size = tuple(self.config.img_size) t = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0) t = F.interpolate(t, size=img_size, mode="bilinear", align_corners=False) tmin, tmax = t.min(), t.max() if tmax - tmin > 1e-6: t = 2.0 * (t - tmin) / (tmax - tmin) - 1.0 return t.to(self.device), raw_min, raw_max @torch.no_grad() def _sample_once(self, condition, num_steps=50, eta=1.0, seed=None): if seed is None: seed = py_random.randint(0, 2**31 - 1) torch.manual_seed(seed) if self.device.type == "cuda": torch.cuda.manual_seed(seed) batch, _, height, width = condition.shape ddpm_scheduler = DDPMScheduler( num_timesteps=self.config.num_timesteps, beta_start=self.config.beta_start, beta_end=self.config.beta_end, ).to(self.device) ddim = DDIMScheduler(ddpm_scheduler, num_inference_steps=num_steps).to(self.device) x = torch.randn(batch, 1, height, width, device=self.device) timesteps = ddim.timesteps.long() self.eval() for i, t in enumerate(timesteps): t_prev = timesteps[i + 1] if i + 1 < len(timesteps) else torch.tensor(-1, device=self.device) t_batch = t.expand(batch).to(self.device) pred_noise = self.forward(x, t_batch, condition) x = ddim.step(pred_noise, t, t_prev, x, eta=eta) return x def predict(self, density, num_steps=50, eta=1.0, seed=None): orig_h, orig_w = density.shape condition, raw_min, raw_max = self._prepare_condition(density) x = self._sample_once(condition, num_steps=num_steps, eta=eta, seed=seed) pred_np = x[0, 0].cpu().float().numpy() pred_phys = (pred_np + 1.0) / 2.0 * (raw_max - raw_min) + raw_min pred_t = torch.from_numpy(pred_phys[None, None].astype(np.float32)) pred_t = F.interpolate(pred_t, size=(orig_h, orig_w), mode="bilinear", align_corners=False) return pred_t[0, 0].numpy() def predict_ensemble(self, density, n_samples=5, num_steps=50, eta=1.0): orig_h, orig_w = density.shape condition, raw_min, raw_max = self._prepare_condition(density) members = [] for _ in range(n_samples): x = self._sample_once(condition, num_steps=num_steps, eta=eta, seed=None) pred_np = x[0, 0].cpu().float().numpy() pred_phys = (pred_np + 1.0) / 2.0 * (raw_max - raw_min) + raw_min members.append(pred_phys) stack = np.stack(members) def resize(arr): t = torch.from_numpy(arr[None, None].astype(np.float32)) t = F.interpolate(t, size=(orig_h, orig_w), mode="bilinear", align_corners=False) return t[0, 0].numpy() return resize(stack.mean(axis=0)), resize(stack.std(axis=0))