Casa-AI / app.py
LN1996's picture
Update app.py
35a15c6 verified
raw
history blame
5.42 kB
import modal
import gradio as gr
import numpy as np
f = modal.Cls.lookup("casa-interior-hf-v3", "DesignModel")
f_gc = modal.Cls.lookup("casa-interior-gc-v1", "GetProduct")
# def casa_ai_run_tab1(image=None, room_type=None, interior_style=None, objects=None):
# if image is None:
# print('Please provide image of empty room to design')
# return None
# if (room_type is None or interior_style is None or objects is None):
# print('Please select all three (Room type, Interior Style and Objects')
# return None
# objects_str = ", ".join(objects)
# text = f"A {room_type} in style of {interior_style} with {objects_str}."
# result_image = f.inference.remote("tab1", image, text)
# return result_image
def casa_ai_run_tab1(image=None, text=None):
if image is None:
print('Please provide image of empty room to design')
return None
if text is None:
print('Please provide a text prompt')
return None
result_image = f.inference.remote("tab1", image, text)
return result_image
def casa_ai_run_tab2(dict=None, text=None):
image = dict["background"].convert("RGB")
mask = dict["layers"][0].convert('L')
if np.sum(np.array(mask)) == 0:
mask = None
if mask is None:
print('Please provide a mask over the object you want to generate again.')
if image is None and text is None:
print('Please provide context in form of image, text')
return None
result_image = f.inference.remote("tab2", image, text, mask)
return result_image
def casa_ai_run_tab3(dict=None):
selected_crop = dict["layers"][0].convert("RGB")
if selected_crop is None:
print('Please provide cropped object')
return None
result_image = f_gc.inference.remote(selected_crop)
return result_image
with gr.Blocks() as casa:
title = "Casa-AI Demo"
description = "A Gradio interface to use CasaAI for virtual staging"
with gr.Tab("Reimagine"):
with gr.Row():
with gr.Column():
inputs = [
gr.Image(sources='upload', type="pil", label="Upload"),
gr.Textbox(label="Room description.")
# gr.Dropdown(
# label="Room type",
# choices=["Living Room", "Bedroom", "Home Office", "Dining Room", "Kitchen", "Bathroom"],
# multiselect=False
# ),
# gr.Dropdown(
# label="Interior Style",
# choices=["Modern", "Traditional", "Contemporary", "Minimalist", "Rustic", "Industrial", "Scandinavian"],
# multiselect=False
# ),
# gr.Dropdown(
# label="Objects",
# choices=["Sofa", "Bed", "Table", "Chair", "Fridge", "Dinning table", "Coffee table"],
# multiselect=True
# ),
]
with gr.Column():
outputs = [gr.Image(label="Generated room image")]
submit_btn = gr.Button("Generate!")
submit_btn.click(casa_ai_run_tab1, inputs=inputs, outputs=outputs)
with gr.Tab("Redesign"):
with gr.Row():
with gr.Column():
inputs = [
gr.ImageEditor(sources='upload', brush=gr.Brush(colors=["#FFFFFF"]), elem_id="image_upload", type="pil", label="Upload", layers=False, eraser=True, transforms=[]),
gr.Textbox(label="Description for redesigning masked object")]
with gr.Column():
outputs = [gr.Image(label="Image with new designed object")]
submit_btn = gr.Button("Redesign!")
submit_btn.click(casa_ai_run_tab2, inputs=inputs, outputs=outputs)
with gr.Tab("Recommendation"):
with gr.Row():
with gr.Column():
inputs = [
gr.ImageEditor(sources='upload', elem_id="image_upload", type="pil", label="Upload", layers=False, eraser=False, brush=False, transforms=['crop']),
]
with gr.Column():
outputs = [gr.Gallery(label="Similar products")]
submit_btn = gr.Button("Find similar products!")
submit_btn.click(casa_ai_run_tab3, inputs=inputs, outputs=outputs)
casa.launch()
# demo = gr.Interface(casa_ai_run,
# inputs = [gr.ImageEditor(sources='upload', brush=gr.Brush(colors=["#FFFFFF"]), elem_id="image_upload", type="pil", label="Upload", layers=False, eraser=True, transforms=[]),
# gr.Textbox(label="Prompt to design room"),
# ],
# outputs = [
# gr.Image(label="Generated room image"),
# ],
# title = title,
# description = description,
# )
# demo.launch()