FluxGenerator.ipynb

fluxgenerator.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNbKXjaOVJMqD4iurNcW2WY",
      "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/219e65739300088ee26027dcc317036b/fluxgenerator.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y2NZKkvDE_V-"
      },
      "outputs": [],
      "source": [
        "!pip install bitsandbytes accelerate"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "os.environ[\"HF_TOKEN\"] = userdata.get('HF_TOKEN')"
      ],
      "metadata": {
        "id": "41vRhg9tGQ1j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from diffusers import FluxTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler\n",
        "from transformers import CLIPTextModel, CLIPTokenizer, BitsAndBytesConfig\n",
        "import bitsandbytes as bnb\n",
        "import torch\n",
        "import gc"
      ],
      "metadata": {
        "id": "vDB5I9gTFpZW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_id = \"black-forest-labs/FLUX.1-schnell\""
      ],
      "metadata": {
        "id": "tnvJtE3M2SDI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch_dtype = torch.float16"
      ],
      "metadata": {
        "id": "6oO_ZHxv2bVa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = CLIPTokenizer.from_pretrained(\n",
        "    model_id, subfolder=\"tokenizer\",\n",
        "    torch_dtype=torch_dtype, low_cpu_mem_usage=True\n",
        ")\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "02NsLHyj2Mkn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import bitsandbytes\n",
        "print(f\"bitsandbytes version: {bitsandbytes.__version__}\")"
      ],
      "metadata": {
        "id": "5rwE9RDT7bvp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from accelerate import Accelerator, load_checkpoint_and_dispatch\n",
        "accelerator = Accelerator(mixed_precision=\"fp16\")"
      ],
      "metadata": {
        "id": "J_gJajH87mUN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "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": "hyIhPyZA_Fbj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "text_encoder = CLIPTextModel.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"text_encoder\",\n",
        "    torch_dtype=torch_dtype,\n",
        "    low_cpu_mem_usage=True,\n",
        "    quantization_config=bnb_config\n",
        ")\n",
        "text_encoder = accelerator.prepare(text_encoder)\n",
        "text_encoder.to(\"cuda\")\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "1alNbWoK5a-g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vae = AutoencoderKL.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"vae\",\n",
        "    #torch_dtype=torch_dtype,\n",
        "    low_cpu_mem_usage=True,\n",
        ")\n",
        "vae = accelerator.prepare(vae)\n",
        "vae.to(\"cuda\")\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "k0ACYIhb5cR1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer = FluxTransformer2DModel.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"transformer\",\n",
        "    torch_dtype=torch_dtype,\n",
        "    low_cpu_mem_usage=True,\n",
        "    quantization_config=bnb_config\n",
        ")"
      ],
      "metadata": {
        "id": "_lXIYaI_5sMD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transformer = accelerator.prepare(transformer)\n",
        "transformer.to(\"cuda\")\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "fCXF1biEEXgV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import T5TokenizerFast\n",
        "tokenizer_2 = T5TokenizerFast.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"tokenizer_2\",\n",
        "    torch_dtype=torch_dtype,\n",
        "    low_cpu_mem_usage=True\n",
        ")\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "TdLmiK56Yucp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import T5EncoderModel\n",
        "text_encoder_2 = T5EncoderModel.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"text_encoder_2\",\n",
        "    torch_dtype=torch_dtype,\n",
        "    low_cpu_mem_usage=True,\n",
        "    quantization_config=bnb_config\n",
        ")\n",
        "text_encoder_2 = accelerator.prepare(text_encoder_2)\n",
        "text_encoder_2.to(\"cuda\")"
      ],
      "metadata": {
        "id": "b45zWPqRY4Jw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()"
      ],
      "metadata": {
        "id": "tRnMBuGWf0E8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"scheduler\",\n",
        "    torch_dtype=torch.float16\n",
        ")"
      ],
      "metadata": {
        "id": "2bYuO-4HntEw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from diffusers import DiffusionPipeline\n",
        "\n",
        "pipe = DiffusionPipeline.from_pretrained(\n",
        "    model_id,\n",
        "    transformer=transformer,\n",
        "    text_encoder=text_encoder,\n",
        "    text_encoder_2=text_encoder_2,\n",
        "    tokenizer=tokenizer,\n",
        "    tokenizer_2=tokenizer_2,\n",
        "    vae=vae,\n",
        "    torch_dtype=torch_dtype,\n",
        "    use_safety_checker=False,\n",
        "    low_cpu_mem_usage=True,\n",
        "    scheduler=scheduler,\n",
        ")\n",
        "pipe.to(\"cuda\")\n",
        "print(\"Loaded\")"
      ],
      "metadata": {
        "id": "Nb_kh84KUohQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"\"\"\n",
        "Photo from a table with a fluffy, black cat holding a sign text: \"I can haz cheezburger?\"\n",
        "\"\"\"\n",
        "image = pipe(\n",
        "    prompt=prompt,\n",
        "    height=768,\n",
        "    width=768,\n",
        "    num_inference_steps=4,\n",
        "    guidance_scale=3.5,\n",
        ").images[0]"
      ],
      "metadata": {
        "id": "u8odauTLVchW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image.save(\"i1.jpg\")"
      ],
      "metadata": {
        "id": "C0NI-vgRtM5_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from IPython.display import Image, display\n",
        "display(Image(\"i1.jpg\"))"
      ],
      "metadata": {
        "id": "1CIGc-b6GDc1"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}