Flux lora update (#237)

* update flux lora

---------

Co-authored-by: tc2000731 <tc2000731@163.com>

diffsynth/schedulers/continuous_ode.py

@@ -10,7 +10,7 @@ class ContinuousODEScheduler():
         self.set_timesteps(num_inference_steps)
 
 
-    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0):
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, **kwargs):
         ramp = torch.linspace(1-denoising_strength, 1, num_inference_steps)
         min_inv_rho = torch.pow(torch.tensor((self.sigma_min,)), (1 / self.rho))
         max_inv_rho = torch.pow(torch.tensor((self.sigma_max,)), (1 / self.rho))

diffsynth/schedulers/ddim.py

@@ -38,7 +38,7 @@ class EnhancedDDIMScheduler():
         return alphas_bar
 
 
-    def set_timesteps(self, num_inference_steps, denoising_strength=1.0):
+    def set_timesteps(self, num_inference_steps, denoising_strength=1.0, **kwargs):
         # The timesteps are aligned to 999...0, which is different from other implementations,
         # but I think this implementation is more reasonable in theory.
         max_timestep = max(round(self.num_train_timesteps * denoising_strength) - 1, 0)
@@ -99,3 +99,7 @@ class EnhancedDDIMScheduler():
             sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[int(timestep.flatten().tolist()[0])])
             target = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
             return target
+        
+    
+    def training_weight(self, timestep):
+        return 1.0

diffsynth/schedulers/flow_match.py

@@ -12,11 +12,21 @@ class FlowMatchScheduler():
         self.set_timesteps(num_inference_steps)
 
 
-    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0):
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
         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 = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
-        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            self.timesteps = torch.linspace(1000, 0, num_inference_steps)
+            
+            # 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_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
+            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):
@@ -49,3 +59,9 @@ class FlowMatchScheduler():
     def training_target(self, sample, noise, timestep):
         target = noise - sample
         return target
+    
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights

diffsynth/trainers/text_to_image.py

@@ -32,12 +32,15 @@ class LightningModelForT2ILoRA(pl.LightningModule):
         self.pipe.denoising_model().train()
 
     
-    def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out"):
+    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
+        if init_lora_weights == "kaiming":
+            init_lora_weights = True
+            
         lora_config = LoraConfig(
             r=lora_rank,
             lora_alpha=lora_alpha,
-            init_lora_weights="gaussian",
+            init_lora_weights=init_lora_weights,
             target_modules=lora_target_modules.split(","),
         )
         model = inject_adapter_in_model(lora_config, model)
@@ -67,7 +70,8 @@ class LightningModelForT2ILoRA(pl.LightningModule):
             noisy_latents, timestep=timestep, **prompt_emb, **extra_input,
             use_gradient_checkpointing=self.use_gradient_checkpointing
         )
-        loss = torch.nn.functional.mse_loss(noise_pred, training_target)
+        loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
+        loss = loss * self.pipe.scheduler.training_weight(timestep)
 
         # Record log
         self.log("train_loss", loss, prog_bar=True)
@@ -179,6 +183,13 @@ def add_general_parsers(parser):
         default=4.0,
         help="The weight of the LoRA update matrices.",
     )
+    parser.add_argument(
+        "--init_lora_weights",
+        type=str,
+        default="kaiming",
+        choices=["gaussian", "kaiming"],
+        help="The initializing method of LoRA weight.",
+    )
     parser.add_argument(
         "--use_gradient_checkpointing",
         default=False,