MoMA_demo / backupApp_version1.py
Kunpeng Song
random
89092e7
raw
history blame
2 kB
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)