[WIP] feature/PrefGRPO-txt2img-cleaning

.gitignore

@@ -27,8 +27,6 @@ share/python-wheels/
 MANIFEST
 
 # PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
 *.manifest
 *.spec
 
@@ -83,36 +81,12 @@ notebooks/
 profile_default/
 ipython_config.py
 
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
 # pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
 .pdm.toml
 .pdm-python
 .pdm-build/
 
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+# PEP 582
 __pypackages__/
 
 # Celery stuff
@@ -155,13 +129,6 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
 # running logs
 examples/wandb
 outputs/

examples/train_diffusion.py

@@ -5,7 +5,7 @@
 from torch.utils.data import DataLoader
 
 from diffulab.diffuse import Diffuser
-from diffulab.training import Trainer
+from diffulab.training import BaseTrainer
 
 
 @hydra.main(version_base=None, config_path="../configs", config_name="train_mnist_flow_matching")
@@ -53,7 +53,7 @@ def count_parameters(model: torch.nn.Module) -> int:
     )
 
     # TODO: add a run name for wandb
-    trainer = Trainer(
+    trainer = BaseTrainer(
         n_epoch=cfg.trainer.n_epoch,
         gradient_accumulation_step=cfg.trainer.gradient_accumulation_step,
         precision_type=cfg.trainer.precision_type,

examples/train_repa.py

@@ -7,7 +7,7 @@
 from torch.utils.data import DataLoader
 
 from diffulab.diffuse import Diffuser
-from diffulab.training import Trainer
+from diffulab.training import BaseTrainer
 from diffulab.training.losses.repa import RepaLoss
 
 
@@ -77,7 +77,7 @@ def count_parameters(model: torch.nn.Module) -> int:
         + list(repa_loss.resampler.parameters() if repa_loss.resampler else []),
     )
 
-    trainer = Trainer(
+    trainer = BaseTrainer(
         n_epoch=cfg.trainer.n_epoch,
         gradient_accumulation_step=cfg.trainer.gradient_accumulation_step,
         precision_type=cfg.trainer.precision_type,

pyproject.toml

@@ -1,20 +1,23 @@
 [project]
 name = "diffulab"
 version = "0.1.0"
-description = "Add your description here"
+description = "DiffuLab is designed to provide a simple and flexible way to train diffusion models while allowing full customization of its core components"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "accelerate>=1.4.0",
+    "blobfile>=3.1.0",
     "diffusers>=0.33.1",
     "einops>=0.8.1",
     "ema-pytorch>=0.7.7",
     "hydra-core>=1.3.2",
     "jaxtyping>=0.3.0",
+    "loguru>=0.7.3",
     "mosaicml-streaming>=0.12.0",
     "omegaconf>=2.3.0",
-    "open-clip-torch>=2.30.0",
     "pyopenssl==23.2.0",
+    "sentencepiece>=0.2.1",
+    "tiktoken>=0.11.0",
     "torch>=2.6.0",
     "transformers>=4.49.0",
     "wandb>=0.19.6",
@@ -31,6 +34,9 @@ dev = [
 repa = [
     "timm>=1.0.15",
 ]
+prefgrpo = [
+    "qwen-vl-utils>=0.0.11",
+]
 
 [tool.uv.sources]
 diffulab = {workspace = true}
@@ -58,7 +64,7 @@ ignore = [
 combine-as-imports = true
 
 [tool.pyright]
-include = ["src/diffulab", "tests/", "examples/"]
+include = ["src/diffulab", "tests/", "examples/", "local/"]
 strict = ["*"]
 exclude = ["**/__pycache__"]
 reportMissingTypeStubs = "warning"

src/diffulab/__init__.py

@@ -1,7 +1,7 @@
 from .datasets import BaseDataset, CIFAR10Dataset, ImageNetLatentREPA, MNISTDataset
 from .diffuse import Diffuser, Flow, GaussianDiffusion
 from .networks import DCAE, REPA, Denoiser, DinoV2, MMDiT, PerceiverResampler, SD3TextEmbedder, UNetModel, VisionTower
-from .training import LossFunction, RepaLoss, Trainer
+from .training import BaseTrainer, GRPOTrainer, LossFunction, RepaLoss, Trainer
 
 __all__ = [
     "BaseDataset",
@@ -22,5 +22,7 @@
     "VisionTower",
     "LossFunction",
     "RepaLoss",
+    "BaseTrainer",
+    "GRPOTrainer",
     "Trainer",
 ]

src/diffulab/datasets/base.py

@@ -7,12 +7,17 @@
 from torch import Tensor
 from torch.utils.data import Dataset
 
-from diffulab.networks.denoisers.common import ModelInput
+from diffulab.networks.denoisers.common import ExtraInputGRPO, ModelInput, ModelInputGRPO
 
 
 class BatchData(TypedDict, total=False):
     model_inputs: Required[ModelInput]
-    extra: NotRequired[dict[str, Tensor | None]]
+    extra: NotRequired[dict[str, Tensor | list[str] | None]]
+
+
+class BatchDataGRPO(TypedDict, total=False):
+    model_inputs: Required[ModelInputGRPO]
+    extra: Required[ExtraInputGRPO]
 
 
 class BaseDataset(Dataset[BatchData], ABC):

src/diffulab/diffuse/diffuser.py

@@ -5,6 +5,7 @@
 from diffulab.diffuse.modelizations.diffusion import Diffusion
 from diffulab.diffuse.modelizations.flow import Flow
 from diffulab.diffuse.modelizations.gaussian_diffusion import GaussianDiffusion
+from diffulab.diffuse.utils import SamplingOutput
 from diffulab.networks.denoisers.common import Denoiser, ModelInput
 from diffulab.networks.vision_towers.common import VisionTower
 from diffulab.training.losses import LossFunction
@@ -96,71 +97,83 @@ def draw_timesteps(self, batch_size: int) -> Tensor:
     def compute_loss(
         self,
         model_inputs: ModelInput,
-        timesteps: Tensor,
+        timesteps: Tensor | None = None,
         noise: Tensor | None = None,
         extra_args: dict[str, Any] = {},
+        grpo: bool = False,
+        grpo_args: dict[str, Any] = {},
     ) -> dict[str, Tensor]:
         """
         Compute the loss for the diffusion model using the denoiser and diffusion process.
         This method serves as a bridge between the Diffuser class and the underlying
         diffusion implementation by forwarding the loss computation to the diffusion model.
         Args:
             model_inputs (ModelInput): A dictionary containing the model inputs,
-              including the data tensor keyed as 'x' and any conditional information.
+              including the data tensor keyed as 'x' and any conditional information. For GRPO,
+              x is not needed in the model inputs.
             timesteps (Tensor): A tensor of timesteps for the batch.
             noise (Tensor | None, optional): Pre-defined noise to add to the input.
               If None, random noise will be generated. Defaults to None.
             extra_args (dict[str, Any], optional): Additional arguments for the additional losses computation.
         Returns:
             dict[str, Tensor]: A dictionary containing the loss value and any additional losses
         """
-        return self.diffusion.compute_loss(self.denoiser, model_inputs, timesteps, noise, self.extra_losses, extra_args)
+        if grpo:
+            assert isinstance(self.diffusion, Flow), "GRPO loss computation is only available for Flow-based models"
+            return self.diffusion.compute_loss_grpo(
+                self.denoiser,
+                model_inputs,  # type: ignore[reportArgumentType]
+                **grpo_args,
+            )
+        assert timesteps is not None, "timesteps must be provided for loss computation"
+        return self.diffusion.compute_loss(self.denoiser, model_inputs, timesteps, noise, self.extra_losses, extra_args)  # type: ignore[reportArgumentType]
 
-    def set_steps(self, n_steps: int, **extra_args: dict[str, Any]) -> None:
+    def set_steps(self, n_steps: int, schedule: str = "linear") -> None:
         """
         Update the number of diffusion steps and related parameters.
         This method allows changing the number of steps used in the diffusion process
         after the Diffuser has been initialized. It delegates to the underlying diffusion
         model's set_steps method.
         Args:
             n_steps (int): The new number of diffusion steps to use.
-            **extra_args (dict[str, Any]): Additional arguments to pass to the diffusion
-                model's set_steps method. These may include parameters like 'schedule'
-                or 'section_counts' depending on the diffusion model implementation.
+            schedule (str, optional): The schedule to use for the timesteps. Defaults to "linear".
         Example:
             ```
             diffuser = Diffuser(denoiser, sampling_method="ddpm", n_steps=1000)
             # Later, change to use fewer steps for faster sampling
             diffuser.set_steps(100, schedule="ddim")
             ```
         """
-        self.diffusion.set_steps(n_steps, **extra_args)  # type: ignore
+        self.diffusion.set_steps(n_steps, schedule=schedule)
 
     def generate(
         self,
-        data_shape: tuple[int, ...],
         model_inputs: ModelInput,
+        data_shape: tuple[int, ...] | None = None,
         use_tqdm: bool = True,
         clamp_x: bool = False,
         guidance_scale: float = 0,
-        **kwargs: dict[str, Any],
-    ) -> Tensor:
+        sampler_args: dict[str, Any] = {},
+        return_intermediates: bool = False,
+        return_latents: bool = False,
+    ) -> SamplingOutput:
         """
         Generates a new sample using the diffusion model.
         This method delegates to the underlying diffusion model's denoise method to
         generate a sample from the diffusion process. It can handle conditional generation
         with guidance when appropriate settings are provided.
         Args:
-            data_shape (tuple[int, ...]): Shape of the data to generate, typically (batch_size, channels, height, width).
-                If a vision tower is used, this should be the shape of the latent space.
             model_inputs (ModelInput): A dictionary containing inputs for the model, such as initial noise,
                 conditional information, or labels. If 'x' is not provided, random noise will be generated.
+            data_shape (tuple[int, ...], optional): The shape of the data to generate. If a vision tower is used,
+                this shape should correspond to the latent space shape.
+                Required if 'x' is not in model_inputs. Defaults to None.
             use_tqdm (bool, optional): Whether to display a progress bar during generation. Defaults to True.
             clamp_x (bool, optional): Whether to clamp the generated values to [-1, 1] range. Defaults to False.
             guidance_scale (float, optional): Scale for classifier or classifier-free guidance.
                 Values greater than 0 enable guidance. Defaults to 0.
-            **kwargs (dict[str, Any]): Additional arguments to pass to the diffusion model's denoise method.
-                These may include parameters like 'classifier', 'classifier_free', etc.
+            return_latents (bool, optional): Whether to return the latent representation when using a
+                vision tower instead of decoded data. Defaults to False.
         Returns:
             Tensor: The generated data tensor.
         Example:
@@ -177,24 +190,27 @@ def generate(
             ```
         """
         if self.vision_tower:
-            z = self.diffusion.denoise(
+            sampling_output = self.diffusion.denoise(
                 self.denoiser,
-                data_shape,
-                model_inputs,
+                model_inputs=model_inputs,
+                data_shape=data_shape,
                 use_tqdm=use_tqdm,
                 clamp_x=clamp_x,
                 guidance_scale=guidance_scale,
-                **kwargs,
+                sampler_args=sampler_args,
+                return_intermediates=return_intermediates,
             )
-            z = z / self.latent_scale
-            return self.vision_tower.decode(z)
+            if not return_latents:
+                sampling_output["x"] = self.vision_tower.decode(sampling_output["x"] / self.latent_scale)
+            return sampling_output
 
         return self.diffusion.denoise(
             self.denoiser,
-            data_shape,
-            model_inputs,
+            model_inputs=model_inputs,
+            data_shape=data_shape,
             use_tqdm=use_tqdm,
             clamp_x=clamp_x,
             guidance_scale=guidance_scale,
-            **kwargs,
+            sampler_args=sampler_args,
+            return_intermediates=return_intermediates,
         )