fix

app.py

@@ -8,6 +8,8 @@ from onnx_export import convert
 from huggingface_hub import HfApi, Repository
 
 
+# TODO: save stats about the Space?
+"""
 DATASET_REPO_URL = "https://huggingface.co/datasets/safetensors/conversions"
 DATA_FILENAME = "data.csv"
 DATA_FILE = os.path.join("data", DATA_FILENAME)
@@ -17,7 +19,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 repo: Optional[Repository] = None
 if HF_TOKEN:
     repo = Repository(local_dir="data", clone_from=DATASET_REPO_URL, token=HF_TOKEN)
-
+"""
 
 def onnx_export(token: str, model_id: str, task: str) -> str:
     if token == "" or model_id == "":
@@ -28,9 +30,16 @@ def onnx_export(token: str, model_id: str, task: str) -> str:
         """
     try:
         api = HfApi(token=token)
-        commit_info = convert(api=api, model_id=model_id, task=task)
+
+        error, commit_info = convert(api=api, model_id=model_id, task=task)
+        if error != "":
+            return error
+        
         print("[commit_info]", commit_info)
 
+        return f"### Success 🔥 Yay! This model was successfully converted and a PR was open using your token, here: {commit_info.pr_url}]({commit_info.pr_url})"
+    except Exception as e:
+        return f"### Error: {e}"
         """
         # save in a private dataset:
         if repo is not None:
@@ -51,8 +60,6 @@ def onnx_export(token: str, model_id: str, task: str) -> str:
 
         return f"### Success 🔥 Yay! This model was successfully converted and a PR was open using your token, here: {commit_info.pr_url}]({commit_info.pr_url})"
         """
-    except Exception as e:
-        return f"### Error: {e}"
 
 
 DESCRIPTION = """
@@ -65,14 +72,14 @@ The steps are the following:
 """
 
 demo = gr.Interface(
-    title="Convert any model to Safetensors and open a PR",
+    title="POC to convert any PyTorch model to ONNX",
     description=DESCRIPTION,
     allow_flagging="never",
     article="Check out the [Optimum repo on GitHub](https://github.com/huggingface/optimum)",
     inputs=[
-        gr.Text(max_lines=1, label="your_hf_token"),
+        gr.Text(max_lines=1, label="Hugging Face token"),
         gr.Text(max_lines=1, label="model_id"),
-        gr.Text(max_lines=1, label="task")
+        gr.Text(value="auto", max_lines=1, label="task (can be left blank, will auto-infer)")
     ],
     outputs=[gr.Markdown(label="output")],
     fn=onnx_export,

onnx_export.py

@@ -12,7 +12,7 @@ import os
 import shutil
 import argparse
 
-from typing import Optional
+from typing import Optional, Tuple, List
 
 from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, get_repo_discussions
 from huggingface_hub.file_download import repo_folder_name
@@ -26,7 +26,7 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
         if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title:
             return discussion
 
-def convert_onnx(model_id: str, task: str, folder: str):
+def convert_onnx(model_id: str, task: str, folder: str) -> List:
 
     # Allocate the model
     model = TasksManager.get_model_from_task(task, model_id, framework="pt")
@@ -77,11 +77,17 @@ def convert_onnx(model_id: str, task: str, folder: str):
     return operations
 
 
-def convert(api: "HfApi", model_id: str, task:str, force: bool=False) -> Optional["CommitInfo"]:
+def convert(api: "HfApi", model_id: str, task: str, force: bool = False) -> Tuple[int, "CommitInfo"]:
     pr_title = "Adding ONNX file of this model"
     info = api.model_info(model_id)
     filenames = set(s.rfilename for s in info.siblings)
 
+    if task == "auto":
+        try:
+            task = TasksManager.infer_task_from_model(model_id)
+        except Exception as e:
+            return f"### Error: {e}. Please pass explicitely the task as it could not be infered.", None
+
     with TemporaryDirectory() as d:
         folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
         os.makedirs(folder)
@@ -105,7 +111,7 @@ def convert(api: "HfApi", model_id: str, task:str, force: bool=False) -> Optiona
                 )
         finally:
             shutil.rmtree(folder)
-        return new_pr
+        return "0", new_pr
 
 
 if __name__ == "__main__":