Update app.py

app.py

@@ -91,7 +91,7 @@ class TextToImageTool(Tool):
     def forward(self, prompt):
         return self.client.text_to_image(prompt)
 '''
-
+ '''
 class TextToImageTool(Tool):
     description = "This tool creates an image according to a prompt. Add details like 'high-res, photorealistic'."
     name = "image_generator"
@@ -116,12 +116,19 @@ class TextToImageTool(Tool):
     def forward(self, prompt):
         image = self.pipe(prompt).images[0]
         return image
+'''
+
+image_generation_tool = Tool.from_space(
+    "black-forest-labs/FLUX.1-schnell",
+    name="image_generator",
+    description="Generate an image from a prompt"
+)
 
 
 # =========================================================
 # Tool and Agent Initialization
 # =========================================================
-image_generation_tool= TextToImageTool()
+#image_generation_tool= TextToImageTool()
 #image_generation_tool = WrappedTextToImageTool()
 search_tool = DuckDuckGoSearchTool()
 llm_engine = InferenceClientModel("Qwen/Qwen2.5-72B-Instruct")