import subprocess
import sys

# Install required packages
packages = {
    "diffusers": "0.30.2",
    "torch": "2.6.0",
    "gradio": "5.17.1",
    "transformers": "4.49.0"
}

def install_packages():
    for package, version in packages.items():
        subprocess.check_call([sys.executable, "-m", "pip", "install", f"{package}=={version}"])

# Install packages before imports
install_packages()

import gradio as gr
from control import main
from diffusers.utils import load_image
from mask import create_mask

def process_image(image, prompt):
    try:
        # Create mask from input image
        mask = create_mask(image)
        
        # First show the generated mask
        yield mask
        
        # Then process image with mask
        result = main(image, mask, prompt)
        yield result
        
    except Exception as e:
        yield str(e)

# Create Gradio interface
demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(label="Input Image", type="pil"),
        gr.Textbox(label="Prompt"),
    ],
    outputs=[
        gr.Image(label="Generated Mask"),
        gr.Image(label="Generated Image")
    ],
    title="Image Inpainting with FLUX ControlNet",
    description="Upload an image and provide a prompt. The system will first generate a mask and then create the inpainted result.",
)

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