TryOn / app.py
Warlord-K's picture
Initial Commit
ce9d0da
raw
history blame
1.4 kB
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()