[Quant]stable_diffusion_3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkxxpTPgpFvX"
      },
      "source": [
        "## Setup and Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U4PgrL5Iekus"
      },
      "outputs": [],
      "source": [
        "!pip install -Uq diffusers\n",
        "!pip install -Uq transformers\n",
        "!pip install -Uq bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sUOkjzQReoYm"
      },
      "outputs": [],
      "source": [
        "from transformers import T5EncoderModel, CLIPTextModelWithProjection\n",
        "from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig\n",
        "\n",
        "from diffusers import SD3Transformer2DModel, AutoencoderKL\n",
        "from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "notebook_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vzFt8ATxizSE"
      },
      "source": [
        "# Text Encoders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bDiDj3zji0zH"
      },
      "outputs": [],
      "source": [
        "model_id = \"stabilityai/stable-diffusion-3-medium-diffusers\"\n",
        "\n",
        "nf4_config = TransformersBitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        ")\n",
        "\n",
        "text_encoder = CLIPTextModelWithProjection.from_pretrained(\n",
        "    model_id, subfolder=\"text_encoder\",\n",
        "    quantization_config=nf4_config,\n",
        "    torch_dtype=torch.float16\n",
        ")\n",
        "\n",
        "text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(\n",
        "    model_id, subfolder=\"text_encoder_2\",\n",
        "    quantization_config=nf4_config,\n",
        "    torch_dtype=torch.float16\n",
        ")\n",
        "\n",
        "text_encoder_3 = T5EncoderModel.from_pretrained(\n",
        "    model_id, subfolder=\"text_encoder_3\",\n",
        "    quantization_config=nf4_config,\n",
        "    torch_dtype=torch.float16\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N5VaDZa1jZy7"
      },
      "outputs": [],
      "source": [
        "text_encoder.save_pretrained(\"text_encoder\")\n",
        "text_encoder_2.save_pretrained(\"text_encoder_2\")\n",
        "text_encoder_3.save_pretrained(\"text_encoder_3\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cJAEOmxGkpD8"
      },
      "source": [
        "## Diffusion Transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "stm1Vy4ceR9m"
      },
      "outputs": [],
      "source": [
        "model_id = \"stabilityai/stable-diffusion-3-medium-diffusers\"\n",
        "\n",
        "nf4_config = DiffusersBitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.float16,\n",
        ")\n",
        "\n",
        "transformer = SD3Transformer2DModel.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"transformer\",\n",
        "    quantization_config=nf4_config,\n",
        "    torch_dtype=torch.float16\n",
        ")\n",
        "vae = AutoencoderKL.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"vae\",\n",
        "    quantization_config=nf4_config,\n",
        "    torch_dtype=torch.float16\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EMbNLvZRk_Nw"
      },
      "outputs": [],
      "source": [
        "transformer.save_pretrained(\"transformer\")\n",
        "vae.save_pretrained(\"vae\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ciRiFic_pP7z"
      },
      "source": [
        "## Upload to the Hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3yzoziepow1Z"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import create_repo, upload_folder\n",
        "\n",
        "repo_id = create_repo(repo_id=\"ariG23498/sd3-medium-nf4-pkg\", private=True, repo_type=\"model\").repo_id\n",
        "\n",
        "upload_folder(repo_id=repo_id, folder_path=\"text_encoder\", path_in_repo=\"text_encoder\")\n",
        "upload_folder(repo_id=repo_id, folder_path=\"text_encoder_2\", path_in_repo=\"text_encoder_2\")\n",
        "upload_folder(repo_id=repo_id, folder_path=\"text_encoder_3\", path_in_repo=\"text_encoder_3\")\n",
        "\n",
        "upload_folder(repo_id=repo_id, folder_path=\"transformer\", path_in_repo=\"transformer\")\n",
        "upload_folder(repo_id=repo_id, folder_path=\"vae\", path_in_repo=\"vae\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}