hfimport

#!/usr/bin/env bash
# hfimport — install a local model directory into the Hugging Face hub cache.
#
# For directories that have nothing to do with git or the Hub, e.g. a model
# you created or converted on this machine:
#   hfimport ./my-qwen-finetune Qwen/Qwen3-0.6B-ft
#   hfimport ./data MyOrg/my-dataset --repo-type dataset
#
# Result lands in $HF_HUB_CACHE/<type>s--<org>--<name>/snapshots/<id>/ with
# refs/<revision> pointing at <id> — exactly what transformers/vllm/etc. read.
# <id> is a sha1 over the directory contents, so it is deterministic and 40-hex
# like a real commit. Prints the snapshot path on stdout (like `hf download`).

set -euo pipefail

# ---------------------------------------------------------------- helpers ----
die()  { printf 'hfimport: %s\n' "$*" >&2; exit 1; }
info() { [[ -n "$QUIET" ]] || printf '\033[1;34m==>\033[0m %s\n' "$*" >&2; }

usage() {
  sed -n '2,12p' "$0" | sed 's/^# \{0,1\}//'
  cat <<'EOF'

Usage: hfimport <src-dir> <org>/<name> [options]

Options:
  --repo-type TYPE     model (default) | dataset | space
  --revision NAME      ref name to write under refs/ (default: main)
  --cache-dir DIR      override cache root (default: $HF_HUB_CACHE or $HF_HOME/hub)
  --force              replace an existing snapshot for this id
  -q, --quiet          only print the final path
  -h, --help
EOF
  exit 0
}

# ------------------------------------------------------------ arg parsing ----
[[ $# -ge 1 ]] || usage

SRC=""; REPO=""; REPO_TYPE="model"; REVISION="main"; CACHE_DIR=""; FORCE=""; QUIET=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --repo-type)  REPO_TYPE="$2"; shift 2 ;;
    --revision)   REVISION="$2";  shift 2 ;;
    --cache-dir)  CACHE_DIR="$2";  shift 2 ;;
    --force)      FORCE=1;         shift ;;
    -q|--quiet)   QUIET=1;         shift ;;
    -h|--help)    usage ;;
    --*)          die "unknown option '$1'" ;;
    *)  if   [[ -z "$SRC"  ]]; then SRC="$1"
        elif [[ -z "$REPO" ]]; then REPO="$1"
        else die "unexpected argument '$1'"; fi; shift ;;
  esac
done

[[ -n "$SRC" ]]  || die "missing source directory"
[[ -d "$SRC" ]]  || die "$SRC is not a directory"
[[ -n "$REPO" ]] || die "missing repo id"
[[ "$REPO" == */* ]] || die "repo id must be 'org/name' (got '$REPO')"
[[ "$REVISION" != */* ]] || die "--revision must be a plain name (got '$REVISION')"

case "$REPO_TYPE" in
  model)   DIR_PREFIX="models--" ;;
  dataset) DIR_PREFIX="datasets--" ;;
  space)   DIR_PREFIX="spaces--" ;;
  *) die "--repo-type must be model, dataset, or space" ;;
esac

[[ -n "$CACHE_DIR" ]] || CACHE_DIR="${HF_HUB_CACHE:-${HF_HOME:-$HOME/.cache/huggingface}/hub}"
STORAGE="$CACHE_DIR/${DIR_PREFIX}${REPO//\//--}"

# ---------------------------------------------------------- snapshot id ------
# sha1 over "<sha1 of file>  <relative path>" lines for every regular file, in
# sorted path order. Deterministic across runs and machines; any content or
# path change yields a new id. Uses NUL separators so odd filenames are safe.
snapshot_id() {  # snapshot_id <dir>
  ( cd "$1" && find . -type f -print0 | LC_ALL=C sort -z | xargs -0 shasum -a 1 ) \
    | shasum -a 1 | cut -c1-40
}

# -------------------------------------------------------------- import -------
command -v shasum >/dev/null || die "shasum not found"

info "Hashing $SRC"
ID=$(snapshot_id "$SRC")
SNAP="$STORAGE/snapshots/$ID"
mkdir -p "$STORAGE/refs" "$STORAGE/snapshots"

# Fast path: identical content already installed → only (re)point the ref.
if [[ -z "$FORCE" && -d "$SNAP" ]]; then
  info "Already cached: $SNAP"
  printf '%s' "$ID" > "$STORAGE/refs/$REVISION"
  printf '%s\n' "$SNAP"
  exit 0
fi

TMP=$(mktemp -d "$STORAGE/.hfimport-tmp.XXXXXX")   # same filesystem → final mv is atomic
trap 'rm -rf "$TMP"' EXIT

info "Copying $SRC → $SNAP"
cp -a "$SRC/." "$TMP/repo/"

# --force replaces the snapshot wholesale; the rename below is atomic only if
# the destination does not exist.
[[ -n "$FORCE" && -d "$SNAP" ]] && rm -rf "$SNAP"
mv "$TMP/repo" "$SNAP"

# refs/<name> → id, no trailing newline (huggingface_hub reads it raw)
printf '%s' "$ID" > "$STORAGE/refs/$REVISION"

info "Done: $REPO @ ${ID:0:12}"
printf '%s\n' "$SNAP"