import gradio as gr
import cv2
import torch
import numpy as np
from torchvision import transforms
import torch
from pytorch_lightning import seed_everything
from torchvision.utils import save_image
from model_lib.modules import MoMA_main_modal
from model_lib.utils import parse_args
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"

title = "MoMA"
description = "This model has to run on GPU"
article = "<p style='text-align: center'><a href='https://news.machinelearning.sg/posts/beautiful_profile_pics_remove_background_image_with_deeplabv3/'>Blog</a> | <a href='https://github.com/eugenesiow/practical-ml'>Github Repo</a></p>"

def MoMA_demo(rgb, mask, subject, prompt):
    # move the input and model to GPU for speed if available
    with torch.no_grad():
        generated_image = model.generate_images(rgb, mask, subject, prompt, strength=1.0, seed=2)
    return generated_image

def inference(rgb, mask, subject, prompt):
    result = MoMA_demo(rgb, mask, subject, prompt)
    return result

seed_everything(0)
args = parse_args()
#load MoMA from HuggingFace. Auto download
model = MoMA_main_modal(args).to(args.device, dtype=torch.float16)


################ change texture ##################
# prompt = "A wooden sculpture of a car on the table."
# generated_image = model.generate_images(rgb_path, mask_path, subject, prompt, strength=0.4, seed=4, return_mask=True)  # set strength to 0.4 for better prompt fidelity
# save_image(generated_image,f"{args.output_path}/{subject}_{prompt}.jpg")


gr.Interface(
    inference,
    [gr.Image(type="pil", label="Input RGB"),
     gr.Image(type="pil", label="Input Mask"),
     gr.Textbox(lines=1, label="subject"),
     gr.Textbox(lines=5, label="Prompt")],
    gr.Image(type="pil", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[["example_images/newImages/3.jpg",'example_images/newImages/3_mask.jpg','car','A car in autumn with falling leaves.']],
    # enable_queue=True
).launch(debug=False)