Update app.py

app.py

@@ -118,20 +118,20 @@ class ModelConverter:
         except Exception as e:
             return False, str(e)
 
-    def upload_model(self, input_model_id: str) -> Optional[str]:
-        """Upload the ONNX folder to the same model path on Hugging Face."""
+    def upload_model(self, input_model_id: str, output_model_id: str) -> Optional[str]:
+        """Upload the converted model to Hugging Face under the 'onnx/' subfolder."""
         try:
-            # Define the ONNX folder path
-            onnx_folder_path = self.config.repo_path / "models" / input_model_id / "onnx"
+            # Ensure the repository exists (no need to create a new one)
+            self.api.repo_exists(output_model_id, repo_type="model")
 
-            # Ensure the ONNX folder exists
-            if not onnx_folder_path.exists():
-                return "ONNX folder not found. Conversion may have failed."
+            # Define the path to the converted model folder
+            model_folder_path = self.config.repo_path / "models" / input_model_id
 
-            # Upload the ONNX folder to the same model path
+            # Upload the folder to the 'onnx/' subfolder in the existing repository
             self.api.upload_folder(
-                folder_path=str(onnx_folder_path), 
-                repo_id=input_model_id
+                folder_path=str(model_folder_path),
+                repo_id=output_model_id,
+                path_in_repo="onnx",  # Upload to the 'onnx/' subfolder
             )
             return None
         except Exception as e:
@@ -139,8 +139,7 @@ class ModelConverter:
         finally:
             import shutil
 
-            # Clean up the ONNX folder after upload
-            shutil.rmtree(onnx_folder_path, ignore_errors=True)
+            shutil.rmtree(model_folder_path, ignore_errors=True)
 
 
 def main():
@@ -165,10 +164,12 @@ def main():
             key="user_hf_token",
         )
 
-        output_model_url = f"{config.hf_base_url}/{input_model_id}"
+        # Use the same model ID for the output (no need to append '-ONNX')
+        output_model_id = input_model_id
+        output_model_url = f"{config.hf_base_url}/{output_model_id}"
 
-        st.write(f"URL where the ONNX folder will be uploaded:")
-        st.code(output_model_url, language="plaintext")
+        st.write(f"URL where the ONNX model will be uploaded to:")
+        st.code(f"{output_model_url}/tree/onnx", language="plaintext")
 
         if not st.button(label="Proceed", type="primary"):
             return
@@ -182,15 +183,19 @@ def main():
             st.success("Conversion successful!")
             st.code(stderr)
 
-        with st.spinner("Uploading ONNX folder..."):
-            error = converter.upload_model(input_model_id)
+        with st.spinner("Uploading model..."):
+            error = converter.upload_model(input_model_id, output_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")
+            st.write("You can now go and view the ONNX model on Hugging Face!")
+            st.link_button(
+                f"Go to {output_model_id}/onnx",
+                f"{output_model_url}/tree/onnx",
+                type="primary",
+            )
 
     except Exception as e:
         logger.exception("Application error")