Update app.py

app.py

@@ -6,24 +6,48 @@ from tempfile import TemporaryDirectory
 import gradio as gr
 
 
-def export_to_onnx(model_id: str, task: str, opset: int, progress=gr.Progress()):
+def export_to_onnx(
+    model_id: str,
+    task: str,
+    opset: int,
+    push_to_hub: bool,
+    repo_id: str,
+    private_repo: bool,
+    progress=gr.Progress(),
+):
     model_id = model_id.strip()
+    repo_id = repo_id.strip()
 
     if not model_id:
-        raise gr.Error("Enter a Hugging Face model ID, for example: distilbert/distilbert-base-uncased.")
+        raise gr.Error(
+            "Enter a Hugging Face model ID, for example: "
+            "distilbert/distilbert-base-uncased."
+        )
+
+    if push_to_hub and not repo_id:
+        raise gr.Error(
+            "Enter a destination model repository, for example: "
+            "your-username/my-model-onnx."
+        )
+
+    if push_to_hub and not os.getenv("HF_TOKEN"):
+        raise gr.Error(
+            "This Space does not have an HF_TOKEN secret configured. "
+            "Add a Hugging Face write token in Settings → Variables and secrets."
+        )
 
     if not task:
         task = "auto"
 
-    progress(0.05, desc="Preparing export...")
-
     try:
-        # Import here so a dependency issue produces a UI error instead of
-        # preventing the Space from starting.
+        from huggingface_hub import HfApi
         from optimum.exporters.onnx import main_export
 
+        progress(0.05, desc="Preparing export...")
+
         with TemporaryDirectory() as temp_dir:
-            output_dir = Path(temp_dir) / "onnx"
+            temp_dir = Path(temp_dir)
+            output_dir = temp_dir / "onnx"
 
             progress(0.15, desc="Downloading model and exporting to ONNX...")
 
@@ -34,38 +58,62 @@ def export_to_onnx(model_id: str, task: str, opset: int, progress=gr.Progress())
                 "opset": int(opset),
             }
 
-            # "auto" is handled by not supplying task.
             if task == "auto":
                 export_args.pop("task")
 
             main_export(**export_args)
 
-            progress(0.85, desc="Creating download archive...")
+            hub_url = None
+
+            if push_to_hub:
+                progress(0.80, desc="Creating or updating Hub repository...")
+
+                api = HfApi(token=os.environ["HF_TOKEN"])
+
+                api.create_repo(
+                    repo_id=repo_id,
+                    repo_type="model",
+                    private=private_repo,
+                    exist_ok=True,
+                )
+
+                api.upload_folder(
+                    folder_path=str(output_dir),
+                    repo_id=repo_id,
+                    repo_type="model",
+                    commit_message=(
+                        f"Export {model_id} to ONNX "
+                        f"(task={task}, opset={int(opset)})"
+                    ),
+                )
+
+                hub_url = f"https://huggingface.co/{repo_id}"
+
+            progress(0.90, desc="Creating download archive...")
 
             safe_name = model_id.replace("/", "--").replace("\\", "--")
-            archive_base = Path(temp_dir) / f"{safe_name}-onnx"
+            archive_base = temp_dir / f"{safe_name}-onnx"
+
             zip_path = shutil.make_archive(
                 base_name=str(archive_base),
                 format="zip",
                 root_dir=str(output_dir),
             )
 
-            # TemporaryDirectory is deleted after this function returns,
-            # so copy the ZIP to a persistent Gradio temp location.
             final_path = Path("/tmp") / f"{safe_name}-onnx.zip"
             shutil.copy2(zip_path, final_path)
 
-        progress(1.0, desc="Finished.")
+            progress(1.0, desc="Finished.")
 
-        return (
-            f"Successfully exported `{model_id}` to ONNX.",
-            str(final_path),
-        )
+            message = f"Successfully exported `{model_id}` to ONNX."
+
+            if hub_url:
+                message += f"\n\nPushed to [{repo_id}]({hub_url})."
+
+            return message, str(final_path)
 
     except Exception as error:
-        raise gr.Error(
-            f"Export failed: {type(error).__name__}: {error}"
-        )
+        raise gr.Error(f"Export failed: {type(error).__name__}: {error}")
 
 
 with gr.Blocks(title="Hugging Face to ONNX") as demo:
@@ -73,9 +121,11 @@ with gr.Blocks(title="Hugging Face to ONNX") as demo:
         """
 # Hugging Face Model to ONNX
 
-Enter a public Hugging Face model ID, choose a task, and download its ONNX export.
+Enter a public Hugging Face model ID, choose a task, then download the ONNX
+export or optionally push it to a Hugging Face model repository.
 
-**Examples:** `distilbert/distilbert-base-uncased`, `google-bert/bert-base-uncased`, or `gpt2`
+**Examples:** `distilbert/distilbert-base-uncased`,
+`google-bert/bert-base-uncased`, or `gpt2`
 """
     )
 
@@ -113,6 +163,27 @@ Enter a public Hugging Face model ID, choose a task, and download its ONNX expor
             scale=1,
         )
 
+    gr.Markdown("## Optional Hub upload")
+
+    with gr.Row():
+        push_to_hub = gr.Checkbox(
+            label="Push ONNX export to Hugging Face Hub",
+            value=False,
+            scale=1,
+        )
+
+        repo_id = gr.Textbox(
+            label="Destination model repository",
+            placeholder="your-username/model-name-onnx",
+            scale=3,
+        )
+
+        private_repo = gr.Checkbox(
+            label="Private repository",
+            value=False,
+            scale=1,
+        )
+
     export_button = gr.Button("Export to ONNX", variant="primary")
 
     status = gr.Markdown()
@@ -120,18 +191,60 @@ Enter a public Hugging Face model ID, choose a task, and download its ONNX expor
 
     export_button.click(
         fn=export_to_onnx,
-        inputs=[model_id, task, opset],
+        inputs=[
+            model_id,
+            task,
+            opset,
+            push_to_hub,
+            repo_id,
+            private_repo,
+        ],
         outputs=[status, download],
     )
 
     gr.Examples(
         examples=[
-            ["distilbert/distilbert-base-uncased", "feature-extraction", 17],
-            ["google-bert/bert-base-uncased", "feature-extraction", 17],
-            ["distilbert/distilbert-base-cased-distilled-squad", "question-answering", 17],
-            ["gpt2", "text-generation", 17],
+            [
+                "distilbert/distilbert-base-uncased",
+                "feature-extraction",
+                17,
+                False,
+                "",
+                False,
+            ],
+            [
+                "google-bert/bert-base-uncased",
+                "feature-extraction",
+                17,
+                False,
+                "",
+                False,
+            ],
+            [
+                "distilbert/distilbert-base-cased-distilled-squad",
+                "question-answering",
+                17,
+                False,
+                "",
+                False,
+            ],
+            [
+                "gpt2",
+                "text-generation",
+                17,
+                False,
+                "",
+                False,
+            ],
+        ],
+        inputs=[
+            model_id,
+            task,
+            opset,
+            push_to_hub,
+            repo_id,
+            private_repo,
         ],
-        inputs=[model_id, task, opset],
     )
 
 if __name__ == "__main__":