Skip to content

Commit

Permalink
Add __call__ for imagesynthesisnode (#67)
Browse files Browse the repository at this point in the history
  • Loading branch information
qbc2016 authored Nov 15, 2024
1 parent 19b168a commit a72c4cd
Showing 1 changed file with 13 additions and 7 deletions.
20 changes: 13 additions & 7 deletions src/agentscope/web/workstation/workflow_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -860,7 +860,10 @@ class ImageSynthesisNode(WorkflowNode):

def _post_init(self) -> None:
super()._post_init()
self.pipeline = partial(image_synthesis, **self.opt_kwargs)
self.function_executor = partial(image_synthesis, **self.opt_kwargs)

def __call__(self, x: dict = None) -> dict:
return self.function_executor(x)

def compile(self) -> dict:
return {
Expand All @@ -883,12 +886,15 @@ class ImageCompositionNode(WorkflowNode):

def _post_init(self) -> None:
super()._post_init()
self.pipeline = partial(stitch_images_with_grid, **self.opt_kwargs)
self.function_executor = partial(
stitch_images_with_grid,
**self.opt_kwargs,
)

def __call__(self, x: list = None) -> dict:
if isinstance(x, dict):
x = [x]
return self.pipeline(x)
return self.function_executor(x)

def compile(self) -> dict:
return {
Expand All @@ -911,13 +917,13 @@ class ImageMotionNode(WorkflowNode):

def _post_init(self) -> None:
super()._post_init()
self.pipeline = partial(
self.function_executor = partial(
create_video_or_gif_from_image,
**self.opt_kwargs,
)

def __call__(self, x: dict = None) -> dict:
return self.pipeline(x)
return self.function_executor(x)

def compile(self) -> dict:
return {
Expand All @@ -941,10 +947,10 @@ class VideoCompositionNode(WorkflowNode):

def _post_init(self) -> None:
super()._post_init()
self.pipeline = partial(merge_videos, **self.opt_kwargs)
self.function_executor = partial(merge_videos, **self.opt_kwargs)

def __call__(self, x: dict = None) -> dict:
return self.pipeline(x)
return self.function_executor(x)

def compile(self) -> dict:
return {
Expand Down

0 comments on commit a72c4cd

Please sign in to comment.