Merge pull request #79 from okotaku/feat/debias_estimation_loss

[Feature] Support DeBias Estimation Loss

README.md

@@ -154,6 +154,9 @@ For detailed user guides and advanced guides, please refer to our [Documentation
       <td>
         <b>DeepFloyd IFs</b>
       </td>
+      <td>
+        <b>Others</b>
+      </td>
     </tr>
     <tr valign="top">
       <td>
@@ -183,6 +186,12 @@ For detailed user guides and advanced guides, please refer to our [Documentation
             <li><a href="configs/deepfloyd_if_dreambooth/README.md">DreamBooth (CVPR'2023)</a></li>
       </ul>
       </td>
+      <td>
+        <ul>
+            <li><a href="configs/min_snr_loss/README.md">Min-SNR Loss (ICCV'2023)</a></li>
+            <li><a href="configs/debias_estimation_loss/README.md">DeBias Estimation Loss (2023)</a></li>
+      </ul>
+      </td>
     </tr>
 </td>
     </tr>

configs/debias_estimation_loss/README.md

@@ -0,0 +1,40 @@
+# Debias the Training of Diffusion Models
+
+[Debias the Training of Diffusion Models](https://arxiv.org/abs/2310.08442)
+
+## Abstract
+
+Diffusion models have demonstrated compelling generation quality by optimizing the variational lower bound through a simple denoising score matching loss. In this paper, we provide theoretical evidence that the prevailing practice of using a constant loss weight strategy in diffusion models leads to biased estimation during the training phase. Simply optimizing the denoising network to predict Gaussian noise with constant weighting may hinder precise estimations of original images. To address the issue, we propose an elegant and effective weighting strategy grounded in the theoretically unbiased principle. Moreover, we conduct a comprehensive and systematic exploration to dissect the inherent bias problem deriving from constant weighting loss from the perspectives of its existence, impact and reasons. These analyses are expected to advance our understanding and demystify the inner workings of diffusion models. Through empirical evaluation, we demonstrate that our proposed debiased estimation method significantly enhances sample quality without the reliance on complex techniques, and exhibits improved efficiency compared to the baseline method both in training and sampling processes.
+
+<div align=center>
+<img src="https://github.com/okotaku/diffengine/assets/24734142/79b19ec5-d612-44b8-88d2-7d8677b80af5"/>
+</div>
+
+## Citation
+
+```
+```
+
+## Run Training
+
+Run Training
+
+```
+# single gpu
+$ mim train diffengine ${CONFIG_FILE}
+# multi gpus
+$ mim train diffengine ${CONFIG_FILE} --gpus 2 --launcher pytorch
+
+# Example.
+$ mim train diffengine configs/debias_estimation_loss/stable_diffusion_xl_pokemon_blip_debias_estimation_loss.py
+```
+
+## Inference with diffusers
+
+You can see more details on [`docs/source/run_guides/run_xl.md`](../../docs/source/run_guides/run_xl.md#inference-with-diffusers).
+
+## Results Example
+
+#### stable_diffusion_xl_pokemon_blip_debias_estimation_loss
+
+![example1](https://github.com/okotaku/diffengine/assets/24734142/e91c527b-2cae-4ef2-b774-a14de45e2143)

configs/debias_estimation_loss/stable_diffusion_xl_pokemon_blip_debias_estimation_loss.py

@@ -0,0 +1,12 @@
+_base_ = [
+    "../_base_/models/stable_diffusion_xl.py",
+    "../_base_/datasets/pokemon_blip_xl.py",
+    "../_base_/schedules/stable_diffusion_xl_50e.py",
+    "../_base_/default_runtime.py",
+]
+
+model = dict(loss=dict(type="DeBiasEstimationLoss", loss_weight=1.0))
+
+train_dataloader = dict(batch_size=1)
+
+optim_wrapper_cfg = dict(accumulative_counts=4)  # update every four times

configs/min_snr_loss/README.md

@@ -0,0 +1,44 @@
+# Efficient Diffusion Training via Min-SNR Weighting Strategy
+
+[Efficient Diffusion Training via Min-SNR Weighting Strategy](https://arxiv.org/abs/2303.09556)
+
+## Abstract
+
+Denoising diffusion models have been a mainstream approach for image generation, however, training these models often suffers from slow convergence. In this paper, we discovered that the slow convergence is partly due to conflicting optimization directions between timesteps. To address this issue, we treat the diffusion training as a multi-task learning problem, and introduce a simple yet effective approach referred to as Min-SNR-γ. This method adapts loss weights of timesteps based on clamped signal-to-noise ratios, which effectively balances the conflicts among timesteps. Our results demonstrate a significant improvement in converging speed, 3.4× faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet 256×256 benchmark using smaller architectures than that employed in previous state-of-the-art.
+
+## Citation
+
+```
+@InProceedings{Hang_2023_ICCV,
+    author    = {Hang, Tiankai and Gu, Shuyang and Li, Chen and Bao, Jianmin and Chen, Dong and Hu, Han and Geng, Xin and Guo, Baining},
+    title     = {Efficient Diffusion Training via Min-SNR Weighting Strategy},
+    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
+    month     = {October},
+    year      = {2023},
+    pages     = {7441-7451}
+}
+```
+
+## Run Training
+
+Run Training
+
+```
+# single gpu
+$ mim train diffengine ${CONFIG_FILE}
+# multi gpus
+$ mim train diffengine ${CONFIG_FILE} --gpus 2 --launcher pytorch
+
+# Example.
+$ mim train diffengine configs/min_snr_loss/stable_diffusion_v15_snr_pokemon_blip.py
+```
+
+## Inference with diffusers
+
+You can see more details on [`docs/source/run_guides/run_sd.md`](../../docs/source/run_guides/run_sd.md#inference-with-diffusers).
+
+## Results Example
+
+#### stable_diffusion_v15_snr_pokemon_blip
+
+![example4](https://github.com/okotaku/diffengine/assets/24734142/b98e887a-d3af-49bb-ad15-9e8250c09578)

configs/stable_diffusion/README.md

@@ -65,7 +65,3 @@ You can see more details on [`docs/source/run_guides/run_sd.md`](../../docs/sour
 #### stable_diffusion_v15_ema_pokemon_blip
 
 ![example3](https://github.com/okotaku/diffengine/assets/24734142/9f3181e2-d244-4aa3-b116-4a07eec1aa5a)
-
-#### stable_diffusion_v15_snr_pokemon_blip
-
-![example4](https://github.com/okotaku/diffengine/assets/24734142/b98e887a-d3af-49bb-ad15-9e8250c09578)

diffengine/models/editors/deepfloyd_if/deepfloyd_if.py

@@ -14,7 +14,6 @@
 from transformers import T5EncoderModel, T5Tokenizer
 
 from diffengine.models.archs import set_text_encoder_lora, set_unet_lora
-from diffengine.models.losses.snr_l2_loss import SNRL2Loss
 from diffengine.registry import MODELS
 
 
@@ -36,10 +35,14 @@ class DeepFloydIF(BaseModel):
             The weight of noise offset introduced in
             https://www.crosslabs.org/blog/diffusion-with-offset-noise
             Defaults to 0.
-        data_preprocessor (dict, optional): The pre-process config of
-            :class:`SDDataPreprocessor`.
         tokenizer_max_length (int): The max length of tokenizer.
             Defaults to 77.
+        prediction_type (str): The prediction_type that shall be used for
+            training. Choose between 'epsilon' or 'v_prediction' or leave
+            `None`. If left to `None` the default prediction type of the
+            scheduler: `noise_scheduler.config.prediciton_type` is chosen.
+        data_preprocessor (dict, optional): The pre-process config of
+            :class:`SDDataPreprocessor`.
         finetune_text_encoder (bool, optional): Whether to fine-tune text
             encoder. Defaults to False.
         gradient_checkpointing (bool): Whether or not to use gradient
@@ -54,8 +57,9 @@ def __init__(
         lora_config: dict | None = None,
         prior_loss_weight: float = 1.,
         noise_offset_weight: float = 0,
-        data_preprocessor: dict | nn.Module | None = None,
         tokenizer_max_length: int = 77,
+        prediction_type: str | None = None,
+        data_preprocessor: dict | nn.Module | None = None,
         *,
         finetune_text_encoder: bool = False,
         gradient_checkpointing: bool = False,
@@ -78,6 +82,8 @@ def __init__(
 
         self.enable_noise_offset = noise_offset_weight > 0
         self.noise_offset_weight = noise_offset_weight
+        assert prediction_type in [None, "epsilon", "v_prediction"]
+        self.prediction_type = prediction_type
 
         self.tokenizer = T5Tokenizer.from_pretrained(
             model, subfolder="tokenizer")
@@ -161,6 +167,10 @@ def infer(self,
             torch_dtype=(torch.float16 if self.device != torch.device("cpu")
                          else torch.float32),
         )
+        if self.prediction_type is not None:
+            # set prediction_type of scheduler if defined
+            pipeline.scheduler.register_to_config(
+                prediction_type=self.prediction_type)
         pipeline.set_progress_bar_config(disable=True)
         images = []
         for p in prompt:
@@ -198,6 +208,42 @@ def test_step(
         msg = "test_step is not implemented now, please use infer."
         raise NotImplementedError(msg)
 
+    def loss(self,
+             model_pred: torch.Tensor,
+             noise: torch.Tensor,
+             latents: torch.Tensor,
+             timesteps: torch.Tensor,
+             weight: torch.Tensor | None = None) -> dict[str, torch.Tensor]:
+        """Calculate loss."""
+        if self.prediction_type is not None:
+            # set prediction_type of scheduler if defined
+            self.scheduler.register_to_config(
+                prediction_type=self.prediction_type)
+
+        if self.scheduler.config.prediction_type == "epsilon":
+            gt = noise
+        elif self.scheduler.config.prediction_type == "v_prediction":
+            gt = self.scheduler.get_velocity(latents, noise, timesteps)
+        else:
+            msg = f"Unknown prediction type {self.scheduler.config.prediction_type}"
+            raise ValueError(msg)
+
+        loss_dict = {}
+        # calculate loss in FP32
+        if self.loss_module.use_snr:
+            loss = self.loss_module(
+                model_pred.float(),
+                gt.float(),
+                timesteps,
+                self.scheduler.alphas_cumprod,
+                self.scheduler.config.prediction_type,
+                weight=weight)
+        else:
+            loss = self.loss_module(
+                model_pred.float(), gt.float(), weight=weight)
+        loss_dict["loss"] = loss
+        return loss_dict
+
     def forward(
             self,
             inputs: torch.Tensor,
@@ -260,32 +306,11 @@ def forward(
         encoder_hidden_states = self.text_encoder(
             inputs["text"], attention_mask=inputs["attention_mask"])[0]
 
-        if self.scheduler.config.prediction_type == "epsilon":
-            gt = noise
-        elif self.scheduler.config.prediction_type == "v_prediction":
-            gt = self.scheduler.get_velocity(model_input, noise, timesteps)
-        else:
-            msg = f"Unknown prediction type {self.scheduler.config.prediction_type}"
-            raise ValueError(msg)
-
         model_pred = self.unet(
             noisy_model_input,
             timesteps,
             encoder_hidden_states=encoder_hidden_states).sample
 
         model_pred, _ = torch.chunk(model_pred, 2, dim=1)
 
-        loss_dict = {}
-        # calculate loss in FP32
-        if isinstance(self.loss_module, SNRL2Loss):
-            loss = self.loss_module(
-                model_pred.float(),
-                gt.float(),
-                timesteps,
-                self.scheduler.alphas_cumprod,
-                weight=weight)
-        else:
-            loss = self.loss_module(
-                model_pred.float(), gt.float(), weight=weight)
-        loss_dict["loss"] = loss
-        return loss_dict
+        return self.loss(model_pred, noise, model_input, timesteps, weight)

diffengine/models/editors/distill_sd/distill_sd_xl.py

@@ -6,7 +6,6 @@
 import torch
 
 from diffengine.models.editors.stable_diffusion_xl import StableDiffusionXL
-from diffengine.models.losses.snr_l2_loss import SNRL2Loss
 from diffengine.registry import MODELS
 
 
@@ -199,6 +198,11 @@ def forward(
             "text_embeds": pooled_prompt_embeds,
         }
 
+        if self.prediction_type is not None:
+            # set prediction_type of scheduler if defined
+            self.scheduler.register_to_config(
+                prediction_type=self.prediction_type)
+
         if self.scheduler.config.prediction_type == "epsilon":
             gt = noise
         elif self.scheduler.config.prediction_type == "v_prediction":
@@ -222,7 +226,7 @@ def forward(
 
         loss_dict = {}
         # calculate loss in FP32
-        if isinstance(self.loss_module, SNRL2Loss):
+        if self.loss_module.use_snr:
             loss_features = 0
             num_blocks = (
                 self.num_blocks

diffengine/models/editors/esd/esd_xl.py

@@ -5,7 +5,6 @@
 from torch import nn
 
 from diffengine.models.editors.stable_diffusion_xl import StableDiffusionXL
-from diffengine.models.losses.snr_l2_loss import SNRL2Loss
 from diffengine.registry import MODELS
 
 
@@ -30,6 +29,7 @@ def __init__(self,
                  width: int = 1024,
                  negative_guidance: float = 1.0,
                  train_method: str = "full",
+                 prediction_type: str | None = None,
                  data_preprocessor: dict | nn.Module | None = None,
                  **kwargs) -> None:
         if data_preprocessor is None:
@@ -39,6 +39,8 @@ def __init__(self,
         assert pre_compute_text_embeddings, \
             "`pre_compute_text_embeddings` should be True when training ESDXL"
         assert train_method in ["full", "xattn", "noxattn", "selfattn"]
+        assert prediction_type is None, \
+            "`prediction_type` should be None when training ESDXL"
 
         self.height = height
         self.width = width
@@ -49,6 +51,7 @@ def __init__(self,
             *args,
             finetune_text_encoder=finetune_text_encoder,
             pre_compute_text_embeddings=pre_compute_text_embeddings,
+            prediction_type=prediction_type,
             data_preprocessor=data_preprocessor,
             **kwargs)  # type: ignore[misc]
 
@@ -178,9 +181,10 @@ def forward(
         orig_model_pred.requires_grad = False
         gt = null_model_pred - self.negative_guidance * (
             orig_model_pred - null_model_pred)
-        if isinstance(self.loss_module, SNRL2Loss):
+        if self.loss_module.use_snr:
             loss = self.loss_module(model_pred.float(), gt.float(), timesteps,
-                                    self.scheduler.alphas_cumprod)
+                                    self.scheduler.alphas_cumprod,
+                                    self.scheduler.config.prediction_type)
         else:
             loss = self.loss_module(model_pred.float(), gt.float())
         loss_dict["loss"] = loss