| | import torch |
| | import nltk |
| | import io |
| | import base64 |
| | from torchvision import transforms |
| | from pytorch_pretrained_biggan import BigGAN, one_hot_from_names, truncated_noise_sample |
| | class PreTrainedPipeline(): |
| | def __init__(self, path=""): |
| | """ |
| | Initialize model |
| | """ |
| | nltk.download('wordnet') |
| | self.model = BigGAN.from_pretrained(path) |
| | self.truncation = 0.1 |
| | def __call__(self, inputs: str): |
| | """ |
| | Args: |
| | inputs (:obj:`str`): |
| | a string containing some text |
| | Return: |
| | A :obj:`PIL.Image` with the raw image representation as PIL. |
| | """ |
| | class_vector = one_hot_from_names([inputs], batch_size=1) |
| | if type(class_vector) == type(None): |
| | raise ValueError("Input is not in ImageNet") |
| | noise_vector = truncated_noise_sample(truncation=self.truncation, batch_size=1) |
| | noise_vector = torch.from_numpy(noise_vector) |
| | class_vector = torch.from_numpy(class_vector) |
| | with torch.no_grad(): |
| | output = self.model(noise_vector, class_vector, self.truncation) |
| | |
| | img = output[0] |
| | img = (img + 1) / 2.0 |
| | img = transforms.ToPILImage()(img) |
| | return img |
