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