mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-22 16:50:47 +00:00
support flux training
This commit is contained in:
Artiprocher
@@ -8,10 +8,10 @@ We have implemented a training framework for text-to-image Diffusion models, ena
|
||||
|
||||
Image Examples of fine-tuned LoRA. The prompt is "一只小狗蹦蹦跳跳，周围是姹紫嫣红的鲜花，远处是山脉" (for Chinese models) or "a dog is jumping, flowers around the dog, the background is mountains and clouds" (for English models).
|
||||
|
||||
||Kolors|Stable Diffusion 3|Hunyuan-DiT|
|
||||
|-|-|-|-|
|
||||
|Without LoRA||||
|
||||
|With LoRA||||
|
||||
||FLUX.1-dev|Kolors|Stable Diffusion 3|Hunyuan-DiT|
|
||||
|-|-|-|-|-|
|
||||
|Without LoRA|||||
|
||||
|With LoRA|||||
|
||||
|
||||
## Install additional packages
|
||||
|
||||
@@ -99,6 +99,78 @@ General options:
|
||||
Access key on ModelScope (https://www.modelscope.cn/). Required if you want to upload the model to ModelScope.
|
||||
```
|
||||
|
||||
### FLUX
|
||||
|
||||
The following files will be used for constructing FLUX. You can download them from [huggingface](https://huggingface.co/black-forest-labs/FLUX.1-dev) or [modelscope](https://www.modelscope.cn/models/ai-modelscope/flux.1-dev). You can use the following code to download these files:
|
||||
|
||||
```python
|
||||
from diffsynth import download_models
|
||||
|
||||
download_models(["FLUX.1-dev"])
|
||||
```
|
||||
|
||||
```
|
||||
models/FLUX/
|
||||
└── FLUX.1-dev
|
||||
├── ae.safetensors
|
||||
├── flux1-dev.safetensors
|
||||
├── text_encoder
|
||||
│ └── model.safetensors
|
||||
└── text_encoder_2
|
||||
├── config.json
|
||||
├── model-00001-of-00002.safetensors
|
||||
├── model-00002-of-00002.safetensors
|
||||
└── model.safetensors.index.json
|
||||
```
|
||||
|
||||
Launch the training task using the following command:
|
||||
|
||||
```
|
||||
CUDA_VISIBLE_DEVICES="0" python examples/train/flux/train_flux_lora.py \
|
||||
--pretrained_text_encoder_path models/FLUX/FLUX.1-dev/text_encoder/model.safetensors \
|
||||
--pretrained_text_encoder_2_path models/FLUX/FLUX.1-dev/text_encoder_2 \
|
||||
--pretrained_dit_path models/FLUX/FLUX.1-dev/flux1-dev.safetensors \
|
||||
--pretrained_vae_path models/FLUX/FLUX.1-dev/ae.safetensors \
|
||||
--dataset_path data/dog \
|
||||
--output_path ./models \
|
||||
--max_epochs 1 \
|
||||
--steps_per_epoch 500 \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--center_crop \
|
||||
--precision "bf16" \
|
||||
--learning_rate 1e-4 \
|
||||
--lora_rank 4 \
|
||||
--lora_alpha 4 \
|
||||
--use_gradient_checkpointing
|
||||
```
|
||||
|
||||
For more information about the parameters, please use `python examples/train/flux/train_flux_lora.py -h` to see the details.
|
||||
|
||||
After training, use `model_manager.load_lora` to load the LoRA for inference.
|
||||
|
||||
```python
|
||||
from diffsynth import ModelManager, FluxImagePipeline
|
||||
import torch
|
||||
|
||||
model_manager = ModelManager(torch_dtype=torch.float16, device="cuda",
|
||||
file_path_list=[
|
||||
"models/FLUX/FLUX.1-dev/text_encoder/model.safetensors",
|
||||
"models/FLUX/FLUX.1-dev/text_encoder_2",
|
||||
"models/FLUX/FLUX.1-dev/ae.safetensors",
|
||||
"models/FLUX/FLUX.1-dev/flux1-dev.safetensors"
|
||||
])
|
||||
model_manager.load_lora("models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt", lora_alpha=1.0)
|
||||
pipe = SDXLImagePipeline.from_model_manager(model_manager)
|
||||
|
||||
torch.manual_seed(0)
|
||||
image = pipe(
|
||||
prompt=prompt,
|
||||
num_inference_steps=30, embedded_guidance=3.5
|
||||
)
|
||||
image.save("image_with_lora.jpg")
|
||||
```
|
||||
|
||||
### Kolors
|
||||
|
||||
The following files will be used for constructing Kolors. You can download Kolors from [huggingface](https://huggingface.co/Kwai-Kolors/Kolors) or [modelscope](https://modelscope.cn/models/Kwai-Kolors/Kolors). Due to precision overflow issues, we need to download an additional VAE model (from [huggingface](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) or [modelscope](https://modelscope.cn/models/AI-ModelScope/sdxl-vae-fp16-fix)). You can use the following code to download these files:
|
||||
|
||||
77
examples/train/flux/train_flux_lora.py
Normal file
77
examples/train/flux/train_flux_lora.py
Normal file
@@ -0,0 +1,77 @@
|
||||
from diffsynth import ModelManager, FluxImagePipeline
|
||||
from diffsynth.trainers.text_to_image import LightningModelForT2ILoRA, add_general_parsers, launch_training_task
|
||||
import torch, os, argparse
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "True"
|
||||
|
||||
|
||||
class LightningModel(LightningModelForT2ILoRA):
|
||||
def __init__(
|
||||
self,
|
||||
torch_dtype=torch.float16, pretrained_weights=[],
|
||||
learning_rate=1e-4, use_gradient_checkpointing=True,
|
||||
lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out"
|
||||
):
|
||||
super().__init__(learning_rate=learning_rate, use_gradient_checkpointing=use_gradient_checkpointing)
|
||||
# Load models
|
||||
model_manager = ModelManager(torch_dtype=torch_dtype, device=self.device)
|
||||
model_manager.load_models(pretrained_weights)
|
||||
self.pipe = FluxImagePipeline.from_model_manager(model_manager)
|
||||
self.pipe.scheduler.set_timesteps(1000)
|
||||
|
||||
self.freeze_parameters()
|
||||
self.add_lora_to_model(self.pipe.denoising_model(), lora_rank=lora_rank, lora_alpha=lora_alpha, lora_target_modules=lora_target_modules)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
||||
parser.add_argument(
|
||||
"--pretrained_text_encoder_path",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
help="Path to pretrained text encoder model. For example, `models/FLUX/FLUX.1-dev/text_encoder/model.safetensors`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pretrained_text_encoder_2_path",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
help="Path to pretrained t5 text encoder model. For example, `models/FLUX/FLUX.1-dev/text_encoder_2`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pretrained_dit_path",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
help="Path to pretrained dit model. For example, `models/FLUX/FLUX.1-dev/flux1-dev.safetensors`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pretrained_vae_path",
|
||||
type=str,
|
||||
default=None,
|
||||
required=True,
|
||||
help="Path to pretrained vae model. For example, `models/FLUX/FLUX.1-dev/ae.safetensors`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora_target_modules",
|
||||
type=str,
|
||||
default="a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp",
|
||||
help="Layers with LoRA modules.",
|
||||
)
|
||||
parser = add_general_parsers(parser)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
model = LightningModel(
|
||||
torch_dtype={"32": torch.float32, "bf16": torch.bfloat16}.get(args.precision, torch.float16),
|
||||
pretrained_weights=[args.pretrained_text_encoder_path, args.pretrained_text_encoder_2_path, args.pretrained_dit_path, args.pretrained_vae_path],
|
||||
learning_rate=args.learning_rate,
|
||||
use_gradient_checkpointing=args.use_gradient_checkpointing,
|
||||
lora_rank=args.lora_rank,
|
||||
lora_alpha=args.lora_alpha,
|
||||
lora_target_modules=args.lora_target_modules
|
||||
)
|
||||
launch_training_task(model, args)
|
||||
Reference in New Issue
Block a user