import gradio as gr
from PIL import Image
from models.segmentation import SamSegmentationModel
from models.inpainting import KandingskyInpaintingModel
from models.product import ProductBackgroundModifier
import torch

def generate(image: Image.Image, prompt: str):
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model = ProductBackgroundModifier(
        segmentation_model=SamSegmentationModel(
            model_type="vit_h", 
            checkpoint_path="model_checkpoints/sam_vit.pth", 
            device=device,
        ),
        inpainting_model=KandingskyInpaintingModel(),
        device=device
    )
    generated = model.generate(image=image, prompt=prompt)
    return generated

gr.Interface(
    fn=generate, 
    inputs=[
        gr.Image(type="pil"),
        gr.Text()
    ], 
    outputs=gr.Image(type="pil"),
).launch()