File size: 1,417 Bytes
e7ea92d
6aae614
9b5b26a
7144088
 
 
9b5b26a
6aae614
ae7a494
e121372
7144088
 
 
13d500a
8c01ffb
8fd45d9
5ee0261
faf2f77
 
8fd45d9
8c01ffb
861422e
 
7144088
6fffb2f
 
7144088
6fffb2f
 
 
 
 
 
 
 
8fe992b
7144088
 
 
 
 
 
 
 
 
 
 
 
 
 
9b5b26a
7144088
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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
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(
    "black-forest-labs/FLUX.1-schnell",
    name="image_generator",
    description="Generate an image 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="Enter Prompt")
    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()