Reorg OSS Diffusion Components to diffusion_labs folder

Summary:
Based on this [proposal](https://docs.google.com/document/d/1GtN2urD8PiRr1X4COzvbVbNE8LWRrcoYkAO4v2aogO8/edit) to reorganize diffusion components and models under a new `diffusion_labs`. This is the first in a stack of diffs. This one only reorganizes what's already been moved to OSS.

This is primarily moving files with a couple of changes based on the proposal:
- predictors.py is split into a separate file per predictor
- adm is moved out of dalle2 to be it's own model adm_unet
- Dalle2ImageTransform is moved to dalle2 out of transforms
- schedule.py is renamed to discrete_guassian_schedule.py and an abstract DIffusionSchedule class was added
- An abstract adapter class was added to be a generic type and enforce the `forward` signature
- An abstract sampler class was added to be a generic type and enforce the `forward` and `generator` signature
- A new dalle2_model unit test was added

Reviewed By: abhinavarora

Differential Revision: D49790849

tests/modules/diffusion/test_cfguidance.py → tests/diffusion_labs/test_adapter_cfguidance.py

@@ -9,8 +9,8 @@
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
 from torch import nn
-from torchmultimodal.modules.diffusion.cfguidance import CFGuidance
-from torchmultimodal.utils.diffusion_utils import DiffusionOutput
+from torchmultimodal.diffusion_labs.modules.adapters.cfguidance import CFGuidance
+from torchmultimodal.diffusion_labs.utils.common import DiffusionOutput
 
 
 @pytest.fixture(autouse=True)

tests/models/dalle2/test_adm.py → tests/diffusion_labs/test_adm.py

@@ -10,9 +10,11 @@
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
 from torch import nn
-from torchmultimodal.models.dalle2.adm.adm import ADM, ADMStack, ADMUNet
-from torchmultimodal.models.dalle2.adm.attention_block import ADMAttentionBlock
-from torchmultimodal.models.dalle2.adm.res_block import ADMResBlock
+from torchmultimodal.diffusion_labs.models.adm_unet.adm import ADM, ADMStack, ADMUNet
+from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
+    ADMAttentionBlock,
+)
+from torchmultimodal.diffusion_labs.models.adm_unet.res_block import ADMResBlock
 
 
 @pytest.fixture(autouse=True)
@@ -116,7 +118,7 @@ def test_predict_variance_value_incorrect_channel_dim_error(
 
 
 # All expected values come after first testing the ADMUNet has the exact output
-# as the corresponding UNet class in d2go, then simply forward passing
+# as the corresponding author UNet implementation, then simply forward passing
 # ADMUNet with params, random seed, and initialization order in this file.
 class TestADMUNet:
     @pytest.fixture

tests/models/dalle2/test_adm_blocks.py → tests/diffusion_labs/test_adm_blocks.py

@@ -10,8 +10,10 @@
 
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.models.dalle2.adm.attention_block import ADMAttentionBlock
-from torchmultimodal.models.dalle2.adm.res_block import (
+from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
+    ADMAttentionBlock,
+)
+from torchmultimodal.diffusion_labs.models.adm_unet.res_block import (
     adm_res_block,
     adm_res_downsample_block,
     adm_res_upsample_block,
@@ -47,7 +49,7 @@ def t(params):
 
 
 # All expected values come after first testing the ADMResBlock has the exact output
-# as the corresponding residual block class in d2go, then simply forward passing
+# as the corresponding residual block class from ADM authors, then simply forward passing
 # ADMResBlock with params, random seed, and initialization order in this file.
 class TestADMResBlock:
     @pytest.fixture

tests/models/dalle2/test_adm_crossattention.py → 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.models.dalle2.adm.attention_block import (
+from torchmultimodal.diffusion_labs.models.adm_unet.attention_block import (
     adm_attention,
     ADMCrossAttention,
 )
@@ -44,7 +44,7 @@ def c(params):
 
 
 # All expected values come after first testing that ADMCrossAttention has
-# the exact output as the corresponding QKVAttention class in d2go, then simply forward passing
+# the exact output as the corresponding QKVAttention class in ADM, then simply forward passing
 # ADMCrossAttention with params, random seed, and initialization order in this file.
 class TestADMCrossAttention:
     @pytest.fixture

tests/utils/test_diffusion_utils.py → tests/diffusion_labs/test_common_util.py

@@ -6,7 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from PIL import Image
-from torchmultimodal.utils.diffusion_utils import cascaded_resize
+from torchmultimodal.diffusion_labs.utils.common import cascaded_resize
 
 
 def test_cascaded_resize():

tests/diffusion_labs/test_dalle2.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env fbpython
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from PIL import Image
+from tests.test_utils import assert_expected, set_rng_seed
+from torchmultimodal.diffusion_labs.models.dalle2.dalle2_decoder import dalle2_decoder
+from torchmultimodal.diffusion_labs.models.dalle2.transforms import Dalle2ImageTransform
+
+
+def test_dalle2_model():
+    set_rng_seed(4)
+    model = dalle2_decoder(
+        timesteps=1,
+        time_embed_dim=1,
+        cond_embed_dim=1,
+        clip_embed_dim=1,
+        clip_embed_name="clip_image",
+        predict_variance_value=True,
+        image_channels=1,
+        depth=32,
+        num_resize=1,
+        num_res_per_layer=1,
+        use_cf_guidance=True,
+        clip_image_guidance_dropout=0.1,
+        guidance_strength=7.0,
+        learn_null_emb=True,
+    )
+    model.eval()
+    x = torch.randn(1, 1, 4, 4)
+    c = torch.ones((1, 1))
+    with torch.no_grad():
+        actual = model(x, conditional_inputs={"clip_image": c}).mean()
+    expected = torch.as_tensor(0.12768)
+    assert_expected(actual, expected, rtol=0, atol=1e-4)
+
+
+def test_dalle2_image_transform():
+    img_size = 5
+    transform = Dalle2ImageTransform(image_size=img_size, image_min=-1, image_max=1)
+    image = Image.new("RGB", size=(20, 20), color=(128, 0, 0))
+    actual = transform(image).sum()
+    normalized128 = 128 / 255 * 2 - 1
+    normalized0 = -1
+    expected = torch.tensor(
+        normalized128 * img_size**2 + 2 * normalized0 * img_size**2
+    )
+    assert_expected(actual, expected, rtol=0, atol=1e-4)

tests/modules/losses/test_diffusion_loss.py → tests/diffusion_labs/test_diffusion_losses.py

@@ -9,12 +9,15 @@
 import torch
 import torch.nn as nn
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.modules.diffusion.schedules import (
-    DiffusionSchedule,
+from torchmultimodal.diffusion_labs.modules.losses.diffusion_hybrid_loss import (
+    DiffusionHybridLoss,
+)
+from torchmultimodal.diffusion_labs.modules.losses.vlb_loss import VLBLoss
+from torchmultimodal.diffusion_labs.schedules.discrete_gaussian_schedule import (
+    DiscreteGaussianSchedule,
     linear_beta_schedule,
 )
-from torchmultimodal.modules.losses.diffusion import DiffusionHybridLoss, VLBLoss
-from torchmultimodal.utils.diffusion_utils import DiffusionOutput
+from torchmultimodal.diffusion_labs.utils.common import DiffusionOutput
 
 
 @pytest.fixture(autouse=True)
@@ -24,7 +27,7 @@ def set_seed():
 
 @pytest.fixture
 def schedule():
-    return DiffusionSchedule(linear_beta_schedule(1000))
+    return DiscreteGaussianSchedule(linear_beta_schedule(1000))
 
 
 @pytest.fixture
@@ -44,7 +47,7 @@ def target():
 
 
 # All expected values come after first testing the HybridLoss has the exact output
-# as the corresponding p_losses in D2Go Guassian Diffusion
+# as the corresponding p_losses in Guassian Diffusion
 class TestDiffusionHybridLoss:
     @pytest.fixture
     def loss(self, schedule):

tests/transforms/test_diffusion_transforms.py → tests/diffusion_labs/test_diffusion_transform.py

@@ -6,10 +6,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from PIL import Image
-from tests.test_utils import assert_expected
-from torchmultimodal.transforms.diffusion_transforms import (
-    Dalle2ImageTransform,
+from torchmultimodal.diffusion_labs.transforms.diffusion_transform import (
     RandomDiffusionSteps,
 )
 
@@ -30,16 +27,3 @@ def test_random_diffusion_steps():
     actual = len(transform(torch.ones(1)))
     expected = 4
     assert actual == expected, "Transform not returning correct keys"
-
-
-def test_dalle_image_transform():
-    img_size = 5
-    transform = Dalle2ImageTransform(image_size=img_size, image_min=-1, image_max=1)
-    image = Image.new("RGB", size=(20, 20), color=(128, 0, 0))
-    actual = transform(image).sum()
-    normalized128 = 128 / 255 * 2 - 1
-    normalized0 = -1
-    expected = torch.tensor(
-        normalized128 * img_size**2 + 2 * normalized0 * img_size**2
-    )
-    assert_expected(actual, expected, rtol=0, atol=1e-4)

tests/modules/diffusion/test_schedule.py → tests/diffusion_labs/test_discrete_schedule.py

@@ -8,9 +8,9 @@
 import pytest
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.modules.diffusion.schedules import (
+from torchmultimodal.diffusion_labs.schedules.discrete_gaussian_schedule import (
     cosine_beta_schedule,
-    DiffusionSchedule,
+    DiscreteGaussianSchedule,
     linear_beta_schedule,
     quadratic_beta_schedule,
     sigmoid_beta_schedule,
@@ -23,11 +23,11 @@ def set_seed():
 
 
 # All expected values come after first testing the Schedule has the exact output
-# as the corresponding q methods from GaussianDiffusion in D2Go
+# as the corresponding q methods from GaussianDiffusion
 class TestDiffusionSchedule:
     @pytest.fixture
     def module(self):
-        schedule = DiffusionSchedule(linear_beta_schedule(1000))
+        schedule = DiscreteGaussianSchedule(linear_beta_schedule(1000))
         return schedule
 
     @pytest.fixture

tests/modules/diffusion/test_predictors.py → tests/diffusion_labs/test_predictors.py

@@ -8,9 +8,10 @@
 import pytest
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.modules.diffusion.predictors import NoisePredictor, TargetPredictor
-from torchmultimodal.modules.diffusion.schedules import (
-    DiffusionSchedule,
+from torchmultimodal.diffusion_labs.predictors.noise_predictor import NoisePredictor
+from torchmultimodal.diffusion_labs.predictors.target_predictor import TargetPredictor
+from torchmultimodal.diffusion_labs.schedules.discrete_gaussian_schedule import (
+    DiscreteGaussianSchedule,
     linear_beta_schedule,
 )
 
@@ -30,11 +31,11 @@ def input():
 
 
 # All expected values come after first testing the Schedule has the exact output
-# as the corresponding q methods from GaussianDiffusion in D2Go
+# as the corresponding q methods from GaussianDiffusion
 class TestNoisePredictor:
     @pytest.fixture
     def module(self):
-        schedule = DiffusionSchedule(linear_beta_schedule(1000))
+        schedule = DiscreteGaussianSchedule(linear_beta_schedule(1000))
         predictor = NoisePredictor(schedule, None)
         return predictor
 
@@ -52,7 +53,7 @@ def test_predict_noise(self, module, input):
 class TestTargetPredictor:
     @pytest.fixture
     def module(self):
-        schedule = DiffusionSchedule(linear_beta_schedule(1000))
+        schedule = DiscreteGaussianSchedule(linear_beta_schedule(1000))
         predictor = TargetPredictor(schedule, None)
         return predictor
 

tests/modules/diffusion/test_ddim.py → tests/diffusion_labs/test_sampler_ddim.py

@@ -9,13 +9,13 @@
 import torch
 import torch.nn as nn
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.modules.diffusion.ddim import DDIModule
-from torchmultimodal.modules.diffusion.predictors import NoisePredictor
-from torchmultimodal.modules.diffusion.schedules import (
-    DiffusionSchedule,
+from torchmultimodal.diffusion_labs.predictors.noise_predictor import NoisePredictor
+from torchmultimodal.diffusion_labs.samplers.ddim import DDIModule
+from torchmultimodal.diffusion_labs.schedules.discrete_gaussian_schedule import (
+    DiscreteGaussianSchedule,
     linear_beta_schedule,
 )
-from torchmultimodal.utils.diffusion_utils import DiffusionOutput
+from torchmultimodal.diffusion_labs.utils.common import DiffusionOutput
 
 
 class DummyUNet(nn.Module):
@@ -38,7 +38,7 @@ class TestDDIModule:
     @pytest.fixture
     def module(self):
         model = DummyUNet(True)
-        schedule = DiffusionSchedule(linear_beta_schedule(1000))
+        schedule = DiscreteGaussianSchedule(linear_beta_schedule(1000))
         predictor = NoisePredictor(schedule)
         model = DDIModule(model, schedule, predictor)
         return model

tests/modules/diffusion/test_ddpm.py → tests/diffusion_labs/test_sampler_ddpm.py

@@ -10,13 +10,13 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from tests.test_utils import assert_expected, set_rng_seed
-from torchmultimodal.modules.diffusion.ddpm import DDPModule
-from torchmultimodal.modules.diffusion.predictors import NoisePredictor
-from torchmultimodal.modules.diffusion.schedules import (
-    DiffusionSchedule,
+from torchmultimodal.diffusion_labs.predictors.noise_predictor import NoisePredictor
+from torchmultimodal.diffusion_labs.samplers.ddpm import DDPModule
+from torchmultimodal.diffusion_labs.schedules.discrete_gaussian_schedule import (
+    DiscreteGaussianSchedule,
     linear_beta_schedule,
 )
-from torchmultimodal.utils.diffusion_utils import DiffusionOutput
+from torchmultimodal.diffusion_labs.utils.common import DiffusionOutput
 
 
 class DummyUNet(nn.Module):
@@ -36,13 +36,13 @@ def forward(self, x, t, c):
 
 
 # All expected values come after first testing the Schedule has the exact output
-# as the corresponding p methods from GaussianDiffusion in D2Go
+# as the corresponding p methods from GaussianDiffusion
 class TestDDPModule:
     @pytest.fixture
     def module(self):
         set_rng_seed(4)
         model = DummyUNet(True)
-        schedule = DiffusionSchedule(linear_beta_schedule(1000))
+        schedule = DiscreteGaussianSchedule(linear_beta_schedule(1000))
         predictor = NoisePredictor(schedule)
         eval_steps = torch.arange(0, 1000, 50)
         model = DDPModule(model, schedule, predictor, eval_steps)

torchmultimodal/diffusion_labs/models/adm_unet/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.