diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 0eb623399..b4983534f 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/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( diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index c164dc0bc..c22c540b6 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/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, diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b91cd8af8..7bf97a231 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/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: diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 9bae83cde..62b2748a1 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/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 diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 17fbcc0d6..19d4ef0eb 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/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 diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index b072bf4ba..28781354a 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/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: diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 3069b4884..b06f47ae2 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/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 diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a8a7be8f6..cb5c0c37f 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/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) diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 0c9ea6da0..20701de2d 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/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 diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index de5dc04e8..170d3d1d3 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/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 diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 8cf63ac07..0c843062f 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/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 diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 61cb1e457..48bcaf171 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/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] diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index 48cd05942..54055d658 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/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], diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b822d7304..de1f39cec 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/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)) diff --git a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py b/modules/modelSetup/mixin/ModelSetupNoiseMixin.py index 1285caba8..345bb32e9 100644 --- a/modules/modelSetup/mixin/ModelSetupNoiseMixin.py +++ b/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, diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index bcca11ae9..f31c320f0 100644 --- a/modules/ui/TrainingTab.py +++ b/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 diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index 1c940da66..e0eaea7c8 100644 --- a/modules/util/config/TrainConfig.py +++ b/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))