Update app.py

app.py

@@ -1,10 +1,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
 
-model = GeneratorUNet.from_pretrained('huggan/pix2pix-edge2shoes')
+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 dummy_fn(text):
-    return "Looks like the model is loaded!"
+def predict_fn(img):
+    inp = transform(img).unsqueeze(0)
+    out = model(inp)
+    save_image(out, 'out.png', normalize=True)
+    return 'out.png'
 
-gr.Interface(dummy_fn, inputs='text', outputs='text').launch()
+gr.Interface(predict_fn, inputs=gr.inputs.Image(type='pil'), outputs='image').launch()