update app.py

.gitignore

@@ -17,5 +17,4 @@ dist
 *egg-info
 *.so
 run.sh
-*.log
-gradio_cached_examples/
+*.log

README.md

@@ -1,6 +1,6 @@
 ---
 title:  Unconstrained & Detailed Clothed Human Digitization (ECON + ControlNet)
-metaTitle: Avatarify from Photo
+metaTitle: ECON-Avatarify from Photo
 emoji: 🤼
 colorFrom: green
 colorTo: pink

app.py

@@ -6,21 +6,14 @@ import os
 
 import subprocess
 
-curr_dir = os.path.dirname(__file__)
-
 if os.getenv('SYSTEM') == 'spaces':
     # subprocess.run('pip install pyembree'.split())
     subprocess.run(
         'pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu116_pyt1130/download.html'
         .split()
     )
-    subprocess.run(
-        f"cd {curr_dir}/lib/common/libmesh && python setup.py build_ext --inplace".split()
-    )
-    subprocess.run(
-        f"cd {curr_dir}/lib/common/libvoxelize && python setup.py build_ext --inplace".split()
-    )
-    subprocess.run(f"cd {curr_dir}".split())
+    subprocess.run("python setup.py build_ext --inplace".split(), cwd="./lib/common/libmesh/")
+    subprocess.run("python setup.py build_ext --inplace".split(), cwd="./lib/common/libvoxelize/")
 
 from apps.infer import generate_model, generate_video
 
@@ -134,6 +127,8 @@ async (image_in_img, prompt, image_file_live_opt, live_conditioning) => {
 # Constants
 low_threshold = 100
 high_threshold = 200
+default_step = 50
+cached = False
 
 # Models
 pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
@@ -162,8 +157,8 @@ hint_prompts = '''
 <strong>Hints</strong>: <br>
 best quality, extremely detailed, solid color background, 
 super detail, high detail, edge lighting, soft focus, 
-light and dark contrast, 8k, high detail, edge lighting, 
-3d, c4d, blender, oc renderer, ultra high definition, 3d rendering
+light and dark contrast, 8k, edge lighting, 3d, c4d, 
+blender, oc renderer, ultra high definition, 3d rendering
 '''
 
 
@@ -213,8 +208,6 @@ def toggle(choice):
 examples_pose = glob.glob('examples/pose/*')
 examples_cloth = glob.glob('examples/cloth/*')
 
-default_step = 50
-
 with gr.Blocks() as demo:
     gr.Markdown(description)
 
@@ -255,7 +248,7 @@ with gr.Blocks() as demo:
                     gallery_cache = gr.State()
                     inp = gr.Image(type="filepath", label="Input Image for ECON")
                     fitting_step = gr.inputs.Slider(
-                        10, 100, step=10, label='Fitting steps', default=default_step
+                        10, 100, step=10, label='Fitting steps (Slower yet Better-aligned SMPL-X)', default=default_step
                     )
 
             with gr.Row():
@@ -283,34 +276,34 @@ with gr.Blocks() as demo:
                 gr.Examples(
                     examples=list(examples_pose),
                     inputs=[inp],
-                    cache_examples=False,
+                    cache_examples=cached,
                     fn=generate_model,
                     outputs=out_lst,
-                    label="Hard Pose Exampels"
+                    label="Hard Pose Examples"
                 )
+                
                 gr.Examples(
                     examples=list(examples_cloth),
                     inputs=[inp],
-                    cache_examples=False,
+                    cache_examples=cached,
                     fn=generate_model,
                     outputs=out_lst,
-                    label="Loose Cloth Exampels"
+                    label="Loose Cloth Examples"
                 )
+                
+            out_vid = gr.Video(label="Shared on Twitter with #ECON")
 
         with gr.Column():
-            overlap_inp = gr.Image(type="filepath", label="Image Normal Overlap")
-            with gr.Row():
-                out_final = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Clothed human")
-                out_smpl = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="SMPL-X body")
+            overlap_inp = gr.Image(type="filepath", label="Image Normal Overlap").style(height=400)
+            out_final = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="Clothed human", elem_id="avatar")
+            out_smpl = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="SMPL-X body", elem_id="avatar")
 
             out_final_obj = gr.State()
             vis_tensor_path = gr.State()
 
             with gr.Row():
                 btn_video = gr.Button("Generate Video (~2min)")
-            with gr.Row():
-                out_vid = gr.Video(label="Shared on Twitter with #ECON")
-
+    
     # with gr.Row():
     #     btn_texture = gr.Button("Generate Full-texture")
 
@@ -345,12 +338,13 @@ with gr.Blocks() as demo:
     )
 
     btn_submit.click(fn=generate_model, inputs=[inp, fitting_step], outputs=out_lst)
+    
     # btn_texture.click(
     #     fn=generate_texture,
     #     inputs=[out_final_obj, prompt, seed, guidance_scale],
     #     outputs=[viewpoint_images, result_video, output_file, progress_text]
     # )
-
+    
     demo.load(None, None, None, _js=load_js)
 
 if __name__ == "__main__":
@@ -359,4 +353,5 @@ if __name__ == "__main__":
     #             auth=(os.environ['USER'], os.environ['PASSWORD']),
     #             auth_message="Register at icon.is.tue.mpg.de to get HuggingFace username and password.")
 
+    demo.queue(concurrency_count=1)
     demo.launch(debug=True, enable_queue=True)

apps/infer.py

@@ -28,6 +28,7 @@ import torch
 import torchvision
 import trimesh
 from pytorch3d.ops import SubdivideMeshes
+from huggingface_hub import hf_hub_download
 from termcolor import colored
 from tqdm.auto import tqdm
 
@@ -47,6 +48,7 @@ from lib.net.geometry import rot6d_to_rotmat, rotation_matrix_to_angle_axis
 
 torch.backends.cudnn.benchmark = True
 
+
 def generate_video(vis_tensor_path):
 
     in_tensor = torch.load(vis_tensor_path)
@@ -60,13 +62,14 @@ def generate_video(vis_tensor_path):
     # self-rotated video
     tmp_path = vis_tensor_path.replace("_in_tensor.pt", "_tmp.mp4")
     out_path = vis_tensor_path.replace("_in_tensor.pt", ".mp4")
-    
+
     render.load_meshes(verts_lst, faces_lst)
     render.get_rendered_video_multi(in_tensor, tmp_path)
-    
-    os.system(f'ffmpeg -y -loglevel quiet -stats -i {tmp_path} -c:v libx264 {out_path}')
-    
-    return out_path, out_path
+
+    os.system(f"ffmpeg -y -loglevel quiet -stats -i {tmp_path} -vcodec libx264 {out_path}")
+
+    return out_path
+
 
 def generate_model(in_path, fitting_step=50):
 
@@ -87,7 +90,12 @@ def generate_model(in_path, fitting_step=50):
 
     # load normal model
     normal_net = Normal.load_from_checkpoint(
-        cfg=cfg, checkpoint_path=cfg.normal_path, map_location=device, strict=False
+        cfg=cfg,
+        checkpoint_path=hf_hub_download(
+            repo_id="Yuliang/ICON", use_auth_token=os.environ["ICON"], filename=cfg.normal_path
+        ),
+        map_location=device,
+        strict=False
     )
     normal_net = normal_net.to(device)
     normal_net.netG.eval()
@@ -111,7 +119,12 @@ def generate_model(in_path, fitting_step=50):
     if cfg.bni.use_ifnet:
         # load IFGeo model
         ifnet = IFGeo.load_from_checkpoint(
-            cfg=cfg, checkpoint_path=cfg.ifnet_path, map_location=device, strict=False
+            cfg=cfg,
+            checkpoint_path=hf_hub_download(
+                repo_id="Yuliang/ICON", use_auth_token=os.environ["ICON"], filename=cfg.ifnet_path
+            ),
+            map_location=device,
+            strict=False
         )
         ifnet = ifnet.to(device)
         ifnet.netG.eval()
@@ -644,15 +657,13 @@ def generate_model(in_path, fitting_step=50):
     overlap_path = img_overlap_path
     vis_tensor_path = osp.join(out_dir, cfg.name, f"vid/{data['name']}_in_tensor.pt")
 
-    # clean all the variables
-    for element in dir():
-        if 'path' not in element:
-            del locals()[element]
+    # # clean all the variables
+    # for element in dir():
+    #     if 'path' not in element:
+    #         del locals()[element]
 
-    import gc
-    gc.collect()
-    torch.cuda.empty_cache()
+    # import gc
+    # gc.collect()
+    # torch.cuda.empty_cache()
 
-    return [
-        smpl_glb_path, refine_glb_path, refine_obj_path, overlap_path, vis_tensor_path
-    ]
+    return [smpl_glb_path, refine_glb_path, refine_obj_path, overlap_path, vis_tensor_path]

configs/econ.yaml

@@ -1,7 +1,7 @@
 name: econ
 ckpt_dir: "./data/ckpt/"
-normal_path: "./data/ckpt/normal.ckpt"
-ifnet_path: "./data/ckpt/ifnet.ckpt"
+normal_path: "normal.ckpt"
+ifnet_path: "ifnet.ckpt"
 results_path: "./results"
 
 net:

gradio_cached_examples/13/log.csv

@@ -0,0 +1,2 @@
+flag,username,timestamp
+,,2023-04-15 18:15:46.412679

gradio_cached_examples/25/log.csv

@@ -0,0 +1,6 @@
+flag,username,timestamp
+,,2023-04-16 10:48:00.715491
+,,2023-04-16 10:50:02.250539
+,,2023-04-16 10:52:15.683112
+,,2023-04-16 10:54:18.253116
+,,2023-04-16 10:56:22.892765

lib/common/libmesh/inside_mesh.py

@@ -1,5 +1,4 @@
 import numpy as np
-
 from .triangle_hash import TriangleHash as _TriangleHash
 
 

lib/common/render.py

@@ -38,6 +38,7 @@ from pytorch3d.renderer import (
 )
 from pytorch3d.renderer.mesh import TexturesVertex
 from pytorch3d.structures import Meshes
+import torch.nn.functional as F
 from termcolor import colored
 from tqdm import tqdm
 
@@ -305,6 +306,9 @@ class Render:
 
         height, width = data["img_raw"].shape[2:]
 
+        width = int(width / (height / 256.0))
+        height = 256
+
         fourcc = cv2.VideoWriter_fourcc(*"mp4v")
         video = cv2.VideoWriter(
             save_path,
@@ -351,9 +355,12 @@ class Render:
                                       data)
             img_cloth = blend_rgb_norm((torch.stack(mesh_renders)[num_obj:, cam_id] - 0.5) * 2.0,
                                        data)
-            final_img = torch.cat([img_raw, img_smpl, img_cloth],
-                                  dim=-1).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
+            final_img = torch.cat([img_raw, img_smpl, img_cloth], dim=-1).squeeze(0)
+
+            final_img_rescale = F.interpolate(
+                final_img, size=(height, width), mode="bilinear", align_corners=False
+            ).squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
 
-            video.write(final_img[:, :, ::-1])
+            video.write(final_img_rescale[:, :, ::-1])
 
         video.release()