Update app.py

app.py

@@ -40,6 +40,46 @@ def get_random_joke() -> str:
     data = response.json()
     return f"{data.get('setup')} - {data.get('punchline')}"
 
+@tool
+def generate_flux_image(prompt: str, width: int = 1024, height: int = 1024, guidance_scale: float = 3.5, num_inference_steps: int = 28) -> str:
+    """Generates an image using FLUX.1 text-to-image model.
+    
+    Args:
+        prompt: Text description of the image to generate
+        width: Width of the generated image (default: 1024)
+        height: Height of the generated image (default: 1024)
+        guidance_scale: How closely the image should follow the prompt (default: 3.5)
+        num_inference_steps: Number of denoising steps (default: 28)
+    """
+    try:
+        from gradio_client import Client
+        import tempfile
+        import os
+        
+        # Create a client for the FLUX model
+        client = Client("black-forest-labs/FLUX.1-dev")
+        
+        # Call the model to generate an image
+        result = client.predict(
+            prompt=prompt,
+            seed=0,
+            randomize_seed=True,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            api_name="/infer"
+        )
+        
+        # The result is typically a path to an image
+        image_path = result
+        
+        # You could return the path or handle the image as needed
+        return f"Image successfully generated based on prompt: '{prompt}'. Image path: {image_path}"
+        
+    except Exception as e:
+        return f"Error generating image: {str(e)}"
+
 final_answer = FinalAnswerTool()
 
 # If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
@@ -55,14 +95,14 @@ custom_role_conversions=None,
 
 
 # Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+#image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer,get_current_time_in_timezone,get_random_joke,image_generation_tool], ## add your tools here (don't remove final answer)
+    tools=[final_answer,get_current_time_in_timezone,get_random_joke,generate_flux_image], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,