Update app.py

app.py

@@ -2,10 +2,11 @@ from PIL import Image, ImageDraw, ImageFont
 import tempfile
 import gradio as gr
 from smolagents import CodeAgent, InferenceClientModel
-from smolagents import DuckDuckGoSearchTool, Tool
+from smolagents import DuckDuckGoSearchTool #Tool
 from huggingface_hub import InferenceClient
 from diffusers import DiffusionPipeline
 import torch
+from autoagents.core import Tool 
 
 # =========================================================
 # Utility functions
@@ -91,7 +92,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,7 +117,32 @@ class TextToImageTool(Tool):
     def forward(self, prompt):
         image = self.pipe(prompt).images[0]
         return image
+'''
+
+# Use a public model (you can swap for a gated model with token later)
+pipe = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1",
+    torch_dtype=torch.float16
+).to("cuda" if torch.cuda.is_available() else "cpu")
+
+# --------------------------------------------------------
+# Define a simple image generation function
+# --------------------------------------------------------
+
+def generate_image(prompt: str) -> Image.Image:
+    """Generate an image from a text prompt using a diffusion model."""
+    image = pipe(prompt).images[0]
+    return image
+
+# --------------------------------------------------------
+# Wrap it into a Tool using from_function
+# --------------------------------------------------------
 
+image_generation_tool = Tool.from_function(
+    name="image_generator",
+    description="Generate an image from a prompt, e.g. 'a dog on Mars in high-res'.",
+    func=generate_image
+)
 
 # =========================================================
 # Tool and Agent Initialization