File size: 13,356 Bytes
40440d6 | 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 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 | """
State-of-the-art sampling techniques (2024-2025).
Includes Flow Matching with Optimal Transport and Consistency Models.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Callable
import math
class FlowMatching:
"""
Flow Matching with Optimal Transport for improved trajectory learning.
Based on recent advances in continuous normalizing flows.
"""
def __init__(self, sigma_min: float = 0.001):
"""
Args:
sigma_min: Minimum noise level for numerical stability
"""
self.sigma_min = sigma_min
def compute_optimal_transport_path(self, x0: torch.Tensor, x1: torch.Tensor,
t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""
Compute optimal transport path from noise (x0) to data (x1).
Args:
x0: Source distribution (noise) [B, ...]
x1: Target distribution (data) [B, ...]
t: Time steps [B] in range [0, 1]
Returns:
x_t: Interpolated samples
v_t: Vector field (velocity)
"""
t = t.view(-1, *([1] * (x0.dim() - 1)))
# Optimal transport interpolation (linear for Gaussian)
x_t = (1 - t) * x0 + t * x1
# Velocity is the derivative: v = dx/dt = x1 - x0
v_t = x1 - x0
return x_t, v_t
def compute_conditional_flow(self, x0: torch.Tensor, x1: torch.Tensor,
t: torch.Tensor, sigma: Optional[float] = None) -> tuple[torch.Tensor, torch.Tensor]:
"""
Compute conditional flow with added noise for robustness.
Args:
x0: Noise [B, ...]
x1: Data [B, ...]
t: Time [B]
sigma: Optional noise level (defaults to sigma_min * (1-t))
Returns:
x_t: Noisy interpolation
v_t: Target velocity
"""
if sigma is None:
sigma = self.sigma_min * (1 - t.view(-1, *([1] * (x0.dim() - 1))))
else:
sigma = sigma * torch.ones_like(t).view(-1, *([1] * (x0.dim() - 1)))
# Conditional path
x_t, v_t = self.compute_optimal_transport_path(x0, x1, t)
# Add Gaussian noise for conditional flow matching
eps = torch.randn_like(x_t)
x_t = x_t + sigma * eps
return x_t, v_t
def sample(self, model: Callable, x_init: torch.Tensor,
num_steps: int = 50, **model_kwargs) -> torch.Tensor:
"""
Sample using Euler method for ODE solving.
Args:
model: Velocity prediction model
x_init: Initial noise [B, ...]
num_steps: Number of integration steps
**model_kwargs: Additional arguments for model
Returns:
x_final: Generated samples
"""
x = x_init
dt = 1.0 / num_steps
for step in range(num_steps):
t_current = torch.full((x.shape[0],), step / num_steps,
device=x.device, dtype=x.dtype)
# Predict velocity
with torch.no_grad():
v = model(x, timesteps=t_current, **model_kwargs)
# Euler step
x = x + v * dt
return x
def sample_heun(self, model: Callable, x_init: torch.Tensor,
num_steps: int = 50, **model_kwargs) -> torch.Tensor:
"""
Higher-order sampling using Heun's method (2nd order).
More accurate but requires 2x model evaluations.
Args:
model: Velocity prediction model
x_init: Initial noise
num_steps: Number of steps
**model_kwargs: Model arguments
Returns:
x_final: Generated samples
"""
x = x_init
dt = 1.0 / num_steps
for step in range(num_steps):
t_current = torch.full((x.shape[0],), step / num_steps,
device=x.device, dtype=x.dtype)
# First stage
with torch.no_grad():
v1 = model(x, timesteps=t_current, **model_kwargs)
x_temp = x + v1 * dt
# Second stage
t_next = torch.full((x.shape[0],), (step + 1) / num_steps,
device=x.device, dtype=x.dtype)
with torch.no_grad():
v2 = model(x_temp, timesteps=t_next, **model_kwargs)
# Heun's method: average of two velocities
x = x + (v1 + v2) * 0.5 * dt
return x
class ConsistencyModel:
"""
Consistency Models for fast, single-step generation.
Based on "Consistency Models" (Song et al., 2023).
"""
def __init__(self, sigma_min: float = 0.002, sigma_max: float = 80.0,
rho: float = 7.0):
"""
Args:
sigma_min: Minimum noise level
sigma_max: Maximum noise level
rho: Schedule parameter
"""
self.sigma_min = sigma_min
self.sigma_max = sigma_max
self.rho = rho
def get_schedule(self, num_steps: int, device: torch.device) -> torch.Tensor:
"""
Get noise schedule for consistency training.
Args:
num_steps: Number of discretization steps
device: Device for tensors
Returns:
sigmas: Noise levels [num_steps + 1]
"""
rho_inv = 1.0 / self.rho
steps = torch.arange(num_steps + 1, device=device)
sigmas = (self.sigma_max ** rho_inv +
steps / num_steps * (self.sigma_min ** rho_inv - self.sigma_max ** rho_inv)) ** self.rho
return sigmas
def consistency_function(self, model: Callable, x: torch.Tensor,
sigma: torch.Tensor, **model_kwargs) -> torch.Tensor:
"""
Apply consistency function: maps (x, sigma) to denoised output.
Args:
model: Denoising model
x: Noisy input
sigma: Noise level
**model_kwargs: Model arguments
Returns:
denoised: Consistency function output
"""
# Skip connection for sigma_min
c_skip = self.sigma_min ** 2 / (sigma ** 2 + self.sigma_min ** 2)
c_out = sigma * self.sigma_min / (sigma ** 2 + self.sigma_min ** 2).sqrt()
# Predict denoised output
with torch.no_grad():
# Convert sigma to timestep for model
t = sigma / self.sigma_max
model_out = model(x, timesteps=t, **model_kwargs)
# Apply consistency function
denoised = c_skip * x + c_out * model_out
return denoised
def sample_onestep(self, model: Callable, shape: tuple,
device: torch.device, dtype: torch.dtype,
**model_kwargs) -> torch.Tensor:
"""
Single-step generation from noise to data.
Args:
model: Consistency model
shape: Output shape
device: Device
dtype: Data type
**model_kwargs: Model arguments
Returns:
samples: Generated samples
"""
# Sample from prior (high noise)
x = torch.randn(*shape, device=device, dtype=dtype) * self.sigma_max
sigma = torch.full((shape[0],), self.sigma_max, device=device, dtype=dtype)
# One-step denoising
samples = self.consistency_function(model, x, sigma, **model_kwargs)
return samples
def sample_multistep(self, model: Callable, shape: tuple,
num_steps: int, device: torch.device,
dtype: torch.dtype, **model_kwargs) -> torch.Tensor:
"""
Multi-step generation for better quality.
Args:
model: Consistency model
shape: Output shape
num_steps: Number of steps
device: Device
dtype: Data type
**model_kwargs: Model arguments
Returns:
samples: Generated samples
"""
# Get noise schedule
sigmas = self.get_schedule(num_steps, device)
# Start from maximum noise
x = torch.randn(*shape, device=device, dtype=dtype) * self.sigma_max
# Iterative refinement
for i in range(num_steps, 0, -1):
sigma_current = sigmas[i]
sigma_next = sigmas[i - 1]
# Denoise to current level
x = self.consistency_function(
model, x,
torch.full((shape[0],), sigma_current, device=device, dtype=dtype),
**model_kwargs
)
# Add noise for next iteration (except last)
if i > 1:
noise = torch.randn_like(x)
x = x + noise * (sigma_next ** 2 - self.sigma_min ** 2).sqrt()
return x
class ClassifierFreeGuidance:
"""
Enhanced Classifier-Free Guidance for better conditioning.
"""
def __init__(self, guidance_scale: float = 7.5, min_scale: float = 1.0):
"""
Args:
guidance_scale: Guidance strength
min_scale: Minimum guidance scale
"""
self.guidance_scale = guidance_scale
self.min_scale = min_scale
def apply_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
guidance_scale: Optional[float] = None) -> torch.Tensor:
"""
Apply classifier-free guidance to model outputs.
Args:
cond_out: Conditional model output
uncond_out: Unconditional model output
guidance_scale: Optional override for guidance scale
Returns:
guided_out: Guided output
"""
scale = guidance_scale if guidance_scale is not None else self.guidance_scale
scale = max(scale, self.min_scale)
# CFG formula: out = uncond + scale * (cond - uncond)
guided_out = uncond_out + scale * (cond_out - uncond_out)
return guided_out
def apply_dynamic_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
t: torch.Tensor, min_t: float = 0.0,
max_t: float = 1.0) -> torch.Tensor:
"""
Apply time-dependent guidance that decreases over time.
Args:
cond_out: Conditional output
uncond_out: Unconditional output
t: Time steps [B]
min_t: Time to start decreasing guidance
max_t: Time to reach minimum guidance
Returns:
guided_out: Dynamically guided output
"""
# Linearly decrease guidance scale from max to min over time
t_normalized = (t - min_t) / (max_t - min_t)
t_normalized = t_normalized.clamp(0, 1)
# Scale goes from guidance_scale to min_scale
scale = self.guidance_scale + (self.min_scale - self.guidance_scale) * t_normalized
scale = scale.view(-1, *([1] * (cond_out.dim() - 1)))
guided_out = uncond_out + scale * (cond_out - uncond_out)
return guided_out
class AdaptiveSampler:
"""
Adaptive sampling that combines multiple techniques for optimal quality/speed.
"""
def __init__(self, use_flow_matching: bool = True,
use_consistency: bool = False,
use_cfg: bool = True):
self.use_flow_matching = use_flow_matching
self.use_consistency = use_consistency
self.use_cfg = use_cfg
if use_flow_matching:
self.flow_matching = FlowMatching()
if use_consistency:
self.consistency_model = ConsistencyModel()
if use_cfg:
self.cfg = ClassifierFreeGuidance()
def sample(self, model: Callable, x_init: torch.Tensor,
num_steps: int = 50, method: str = "heun",
cfg_scale: Optional[float] = None,
**model_kwargs) -> torch.Tensor:
"""
Adaptive sampling with multiple techniques.
Args:
model: Model for sampling
x_init: Initial noise
num_steps: Number of steps
method: Sampling method ("euler", "heun", "consistency")
cfg_scale: Optional CFG scale
**model_kwargs: Model arguments
Returns:
samples: Generated samples
"""
if method == "consistency" and self.use_consistency:
return self.consistency_model.sample_multistep(
model, x_init.shape, num_steps,
x_init.device, x_init.dtype, **model_kwargs
)
elif method == "heun" and self.use_flow_matching:
return self.flow_matching.sample_heun(
model, x_init, num_steps, **model_kwargs
)
else: # euler
return self.flow_matching.sample(
model, x_init, num_steps, **model_kwargs
)
|