theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-22 16:50:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

align flux lora format (#204)

Browse Source
This commit is contained in:
Zhongjie Duan
2024-09-12 16:01:27 +08:00
committed by GitHub
parent 0b9e673fa2
commit a8cb4a21d1
4 changed files with 71 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
examples/train/README.md
View File

@@ -142,9 +142,12 @@ CUDA_VISIBLE_DEVICES="0" python examples/train/flux/train_flux_lora.py \
--learning_rate 1e-4 \
--lora_rank 4 \
--lora_alpha 4 \
--use_gradient_checkpointing
--use_gradient_checkpointing \
--align_to_opensource_format
```
**`--align_to_opensource_format` means that this script will export the LoRA weights in the opensource format. This format can be loaded in both DiffSynth-Studio and other codebases.**
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.
@@ -165,7 +168,7 @@ pipe = SDXLImagePipeline.from_model_manager(model_manager)
torch.manual_seed(0)
image = pipe(
prompt=prompt,
prompt="a dog is jumping, flowers around the dog, the background is mountains and clouds",
num_inference_steps=30, embedded_guidance=3.5
)
image.save("image_with_lora.jpg")

15
examples/train/flux/train_flux_lora.py
View File

@@ -1,5 +1,6 @@
from diffsynth import ModelManager, FluxImagePipeline
from diffsynth.trainers.text_to_image import LightningModelForT2ILoRA, add_general_parsers, launch_training_task
from diffsynth.models.lora import FluxLoRAConverter
import torch, os, argparse
os.environ["TOKENIZERS_PARALLELISM"] = "True"
@@ -9,9 +10,10 @@ class LightningModel(LightningModelForT2ILoRA):
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"
lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out",
state_dict_converter=None,
):
super().__init__(learning_rate=learning_rate, use_gradient_checkpointing=use_gradient_checkpointing)
super().__init__(learning_rate=learning_rate, use_gradient_checkpointing=use_gradient_checkpointing, state_dict_converter=state_dict_converter)
# Load models
model_manager = ModelManager(torch_dtype=torch_dtype, device=self.device)
model_manager.load_models(pretrained_weights)
@@ -58,6 +60,12 @@ def parse_args():
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_argument(
"--align_to_opensource_format",
default=False,
action="store_true",
help="Whether to export lora files aligned with other opensource format.",
)
parser = add_general_parsers(parser)
args = parser.parse_args()
return args
@@ -72,6 +80,7 @@ if __name__ == '__main__':
use_gradient_checkpointing=args.use_gradient_checkpointing,
lora_rank=args.lora_rank,
lora_alpha=args.lora_alpha,
lora_target_modules=args.lora_target_modules
lora_target_modules=args.lora_target_modules,
state_dict_converter=FluxLoRAConverter().align_to_opensource_format if args.align_to_opensource_format else None,
)
launch_training_task(model, args)