Gyaneshere commited on
Commit
ab6964a
·
verified ·
1 Parent(s): 9bae9c8

Create tool.py

Browse files
Files changed (1) hide show
  1. tool.py +15 -0
tool.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from smolagents import Tool
2
+ from huggingface_hub import InferenceClient
3
+
4
+
5
+ class TextToImageTool(Tool):
6
+ description = "This tool creates an image according to a prompt, which is a text description."
7
+ name = "image_generator"
8
+ 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."}}
9
+ output_type = "image"
10
+ model_sdxl = "black-forest-labs/FLUX.1-schnell"
11
+ client = InferenceClient(model_sdxl)
12
+
13
+
14
+ def forward(self, prompt):
15
+ return self.client.text_to_image(prompt)