[tests] test encode_prompt() in isolation (huggingface#10438)

* poc encode_prompt() tests

* fix

* updates.

* fixes

* fixes

* updates

* updates

* updates

* revert

* updates

* updates

* updates

* updates

* remove SDXLOptionalComponentsTesterMixin.

* remove tests that directly leveraged encode_prompt() in some way or the other.

* fix imports.

* remove _save_load

* fixes

* fixes

* fixes

* fixes

Author: sayakpaul

src/diffusers/pipelines/pag/pipeline_pag_sana.py

@@ -268,7 +268,8 @@ def encode_prompt(
         else:
             batch_size = prompt_embeds.shape[0]
 
-        self.tokenizer.padding_side = "right"
+        if getattr(self, "tokenizer", None) is not None:
+            self.tokenizer.padding_side = "right"
 
         # See Section 3.1. of the paper.
         max_length = max_sequence_length

src/diffusers/pipelines/sana/pipeline_sana.py

@@ -312,7 +312,8 @@ def encode_prompt(
         else:
             batch_size = prompt_embeds.shape[0]
 
-        self.tokenizer.padding_side = "right"
+        if getattr(self, "tokenizer", None) is not None:
+            self.tokenizer.padding_side = "right"
 
         # See Section 3.1. of the paper.
         max_length = max_sequence_length

src/diffusers/utils/source_code_parsing_utils.py

@@ -0,0 +1,52 @@
+import ast
+import importlib
+import inspect
+import textwrap
+
+
+class ReturnNameVisitor(ast.NodeVisitor):
+    """Thanks to ChatGPT for pairing."""
+
+    def __init__(self):
+        self.return_names = []
+
+    def visit_Return(self, node):
+        # Check if the return value is a tuple.
+        if isinstance(node.value, ast.Tuple):
+            for elt in node.value.elts:
+                if isinstance(elt, ast.Name):
+                    self.return_names.append(elt.id)
+                else:
+                    try:
+                        self.return_names.append(ast.unparse(elt))
+                    except Exception:
+                        self.return_names.append(str(elt))
+        else:
+            if isinstance(node.value, ast.Name):
+                self.return_names.append(node.value.id)
+            else:
+                try:
+                    self.return_names.append(ast.unparse(node.value))
+                except Exception:
+                    self.return_names.append(str(node.value))
+        self.generic_visit(node)
+
+    def _determine_parent_module(self, cls):
+        from diffusers import DiffusionPipeline
+        from diffusers.models.modeling_utils import ModelMixin
+
+        if issubclass(cls, DiffusionPipeline):
+            return "pipelines"
+        elif issubclass(cls, ModelMixin):
+            return "models"
+        else:
+            raise NotImplementedError
+
+    def get_ast_tree(self, cls, attribute_name="encode_prompt"):
+        parent_module_name = self._determine_parent_module(cls)
+        main_module = importlib.import_module(f"diffusers.{parent_module_name}")
+        current_cls_module = getattr(main_module, cls.__name__)
+        source_code = inspect.getsource(getattr(current_cls_module, attribute_name))
+        source_code = textwrap.dedent(source_code)
+        tree = ast.parse(source_code)
+        return tree

tests/pipelines/animatediff/test_animatediff.py

@@ -548,6 +548,14 @@ def test_xformers_attention_forwardGenerator_pass(self):
     def test_vae_slicing(self):
         return super().test_vae_slicing(image_count=2)
 
+    def test_encode_prompt_works_in_isolation(self):
+        extra_required_param_value_dict = {
+            "device": torch.device(torch_device).type,
+            "num_images_per_prompt": 1,
+            "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
+        }
+        return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
+
 
 @slow
 @require_torch_accelerator

tests/pipelines/animatediff/test_animatediff_controlnet.py

@@ -517,3 +517,11 @@ def test_vae_slicing(self, video_count=2):
         output_2 = pipe(**inputs)
 
         assert np.abs(output_2[0].flatten() - output_1[0].flatten()).max() < 1e-2
+
+    def test_encode_prompt_works_in_isolation(self):
+        extra_required_param_value_dict = {
+            "device": torch.device(torch_device).type,
+            "num_images_per_prompt": 1,
+            "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
+        }
+        return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)

tests/pipelines/animatediff/test_animatediff_sdxl.py

@@ -21,7 +21,6 @@
     IPAdapterTesterMixin,
     PipelineTesterMixin,
     SDFunctionTesterMixin,
-    SDXLOptionalComponentsTesterMixin,
 )
 
 
@@ -36,7 +35,6 @@ class AnimateDiffPipelineSDXLFastTests(
     IPAdapterTesterMixin,
     SDFunctionTesterMixin,
     PipelineTesterMixin,
-    SDXLOptionalComponentsTesterMixin,
     unittest.TestCase,
 ):
     pipeline_class = AnimateDiffSDXLPipeline
@@ -250,33 +248,6 @@ def test_to_dtype(self):
         model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")]
         self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes))
 
-    def test_prompt_embeds(self):
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.set_progress_bar_config(disable=None)
-        pipe.to(torch_device)
-
-        inputs = self.get_dummy_inputs(torch_device)
-        prompt = inputs.pop("prompt")
-
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-            pooled_prompt_embeds,
-            negative_pooled_prompt_embeds,
-        ) = pipe.encode_prompt(prompt)
-
-        pipe(
-            **inputs,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-        )
-
-    def test_save_load_optional_components(self):
-        self._test_save_load_optional_components()
-
     @unittest.skipIf(
         torch_device != "cuda" or not is_xformers_available(),
         reason="XFormers attention is only available with CUDA and `xformers` installed",
@@ -305,3 +276,11 @@ def test_xformers_attention_forwardGenerator_pass(self):
 
         max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
         self.assertLess(max_diff, 1e-4, "XFormers attention should not affect the inference results")
+
+    @unittest.skip("Test currently not supported.")
+    def test_encode_prompt_works_in_isolation(self):
+        pass
+
+    @unittest.skip("Functionality is tested elsewhere.")
+    def test_save_load_optional_components(self):
+        pass

tests/pipelines/animatediff/test_animatediff_sparsectrl.py

@@ -484,3 +484,11 @@ def test_free_init_with_schedulers(self):
 
     def test_vae_slicing(self):
         return super().test_vae_slicing(image_count=2)
+
+    def test_encode_prompt_works_in_isolation(self):
+        extra_required_param_value_dict = {
+            "device": torch.device(torch_device).type,
+            "num_images_per_prompt": 1,
+            "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
+        }
+        return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)

tests/pipelines/animatediff/test_animatediff_video2video.py

@@ -544,3 +544,11 @@ def test_free_noise_multi_prompt(self):
             inputs["strength"] = 0.5
             inputs["prompt"] = {0: "Caterpillar on a leaf", 10: "Butterfly on a leaf", 42: "Error on a leaf"}
             pipe(**inputs).frames[0]
+
+    def test_encode_prompt_works_in_isolation(self):
+        extra_required_param_value_dict = {
+            "device": torch.device(torch_device).type,
+            "num_images_per_prompt": 1,
+            "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
+        }
+        return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)

tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py

@@ -533,3 +533,11 @@ def test_free_noise_multi_prompt(self):
             inputs["strength"] = 0.5
             inputs["prompt"] = {0: "Caterpillar on a leaf", 10: "Butterfly on a leaf", 42: "Error on a leaf"}
             pipe(**inputs).frames[0]
+
+    def test_encode_prompt_works_in_isolation(self):
+        extra_required_param_value_dict = {
+            "device": torch.device(torch_device).type,
+            "num_images_per_prompt": 1,
+            "do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
+        }
+        return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)

tests/pipelines/audioldm2/test_audioldm2.py

@@ -508,9 +508,14 @@ def test_to_dtype(self):
         model_dtypes = {key: component.dtype for key, component in components.items() if hasattr(component, "dtype")}
         self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes.values()))
 
+    @unittest.skip("Test not supported.")
     def test_sequential_cpu_offload_forward_pass(self):
         pass
 
+    @unittest.skip("Test not supported for now because of the use of `projection_model` in `encode_prompt()`.")
+    def test_encode_prompt_works_in_isolation(self):
+        pass
+
 
 @nightly
 class AudioLDM2PipelineSlowTests(unittest.TestCase):