denoising-diffusion-pytorch 2.2.3__tar.gz → 2.2.5__tar.gz

{denoising_diffusion_pytorch-2.2.3 → denoising_diffusion_pytorch-2.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: denoising-diffusion-pytorch
-Version: 2.2.3
+Version: 2.2.5
 Summary: Denoising Diffusion Probabilistic Models - Pytorch
 Home-page: https://github.com/lucidrains/denoising-diffusion-pytorch
 Author: Phil Wang

{denoising_diffusion_pytorch-2.2.3 → denoising_diffusion_pytorch-2.2.5}/denoising_diffusion_pytorch/attend.py

@@ -7,10 +7,11 @@ from torch import nn, einsum
 import torch.nn.functional as F
 
 from einops import rearrange
+from torch.nn.attention import SDPBackend
 
 # constants
 
-AttentionConfig = namedtuple('AttentionConfig', ['enable_flash', 'enable_math', 'enable_mem_efficient'])
+AttentionConfig = namedtuple('AttentionConfig', ['backends'])
 
 # helpers
 
@@ -52,7 +53,7 @@ class Attend(nn.Module):
 
         # determine efficient attention configs for cuda and cpu
 
-        self.cpu_config = AttentionConfig(True, True, True)
+        self.cpu_config = AttentionConfig([SDPBackend.FLASH_ATTENTION, SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION])
         self.cuda_config = None
 
         if not torch.cuda.is_available() or not flash:
@@ -64,10 +65,10 @@ class Attend(nn.Module):
 
         if device_version > version.parse('8.0'):
             print_once('A100 GPU detected, using flash attention if input tensor is on cuda')
-            self.cuda_config = AttentionConfig(True, False, False)
+            self.cuda_config = AttentionConfig([SDPBackend.FLASH_ATTENTION])
         else:
             print_once('Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda')
-            self.cuda_config = AttentionConfig(False, True, True)
+            self.cuda_config = AttentionConfig([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION])
 
     def flash_attn(self, q, k, v):
         _, heads, q_len, _, k_len, is_cuda, device = *q.shape, k.shape[-2], q.is_cuda, q.device
@@ -84,7 +85,7 @@ class Attend(nn.Module):
 
         # pytorch 2.0 flash attn: q, k, v, mask, dropout, causal, softmax_scale
 
-        with torch.backends.cuda.sdp_kernel(**config._asdict()):
+        with torch.nn.attention.sdpa_kernel(**config._asdict()):
             out = F.scaled_dot_product_attention(
                 q, k, v,
                 dropout_p = self.dropout if self.training else 0.

{denoising_diffusion_pytorch-2.2.3 → denoising_diffusion_pytorch-2.2.5}/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -734,14 +734,15 @@ class GaussianDiffusion1D(Module):
             raise ValueError(f'unknown objective {self.objective}')
 
         loss = F.mse_loss(model_out, target, reduction = 'none')
+
+        if not return_reduced_loss:
+            return loss * extract(self.loss_weight, t, loss.shape)
+
         loss = reduce(loss, 'b ... -> b', 'mean')
 
         loss = loss * extract(self.loss_weight, t, loss.shape)
 
-        if return_reduced_loss:
-            loss = loss.mean()
-
-        return loss
+        return loss.mean()
 
     def forward(self, img, *args, **kwargs):
         b, n, device, seq_length, = img.shape[0], img.shape[self.seq_index], img.device, self.seq_length

denoising_diffusion_pytorch-2.2.5/denoising_diffusion_pytorch/version.py

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

{denoising_diffusion_pytorch-2.2.3 → denoising_diffusion_pytorch-2.2.5}/denoising_diffusion_pytorch.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: denoising-diffusion-pytorch
-Version: 2.2.3
+Version: 2.2.5
 Summary: Denoising Diffusion Probabilistic Models - Pytorch
 Home-page: https://github.com/lucidrains/denoising-diffusion-pytorch
 Author: Phil Wang

denoising_diffusion_pytorch-2.2.3/denoising_diffusion_pytorch/version.py

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