from utils.model import load_seg, load_inpainting, generate_with_mask, generate
from utils.scraper import extract_link
import gradio as gr

extractor, model = load_seg()
prompt_pipe = load_inpainting(using_prompt = True, fast=True)
cloth_pipe = load_inpainting(fast=True)

def generate_with_mask_(image_path: str, cloth_path: str = None, prompt: str = None):
    """
    Generate Image.

    Request Body
    request = {
        "image" : Input Image URL
        "cloth" : Cloth Image URL
        "prompt" : Prompt, In case example image is not provided
    }

    Return Body:
    {
    gen: Generated Image
    }
    """
    using_prompt = True if prompt else False
    image_url = extract_link(image_path)
    cloth_url = extract_link(cloth_path)
    image_path = image_url if image_url else image_path
    cloth_path = cloth_url if cloth_url else cloth_path
    if using_prompt:
        gen = generate(image_path, extractor, model, prompt_pipe, cloth_path, prompt)
    else:
        gen = generate_with_mask(image_path, extractor, model, cloth_pipe, cloth_path, prompt)
    return gen


with gr.Blocks() as demo:
    image = gr.inputs.Image()
    cloth = gr.inputs.Image()
    prompt = gr.inputs.Textbox(lines=5, label="Editing Prompt")
    output = gr.outputs.Image(label="Generated Image")
    gr.Interface(generate_with_mask_, inputs=[image, cloth, prompt], outputs=output).launch()

demo.launch()