import modal

f = modal.Cls.lookup("casa-interior-hf-v1", "DesignModel")

def casa_ai_run(dict=None, text=None):

    image = dict["background"].convert("RGB")
    mask = dict["layers"][0].convert('L')
    
    if image is None and text is None: 
        print('Please provide context in form of image, audio, text')
        return None
    
    result_image = f.inference.remote(image, text)
    return result_image

import gradio as gr 

title = "Casa-AI Demo"
description = "A Gradio interface to use CasaAI for virtual staging"

demo = gr.Interface(casa_ai_run,
                    inputs = [gr.ImageEditor(sources='upload', brush=gr.Brush(colors=["#000000"]), 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()