Spaces:
Running
Running
File size: 1,340 Bytes
5ab5cab |
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 torch
import torch.nn as nn
import yaml
import transformers
class UnetWrapper(nn.Module):
def __init__(self, Unet: nn.Module, config_path: str,
cond_encoder = None):
super().__init__()
with open(config_path, "r") as file:
config = yaml.safe_load(file)['unet']
self.add_module('network', Unet(**config))
# ConditionalEncoder
self.add_module('cond_encoder', cond_encoder)
def forward(self, x, t, y=None, cond_drop_all:bool = False):
if t.dim() == 0:
t = x.new_full((x.size(0), ), t, dtype = torch.int, device = x.device)
if y is not None:
assert self.cond_encoder is not None, 'You need to set ConditionalEncoder for conditional sampling.'
if isinstance(y, str) or isinstance(y, transformers.tokenization_utils_base.BatchEncoding):
y = self.cond_encoder(y, cond_drop_all=cond_drop_all).to(x.device)
else:
if torch.is_tensor(y) == False:
y = torch.tensor([y], device=x.device)
y = self.cond_encoder(y, cond_drop_all=cond_drop_all).squeeze()
if y.size(0) != x.size(0):
y = y.repeat(x.size(0), 1)
return self.network(x, t, y)
else:
return self.network(x, t) |