mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-05-15 05:11:10 +08:00
50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
import math
|
||
import torch
|
||
|
||
class Diffusion:
|
||
def __init__(self, device, T=100):
|
||
self.T = T
|
||
self.device = device
|
||
|
||
# cosine schedule（推荐）
|
||
steps = torch.arange(T + 1, dtype=torch.float32)
|
||
s = 0.008
|
||
f = torch.cos(((steps / T) + s) / (1 + s) * math.pi * 0.5) ** 2
|
||
alpha_bar = f / f[0]
|
||
|
||
self.alpha_bar = alpha_bar.to(device)
|
||
self.sqrt_ab = torch.sqrt(self.alpha_bar)
|
||
self.sqrt_one_minus_ab = torch.sqrt(1 - self.alpha_bar)
|
||
|
||
def q_sample(self, x0, t, noise):
|
||
"""
|
||
前向加噪
|
||
"""
|
||
return (
|
||
self.sqrt_ab[t][:, None, None, None] * x0
|
||
+ self.sqrt_one_minus_ab[t][:, None, None, None] * noise
|
||
)
|
||
|
||
def sample(self, model, cond: torch.Tensor, steps=20):
|
||
B = cond.shape[0]
|
||
x = torch.randn_like(cond).to(cond.device)
|
||
|
||
step_size = self.T // steps
|
||
|
||
for i in reversed(range(0, self.T, step_size)):
|
||
t = torch.full((B,), i, device=cond.device)
|
||
|
||
pred_noise = model(x, cond, t)
|
||
|
||
alpha = self.alpha_bar[i]
|
||
alpha_prev = self.alpha_bar[max(i - step_size, 0)]
|
||
|
||
x0_pred = (x - torch.sqrt(1 - alpha) * pred_noise) / torch.sqrt(alpha)
|
||
|
||
x = (
|
||
torch.sqrt(alpha_prev) * x0_pred
|
||
+ torch.sqrt(1 - alpha_prev) * pred_noise
|
||
)
|
||
|
||
return x
|
||
|