[diffusion_labs] Move attention and res blocks under modules

ghstack-source-id: b9ab08988a1272c658b37b4dcf1cddf48cc20270 Pull Request resolved: #490
facebookresearch · Oct 16, 2023 · 011c1a0 · 011c1a0
1 parent a60a47e
commit 011c1a0
Show file tree

Hide file tree

Showing 6 changed files with 7 additions and 9 deletions.
diff --git a/tests/diffusion_labs/test_adm.py b/tests/diffusion_labs/test_adm.py
@@ -11,10 +11,10 @@
 from tests.test_utils import assert_expected, set_rng_seed
 from torch import nn
 from torchmultimodal.diffusion_labs.models.adm_unet.adm import ADMStack, ADMUNet
-from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
+from torchmultimodal.diffusion_labs.modules.layers.attention_block import (
     ADMAttentionBlock,
 )
-from torchmultimodal.diffusion_labs.models.adm_unet.res_block import ADMResBlock
+from torchmultimodal.diffusion_labs.modules.layers.res_block import ADMResBlock
 
 
 @pytest.fixture(autouse=True)

diff --git a/tests/diffusion_labs/test_adm_blocks.py b/tests/diffusion_labs/test_adm_blocks.py
@@ -10,10 +10,10 @@
 
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
+from torchmultimodal.diffusion_labs.modules.layers.attention_block import (
     ADMAttentionBlock,
 )
-from torchmultimodal.diffusion_labs.models.adm_unet.res_block import (
+from torchmultimodal.diffusion_labs.modules.layers.res_block import (
     adm_res_block,
     adm_res_downsample_block,
     adm_res_upsample_block,

diff --git a/tests/diffusion_labs/test_adm_crossattention.py b/tests/diffusion_labs/test_adm_crossattention.py
@@ -8,7 +8,7 @@
 
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
+from torchmultimodal.diffusion_labs.modules.layers.attention_block import (
     adm_attention,
     ADMCrossAttention,
 )

diff --git a/torchmultimodal/diffusion_labs/models/adm_unet/adm.py b/torchmultimodal/diffusion_labs/models/adm_unet/adm.py
@@ -9,10 +9,8 @@
 
 import torch
 from torch import nn, Tensor
-from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
-    adm_attn_block,
-)
-from torchmultimodal.diffusion_labs.models.adm_unet.res_block import (
+from torchmultimodal.diffusion_labs.modules.layers.attention_block import adm_attn_block
+from torchmultimodal.diffusion_labs.modules.layers.res_block import (
     adm_res_block,
     adm_res_downsample_block,
     adm_res_upsample_block,

diff --git a/...n_labs/models/adm_unet/attention_block.py → ...on_labs/modules/layers/attention_block.py b/...n_labs/models/adm_unet/attention_block.py → ...on_labs/modules/layers/attention_block.py
diff --git a/...ffusion_labs/models/adm_unet/res_block.py → ...iffusion_labs/modules/layers/res_block.py b/...ffusion_labs/models/adm_unet/res_block.py → ...iffusion_labs/modules/layers/res_block.py