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 'neuromancer.ai.lover@gmail.com'."
 
 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,