| 8 | # A wrapper model for Classifier-free guidance **SAMPLING** only |
| 9 | # https://arxiv.org/abs/2207.12598 |
| 10 | class ClassifierFreeSampleModel(nn.Module): |
| 11 | |
| 12 | def __init__(self, model): |
| 13 | super().__init__() |
| 14 | self.model = model # model is the actual model to run |
| 15 | |
| 16 | assert self.model.cond_mask_prob > 0, 'Cannot run a guided diffusion on a model that has not been trained with no conditions' |
| 17 | |
| 18 | # pointers to inner model |
| 19 | self.rot2xyz = self.model.rot2xyz |
| 20 | self.translation = self.model.translation |
| 21 | self.njoints = self.model.njoints |
| 22 | self.nfeats = self.model.nfeats |
| 23 | self.data_rep = self.model.data_rep |
| 24 | self.cond_mode = self.model.cond_mode |
| 25 | self.encode_text = self.model.encode_text |
| 26 | |
| 27 | def forward(self, x, timesteps, y=None): |
| 28 | cond_mode = self.model.cond_mode |
| 29 | assert cond_mode in ['text', 'action'] |
| 30 | y_uncond = deepcopy(y) |
| 31 | y_uncond['uncond'] = True |
| 32 | out = self.model(x, timesteps, y) |
| 33 | out_uncond = self.model(x, timesteps, y_uncond) |
| 34 | return out_uncond + (y['scale'].view(-1, 1, 1, 1) * (out - out_uncond)) |
| 35 | |
| 36 | def __getattr__(self, name, default=None): |
| 37 | # this method is reached only if name is not in self.__dict__. |
| 38 | return wrapped_getattr(self, name, default=None) |
| 39 | |
| 40 | |
| 41 | class AutoRegressiveSampler(): |