Update app.py

app.py

@@ -1,47 +1,14 @@
-import onnxruntime
-print(onnxruntime.get_device())
-import os
-import gradio as gr
-os.system("pip install gdown")
-os.system("gdown https://drive.google.com/uc?id=1riNxV1BWMAXfmWZ3LrQbEkvzV8f7lOCp")
-
-opts = onnxruntime.SessionOptions()
-opts.intra_op_num_threads = 32
-onnx_session = onnxruntime.InferenceSession("face_paint_512_v2_0.onnx",sess_options=opts)
-
-input_name = onnx_session.get_inputs()[0].name
-output_name = onnx_session.get_outputs()[0].name
-
-side_length = 512
-
-import cv2 as cv
-import numpy as np
 from PIL import Image
+import torch
 
-def inference(img):
-    image = np.array(img) 
-    image = image[:, :, ::-1].copy() 
-    image = cv.resize(image, dsize=(side_length, side_length))
-    x = cv.cvtColor(image, cv.COLOR_BGR2RGB)
-    
-    x = np.array(x, dtype=np.float32)
-    x = x.transpose(2, 0, 1)
-    x = x * 2 - 1
-    x = x.reshape(-1, 3, side_length, side_length)
-    
-    onnx_result = onnx_session.run([output_name], {input_name: x})
-    
-    onnx_result = np.array(onnx_result).squeeze()
-    onnx_result = (onnx_result * 0.5 + 0.5).clip(0, 1)
-    onnx_result = onnx_result * 255
-    
-    onnx_result = onnx_result.transpose(1, 2, 0).astype('uint8')
-    onnx_result = cv.cvtColor(onnx_result, cv.COLOR_RGB2BGR)
-    
+face2paint = torch.hub.load(
+    'bryandlee/animegan2-pytorch:main', 'face2paint', 
+    size=512, device="cuda"
+)
 
-    img = cv.cvtColor(onnx_result, cv.COLOR_BGR2RGB)
-    im_pil = Image.fromarray(img)
-    return im_pil
+def inference(img):
+    out = face2paint(model, img)
+    return out
       
   
 title = "Animeganv2"