from smolagents import CodeAgent,Tool,HfApiModel,load_tool
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
from PIL import Image
import yaml,os
import gradio as gr

final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
)

comics_generation_tool = Tool.from_space(
    #"jbilcke-hf/ai-comic-factory",
    "black-forest-labs/FLUX.1-schnell",
    name="ai_comics_generator",
    description="Generate comics from a prompt"
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[final_answer, comics_generation_tool],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

def final_answer(image_path):
    return Image.open(image_path)

def generate_image(prompt):
    result = agent.run(prompt)
    if not result or not os.path.exists(result):
        return "Failed to generate image."
    return final_answer(result)

with gr.Blocks() as demo:
    prompt_input = gr.Textbox(label="Describe your dream")
    image_output = gr.Image(label="Generated Image", type="pil")
    generate_button = gr.Button("Generate Image")
    generate_button.click(generate_image, inputs=prompt_input, outputs=image_output)

demo.launch()