File size: 1,185 Bytes
788ee56 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
from agents import Tool
class SpaceToolWrapper(Tool):
name = "face_swapper"
description = """Tool that puts the face shown on the first image on the second image. It taks paths to images as arguments (can be relative path), or PIL images directly."""
inputs = {"source_file":{"type":"any","description":""},"target_file":{"type":"any","description":""}}
output_type = "image"
def forward(self, *args, **kwargs):
# Preprocess args and kwargs:
args = list(args)
for i, arg in enumerate(args):
args[i] = self.sanitize_argument_for_prediction(arg)
for arg_name, arg in kwargs.items():
kwargs[arg_name] = self.sanitize_argument_for_prediction(arg)
output = self.client.predict(*args, api_name=self.api_name, **kwargs)
if isinstance(output, tuple) or isinstance(output, list):
return output[
0
] # Sometime the space also returns the generation seed, in which case the result is at index 0
return output |