File size: 677 Bytes
a3e0416
7d0d1dc
 
 
a3e0416
a820a56
 
7d0d1dc
 
 
 
 
 
 
 
a3e0416
7d0d1dc
 
 
 
 
a3e0416
7d0d1dc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import gradio as gr
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from PIL import Image
from torchvision.utils import save_image

from huggan.pytorch.pix2pix.modeling_pix2pix import GeneratorUNet

transform = Compose(
    [
        Resize((256, 256), Image.BICUBIC),
        ToTensor(),
        Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
    ]
)
model = GeneratorUNet.from_pretrained('huggan/pix2pix-cityscapes')

def predict_fn(img):
    inp = transform(img).unsqueeze(0)
    out = model(inp)
    save_image(out, 'out.png', normalize=True)
    return 'out.png'

gr.Interface(predict_fn, inputs=gr.inputs.Image(type='pil'), outputs='image').launch()