Skip to main content

FLUX.1 TPU

FLUX.1 is the most advanced open-source image generation model, developed by Black Forest Labs, the creators of Stable Diffusion. It delivers fast, high-quality outputs based on user prompts, revolutionizing the way digital images are generated. One of FLUX.1’s main advantages is its ability to accurately match user prompts, whether they are simple or complex. This makes it a versatile tool suitable for a wide range of applications, from artistic projects to commercial uses. The model's architecture incorporates advanced technologies such as rectified transformers and parallel attention layers, allowing for the generation of efficient, realistic, and highly detailed images.

FLUX.1 comes in three versions—FLUX.1[Schnell], FLUX.1[Dev], and FLUX.1[Pro]—each designed to meet different needs. FLUX.1[Schnell] prioritizes speed, generating images up to 10 times faster with moderate quality. FLUX.1[Dev] offers advanced features such as image-to-image generation for developers, making it ideal for technical and creative projects. FLUX.1[Pro] is the most powerful version, with 12 billion parameters, providing top-tier image quality via an API. For more technical details on FLUX.1, please refer to the official website.

FLUX.1_TPU ports the Schnell and Dev versions to the SG2300X chip series for local TPU hardware-accelerated inference through the Sophon SDK.

TPU Setup​

Recommended TPU memory settings: NPU->7615MB, VPU->2360MB, VPP->2360MB. How to modify?

Application Deployment​

  • Clone the repository

    git clone https://github.com/zifeng-radxa/flux.1_tpu.git

  • Use git LFS to download the FLUX.1 bmodel from ModelScope

    cd flux.1_tpu
    git clone https://www.modelscope.cn/tpu-mlir/FLUX.1_TPU.git
    mv FLUX.1_TPU/ models/

  • Install dependencies

    It is recommended to create a virtual environment to avoid affecting other applications. For instructions on using virtual environments, refer to here

    pip3 install -r requirements.txt

  • Start the program

    • (Recommend) gradio

      python3 gr.py

      Access the Airbox IP address on port 8999 via a browser

    • (Optional) CLI

      # schnell
      python3 flux_pipeline.py --models schnell
      # dev
      python3 flux_pipeline.py -m dev -s 28

      Parameter explanation:

      -m, --models, model selection, required, choose from [schnell, dev]

      -s, --steps, inference steps, optional, default is 4, dev recommends 28

      -g, --cfg, CFG value, optional, default is 0.0

      -r, --seed, random seed, optional, default is None

Application Showcase​

Text-to-Image​

Prompt: high detailed texture, photograph, realistic, RAW photo of young lady, emo, dark hair, makeup, depth of field, soft split ambient window light, soft realistic shadows, perfect composition

flux_1_3

Model Conversion​

TODO