app.py

import argparse
import json
import os
import shutil
from collections import defaultdict
from inspect import signature
from tempfile import TemporaryDirectory
from typing import Dict, List, Optional, Set

import torch

from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download
from huggingface_hub.file_download import repo_folder_name
from safetensors.torch import load_file, save_file
from transformers import AutoConfig
from transformers.pipelines.base import infer_framework_load_model

import csv
from datetime import datetime
import os
from typing import Optional
from huggingface_hub import HfApi, Repository

import gradio as gr

class AlreadyExists(Exception):
    pass


def shared_pointers(tensors):
    ptrs = defaultdict(list)
    for k, v in tensors.items():
        ptrs[v.data_ptr()].append(k)
    failing = []
    for ptr, names in ptrs.items():
        if len(names) > 1:
            failing.append(names)
    return failing


def check_file_size(sf_filename: str, pt_filename: str):
    sf_size = os.stat(sf_filename).st_size
    pt_size = os.stat(pt_filename).st_size

    if (sf_size - pt_size) / pt_size > 0.01:
        raise RuntimeError(
            f"""The file size different is more than 1%:
         - {sf_filename}: {sf_size}
         - {pt_filename}: {pt_size}
         """
        )


def rename(pt_filename: str) -> str:
    filename, ext = os.path.splitext(pt_filename)
    local = f"{filename}.safetensors"
    local = local.replace("pytorch_model", "model")
    return local


def convert_multi(model_id: str, folder: str) -> List["CommitOperationAdd"]:
    filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin.index.json")
    with open(filename, "r") as f:
        data = json.load(f)

    filenames = set(data["weight_map"].values())
    local_filenames = []
    for filename in filenames:
        pt_filename = hf_hub_download(repo_id=model_id, filename=filename)

        sf_filename = rename(pt_filename)
        sf_filename = os.path.join(folder, sf_filename)
        convert_file(pt_filename, sf_filename)
        local_filenames.append(sf_filename)

    index = os.path.join(folder, "model.safetensors.index.json")
    with open(index, "w") as f:
        newdata = {k: v for k, v in data.items()}
        newmap = {k: rename(v) for k, v in data["weight_map"].items()}
        newdata["weight_map"] = newmap
        json.dump(newdata, f, indent=4)
    local_filenames.append(index)

    operations = [
        CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames
    ]

    return operations


def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:
    pt_filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")

    sf_name = "model.safetensors"
    sf_filename = os.path.join(folder, sf_name)
    convert_file(pt_filename, sf_filename)
    operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)]
    return operations


def convert_file(
    pt_filename: str,
    sf_filename: str,
):
    loaded = torch.load(pt_filename, map_location="cpu")
    if "state_dict" in loaded:
        loaded = loaded["state_dict"]
    shared = shared_pointers(loaded)
    for shared_weights in shared:
        for name in shared_weights[1:]:
            loaded.pop(name)

    # For tensors to be contiguous
    loaded = {k: v.contiguous() for k, v in loaded.items()}

    dirname = os.path.dirname(sf_filename)
    os.makedirs(dirname, exist_ok=True)
    save_file(loaded, sf_filename, metadata={"format": "pt"})
    check_file_size(sf_filename, pt_filename)
    reloaded = load_file(sf_filename)
    for k in loaded:
        pt_tensor = loaded[k]
        sf_tensor = reloaded[k]
        if not torch.equal(pt_tensor, sf_tensor):
            raise RuntimeError(f"The output tensors do not match for key {k}")


def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]]) -> str:
    errors = []
    for key in ["missing_keys", "mismatched_keys", "unexpected_keys"]:
        pt_set = set(pt_infos[key])
        sf_set = set(sf_infos[key])

        pt_only = pt_set - sf_set
        sf_only = sf_set - pt_set

        if pt_only:
            errors.append(f"{key} : PT warnings contain {pt_only} which are not present in SF warnings")
        if sf_only:
            errors.append(f"{key} : SF warnings contain {sf_only} which are not present in PT warnings")
    return "\n".join(errors)

def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
    try:
        discussions = api.get_repo_discussions(repo_id=model_id)
    except Exception:
        return None
    for discussion in discussions:
        if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title:
            return discussion


def convert_generic(model_id: str, folder: str, filenames: Set[str]) -> List["CommitOperationAdd"]:
    operations = []

    extensions = set([".bin", ".ckpt"])
    for filename in filenames:
        prefix, ext = os.path.splitext(filename)
        if ext in extensions:
            pt_filename = hf_hub_download(model_id, filename=filename)
            dirname, raw_filename = os.path.split(filename)
            if raw_filename == "pytorch_model.bin":
                # XXX: This is a special case to handle `transformers` and the
                # `transformers` part of the model which is actually loaded by `transformers`.
                sf_in_repo = os.path.join(dirname, "model.safetensors")
            else:
                sf_in_repo = f"{prefix}.safetensors"
            sf_filename = os.path.join(folder, sf_in_repo)
            convert_file(pt_filename, sf_filename)
    return sf_filename


def convert(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:
    pr_title = "Adding `safetensors` variant of this model"
    info = api.model_info(model_id)

    def is_valid_filename(filename):
        return len(filename.split("/")) > 1 or filename in ["pytorch_model.bin", "diffusion_pytorch_model.bin"]
    filenames = set(s.rfilename for s in info.siblings if is_valid_filename(s.rfilename))

    print(filenames)


    folder = os.path.join("./", repo_folder_name(repo_id=model_id, repo_type="models"))
    os.makedirs(folder)
    print(folder)
    new_pr = None
    try:
        operations = None
        pr = previous_pr(api, model_id, pr_title)

        library_name = getattr(info, "library_name", None)
        if any(filename.endswith(".safetensors") for filename in filenames) and not force:
            raise AlreadyExists(f"Model {model_id} is already converted, skipping..")
        elif pr is not None and not force:
            url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
            new_pr = pr
            raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}")
        else:
            print("Convert generic")
            operations = convert_generic(model_id, folder, filenames)

    finally:
        print(folder)
    return folder






DATASET_REPO_URL = "https://huggingface.co/datasets/safetensors/conversions"
DATA_FILENAME = "data.csv"
DATA_FILE = os.path.join("data", DATA_FILENAME)

HF_TOKEN = os.environ.get("HF_TOKEN")

repo: Optional[Repository] = None
if HF_TOKEN:
    repo = Repository(local_dir="data", clone_from=DATASET_REPO_URL, token=HF_TOKEN)


def run(token: str, model_id: str) -> str:
    if token == "" or model_id == "":
        return """
        ### Invalid input 🐞

        Please fill a token and model_id.
        """
    try:
        api = HfApi(token=token)
        is_private = api.model_info(repo_id=model_id).private
        folder = convert(api=api, model_id=model_id, force=True)

        return folder

    except Exception as e:
        return f"""
        ### Error 😢😢😢

        {e}
        """

def conversion(hf_token, Model, Username, Repo_name):
  repo_id = Username + "/" + Repo_name
  folder = run(hf_token, Model)

  api = HfApi()

  api.create_repo(
      repo_id = repo_id,
      token = hf_token,
      repo_type = "model",
      exist_ok = True
  )

  api.upload_file(
    path_or_fileobj= folder + "/model.safetensors",
    path_in_repo = "model.safetensors",
    token = hf_token,
    repo_id = repo_id,
    repo_type = "model",
  )

  shutil.rmtree(folder)
  return f"Successfully converted to safeTensors and pushed to huggingface hub {repo_id}"


inputs = [gr.Textbox(label="hf_token", elem_classes="inputs"),
gr.Textbox(label="Model_id_to_convert", elem_classes="inputs"),
gr.Textbox(label="hf_username", elem_classes="inputs"),
gr.Textbox(label="Repo_name", elem_classes="inputs")]


desc = "The Hugging Face Model Converter is a powerful tool designed to streamline the conversion process from PyTorch.bin format to SafeTensors." \
 "This Gradio app offers a user-friendly interface where users can effortlessly input their Hugging Face model details," \
 "including the Hugging Face token, model ID, username, and repository name. With just a click of a button, the conversion process is initiated. Afterwards, the" \
 "safeTensor file is pushed to your repository on huggingface"

demo = gr.Interface(fn=conversion,
                    inputs=inputs,
                    outputs=[gr.Textbox(label="Status")],
                    title="Hugging Face Model Converter: PyTorch.bin to SafeTensors",
                    description=desc,
                    theme=gr.Theme.from_hub('HaleyCH/HaleyCH_Theme')
                    )

demo.launch(debug=True)