add sign in with hf button

app.py

@@ -7,7 +7,7 @@ import gradio as gr
 from huggingface_hub import HfApi, Repository
 
 from onnx_export import convert
-
+from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from apscheduler.schedulers.background import BackgroundScheduler
 
 DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/exporters"
@@ -23,20 +23,22 @@ repo: Optional[Repository] = None
 #     repo = Repository(local_dir=DATADIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN)
 
 
-def onnx_export(token: str, model_id: str, task: str, opset: Union[int, str]) -> str:
-    if token == "" or model_id == "":
-        return """
-        ### Invalid input 🐞
+def onnx_export(
+    model_id: str, task: str, opset: Union[int, str], oauth_token: gr.OAuthToken
+) -> str:
+    if oauth_token.token is None:
+        return "You must be logged in to use this space"
+
+    if not model_id:
+        return f"### Invalid input 🐞 Please specify a model name, got {model_id}"
 
-        Please fill a token and model name.
-        """
     try:
         if opset == "":
             opset = None
         else:
             opset = int(opset)
 
-        api = HfApi(token=token)
+        api = HfApi(token=oauth_token.token)
 
         error, commit_info = convert(api=api, model_id=model_id, task=task, opset=opset)
         if error != "0":
@@ -100,22 +102,19 @@ TITLE = """
 
 # for some reason https://huggingface.co/settings/tokens is not showing as a link by default?
 DESCRIPTION = """
-This Space allows you to automatically export 🤗 transformers PyTorch models hosted on the Hugging Face Hub to [ONNX](https://onnx.ai/). It opens a PR on the target model, and it is up to the owner of the original model
+This Space allows you to automatically export 🤗 transformers, diffusers, timm and sentence-transformers PyTorch models hosted on the Hugging Face Hub to [ONNX](https://onnx.ai/). It opens a PR on the target model, and it is up to the owner of the original model
 to merge the PR to allow people to leverage the ONNX standard to share and use the model on a wide range of devices!
 
 Once exported, the model can, for example, be used in the [🤗 Optimum](https://huggingface.co/docs/optimum/) library closely following the transformers API.
 Check out [this guide](https://huggingface.co/docs/optimum/main/en/onnxruntime/usage_guides/models) to see how!
 
-The steps are as following:
-- Paste a read-access token from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens). Read access is enough given that we will open a PR against the source repo.
-- Input a model id from the Hub (for example: [textattack/distilbert-base-cased-CoLA](https://huggingface.co/textattack/distilbert-base-cased-CoLA))
-- Click "Export to ONNX"
-- That's it! You'll get feedback on if the export was successful or not, and if it was, you'll get the URL of the opened PR!
-
 Note: in case the model to export is larger than 2 GB, it will be saved in a subfolder called `onnx/`. To load it from Optimum, the argument `subfolder="onnx"` should be provided.
 """
 
 with gr.Blocks() as demo:
+    gr.Markdown("You must be logged to use this space")
+    gr.LoginButton(min_width=250)
+
     gr.HTML(TTILE_IMAGE)
     gr.HTML(TITLE)
 
@@ -124,14 +123,10 @@ with gr.Blocks() as demo:
             gr.Markdown(DESCRIPTION)
 
         with gr.Column(scale=50):
-            input_token = gr.Textbox(
-                max_lines=1,
-                label="Hugging Face token",
-            )
-            input_model = gr.Textbox(
-                max_lines=1,
-                label="Model name",
-                placeholder="textattack/distilbert-base-cased-CoLA",
+            input_model = HuggingfaceHubSearch(
+                label="Hub model ID",
+                placeholder="Search for model ID on the hub",
+                search_type="model",
             )
             input_task = gr.Textbox(
                 value="auto",
@@ -143,19 +138,20 @@ with gr.Blocks() as demo:
                 max_lines=1,
                 label="ONNX opset (optional, can be left blank)",
             )
-
             btn = gr.Button("Export to ONNX")
             output = gr.Markdown(label="Output")
 
     btn.click(
         fn=onnx_export,
-        inputs=[input_token, input_model, input_task, onnx_opset],
+        inputs=[input_model, input_task, onnx_opset],
         outputs=output,
     )
 
+
 def restart_space():
     HfApi().restart_space(repo_id="onnx/export", token=HF_TOKEN, factory_reboot=True)
 
+
 scheduler = BackgroundScheduler()
 scheduler.add_job(restart_space, "interval", seconds=21600)
 scheduler.start()

onnx_export.py

@@ -82,7 +82,9 @@ def convert(
         try:
             pr = previous_pr(api, model_id, pr_title)
             if "model.onnx" in filenames and not force:
-                raise Exception(f"Model {model_id} is already converted, skipping the export.")
+                raise Exception(
+                    f"Model {model_id} is already converted, skipping the export."
+                )
             elif pr is not None and not force:
                 url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
                 new_pr = pr

requirements.txt

@@ -5,3 +5,5 @@ diffusers
 sentence-transformers
 timm
 APScheduler
+gradio[oauth]>=4.37.2
+gradio_huggingfacehub_search==0.0.6