FLUX1-inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f22f970-f4ce-44df-8027-5fb0805c708b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip list | egrep -w \"diffusers|quanto|transformers\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd0b3d8-3260-4f41-af53-89a77bccc5b6",
   "metadata": {},
   "source": [
    "```\n",
    "diffusers                     0.30.0\n",
    "optimum-quanto                0.2.4\n",
    "sentence-transformers         2.3.1\n",
    "transformers                  4.40.2\n",
    "transformers-stream-generator 0.0.5\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a75a10eb-f541-42a8-a00d-89804e470ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from optimum.quanto import freeze, qfloat8, quantize, qint8, qint4\n",
    "\n",
    "from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL\n",
    "from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel\n",
    "from diffusers.pipelines.flux.pipeline_flux import FluxPipeline\n",
    "from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast\n",
    "\n",
    "# black-forest-labs/FLUX.1-schnell\n",
    "fluxrepo_local = \"/home/g/models/FLUX.1-schnell\"\n",
    "dtype = torch.bfloat16\n",
    "\n",
    "scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"scheduler\"\n",
    ")\n",
    "text_encoder = CLIPTextModel.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"text_encoder\", \n",
    "    torch_dtype=dtype\n",
    ")\n",
    "tokenizer = CLIPTokenizer.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"tokenizer\", \n",
    "    torch_dtype=dtype\n",
    ")\n",
    "text_encoder_2 = T5EncoderModel.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"text_encoder_2\", \n",
    "    torch_dtype=dtype\n",
    ")\n",
    "tokenizer_2 = T5TokenizerFast.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"tokenizer_2\", \n",
    "    torch_dtype=dtype\n",
    ")\n",
    "vae = AutoencoderKL.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"vae\", \n",
    "    torch_dtype=dtype\n",
    ")\n",
    "transformer = FluxTransformer2DModel.from_pretrained(\n",
    "    fluxrepo_local, \n",
    "    subfolder=\"transformer\", \n",
    "    torch_dtype=dtype\n",
    ")\n",
    "\n",
    "## OOM on a 3090 if you don't quantize these two \n",
    "print(\"quantizing transformer ..\")\n",
    "quantize(transformer.to(\"cuda\"), weights=qfloat8)\n",
    "freeze(transformer)\n",
    "\n",
    "print(\"quantizing text_encoder_2 ..\")\n",
    "quantize(text_encoder_2.to(\"cuda\"), weights=qfloat8)\n",
    "freeze(text_encoder_2)\n",
    "\n",
    "pipe = FluxPipeline(\n",
    "    scheduler=scheduler,\n",
    "    text_encoder=text_encoder,\n",
    "    tokenizer=tokenizer,\n",
    "    text_encoder_2=text_encoder_2,\n",
    "    tokenizer_2=tokenizer_2,\n",
    "    vae=vae,\n",
    "    transformer=transformer,\n",
    ")\n",
    "# .to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "486d7ada-aa33-4fad-9c0f-26610d1cac33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# just to check how much VRAM actually used: ~18GB \n",
    "# import gc\n",
    "# gc.collect()\n",
    "# torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afdd737a-41f5-4869-929d-bccc5d3097b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"photo of a cat riding a boeing 747\"\n",
    "seed = 42\n",
    "\n",
    "# SCHNELL gen. params (https://huggingface.co/black-forest-labs/FLUX.1-schnell)\n",
    "image = pipe(\n",
    "    prompt,\n",
    "    guidance_scale=0.0,\n",
    "    num_inference_steps=4,\n",
    "    max_sequence_length=256,\n",
    "    generator=torch.Generator(\"cpu\").manual_seed(seed)\n",
    ").images[0]\n",
    "image.save(\"flux-schnell.png\")\n",
    "\n",
    "# DEV gen. params (https://huggingface.co/black-forest-labs/FLUX.1-dev)\n",
    "# image = pipe(\n",
    "#     prompt,\n",
    "#     height=1024,\n",
    "#     width=1024,\n",
    "#     guidance_scale=3.5,\n",
    "#     num_inference_steps=50,\n",
    "#     max_sequence_length=512,\n",
    "#     generator=torch.Generator(\"cpu\").manual_seed(seed)\n",
    "# ).images[0]\n",
    "# image.save(\"flux-dev.png\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}