forked from huggingface/diffusers
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsource_code_parsing_utils.py
More file actions
52 lines (45 loc) · 1.82 KB
/
source_code_parsing_utils.py
File metadata and controls
52 lines (45 loc) · 1.82 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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