add image_generator tool

.gitignore

@@ -0,0 +1,3 @@
+.venv
+__pycache__
+.gradio

app.py

@@ -6,6 +6,7 @@ import yaml
 from tools.final_answer import FinalAnswerTool
 from tools.visit_webpage import VisitWebpageTool
 from tools.web_search import DuckDuckGoSearchTool
+from tools.image_generator import ImageGeneratorTool
 
 from Gradio_UI import GradioUI
 
@@ -59,7 +60,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[visit_webpage, web_search, final_answer], ## add your tools here (don't remove final answer)
+    tools=[visit_webpage, web_search, final_answer, image_generation_tool], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,

tools/image_generator.py

@@ -0,0 +1,31 @@
+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