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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: denoising-diffusion-pytorch
-Version: 2.2.4
+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.4 → 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.5/denoising_diffusion_pytorch/version.py

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

{denoising_diffusion_pytorch-2.2.4 → 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.4
+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.4/denoising_diffusion_pytorch/version.py

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