Update app.py

app.py

@@ -63,8 +63,7 @@ def image_create(input_img):
             #gaussian=gaussian_radius,
             log_visual_freq=1000,
             input='text',
-        ))
-    device = th.device()
+        )
     generator = create_generator("stylegan2-ffhq-config-f.pt","feng2022/Time-TravelRephotography_stylegan2-ffhq-config-f",args, device)
     latent = torch.randn((1, 512), device=device) 
     img_out, _, _ = generator([latent])
@@ -79,8 +78,11 @@ def main():
     #    ini = "False1"
     #result = subprocess.check_output(['nvidia-smi'])
     args = parse_args()
+    device = th.device()
+    func = functools.partial(image_create, device=device)
+    func = functools.update_wrapper(func, generate_image)
     iface = gr.Interface(
-          image_create, 
+          func, 
           [
                 gr.inputs.Number(default=0, label='Seed'),
                 gr.inputs.Slider(