Update app.py

app.py

@@ -118,20 +118,25 @@ class ModelConverter:
         except Exception as e:
             return False, str(e)
 
-    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."""
+    def upload_model(self, input_model_id: str) -> Optional[str]:
+        """Upload the converted model to the `onnx/` subfolder in the existing model repository."""
         try:
-            # Ensure the repository exists (no need to create a new one)
-            self.api.repo_exists(output_model_id, repo_type="model")
-
-            # Define the path to the converted model folder
             model_folder_path = self.config.repo_path / "models" / input_model_id
+            onnx_folder_path = model_folder_path / "onnx"
+
+            # Create the `onnx` subfolder if it doesn't exist
+            onnx_folder_path.mkdir(exist_ok=True)
+
+            # Move the ONNX files to the `onnx` subfolder
+            for file in model_folder_path.iterdir():
+                if file.is_file() and file.suffix == ".onnx":
+                    file.rename(onnx_folder_path / file.name)
 
-            # Upload the folder to the 'onnx/' subfolder in the existing repository
+            # Upload the `onnx` subfolder to the existing repository
             self.api.upload_folder(
-                folder_path=str(model_folder_path),
-                repo_id=output_model_id,
-                path_in_repo="onnx",  # Upload to the 'onnx/' subfolder
+                folder_path=str(onnx_folder_path),
+                repo_id=input_model_id,
+                path_in_repo="onnx",
             )
             return None
         except Exception as e:
@@ -164,12 +169,7 @@ def main():
             key="user_hf_token",
         )
 
-        # 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 model will be uploaded to:")
-        st.code(f"{output_model_url}/tree/onnx", language="plaintext")
+        output_model_url = f"{config.hf_base_url}/{input_model_id}"
 
         if not st.button(label="Proceed", type="primary"):
             return
@@ -184,18 +184,14 @@ def main():
             st.code(stderr)
 
         with st.spinner("Uploading model..."):
-            error = converter.upload_model(input_model_id, output_model_id)
+            error = converter.upload_model(input_model_id)
             if error:
                 st.error(f"Upload failed: {error}")
                 return
 
             st.success("Upload successful!")
-            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",
-            )
+            st.write("You can now go and view the 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")