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
Recommend TPU Memory: NPU->7615MB, VPU->2360MB, VPP->2360MB. How to modify ?
-
Clone this repository
git clone https://github.com/zifeng-radxa/flux.1_tpu.git
-
Download models from ModelScope via GLF
cd flux.1_tpu git clone https://modelscope.cn/models/tpu-mlir/FLUX.1_TPU mv FLUX.1_TPU/ models/ -
Setup environments
pip3 install -r requirements.txt
-
Compile flux
mkdir build && cd build cmake .. && make -j
-
Run inference
# schnell python3 flux_pipeline.py --models schnell # dev python3 flux_pipeline.py -m dev -s 28
flux_pipeline.pyparametersusage: 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 GFC, --gfc GFC guidance_scale -r SEED, --seed SEED random seed
- Fix o3 compile segmentation fault
- Add gradio demo