526 lines
20 KiB
Markdown
Executable File
526 lines
20 KiB
Markdown
Executable File
# FLUX.1 Full Parameter Training Guide
|
||
|
||
This document provides a complete workflow for full parameter training of FLUX.1 Diffusion Transformer, including environment configuration, data preparation, distributed training, and inference testing.
|
||
|
||
---
|
||
|
||
## Table of Contents
|
||
- [1. Environment Configuration](#1-environment-configuration)
|
||
- [2. Data Preparation](#2-data-preparation)
|
||
- [2.1 Quick Test Dataset](#21-quick-test-dataset)
|
||
- [2.2 Dataset Structure](#22-dataset-structure)
|
||
- [2.3 metadata.json Format](#23-metadatajson-format)
|
||
- [2.4 Relative vs Absolute Path Usage](#24-relative-vs-absolute-path-usage)
|
||
- [3. Full Parameter Training](#3-full-parameter-training)
|
||
- [3.1 Download Pretrained Model](#31-download-pretrained-model)
|
||
- [3.2 Quick Start (DeepSpeed-Zero-2)](#32-quick-start-deepspeed-zero-2)
|
||
- [3.3 Common Training Parameters](#33-common-training-parameters)
|
||
- [3.4 Training Validation](#34-training-validation)
|
||
- [3.5 Training with FSDP](#35-training-with-fsdp)
|
||
- [3.6 Other Backends](#36-other-backends)
|
||
- [3.7 Multi-Machine Distributed Training](#37-multi-machine-distributed-training)
|
||
- [4. Inference Testing](#4-inference-testing)
|
||
- [4.1 Inference Parameters](#41-inference-parameters)
|
||
- [4.2 Single GPU Inference](#42-single-gpu-inference)
|
||
- [4.3 Multi-GPU Parallel Inference](#43-multi-gpu-parallel-inference)
|
||
- [5. Additional Resources](#5-additional-resources)
|
||
|
||
---
|
||
|
||
## 1. Environment Configuration
|
||
|
||
**Method 1: Using requirements.txt**
|
||
|
||
```bash
|
||
pip install -r requirements.txt
|
||
```
|
||
|
||
**Method 2: Manual Dependency Installation**
|
||
|
||
```bash
|
||
pip install Pillow einops safetensors timm tomesd librosa "torch>=2.1.2" torchdiffeq torchsde decord datasets numpy scikit-image
|
||
pip install omegaconf SentencePiece imageio[ffmpeg] imageio[pyav] tensorboard beautifulsoup4 ftfy func_timeout onnxruntime
|
||
pip install "peft>=0.17.0" "accelerate>=0.25.0" "gradio>=3.41.2" "diffusers>=0.30.1" "transformers>=4.46.2"
|
||
pip install yunchang xfuser modelscope openpyxl
|
||
pip uninstall opencv-python opencv-contrib-python opencv-python-headless -y
|
||
pip install opencv-python-headless
|
||
pip install deepspeed==0.17.0 numpy==1.26.4
|
||
```
|
||
|
||
**Method 3: Using Docker**
|
||
|
||
When using Docker, please ensure that the GPU driver and CUDA environment are correctly installed on your machine, then execute the following commands:
|
||
|
||
```
|
||
# pull image
|
||
docker pull mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||
|
||
# enter image
|
||
docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unconfined --shm-size 200g mybigpai-public-registry.cn-beijing.cr.aliyuncs.com/easycv/torch_cuda:cogvideox_fun
|
||
```
|
||
|
||
---
|
||
|
||
## 2. Data Preparation
|
||
|
||
### 2.1 Quick Test Dataset
|
||
|
||
We provide a test dataset containing several training samples.
|
||
|
||
```bash
|
||
# Download official demo dataset
|
||
modelscope download --dataset PAI/X-Fun-Images-Demo --local_dir ./datasets/X-Fun-Images-Demo
|
||
```
|
||
|
||
### 2.2 Dataset Structure
|
||
|
||
```
|
||
📦 datasets/
|
||
├── 📂 my_dataset/
|
||
│ ├── 📂 train/
|
||
│ │ ├── 📄 image001.jpg
|
||
│ │ ├── 📄 image002.png
|
||
│ │ └── 📄 ...
|
||
│ └── 📄 metadata.json
|
||
```
|
||
|
||
### 2.3 metadata.json Format
|
||
|
||
**Relative Path Format** (example):
|
||
```json
|
||
[
|
||
{
|
||
"file_path": "train/image001.jpg",
|
||
"text": "A beautiful sunset over the ocean, golden hour lighting",
|
||
"width": 1024,
|
||
"height": 1024
|
||
},
|
||
{
|
||
"file_path": "train/image002.png",
|
||
"text": "Portrait of a young woman, studio lighting, high quality",
|
||
"width": 1328,
|
||
"height": 1328
|
||
}
|
||
]
|
||
```
|
||
|
||
**Absolute Path Format**:
|
||
```json
|
||
[
|
||
{
|
||
"file_path": "/mnt/data/images/sunset.jpg",
|
||
"text": "A beautiful sunset over the ocean",
|
||
"width": 1024,
|
||
"height": 1024
|
||
}
|
||
]
|
||
```
|
||
|
||
**Key Fields Description**:
|
||
- `file_path`: Image path (relative or absolute)
|
||
- `text`: Image description (English prompt)
|
||
- `width` / `height`: Image dimensions (**recommended** to provide for bucket training; if not provided, they will be automatically read during training, which may slow down training when data is stored on slow systems like OSS)
|
||
- You can use `scripts/process_json_add_width_and_height.py` to add width and height fields to JSON files without these fields, supporting both images and videos
|
||
- Usage: `python scripts/process_json_add_width_and_height.py --input_file datasets/X-Fun-Images-Demo/metadata.json --output_file datasets/X-Fun-Images-Demo/metadata_add_width_height.json`
|
||
|
||
### 2.4 Relative vs Absolute Path Usage
|
||
|
||
**Relative Paths**:
|
||
|
||
If your data uses relative paths, configure the training script as follows:
|
||
|
||
```bash
|
||
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
|
||
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
|
||
```
|
||
|
||
**Absolute Paths**:
|
||
|
||
If your data uses absolute paths, configure the training script as follows:
|
||
|
||
```bash
|
||
export DATASET_NAME=""
|
||
export DATASET_META_NAME="/mnt/data/metadata_add_width_height.json"
|
||
```
|
||
|
||
> 💡 **Recommendation**: If the dataset is small and stored locally, use relative paths. If the dataset is stored on external storage (e.g., NAS, OSS) or shared across multiple machines, use absolute paths.
|
||
|
||
---
|
||
|
||
## 3. Full Parameter Training
|
||
|
||
### 3.1 Download Pretrained Model
|
||
|
||
```bash
|
||
# Create model directory
|
||
mkdir -p models/Diffusion_Transformer
|
||
|
||
# Download FLUX.1 official weights
|
||
modelscope download --model black-forest-labs/FLUX.1-dev --local_dir models/Diffusion_Transformer/FLUX.1-dev
|
||
```
|
||
|
||
### 3.2 Quick Start (DeepSpeed-Zero-2)
|
||
|
||
If you have downloaded the data as per **2.1 Quick Test Dataset** and the weights as per **3.1 Download Pretrained Model**, you can directly copy and run the quick start command.
|
||
|
||
DeepSpeed-Zero-2 and FSDP are recommended for training. Here we use DeepSpeed-Zero-2 as an example.
|
||
|
||
The difference between DeepSpeed-Zero-2 and FSDP lies in whether the model weights are sharded. **If VRAM is insufficient when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP.
|
||
|
||
```bash
|
||
export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev"
|
||
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
|
||
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
|
||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||
# export NCCL_IB_DISABLE=1
|
||
# export NCCL_P2P_DISABLE=1
|
||
NCCL_DEBUG=INFO
|
||
|
||
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/flux/train.py \
|
||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||
--train_data_dir=$DATASET_NAME \
|
||
--train_data_meta=$DATASET_META_NAME \
|
||
--train_batch_size=1 \
|
||
--image_sample_size=1024 \
|
||
--gradient_accumulation_steps=1 \
|
||
--dataloader_num_workers=8 \
|
||
--num_train_epochs=100 \
|
||
--checkpointing_steps=50 \
|
||
--learning_rate=2e-05 \
|
||
--lr_scheduler="constant_with_warmup" \
|
||
--lr_warmup_steps=100 \
|
||
--seed=42 \
|
||
--output_dir="output_dir_flux" \
|
||
--gradient_checkpointing \
|
||
--mixed_precision="bf16" \
|
||
--adam_weight_decay=3e-2 \
|
||
--adam_epsilon=1e-10 \
|
||
--vae_mini_batch=1 \
|
||
--max_grad_norm=0.05 \
|
||
--enable_bucket \
|
||
--uniform_sampling \
|
||
--trainable_modules "."
|
||
```
|
||
|
||
### 3.3 Common Training Parameters
|
||
|
||
**Key Parameters Description**:
|
||
|
||
| Parameter | Description | Example Value |
|
||
|-----|------|-------|
|
||
| `--pretrained_model_name_or_path` | Pretrained model path | `models/Diffusion_Transformer/FLUX.1-dev` |
|
||
| `--train_data_dir` | Training data directory | `datasets/X-Fun-Images-Demo/` |
|
||
| `--train_data_meta` | Training data metadata file | `datasets/X-Fun-Images-Demo/metadata_add_width_height.json` |
|
||
| `--train_batch_size` | Batch size per device | 1 |
|
||
| `--image_sample_size` | Maximum training resolution (auto bucketing) | 1024 |
|
||
| `--gradient_accumulation_steps` | Gradient accumulation steps (effective batch size increase) | 1 |
|
||
| `--dataloader_num_workers` | DataLoader subprocess count | 8 |
|
||
| `--num_train_epochs` | Number of training epochs | 100 |
|
||
| `--checkpointing_steps` | Save checkpoint every N steps | 50 |
|
||
| `--learning_rate` | Initial learning rate | 2e-05 |
|
||
| `--lr_scheduler` | Learning rate scheduler | `constant_with_warmup` |
|
||
| `--lr_warmup_steps` | Learning rate warmup steps | 100 |
|
||
| `--seed` | Random seed | 42 |
|
||
| `--output_dir` | Output directory | `output_dir_flux` |
|
||
| `--gradient_checkpointing` | Enable gradient checkpointing | - |
|
||
| `--mixed_precision` | Mixed precision: `fp16/bf16` | `bf16` |
|
||
| `--adam_weight_decay` | AdamW weight decay | 3e-2 |
|
||
| `--adam_epsilon` | AdamW epsilon value | 1e-10 |
|
||
| `--vae_mini_batch` | Mini batch size for VAE encoding | 1 |
|
||
| `--max_grad_norm` | Gradient clipping threshold | 0.05 |
|
||
| `--enable_bucket` | Enable bucket training (no center crop, train full images grouped by resolution) | - |
|
||
| `--random_hw_adapt` | Auto-scale images to random sizes in `[512, image_sample_size]` range | - |
|
||
| `--resume_from_checkpoint` | Resume training path, use `"latest"` to auto-select latest checkpoint | None |
|
||
| `--uniform_sampling` | Uniform timestep sampling | - |
|
||
| `--trainable_modules` | Trainable modules (`"."` means all modules) | `"."` |
|
||
| `--validation_steps` | Execute validation every N steps | 100 |
|
||
| `--validation_epochs` | Execute validation every N epochs | 100 |
|
||
| `--validation_prompts` | Prompts used during validation | `"1girl, black_hair, ..."` |
|
||
|
||
### 3.4 Training Validation
|
||
|
||
You can configure validation parameters to periodically generate test images during training, allowing you to monitor training progress and model quality.
|
||
|
||
**Validation Parameters**:
|
||
|
||
| Parameter | Description | Recommended Value |
|
||
|-----------|-------------|-------------------|
|
||
| `--validation_steps` | Execute validation every N steps | 100 |
|
||
| `--validation_epochs` | Execute validation every N epochs | 100 |
|
||
| `--validation_prompts` | Prompt for validation image generation. Use multiple space-separated prompt strings | Space-separated prompt strings |
|
||
|
||
**Example**:
|
||
|
||
```bash
|
||
--validation_steps=100 \
|
||
--validation_epochs=100 \
|
||
--validation_prompts="1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||
```
|
||
|
||
**Notes**:
|
||
- Validation images will be saved to the `output_dir` directory
|
||
- For multi-prompt validation, use: `--validation_prompts "prompt1" "prompt2" "prompt3"`
|
||
|
||
### 3.5 Training with FSDP
|
||
|
||
**If VRAM is insufficient when using multiple GPUs with DeepSpeed-Zero-2**, you can switch to FSDP.
|
||
|
||
```sh
|
||
export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev"
|
||
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
|
||
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
|
||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||
# export NCCL_IB_DISABLE=1
|
||
# export NCCL_P2P_DISABLE=1
|
||
NCCL_DEBUG=INFO
|
||
|
||
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap FluxSingleTransformerBlock,FluxTransformerBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/flux/train.py \
|
||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||
--train_data_dir=$DATASET_NAME \
|
||
--train_data_meta=$DATASET_META_NAME \
|
||
--train_batch_size=1 \
|
||
--image_sample_size=1024 \
|
||
--gradient_accumulation_steps=1 \
|
||
--dataloader_num_workers=8 \
|
||
--num_train_epochs=100 \
|
||
--checkpointing_steps=50 \
|
||
--learning_rate=2e-05 \
|
||
--lr_scheduler="constant_with_warmup" \
|
||
--lr_warmup_steps=100 \
|
||
--seed=42 \
|
||
--output_dir="output_dir_flux" \
|
||
--gradient_checkpointing \
|
||
--mixed_precision="bf16" \
|
||
--adam_weight_decay=3e-2 \
|
||
--adam_epsilon=1e-10 \
|
||
--vae_mini_batch=1 \
|
||
--max_grad_norm=0.05 \
|
||
--enable_bucket \
|
||
--uniform_sampling \
|
||
--trainable_modules "."
|
||
```
|
||
|
||
### 3.6 Training Without DeepSpeed or FSDP
|
||
|
||
**This approach is not recommended as there is no memory-saving backend, which may cause insufficient VRAM.** This is provided for reference only.
|
||
|
||
```sh
|
||
export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev"
|
||
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
|
||
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
|
||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||
# export NCCL_IB_DISABLE=1
|
||
# export NCCL_P2P_DISABLE=1
|
||
NCCL_DEBUG=INFO
|
||
|
||
accelerate launch --mixed_precision="bf16" scripts/flux/train.py \
|
||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||
--train_data_dir=$DATASET_NAME \
|
||
--train_data_meta=$DATASET_META_NAME \
|
||
--train_batch_size=1 \
|
||
--image_sample_size=1024 \
|
||
--gradient_accumulation_steps=1 \
|
||
--dataloader_num_workers=8 \
|
||
--num_train_epochs=100 \
|
||
--checkpointing_steps=50 \
|
||
--learning_rate=2e-05 \
|
||
--lr_scheduler="constant_with_warmup" \
|
||
--lr_warmup_steps=100 \
|
||
--seed=42 \
|
||
--output_dir="output_dir_flux" \
|
||
--gradient_checkpointing \
|
||
--mixed_precision="bf16" \
|
||
--adam_weight_decay=3e-2 \
|
||
--adam_epsilon=1e-10 \
|
||
--vae_mini_batch=1 \
|
||
--max_grad_norm=0.05 \
|
||
--enable_bucket \
|
||
--uniform_sampling \
|
||
--trainable_modules "."
|
||
```
|
||
|
||
### 3.7 Multi-Machine Distributed Training
|
||
|
||
**Suitable for**: Large-scale datasets, faster training speed
|
||
|
||
#### 3.7.1 Environment Configuration
|
||
|
||
Assume 2 machines, each with 8 GPUs:
|
||
|
||
**Machine 0 (Master)**:
|
||
```bash
|
||
export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev"
|
||
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
|
||
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
|
||
export MASTER_ADDR="192.168.1.100" # Master machine IP
|
||
export MASTER_PORT=10086
|
||
export WORLD_SIZE=2 # Total number of machines
|
||
export NUM_PROCESS=16 # Total processes = machines × 8
|
||
export RANK=0 # Current machine rank (0 or 1)
|
||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||
# export NCCL_IB_DISABLE=1
|
||
# export NCCL_P2P_DISABLE=1
|
||
NCCL_DEBUG=INFO
|
||
|
||
accelerate launch --mixed_precision="bf16" --main_process_ip=$MASTER_ADDR --main_process_port=$MASTER_PORT --num_machines=$WORLD_SIZE --num_processes=$NUM_PROCESS --machine_rank=$RANK --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/flux/train.py \
|
||
--pretrained_model_name_or_path=$MODEL_NAME \
|
||
--train_data_dir=$DATASET_NAME \
|
||
--train_data_meta=$DATASET_META_NAME \
|
||
--train_batch_size=1 \
|
||
--image_sample_size=1024 \
|
||
--gradient_accumulation_steps=1 \
|
||
--dataloader_num_workers=8 \
|
||
--num_train_epochs=100 \
|
||
--checkpointing_steps=50 \
|
||
--learning_rate=2e-05 \
|
||
--lr_scheduler="constant_with_warmup" \
|
||
--lr_warmup_steps=100 \
|
||
--seed=42 \
|
||
--output_dir="output_dir_flux" \
|
||
--gradient_checkpointing \
|
||
--mixed_precision="bf16" \
|
||
--adam_weight_decay=3e-2 \
|
||
--adam_epsilon=1e-10 \
|
||
--vae_mini_batch=1 \
|
||
--max_grad_norm=0.05 \
|
||
--enable_bucket \
|
||
--uniform_sampling \
|
||
--trainable_modules "."
|
||
```
|
||
|
||
**Machine 1 (Worker)**:
|
||
```bash
|
||
export MODEL_NAME="models/Diffusion_Transformer/FLUX.1-dev"
|
||
export DATASET_NAME="datasets/X-Fun-Images-Demo/"
|
||
export DATASET_META_NAME="datasets/X-Fun-Images-Demo/metadata_add_width_height.json"
|
||
export MASTER_ADDR="192.168.1.100" # Same as Master
|
||
export MASTER_PORT=10086
|
||
export WORLD_SIZE=2
|
||
export NUM_PROCESS=16
|
||
export RANK=1 # Note this is 1
|
||
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
|
||
# export NCCL_IB_DISABLE=1
|
||
# export NCCL_P2P_DISABLE=1
|
||
NCCL_DEBUG=INFO
|
||
|
||
# Use the same accelerate launch command as machine 0
|
||
```
|
||
|
||
#### 3.7.2 Multi-Machine Training Notes
|
||
|
||
- **Network Requirements**:
|
||
- RDMA/InfiniBand recommended (high performance)
|
||
- Without RDMA, add environment variables:
|
||
```bash
|
||
export NCCL_IB_DISABLE=1
|
||
export NCCL_P2P_DISABLE=1
|
||
```
|
||
|
||
- **Data Synchronization**: All machines must have access to the same data paths (NFS/shared storage)
|
||
|
||
## 4. Inference Testing
|
||
|
||
### 4.1 Inference Parameters
|
||
|
||
**Key Parameters Description**:
|
||
|
||
| Parameter | Description | Example Value |
|
||
|------|------|-------|
|
||
| `GPU_memory_mode` | GPU memory management mode, see table below | `model_cpu_offload_and_qfloat8` |
|
||
| `ulysses_degree` | Head dimension parallelism degree, 1 for single GPU | 1 |
|
||
| `ring_degree` | Sequence dimension parallelism degree, 1 for single GPU | 1 |
|
||
| `fsdp_dit` | Use FSDP for Transformer in multi-GPU inference to save VRAM | `False` |
|
||
| `fsdp_text_encoder` | Use FSDP for text encoder in multi-GPU inference | `False` |
|
||
| `compile_dit` | Compile Transformer for faster inference (effective at fixed resolution) | `False` |
|
||
| `model_name` | Model path | `models/Diffusion_Transformer/FLUX.1-dev` |
|
||
| `sampler_name` | Sampler type: `Flow`, `Flow_Unipc`, `Flow_DPM++` | `Flow` |
|
||
| `transformer_path` | Path to trained Transformer weights | `None` |
|
||
| `vae_path` | Path to trained VAE weights | `None` |
|
||
| `lora_path` | LoRA weights path | `None` |
|
||
| `sample_size` | Generated image resolution `[height, width]` | `[1344, 768]` |
|
||
| `weight_dtype` | Model weight precision, use `torch.float16` for GPUs without bf16 support | `torch.bfloat16` |
|
||
| `prompt` | Positive prompt describing the content | `"1girl, black_hair..."` |
|
||
| `negative_prompt` | Negative prompt for content to avoid | `" "` |
|
||
| `guidance_scale` | Guidance strength | 1.0 |
|
||
| `seed` | Random seed for reproducibility | 43 |
|
||
| `num_inference_steps` | Inference steps | 50 |
|
||
| `lora_weight` | LoRA weight strength | 0.55 |
|
||
| `save_path` | Generated image save path | `samples/flux-t2i` |
|
||
|
||
**GPU Memory Management Modes**:
|
||
|
||
| Mode | Description | VRAM Usage |
|
||
|------|------|---------|
|
||
| `model_full_load` | Load entire model to GPU | Highest |
|
||
| `model_full_load_and_qfloat8` | Full load + FP8 quantization | High |
|
||
| `model_cpu_offload` | Offload model to CPU after use | Medium |
|
||
| `model_cpu_offload_and_qfloat8` | CPU offload + FP8 quantization | Medium-Low |
|
||
| `model_group_offload` | Layer groups switch between CPU/CUDA | Low |
|
||
| `sequential_cpu_offload` | Layer-by-layer offload (slowest) | Lowest |
|
||
|
||
### 4.2 Single GPU Inference
|
||
|
||
Run single GPU inference with the following command:
|
||
|
||
```bash
|
||
python examples/flux/predict_t2i.py
|
||
```
|
||
|
||
Edit `examples/flux/predict_t2i.py` according to your needs. For first-time inference, focus on these parameters. For other parameters, see the inference parameter section above.
|
||
|
||
```python
|
||
# Choose based on GPU VRAM
|
||
GPU_memory_mode = "model_cpu_offload_and_qfloat8"
|
||
# Based on actual model path
|
||
model_name = "models/Diffusion_Transformer/FLUX.1-dev"
|
||
# Trained weights path, e.g., "output_dir_flux/checkpoint-xxx/diffusion_pytorch_model.safetensors"
|
||
transformer_path = None
|
||
# Write based on content to generate
|
||
prompt = "1girl, black_hair, brown_eyes, earrings, freckles, grey_background, jewelry, lips, long_hair, looking_at_viewer, nose, piercing, realistic, red_lips, solo, upper_body"
|
||
# ...
|
||
```
|
||
|
||
### 4.3 Multi-GPU Parallel Inference
|
||
|
||
**Suitable for**: High-resolution generation, faster inference
|
||
|
||
#### Install Parallel Inference Dependencies
|
||
|
||
```bash
|
||
pip install xfuser==0.4.2 yunchang==0.6.2
|
||
```
|
||
|
||
#### Configure Parallel Strategy
|
||
|
||
Edit `examples/flux/predict_t2i.py`:
|
||
|
||
```python
|
||
# Ensure ulysses_degree × ring_degree = number of GPUs
|
||
# For example, using 2 GPUs:
|
||
ulysses_degree = 2 # Head dimension parallelism
|
||
ring_degree = 1 # Sequence dimension parallelism
|
||
```
|
||
|
||
**Configuration Principles**:
|
||
- `ulysses_degree` must divide the model's head count evenly.
|
||
- `ring_degree` splits on sequence dimension, affecting communication overhead. Avoid using it when head count can be divided.
|
||
|
||
**Example Configurations**:
|
||
|
||
| GPU Count | ulysses_degree | ring_degree | Description |
|
||
|---------|---------------|-------------|------|
|
||
| 1 | 1 | 1 | Single GPU |
|
||
| 4 | 4 | 1 | Head parallelism |
|
||
| 8 | 8 | 1 | Head parallelism |
|
||
| 8 | 4 | 2 | Hybrid parallelism |
|
||
|
||
#### Run Multi-GPU Inference
|
||
|
||
```bash
|
||
torchrun --nproc-per-node=2 examples/flux/predict_t2i.py
|
||
```
|
||
|
||
## 5. Additional Resources
|
||
|
||
- **Official GitHub**: https://github.com/aigc-apps/VideoX-Fun
|