Update app.py

app.py

@@ -55,7 +55,7 @@ def generate_prompts_for_object(object_name):
 # =========================================================
 # Tool wrapper for m-ric/text-to-image
 # =========================================================
-
+'''
 class WrappedTextToImageTool(Tool):
     name = "text_to_image"
     description = "Generates an image from a text prompt using the m-ric/text-to-image tool."
@@ -73,11 +73,28 @@ class WrappedTextToImageTool(Tool):
     def forward(self, prompt):
         return self.client.text_to_image(prompt)
 
+'''
+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)
+
+
 # =========================================================
 # Tool and Agent Initialization
 # =========================================================
-
-image_generation_tool = WrappedTextToImageTool()
+image_generation_tool= TextToImageTool()
+#image_generation_tool = WrappedTextToImageTool()
 search_tool = DuckDuckGoSearchTool()
 llm_engine = InferenceClientModel("Qwen/Qwen2.5-72B-Instruct")