Merge pull request bghira#2592 from bghira/sliders

Sliders

Author: bghira

simpletuner/helpers/data_backend/factory.py

@@ -3198,11 +3198,19 @@ def _create_dataset_and_sampler(
         elif caption_strategy == "instanceprompt":
             use_captions = False
 
+        slider_strength_raw = backend.get("slider_strength", 1.0)
+        try:
+            slider_strength = float(slider_strength_raw) if slider_strength_raw is not None else 1.0
+        except (TypeError, ValueError):
+            logging.warning("Invalid slider_strength %r in backend; defaulting to 1.0", slider_strength_raw)
+            slider_strength = 1.0
+
         init_backend["train_dataset"] = MultiAspectDataset(
             id=init_backend["id"],
             datasets=[init_backend["metadata_backend"]],
             is_regularisation_data=is_regularisation_data,
             is_i2v_data=is_i2v_data,
+            slider_strength=slider_strength,
         )
 
         if "deepfloyd" in self.args.model_type:

simpletuner/helpers/multiaspect/dataset.py

@@ -31,12 +31,14 @@ def __init__(
         print_names: bool = False,
         is_regularisation_data: bool = False,
         is_i2v_data: bool = False,
+        slider_strength: float = 1.0,
     ):
         self.id = id
         self.datasets = datasets
         self.print_names = print_names
         self.is_regularisation_data = is_regularisation_data
         self.is_i2v_data = is_i2v_data
+        self.slider_strength = slider_strength
 
     def __len__(self):
         # Sum the length of all data backends:
@@ -54,6 +56,7 @@ def __getitem__(self, image_tuple: list[dict[str, Any] | TrainingSample]):
             "conditioning_samples": [],
             "is_regularisation_data": self.is_regularisation_data,
             "is_i2v_data": self.is_i2v_data,
+            "slider_strength": self.slider_strength,
         }
         first_aspect_ratio = None
         for sample in image_tuple:

simpletuner/helpers/training/collate.py

@@ -1219,4 +1219,7 @@ def _conditioning_pixel_value_for_example(example_idx: int):
         "is_audio_only": is_audio_only,
         "s2v_audio_paths": s2v_audio_paths if any(s2v_audio_paths) else None,
         "s2v_audio_backend_ids": s2v_audio_backend_ids if any(s2v_audio_backend_ids) else None,
+
+
+        "slider_strength": batch.get("slider_strength")
     }

simpletuner/helpers/training/trainer.py

@@ -2892,6 +2892,20 @@ def _get_trainable_parameters(self):
                 return self.lycoris_wrapped_network.parameters()
         return [param for param in self.model.get_trained_component(unwrap_model=False).parameters() if param.requires_grad]
 
+    def _get_slider_tuner_layers(self):
+        """Return cached list of (id, module) for BaseTunerLayer modules that have scaling dicts."""
+        if hasattr(self, "_slider_tuner_layers_cache"):
+            return self._slider_tuner_layers_cache
+
+        from peft.tuners.tuners_utils import BaseTunerLayer
+
+        result = []
+        for name, module in self.model.get_trained_component().named_modules():
+            if isinstance(module, BaseTunerLayer) and hasattr(module, "scaling"):
+                result.append((id(module), module))
+        self._slider_tuner_layers_cache = result
+        return result
+
     def _ensure_parameter_dtype(self, parameters, target_dtype: torch.dtype, optimizer_name: str | None = None):
         converted = 0
         for param_or_group in parameters:
@@ -5431,6 +5445,28 @@ def train(self):
                             else:
                                 self.model.get_trained_component().enable_lora()
 
+                    # slider
+                    raw_strength = prepared_batch.get("slider_strength", 1.0)
+                    try:
+                        strength = float(raw_strength)
+                    except (TypeError, ValueError):
+                        strength = 1.0
+
+                    slider_original_scaling = None
+                    if self.config.model_type == "lora" and strength != 1.0:
+                        with torch.no_grad():
+                            if self.config.lora_type.lower() == "lycoris":
+                                self.accelerator._lycoris_wrapped_network.set_multiplier(strength)
+                            else:
+                                tuner_layers = self._get_slider_tuner_layers()
+                                slider_original_scaling = {}
+                                for layer_id, module in tuner_layers:
+                                    saved = {}
+                                    for key, val in module.scaling.items():
+                                        saved[key] = val
+                                        module.scaling[key] = val * strength
+                                    slider_original_scaling[layer_id] = (module, saved)
+
                     training_logger.debug("Predicting.")
                     model_pred = self.model_predict(
                         prepared_batch=prepared_batch,
@@ -5601,6 +5637,14 @@ def train(self):
                         ):
                             self.distiller.discriminator_step(prepared_batch=prepared_batch)
                             self.distiller.post_training_step(self.model, step)
+                    if self.config.model_type == "lora" and strength != 1:
+                        with torch.no_grad():
+                            if self.config.lora_type.lower() == "lycoris":
+                                self.accelerator._lycoris_wrapped_network.set_multiplier(1.0)
+                            elif slider_original_scaling is not None:
+                                for module, saved in slider_original_scaling.values():
+                                    for key, val in saved.items():
+                                        module.scaling[key] = val
 
                 # Checks if the accelerator has performed an optimization step behind the scenes
                 wandb_logs = {}