| import torch
|
| import torch.nn as nn
|
| import numpy as np
|
| from typing import Optional, Tuple, Union
|
| import math
|
|
|
|
|
| class BetaScheduler:
|
| """Beta调度器"""
|
| @staticmethod
|
| def linear(num_timesteps: int, beta_start: float = 0.0001, beta_end: float = 0.02) -> np.ndarray:
|
| """线性调度"""
|
| return np.linspace(beta_start, beta_end, num_timesteps, dtype=np.float32)
|
|
|
| @staticmethod
|
| def cosine(num_timesteps: int, s: float = 0.008) -> np.ndarray:
|
| """余弦调度"""
|
| steps = num_timesteps + 1
|
| x = np.linspace(0, num_timesteps, steps)
|
| alphas_cumprod = np.cos(((x / num_timesteps) + s) / (1 + s) * np.pi * 0.5) ** 2
|
| alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
| betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
|
| return np.clip(betas, 0, 0.999)
|
|
|
| @staticmethod
|
| def scaled_linear(num_timesteps: int) -> np.ndarray:
|
| """缩放线性调度（Stable Diffusion默认）"""
|
| beta_start = 0.00085
|
| beta_end = 0.012
|
| return np.linspace(beta_start**0.5, beta_end**0.5, num_timesteps) ** 2
|
|
|
|
|
| class DiffusionProcess:
|
| """扩散过程管理"""
|
| def __init__(self, config: dict):
|
| self.config = config
|
| diff_config = config.get('diffusion', {})
|
|
|
| self.num_train_timesteps = diff_config.get('num_train_timesteps', 1000)
|
| self.num_inference_timesteps = diff_config.get('num_inference_timesteps', 50)
|
| self.beta_start = diff_config.get('beta_start', 0.00085)
|
| self.beta_end = diff_config.get('beta_end', 0.012)
|
| self.beta_schedule = diff_config.get('beta_schedule', 'scaled_linear')
|
| self.prediction_type = diff_config.get('prediction_type', 'epsilon')
|
|
|
|
|
| self._init_schedule()
|
|
|
| def _init_schedule(self):
|
| """初始化扩散调度参数"""
|
|
|
| if self.beta_schedule == "linear":
|
| betas = BetaScheduler.linear(
|
| self.num_train_timesteps,
|
| self.beta_start,
|
| self.beta_end
|
| )
|
| elif self.beta_schedule == "cosine":
|
| betas = BetaScheduler.cosine(self.num_train_timesteps)
|
| elif self.beta_schedule == "scaled_linear":
|
| betas = BetaScheduler.scaled_linear(self.num_train_timesteps)
|
| else:
|
| raise ValueError(f"Unknown beta schedule: {self.beta_schedule}")
|
|
|
| self.betas = torch.from_numpy(betas).float()
|
|
|
|
|
| self.alphas = 1.0 - self.betas
|
| self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
|
| self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0)
|
|
|
|
|
| self.variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
|
|
|
|
|
| self.register_buffer = lambda name, tensor: setattr(self, name, tensor)
|
| self.register_buffer('betas', self.betas)
|
| self.register_buffer('alphas', self.alphas)
|
| self.register_buffer('alphas_cumprod', self.alphas_cumprod)
|
| self.register_buffer('alphas_cumprod_prev', self.alphas_cumprod_prev)
|
| self.register_buffer('variance', self.variance)
|
|
|
|
|
| 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))
|
| self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1.0 - self.alphas_cumprod))
|
| self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1.0 / self.alphas_cumprod))
|
| self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1.0 / self.alphas_cumprod - 1))
|
|
|
| def q_sample(self, x_start: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
|
| """前向扩散过程：加噪"""
|
| if noise is None:
|
| noise = torch.randn_like(x_start)
|
|
|
| sqrt_alphas_cumprod_t = self.extract(self.sqrt_alphas_cumprod, t, x_start.shape)
|
| sqrt_one_minus_alphas_cumprod_t = self.extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
|
|
|
| return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
|
|
|
| def extract(self, a: torch.Tensor, t: torch.Tensor, x_shape: Tuple[int, ...]) -> torch.Tensor:
|
| """从张量a中提取索引t处的值"""
|
| batch_size = t.shape[0]
|
| out = a.gather(-1, t.cpu())
|
| return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
|
|
|
| def get_loss_weight(self, snr: torch.Tensor, gamma: float = 5.0) -> torch.Tensor:
|
| """根据SNR计算损失权重"""
|
| if gamma is None:
|
| return torch.ones_like(snr)
|
|
|
| snr = torch.clamp(snr, min=1e-8)
|
| min_snr = torch.tensor(gamma, device=snr.device)
|
| weight = torch.minimum(snr, min_snr) / snr
|
| return weight
|
|
|
| def compute_snr(self, timesteps: torch.Tensor) -> torch.Tensor:
|
| """计算信噪比(SNR)"""
|
| alphas_cumprod = self.extract(self.alphas_cumprod, timesteps, timesteps.shape)
|
| snr = alphas_cumprod / (1 - alphas_cumprod)
|
| return snr
|
|
|
|
|
| class DDIMScheduler:
|
| """DDIM采样器"""
|
| def __init__(self, diffusion: DiffusionProcess):
|
| self.diffusion = diffusion
|
| self.num_train_timesteps = diffusion.num_train_timesteps
|
| self.num_inference_timesteps = diffusion.num_inference_timesteps
|
|
|
|
|
| self.set_timesteps(self.num_inference_timesteps)
|
|
|
| def set_timesteps(self, num_inference_timesteps: int):
|
| """设置推理时间步"""
|
| self.num_inference_timesteps = num_inference_timesteps
|
|
|
|
|
| if self.num_train_timesteps == self.num_inference_timesteps:
|
| self.timesteps = torch.arange(0, self.num_train_timesteps).long()
|
| else:
|
| step_ratio = self.num_train_timesteps // self.num_inference_timesteps
|
| self.timesteps = torch.arange(0, self.num_train_timesteps, step_ratio).long()
|
|
|
| self.timesteps = self.timesteps.flip(0)
|
|
|
| @torch.no_grad()
|
| def step(self, model_output: torch.Tensor, timestep: int, sample: torch.Tensor, eta: float = 0.0) -> torch.Tensor:
|
| """DDIM单步采样"""
|
|
|
| prev_timestep = timestep - self.num_train_timesteps // self.num_inference_timesteps
|
|
|
|
|
| alpha_prod_t = self.diffusion.extract(self.diffusion.alphas_cumprod, timestep, sample.shape)
|
| alpha_prod_t_prev = self.diffusion.extract(
|
| self.diffusion.alphas_cumprod,
|
| prev_timestep,
|
| sample.shape
|
| ) if prev_timestep >= 0 else torch.ones_like(alpha_prod_t)
|
|
|
|
|
| if self.diffusion.prediction_type == "epsilon":
|
| pred_original_sample = (sample - (1 - alpha_prod_t) ** 0.5 * model_output) / alpha_prod_t ** 0.5
|
| pred_epsilon = model_output
|
| elif self.diffusion.prediction_type == "sample":
|
| pred_original_sample = model_output
|
| pred_epsilon = (sample - alpha_prod_t ** 0.5 * pred_original_sample) / (1 - alpha_prod_t) ** 0.5
|
| elif self.diffusion.prediction_type == "v_prediction":
|
| pred_original_sample = (alpha_prod_t ** 0.5) * sample - (1 - alpha_prod_t) ** 0.5 * model_output
|
| pred_epsilon = (alpha_prod_t ** 0.5) * model_output + (1 - alpha_prod_t) ** 0.5 * sample
|
| else:
|
| raise ValueError(f"Unsupported prediction type: {self.diffusion.prediction_type}")
|
|
|
|
|
| variance = (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
|
| std_dev_t = eta * variance ** 0.5
|
|
|
|
|
| pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon
|
| prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
|
|
|
|
|
| if eta > 0:
|
| noise = torch.randn_like(model_output)
|
| prev_sample = prev_sample + std_dev_t * noise
|
|
|
| return prev_sample
|
|
|
|
|
| class DiffusionModel(nn.Module):
|
| """扩散模型封装"""
|
| def __init__(self, unet: nn.Module, diffusion: DiffusionProcess):
|
| super().__init__()
|
| self.unet = unet
|
| self.diffusion = diffusion
|
| self.scheduler = DDIMScheduler(diffusion)
|
|
|
| def forward(self, x: torch.Tensor, timesteps: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
|
| """前向传播：预测噪声"""
|
| return self.unet(x, timesteps, context)
|
|
|
| def compute_loss(self, x_start: torch.Tensor, context: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
|
| """计算扩散损失"""
|
| if noise is None:
|
| noise = torch.randn_like(x_start)
|
|
|
|
|
| batch_size = x_start.shape[0]
|
| timesteps = torch.randint(
|
| 0, self.diffusion.num_train_timesteps,
|
| (batch_size,), device=x_start.device
|
| ).long()
|
|
|
|
|
| x_noisy = self.diffusion.q_sample(x_start, timesteps, noise)
|
|
|
|
|
| predicted_noise = self.unet(x_noisy, timesteps, context)
|
|
|
|
|
| loss = F.mse_loss(predicted_noise, noise)
|
|
|
| return loss
|
|
|
| @torch.no_grad()
|
| def generate(
|
| self,
|
| context: torch.Tensor,
|
| num_samples: int = 1,
|
| height: int = 512,
|
| width: int = 512,
|
| guidance_scale: float = 7.5
|
| ) -> torch.Tensor:
|
| """生成图像"""
|
|
|
| latents = torch.randn(
|
| (num_samples, self.unet.in_channels, height // 8, width // 8),
|
| device=next(self.unet.parameters()).device
|
| )
|
|
|
|
|
| self.scheduler.set_timesteps(self.diffusion.num_inference_timesteps)
|
|
|
| for t in self.scheduler.timesteps:
|
|
|
| latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents
|
| latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
|
|
|
|
|
| timesteps = torch.full((num_samples,), t, device=latents.device).long()
|
| if guidance_scale > 1.0:
|
| timesteps = torch.cat([timesteps] * 2)
|
|
|
| noise_pred = self.unet(latent_model_input, timesteps, context)
|
|
|
|
|
| if guidance_scale > 1.0:
|
| noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
| noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
|
|
|
|
| latents = self.scheduler.step(noise_pred, t, latents).prev_sample
|
|
|
| return latents |
