update debug

LHM/models/encoders/dinov2_fusion_wrapper.py

@@ -19,9 +19,9 @@ import kornia
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from accelerate.logging import get_logger
+# from accelerate.logging import get_logger
 
-logger = get_logger(__name__)
+# logger = get_logger(__name__)
 
 
 class DPTHead(nn.Module):
@@ -126,7 +126,7 @@ class Dinov2FusionWrapper(nn.Module):
             self._freeze()
 
     def _freeze(self):
-        logger.warning(f"======== Freezing Dinov2FusionWrapper ========")
+        # logger.warning(f"======== Freezing Dinov2FusionWrapper ========")
         self.model.eval()
         for name, param in self.model.named_parameters():
             param.requires_grad = False
@@ -170,7 +170,7 @@ class Dinov2FusionWrapper(nn.Module):
 
         dinov2_hub = import_module(".dinov2.hub.backbones", package=__package__)
         model_fn = getattr(dinov2_hub, model_name)
-        logger.debug(f"Modulation dim for Dinov2 is {modulation_dim}.")
+        # logger.debug(f"Modulation dim for Dinov2 is {modulation_dim}.")
         model = model_fn(modulation_dim=modulation_dim, pretrained=pretrained)
         return model
 

LHM/models/encoders/sapiens_warpper.py

@@ -18,10 +18,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision
-from accelerate.logging import get_logger
+# from accelerate.logging import get_logger
 from tqdm import tqdm
 
-logger = get_logger(__name__)
+# logger = get_logger(__name__)
 
 timings = {}
 BATCH_SIZE = 64
@@ -188,7 +188,7 @@ class SapiensWrapper(nn.Module):
     @staticmethod
     def _build_sapiens(model_name: str, pretrained: bool = True):
 
-        logger.debug(f"Using Sapiens model: {model_name}")
+        # logger.debug(f"Using Sapiens model: {model_name}")
         USE_TORCHSCRIPT = "_torchscript" in model_name
 
         # build the model from a checkpoint file
@@ -201,7 +201,7 @@ class SapiensWrapper(nn.Module):
         return model
 
     def _freeze(self):
-        logger.warning(f"======== Freezing Sapiens Model ========")
+        # logger.warning(f"======== Freezing Sapiens Model ========")
         self.model.eval()
         for name, param in self.model.named_parameters():
             param.requires_grad = False

LHM/models/modeling_human_lrm.py

@@ -14,7 +14,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from accelerate.logging import get_logger
+# from accelerate.logging import get_logger
 from diffusers.utils import is_torch_version
 
 from LHM.models.arcface_utils import ResNetArcFace
@@ -29,7 +29,7 @@ from .embedder import CameraEmbedder
 from .rendering.synthesizer import TriplaneSynthesizer
 from .transformer import TransformerDecoder
 
-logger = get_logger(__name__)
+# logger = get_logger(__name__)
 
 
 class ModelHumanLRM(nn.Module):
@@ -212,42 +212,42 @@ class ModelHumanLRM(nn.Module):
         if encoder_type == "dino":
             from .encoders.dino_wrapper import DinoWrapper
 
-            logger.info("Using DINO as the encoder")
+            # logger.info("Using DINO as the encoder")
             return DinoWrapper
         elif encoder_type == "dinov2":
             from .encoders.dinov2_wrapper import Dinov2Wrapper
 
-            logger.info("Using DINOv2 as the encoder")
+            # logger.info("Using DINOv2 as the encoder")
             return Dinov2Wrapper
         elif encoder_type == "dinov2_unet":
             from .encoders.dinov2_unet_wrapper import Dinov2UnetWrapper
 
-            logger.info("Using Dinov2Unet as the encoder")
+            # logger.info("Using Dinov2Unet as the encoder")
             return Dinov2UnetWrapper
         elif encoder_type == "resunet":
             from .encoders.xunet_wrapper import XnetWrapper
 
-            logger.info("Using XnetWrapper as the encoder")
+            # logger.info("Using XnetWrapper as the encoder")
             return XnetWrapper
         elif encoder_type == "dinov2_featup":
             from .encoders.dinov2_featup_wrapper import Dinov2FeatUpWrapper
 
-            logger.info("Using Dinov2FeatUpWrapper as the encoder")
+            # logger.info("Using Dinov2FeatUpWrapper as the encoder")
             return Dinov2FeatUpWrapper
         elif encoder_type == "dinov2_dpt":
             from .encoders.dinov2_dpt_wrapper import Dinov2DPTWrapper
 
-            logger.info("Using Dinov2DPTWrapper as the encoder")
+            # logger.info("Using Dinov2DPTWrapper as the encoder")
             return Dinov2DPTWrapper
         elif encoder_type == "dinov2_fusion":
             from .encoders.dinov2_fusion_wrapper import Dinov2FusionWrapper
 
-            logger.info("Using Dinov2FusionWrapper as the encoder")
+            # logger.info("Using Dinov2FusionWrapper as the encoder")
             return Dinov2FusionWrapper
         elif encoder_type == "sapiens":
             from .encoders.sapiens_warpper import SapiensWrapper
 
-            logger.info("Using Sapiens as the encoder")
+            # logger.info("Using Sapiens as the encoder")
             return SapiensWrapper
 
     def forward_transformer(self, image_feats, camera_embeddings, query_points):
@@ -560,10 +560,10 @@ class ModelHumanLRM(nn.Module):
             },
         ]
 
-        logger.info("======== Weight Decay Parameters ========")
-        logger.info(f"Total: {len(decay_params)}")
-        logger.info("======== No Weight Decay Parameters ========")
-        logger.info(f"Total: {len(no_decay_params)}")
+        # logger.info("======== Weight Decay Parameters ========")
+        # logger.info(f"Total: {len(decay_params)}")
+        # logger.info("======== No Weight Decay Parameters ========")
+        # logger.info(f"Total: {len(no_decay_params)}")
 
         print(f"Total Params: {len(no_decay_params) + len(decay_params)}")
 
@@ -936,10 +936,10 @@ class ModelHumanLRMSapdinoBodyHeadSD3_5(ModelHumanLRMSapdinoBodyHeadSD3):
             },
         ]
 
-        logger.info("======== Weight Decay Parameters ========")
-        logger.info(f"Total: {len(decay_params)}")
-        logger.info("======== No Weight Decay Parameters ========")
-        logger.info(f"Total: {len(no_decay_params)}")
+        # logger.info("======== Weight Decay Parameters ========")
+        # logger.info(f"Total: {len(decay_params)}")
+        # logger.info("======== No Weight Decay Parameters ========")
+        # logger.info(f"Total: {len(no_decay_params)}")
 
         print(f"Total Params: {len(no_decay_params) + len(decay_params)}")
 

LHM/models/rendering/smpl_x_voxel_dense_sampling.py

@@ -318,8 +318,10 @@ class SMPLX_Mesh(object):
         return joint_offset
 
     def get_subdivider(self, subdivide_num):
-        vert = self.layer["neutral"].v_template.float().cuda()
-        face = torch.LongTensor(self.face).cuda()
+        # vert = self.layer["neutral"].v_template.float().cuda()
+        # face = torch.LongTensor(self.face).cuda()
+        vert = self.layer["neutral"].v_template.float()
+        face = torch.LongTensor(self.face)
         mesh = Meshes(vert[None, :, :], face[None, :, :])
 
         if subdivide_num > 0:
@@ -419,7 +421,8 @@ class SMPLX_Mesh(object):
         normal = (
             Meshes(
                 verts=mesh_neutral_pose[None, :, :],
-                faces=torch.LongTensor(self.face_upsampled).cuda()[None, :, :],
+                # faces=torch.LongTensor(self.face_upsampled).cuda()[None, :, :],
+                faces=torch.LongTensor(self.face_upsampled)[None, :, :],
             )
             .verts_normals_packed()
             .reshape(self.vertex_num_upsampled, 3)
@@ -537,11 +540,14 @@ class SMPLXVoxelMeshModel(nn.Module):
     ):
         """Smooth KNN to handle skirt deformation."""
 
-        lbs_weights = lbs_weights.cuda()
+        # lbs_weights = lbs_weights.cuda()
+        lbs_weights = lbs_weights
 
         dist = knn_points(
-            voxel_v.unsqueeze(0).cuda(),
-            template_v.unsqueeze(0).cuda(),
+            # voxel_v.unsqueeze(0).cuda(),
+            # template_v.unsqueeze(0).cuda(),
+            voxel_v.unsqueeze(0),
+            template_v.unsqueeze(0),
             K=1,
             return_nn=True,
         )
@@ -555,8 +561,10 @@ class SMPLXVoxelMeshModel(nn.Module):
         # Smooth Skinning
 
         knn_dis = knn_points(
-            voxel_v.unsqueeze(0).cuda(),
-            voxel_v.unsqueeze(0).cuda(),
+            # voxel_v.unsqueeze(0).cuda(),
+            # voxel_v.unsqueeze(0).cuda(),
+            voxel_v.unsqueeze(0),
+            voxel_v.unsqueeze(0),
             K=smooth_k + 1,
             return_nn=True,
         )
@@ -667,7 +675,7 @@ class SMPLXVoxelMeshModel(nn.Module):
         )
 
         coordinates = coordinates.view(-1, 3).float()
-        coordinates = coordinates.cuda()
+        # coordinates = coordinates.cuda()
 
         if os.path.exists(f"./pretrained_models/voxel_grid/voxel_{voxel_size}.pth"):
             print(f"load voxel_grid voxel_{voxel_size}.pth")
@@ -722,12 +730,15 @@ class SMPLXVoxelMeshModel(nn.Module):
         smpl_x = self.smpl_x
 
         # using KNN to query subdivided mesh
-        dense_pts = self.dense_pts.cuda()
+        # dense_pts = self.dense_pts.cuda()
+        dense_pts = self.dense_pts
         template_verts = self.smplx_layer.v_template
 
         nn_vertex_idxs = knn_points(
-            dense_pts.unsqueeze(0).cuda(),
-            template_verts.unsqueeze(0).cuda(),
+            # dense_pts.unsqueeze(0).cuda(),
+            # template_verts.unsqueeze(0).cuda(),
+            dense_pts.unsqueeze(0),
+            template_verts.unsqueeze(0),
             K=1,
             return_nn=True,
         ).idx
@@ -1046,7 +1057,8 @@ class SMPLXVoxelMeshModel(nn.Module):
         )  # [B, 54, 3]
         # smplx pose-dependent vertex offset
         pose = (
-            axis_angle_to_matrix(pose) - torch.eye(3)[None, None, :, :].float().cuda()
+            # axis_angle_to_matrix(pose) - torch.eye(3)[None, None, :, :].float().cuda()
+            axis_angle_to_matrix(pose) - torch.eye(3)[None, None, :, :].float()
         ).view(batch_size, (self.smpl_x.joint_num - 1) * 9)
         # (B, 54 * 9) x (54*9, V)
 
@@ -1499,7 +1511,8 @@ def read_smplx_param(smplx_data_root, shape_param_file, batch_size=1, device="cu
             osp.join(data_root_path, "cam_params", str(frame_idx) + ".json")
         ) as f:
             cam_param = {
-                k: torch.FloatTensor(v).cuda() for k, v in json.load(f).items()
+                # k: torch.FloatTensor(v).cuda() for k, v in json.load(f).items()
+                k: torch.FloatTensor(v) for k, v in json.load(f).items()
             }
             cam_param_list.append(cam_param)
 
@@ -1668,7 +1681,8 @@ def generate_smplx_point():
         for k, v in data.items():
             if k in smplx_keys:
                 # print(k, v.shape)
-                smplx_params[k] = data[k].unsqueeze(0).cuda()
+                # smplx_params[k] = data[k].unsqueeze(0).cuda()
+                smplx_params[k] = data[k].unsqueeze(0)
         return smplx_params
 
     def sample_one(data):

LHM/models/transformer.py

@@ -11,10 +11,10 @@ from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
-from accelerate.logging import get_logger
+# from accelerate.logging import get_logger
 from diffusers.utils import is_torch_version
 
-logger = get_logger(__name__)
+# logger = get_logger(__name__)
 
 
 class TransformerDecoder(nn.Module):
@@ -106,7 +106,7 @@ class TransformerDecoder(nn.Module):
             ), f"Condition and modulation are not supported for BasicBlock"
             from .block import BasicBlock
 
-            logger.debug(f"Using BasicBlock")
+            # logger.debug(f"Using BasicBlock")
             return partial(BasicBlock, inner_dim=inner_dim)
         elif self.block_type == "cond":
             assert (
@@ -117,10 +117,10 @@ class TransformerDecoder(nn.Module):
             ), f"Modulation dimension is not supported for ConditionBlock"
             from .block import ConditionBlock
 
-            logger.debug(f"Using ConditionBlock")
+            # logger.debug(f"Using ConditionBlock")
             return partial(ConditionBlock, inner_dim=inner_dim, cond_dim=cond_dim)
         elif self.block_type == "mod":
-            logger.error(f"modulation without condition is not implemented")
+            # logger.error(f"modulation without condition is not implemented")
             raise NotImplementedError(
                 f"modulation without condition is not implemented"
             )
@@ -130,7 +130,7 @@ class TransformerDecoder(nn.Module):
             ), f"Condition and modulation dimensions must be specified for ConditionModulationBlock"
             from .block import ConditionModulationBlock
 
-            logger.debug(f"Using ConditionModulationBlock")
+            # logger.debug(f"Using ConditionModulationBlock")
             return partial(
                 ConditionModulationBlock,
                 inner_dim=inner_dim,
@@ -138,25 +138,25 @@ class TransformerDecoder(nn.Module):
                 mod_dim=mod_dim,
             )
         elif self.block_type == "cogvideo_cond":
-            logger.debug(f"Using CogVideoXBlock")
+            # logger.debug(f"Using CogVideoXBlock")
             from LHM.models.transformer_dit import CogVideoXBlock
 
             # assert inner_dim == cond_dim, f"inner_dim:{inner_dim}, cond_dim:{cond_dim}"
             return partial(CogVideoXBlock, dim=inner_dim, attention_bias=True)
         elif self.block_type == "sd3_cond":
-            logger.debug(f"Using SD3JointTransformerBlock")
+            # logger.debug(f"Using SD3JointTransformerBlock")
             from LHM.models.transformer_dit import SD3JointTransformerBlock
 
             return partial(SD3JointTransformerBlock, dim=inner_dim, qk_norm="rms_norm")
         elif self.block_type == "sd3_mm_cond":
-            logger.debug(f"Using SD3MMJointTransformerBlock")
+            # logger.debug(f"Using SD3MMJointTransformerBlock")
             from LHM.models.transformer_dit import SD3MMJointTransformerBlock
 
             return partial(
                 SD3MMJointTransformerBlock, dim=inner_dim, qk_norm="rms_norm"
             )
         elif self.block_type == "sd3_mm_bh_cond":
-            logger.debug(f"Using SD3MMJointTransformerBlock")
+            # logger.debug(f"Using SD3MMJointTransformerBlock")
             from LHM.models.transformer_dit import SD3BodyHeadMMJointTransformerBlock
 
             return partial(