theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-23 17:38:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #239 from modelscope/flux-lora-update

Browse Source 
Flux lora update
This commit is contained in:
Zhongjie Duan
2024-10-14 19:12:33 +08:00
committed by GitHub
parent 22e4ae99e8 540c036988
commit 7d7d72dcfe
3 changed files with 6 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
diffsynth/schedulers/flow_match.py
View File

@@ -16,17 +16,13 @@ class FlowMatchScheduler():
sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
self.timesteps = self.sigmas * self.num_train_timesteps
if training: if training:
self.timesteps = torch.linspace(1000, 0, num_inference_steps) x = self.timesteps
# prepare timestep weights
x = torch.arange(num_inference_steps, dtype=torch.float32)
y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
y_shifted = y - y.min() y_shifted = y - y.min()
bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
self.linear_timesteps_weights = bsmntw_weighing self.linear_timesteps_weights = bsmntw_weighing
else:
self.timesteps = self.sigmas * self.num_train_timesteps
def step(self, model_output, timestep, sample, to_final=False): def step(self, model_output, timestep, sample, to_final=False):

4
diffsynth/trainers/text_to_image.py
View File

@@ -18,6 +18,7 @@ class LightningModelForT2ILoRA(pl.LightningModule):
self.learning_rate = learning_rate self.learning_rate = learning_rate
self.use_gradient_checkpointing = use_gradient_checkpointing self.use_gradient_checkpointing = use_gradient_checkpointing
self.state_dict_converter = state_dict_converter self.state_dict_converter = state_dict_converter
self.lora_alpha = None
def load_models(self): def load_models(self):
@@ -34,6 +35,7 @@ class LightningModelForT2ILoRA(pl.LightningModule):
def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out", init_lora_weights="gaussian"): def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out", init_lora_weights="gaussian"):
# Add LoRA to UNet # Add LoRA to UNet
self.lora_alpha = lora_alpha
if init_lora_weights == "kaiming": if init_lora_weights == "kaiming":
init_lora_weights = True init_lora_weights = True
@@ -94,7 +96,7 @@ class LightningModelForT2ILoRA(pl.LightningModule):
if name in trainable_param_names: if name in trainable_param_names:
lora_state_dict[name] = param lora_state_dict[name] = param
if self.state_dict_converter is not None: if self.state_dict_converter is not None:
lora_state_dict = self.state_dict_converter(lora_state_dict) lora_state_dict = self.state_dict_converter(lora_state_dict, alpha=self.lora_alpha)
checkpoint.update(lora_state_dict) checkpoint.update(lora_state_dict)

2
examples/train/README.md
View File

@@ -16,7 +16,7 @@ Image Examples of fine-tuned LoRA. The prompt is "一只小狗蹦蹦跳跳，周
## Install additional packages ## Install additional packages
``` ```
pip install peft lightning pip install peft lightning pandas
``` ```
## Prepare your dataset ## Prepare your dataset