Image_generator / app.py
Naveen671's picture
Update app.py
04be80a verified
from smolagents import CodeAgent, HfApiModel, load_tool
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
# Initialize the mandatory final answer tool
final_answer = FinalAnswerTool()
# Set up the AI model connection
model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
)
# Import the image generation tool from the Hugging Face Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
# Load the prompt templates from your prompts.yaml file
with open("prompts.yaml", 'r') as stream:
prompt_templates = yaml.safe_load(stream)
# Create the agent and provide it with the correct tools
agent = CodeAgent(
model=model,
# Ensure both the image tool and final answer tool are in this list
tools=[image_generation_tool, final_answer],
max_steps=6,
verbosity_level=1,
prompt_templates=prompt_templates
)
# Launch the user interface
GradioUI(agent).launch()