Shape mismatch in parameter_transform between TorchScript and MHR.from_files

[slnfox]

Hello, thanks for sharing this amazing project. I’ve noticed a difference in the tensor shapes of parameter_transform when comparing a model loaded via MHR.from_files against the provided TorchScript model mhr_mode.pt. When initializing the model normally with lod=1, the shape is [889, 321]. However, the TorchScript model shows a shape of [889, 249]. Is this mismatch intended? Could you share the script or the specific steps used to export the mhr_model.pt?

from pathlib import Path
import torch
from mhr.mhr import MHR

# Load the TorchScript model
ts_model = torch.jit.load('assets/mhr_model.pt')
ts_shape = ts_model.character_torch.parameter_transform.parameter_transform.shape
print(f"TorchScript shape: {ts_shape}") 
# Output: torch.Size([889, 249])

# Load via MHR method
mhr_model = MHR.from_files(Path("assets"), device=torch.device("cpu"), lod=1)
mhr_shape = mhr_model.character_torch.parameter_transform.parameter_transform.shape
print(f"MHR.from_files shape: {mhr_shape}")
# Output: torch.Size([889, 321])

pymomentum installed with pixi add pymomentum.

[JinlongYANG]

Hi @slnfox , thank you very much for pointing that out!

Sorry that we didn't notice that difference between the jit model and the model loaded from the model files. Actually, the jit model we provide is targeted for demo purpose, which requires some changes.

Nonetheless, the mismatch of character_torch.parameter_transform.parameter_transform.shape is not intentional. The shape difference 72 is due to facial expressions: The facial expression parameters are not integrated into parameter_transform. But the jit model has facial expression basis and they work properly in the forward method, just as the model created through MHR.from_files().

Yet as explained, you are encouraged to use MHR.from_files() to load the model for your applications if you are not simply trying the demos.

Does the mismatch you mentioned create any trouble in your application? If yes, we can look into it further to correct it.