Update app.py

app.py

@@ -3,8 +3,14 @@ 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
-import numpy as np
 
+
+def predict_fn(img):
+    inp = transform(img).unsqueeze(0)
+    out = model(inp)
+    save_image(out, 'out.png', normalize=True)
+    return 'out.png'
+    
 transform = Compose(
     [
         Resize((1024, 1024), Image.BICUBIC),
@@ -15,10 +21,4 @@ transform = Compose(
 
 model = GeneratorUNet.from_pretrained('huggan/pix2pix-uavid-15')
 
-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', examples=[['sample.jpg'], ['sample2.jpg'], ['sample3.jpg']]).launch()