| import torch | |
| import torch.nn as nn | |
| import math | |
| import json | |
| from diffusers import UNet2DConditionModel | |
| import sys | |
| import time | |
| import numpy as np | |
| import os | |
| class PositionalEncoding(nn.Module): | |
| def __init__(self, d_model=384, max_len=5000): | |
| super(PositionalEncoding, self).__init__() | |
| pe = torch.zeros(max_len, d_model) | |
| position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) | |
| div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) | |
| pe[:, 0::2] = torch.sin(position * div_term) | |
| pe[:, 1::2] = torch.cos(position * div_term) | |
| pe = pe.unsqueeze(0) | |
| self.register_buffer('pe', pe) | |
| def forward(self, x): | |
| b, seq_len, d_model = x.size() | |
| pe = self.pe[:, :seq_len, :] | |
| x = x + pe.to(x.device) | |
| return x | |
| class UNet(): | |
| def __init__(self, | |
| unet_config, | |
| model_path, | |
| use_float16=False, | |
| ): | |
| with open(unet_config, 'r') as f: | |
| unet_config = json.load(f) | |
| self.model = UNet2DConditionModel(**unet_config) | |
| self.pe = PositionalEncoding(d_model=384) | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.weights = torch.load(model_path) if torch.cuda.is_available() else torch.load(model_path, map_location=self.device) | |
| self.model.load_state_dict(self.weights) | |
| if use_float16: | |
| self.model = self.model.half() | |
| self.model.to(self.device) | |
| if __name__ == "__main__": | |
| unet = UNet() |