remove cuda

app.py

@@ -104,8 +104,8 @@ def gen_image(input_image, seed, scale, step):
     np_imgs = np.concatenate(stage1_images, 1)
     np_xyzs = np.concatenate(stage2_images, 1)
 
-    glb_path, obj_path = generate3d(model, np_imgs, np_xyzs, args.device)
-    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path, obj_path
+    glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
+    return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path#, obj_path
 
 
 parser = argparse.ArgumentParser()
@@ -202,7 +202,7 @@ with gr.Blocks() as demo:
                 interactive=False,
             )
             gr.Markdown("Note: The GLB model shown here has a darker lighting and enlarged UV seams. Download for correct results.")
-            output_obj = gr.File(interactive=False, label="Output OBJ")
+            # output_obj = gr.File(interactive=False, label="Output OBJ")
 
     inputs = [
         processed_image,
@@ -214,7 +214,7 @@ with gr.Blocks() as demo:
         image_output,
         xyz_ouput,
         output_model,
-        output_obj,
+        # output_obj,
     ]
 
 

inference.py

@@ -37,7 +37,7 @@ def generate3d(model, rgb, ccm, device):
     triplane = torch.cat([color,xyz],dim=1).to(device)
     # 3D visualize
     model.eval()
-    glctx = dr.RasterizeGLContext()#dr.RasterizeCudaContext()
+    
 
     if model.denoising == True:
         tnew = 20
@@ -67,29 +67,33 @@ def generate3d(model, rgb, ccm, device):
         
 
     from kiui.mesh_utils import clean_mesh
-    verts, faces = clean_mesh(data_config['verts'].squeeze().cpu().numpy().astype(np.float32), data_config['faces'].squeeze().cpu().numpy().astype(np.int32), repair = False, remesh=False, remesh_size=0.005)
+    verts, faces = clean_mesh(data_config['verts'].squeeze().cpu().numpy().astype(np.float32), data_config['faces'].squeeze().cpu().numpy().astype(np.int32), repair = False, remesh=True, remesh_size=0.005, remesh_iters=1)
     data_config['verts'] = torch.from_numpy(verts).cuda().contiguous()
     data_config['faces'] = torch.from_numpy(faces).cuda().contiguous()
 
     start_time = time.time()
     with torch.no_grad():
-        mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
-        model.export_mesh_wt_uv(glctx, data_config, mesh_path_obj, "", device, res=(1024,1024), tri_fea_2=triplane_feature2)
-
-        mesh = Mesh.load(mesh_path_obj+".obj", bound=0.9, front_dir="+z")
         mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
-        mesh.write(mesh_path_glb+".glb")
+        model.export_mesh(data_config, mesh_path_glb, tri_fea_2 = triplane_feature2)
+
+        # glctx = dr.RasterizeGLContext()#dr.RasterizeCudaContext()
+        # mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
+        # model.export_mesh_wt_uv(glctx, data_config, mesh_path_obj, "", device, res=(1024,1024), tri_fea_2=triplane_feature2)
+
+        # mesh = Mesh.load(mesh_path_obj+".obj", bound=0.9, front_dir="+z")
+        # mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
+        # mesh.write(mesh_path_glb+".glb")
 
-        # mesh_obj2 = trimesh.load(mesh_path_glb+".glb", file_type='glb')
-        # mesh_path_obj2 = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
-        # mesh_obj2.export(mesh_path_obj2+".obj")
+        # # mesh_obj2 = trimesh.load(mesh_path_glb+".glb", file_type='glb')
+        # # mesh_path_obj2 = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
+        # # mesh_obj2.export(mesh_path_obj2+".obj")
 
-        with zipfile.ZipFile(mesh_path_obj+'.zip', 'w') as myzip:
-            myzip.write(mesh_path_obj+'.obj', mesh_path_obj.split("/")[-1]+'.obj')
-            myzip.write(mesh_path_obj+'.png', mesh_path_obj.split("/")[-1]+'.png')
-            myzip.write(mesh_path_obj+'.mtl', mesh_path_obj.split("/")[-1]+'.mtl')
+        # with zipfile.ZipFile(mesh_path_obj+'.zip', 'w') as myzip:
+        #     myzip.write(mesh_path_obj+'.obj', mesh_path_obj.split("/")[-1]+'.obj')
+        #     myzip.write(mesh_path_obj+'.png', mesh_path_obj.split("/")[-1]+'.png')
+        #     myzip.write(mesh_path_obj+'.mtl', mesh_path_obj.split("/")[-1]+'.mtl')
 
     end_time = time.time()
     elapsed_time = end_time - start_time
     print(f"uv takes {elapsed_time}s")
-    return mesh_path_glb+".glb", mesh_path_obj+'.zip'
+    return mesh_path_glb+".obj"

model/crm/model.py

@@ -89,7 +89,7 @@ class CRM(nn.Module):
         _, verts, faces = self.renderer(data, pred_sdf, deformation, tet_verts, tet_indices, weight= weight)
         return verts[0].unsqueeze(0), faces[0].int()
 
-    def export_mesh(self, data, out_dir, ind, device=None, tri_fea_2 = None):
+    def export_mesh(self, data, out_dir, tri_fea_2 = None):
         verts = data['verts']
         faces = data['faces']
 
@@ -98,13 +98,13 @@ class CRM(nn.Module):
         # Expect predicted colors value range from [-1, 1]
         colors = (colors * 0.5 + 0.5).clip(0, 1)
 
-        verts = verts.squeeze().cpu().numpy()
-        faces = faces[..., [2, 1, 0]].squeeze().cpu().numpy()
+        verts = verts[..., [0, 2, 1]].squeeze().cpu().numpy()
+        faces = faces[..., [2, 1, 0]][..., [0, 2, 1]].squeeze().cpu().numpy()#faces[..., [2, 1, 0]].squeeze().cpu().numpy()
 
         # export the final mesh
         with torch.no_grad():
             mesh = trimesh.Trimesh(verts, faces, vertex_colors=colors, process=False) # important, process=True leads to seg fault...
-            mesh.export(out_dir / f'{ind}.obj')
+            mesh.export(f'{out_dir}.obj')
 
     def export_mesh_wt_uv(self, ctx, data, out_dir, ind, device, res, tri_fea_2=None):