Update app.py

app.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Optional, Tuple
 from urllib.request import urlopen, urlretrieve
 
-import gradio as gr
+import streamlit as st
 from huggingface_hub import HfApi, whoami
 
 logging.basicConfig(level=logging.INFO)
@@ -28,9 +28,10 @@ class Config:
     repo_path: Path = Path("./transformers.js")
 
     @classmethod
-    def from_env(cls, user_token: Optional[str] = None) -> "Config":
+    def from_env(cls) -> "Config":
         """Create config from environment variables and secrets."""
-        system_token = os.getenv("HF_TOKEN")
+        system_token = st.secrets.get("HF_TOKEN")
+        user_token = st.session_state.get("user_hf_token")
         if user_token:
             hf_username = whoami(token=user_token)["name"]
         else:
@@ -118,22 +119,19 @@ class ModelConverter:
             return False, str(e)
 
     def upload_model(self, input_model_id: str) -> Optional[str]:
-        """Upload the converted model to Hugging Face under the onnx/ folder."""
+        """Upload the ONNX folder to the same model path on Hugging Face."""
         try:
-            model_folder_path = self.config.repo_path / "models" / input_model_id
-            onnx_folder_path = model_folder_path / "onnx"
+            # Define the ONNX folder path
+            onnx_folder_path = self.config.repo_path / "models" / input_model_id / "onnx"
 
-            # Create the onnx folder if it doesn't exist
-            onnx_folder_path.mkdir(exist_ok=True)
+            # Ensure the ONNX folder exists
+            if not onnx_folder_path.exists():
+                return "ONNX folder not found. Conversion may have failed."
 
-            # Move the converted model files to the onnx folder
-            for file in model_folder_path.iterdir():
-                if file.is_file() and file.suffix == ".onnx":
-                    file.rename(onnx_folder_path / file.name)
-
-            # Upload the onnx folder to the same model path
+            # Upload the ONNX folder to the same model path
             self.api.upload_folder(
-                folder_path=str(onnx_folder_path), repo_id=input_model_id, path_in_repo="onnx"
+                folder_path=str(onnx_folder_path), 
+                repo_id=input_model_id
             )
             return None
         except Exception as e:
@@ -141,48 +139,63 @@ class ModelConverter:
         finally:
             import shutil
 
-            shutil.rmtree(model_folder_path, ignore_errors=True)
+            # Clean up the ONNX folder after upload
+            shutil.rmtree(onnx_folder_path, ignore_errors=True)
+
 
+def main():
+    """Main application entry point."""
+    st.write("## Convert a Hugging Face model to ONNX")
 
-def convert_and_upload_model(input_model_id: str, user_hf_token: str = None):
-    """Function to handle model conversion and upload."""
     try:
-        config = Config.from_env(user_hf_token)
+        config = Config.from_env()
         converter = ModelConverter(config)
         converter.setup_repository()
 
-        if converter.api.repo_exists(input_model_id):
-            with gr.Progress() as progress:
-                progress(0, desc="Converting model...")
-                success, stderr = converter.convert_model(input_model_id)
-                if not success:
-                    return f"Conversion failed: {stderr}"
+        input_model_id = st.text_input(
+            "Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"
+        )
 
-                progress(0.5, desc="Uploading model...")
-                error = converter.upload_model(input_model_id)
-                if error:
-                    return f"Upload failed: {error}"
+        if not input_model_id:
+            return
 
-                return f"Model uploaded to {input_model_id}/onnx"
-        else:
-            return f"Model {input_model_id} does not exist on Hugging Face."
+        st.text_input(
+            f"Optional: Your Hugging Face write token. Fill it if you want to upload the model under your account.",
+            type="password",
+            key="user_hf_token",
+        )
+
+        output_model_url = f"{config.hf_base_url}/{input_model_id}"
+
+        st.write(f"URL where the ONNX folder will be uploaded:")
+        st.code(output_model_url, language="plaintext")
+
+        if not st.button(label="Proceed", type="primary"):
+            return
+
+        with st.spinner("Converting model..."):
+            success, stderr = converter.convert_model(input_model_id)
+            if not success:
+                st.error(f"Conversion failed: {stderr}")
+                return
+
+            st.success("Conversion successful!")
+            st.code(stderr)
+
+        with st.spinner("Uploading ONNX folder..."):
+            error = converter.upload_model(input_model_id)
+            if error:
+                st.error(f"Upload failed: {error}")
+                return
+
+            st.success("Upload successful!")
+            st.write("The ONNX folder has been uploaded to the existing model on Hugging Face!")
+            st.link_button(f"Go to {input_model_id}", output_model_url, type="primary")
 
     except Exception as e:
         logger.exception("Application error")
-        return f"An error occurred: {str(e)}"
-
+        st.error(f"An error occurred: {str(e)}")
 
-# Gradio Interface
-iface = gr.Interface(
-    fn=convert_and_upload_model,
-    inputs=[
-        gr.Textbox(label="Hugging Face Model ID", placeholder="Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"),
-        gr.Textbox(label="Hugging Face Write Token (Optional)", type="password", placeholder="Fill it if you want to upload the model under your account.")
-    ],
-    outputs=gr.Textbox(label="Output"),
-    title="Convert a Hugging Face model to ONNX",
-    description="This tool converts a Hugging Face model to ONNX format and uploads it to the same model path under an `onnx/` folder."
-)
 
 if __name__ == "__main__":
-    iface.launch()
+    main()