add another ldm fast test

src/diffusers/models/unet_ldm.py

@@ -81,61 +81,62 @@ def Normalize(in_channels):
     return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
 
 
-class LinearAttention(nn.Module):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x)
-        q, k, v = rearrange(qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3)
-        k = k.softmax(dim=-1)
-        context = torch.einsum("bhdn,bhen->bhde", k, v)
-        out = torch.einsum("bhde,bhdn->bhen", context, q)
-        out = rearrange(out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w)
-        return self.to_out(out)
-
-
-class SpatialSelfAttention(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
+#class LinearAttention(nn.Module):
+#    def __init__(self, dim, heads=4, dim_head=32):
+#        super().__init__()
+#        self.heads = heads
+#        hidden_dim = dim_head * heads
+#        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+#        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+#
+#    def forward(self, x):
+#        b, c, h, w = x.shape
+#        qkv = self.to_qkv(x)
+#        q, k, v = rearrange(qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3)
+#        import ipdb; ipdb.set_trace()
+#        k = k.softmax(dim=-1)
+#        context = torch.einsum("bhdn,bhen->bhde", k, v)
+#        out = torch.einsum("bhde,bhdn->bhen", context, q)
+#        out = rearrange(out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w)
+#        return self.to_out(out)
+#
+
+#class SpatialSelfAttention(nn.Module):
+#    def __init__(self, in_channels):
+#        super().__init__()
+#        self.in_channels = in_channels
+#
+#        self.norm = Normalize(in_channels)
+#        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+#        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+#        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+#        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+#
+#    def forward(self, x):
+#        h_ = x
+#        h_ = self.norm(h_)
+#        q = self.q(h_)
+#        k = self.k(h_)
+#        v = self.v(h_)
+#
         # compute attention
-        b, c, h, w = q.shape
-        q = rearrange(q, "b c h w -> b (h w) c")
-        k = rearrange(k, "b c h w -> b c (h w)")
-        w_ = torch.einsum("bij,bjk->bik", q, k)
-
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
+#        b, c, h, w = q.shape
+#        q = rearrange(q, "b c h w -> b (h w) c")
+#        k = rearrange(k, "b c h w -> b c (h w)")
+#        w_ = torch.einsum("bij,bjk->bik", q, k)
+#
+#        w_ = w_ * (int(c) ** (-0.5))
+#        w_ = torch.nn.functional.softmax(w_, dim=2)
+#
         # attend to values
-        v = rearrange(v, "b c h w -> b c (h w)")
-        w_ = rearrange(w_, "b i j -> b j i")
-        h_ = torch.einsum("bij,bjk->bik", v, w_)
-        h_ = rearrange(h_, "b c (h w) -> b c h w", h=h)
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
+#        v = rearrange(v, "b c h w -> b c (h w)")
+#        w_ = rearrange(w_, "b i j -> b j i")
+#        h_ = torch.einsum("bij,bjk->bik", v, w_)
+#        h_ = rearrange(h_, "b c (h w) -> b c h w", h=h)
+#        h_ = self.proj_out(h_)
+#
+#        return x + h_
+#
 
 class CrossAttention(nn.Module):
     def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):

tests/test_modeling_utils.py

@@ -511,6 +511,28 @@ class UNetLDMModelTests(ModelTesterMixin, unittest.TestCase):
 
         self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
 
+    def test_output_pretrained_spatial_transformer(self):
+        model = UNetLDMModel.from_pretrained("fusing/unet-ldm-dummy-spatial")
+        model.eval()
+
+        torch.manual_seed(0)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(0)
+
+        noise = torch.randn(1, model.config.in_channels, model.config.image_size, model.config.image_size)
+        context = torch.ones((1, 16, 64), dtype=torch.float32)
+        time_step = torch.tensor([10] * noise.shape[0])
+
+        with torch.no_grad():
+            output = model(noise, time_step, context=context)
+
+        output_slice = output[0, -1, -3:, -3:].flatten()
+        # fmt: off
+        expected_output_slice = torch.tensor([61.3445, 56.9005, 29.4339, 59.5497, 60.7375, 34.1719, 48.1951, 42.6569, 25.0890])
+        # fmt: on
+
+        self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
+
 
 class UNetGradTTSModelTests(ModelTesterMixin, unittest.TestCase):
     model_class = UNetGradTTSModel