denoising-diffusion-pytorch 2.3.0__tar.gz → 2.3.2__tar.gz

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: denoising-diffusion-pytorch
-Version: 2.3.0
+Version: 2.3.2
 Summary: Denoising Diffusion Probabilistic Models - Pytorch
 Project-URL: Homepage, https://github.com/lucidrains/denoising-diffusion-pytorch
 Project-URL: Repository, https://github.com/lucidrains/denoising-diffusion-pytorch
@@ -115,6 +115,39 @@ trainer.train()
 
 Samples and model checkpoints will be logged to `./results` periodically
 
+## Explorative Modeling (Forward XM)
+
+To use <a href="https://arxiv.org/abs/2607.27372">Explorative Modeling (Forward XM)</a> for multi-candidate loss calculation during training, wrap any diffusion model in `XMWrapper`.
+
+```python
+import torch
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion, XMWrapper
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+)
+
+diffusion = GaussianDiffusion(
+    model,
+    image_size = 128,
+    timesteps = 1000
+)
+
+xm = XMWrapper(
+    diffusion,
+    candidates = 4 # generate 4 candidates per sample and pick minimum loss
+)
+
+training_images = torch.rand(8, 3, 128, 128)
+loss = xm(training_images)
+loss.backward()
+
+# sampling works as usual
+
+sampled_images = xm.sample(batch_size = 4)
+```
+
 ## Multi-GPU Training
 
 The `Trainer` class is now equipped with <a href="https://huggingface.co/docs/accelerate/accelerator">🤗 Accelerator</a>. You can easily do multi-gpu training in two steps using their `accelerate` CLI

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/README.md

@@ -87,6 +87,39 @@ trainer.train()
 
 Samples and model checkpoints will be logged to `./results` periodically
 
+## Explorative Modeling (Forward XM)
+
+To use <a href="https://arxiv.org/abs/2607.27372">Explorative Modeling (Forward XM)</a> for multi-candidate loss calculation during training, wrap any diffusion model in `XMWrapper`.
+
+```python
+import torch
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion, XMWrapper
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+)
+
+diffusion = GaussianDiffusion(
+    model,
+    image_size = 128,
+    timesteps = 1000
+)
+
+xm = XMWrapper(
+    diffusion,
+    candidates = 4 # generate 4 candidates per sample and pick minimum loss
+)
+
+training_images = torch.rand(8, 3, 128, 128)
+loss = xm(training_images)
+loss.backward()
+
+# sampling works as usual
+
+sampled_images = xm.sample(batch_size = 4)
+```
+
 ## Multi-GPU Training
 
 The `Trainer` class is now equipped with <a href="https://huggingface.co/docs/accelerate/accelerator">🤗 Accelerator</a>. You can easily do multi-gpu training in two steps using their `accelerate` CLI

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/classifier_free_guidance.py

@@ -411,7 +411,8 @@ class Unet(nn.Module):
         x,
         time,
         classes,
-        cond_drop_prob = None
+        cond_drop_prob = None,
+        cond_keep_mask = None
     ):
         batch, device = x.shape[0], x.device
 
@@ -421,8 +422,10 @@ class Unet(nn.Module):
 
         classes_emb = self.classes_emb(classes)
 
-        if cond_drop_prob > 0:
-            keep_mask = prob_mask_like((batch,), 1 - cond_drop_prob, device = device)
+        if cond_drop_prob > 0 or exists(cond_keep_mask):
+            keep_mask = default(cond_keep_mask, lambda: prob_mask_like((batch,), 1 - cond_drop_prob, device = device))
+            assert keep_mask.shape == (batch,)
+            keep_mask = keep_mask.to(device = device, dtype = torch.bool)
             null_classes_emb = repeat(self.null_classes_emb, 'd -> b d', b = batch)
 
             classes_emb = torch.where(
@@ -606,6 +609,20 @@ class GaussianDiffusion(nn.Module):
     def device(self):
         return self.betas.device
 
+    def random_times(self, batch_size):
+        return torch.randint(0, self.num_timesteps, (batch_size,), device = self.device).long()
+
+    def random_cond_keep_mask(self, batch_size):
+        keep_prob = 1. - self.model.cond_drop_prob
+        return prob_mask_like((batch_size,), keep_prob, device = self.device)
+
+    def xm_shared_random_kwargs(self, batch_size):
+        # condition dropping is shared across Forward XM candidates
+
+        return dict(
+            cond_keep_mask = self.random_cond_keep_mask(batch_size)
+        )
+
     def predict_start_from_noise(self, x_t, t, noise):
         return (
             extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
@@ -774,7 +791,7 @@ class GaussianDiffusion(nn.Module):
             extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
-    def p_losses(self, x_start, t, *, classes, noise = None, loss_reduction = 'mean'):
+    def p_losses(self, x_start, t, *, classes, noise = None, cond_keep_mask = None, loss_reduction = 'mean'):
         b, c, h, w = x_start.shape
         noise = default(noise, lambda: torch.randn_like(x_start))
 
@@ -784,7 +801,8 @@ class GaussianDiffusion(nn.Module):
 
         # predict and take gradient step
 
-        model_out = self.model(x, t, classes)
+        model_kwargs = dict(cond_keep_mask = cond_keep_mask) if exists(cond_keep_mask) else dict()
+        model_out = self.model(x, t, classes, **model_kwargs)
 
         if self.objective == 'pred_noise':
             target = noise
@@ -806,13 +824,13 @@ class GaussianDiffusion(nn.Module):
 
         return loss.mean()
 
-    def forward(self, img, *args, loss_reduction = 'mean', **kwargs):
+    def forward(self, img, *args, times = None, cond_keep_mask = None, loss_reduction = 'mean', **kwargs):
         b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
         assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
-        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
+        times = default(times, lambda: self.random_times(b))
 
         img = normalize_to_neg_one_to_one(img)
-        return self.p_losses(img, t, *args, loss_reduction = loss_reduction, **kwargs)
+        return self.p_losses(img, times, *args, cond_keep_mask = cond_keep_mask, loss_reduction = loss_reduction, **kwargs)
 
 # example
 

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -261,7 +261,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
 
         if self.min_snr_loss_weight:
             snr = log_snr.exp()
-            loss_weight = snr.clamp(min = self.min_snr_gamma) / snr
+            loss_weight = snr.clamp(max = self.min_snr_gamma) / snr
             losses = losses * loss_weight
 
         if loss_reduction == 'none':
@@ -269,10 +269,10 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
 
         return losses.mean()
 
-    def forward(self, img, *args, loss_reduction = 'mean', **kwargs):
+    def forward(self, img, *args, times = None, loss_reduction = 'mean', **kwargs):
         b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
         assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
 
-        times = self.random_times(b)
+        times = default(times, self.random_times(b))
         img = normalize_to_neg_one_to_one(img)
         return self.p_losses(img, times, *args, loss_reduction = loss_reduction, **kwargs)

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -836,13 +836,17 @@ class GaussianDiffusion(Module):
 
         return loss.mean()
 
-    def forward(self, img, *args, loss_reduction = 'mean', **kwargs):
+    def random_times(self, batch_size):
+        return torch.randint(0, self.num_timesteps, (batch_size,), device = self.device).long()
+
+    def forward(self, img, *args, times = None, loss_reduction = 'mean', **kwargs):
         b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
         assert h == img_size[0] and w == img_size[1], f'height and width of image must be {img_size}'
-        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
+
+        times = default(times, lambda: self.random_times(b))
 
         img = self.normalize(img)
-        return self.p_losses(img, t, *args, loss_reduction = loss_reduction, **kwargs)
+        return self.p_losses(img, times, *args, loss_reduction = loss_reduction, **kwargs)
 
 # dataset classes
 

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -545,6 +545,10 @@ class GaussianDiffusion1D(Module):
         posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
+    @property
+    def device(self):
+        return self.betas.device
+
     def model_predictions(self, x, t, x_self_cond = None, clip_x_start = False, rederive_pred_noise = False, model_forward_kwargs: dict = dict()):
 
         if exists(x_self_cond):
@@ -747,14 +751,18 @@ class GaussianDiffusion1D(Module):
 
         return loss.mean()
 
-    def forward(self, img, *args, loss_reduction = 'mean', **kwargs):
+    def random_times(self, batch_size):
+        return torch.randint(0, self.num_timesteps, (batch_size,), device = self.device).long()
+
+    def forward(self, img, *args, times = None, loss_reduction = 'mean', **kwargs):
         b, n, device, seq_length, = img.shape[0], img.shape[self.seq_index], img.device, self.seq_length
 
         assert n == seq_length, f'seq length must be {seq_length}'
-        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
+
+        times = default(times, lambda: self.random_times(b))
 
         img = self.normalize(img)
-        return self.p_losses(img, t, *args, loss_reduction = loss_reduction, **kwargs)
+        return self.p_losses(img, times, *args, loss_reduction = loss_reduction, **kwargs)
 
 # trainer class
 

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/elucidated_diffusion.py

@@ -244,7 +244,7 @@ class ElucidatedDiffusion(nn.Module):
     def noise_distribution(self, batch_size):
         return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
 
-    def forward(self, images, loss_reduction = 'mean'):
+    def forward(self, images, *args, sigma = None, loss_reduction = 'mean', **kwargs):
         batch_size, c, h, w, device, image_size, channels = *images.shape, images.device, self.image_size, self.channels
 
         assert h == image_size and w == image_size, f'height and width of image must be {image_size}'
@@ -252,7 +252,7 @@ class ElucidatedDiffusion(nn.Module):
 
         images = normalize_to_neg_one_to_one(images)
 
-        sigmas = self.noise_distribution(batch_size)
+        sigmas = default(sigma, lambda: self.noise_distribution(batch_size))
         padded_sigmas = rearrange(sigmas, 'b -> b 1 1 1')
 
         noise = torch.randn_like(images)

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py

@@ -183,10 +183,10 @@ class VParamContinuousTimeGaussianDiffusion(nn.Module):
 
         return losses.mean()
 
-    def forward(self, img, *args, loss_reduction = 'mean', **kwargs):
+    def forward(self, img, *args, times = None, loss_reduction = 'mean', **kwargs):
         b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
         assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
 
-        times = self.random_times(b)
+        times = default(times, self.random_times(b))
         img = normalize_to_neg_one_to_one(img)
         return self.p_losses(img, times, *args, loss_reduction = loss_reduction, **kwargs)

denoising_diffusion_pytorch-2.3.2/denoising_diffusion_pytorch/version.py

@@ -0,0 +1 @@
+__version__ = '2.3.2'

{denoising_diffusion_pytorch-2.3.0 → denoising_diffusion_pytorch-2.3.2}/denoising_diffusion_pytorch/xm.py

@@ -1,5 +1,6 @@
 from inspect import signature
 
+import torch
 from torch.nn import Module
 from torch import cat, is_tensor
 from torch.utils._pytree import tree_map, tree_flatten
@@ -30,7 +31,9 @@ class XMWrapper(Module):
         self,
         flow_model: Module,
         candidates = 1,
-        max_batch_size = None
+        max_batch_size = None,
+        random_time_method = 'random_times',
+        random_time_kwarg = 'times'
     ):
         super().__init__()
         self.flow_model = flow_model
@@ -40,6 +43,9 @@ class XMWrapper(Module):
         self.max_batch_size = max_batch_size
         self.has_loss_reduction = 'loss_reduction' in signature(flow_model.forward).parameters
 
+        self.random_time_method = random_time_method
+        self.random_time_kwarg = random_time_kwarg
+
     @property
     def data_shape(self):
         return getattr(self.flow_model, 'data_shape', None)
@@ -71,6 +77,20 @@ class XMWrapper(Module):
         first_tensor = next(t for t in leaves if is_tensor(t))
         batch = first_tensor.shape[0]
 
+        # sample anything that must be shared before expanding candidates
+
+        if hasattr(self.flow_model, 'xm_shared_random_kwargs'):
+            random_kwargs = self.flow_model.xm_shared_random_kwargs(batch)
+
+            for key, value in random_kwargs.items():
+                if key not in kwargs:
+                    kwargs[key] = value
+
+        if self.random_time_kwarg not in kwargs:
+            assert hasattr(self.flow_model, self.random_time_method), f'flow_model must have a {self.random_time_method} method'
+            fn = getattr(self.flow_model, self.random_time_method)
+            kwargs[self.random_time_kwarg] = fn(batch)
+
         # repeat inputs K candidates times
 
         args_K, kwargs_K = tree_map(

denoising_diffusion_pytorch-2.3.0/denoising_diffusion_pytorch/version.py

@@ -1 +0,0 @@
-__version__ = '2.3.0'