update

infer_api.py

@@ -107,7 +107,7 @@ for file in all_files:
     hf_hub_download(repo_id, file, local_dir="./ckpt")
 
 @spaces.GPU
-def set_seed(seed):
+def set_seed22(seed):
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
@@ -174,7 +174,7 @@ def process_image(image, totensor, width, height):
 def inference(validation_pipeline, input_image, vae, feature_extractor, image_encoder, unet, ref_unet, tokenizer,
               text_encoder, pretrained_model_path, validation, val_width, val_height, unet_condition_type,
               use_noise=True, noise_d=256, crop=False, seed=100, timestep=20):
-    set_seed(seed)
+    set_seed2(seed)
     generator = torch.Generator(device=device).manual_seed(seed)
     
     totensor = transforms.ToTensor()
@@ -372,10 +372,10 @@ class InferAPI:
         return infer_multiview_gen(img, seed, num_levels)
 
     def genStage3(self, img):
-        return self.slrm_infer.gen(img)
+        return infer_slrm_gen(img)
 
     def genStage4(self, meshes, imgs):
-        return self.refine_infer.refine(meshes, imgs)
+        return infer_refine(meshes, imgs)
 
 
 ############## Refine ##############
@@ -400,6 +400,7 @@ def srgb_to_linear(c_srgb):
     return c_linear.clip(0, 1.)
 
 
+@spaces.GPU
 def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
     # convert from pytorch3d meshes to trimesh mesh
     vertices = meshes.verts_packed().cpu().float().numpy()
@@ -515,245 +516,240 @@ def get_distract_mask(generator, color_0, color_1, normal_0=None, normal_1=None,
     return distract_mask, distract_bbox, random_sampled_points, final_mask
 
 
-class InferRefineAPI:
-    @spaces.GPU
-    def __init__(self, config):
-        self.sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda()
-        self.generator = SamAutomaticMaskGenerator(
-            model=self.sam,
-            points_per_side=64,
-            pred_iou_thresh=0.80,
-            stability_score_thresh=0.92,
-            crop_n_layers=1,
-            crop_n_points_downscale_factor=2,
-            min_mask_region_area=100,
-        )
-        self.outside_ratio = 0.20
-
-    @spaces.GPU
-    def refine(self, meshes, imgs):
-        fixed_v, fixed_f, fixed_t = None, None, None
-        flow_vert, flow_vector = None, None
-        last_colors, last_normals = None, None
-        last_front_color, last_front_normal = None, None
-        distract_mask = None
-
-        mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
-        mv = mv[[4, 3, 2, 0, 6, 5]]        
-        renderer = NormalsRenderer(mv,proj,(1024,1024))
-
-        results = []
-
-        for name_idx, level in zip([2, 0, 1], [2, 1, 0]):
-            mesh = trimesh.load(meshes[name_idx])
-            new_mesh = mesh.split(only_watertight=False)
-            new_mesh = [ j for j in new_mesh if len(j.vertices) >= 300 ]
-            mesh = trimesh.Scene(new_mesh).dump(concatenate=True)
-            mesh_v, mesh_f = mesh.vertices, mesh.faces
-
-            if last_colors is None:
-                images = renderer.render(
-                    torch.tensor(mesh_v, device='cuda').float(),
-                    torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
-                    torch.tensor(mesh_f, device='cuda'),
-                )
-                mask = (images[..., 3] < 0.9).cpu().numpy()
-
-            colors, normals = [], []
-            for i in range(6):
-                color = np.array(imgs[level]['images'][i])
-                normal = np.array(imgs[level]['normals'][i])
-
-                if last_colors is not None:
-                    offset = calc_horizontal_offset(np.array(last_colors[i]), color)
-                    # print('offset', i, offset)
-                else:
-                    offset = calc_horizontal_offset2(mask[i], color)
-                    # print('init offset', i, offset)
+infer_refine_sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda()
+infer_refine_generator = SamAutomaticMaskGenerator(
+    model=infer_refine_sam,
+    points_per_side=64,
+    pred_iou_thresh=0.80,
+    stability_score_thresh=0.92,
+    crop_n_layers=1,
+    crop_n_points_downscale_factor=2,
+    min_mask_region_area=100,
+)
+infer_refine_outside_ratio = 0.20
 
-                if offset != 0:
-                    color = np.roll(color, offset, axis=1)
-                    normal = np.roll(normal, offset, axis=1)
+@spaces.GPU
+def infer_refine(meshes, imgs):
+    fixed_v, fixed_f, fixed_t = None, None, None
+    flow_vert, flow_vector = None, None
+    last_colors, last_normals = None, None
+    last_front_color, last_front_normal = None, None
+    distract_mask = None
 
-                color = Image.fromarray(color)
-                normal = Image.fromarray(normal)
-                colors.append(color)
-                normals.append(normal)
+    mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
+    mv = mv[[4, 3, 2, 0, 6, 5]]        
+    renderer = NormalsRenderer(mv,proj,(1024,1024))
 
-            if last_front_color is not None and level == 0:
-                original_mask, distract_bbox, _, distract_mask = get_distract_mask(self.generator, last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=self.outside_ratio)
-            else:  
-                distract_mask = None
-                distract_bbox = None
+    results = []
 
-            last_front_color = np.array(colors[0]).astype(np.float32) / 255.0
-            last_front_normal = np.array(normals[0]).astype(np.float32) / 255.0
+    for name_idx, level in zip([2, 0, 1], [2, 1, 0]):
+        mesh = trimesh.load(meshes[name_idx])
+        new_mesh = mesh.split(only_watertight=False)
+        new_mesh = [ j for j in new_mesh if len(j.vertices) >= 300 ]
+        mesh = trimesh.Scene(new_mesh).dump(concatenate=True)
+        mesh_v, mesh_f = mesh.vertices, mesh.faces
+
+        if last_colors is None:
+            images = renderer.render(
+                torch.tensor(mesh_v, device='cuda').float(),
+                torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
+                torch.tensor(mesh_f, device='cuda'),
+            )
+            mask = (images[..., 3] < 0.9).cpu().numpy()
 
-            if last_colors is None:
-                from copy import deepcopy
-                last_colors, last_normals = deepcopy(colors), deepcopy(normals)
+        colors, normals = [], []
+        for i in range(6):
+            color = np.array(imgs[level]['images'][i])
+            normal = np.array(imgs[level]['normals'][i])
 
-            # my mesh flow weight by nearest vertexs
-            if fixed_v is not None and fixed_f is not None and level == 1:
-                t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
+            if last_colors is not None:
+                offset = calc_horizontal_offset(np.array(last_colors[i]), color)
+                # print('offset', i, offset)
+            else:
+                offset = calc_horizontal_offset2(mask[i], color)
+                # print('init offset', i, offset)
+
+            if offset != 0:
+                color = np.roll(color, offset, axis=1)
+                normal = np.roll(normal, offset, axis=1)
+
+            color = Image.fromarray(color)
+            normal = Image.fromarray(normal)
+            colors.append(color)
+            normals.append(normal)
+
+        if last_front_color is not None and level == 0:
+            original_mask, distract_bbox, _, distract_mask = get_distract_mask(infer_refine_generator, last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=infer_refine_outside_ratio)
+        else:  
+            distract_mask = None
+            distract_bbox = None
+
+        last_front_color = np.array(colors[0]).astype(np.float32) / 255.0
+        last_front_normal = np.array(normals[0]).astype(np.float32) / 255.0
+
+        if last_colors is None:
+            from copy import deepcopy
+            last_colors, last_normals = deepcopy(colors), deepcopy(normals)
+
+        # my mesh flow weight by nearest vertexs
+        if fixed_v is not None and fixed_f is not None and level == 1:
+            t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
+
+            fixed_v_cpu = fixed_v.cpu().numpy()
+            kdtree_anchor = KDTree(fixed_v_cpu)
+            kdtree_mesh_v = KDTree(mesh_v)
+            _, idx_anchor = kdtree_anchor.query(mesh_v, k=1)
+            _, idx_mesh_v = kdtree_mesh_v.query(mesh_v, k=25)
+            idx_anchor = idx_anchor.squeeze()
+            neighbors = torch.tensor(mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
+            # calculate the distances neighbors [V, 25, 3]; mesh_v [V, 3] -> [V, 25]
+            neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v).cuda()[:, None], dim=-1)
+            neighbor_dists[neighbor_dists > 0.06] = 114514.
+            neighbor_weights = torch.exp(-neighbor_dists * 1.)
+            neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
+            anchors = fixed_v[idx_anchor]  # V, 3
+            anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor]  # V, 3
+            dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
+            vec_anchor = dis_anchor[:, None] * anchor_normals  # V, 3
+            vec_anchor = vec_anchor[idx_mesh_v]  # V, 25, 3
+            weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1)  # V, 3
+            mesh_v += weighted_vec_anchor.cpu().numpy()
+
+            t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
+
+        mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
+        mesh_f = torch.tensor(mesh_f, device='cuda')
+
+        new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask, distract_bbox=distract_bbox)
+
+        # my mesh flow weight by nearest vertexs
+        try:
+            if fixed_v is not None and fixed_f is not None and level != 0:
+                new_mesh_v = new_mesh.verts_packed().cpu().numpy()
 
                 fixed_v_cpu = fixed_v.cpu().numpy()
                 kdtree_anchor = KDTree(fixed_v_cpu)
-                kdtree_mesh_v = KDTree(mesh_v)
-                _, idx_anchor = kdtree_anchor.query(mesh_v, k=1)
-                _, idx_mesh_v = kdtree_mesh_v.query(mesh_v, k=25)
+                kdtree_mesh_v = KDTree(new_mesh_v)
+                _, idx_anchor = kdtree_anchor.query(new_mesh_v, k=1)
+                _, idx_mesh_v = kdtree_mesh_v.query(new_mesh_v, k=25)
                 idx_anchor = idx_anchor.squeeze()
-                neighbors = torch.tensor(mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
-                # calculate the distances neighbors [V, 25, 3]; mesh_v [V, 3] -> [V, 25]
-                neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v).cuda()[:, None], dim=-1)
+                neighbors = torch.tensor(new_mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
+                # calculate the distances neighbors [V, 25, 3]; new_mesh_v [V, 3] -> [V, 25]
+                neighbor_dists = torch.norm(neighbors - torch.tensor(new_mesh_v).cuda()[:, None], dim=-1)
                 neighbor_dists[neighbor_dists > 0.06] = 114514.
                 neighbor_weights = torch.exp(-neighbor_dists * 1.)
                 neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
                 anchors = fixed_v[idx_anchor]  # V, 3
                 anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor]  # V, 3
-                dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
+                dis_anchor = torch.clamp(((anchors - torch.tensor(new_mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
                 vec_anchor = dis_anchor[:, None] * anchor_normals  # V, 3
                 vec_anchor = vec_anchor[idx_mesh_v]  # V, 25, 3
                 weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1)  # V, 3
-                mesh_v += weighted_vec_anchor.cpu().numpy()
-
-                t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)
-
-            mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
-            mesh_f = torch.tensor(mesh_f, device='cuda')
-
-            new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask, distract_bbox=distract_bbox)
-
-            # my mesh flow weight by nearest vertexs
-            try:
-                if fixed_v is not None and fixed_f is not None and level != 0:
-                    new_mesh_v = new_mesh.verts_packed().cpu().numpy()
-
-                    fixed_v_cpu = fixed_v.cpu().numpy()
-                    kdtree_anchor = KDTree(fixed_v_cpu)
-                    kdtree_mesh_v = KDTree(new_mesh_v)
-                    _, idx_anchor = kdtree_anchor.query(new_mesh_v, k=1)
-                    _, idx_mesh_v = kdtree_mesh_v.query(new_mesh_v, k=25)
-                    idx_anchor = idx_anchor.squeeze()
-                    neighbors = torch.tensor(new_mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
-                    # calculate the distances neighbors [V, 25, 3]; new_mesh_v [V, 3] -> [V, 25]
-                    neighbor_dists = torch.norm(neighbors - torch.tensor(new_mesh_v).cuda()[:, None], dim=-1)
-                    neighbor_dists[neighbor_dists > 0.06] = 114514.
-                    neighbor_weights = torch.exp(-neighbor_dists * 1.)
-                    neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
-                    anchors = fixed_v[idx_anchor]  # V, 3
-                    anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor]  # V, 3
-                    dis_anchor = torch.clamp(((anchors - torch.tensor(new_mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
-                    vec_anchor = dis_anchor[:, None] * anchor_normals  # V, 3
-                    vec_anchor = vec_anchor[idx_mesh_v]  # V, 25, 3
-                    weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1)  # V, 3
-                    new_mesh_v += weighted_vec_anchor.cpu().numpy()
-
-                    # replace new_mesh verts with new_mesh_v
-                    new_mesh = Meshes(verts=[torch.tensor(new_mesh_v, device='cuda')], faces=new_mesh.faces_list(), textures=new_mesh.textures)
-
-            except Exception as e:
-                pass
-
-            notsimp_v, notsimp_f, notsimp_t = new_mesh.verts_packed(), new_mesh.faces_packed(), new_mesh.textures.verts_features_packed()
-
-            if fixed_v is None:
-                fixed_v, fixed_f = simp_v, simp_f
-                complete_v, complete_f, complete_t = notsimp_v, notsimp_f, notsimp_t
-            else:
-                fixed_f = torch.cat([fixed_f, simp_f + fixed_v.shape[0]], dim=0)
-                fixed_v = torch.cat([fixed_v, simp_v], dim=0)
-                
-                complete_f = torch.cat([complete_f, notsimp_f + complete_v.shape[0]], dim=0)
-                complete_v = torch.cat([complete_v, notsimp_v], dim=0)
-                complete_t = torch.cat([complete_t, notsimp_t], dim=0)
-            
-            if level == 2:
-                new_mesh = Meshes(verts=[new_mesh.verts_packed()], faces=[new_mesh.faces_packed()], textures=pytorch3d.renderer.mesh.textures.TexturesVertex(verts_features=[torch.ones_like(new_mesh.textures.verts_features_packed(), device=new_mesh.verts_packed().device)*0.5]))
-
-            save_py3dmesh_with_trimesh_fast(new_mesh, meshes[name_idx].replace('.obj', '_refined.obj'), apply_sRGB_to_LinearRGB=False)
-            results.append(meshes[name_idx].replace('.obj', '_refined.obj'))
-
-        # save whole mesh
-        save_py3dmesh_with_trimesh_fast(Meshes(verts=[complete_v], faces=[complete_f], textures=pytorch3d.renderer.mesh.textures.TexturesVertex(verts_features=[complete_t])), meshes[name_idx].replace('.obj', '_refined_whole.obj'), apply_sRGB_to_LinearRGB=False)
-        results.append(meshes[name_idx].replace('.obj', '_refined_whole.obj'))
-
-        return results
-
-
-class InferSlrmAPI:
-    @spaces.GPU
-    def __init__(self, config):
-        self.config_path = config['config_path']
-        self.config = OmegaConf.load(self.config_path)
-        self.config_name = os.path.basename(self.config_path).replace('.yaml', '')
-        self.model_config = self.config.model_config
-        self.infer_config = self.config.infer_config
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.model = instantiate_from_config(self.model_config)
-        state_dict = torch.load(self.infer_config.model_path, map_location='cpu')
-        self.model.load_state_dict(state_dict, strict=False)
-        self.model = self.model.to(self.device)
-        self.model.init_flexicubes_geometry(self.device, fovy=30.0, is_ortho=self.model.is_ortho)
-        self.model = self.model.eval()
-
-    @spaces.GPU
-    def gen(self, imgs):
-        imgs = [ cv2.imread(img[0])[:, :, ::-1] for img in imgs ]
-        imgs = np.stack(imgs, axis=0).astype(np.float32) / 255.0
-        imgs = torch.from_numpy(np.array(imgs)).permute(0, 3, 1, 2).contiguous().float()   # (6, 3, 1024, 1024)
-        mesh_glb_fpaths = self.make3d(imgs)
-        return mesh_glb_fpaths[1:4] + mesh_glb_fpaths[0:1]
-
-    @spaces.GPU
-    def make3d(self, images):
-        input_cameras = torch.tensor(np.load('slrm/cameras.npy')).to(device)
-
-        images = images.unsqueeze(0).to(device)
-        images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1)
-
-        mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name
-        print(mesh_fpath)
-        mesh_basename = os.path.basename(mesh_fpath).split('.')[0]
-        mesh_dirname = os.path.dirname(mesh_fpath)
-
-        with torch.no_grad():
-            # get triplane
-            planes = self.model.forward_planes(images, input_cameras.float())
-
-            # get mesh
-            mesh_glb_fpaths = []
-            for j in range(4):
-                mesh_glb_fpath = self.make_mesh(mesh_fpath.replace(mesh_fpath[-4:], f'_{j}{mesh_fpath[-4:]}'), planes, level=[0, 3, 4, 2][j])
-                mesh_glb_fpaths.append(mesh_glb_fpath)
-
-        return mesh_glb_fpaths
-
-    @spaces.GPU
-    def make_mesh(self, mesh_fpath, planes, level=None):
-        mesh_basename = os.path.basename(mesh_fpath).split('.')[0]
-        mesh_dirname = os.path.dirname(mesh_fpath)
-        mesh_glb_fpath = os.path.join(mesh_dirname, f"{mesh_basename}.glb")
-            
-        with torch.no_grad():
-            # get mesh
-            mesh_out = self.model.extract_mesh(
-                planes,
-                use_texture_map=False,
-                levels=torch.tensor([level]).to(device),
-                **self.infer_config,
-            )
+                new_mesh_v += weighted_vec_anchor.cpu().numpy()
+
+                # replace new_mesh verts with new_mesh_v
+                new_mesh = Meshes(verts=[torch.tensor(new_mesh_v, device='cuda')], faces=new_mesh.faces_list(), textures=new_mesh.textures)
 
-            vertices, faces, vertex_colors = mesh_out
-            vertices = vertices[:, [1, 2, 0]]
+        except Exception as e:
+            pass
 
-            if level == 2:
-                # fill all vertex_colors with 127
-                vertex_colors = np.ones_like(vertex_colors) * 127
+        notsimp_v, notsimp_f, notsimp_t = new_mesh.verts_packed(), new_mesh.faces_packed(), new_mesh.textures.verts_features_packed()
+
+        if fixed_v is None:
+            fixed_v, fixed_f = simp_v, simp_f
+            complete_v, complete_f, complete_t = notsimp_v, notsimp_f, notsimp_t
+        else:
+            fixed_f = torch.cat([fixed_f, simp_f + fixed_v.shape[0]], dim=0)
+            fixed_v = torch.cat([fixed_v, simp_v], dim=0)
             
-            save_obj(vertices, faces, vertex_colors, mesh_fpath)
+            complete_f = torch.cat([complete_f, notsimp_f + complete_v.shape[0]], dim=0)
+            complete_v = torch.cat([complete_v, notsimp_v], dim=0)
+            complete_t = torch.cat([complete_t, notsimp_t], dim=0)
+        
+        if level == 2:
+            new_mesh = Meshes(verts=[new_mesh.verts_packed()], faces=[new_mesh.faces_packed()], textures=pytorch3d.renderer.mesh.textures.TexturesVertex(verts_features=[torch.ones_like(new_mesh.textures.verts_features_packed(), device=new_mesh.verts_packed().device)*0.5]))
+
+        save_py3dmesh_with_trimesh_fast(new_mesh, meshes[name_idx].replace('.obj', '_refined.obj'), apply_sRGB_to_LinearRGB=False)
+        results.append(meshes[name_idx].replace('.obj', '_refined.obj'))
+
+    # save whole mesh
+    save_py3dmesh_with_trimesh_fast(Meshes(verts=[complete_v], faces=[complete_f], textures=pytorch3d.renderer.mesh.textures.TexturesVertex(verts_features=[complete_t])), meshes[name_idx].replace('.obj', '_refined_whole.obj'), apply_sRGB_to_LinearRGB=False)
+    results.append(meshes[name_idx].replace('.obj', '_refined_whole.obj'))
+
+    return results
+
+config_slrm = {
+    'config_path': './configs/mesh-slrm-infer.yaml'
+}
+infer_slrm_config_path = config_slrm['config_path']
+infer_slrm_config = OmegaConf.load(infer_slrm_config_path)
+infer_slrm_config_name = os.path.basename(infer_slrm_config_path).replace('.yaml', '')
+infer_slrm_model_config = infer_slrm_config.model_config
+infer_slrm_infer_config = infer_slrm_config.infer_config
+infer_slrm_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+infer_slrm_model = instantiate_from_config(infer_slrm_model_config)
+state_dict = torch.load(infer_slrm_infer_config.model_path, map_location='cpu')
+infer_slrm_model.load_state_dict(state_dict, strict=False)
+infer_slrm_model = infer_slrm_model.to(infer_slrm_device)
+infer_slrm_model.init_flexicubes_geometry(infer_slrm_device, fovy=30.0, is_ortho=infer_slrm_model.is_ortho)
+infer_slrm_model = infer_slrm_model.eval()
+
+@spaces.GPU
+def infer_slrm_gen(imgs):
+    imgs = [ cv2.imread(img[0])[:, :, ::-1] for img in imgs ]
+    imgs = np.stack(imgs, axis=0).astype(np.float32) / 255.0
+    imgs = torch.from_numpy(np.array(imgs)).permute(0, 3, 1, 2).contiguous().float()   # (6, 3, 1024, 1024)
+    mesh_glb_fpaths = infer_slrm_make3d(imgs)
+    return mesh_glb_fpaths[1:4] + mesh_glb_fpaths[0:1]
+
+@spaces.GPU
+def infer_slrm_make3d(images):
+    input_cameras = torch.tensor(np.load('slrm/cameras.npy')).to(device)
+
+    images = images.unsqueeze(0).to(device)
+    images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1)
+
+    mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name
+    print(mesh_fpath)
+    mesh_basename = os.path.basename(mesh_fpath).split('.')[0]
+    mesh_dirname = os.path.dirname(mesh_fpath)
+
+    with torch.no_grad():
+        # get triplane
+        planes = infer_slrm_model.forward_planes(images, input_cameras.float())
+
+        # get mesh
+        mesh_glb_fpaths = []
+        for j in range(4):
+            mesh_glb_fpath = infer_slrm_make_mesh(mesh_fpath.replace(mesh_fpath[-4:], f'_{j}{mesh_fpath[-4:]}'), planes, level=[0, 3, 4, 2][j])
+            mesh_glb_fpaths.append(mesh_glb_fpath)
+
+    return mesh_glb_fpaths
+
+@spaces.GPU
+def infer_slrm_make_mesh(mesh_fpath, planes, level=None):
+    mesh_basename = os.path.basename(mesh_fpath).split('.')[0]
+    mesh_dirname = os.path.dirname(mesh_fpath)
+        
+    with torch.no_grad():
+        # get mesh
+        mesh_out = infer_slrm_model.extract_mesh(
+            planes,
+            use_texture_map=False,
+            levels=torch.tensor([level]).to(device),
+            **infer_slrm_infer_config,
+        )
+
+        vertices, faces, vertex_colors = mesh_out
+        vertices = vertices[:, [1, 2, 0]]
+
+        if level == 2:
+            # fill all vertex_colors with 127
+            vertex_colors = np.ones_like(vertex_colors) * 127
+        
+        save_obj(vertices, faces, vertex_colors, mesh_fpath)
 
-        return mesh_fpath
+    return mesh_fpath
 
 
 parser = argparse.ArgumentParser()