Update app.py

app.py

@@ -92,23 +92,39 @@ class ModelConverter:
             extracted_folder = next(Path(tmp_dir).iterdir())
             extracted_folder.rename(self.config.repo_path)
 
-    def convert_model(self, input_model_id: str) -> Tuple[bool, Optional[str]]:
+    def convert_model(self, input_model_id: str, trust_remote_code=False) -> Tuple[bool, Optional[str]]:
         """Convert the model to ONNX format."""
         try:
-            result = subprocess.run(
-                [
-                    sys.executable,
-                    "-m",
-                    "scripts.convert",
-                    "--quantize",
-                    "--model_id",
-                    input_model_id,
-                ],
-                cwd=self.config.repo_path,
-                capture_output=True,
-                text=True,
-                env={},
-            )
+            if trust_remote_code:
+                result = subprocess.run(
+                    [
+                        sys.executable,
+                        "-m",
+                        "scripts.convert",
+                        "--quantize",
+                        "--model_id",
+                        input_model_id,
+                    ],
+                    cwd=self.config.repo_path,
+                    capture_output=True,
+                    text=True,
+                    env={},
+                )
+            else:
+                result = subprocess.run(
+                    [
+                        sys.executable,
+                        "-m",
+                        "scripts.convert",
+                        "--quantize",
+                        "--model_id",
+                        input_model_id,
+                    ],
+                    cwd=self.config.repo_path,
+                    capture_output=True,
+                    text=True,
+                    env={},
+                )
 
             if result.returncode != 0:
                 return False, result.stderr
@@ -177,6 +193,9 @@ def main():
             type="password",
             key="user_hf_token",
         )
+        trust_remote_code = st.toggle("Trust Remote Code?")
+        if trust_remote_code:
+            st.warning("Remote code could be used for malicious purposes. Make sure you trust the code fully."
 
         if config.hf_username == input_model_id.split("/")[0]:
             same_repo = st.checkbox(
@@ -206,7 +225,7 @@ def main():
             return
 
         with st.spinner("Converting model..."):
-            success, stderr = converter.convert_model(input_model_id)
+            success, stderr = converter.convert_model(input_model_id, trust_remote_code=trust_remote_code)
             if not success:
                 st.error(f"Conversion failed: {stderr}")
                 return