app.py

from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Подключаем модели
text_model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # Модель для обработки текста
    custom_role_conversions=None,
)

image_model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id="black-forest-labs/FLUX.1-dev",  # Модель для генерации изображений
    custom_role_conversions=None,
)

@tool
def generate_image_from_text(prompt: str) -> str:
    """Инструмент для генерации изображения на основе текстового описания."""
    try:
        # Загружаем инструмент для генерации изображений с Hugging Face
        image_generation_tool = load_tool("multimodalart/flux.1-dev", trust_remote_code=True)

        # Генерация изображения
        generated_image = image_generation_tool(prompt=prompt)
        
        # Возвращаем путь или URL изображения
        return generated_image
    except Exception as e:
        return f"Error generating image: {str(e)}"


 


final_answer = FinalAnswerTool()

 

# Load prompt templates
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
# Инициализируем агента, передавая обе модели
agent = CodeAgent(
    model=text_model,  # Здесь указываем модель для обработки текста
    tools=[FinalAnswerTool(), generate_image_from_text],  # Добавляем инструменты
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=None,  # Указывайте шаблоны при необходимости
)

# Launch the AI agent with Gradio
GradioUI(agent).launch()