gradio_demo.py

import gradio as gr
from src.agents.mask_generation_agent import mask_generation_agent, ImageEditDeps
import os
from src.hopter.client import Hopter, Environment
from src.services.generate_mask import GenerateMaskService
from dotenv import load_dotenv
from pydantic_ai.messages import (
    ToolReturnPart
)
from src.agents.mask_generation_agent import EditImageResult
from src.utils import upload_image
load_dotenv()

async def process_edit(image, instruction):
    hopter = Hopter(os.environ.get("HOPTER_API_KEY"), environment=Environment.STAGING)
    mask_service = GenerateMaskService(hopter=hopter)
    image_url = upload_image(image)
    messages = [
        {
            "type": "text",
            "text": instruction
        },
    ]
    if image:
        messages.append(
            {"type": "image_url", "image_url": {"url": image_url}}
        )
    deps = ImageEditDeps(
        edit_instruction=instruction,
        image_url=image_url,
        hopter_client=hopter,
        mask_service=mask_service
    )
    result = await mask_generation_agent.run(
        messages,
        deps=deps
    )
    # Extract the edited image URL from the tool return
    for message in result.new_messages():
        for part in message.parts:
            if isinstance(part, ToolReturnPart) and isinstance(part.content, EditImageResult):
                return part.content.edited_image_url
    return None

async def use_edited_image(edited_image):
    return edited_image

def clear_instruction():
    # Only clear the instruction text.
    return ""

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# PicEdit")
    gr.Markdown("""
    Welcome to PicEdit - an AI-powered image editing tool. 
    Simply upload an image and describe the changes you want to make in natural language.
    """)
    
    with gr.Row():
        # Input image on the left
        input_image = gr.Image(label="Original Image", type="filepath")
        
        with gr.Column():
            # Output image on the right
            output_image = gr.Image(label="Edited Image", type="filepath", interactive=False, scale=3)
            use_edited_btn = gr.Button("👈 Use Edited Image 👈")

    # Text input for editing instructions
    instruction = gr.Textbox(
        label="Editing Instructions",
        placeholder="Describe the changes you want to make to the image..."
    )
    
    # Clear button
    with gr.Row():
        clear_btn = gr.Button("Clear")
        submit_btn = gr.Button("Apply Edit", variant="primary")
    
    # Set up the event handlers
    submit_btn.click(
        fn=process_edit,
        inputs=[input_image, instruction],
        outputs=output_image
    )
    
    use_edited_btn.click(
        fn=use_edited_image,
        inputs=[output_image],
        outputs=[input_image]
    )

    # Bind the clear button's click event to only clear the instruction textbox.
    clear_btn.click(
        fn=clear_instruction,
        inputs=[],
        outputs=[instruction]
    )

    examples = gr.Examples(
        examples=[
            ["https://i.ibb.co/qYwhcc6j/c837c212afbf.jpg", "remove the pole"],
            ["https://i.ibb.co/2Mrxztw/image.png", "replace the cat with a dog"],
            ["https://i.ibb.co/9mT4cvnt/resized-78-B40-C09-1037-4-DD3-9-F48-D73637-EE4-E51.png", "ENHANCE!"]
        ],
        inputs=[input_image, instruction]
    )

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