from smolagents.tools import Tool

from PIL.Image import Image


class ImageGeneratorTool(Tool):
    name = "image_generator"
    description = "Generates an image based on your query."
    inputs = {
        "query": {
            "type": "string",
            "description": "The query to generate an image for.",
        }
    }
    output_type = "any"

    def __init__(self, **kwargs):
        super().__init__()
        try:
            from huggingface_hub import InferenceClient
        except ImportError as e:
            raise ImportError(
                "You must install package `huggingface_hub` to run this tool: for instance run `pip install huggingface_hub`."
            ) from e
        self.client = InferenceClient("black-forest-labs/FLUX.1-dev")

    def forward(self, query: str) -> Image:
        image: Image = self.client.text_to_image(query)
        if image is None:
            raise Exception("No results found! Try a less restrictive/shorter query.")
        return image