finish zerogpu debug

Anymate/.gitignore → .gitignore

@@ -23,4 +23,6 @@ blender-*
 *.csv
 *.tga
 *.png
-*.jpg
+*.jpg
+
+core*

Anymate/model.py

@@ -352,6 +352,22 @@ class EncoderDecoder(nn.Module):
             return cond
     
     def forward(self, data, device='cuda', downsample=False, **kwargs):
+        # data['points_cloud'] = data['points_cloud'].to(self.dtype).to(device)
+        # if 'vertices' in data.keys():
+        #     data['vertices'] = torch.tensor(data['vertices']).to(self.dtype).to(device).unsqueeze(0)
+        # if 'joints' in data.keys():
+        #     data['joints'] = torch.tensor(data['joints'], dtype=torch.float32).to(device).unsqueeze(0)
+        # if 'bones' in data.keys():
+        #     data['bones'] = torch.tensor(data['bones'], dtype=torch.float32).to(device).unsqueeze(0)
+        # if 'joints_mask' in data.keys():
+        #     data['joints_mask'] = torch.tensor(data['joints_mask'], dtype=torch.float32).to(device).unsqueeze(0)
+        # if 'bones_mask' in data.keys():
+        #     data['bones_mask'] = torch.tensor(data['bones_mask'], dtype=torch.float32).to(device).unsqueeze(0)
+
+        for key in data.keys():
+            if key in ['points_cloud', 'vertices', 'joints', 'bones', 'joints_mask', 'bones_mask']:
+                data[key] = data[key].to(torch.float32).to(device).unsqueeze(0)
+
         latents = self.encode(data, device)
         # print('latents shape', latents.shape)
         

Anymate/models/conn.py

@@ -24,6 +24,7 @@ class AttendjointsDecoder_con_combine(nn.Module):
         self.use_checkpoint = use_checkpoint
         self.separate = separate
         self.use_mask = use_mask
+        self.inference = False
         # self.num_latents = num_latents
 
         # self.query = nn.Parameter(torch.randn((num_latents, width), device=device, dtype=dtype) * 0.02)
@@ -104,6 +105,9 @@ class AttendjointsDecoder_con_combine(nn.Module):
             mask = data['joints_mask'].to(device)
             logits = logits.masked_fill(mask.unsqueeze(1) == 0, -1e8)
 
+        if self.inference:
+            logits = torch.argmax(logits, dim=-1).squeeze(0).to('cpu')
+
         return logits
     
 class AttendjointsDecoder_con_token(nn.Module):

Anymate/models/joint.py

@@ -8,6 +8,7 @@ from Anymate.utils.vol_utils import get_co, sample_from_planes, generate_planes
 from einops import repeat
 from sklearn.cluster import DBSCAN
 from Anymate.utils.vol_utils import extract_keypoints
+import numpy as np
 
 class TransformerDecoder(nn.Module):
     def __init__(self, 
@@ -66,7 +67,7 @@ class TransformerDecoder(nn.Module):
             cluster_centers = []
             for cluster in set(clustering.labels_):
                 cluster_centers.append(joints[clustering.labels_ == cluster].mean(axis=0))
-            return cluster_centers
+            return torch.tensor(np.array(cluster_centers)).to('cpu')
         return logits
     
 

Anymate/models/skin.py

@@ -140,6 +140,9 @@ class AttendjointsDecoder_combine(nn.Module):
 
         if downsample and not self.inference:
             return logits, idx
+        
+        if self.inference:
+            logits = logits.softmax(dim=-1).squeeze(0).to('cpu')
 
         return logits
 

Anymate/utils/ui_utils.py

@@ -198,6 +198,17 @@ def vis_skinning(normalized_mesh_file, joints, conns, skins):
     vis_file = visualize_results(mesh_file=normalized_mesh_file, joints=joints, conns=conns, skins=skins)
     return vis_file, vis_file
 
+def vis_all(normalized_mesh_file):
+    if normalized_mesh_file is None:
+        return None, None
+    joints = torch.load(normalized_mesh_file.replace('object.obj', 'joints.pt'))
+    conns = torch.load(normalized_mesh_file.replace('object.obj', 'conns.pt'))
+    skins = torch.load(normalized_mesh_file.replace('object.obj', 'skins.pt'))
+    print("All results loaded")
+    vis_file = visualize_results(mesh_file=normalized_mesh_file, joints=joints, conns=conns, skins=skins)
+    print("Finish Visualization")
+    return vis_file, vis_file
+
 def prepare_blender_file(normalized_mesh_file):
     if normalized_mesh_file is None:
         return None
@@ -245,11 +256,11 @@ def process_input(mesh_file):
 
     vis_file = visualize_results(mesh_file=normalized_mesh_file)
     pc = process_mesh_to_pc(normalized_mesh_file)
-    pc = torch.from_numpy(pc).to(anymate_args.device).to(torch.float32)
+    pc = torch.from_numpy(pc).to(torch.float32)
 
     # print(pc.shape, pc.max(dim=0), pc.min(dim=0))
 
-    return normalized_mesh_file, vis_file, vis_file, None, pc, None, None, None
+    return normalized_mesh_file, vis_file, vis_file, None, pc
 
 
 def get_model(checkpoint):
@@ -265,9 +276,10 @@ def get_result_connectivity(mesh_file, model, pc, joints):
 def get_result_skinning(mesh_file, model, pc, joints, conns):
     # mesh = trimesh.load(mesh_file)
     mesh = obj2mesh(mesh_file)
-    vertices = torch.from_numpy(mesh.vertices).to(anymate_args.device).to(torch.float32)
-    vertex_normals = torch.from_numpy(mesh.vertex_normals).to(anymate_args.device).to(torch.float32)
-    vertices = torch.cat([vertices, vertex_normals], dim=-1)
+    # vertices = torch.from_numpy(mesh.vertices, device='cpu')
+    # vertex_normals = torch.from_numpy(mesh.vertex_normals, device='cpu')
+    # vertices = torch.cat([vertices, vertex_normals], dim=-1)
+    vertices = torch.tensor(np.concatenate([mesh.vertices, mesh.vertex_normals], axis=-1))
     return get_skinning(pc, joints, conns, model, vertices=vertices, device=anymate_args.device, save=mesh_file.replace('object.obj', 'skins.pt'))
 
 def get_all_models(checkpoint_joint, checkpoint_conn, checkpoint_skin):

Anymate/utils/utils.py

@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 from Anymate.model import EncoderDecoder
 from sklearn.cluster import DBSCAN
 
@@ -25,26 +26,28 @@ def load_checkpoint(path, device, num_joints):
 
 def get_joint(pc, model, device='cuda', save=None, vox=None, eps=0.03, min_samples=1):
     model.eval()
-    data = {'points_cloud': pc.unsqueeze(0)}
+    data = {'points_cloud': pc}
     if vox is not None:
         data['vox'] = vox.unsqueeze(0)
     with torch.no_grad():
         model.decoder.inference_mode(eps=eps, min_samples=min_samples)
         joints = model(data, device=device)
-        joints = torch.tensor(joints, dtype=torch.float32).to(device)
+        # joints = torch.tensor(joints, dtype=torch.float32).to(device)
 
         if save is not None:
             torch.save(joints, save)
 
         return joints
-    
+
 def get_connectivity(pc, joints, model, device='cuda',return_prob=False, save=None):
     model.eval()
-    data = {'points_cloud': pc.unsqueeze(0), 'joints': joints.unsqueeze(0), 'joints_num': torch.tensor([joints.shape[0]]), 
-            'joints_mask': torch.ones(joints.shape[0], device=device).unsqueeze(0)}
+    data = {'points_cloud': pc, 'joints': joints, 'joints_num': torch.tensor([len(joints)]), 
+            'joints_mask': torch.ones(len(joints))}
     with torch.no_grad():
-        conns = model(data, device=device).softmax(dim=-1)
-        conns = conns.squeeze(0) if return_prob else torch.argmax(conns, dim=-1).squeeze(0)
+        if not return_prob:
+            model.decoder.inference = True
+        conns = model(data, device=device)
+        # conns = conns.squeeze(0) if return_prob else torch.argmax(conns, dim=-1).squeeze(0)
 
         if save is not None:
             torch.save(conns, save)
@@ -56,20 +59,20 @@ def get_skinning(pc, joints, conns, model, vertices=None, bones=None, device='cu
     
     if bones is None:
         bones = []
-        for i in range(joints.shape[0]):
+        for i in range(len(joints)):
             if conns[i] != i:
-                bones.append(torch.cat((joints[conns[i]], joints[i]), dim=-1))
-        bones = torch.stack(bones, dim=0)
+                bones.append(np.concatenate((joints[conns[i]], joints[i]), axis=-1))
+        bones = torch.tensor(np.array(bones))
 
-    data = {'points_cloud': pc.unsqueeze(0), 'bones': bones.unsqueeze(0), 'bones_num': torch.tensor([bones.shape[0]]),
-            'bones_mask': torch.ones(bones.shape[0], device=device).unsqueeze(0)}
-    
-    if vertices is not None:
-        data['vertices'] = vertices.unsqueeze(0)
-        model.decoder.inference = True
+    data = {'points_cloud': pc, 'bones': bones, 'bones_num': torch.tensor([len(bones)]),
+            'bones_mask': torch.ones(len(bones))}
 
     with torch.no_grad():
-        skins = model(data, device=device).softmax(dim=-1).squeeze(0)
+        if vertices is not None:
+            data['vertices'] = vertices
+            model.decoder.inference = True
+
+        skins = model(data, device=device)
         
         if save is not None:
             torch.save(skins, save)

app.py

@@ -1,8 +1,32 @@
+import spaces
 import gradio as gr
 import os
-from Anymate.args import ui_args
-from Anymate.utils.ui_utils import process_input, vis_joint, vis_connectivity, vis_skinning, prepare_blender_file
-from Anymate.utils.ui_utils import get_model, get_result_joint, get_result_connectivity, get_result_skinning, get_all_models, get_all_results
+from Anymate.args import ui_args, anymate_args
+from Anymate.utils.ui_utils import process_input, vis_joint, vis_connectivity, vis_skinning, vis_all, prepare_blender_file
+from Anymate.utils.ui_utils import get_result_joint, get_result_connectivity, get_result_skinning
+
+from Anymate.utils.utils import load_checkpoint
+
+# Check if checkpoints exist, if not download them
+if not (os.path.exists(ui_args.checkpoint_joint) and 
+        os.path.exists(ui_args.checkpoint_conn) and 
+        os.path.exists(ui_args.checkpoint_skin)):
+    print("Missing checkpoints, downloading them...")
+    os.system("bash Anymate/get_checkpoints.sh")
+
+model_joint = load_checkpoint(ui_args.checkpoint_joint, 'cpu', anymate_args.num_joints).to(anymate_args.device)
+model_connectivity = load_checkpoint(ui_args.checkpoint_conn, 'cpu', anymate_args.num_joints).to(anymate_args.device)
+model_skinning = load_checkpoint(ui_args.checkpoint_skin, 'cpu', anymate_args.num_joints).to(anymate_args.device)
+
+@spaces.GPU
+def get_all_results(mesh_file, pc, eps=0.03, min_samples=1):
+    # pc = pc.to(anymate_args.device)
+    joints = get_result_joint(mesh_file, model_joint, pc, eps=eps, min_samples=min_samples)
+    conns = get_result_connectivity(mesh_file, model_connectivity, pc, joints)
+    skins = get_result_skinning(mesh_file, model_skinning, pc, joints, conns)
+    print("Finish Inference")
+    return
+
 
 with gr.Blocks() as demo:
     gr.Markdown("""
@@ -12,13 +36,13 @@ with gr.Blocks() as demo:
     pc = gr.State(value=None)
     normalized_mesh_file = gr.State(value=None)
 
-    result_joint = gr.State(value=None)
-    result_connectivity = gr.State(value=None)
-    result_skinning = gr.State(value=None)
+    # result_joint = gr.State(value=None)
+    # result_connectivity = gr.State(value=None)
+    # result_skinning = gr.State(value=None)
 
-    model_joint = gr.State(value=None)
-    model_connectivity = gr.State(value=None)
-    model_skinning = gr.State(value=None)
+    # model_joint = gr.State(value=model_joint)
+    # model_connectivity = gr.State(value=model_connectivity)
+    # model_skinning = gr.State(value=model_skinning)
     
     with gr.Row():
         with gr.Column():
@@ -36,7 +60,8 @@ with gr.Blocks() as demo:
             sample_dropdown = gr.Dropdown(
                 label="Select Sample Object",
                 choices=sample_objects,
-                interactive=True
+                interactive=True,
+                value=sample_objects[0]
             )
             
             load_sample_btn = gr.Button("Load Sample")
@@ -54,51 +79,51 @@ with gr.Blocks() as demo:
             blender_file = gr.File(label="Output Blender File", scale=1)
 
     # Checkpoint paths
-    joint_models_dir = 'Anymate/checkpoints/joint'
-    joint_models = [os.path.join(joint_models_dir, f) for f in os.listdir(joint_models_dir) 
-                    if os.path.isfile(os.path.join(joint_models_dir, f))]
-    with gr.Row():
-        joint_checkpoint = gr.Dropdown(
-            label="Joint Checkpoint",
-            choices=joint_models,
-            value=ui_args.checkpoint_joint,
-            interactive=True
-        )
-        joint_status = gr.Checkbox(label="Joint Model Status", value=False, interactive=False, scale=0.3)
+    # joint_models_dir = 'Anymate/checkpoints/joint'
+    # joint_models = [os.path.join(joint_models_dir, f) for f in os.listdir(joint_models_dir) 
+    #                 if os.path.isfile(os.path.join(joint_models_dir, f))]
+    # with gr.Row():
+    #     joint_checkpoint = gr.Dropdown(
+    #         label="Joint Checkpoint",
+    #         choices=joint_models,
+    #         value=ui_args.checkpoint_joint,
+    #         interactive=True
+    #     )
+    #     joint_status = gr.Checkbox(label="Joint Model Status", value=False, interactive=False, scale=0.3)
         # with gr.Column():
         #     with gr.Row():
         #         load_joint_btn = gr.Button("Load", scale=0.3)
                 
         #     process_joint_btn = gr.Button("Process", scale=0.3)
 
-    conn_models_dir = 'Anymate/checkpoints/conn' 
-    conn_models = [os.path.join(conn_models_dir, f) for f in os.listdir(conn_models_dir)
-                    if os.path.isfile(os.path.join(conn_models_dir, f))]
-    with gr.Row():
-        conn_checkpoint = gr.Dropdown(
-            label="Connection Checkpoint",
-            choices=conn_models,
-            value=ui_args.checkpoint_conn,
-            interactive=True
-        )
-        conn_status = gr.Checkbox(label="Connectivity Model Status", value=False, interactive=False, scale=0.3)
+    # conn_models_dir = 'Anymate/checkpoints/conn' 
+    # conn_models = [os.path.join(conn_models_dir, f) for f in os.listdir(conn_models_dir)
+    #                 if os.path.isfile(os.path.join(conn_models_dir, f))]
+    # with gr.Row():
+    #     conn_checkpoint = gr.Dropdown(
+    #         label="Connection Checkpoint",
+    #         choices=conn_models,
+    #         value=ui_args.checkpoint_conn,
+    #         interactive=True
+    #     )
+    #     conn_status = gr.Checkbox(label="Connectivity Model Status", value=False, interactive=False, scale=0.3)
         # with gr.Column():
         #     with gr.Row():
         #         load_conn_btn = gr.Button("Load", scale=0.3)
                 
         #     process_conn_btn = gr.Button("Process", scale=0.3)
 
-    skin_models_dir = 'Anymate/checkpoints/skin'
-    skin_models = [os.path.join(skin_models_dir, f) for f in os.listdir(skin_models_dir)
-                    if os.path.isfile(os.path.join(skin_models_dir, f))]
-    with gr.Row():
-        skin_checkpoint = gr.Dropdown(
-            label="Skin Checkpoint", 
-            choices=skin_models,
-            value=ui_args.checkpoint_skin,
-            interactive=True
-        )
-        skin_status = gr.Checkbox(label="Skinning Model Status", value=False, interactive=False, scale=0.3)
+    # skin_models_dir = 'Anymate/checkpoints/skin'
+    # skin_models = [os.path.join(skin_models_dir, f) for f in os.listdir(skin_models_dir)
+    #                 if os.path.isfile(os.path.join(skin_models_dir, f))]
+    # with gr.Row():
+    #     skin_checkpoint = gr.Dropdown(
+    #         label="Skin Checkpoint", 
+    #         choices=skin_models,
+    #         value=ui_args.checkpoint_skin,
+    #         interactive=True
+    #     )
+    #     skin_status = gr.Checkbox(label="Skinning Model Status", value=False, interactive=False, scale=0.3)
         # with gr.Column():
         #     with gr.Row():
         #         load_skin_btn = gr.Button("Load", scale=0.3)
@@ -106,7 +131,7 @@ with gr.Blocks() as demo:
         #     process_skin_btn = gr.Button("Process", scale=0.3)
 
     with gr.Row():
-        load_all_btn = gr.Button("Load all models", scale=1)
+        # load_all_btn = gr.Button("Load all models", scale=1)
         process_all_btn = gr.Button("Run all models", scale=1)
         # download_btn = gr.DownloadButton("Blender File Not Ready", scale=0.3)
         # blender_file = gr.File(label="Blender File", scale=1)
@@ -117,7 +142,7 @@ with gr.Blocks() as demo:
     mesh_input.change(
         process_input,
         inputs=mesh_input,
-        outputs=[normalized_mesh_file, mesh_output, mesh_output2, blender_file, pc, result_joint, result_connectivity, result_skinning]
+        outputs=[normalized_mesh_file, mesh_output, mesh_output2, blender_file, pc]
     )
 
     load_sample_btn.click(
@@ -127,7 +152,7 @@ with gr.Blocks() as demo:
     ).then(
         process_input,
         inputs=mesh_input,
-        outputs=[normalized_mesh_file, mesh_output, mesh_output2, blender_file, pc, result_joint, result_connectivity, result_skinning]
+        outputs=[normalized_mesh_file, mesh_output, mesh_output2, blender_file, pc]
     )
 
     normalized_mesh_file.change(
@@ -136,95 +161,43 @@ with gr.Blocks() as demo:
         outputs=mesh_input
     )
 
-    result_joint.change(
-        vis_joint,
-        inputs=[normalized_mesh_file, result_joint],
-        outputs=[mesh_output, mesh_output2]
-    )
-
-    result_connectivity.change(
-        vis_connectivity,
-        inputs=[normalized_mesh_file, result_joint, result_connectivity],
-        outputs=[mesh_output, mesh_output2]
-    )
-
-    result_skinning.change(
-        vis_skinning,
-        inputs=[normalized_mesh_file, result_joint, result_connectivity, result_skinning],
-        outputs=[mesh_output, mesh_output2]
-    )
-
-    result_skinning.change(
-        prepare_blender_file,
-        inputs=[normalized_mesh_file],
-        outputs=blender_file
-    )
-
-    joint_checkpoint.change(
-        get_model,
-        inputs=joint_checkpoint,
-        outputs=[model_joint, joint_status]
-    )
-    
-    conn_checkpoint.change(
-        get_model,
-        inputs=conn_checkpoint,
-        outputs=[model_connectivity, conn_status]
-    )
-    
-    skin_checkpoint.change(
-        get_model,
-        inputs=skin_checkpoint,
-        outputs=[model_skinning, skin_status]
-    )
-
-    load_all_btn.click(
-        get_all_models,
-        inputs=[joint_checkpoint, conn_checkpoint, skin_checkpoint],
-        outputs=[model_joint, model_connectivity, model_skinning, joint_status, conn_status, skin_status]
-    )
-
-    process_all_btn.click(
-        get_all_results,
-        inputs=[normalized_mesh_file, model_joint, model_connectivity, model_skinning, pc, eps, min_samples],
-        outputs=[result_joint, result_connectivity, result_skinning]
-    )
-
-    # load_joint_btn.click(
-    #     fn=get_model,
-    #     inputs=joint_checkpoint,
-    #     outputs=[model_joint, joint_status]
+    # result_joint.change(
+    #     vis_joint,
+    #     inputs=[normalized_mesh_file, result_joint],
+    #     outputs=[mesh_output, mesh_output2]
     # )
 
-    # load_conn_btn.click(
-    #     fn=get_model,
-    #     inputs=conn_checkpoint,
-    #     outputs=[model_connectivity, conn_status]
+    # result_connectivity.change(
+    #     vis_connectivity,
+    #     inputs=[normalized_mesh_file, result_joint, result_connectivity],
+    #     outputs=[mesh_output, mesh_output2]
     # )
 
-    # load_skin_btn.click(
-    #     fn=get_model,
-    #     inputs=skin_checkpoint,
-    #     outputs=[model_skinning, skin_status]
+    # result_skinning.change(
+    #     vis_skinning,
+    #     inputs=[normalized_mesh_file, result_joint, result_connectivity, result_skinning],
+    #     outputs=[mesh_output, mesh_output2]
     # )
 
-    # process_joint_btn.click(
-    #     fn=get_result_joint,
-    #     inputs=[normalized_mesh_file, model_joint, pc, eps, min_samples],
-    #     outputs=result_joint
-    # )
-    
-    # process_conn_btn.click(
-    #     fn=get_result_connectivity,
-    #     inputs=[normalized_mesh_file, model_connectivity, pc, result_joint],
-    #     outputs=result_connectivity
-    # )
-    
-    # process_skin_btn.click(
-    #     fn=get_result_skinning,
-    #     inputs=[normalized_mesh_file, model_skinning, pc, result_joint, result_connectivity],
-    #     outputs=result_skinning
+    # result_skinning.change(
+    #     prepare_blender_file,
+    #     inputs=[normalized_mesh_file],
+    #     outputs=blender_file
     # )
 
+    process_all_btn.click(
+        get_all_results,
+        inputs=[normalized_mesh_file, pc, eps, min_samples],
+        outputs=[]
+    ).then(
+        vis_all,
+        inputs=[normalized_mesh_file],
+        outputs=[mesh_output, mesh_output2]
+    ).then(
+        prepare_blender_file,
+        inputs=[normalized_mesh_file],
+        outputs=blender_file
+    )
+
 if __name__ == "__main__":
     demo.launch()