Update app.py

app.py

@@ -4,6 +4,8 @@ import gradio as gr
 from smolagents import CodeAgent, InferenceClientModel
 from smolagents import DuckDuckGoSearchTool, Tool
 from huggingface_hub import InferenceClient
+from diffusers import DiffusionPipeline
+import torch
 
 # =========================================================
 # Utility functions
@@ -76,7 +78,7 @@ class WrappedTextToImageTool(Tool):
 '''
 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"
@@ -88,6 +90,32 @@ 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"
+    inputs = {
+        "prompt": {
+            "type": "string",
+            "description": "The image generation prompt"
+        }
+    }
+    output_type = "image"
+
+    def __init__(self):
+        super().__init__()
+        dtype = torch.bfloat16
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {device}")
+        self.pipe = DiffusionPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-schnell", 
+            torch_dtype=dtype
+        ).to(device)
+
+    def forward(self, prompt):
+        image = self.pipe(prompt).images[0]
+        return image
 
 
 # =========================================================