FluxReduxQuantized.ipynb

fluxreduxquantized.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNXonFAYFT5RiE8pcbS2p92",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/hathibelagal-dev/12afe041d5898b93fc1026c1faf6c871/fluxreduxquantized.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install bitsandbytes accelerate"
      ],
      "metadata": {
        "id": "XC9lWimuojdR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformer-engine"
      ],
      "metadata": {
        "id": "U3SW5ywrqokx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import BitsAndBytesConfig"
      ],
      "metadata": {
        "id": "E8YifYr-op-s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from accelerate import Accelerator, load_checkpoint_and_dispatch\n",
        "accelerator = Accelerator(mixed_precision=\"bf16\")"
      ],
      "metadata": {
        "id": "KdPSv7vPo3zD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "torch_dtype = torch.bfloat16\n",
        "bnb_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_compute_dtype=torch_dtype,\n",
        "    bnb_4bit_use_double_quant=True\n",
        ")"
      ],
      "metadata": {
        "id": "d8C7C5huoq0E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from diffusers import FluxPriorReduxPipeline, FluxPipeline, FluxTransformer2DModel\n",
        "from diffusers.utils import load_image"
      ],
      "metadata": {
        "id": "9hgJjs8zndaM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = \"black-forest-labs/FLUX.1-Redux-dev\""
      ],
      "metadata": {
        "id": "ngzB9T3gne7F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(\n",
        "    model_name, torch_dtype=torch.bfloat16,\n",
        ")"
      ],
      "metadata": {
        "id": "whPYzg46ngHG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer = FluxTransformer2DModel.from_pretrained(\n",
        "    \"black-forest-labs/FLUX.1-schnell\",\n",
        "    subfolder=\"transformer\",\n",
        "    torch_dtype=torch_dtype,\n",
        "    quantization_config=bnb_config\n",
        ")"
      ],
      "metadata": {
        "id": "lnIJVB1Rs73l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sCK23VycjsA2"
      },
      "outputs": [],
      "source": [
        "pipe = FluxPipeline.from_pretrained(\n",
        "    \"black-forest-labs/FLUX.1-schnell\",\n",
        "    text_encoder=None,\n",
        "    text_encoder_2=None,\n",
        "    transformer=transformer,\n",
        "    torch_dtype=torch.bfloat16,\n",
        ")\n",
        "pipe = accelerator.prepare(pipe)\n",
        "pipe.enable_model_cpu_offload()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pipe_prior_redux = accelerator.prepare(pipe_prior_redux)\n",
        "pipe_prior_redux.enable_model_cpu_offload()"
      ],
      "metadata": {
        "id": "GFwaC7bSn71l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "import gc\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "ZBmcXTaQxN_p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image = load_image(\"0.jpg\")\n",
        "pipe_prior_output = pipe_prior_redux(image)\n",
        "torch.cuda.empty_cache()\n",
        "images = pipe(\n",
        "    guidance_scale=2.5,\n",
        "    num_inference_steps=1,\n",
        "    generator=torch.Generator(\"cpu\").manual_seed(0),\n",
        "    **pipe_prior_output,\n",
        ").images\n",
        "images[0].save(\"1.jpg\")"
      ],
      "metadata": {
        "id": "ASMpDdVWnPky"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}