Update app.py

Added image gen

app.py

@@ -4,19 +4,49 @@ import requests
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
-
+from diffusers import StableDiffusionPipeline
+import torch
+from io import BytesIO
+import base64
 from Gradio_UI import GradioUI
 
+class ImageGenerator:
+    def __init__(self, model_id="runwayml/stable-diffusion-v1-5", device="cuda" if torch.cuda.is_available() else "cpu"):
+        self.pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
+        self.device = device
+
+    def generate_image(self, prompt, num_inference_steps=25, guidance_scale=7.5):
+        """Generates an image from a text prompt."""
+        image = self.pipeline(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
+        return image
+
+    def generate_base64_image(self, prompt, num_inference_steps=25, guidance_scale=7.5):
+        """Generates a base64 encoded image from a text prompt."""
+        image = self.generate_image(prompt, num_inference_steps, guidance_scale)
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        return img_str
+
+def generate_image_tool(image_generator):
+    """Creates a tool function for image generation."""
+    def image_generation_tool(prompt):
+        """Generates an image from a prompt."""
+        return image_generator.generate_base64_image(prompt)
+    return image_generation_tool
+
+# Initialize the ImageGenerator and tool
+image_generator = ImageGenerator()
+image_generation_tool_function = generate_image_tool(image_generator)
+
 # Below is an example of a tool that does nothing. Amaze us with your creativity !
 @tool
-def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type
-    #Keep this format for the description / args / args description but feel free to modify the tool
-    """A tool that does nothing yet 
+def generate_image_from_prompt(prompt: str) -> str:
+    """Generates an image from a text prompt.
     Args:
-        arg1: the first argument
-        arg2: the second argument
+        prompt: The text prompt to generate the image from.
     """
-    return "What magic will you build ?"
+    return image_generation_tool_function(prompt)
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str: