Add trust remote code option and improve HF Token handling in model conversion

app.py

@@ -2,9 +2,12 @@ import logging
 import os
 import subprocess
 import sys
+import tempfile
+import tarfile
+import shutil
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 from urllib.request import urlopen, urlretrieve
 
 import streamlit as st
@@ -20,6 +23,7 @@ class Config:
 
     hf_token: str
     hf_username: str
+    is_using_user_token: bool
     transformers_version: str = "3.5.0"
     hf_base_url: str = "https://huggingface.co"
     transformers_base_url: str = (
@@ -32,18 +36,32 @@ class Config:
         """Create config from environment variables and secrets."""
         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:
             hf_username = (
                 os.getenv("SPACE_AUTHOR_NAME") or whoami(token=system_token)["name"]
             )
+
         hf_token = user_token or system_token
 
         if not hf_token:
-            raise ValueError("HF_TOKEN must be set")
+            raise ValueError(
+                "When the user token is not provided, the system token must be set."
+            )
 
-        return cls(hf_token=hf_token, hf_username=hf_username)
+        return cls(
+            hf_token=hf_token,
+            hf_username=hf_username,
+            is_using_user_token=bool(user_token),
+        )
+
+    def get_env_vars(self) -> Dict[str, str]:
+        """Get environment variables with HF_TOKEN set appropriately."""
+        env_vars = os.environ.copy()
+        env_vars["HF_TOKEN"] = self.hf_token
+        return env_vars
 
 
 class ModelConverter:
@@ -82,9 +100,6 @@ class ModelConverter:
 
     def _extract_archive(self, archive_path: Path) -> None:
         """Extract the downloaded archive."""
-        import tarfile
-        import tempfile
-
         with tempfile.TemporaryDirectory() as tmp_dir:
             with tarfile.open(archive_path, "r:gz") as tar:
                 tar.extractall(tmp_dir)
@@ -92,43 +107,46 @@ class ModelConverter:
             extracted_folder = next(Path(tmp_dir).iterdir())
             extracted_folder.rename(self.config.repo_path)
 
-    def convert_model(self, input_model_id: str, trust_remote_code=False) -> Tuple[bool, Optional[str]]:
+    def _run_conversion_subprocess(
+        self, input_model_id: str, extra_args: List[str] = None
+    ) -> subprocess.CompletedProcess:
+        """Run the conversion subprocess with the given arguments."""
+        cmd = [
+            sys.executable,
+            "-m",
+            "scripts.convert",
+            "--quantize",
+            "--model_id",
+            input_model_id,
+        ]
+
+        if extra_args:
+            cmd.extend(extra_args)
+
+        return subprocess.run(
+            cmd,
+            cwd=self.config.repo_path,
+            capture_output=True,
+            text=True,
+            env=self.config.get_env_vars(),
+        )
+
+    def convert_model(
+        self, input_model_id: str, trust_remote_code=False
+    ) -> Tuple[bool, Optional[str]]:
         """Convert the model to ONNX format."""
         try:
             if trust_remote_code:
-                if st.session_state.get("user_hf_token") != "":
-                    result = subprocess.run(
-                        [
-                            sys.executable,
-                            "-m",
-                            "scripts.convert",
-                            "--quantize",
-                            "--trust_remote_code",
-                            "--model_id",
-                            input_model_id,
-                        ],
-                        cwd=self.config.repo_path,
-                        capture_output=True,
-                        text=True,
-                        env={},
+                if not self.config.is_using_user_token:
+                    raise Exception(
+                        "Trust Remote Code requires your own HuggingFace token."
                     )
-                else:
-                    raise Exception("Trust Remote Code requires your own HuggingFace token.")
-            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={},
+
+                result = self._run_conversion_subprocess(
+                    input_model_id, extra_args=["--trust_remote_code"]
                 )
+            else:
+                result = self._run_conversion_subprocess(input_model_id)
 
             if result.returncode != 0:
                 return False, result.stderr
@@ -158,8 +176,6 @@ class ModelConverter:
         except Exception as e:
             return str(e)
         finally:
-            import shutil
-
             shutil.rmtree(model_folder_path, ignore_errors=True)
 
     def generate_readme(self, imi: str):
@@ -197,9 +213,11 @@ def main():
             type="password",
             key="user_hf_token",
         )
-        trust_remote_code = st.toggle("Trust Remote Code?")
+        trust_remote_code = st.toggle("Optional: Trust Remote Code.")
         if trust_remote_code:
-            st.warning("Remote code could be used for malicious purposes. Make sure you trust the code fully. You must use your own Hugging Face write token.")
+            st.warning(
+                "This option should only be enabled for repositories you trust and in which you have read the code, as it will execute arbitrary code present in the model repository. When this option is enabled, you must use your own Hugging Face write token."
+            )
 
         if config.hf_username == input_model_id.split("/")[0]:
             same_repo = st.checkbox(
@@ -229,7 +247,9 @@ def main():
             return
 
         with st.spinner("Converting model..."):
-            success, stderr = converter.convert_model(input_model_id, trust_remote_code=trust_remote_code)
+            success, stderr = converter.convert_model(
+                input_model_id, trust_remote_code=trust_remote_code
+            )
             if not success:
                 st.error(f"Conversion failed: {stderr}")
                 return
@@ -253,4 +273,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()