CIOP: Coordinated Input-Output Perturbation regularization

modules/modelSetup/BaseChromaSetup.py

@@ -205,6 +205,15 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image,
+                    latent_noise,
+                    config,
+                    generator,
+                    rand
+                )
+
             latent_input = scaled_noisy_latent_image
 
             text_ids = torch.zeros(

modules/modelSetup/BaseErnieSetup.py

@@ -119,6 +119,15 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image,
+                    latent_noise,
+                    config,
+                    generator,
+                    rand
+                )
+
             predicted_flow = model.transformer(
                 hidden_states=scaled_noisy_latent_image.to(dtype=model.train_dtype.torch_dtype()),
                 timestep=timestep,

modules/modelSetup/BaseFlux2Setup.py

@@ -127,6 +127,16 @@ def predict(
                 timestep,
                 model.noise_scheduler.timesteps,
             )
+
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image,
+                    latent_noise,
+                    config,
+                    generator,
+                    rand
+                )
+
             latent_input = scaled_noisy_latent_image
 
             if model.transformer.config.guidance_embeds:

modules/modelSetup/BaseFluxSetup.py

@@ -258,6 +258,15 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image,
+                    latent_noise,
+                    config,
+                    generator,
+                    rand
+                )
+
             if config.model_type.has_mask_input() and config.model_type.has_conditioning_image_input():
                 latent_input = torch.concat(
                     [scaled_noisy_latent_image, scaled_latent_conditioning_image, batch['latent_mask']], 1

modules/modelSetup/BaseHiDreamSetup.py

@@ -357,6 +357,15 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image,
+                    latent_noise,
+                    config,
+                    generator,
+                    rand
+                )
+
             if config.model_type.has_mask_input() and config.model_type.has_conditioning_image_input():
                 latent_input = torch.concat(
                     [scaled_noisy_latent_image, scaled_latent_conditioning_image, batch['latent_mask']], 1

modules/modelSetup/BaseHunyuanVideoSetup.py

@@ -252,6 +252,15 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image,
+                    latent_noise,
+                    config,
+                    generator,
+                    rand
+                )
+
             latent_input = scaled_noisy_latent_image
 
             if model.transformer.config.guidance_embeds:

modules/modelSetup/BasePixArtAlphaSetup.py

@@ -202,6 +202,11 @@ def predict(
                 model.noise_scheduler.betas,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             if config.model_type.has_mask_input() and config.model_type.has_conditioning_image_input():
                 latent_input = torch.concat(
                     [scaled_noisy_latent_image, batch['latent_mask'], scaled_latent_conditioning_image], 1

modules/modelSetup/BaseQwenSetup.py

@@ -123,6 +123,11 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             latent_input = scaled_noisy_latent_image
             packed_latent_input = model.pack_latents(latent_input)
 

modules/modelSetup/BaseSanaSetup.py

@@ -212,6 +212,11 @@ def predict(
                 model.noise_scheduler.betas,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             if config.model_type.has_mask_input() and config.model_type.has_conditioning_image_input():
                 latent_input = torch.concat(
                     [scaled_noisy_latent_image, batch['latent_mask'], scaled_latent_conditioning_image], 1

modules/modelSetup/BaseStableDiffusion3Setup.py

@@ -309,6 +309,11 @@ def predict(
                 model.noise_scheduler.timesteps,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             if config.model_type.has_mask_input() and config.model_type.has_conditioning_image_input():
                 latent_input = torch.concat(
                     [scaled_noisy_latent_image, batch['latent_mask'], scaled_latent_conditioning_image], 1

modules/modelSetup/BaseStableDiffusionSetup.py

@@ -199,6 +199,11 @@ def predict(
                 model.noise_scheduler.betas,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             if config.model_type.has_mask_input() and config.model_type.has_conditioning_image_input():
                 latent_input = torch.concat(
                     [scaled_noisy_latent_image, batch['latent_mask'], scaled_latent_conditioning_image], 1

modules/modelSetup/BaseStableDiffusionXLSetup.py

@@ -250,6 +250,11 @@ def predict(
                 model.noise_scheduler.betas,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             # original size of the image
             original_height = batch['original_resolution'][0]
             original_width = batch['original_resolution'][1]

modules/modelSetup/BaseWuerstchenSetup.py

@@ -236,6 +236,11 @@ def predict(
                 self.__alpha_cumprod,
             )
 
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             text_embedding, pooled_text_text_embedding = model.encode_text(
                 train_device=self.train_device,
                 batch_size=batch['latent_image'].shape[0],

modules/modelSetup/BaseZImageSetup.py

@@ -122,6 +122,12 @@ def predict(
                 timestep,
                 model.noise_scheduler.timesteps,
             )
+
+            if not deterministic:
+                scaled_noisy_latent_image, latent_noise = self._apply_ciop(
+                    scaled_noisy_latent_image, latent_noise, config, generator, rand
+                )
+
             latent_input = scaled_noisy_latent_image.unsqueeze(2).to(dtype=model.train_dtype.torch_dtype())
             latent_input_list = list(latent_input.unbind(dim=0))
 

modules/modelSetup/mixin/ModelSetupNoiseMixin.py

@@ -1,5 +1,6 @@
 import math
 from abc import ABCMeta
+from random import Random
 
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.TimestepDistribution import TimestepDistribution
@@ -118,6 +119,46 @@ def _create_noise(
 
         return noise
 
+    def _apply_ciop(
+            self,
+            noisy_latent: Tensor,
+            target_noise: Tensor,
+            config: TrainConfig,
+            generator: Generator,
+            rand: Random
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Applies Coordinated Input-Output Perturbation (CIOP) as per Equation (6) & (7).
+        Paper: "Bidirectional Noise Injection: Enhancing Diffusion Models via Coordinated Input-Output Perturbation"
+
+        With probability p, injects independent Gaussian noise to both the noisy latent
+        and the target noise.
+        """
+        ciop_noise_weight = config.ciop_noise_weight
+        if ciop_noise_weight != 0:
+            prob_threshold = torch.rand(1, generator=generator, device=noisy_latent.device)
+
+            apply_mask = (prob_threshold < config.ciop_p).to(noisy_latent.dtype)
+
+            eps_in = torch.randn(
+                noisy_latent.shape,
+                generator=generator,
+                device=noisy_latent.device,
+                dtype=noisy_latent.dtype
+            ) * (ciop_noise_weight * apply_mask)
+
+            eps_out = torch.randn(
+                target_noise.shape,
+                generator=generator,
+                device=target_noise.device,
+                dtype=target_noise.dtype
+            ) * (ciop_noise_weight * apply_mask)
+
+            noisy_latent = noisy_latent + eps_in
+            target_noise = target_noise + eps_out
+
+        return noisy_latent, target_noise
+
     def _get_timestep_discrete(
             self,
             num_train_timesteps: int,

modules/ui/TrainingTab.py

@@ -642,60 +642,84 @@ def __create_noise_frame(self, master, row, supports_generalized_offset_noise: b
         frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew")
         frame.grid_columnconfigure(0, weight=1)
 
+        row = 0
+
         # offset noise weight
-        components.label(frame, 0, 0, "Offset Noise Weight",
+        components.label(frame, row, 0, "Offset Noise Weight",
                          tooltip="The weight of offset noise added to each training step")
-        components.entry(frame, 0, 1, self.ui_state, "offset_noise_weight")
+        components.entry(frame, row, 1, self.ui_state, "offset_noise_weight")
+        row += 1
 
         if supports_generalized_offset_noise:
             # generalized offset noise weight
-            generalised_offset_label = components.label(frame, 1, 0, "Generalized Offset Noise",
+            generalised_offset_label = components.label(frame, row, 0, "Generalized Offset Noise",
                             tooltip="Per-timestep 'brightness knob' instead of a fixed offset - steadier training, better starts, and improved very dark/bright images. Compatible with V-pred and Eps-pred. Start with 0.02 and adjust as needed.")
             generalised_offset_label.configure(wraplength=130, justify="left")
-            components.switch(frame, 1, 1, self.ui_state, "generalized_offset_noise")
+            components.switch(frame, row, 1, self.ui_state, "generalized_offset_noise")
+            row += 1
 
         # perturbation noise weight
-        components.label(frame, 2, 0, "Perturbation Noise Weight",
+        components.label(frame, row, 0, "Perturbation Noise Weight",
                          tooltip="The weight of perturbation noise added to each training step")
-        components.entry(frame, 2, 1, self.ui_state, "perturbation_noise_weight")
+        components.entry(frame, row, 1, self.ui_state, "perturbation_noise_weight")
+        row += 1
+
+        # CIOP noise weight
+        components.label(frame, row, 0, "I/O Noise Weight",
+                         tooltip="Applies Coordinated Input-Output Perturbation (CIOP) to each training step.")
+        components.entry(frame, row, 1, self.ui_state, "ciop_noise_weight")
+        row += 1
+
+        # CIOP noise probability
+        components.label(frame, row, 0, "I/O Noise Probability",
+                         tooltip="The probability of I/O perturbation noise for each training step")
+        components.entry(frame, row, 1, self.ui_state, "ciop_p")
+        row += 1
 
         # timestep distribution
-        components.label(frame, 3, 0, "Timestep Distribution",
+        components.label(frame, row, 0, "Timestep Distribution",
                          tooltip="Selects the function to sample timesteps during training",
                          wide_tooltip=True)
-        components.options_adv(frame, 3, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution",
+        components.options_adv(frame, row, 1, [str(x) for x in list(TimestepDistribution)], self.ui_state, "timestep_distribution",
                                adv_command=self.__open_timestep_distribution_window)
+        row += 1
 
         # min noising strength
-        components.label(frame, 4, 0, "Min Noising Strength",
+        components.label(frame, row, 0, "Min Noising Strength",
                          tooltip="Specifies the minimum noising strength used during training. This can help to improve composition, but prevents finer details from being trained")
-        components.entry(frame, 4, 1, self.ui_state, "min_noising_strength", required=True)
+        components.entry(frame, row, 1, self.ui_state, "min_noising_strength", required=True)
+        row += 1
 
         # max noising strength
-        components.label(frame, 5, 0, "Max Noising Strength",
+        components.label(frame, row, 0, "Max Noising Strength",
                          tooltip="Specifies the maximum noising strength used during training. This can be useful to reduce overfitting, but also reduces the impact of training samples on the overall image composition")
-        components.entry(frame, 5, 1, self.ui_state, "max_noising_strength", required=True)
+        components.entry(frame, row, 1, self.ui_state, "max_noising_strength", required=True)
+        row += 1
 
         # noising weight
-        components.label(frame, 6, 0, "Noising Weight",
+        components.label(frame, row, 0, "Noising Weight",
                          tooltip="Controls the weight parameter of the timestep distribution function. Use the preview to see more details.")
-        components.entry(frame, 6, 1, self.ui_state, "noising_weight", required=True)
+        components.entry(frame, row, 1, self.ui_state, "noising_weight", required=True)
+        row += 1
 
         # noising bias
-        components.label(frame, 7, 0, "Noising Bias",
+        components.label(frame, row, 0, "Noising Bias",
                          tooltip="Controls the bias parameter of the timestep distribution function. Use the preview to see more details.")
-        components.entry(frame, 7, 1, self.ui_state, "noising_bias", required=True)
+        components.entry(frame, row, 1, self.ui_state, "noising_bias", required=True)
+        row += 1
 
         # timestep shift
-        components.label(frame, 8, 0, "Timestep Shift",
+        components.label(frame, row, 0, "Timestep Shift",
                          tooltip="Shift the timestep distribution. Use the preview to see more details.")
-        components.entry(frame, 8, 1, self.ui_state, "timestep_shift", required=True)
+        components.entry(frame, row, 1, self.ui_state, "timestep_shift", required=True)
+        row += 1
 
         if supports_dynamic_timestep_shifting:
             # dynamic timestep shifting
-            components.label(frame, 9, 0, "Dynamic Timestep Shifting",
+            components.label(frame, row, 0, "Dynamic Timestep Shifting",
                              tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image and Flux2, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True)
-            components.switch(frame, 9, 1, self.ui_state, "dynamic_timestep_shifting")
+            components.switch(frame, row, 1, self.ui_state, "dynamic_timestep_shifting")
+            row += 1
 
 
 

modules/util/config/TrainConfig.py

@@ -435,6 +435,8 @@ class TrainConfig(BaseConfig):
     offset_noise_weight: float
     generalized_offset_noise: bool
     perturbation_noise_weight: float
+    ciop_noise_weight: float
+    ciop_p: float
     rescale_noise_scheduler_to_zero_terminal_snr: bool
     force_v_prediction: bool
     force_epsilon_prediction: bool
@@ -1014,6 +1016,8 @@ def default_values() -> 'TrainConfig':
         data.append(("offset_noise_weight", 0.0, float, False))
         data.append(("generalized_offset_noise", False, bool, False))
         data.append(("perturbation_noise_weight", 0.0, float, False))
+        data.append(("ciop_noise_weight", 0.0, float, False))
+        data.append(("ciop_p", 0.8, float, False))
         data.append(("rescale_noise_scheduler_to_zero_terminal_snr", False, bool, False))
         data.append(("force_v_prediction", False, bool, False))
         data.append(("force_epsilon_prediction", False, bool, False))