tests

Author: williamberman

src/diffusers/models/consistency_decoder_vae.py

@@ -58,6 +58,7 @@ class ConsistencyDecoderVae(ModelMixin, ConfigMixin):
         >>> import torch
         >>> from diffusers import DiffusionPipeline, ConsistencyDecoderVae
 
+        >>> # TODO - is this going to be where the model is uploaded?
         >>> vae = ConsistencyDecoderVae.from_pretrained("openai/consistency-decoder", torch_dtype=pipe.torch_dtype)
         >>> pipe = StableDiffusionPipeline.from_pretrained(
         ...     "runwayml/stable-diffusion-v1-5", vae=vae, torch_dtype=torch.float16
@@ -242,7 +243,9 @@ def decode(
         num_inference_steps=2,
     ) -> Union[DecoderOutput, torch.FloatTensor]:
         z = (z - self.means) / self.stds
-        z = F.interpolate(z, mode="nearest", scale_factor=8)
+
+        scale_factor = 2 ** (len(self.config.block_out_channels) - 1)
+        z = F.interpolate(z, mode="nearest", scale_factor=scale_factor)
 
         batch_size, _, height, width = z.shape
 
@@ -334,7 +337,6 @@ def tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> Consis
 
         return ConsistencyDecoderVaeOutput(latent_dist=posterior)
 
-    # Copied from diffusers.models.autoencoder_kl.AutoencoderKL.forward
     def forward(
         self,
         sample: torch.FloatTensor,
@@ -356,7 +358,7 @@ def forward(
             z = posterior.sample(generator=generator)
         else:
             z = posterior.mode()
-        dec = self.decode(z).sample
+        dec = self.decode(z, generator=generator).sample
 
         if not return_dict:
             return (dec,)

tests/models/test_modeling_common.py

@@ -196,11 +196,15 @@ def test_forward_with_norm_groups(self):
 class ModelTesterMixin:
     main_input_name = None  # overwrite in model specific tester class
     base_precision = 1e-3
+    forward_requires_fresh_args = False
 
     def test_from_save_pretrained(self, expected_max_diff=5e-5):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        if self.forward_requires_fresh_args:
+            model = self.model_class(**self.init_dict)
+        else:
+            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            model = self.model_class(**init_dict)
 
-        model = self.model_class(**init_dict)
         if hasattr(model, "set_default_attn_processor"):
             model.set_default_attn_processor()
         model.to(torch_device)
@@ -214,11 +218,18 @@ def test_from_save_pretrained(self, expected_max_diff=5e-5):
             new_model.to(torch_device)
 
         with torch.no_grad():
-            image = model(**inputs_dict)
+            if self.forward_requires_fresh_args:
+                image = model(**self.inputs_dict(0))
+            else:
+                image = model(**inputs_dict)
+
             if isinstance(image, dict):
                 image = image.to_tuple()[0]
 
-            new_image = new_model(**inputs_dict)
+            if self.forward_requires_fresh_args:
+                new_image = new_model(**self.inputs_dict(0))
+            else:
+                new_image = new_model(**inputs_dict)
 
             if isinstance(new_image, dict):
                 new_image = new_image.to_tuple()[0]
@@ -275,8 +286,11 @@ def test_getattr_is_correct(self):
     )
     def test_set_xformers_attn_processor_for_determinism(self):
         torch.use_deterministic_algorithms(False)
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
+        if self.forward_requires_fresh_args:
+            model = self.model_class(**self.init_dict)
+        else:
+            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            model = self.model_class(**init_dict)
         model.to(torch_device)
 
         if not hasattr(model, "set_attn_processor"):
@@ -286,17 +300,26 @@ def test_set_xformers_attn_processor_for_determinism(self):
         model.set_default_attn_processor()
         assert all(type(proc) == AttnProcessor for proc in model.attn_processors.values())
         with torch.no_grad():
-            output = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output = model(**self.inputs_dict(0))[0]
+            else:
+                output = model(**inputs_dict)[0]
 
         model.enable_xformers_memory_efficient_attention()
         assert all(type(proc) == XFormersAttnProcessor for proc in model.attn_processors.values())
         with torch.no_grad():
-            output_2 = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output_2 = model(**self.inputs_dict(0))[0]
+            else:
+                output_2 = model(**inputs_dict)[0]
 
         model.set_attn_processor(XFormersAttnProcessor())
         assert all(type(proc) == XFormersAttnProcessor for proc in model.attn_processors.values())
         with torch.no_grad():
-            output_3 = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output_3 = model(**self.inputs_dict(0))[0]
+            else:
+                output_3 = model(**inputs_dict)[0]
 
         torch.use_deterministic_algorithms(True)
 
@@ -307,8 +330,12 @@ def test_set_xformers_attn_processor_for_determinism(self):
     @require_torch_gpu
     def test_set_attn_processor_for_determinism(self):
         torch.use_deterministic_algorithms(False)
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
+        if self.forward_requires_fresh_args:
+            model = self.model_class(**self.init_dict)
+        else:
+            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            model = self.model_class(**init_dict)
+
         model.to(torch_device)
 
         if not hasattr(model, "set_attn_processor"):
@@ -317,22 +344,34 @@ def test_set_attn_processor_for_determinism(self):
 
         assert all(type(proc) == AttnProcessor2_0 for proc in model.attn_processors.values())
         with torch.no_grad():
-            output_1 = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output_1 = model(**self.inputs_dict(0))[0]
+            else:
+                output_1 = model(**inputs_dict)[0]
 
         model.set_default_attn_processor()
         assert all(type(proc) == AttnProcessor for proc in model.attn_processors.values())
         with torch.no_grad():
-            output_2 = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output_2 = model(**self.inputs_dict(0))[0]
+            else:
+                output_2 = model(**inputs_dict)[0]
 
         model.set_attn_processor(AttnProcessor2_0())
         assert all(type(proc) == AttnProcessor2_0 for proc in model.attn_processors.values())
         with torch.no_grad():
-            output_4 = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output_4 = model(**self.inputs_dict(0))[0]
+            else:
+                output_4 = model(**inputs_dict)[0]
 
         model.set_attn_processor(AttnProcessor())
         assert all(type(proc) == AttnProcessor for proc in model.attn_processors.values())
         with torch.no_grad():
-            output_5 = model(**inputs_dict)[0]
+            if self.forward_requires_fresh_args:
+                output_5 = model(**self.inputs_dict(0))[0]
+            else:
+                output_5 = model(**inputs_dict)[0]
 
         torch.use_deterministic_algorithms(True)
 
@@ -342,9 +381,12 @@ def test_set_attn_processor_for_determinism(self):
         assert torch.allclose(output_2, output_5, atol=self.base_precision)
 
     def test_from_save_pretrained_variant(self, expected_max_diff=5e-5):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        if self.forward_requires_fresh_args:
+            model = self.model_class(**self.init_dict)
+        else:
+            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            model = self.model_class(**init_dict)
 
-        model = self.model_class(**init_dict)
         if hasattr(model, "set_default_attn_processor"):
             model.set_default_attn_processor()
 
@@ -367,11 +409,17 @@ def test_from_save_pretrained_variant(self, expected_max_diff=5e-5):
             new_model.to(torch_device)
 
         with torch.no_grad():
-            image = model(**inputs_dict)
+            if self.forward_requires_fresh_args:
+                image = model(**self.inputs_dict(0))
+            else:
+                image = model(**inputs_dict)
             if isinstance(image, dict):
                 image = image.to_tuple()[0]
 
-            new_image = new_model(**inputs_dict)
+            if self.forward_requires_fresh_args:
+                new_image = new_model(**self.inputs_dict(0))
+            else:
+                new_image = new_model(**inputs_dict)
 
             if isinstance(new_image, dict):
                 new_image = new_image.to_tuple()[0]
@@ -405,17 +453,26 @@ def test_from_save_pretrained_dtype(self):
                 assert new_model.dtype == dtype
 
     def test_determinism(self, expected_max_diff=1e-5):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
+        if self.forward_requires_fresh_args:
+            model = self.model_class(**self.init_dict)
+        else:
+            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            model = self.model_class(**init_dict)
         model.to(torch_device)
         model.eval()
 
         with torch.no_grad():
-            first = model(**inputs_dict)
+            if self.forward_requires_fresh_args:
+                first = model(**self.inputs_dict(0))
+            else:
+                first = model(**inputs_dict)
             if isinstance(first, dict):
                 first = first.to_tuple()[0]
 
-            second = model(**inputs_dict)
+            if self.forward_requires_fresh_args:
+                second = model(**self.inputs_dict(0))
+            else:
+                second = model(**inputs_dict)
             if isinstance(second, dict):
                 second = second.to_tuple()[0]
 
@@ -548,15 +605,22 @@ def recursive_check(tuple_object, dict_object):
                     ),
                 )
 
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        if self.forward_requires_fresh_args:
+            model = self.model_class(**self.init_dict)
+        else:
+            init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+            model = self.model_class(**init_dict)
 
-        model = self.model_class(**init_dict)
         model.to(torch_device)
         model.eval()
 
         with torch.no_grad():
-            outputs_dict = model(**inputs_dict)
-            outputs_tuple = model(**inputs_dict, return_dict=False)
+            if self.forward_requires_fresh_args:
+                outputs_dict = model(**self.inputs_dict(0))
+                outputs_tuple = model(**self.inputs_dict(0), return_dict=False)
+            else:
+                outputs_dict = model(**inputs_dict)
+                outputs_tuple = model(**inputs_dict, return_dict=False)
 
         recursive_check(outputs_tuple, outputs_dict)
 

tests/models/test_models_vae.py

@@ -19,7 +19,7 @@
 import torch
 from parameterized import parameterized
 
-from diffusers import AsymmetricAutoencoderKL, AutoencoderKL, AutoencoderTiny
+from diffusers import AsymmetricAutoencoderKL, AutoencoderKL, AutoencoderTiny, ConsistencyDecoderVae
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
     enable_full_determinism,
@@ -30,6 +30,7 @@
     torch_all_close,
     torch_device,
 )
+from diffusers.utils.torch_utils import randn_tensor
 
 from .test_modeling_common import ModelTesterMixin, UNetTesterMixin
 
@@ -269,6 +270,79 @@ def test_outputs_equivalence(self):
         pass
 
 
+class ConsistencyDecoderVaeTests(ModelTesterMixin, unittest.TestCase):
+    model_class = ConsistencyDecoderVae
+    main_input_name = "sample"
+    base_precision = 1e-2
+    forward_requires_fresh_args = True
+
+    def inputs_dict(self, seed=None):
+        generator = torch.Generator("cpu")
+        if seed is not None:
+            generator.manual_seed(0)
+        image = randn_tensor((4, 3, 32, 32), generator=generator, device=torch.device(torch_device))
+
+        return {"sample": image, "generator": generator}
+
+    @property
+    def input_shape(self):
+        return (3, 32, 32)
+
+    @property
+    def output_shape(self):
+        return (3, 32, 32)
+
+    @property
+    def init_dict(self):
+        return {
+            "encoder_args": {
+                "block_out_channels": [32, 64],
+                "in_channels": 3,
+                "out_channels": 4,
+                "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
+            },
+            "decoder_args": {
+                "act_fn": "silu",
+                "add_attention": False,
+                "block_out_channels": [32, 64],
+                "down_block_types": [
+                    "ResnetDownsampleBlock2D",
+                    "ResnetDownsampleBlock2D",
+                ],
+                "downsample_padding": 1,
+                "downsample_type": "conv",
+                "dropout": 0.0,
+                "in_channels": 7,
+                "layers_per_block": 1,
+                "norm_eps": 1e-05,
+                "norm_num_groups": 32,
+                "num_train_timesteps": 1024,
+                "out_channels": 6,
+                "resnet_time_scale_shift": "scale_shift",
+                "time_embedding_type": "learned",
+                "up_block_types": [
+                    "ResnetUpsampleBlock2D",
+                    "ResnetUpsampleBlock2D",
+                ],
+                "upsample_type": "conv",
+            },
+            "scaling_factor": 1,
+            "block_out_channels": [32, 64],
+            "latent_channels": 4,
+        }
+
+    def prepare_init_args_and_inputs_for_common(self):
+        return self.init_dict, self.inputs_dict()
+
+    @unittest.skip
+    def test_training(self):
+        ...
+
+    @unittest.skip
+    def test_ema_training(self):
+        ...
+
+
 @slow
 class AutoencoderTinyIntegrationTests(unittest.TestCase):
     def tearDown(self):