Added new model

app.py

@@ -18,24 +18,41 @@ logging.basicConfig(
 MODEL_IMG_SIZE = 256
 def load_model():
     REPO_ID = "Podtekatel/ARCNEGAN"
-    FILENAME = "arcane_exp_203_ep_399.onnx"
+    FILENAME_OLD = "arcane_exp_203_ep_399.onnx"
+    FILENAME_NEW = "arcane_exp_206_ep_138.onnx"
 
-    global model
-    global pipeline
+    global model_old
+    global model_new
+    global pipeline_old
+    global pipeline_new
 
+    # Old model
     model_path = cached_download(
-        hf_hub_url(REPO_ID, FILENAME), use_auth_token=os.getenv('HF_TOKEN')
+        hf_hub_url(REPO_ID, FILENAME_OLD), use_auth_token=os.getenv('HF_TOKEN')
     )
-    model = ONNXModel(model_path)
+    model_old = ONNXModel(model_path)
 
-    pipeline = VSNetModelPipeline(model, StatRetinaFaceDetector(MODEL_IMG_SIZE), background_resize=1024, no_detected_resize=1024)
-    return model
+    pipeline_old = VSNetModelPipeline(model_old, StatRetinaFaceDetector(MODEL_IMG_SIZE), background_resize=1024, no_detected_resize=1024)
 
+    # New model
+    model_path = cached_download(
+        hf_hub_url(REPO_ID, FILENAME_NEW), use_auth_token=os.getenv('HF_TOKEN')
+    )
+    model_new = ONNXModel(model_path)
+
+    pipeline_new = VSNetModelPipeline(model_new, StatRetinaFaceDetector(MODEL_IMG_SIZE), background_resize=1024,
+                                      no_detected_resize=1024)
+
+    return model_old, model_new
 load_model()
 
-def inference(img):
+def inference(img, ver):
     img = np.array(img)
-    out_img = pipeline(img)
+    if ver == 'version 2':
+        out_img = pipeline_new(img)
+    else:
+        out_img = pipeline_old(img)
+
     out_img = Image.fromarray(out_img)
     return out_img
 
@@ -52,11 +69,11 @@ article = "This is one of my successful experiments on style transfer. I've buil
           "If you want to use this app or integrate this model into yours, please contact me at email '[email protected]'."
 
 imgs_folder = 'demo'
-examples = [[os.path.join(imgs_folder, img_filename)] for img_filename in sorted(os.listdir(imgs_folder))]
+examples = [[os.path.join(imgs_folder, img_filename), version] for img_filename in sorted(os.listdir(imgs_folder)) for version in ['version 2']]
 
 demo = gr.Interface(
     fn=inference,
-    inputs=[gr.inputs.Image(type="pil")],
+    inputs=[gr.inputs.Image(type="pil"), gr.inputs.Radio(['version 1', 'version 2'], type="value", default='version 2', label='version')],
     outputs=gr.outputs.Image(type="pil"),
     title=title,
     description=description,