| from smolagents import Tool |
| from huggingface_hub import InferenceClient |
|
|
|
|
| class TextToImageTool(Tool): |
| description = "This tool creates an image according to a prompt, which is a text description." |
| name = "image_generator" |
| inputs = {"prompt": {"type": "string", "description": "The image generator prompt. Don't hesitate to add details in the prompt to make the image look better, like 'high-res, photorealistic', etc."}} |
| output_type = "image" |
| model_sdxl = "black-forest-labs/FLUX.1-schnell" |
| client = InferenceClient(model_sdxl) |
|
|
|
|
| def forward(self, prompt): |
| return self.client.text_to_image(prompt) |
|
|
