fixes

Author: sayakpaul

src/diffusers/utils/source_code_parsing_utils.py

@@ -1,14 +1,15 @@
 import ast
 import importlib
-import inspect 
+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):
@@ -37,7 +38,7 @@ def _determine_parent_module(self, cls):
         if issubclass(cls, DiffusionPipeline):
             return "pipelines"
         elif issubclass(cls, ModelMixin):
-            return "models" 
+            return "models"
         else:
             raise NotImplementedError
 
@@ -48,4 +49,4 @@ def get_ast_tree(self, cls, attribute_name="encode_prompt"):
         source_code = inspect.getsource(getattr(current_cls_module, attribute_name))
         source_code = textwrap.dedent(source_code)
         tree = ast.parse(source_code)
-        return tree
+        return tree

tests/pipelines/test_pipelines_common.py

@@ -5,8 +5,6 @@
 import tempfile
 import unittest
 import uuid
-import textwrap
-import ast
 from typing import Any, Callable, Dict, Union
 
 import numpy as np
@@ -16,7 +14,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,
@@ -43,6 +41,7 @@
 from diffusers.schedulers import KarrasDiffusionSchedulers
 from diffusers.utils import logging
 from diffusers.utils.import_utils import is_xformers_available
+from diffusers.utils.source_code_parsing_utils import ReturnNameVisitor
 from diffusers.utils.testing_utils import (
     CaptureLogger,
     require_accelerate_version_greater,
@@ -53,7 +52,6 @@
     skip_mps,
     torch_device,
 )
-from diffusers.utils.source_code_parsing_utils import ReturnNameVisitor
 
 from ..models.autoencoders.vae import (
     get_asym_autoencoder_kl_config,