{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3eDO2Vp7y94X"
      },
      "outputs": [],
      "source": [
        "!pip install huggingface_hub hf_transfer"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!export HF_HUB_ENABLE_HF_TRANSFER=1"
      ],
      "metadata": {
        "id": "j4X2v5yh0qJC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from huggingface_hub import notebook_login, HfApi\n",
        "from safetensors.torch import load_file, save_file"
      ],
      "metadata": {
        "id": "kPX5fBF5zHYN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "notebook_login() # WRITE token"
      ],
      "metadata": {
        "id": "50ZzcvkI1F0W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "api = HfApi()"
      ],
      "metadata": {
        "id": "zTFJTeRHzJ6a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def download(repo_id: str, filename: str, save_name: str, dtype = torch.float8_e4m3fn):\n",
        "  ckpt_path = api.hf_hub_download(\n",
        "      repo_id,\n",
        "      filename=filename,\n",
        "  )\n",
        "\n",
        "  state_dict = load_file(ckpt_path)\n",
        "  for key, value in state_dict.items():\n",
        "      state_dict[key] = value.to(dtype)\n",
        "\n",
        "  save_file(state_dict, f\"./{save_name}\")\n",
        "  return f\"./{save_name}\""
      ],
      "metadata": {
        "id": "IDgia3bzzJBJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def push(repo_id: str, local_file: str):\n",
        "  api.upload_file(\n",
        "      path_or_fileobj=local_file,\n",
        "      path_in_repo=local_file.split(\"/\")[-1],\n",
        "      repo_id=repo_id,\n",
        "      repo_type=\"model\",\n",
        "  )"
      ],
      "metadata": {
        "id": "o1-kDCqY1yZY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_path = download(\n",
        "    \"OnomaAIResearch/Illustrious-xl-early-release-v0\",\n",
        "    \"Illustrious-XL-v0.1.safetensors\",\n",
        "    \"Illustrious-XL-v0.1.float8_e4m3fn.safetensors\",\n",
        "    torch.float8_e4m3fn, # float8_e5m2\n",
        ")"
      ],
      "metadata": {
        "id": "dOBdDRTpzLcV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "push(YOUR_REPO_NAME, model_path)"
      ],
      "metadata": {
        "id": "jYGONemw2Lj_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}