app.py

import logging
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple
from urllib.request import urlopen, urlretrieve

import gradio as gr
from huggingface_hub import HfApi, whoami

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class Config:
    """Application configuration."""

    hf_token: str
    hf_username: str
    transformers_version: str = "3.0.0"
    hf_base_url: str = "https://huggingface.co"
    transformers_base_url: str = (
        "https://github.com/xenova/transformers.js/archive/refs"
    )
    repo_path: Path = Path("./transformers.js")

    @classmethod
    def from_env(cls) -> "Config":
        """Create config from environment variables and secrets."""
        system_token = os.getenv("HF_TOKEN")
        user_token = gr.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")

        return cls(hf_token=hf_token, hf_username=hf_username)


class ModelConverter:
    """Handles model conversion and upload operations."""

    def __init__(self, config: Config):
        self.config = config
        self.api = HfApi(token=config.hf_token)

    def _get_ref_type(self) -> str:
        """Determine the reference type for the transformers repository."""
        url = f"{self.config.transformers_base_url}/tags/{self.config.transformers_version}.tar.gz"
        try:
            return "tags" if urlopen(url).getcode() == 200 else "heads"
        except Exception as e:
            logger.warning(f"Failed to check tags, defaulting to heads: {e}")
            return "heads"

    def setup_repository(self) -> None:
        """Download and setup transformers repository if needed."""
        if self.config.repo_path.exists():
            return

        ref_type = self._get_ref_type()
        archive_url = f"{self.config.transformers_base_url}/{ref_type}/{self.config.transformers_version}.tar.gz"
        archive_path = Path(f"./transformers_{self.config.transformers_version}.tar.gz")

        try:
            urlretrieve(archive_url, archive_path)
            self._extract_archive(archive_path)
            logger.info("Repository downloaded and extracted successfully")
        except Exception as e:
            raise RuntimeError(f"Failed to setup repository: {e}")
        finally:
            archive_path.unlink(missing_ok=True)

    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)

            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]]:
        """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 result.returncode != 0:
                return False, result.stderr

            return True, result.stderr

        except Exception as e:
            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."""
        try:
            model_folder_path = self.config.repo_path / "models" / input_model_id
            onnx_folder_path = model_folder_path / "onnx"

            # Create the onnx folder if it doesn't exist
            onnx_folder_path.mkdir(exist_ok=True)

            # 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
            self.api.upload_folder(
                folder_path=str(onnx_folder_path), repo_id=input_model_id, path_in_repo="onnx"
            )
            return None
        except Exception as e:
            return str(e)
        finally:
            import shutil

            shutil.rmtree(model_folder_path, ignore_errors=True)


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()
        if user_hf_token:
            gr.session_state["user_hf_token"] = user_hf_token
            config = Config.from_env()

        converter = ModelConverter(config)
        converter.setup_repository()

        if converter.api.repo_exists(input_model_id):
            with gr.spinner("Converting model..."):
                success, stderr = converter.convert_model(input_model_id)
                if not success:
                    return f"Conversion failed: {stderr}"

                gr.Info("Conversion successful!")
                gr.Code(stderr)

            with gr.spinner("Uploading model..."):
                error = converter.upload_model(input_model_id)
                if error:
                    return f"Upload failed: {error}"

                gr.Info("Upload successful!")
                return f"Model uploaded to {input_model_id}/onnx"
        else:
            return f"Model {input_model_id} does not exist on Hugging Face."

    except Exception as e:
        logger.exception("Application error")
        return 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()