updated app

app.py

@@ -1,29 +1,25 @@
-import os, sys
+import gradio as gr
+from rome.infer import Infer
+import pickle
+from easydict import EasyDict as edict
+from huggingface_hub import hf_hub_download
+from rome.src.utils.processing import process_black_shape, tensor2image
+from rome.src.utils import args as args_utils
+import sys
 import torch
-import argparse
-
-import numpy as np
 import torch
-import matplotlib.pyplot as plt
-from PIL import Image
-
-print(torch.__version__)
-print(torch.version.cuda)
 
 sys.path.append("./rome/")
 sys.path.append('./DECA')
 
-from rome.src.utils import args as args_utils
-from rome.src.utils.processing import process_black_shape, tensor2image
 
 # loading models ---- create model repo
-from huggingface_hub import hf_hub_download
 
-default_modnet_path = hf_hub_download('Pie31415/rome','modnet_photographic_portrait_matting.ckpt')
-default_model_path = hf_hub_download('Pie31415/rome','rome.pth')
+default_modnet_path = hf_hub_download(
+    'Pie31415/rome', 'modnet_photographic_portrait_matting.ckpt')
+default_model_path = hf_hub_download('Pie31415/rome', 'rome.pth')
 
 # parser configurations
-from easydict import EasyDict as edict
 
 args = edict({
     "save_dir": ".",
@@ -93,8 +89,8 @@ args = edict({
     "num_vertex": 5023,
     "train_basis": True,
     "path_to_deca": "DECA",
-    "path_to_linear_hair_model": "data/linear_hair.pth", # N/A
-    "path_to_mobile_model": "data/disp_model.pth", # N/A
+    "path_to_linear_hair_model": "data/linear_hair.pth",  # N/A
+    "path_to_mobile_model": "data/disp_model.pth",  # N/A
     "n_scalp": 60,
     "use_distill": False,
     "use_mobile_version": False,
@@ -109,40 +105,49 @@ args = edict({
 })
 
 # download FLAME and DECA pretrained
-generic_model_path = hf_hub_download('Pie31415/rome','generic_model.pkl')
-deca_model_path = hf_hub_download('Pie31415/rome','deca_model.tar')
-
-import pickle
+generic_model_path = hf_hub_download('Pie31415/rome', 'generic_model.pkl')
+deca_model_path = hf_hub_download('Pie31415/rome', 'deca_model.tar')
 
 with open(generic_model_path, 'rb') as f:
-  ss = pickle.load(f, encoding='latin1')
+    ss = pickle.load(f, encoding='latin1')
 
-  with open('./DECA/data/generic_model.pkl', 'wb') as out:
-    pickle.dump(ss, out)
+    with open('./DECA/data/generic_model.pkl', 'wb') as out:
+        pickle.dump(ss, out)
 
 with open(deca_model_path, "rb") as input:
-  with open('./DECA/data/deca_model.tar', "wb") as out:
-    for line in input:
-      out.write(line)
+    with open('./DECA/data/deca_model.tar', "wb") as out:
+        for line in input:
+            out.write(line)
 
 # load ROME inference model
-from rome.infer import Infer
 infer = Infer(args)
 
-def predict(source_img, driver_img):
+def image_inference(
+    source_img: gr.inputs.Image = None,
+    driver_img: gr.inputs.Image = None
+):
     out = infer.evaluate(source_img, driver_img, crop_center=False)
     res = tensor2image(torch.cat([out['source_information']['data_dict']['source_img'][0].cpu(),
-                              out['source_information']['data_dict']['target_img'][0].cpu(),
-                        out['render_masked'].cpu(), out['pred_target_shape_img'][0].cpu()], dim=2))
+                                  out['source_information']['data_dict']['target_img'][0].cpu(
+    ),
+        out['render_masked'].cpu(), out['pred_target_shape_img'][0].cpu()], dim=2))
     return res[..., ::-1]
 
-import gradio as gr
+def folder_inference():
+    pass
+
+with gr.Blocks() as demo:
+    with gr.Tab("Image Inference"):
+        image_input = [gr.Image(type="pil"), gr.Image(type="pil")]
+        image_output = gr.Image()
+        image_button = gr.Button("Predict")
+    with gr.Tab("Inference Over Folder"):
+        pass
+    with gr.Tab("Video Inference"):
+        pass
+
+    image_button.click(image_inference, inputs=image_input, outputs=image_output)
+    title = "ROME: Realistic one-shot mesh-based head avatars"
+    examples = gr.Examples(["examples/lincoln.jpg", "examples/tars2.jpg"])
 
-gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Image(type="pil"),
-        gr.Image(type="pil")
-    ],
-    outputs=gr.Image(),
-    examples=[]).launch()
+demo.launch()