updates.

Author: sayakpaul

src/diffusers/utils/source_code_parsing_utils.py

@@ -0,0 +1,51 @@
+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/test_pipelines_common.py

@@ -5,6 +5,8 @@
 import tempfile
 import unittest
 import uuid
+import textwrap
+import ast
 from typing import Any, Callable, Dict, Union
 
 import numpy as np
@@ -14,7 +16,7 @@
 from huggingface_hub import ModelCard, delete_repo
 from huggingface_hub.utils import is_jinja_available
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
-
+import importlib
 import diffusers
 from diffusers import (
     AsymmetricAutoencoderKL,
@@ -51,6 +53,7 @@
     skip_mps,
     torch_device,
 )
+from diffusers.utils.source_code_parsing_utils import ReturnNameVisitor
 
 from ..models.autoencoders.vae import (
     get_asym_autoencoder_kl_config,
@@ -1997,26 +2000,27 @@ def test_encode_prompt_works_in_isolation(self):
                 components_with_text_encoders[k] = components[k]
             else:
                 components_with_text_encoders[k] = None
-        pipe = self.pipeline_class(**components_with_text_encoders)
-        pipe = pipe.to(torch_device)
+        pipe_with_just_text_encoder = self.pipeline_class(**components_with_text_encoders)
+        pipe_with_just_text_encoder = pipe_with_just_text_encoder.to(torch_device)
 
         inputs = self.get_dummy_inputs(torch_device)
-        encode_prompt_signature = inspect.signature(pipe.encode_prompt)
+        encode_prompt_signature = inspect.signature(pipe_with_just_text_encoder.encode_prompt)
         encode_prompt_parameters = list(encode_prompt_signature.parameters.values())
 
-        # Required parameters in encode_prompt = those with no default
+        # Required parameters in encode_prompt with those with no default
         required_params = []
         for param in encode_prompt_parameters:
-            if param.name == "self":
+            if param.name == "self" or param.name == "kwargs":
                 continue
             if param.default is inspect.Parameter.empty:
                 required_params.append(param.name)
 
+        # Craft inputs for the `encode_prompt()` method to run in isolation.
         encode_prompt_param_names = [p.name for p in encode_prompt_parameters if p.name != "self"]
         input_keys = list(inputs.keys())
         encode_prompt_inputs = {k: inputs.pop(k) for k in input_keys if k in encode_prompt_param_names}
 
-        pipe_call_signature = inspect.signature(pipe.__call__)
+        pipe_call_signature = inspect.signature(pipe_with_just_text_encoder.__call__)
         pipe_call_parameters = pipe_call_signature.parameters
 
         # For each required param in encode_prompt, check if it's missing
@@ -2034,28 +2038,44 @@ def test_encode_prompt_works_in_isolation(self):
                         f"encode_prompt has no default in either encode_prompt or __call__."
                     )
 
+        # Compute `encode_prompt()`.
         with torch.no_grad():
-            encoded_prompt_outputs = pipe.encode_prompt(**encode_prompt_inputs)
-
-        prompt_embeds_kwargs = dict(zip(self.prompt_embed_kwargs, encoded_prompt_outputs))
+            encoded_prompt_outputs = pipe_with_just_text_encoder.encode_prompt(**encode_prompt_inputs)
+
+        # Programatically determine the reutrn names of `encode_prompt.`
+        ast_vistor = ReturnNameVisitor()
+        encode_prompt_tree = ast_vistor.get_ast_tree(cls=self.pipeline_class)
+        ast_vistor.visit(encode_prompt_tree)
+        prompt_embed_kwargs = ast_vistor.return_names
+        print(f"{prompt_embed_kwargs=}")
+        prompt_embeds_kwargs = dict(zip(prompt_embed_kwargs, encoded_prompt_outputs))
+        # Pack the outputs of `encode_prompt`.
         adapted_prompt_embeds_kwargs = {
             k: prompt_embeds_kwargs.pop(k) for k in list(prompt_embeds_kwargs.keys()) if k in pipe_call_parameters
         }
 
-        # now initialize a pipeline without text encoders
+        # now initialize a pipeline without text encoders and compute outputs with the
+        # `encode_prompt()` outputs and other relevant inputs.
         components_with_text_encoders = {}
         for k in components:
             if "text" in k or "tokenizer" in k:
                 components_with_text_encoders[k] = None
             else:
                 components_with_text_encoders[k] = components[k]
         pipe_without_text_encoders = self.pipeline_class(**components_with_text_encoders).to(torch_device)
-        pipe_out = pipe_without_text_encoders(**inputs, **adapted_prompt_embeds_kwargs)[0]
 
+        # Set `negative_prompt` to None as we have already calculated its embeds
+        # if it was present in `inputs`. This is because otherwise we will interfere wrongly
+        # for non-None `negative_prompt` values as defaults (PixArt for example).
+        pipe_without_tes_inputs = {**inputs, **adapted_prompt_embeds_kwargs}
+        if pipe_call_parameters.get("negative_prompt", None) is not None:
+            pipe_without_tes_inputs.update({"negative_prompt": None})
+        pipe_out = pipe_without_text_encoders(**pipe_without_tes_inputs)[0]
+
+        # Compare against regular pipeline outputs.
         full_pipe = self.pipeline_class(**components).to(torch_device)
         inputs = self.get_dummy_inputs(torch_device)
         pipe_out_2 = full_pipe(**inputs)[0]
-
         self.assertTrue(np.allclose(pipe_out, pipe_out_2, atol=1e-4, rtol=1e-4))
 
     def test_StableDiffusionMixin_component(self):