Anonymous941
Update app.py and pipeline.py
61f9504
import gradio as gr
from datasets import load_dataset, Image
import torch
import nltk
import io
import base64
import shutil
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)
# Scale image
img = output[0]
img = (img + 1) / 2.0
img = transforms.ToPILImage()(img)
dataset = load_dataset("botmaster/mother-2-battle-sprites", split="train")
gr.Interface.load("models/templates/text-to-image").launch()