Enable image output for Tool.from_space by aymeric-roucher · Pull Request #1510 · huggingface/smolagents (original) (raw)

if isinstance(output, tuple) or isinstance(output, list):

return output[

if isinstance(output[1], str):

raise ValueError("The space returned this message: " + output[1])

output = output[

0

] # Sometime the space also returns the generation seed, in which case the result is at index 0

IMAGE_EXTENTIONS = [".png", ".jpg", ".jpeg", ".gif", ".webp"]

AUDIO_EXTENTIONS = [".mp3", ".wav", ".ogg", ".m4a", ".flac"]

if isinstance(output, str) and any([output.endswith(ext) for ext in IMAGE_EXTENTIONS]):

output = AgentImage(output)

elif isinstance(output, str) and any([output.endswith(ext) for ext in AUDIO_EXTENTIONS]):

output = AgentAudio(output)

return output