import modal
import gradio as gr 


f = modal.Cls.lookup("casa-interior-hf-v2", "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, mask)
    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.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")]
            with gr.Column():
                outputs = [gr.Image(label="Generated room image")]
                
        submit_btn = gr.Button("Submit")
        submit_btn.click(casa_ai_run, inputs=inputs, outputs=outputs)

        
    with gr.Tab("ReDesign"):
        pass

    with gr.Tab("Recommendation"):
        pass

    


# 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,
#                     )
casa.launch()