VQ4DiT_quantization_implementation.ipynb

{
  "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": "wXmurl0wSa6E"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModel, AutoTokenizer\n",
        "from transformers import HfApi, HfFolder\n",
        "import torch\n",
        "import numpy as np\n",
        "from sklearn.cluster import KMeans\n",
        "\n",
        "# Load a Diffusion Transformer Model from Hugging Face Hub\n",
        "model_name = \"CompVis/stable-diffusion-v1-4\"  # Example model; replace with a DiT model\n",
        "model = AutoModel.from_pretrained(model_name)\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "# VQ4DiT quantization function\n",
        "def vq4dit_quantize(model):\n",
        "    for name, param in model.named_parameters():\n",
        "        if param.requires_grad and len(param.shape) > 1:  # Quantize only weight matrices\n",
        "            weight = param.data.cpu().numpy()\n",
        "            n_clusters = min(256, weight.shape[0] * weight.shape[1])  # Adjust as per paper\n",
        "\n",
        "            # K-Means Clustering to create codebook\n",
        "            kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(weight.reshape(-1, 1))\n",
        "            codebook = kmeans.cluster_centers_.flatten()\n",
        "            assignments = kmeans.labels_.reshape(weight.shape)\n",
        "\n",
        "            # Reconstruct quantized weight\n",
        "            quantized_weight = codebook[assignments]\n",
        "            param.data = torch.from_numpy(quantized_weight).to(param.device)\n",
        "\n",
        "    return model\n",
        "\n",
        "# Apply VQ4DiT quantization to the model\n",
        "quantized_model = vq4dit_quantize(model)\n",
        "\n",
        "# Save the quantized model\n",
        "quantized_model.save_pretrained(\"./quantized_model\")\n",
        "\n",
        "# Push the quantized model to Hugging Face Hub\n",
        "api = HfApi()\n",
        "username = api.whoami()[\"name\"]  # Assumes that you are already logged in with `huggingface-cli login`\n",
        "repo_id = f\"{username}/quantized-{model_name.split('/')[-1]}\"\n",
        "api.create_repo(repo_id, private=False)\n",
        "quantized_model.push_to_hub(repo_id, use_auth_token=HfFolder.get_token())\n",
        "tokenizer.push_to_hub(repo_id, use_auth_token=HfFolder.get_token())\n",
        "\n",
        "print(\"Quantized model successfully pushed to Hugging Face Hub.\")\n"
      ]
    }
  ]
}