theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-19 06:39:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

support sd35-lora

Browse Source
This commit is contained in:
Artiprocher
2024-11-29 11:45:40 +08:00
parent 7f2a5424d4
commit 9d09121fbc
2 changed files with 90 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
examples/train/stable_diffusion_3/train_sd3_lora.py
View File

@@ -7,7 +7,7 @@ os.environ["TOKENIZERS_PARALLELISM"] = "True"
class LightningModel(LightningModelForT2ILoRA):
def __init__(
self,
torch_dtype=torch.float16, pretrained_weights=[],
torch_dtype=torch.float16, pretrained_weights=[], preset_lora_path=None,
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", init_lora_weights="gaussian",
):
@@ -16,7 +16,12 @@ class LightningModel(LightningModelForT2ILoRA):
model_manager = ModelManager(torch_dtype=torch_dtype, device=self.device)
model_manager.load_models(pretrained_weights)
self.pipe = SD3ImagePipeline.from_model_manager(model_manager)
self.pipe.scheduler.set_timesteps(1000)
self.pipe.scheduler.set_timesteps(1000, training=True)
if preset_lora_path is not None:
preset_lora_path = preset_lora_path.split(",")
for path in preset_lora_path:
model_manager.load_lora(path)
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, init_lora_weights=init_lora_weights)
@@ -29,14 +34,26 @@ def parse_args():
type=str,
default=None,
required=True,
help="Path to pretrained model. For example, `models/stable_diffusion_3/sd3_medium_incl_clips.safetensors` or `models/stable_diffusion_3/sd3_medium_incl_clips_t5xxlfp16.safetensors`.",
help="Path to pretrained models, seperated by comma. For example, SD3: `models/stable_diffusion_3/sd3_medium_incl_clips_t5xxlfp16.safetensors`, SD3.5-large: `models/stable_diffusion_3/text_encoders/clip_g.safetensors,models/stable_diffusion_3/text_encoders/clip_l.safetensors,models/stable_diffusion_3/text_encoders/t5xxl_fp16.safetensors,models/stable_diffusion_3/sd3.5_large.safetensors`",
)
parser.add_argument(
"--lora_target_modules",
type=str,
default="a_to_qkv,b_to_qkv",
default="a_to_qkv,b_to_qkv,norm_1_a.linear,norm_1_b.linear,a_to_out,b_to_out,ff_a.0,ff_a.2,ff_b.0,ff_b.2",
help="Layers with LoRA modules.",
)
parser.add_argument(
"--preset_lora_path",
type=str,
default=None,
help="Preset LoRA path.",
)
parser.add_argument(
"--num_timesteps",
type=int,
default=1000,
help="Number of total timesteps. For turbo models, please set this parameter to the number of expected number of inference steps.",
)
parser = add_general_parsers(parser)
args = parser.parse_args()
return args
@@ -46,7 +63,8 @@ if __name__ == '__main__':
args = parse_args()
model = LightningModel(
torch_dtype=torch.float32 if args.precision == "32" else torch.float16,
pretrained_weights=[args.pretrained_path],
pretrained_weights=args.pretrained_path.split(","),
preset_lora_path=args.preset_lora_path,
learning_rate=args.learning_rate,
use_gradient_checkpointing=args.use_gradient_checkpointing,
lora_rank=args.lora_rank,