FLUX.1 [schnell/dev] on SOPHON TPU

This repo contains minimal inference code to run text-to-image with Flux latent rectified flow transformers on SOPHON 2300x TPU.

FLUX.1 dev is an open-weight, guidance-distilled model for non-commercial applications. Directly distilled from FLUX.1 [pro], FLUX.1 [dev] obtains similar quality and prompt adherence capabilities, while being more efficient than a standard model of the same size.

FLUX.1 schnell model is tailored for local development and personal use. FLUX.1 schnell is openly available under an Apache2.0 license.

FLUX.1 [schnell/dev] models by Black Forest Labs: https://blackforestlabs.ai

---

Tested Devices

• radxa Fogwise® AirBox

TPU Setting

Recommend TPU Memory: NPU->7615MB, VPU->2360MB, VPP->2360MB. How to modify ?

---

Usage

• Clone this repository

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

• Download models from ModelScope via git LFS

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

• Setup environments

  pip3 install -r requirements.txt

• Compile flux

  mkdir build && cd build
  cmake .. && make -j

• Run inference

  • gradio

    python3 gr.py

    Access the Airbox IP address on port 8999 via a browser
  • or CLI

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

    flux_pipeline.py parameters

    usage: flux_pipeline.py [-h] -m {dev,schnell} [-s STEPS] [-g GFC] [-r SEED]
    
    inference full flux pipline on sophon 2300x
    
    optional arguments:
      -h, --help            show this help message and exit
      -m {dev,schnell}, --models {dev,schnell} model choices in ['dev', 'schnell']
      -s STEPS, --steps STEPS steps
      -g CFG, --cfg CFG     guidance_scale
      -r SEED, --seed SEED  random seed

---

TODO

• Fix o3 compile segmentation fault
• Add gradio demo