image_edit_chat.py

import gradio as gr
from src.agents.image_edit_agent import image_edit_agent, ImageEditDeps, EditImageResult
import os
from src.hopter.client import Hopter, Environment
from src.services.generate_mask import GenerateMaskService
from dotenv import load_dotenv
from src.utils import upload_image
from pydantic_ai.messages import ToolCallPart, ToolReturnPart
from pydantic_ai.models.openai import OpenAIModel

model = OpenAIModel(
    "gpt-4o",
    api_key=os.environ.get("OPENAI_API_KEY"),
)

INTRO = """
# Image Editing Assistant

### Experience seamless image editing using natural language in a chat-based interface.

With this demo, you can:
- Enhance or upscale an image
- Remove objects from an image
- Replace elements within an image
- Change the background
"""

EXAMPLES = [
    {
        "text": "Replace the background to the space with stars and planets",
        "files": [
            "https://cdn.prod.website-files.com/66f230993926deadc0ac3a44/66f370d65f158cbbcfbcc532_Crossed%20Arms%20Levi%20Meir%20Clancy.jpg"
        ],
    },
    {
        "text": "Change all the balloons to red in the image",
        "files": [
            "https://www.apple.com/tv-pr/articles/2024/10/apple-tv-unveils-severance-season-two-teaser-ahead-of-the-highly-anticipated-return-of-the-emmy-and-peabody-award-winning-phenomenon/images/big-image/big-image-01/1023024_Severance_Season_Two_Official_Trailer_Big_Image_01_big_image_post.jpg.large_2x.jpg"
        ],
    },
    {
        "text": "Change coffee to a glass of water",
        "files": [
            "https://previews.123rf.com/images/vadymvdrobot/vadymvdrobot1812/vadymvdrobot181201149/113217373-image-of-smiling-woman-holding-takeaway-coffee-in-paper-cup-and-taking-selfie-while-walking-through.jpg"
        ],
    },
    {
        "text": "ENHANCE!",
        "files": [
            "https://m.media-amazon.com/images/M/MV5BNzM3ODc5NzEtNzJkOC00MDM4LWI0MTYtZTkyNmY3ZTBhYzkxXkEyXkFqcGc@._V1_QL75_UX1000_CR0,52,1000,563_.jpg"
        ],
    },
]

load_dotenv()


def build_user_message(chat_input):
    text = chat_input["text"]
    images = chat_input["files"]
    messages = [{"role": "user", "content": text}]
    if images:
        messages.extend(
            [{"role": "user", "content": {"path": image}} for image in images]
        )
    return messages


def build_messages_for_agent(chat_input, past_messages):
    # filter out image messages from past messages to save on tokens
    messages = past_messages

    # add the user's text message
    if chat_input["text"]:
        messages.append({"type": "text", "text": chat_input["text"]})

    # add the user's image message
    files = chat_input.get("files", [])
    image_url = upload_image(files[0]) if files else None
    if image_url:
        messages.append({"type": "image_url", "image_url": {"url": image_url}})

    return messages


def select_example(x: gr.SelectData, chat_input):
    chat_input["text"] = x.value["text"]
    chat_input["files"] = x.value["files"]
    return chat_input


async def stream_from_agent(chat_input, chatbot, past_messages, current_image):
    # Prepare messages for the UI
    chatbot.extend(build_user_message(chat_input))
    yield {"text": "", "files": []}, chatbot, gr.skip, gr.skip()

    # Prepare messages for the agent
    text = chat_input["text"]
    files = chat_input.get("files", [])
    image_url = upload_image(files[0]) if files else None
    messages = [
        {"type": "text", "text": text},
    ]
    if image_url:
        messages.append({"type": "image_url", "image_url": {"url": image_url}})
        current_image = image_url

    # Dependencies
    hopter = Hopter(os.environ.get("HOPTER_API_KEY"), environment=Environment.STAGING)
    mask_service = GenerateMaskService(hopter=hopter)
    deps = ImageEditDeps(
        edit_instruction=text,
        image_url=current_image,
        hopter_client=hopter,
        mask_service=mask_service,
    )

    # Run the agent
    async with image_edit_agent.run_stream(
        messages, deps=deps, message_history=past_messages
    ) as result:
        for message in result.new_messages():
            for call in message.parts:
                if isinstance(call, ToolCallPart):
                    call_args = (
                        call.args.args_json
                        if hasattr(call.args, "args_json")
                        else call.args
                    )
                    metadata = {
                        "title": f"🛠️ Using {call.tool_name}",
                    }
                    # set the tool call id so that when the tool returns
                    # we can find this message and update with the result
                    if call.tool_call_id is not None:
                        metadata["id"] = call.tool_call_id

                    # Create a tool call message to show on the UI
                    gr_message = {
                        "role": "assistant",
                        "content": "Parameters: " + call_args,
                        "metadata": metadata,
                    }
                    chatbot.append(gr_message)
                if isinstance(call, ToolReturnPart):
                    for gr_message in chatbot:
                        # Skip messages without metadata
                        if not gr_message.get("metadata"):
                            continue

                        if gr_message["metadata"].get("id", "") == call.tool_call_id:
                            if isinstance(call.content, EditImageResult):
                                chatbot.append(
                                    {
                                        "role": "assistant",
                                        "content": gr.Image(
                                            call.content.edited_image_url
                                        ),
                                        "files": [call.content.edited_image_url],
                                    }
                                )
                                current_image = call.content.edited_image_url
                            else:
                                gr_message["content"] += f"\nOutput: {call.content}"
                yield gr.skip(), chatbot, gr.skip(), gr.skip()

        chatbot.append({"role": "assistant", "content": ""})
        async for message in result.stream_text():
            chatbot[-1]["content"] = message
            yield gr.skip(), chatbot, gr.skip(), gr.skip()
        past_messages = result.all_messages()

        yield gr.Textbox(interactive=True), gr.skip(), past_messages, current_image


with gr.Blocks() as demo:
    gr.Markdown(INTRO)

    current_image = gr.State(None)
    past_messages = gr.State([])
    chatbot = gr.Chatbot(
        elem_id="chatbot",
        label="Image Editing Assistant",
        type="messages",
        avatar_images=(None, "https://ai.pydantic.dev/img/logo-white.svg"),
        examples=EXAMPLES,
    )

    with gr.Row():
        chat_input = gr.MultimodalTextbox(
            interactive=True,
            file_count="single",
            show_label=False,
            placeholder="How would you like to edit this image?",
            sources=["upload"],
        )
    generation = chat_input.submit(
        stream_from_agent,
        inputs=[chat_input, chatbot, past_messages, current_image],
        outputs=[chat_input, chatbot, past_messages, current_image],
    )

    chatbot.example_select(
        select_example,
        inputs=[chat_input],
        outputs=[chat_input],
    ).then(
        stream_from_agent,
        inputs=[chat_input, chatbot, past_messages, current_image],
        outputs=[chat_input, chatbot, past_messages, current_image],
    )

    examples = gr.Examples(
        examples=EXAMPLES,
        inputs=[chat_input],
        outputs=[chat_input],
    )


if __name__ == "__main__":
    demo.launch()