FLUX_FP8_Demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "machine_shape": "hm",
      "private_outputs": true,
      "accelerator": "GPU"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "intro"
      },
      "source": [
        "# FLUX.1-dev-ControlNet-Union-Pro-2.0 (FP8 Quantized) Demo\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ABDALLALSWAITI/FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8/blob/main/FLUX_FP8_Demo.ipynb)\n\nThis notebook demonstrates how to use the FP8 quantized version of the FLUX.1-dev-ControlNet-Union-Pro-2.0 model. This is a direct quantization of the original model to FP8 format."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "setup"
      },
      "source": [
        "# Install dependencies\n!pip install -q diffusers transformers accelerate controlnet_aux opencv-python\n\n# Clone the repository\n!git clone https://huggingface.co/ABDALLALSWAITI/FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8\n!cp FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8/pipeline_flux_controlnet.py .\n!cp FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8/controlnet_flux.py ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "load_model"
      },
      "source": [
        "import torch\nfrom pipeline_flux_controlnet import FluxControlNetPipeline\nfrom controlnet_flux import FluxControlNetModel\n\n# Check if FP8 is supported\nfp8_supported = False\ntry:\n    test = torch.tensor([1.0], dtype=torch.float8_e4m3fn)\n    fp8_supported = True\n    print('FP8 is supported!')\nexcept:\n    print('FP8 not supported, using BF16 instead')\n\n# Load the model\ncontrolnet = FluxControlNetModel.from_pretrained(\n    'FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8',\n    torch_dtype=torch.bfloat16)\n\nprint('Model loaded successfully!')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}