44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
import torch
|
|
|
|
|
|
class RF:
|
|
def __init__(self, model, ln=True):
|
|
self.model = model
|
|
self.ln = ln
|
|
|
|
def forward(self, x0, x1, cond):
|
|
b = x0.size(0)
|
|
if self.ln:
|
|
nt = torch.randn((b,)).to(x0.device)
|
|
t = torch.sigmoid(nt)
|
|
else:
|
|
t = torch.rand((b,)).to(x0.device)
|
|
texp = t.view([b, *([1] * len(x0.shape[1:]))])
|
|
zt = (1 - texp) * x0 + texp * x1
|
|
vtheta = self.model(zt, t, cond)
|
|
batchwise_mse = ((x1 - x0 - vtheta) ** 2).mean(
|
|
dim=list(range(1, len(x0.shape)))
|
|
)
|
|
tlist = batchwise_mse.detach().cpu().reshape(-1).tolist()
|
|
ttloss = [(tv, tloss) for tv, tloss in zip(t, tlist)]
|
|
return batchwise_mse.mean(), ttloss
|
|
|
|
@torch.no_grad()
|
|
def sample(self, z, cond, null_cond=None, sample_steps=50, cfg=2.0):
|
|
b = z.size(0)
|
|
dt = 1.0 / sample_steps
|
|
dt = torch.tensor([dt] * b).to(z.device).view([b, *([1] * len(z.shape[1:]))])
|
|
images = [z]
|
|
for i in range(sample_steps, 0, -1):
|
|
t = i / sample_steps
|
|
t = torch.tensor([t] * b).to(z.device)
|
|
|
|
vc = self.model(z, t, cond)
|
|
if null_cond is not None:
|
|
vu = self.model(z, t, null_cond)
|
|
vc = vu + cfg * (vc - vu)
|
|
|
|
z = z - dt * vc
|
|
images.append(z)
|
|
return images
|