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)
cloth_pipe = load_inpainting()

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:
    gr.Markdown('# Try On Clothes Online!')
    gr.Markdown('## Add Your Image via an Image URL or By Simply Uploading it')
    gr.Markdown('## Paste in a link of a Product from Flipkart, Amazon or Myntra and See the Preview! You can also paste in any Image link!')
    gr.Markdown('## Optionally you can add a prompt to generate an image instead of using an example image.')
    with gr.Row():
        with gr.Column():
            image = gr.inputs.Image(type = "filepath", label = "Input Image")
            # image_url = gr.inputs.Textbox(label = "Input Image URL")
            cloth_url = gr.inputs.Textbox(label = "Cloth Image URL")
            prompt = gr.inputs.Textbox(label="Optional Prompt")
        output = gr.outputs.Image(type = "pil", label="Generated Image")
    run = gr.Button(label="Generate Preview")
    gr.Markdown('## Made By [Yatharth Gupta](https://www.linkedin.com/in/yatharth-g/)')
    run.click(generate_with_mask_, inputs=[image, cloth_url, prompt], outputs=output)

demo.launch()

demo.launch()