/
cfg_sampler.py
33 lines (27 loc) · 1.18 KB
/
cfg_sampler.py
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
import numpy as np
import torch
import torch.nn as nn
from copy import deepcopy
# A wrapper model for Classifier-free guidance **SAMPLING** only
# https://arxiv.org/abs/2207.12598
class ClassifierFreeSampleModel(nn.Module):
def __init__(self, model):
super().__init__()
self.model = model # model is the actual model to run
assert self.model.cond_mask_prob > 0, 'Cannot run a guided diffusion on a model that has not been trained with no conditions'
# pointers to inner model
self.rot2xyz = self.model.rot2xyz
self.translation = self.model.translation
self.njoints = self.model.njoints
self.nfeats = self.model.nfeats
self.data_rep = self.model.data_rep
self.cond_mode = self.model.cond_mode
self.encode_text = self.model.encode_text
def forward(self, x, timesteps, y=None):
cond_mode = self.model.cond_mode
assert cond_mode in ['text', 'action']
y_uncond = deepcopy(y)
y_uncond['uncond'] = True
out = self.model(x, timesteps, y)
out_uncond = self.model(x, timesteps, y_uncond)
return out_uncond + (y['scale'].view(-1, 1, 1, 1) * (out - out_uncond))