File size: 6,538 Bytes
1156de8 b9255c1 1156de8 4a95232 1156de8 b9255c1 1156de8 b9255c1 1156de8 4a95232 b9255c1 1156de8 b9255c1 1156de8 4a95232 b9255c1 1156de8 b9255c1 1156de8 b9255c1 1156de8 4a95232 b9255c1 1156de8 4a95232 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """
All-in-one utilities for Ambient Diffusion (random-inpainting variant).
Implements
----------
sample_inpainting_mask - A (Eq. 3.1)
further_corrupt - ~A (Eq. 3.2)
make_ambient_batch - build (~A x_t , A) for training
ambient_loss - J_corr (Eq. 3.2)
AmbientDDPMPipeline - fixed-mask sampler (Eq. 3.3)
"""
from __future__ import annotations
import torch
from torch import Tensor
from typing import Tuple, Optional
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
from diffusers import DDPMPipeline
# ------------------------------------------------------------------
# Corruptions
# ------------------------------------------------------------------
def sample_inpainting_mask(
shape: Tuple[int, int, int, int],
p: float = 0.9,
device: torch.device | str | None = None,
) -> Tensor:
"""Diagonal Bernoulli mask A: 1 = keep, 0 = erase."""
B, C, H, W = shape
return torch.bernoulli(torch.full((B, 1, H, W), p, device=device)).expand(-1, C, -1, -1)
def further_corrupt(A: Tensor, delta: float = 0.05) -> Tensor:
"""Sample ~A = B A by turning surviving pixels off with prob delta."""
B_mat = torch.bernoulli(torch.full_like(A, 1.0 - delta))
return B_mat * A
# ------------------------------------------------------------------
# Training helpers
# ------------------------------------------------------------------
def make_ambient_batch(
clean: Tensor,
noise_scheduler: DDPMScheduler,
timesteps: Tensor,
p: float = 0.9,
delta: float = 0.05,
) -> tuple[Tensor, Tensor]:
"""
Returns:
y_t_tilde : ~A x_t (to feed the network)
A_mask : A (needed only for the loss)
"""
B, C, H, W = clean.shape
device = clean.device
# 1. A x₀
A = sample_inpainting_mask((B, C, H, W), p=p, device=device)
# 2. add diffusion noise → x_t , then A x_t
noise = torch.randn_like(clean)
if timesteps.ndim == 0:
timesteps = torch.full((B,), timesteps, device=device, dtype=torch.long)
elif timesteps.shape[0] != B:
raise ValueError(f"Timesteps batch size {timesteps.shape[0]} does not match input batch size {B}")
x_t = noise_scheduler.add_noise(clean, noise, timesteps)
# 3. ~A x_t
A_tilde = further_corrupt(A, delta=delta)
mask_ch = A_tilde[:, :1] # (B,1,H,W)
net_input = torch.cat([A_tilde * x_t, mask_ch], dim=1)
return net_input, A
def ambient_loss(pred, clean, A_mask, snr_weights=None):
"""
Masked L2 loss with optional SNR weights (Eq. 2 in the paper).
"""
diff = A_mask * (pred - clean)
loss = 0.5 * diff.pow(2)
if snr_weights is not None:
loss = snr_weights * loss
return loss.mean()
# ------------------------------------------------------------------
# Sampler
# ------------------------------------------------------------------
class AmbientDDPMPipeline(DDPMPipeline):
"""
Fixed-mask sampler (Eq. 3.3). Drop-in replacement for DDPMPipeline.
"""
def __init__(self, *, unet, scheduler, p_mask: float = 0.9):
# 1) Let base class register the trainable modules
super().__init__(unet=unet, scheduler=scheduler)
# 2) Store hyper‑parameter in the *config* **and** as an attribute
self.register_to_config(p_mask=p_mask) # guarantees serialisation
self.p_mask = p_mask # convenient runtime access
@torch.no_grad()
def __call__(
self,
batch_size: int = 1,
generator: Optional[torch.Generator] = None,
num_inference_steps: int = 250,
class_labels: Optional[torch.Tensor] = None,
guidance_scale: float = 1.5,
output_type: str = "pt",
return_dict: bool = True,
mem = False
):
device = self.device
h = w = self.unet.config.sample_size
# Create a dummy mask that's all ones (no masking)
dummy_mask = torch.ones((batch_size, 1, h, w), device=device)
# Start with random noise
x = torch.randn(
(batch_size, 1, h, w),
generator=generator,
device=device,
)
self.scheduler.set_timesteps(num_inference_steps, device=device)
if mem:
TCNP = torch.empty((batch_size, num_inference_steps), device=self.device)
# Set up conditioning labels
if class_labels is None:
cond_lbls = torch.zeros(
(batch_size, getattr(self.unet.config, "multihot_dim", 1)),
dtype=torch.long if not hasattr(self.unet.config, "multihot_dim") else torch.float,
device=device,
)
else:
cond_lbls = class_labels.to(device)
uncond_lbls = torch.zeros_like(cond_lbls)
for i, t in enumerate(self.scheduler.timesteps):
# Always include mask channel to match training, but mask is all ones
x_with_mask = torch.cat([x, dummy_mask], dim=1)
if guidance_scale == 1.0 or class_labels is None:
# Single conditional/unconditional pass
eps = self.unet(x_with_mask, t, class_labels=cond_lbls).sample
else:
# Classifier-free guidance with duplicate batch
inp = torch.cat([x_with_mask, x_with_mask], dim=0)
lbls = torch.cat([cond_lbls, uncond_lbls], dim=0)
tids = t.expand(2 * batch_size)
eps_cond, eps_uncond = (
self.unet(inp, tids, class_labels=lbls).sample.chunk(2)
)
eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
if mem:
# print(cond_output.squeeze().shape, uncond_output.shape)
TCNP[:, i] = torch.linalg.norm(eps_cond.squeeze() - eps_uncond.squeeze(), dim=[0,1])
# Standard update (same as before)
x0_hat = eps
sigma = self.scheduler._get_variance(t).sqrt()
gamma = sigma ** 2 / (sigma ** 2 + 1)
x = gamma * x + (1 - gamma) * x0_hat
x = self.scheduler.step(eps, t, x).prev_sample
x = (x / 2 + 0.5).clamp(0, 1)
img = x.cpu().permute(0, 2, 3, 1).numpy()
if output_type == "pil":
img = self.numpy_to_pil(img)
if not return_dict:
return (img,)
if mem:
return dict(images=img, TCNP=TCNP)
return dict(images=img) |