to solve gs render error

LHM/models/modeling_human_lrm.py

@@ -1092,7 +1092,7 @@ class ModelHumanLRMSapdinoBodyHeadSD3_5(ModelHumanLRMSapdinoBodyHeadSD3):
         print(f"time elpased(forward gs model):{time.time() -  start_time}")
         return gs_model_list, query_points, smplx_params['transform_mat_neutral_pose']
     
-    
+    @torch.no_grad()
     def animation_infer(self, gs_model_list, query_points, smplx_params, render_c2ws, render_intrs, render_bg_colors):
         '''Inference code avoid repeat forward.
         '''

LHM/models/rendering/gs_renderer.py

@@ -818,6 +818,7 @@ class GS3DRenderer(nn.Module):
     def hyper_step(self, step):
         self.gs_net.hyper_step(step)
 
+    @torch.no_grad()
     def forward_single_view(
         self,
         gs: GaussianModel,
@@ -884,17 +885,17 @@ class GS3DRenderer(nn.Module):
 
         # Rasterize visible Gaussians to image, obtain their radii (on screen).
         # NOTE that dadong tries to regress rgb not shs
-        with torch.autocast(device_type=self.device.type, dtype=torch.float32):
-            rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
-                means3D=means3D.float(),
-                means2D=means2D.float(),
-                shs=shs,
-                colors_precomp=colors_precomp,
-                opacities=opacity.float(),
-                scales=scales.float(),
-                rotations=rotations.float(),
-                cov3D_precomp=cov3D_precomp,
-            )
+        # with torch.autocast(device_type=self.device.type, dtype=torch.float32):
+        rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
+            means3D=means3D.float(),
+            means2D=means2D.float(),
+            shs=shs,
+            colors_precomp=colors_precomp,
+            opacities=opacity.float(),
+            scales=scales.float(),
+            rotations=rotations.float(),
+            cov3D_precomp=cov3D_precomp,
+        )
 
         ret = {
             "comp_rgb": rendered_image.permute(1, 2, 0),  # [H, W, 3]
@@ -934,7 +935,7 @@ class GS3DRenderer(nn.Module):
         #         ret["comp_mask"] = rendered_mask.permute(1, 2, 0)
 
         return ret
-
+    @torch.no_grad()
     def animate_gs_model(
         self, gs_attr: GaussianAppOutput, query_points, smplx_data, debug=False
     ):
@@ -1063,7 +1064,7 @@ class GS3DRenderer(nn.Module):
                 gs_list.append(gs_copy)
 
         return gs_list, cano_gs_list
-
+    @torch.no_grad()
     def forward_gs_attr(self, x, query_points, smplx_data, debug=False, x_fine=None):
         """
         x: [N, C] Float[Tensor, "Np Cp"],
@@ -1080,7 +1081,7 @@ class GS3DRenderer(nn.Module):
         gs_attr: GaussianAppOutput = self.gs_net(x, query_points, x_fine)
 
         return gs_attr
-
+    @torch.no_grad()
     def get_query_points(self, smplx_data, device):
         with torch.no_grad():
             with torch.autocast(device_type=device.type, dtype=torch.float32):
@@ -1112,7 +1113,7 @@ class GS3DRenderer(nn.Module):
             pcl_embed.to(dtype=latent_feat.dtype), latent_feat, extra_info
         )
         return gs_feats
-
+    @torch.no_grad()
     def query_latent_feat(
         self,
         positions: Float[Tensor, "*B N1 3"],
@@ -1137,7 +1138,7 @@ class GS3DRenderer(nn.Module):
             )
 
         return gs_feats, positions, smplx_data
-
+    @torch.no_grad()
     def forward_single_batch(
         self,
         gs_list: list[GaussianModel],
@@ -1286,7 +1287,7 @@ class GS3DRenderer(nn.Module):
                     :, vidx : vidx + 1
                 ]  # e.g. body_pose: [1, N_v, 21, 3] -> [1, 1, 21, 3]
         return smpl_data_single_view
-
+    @torch.no_grad()
     def forward_gs(
         self,
         gs_hidden_features: Float[Tensor, "B Np Cp"],
@@ -1321,7 +1322,7 @@ class GS3DRenderer(nn.Module):
             gs_attr_list.append(gs_attr)
 
         return gs_attr_list, query_points, smplx_data
-
+    @torch.no_grad()
     def forward_animate_gs(
         self,
         gs_attr_list,
@@ -1394,7 +1395,7 @@ class GS3DRenderer(nn.Module):
         )  # [B, NV, H, W, 3] -> [B, NV, 1, H, W]
 
         return out
-
+    @torch.no_grad()
     def forward(
         self,
         gs_hidden_features: Float[Tensor, "B Np Cp"],

LHM/runners/infer/human_lrm.py

@@ -6,6 +6,7 @@
 # @Function      : Inference code for human_lrm model
 
 import argparse
+from omegaconf import OmegaConf
 import os
 import pdb
 import time
@@ -15,7 +16,7 @@ import numpy as np
 import spaces
 import torch
 from accelerate.logging import get_logger
-from omegaconf import OmegaConf
+
 from PIL import Image
 from tqdm.auto import tqdm
 

app.py

@@ -14,14 +14,21 @@
 
 
 import os
+os.system("rm -rf /data-nvme/zerogpu-offload/")
+import cv2
+import time
 from PIL import Image
 import numpy as np
 import gradio as gr
 import base64
 import spaces
+import torch
+torch._dynamo.config.disable = True
 import subprocess
 import os
-
+import argparse
+from omegaconf import OmegaConf
+from rembg import remove
 from engine.pose_estimation.pose_estimator import PoseEstimator
 from LHM.utils.face_detector import VGGHeadDetector
 from LHM.utils.hf_hub import wrap_model_hub
@@ -31,6 +38,29 @@ from LHM.runners.infer.utils import (
     prepare_motion_seqs,
     resize_image_keepaspect_np,
 )
+from engine.SegmentAPI.base import Bbox
+
+def get_bbox(mask):
+    height, width = mask.shape
+    pha = mask / 255.0
+    pha[pha < 0.5] = 0.0
+    pha[pha >= 0.5] = 1.0
+
+    # obtain bbox
+    _h, _w = np.where(pha == 1)
+
+    whwh = [
+        _w.min().item(),
+        _h.min().item(),
+        _w.max().item(),
+        _h.max().item(),
+    ]
+
+    box = Bbox(whwh)
+
+    # scale box to 1.05
+    scale_box = box.scale(1.1, width=width, height=height)
+    return scale_box
 
 def infer_preprocess_image(
     rgb_path,
@@ -244,6 +274,55 @@ def launch_env_not_compile_with_cuda():
     os.system("pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt251/download.html")
 
 
+def animation_infer(renderer, gs_model_list, query_points, smplx_params, render_c2ws, render_intrs, render_bg_colors):
+    '''Inference code avoid repeat forward.
+    '''
+    render_h, render_w = int(render_intrs[0, 0, 1, 2] * 2), int(
+        render_intrs[0, 0, 0, 2] * 2
+    )
+    # render target views
+    render_res_list = []
+    num_views = render_c2ws.shape[1]
+    start_time = time.time()
+
+    # render target views
+    render_res_list = []
+
+    for view_idx in range(num_views):
+        render_res = renderer.forward_animate_gs(
+            gs_model_list,
+            query_points,
+            renderer.get_single_view_smpl_data(smplx_params, view_idx),
+            render_c2ws[:, view_idx : view_idx + 1],
+            render_intrs[:, view_idx : view_idx + 1],
+            render_h,
+            render_w,
+            render_bg_colors[:, view_idx : view_idx + 1],
+        )
+        render_res_list.append(render_res)
+    print(
+        f"time elpased(animate gs model per frame):{(time.time() -  start_time)/num_views}"
+    )
+
+    out = defaultdict(list)
+    for res in render_res_list:
+        for k, v in res.items():
+            if isinstance(v[0], torch.Tensor):
+                out[k].append(v.detach().cpu())
+            else:
+                out[k].append(v)
+    for k, v in out.items():
+        # print(f"out key:{k}")
+        if isinstance(v[0], torch.Tensor):
+            out[k] = torch.concat(v, dim=1)
+            if k in ["comp_rgb", "comp_mask", "comp_depth"]:
+                out[k] = out[k][0].permute(
+                    0, 2, 3, 1
+                )  # [1, Nv, 3, H, W] -> [Nv, 3, H, W] - > [Nv, H, W, 3]
+        else:
+            out[k] = v
+    return out
+
 def assert_input_image(input_image):
     if input_image is None:
         raise gr.Error("No image selected or uploaded!")
@@ -273,7 +352,7 @@ def get_image_base64(path):
     return f"data:image/png;base64,{encoded_string}"
 
 
-def demo_lhm(pose_estimator, face_detector, lhm_model, cfg):
+def demo_lhm(pose_estimator, face_detector, lhm, cfg):
 
     @spaces.GPU
     def core_fn(image: str, video_params, working_dir):
@@ -311,6 +390,7 @@ def demo_lhm(pose_estimator, face_detector, lhm_model, cfg):
 
         print(image_raw, motion_seqs_dir, dump_image_dir, dump_video_path)
 
+        dump_tmp_dir = dump_image_dir
 
         shape_pose = pose_estimator(image_raw)
         assert shape_pose.is_full_body, f"The input image is illegal, {shape_pose.msg}"
@@ -458,7 +538,7 @@ def demo_lhm(pose_estimator, face_detector, lhm_model, cfg):
                         :, batch_i : batch_i + batch_size
                     ].to(device)
 
-                res = self.model.animation_infer(gs_model_list, query_points, batch_smplx_params,
+                res = lhm.animation_infer(gs_model_list, query_points, batch_smplx_params,
                     render_c2ws=motion_seq["render_c2ws"][
                         :, batch_i : batch_i + batch_size
                     ].to(device),
@@ -511,7 +591,6 @@ def demo_lhm(pose_estimator, face_detector, lhm_model, cfg):
             verbose=True,
         )
 
-
         # self.infer_single(
         #     image_path,
         #     motion_seqs_dir=motion_seqs_dir,
@@ -525,8 +604,6 @@ def demo_lhm(pose_estimator, face_detector, lhm_model, cfg):
         #     shape_param=shape_pose.beta,
         # )
 
-
-
         # status = spaces.GPU(infer_impl(
         #     gradio_demo_image=image_raw, 
         #     gradio_motion_file=smplx_params_dir, 
@@ -695,6 +772,7 @@ def launch_gradio_app():
 if __name__ == '__main__':
     # launch_pretrained()
     # launch_env_not_compile_with_cuda()
+    # os.system("rm -rf /data-nvme/zerogpu-offload/")
     launch_gradio_app()
 
 # import gradio as gr

requirements_lhm.txt

@@ -37,8 +37,8 @@ setuptools==74.0.0
 taming_transformers_rom1504==0.0.6
 timm==1.0.15
 
-https://download.pytorch.org/whl/cu121/torch-2.5.1%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=92af92c569de5da937dd1afb45ecfdd598ec1254cf2e49e3d698cb24d71aae14
-https://download.pytorch.org/whl/cu121/torchvision-0.20.1%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=304937b82c933d5155bd04d771f4b187273f67a76050bb4276b521f7e9b4c4e7
+# https://download.pytorch.org/whl/cu121/torch-2.5.1%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=92af92c569de5da937dd1afb45ecfdd598ec1254cf2e49e3d698cb24d71aae14
+# https://download.pytorch.org/whl/cu121/torchvision-0.20.1%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=304937b82c933d5155bd04d771f4b187273f67a76050bb4276b521f7e9b4c4e7
 # https://download.pytorch.org/whl/cu121/xformers-0.0.29.post1-cp310-cp310-manylinux_2_28_x86_64.whl#sha256=e213ff8123e20602bd486739ffee4013338b02f9d2e0e4635a2912750854fdbe
 
 tqdm==4.66.4