Question regarding the return type of geometry_motion_vae.encode in trainers/unet_base.py

[maoxindong123]

Hi there,

First of all, thank you for your amazing work!

While reading the code, I encountered a minor confusion regarding the encode_point_map function in trainers/unet_base.py.

In the AutoencoderKL branch, .latent_dist.mode() is explicitly called to extract the tensor. However, in the SeperateAutoencoderKL and UnifyAutoencoderKL branches, the output of self.geometry_motion_vae.encode is directly assigned to latent, followed by an assert isinstance(latent, torch.Tensor).

elif self.config.vae_type == "SeperateAutoencoderKL":
    latent = self.geometry_motion_vae.encode(pmap_[i : i + chunk_size])
    # ...
elif self.config.vae_type == "UnifyAutoencoderKL":
    latent = self.geometry_motion_vae.encode( ... )
    
assert isinstance(latent, torch.Tensor)

My questions are:

Does self.geometry_motion_vae.encode return a DiagonalGaussianDistribution (or a similar output object) in these two elif branches? If so, the assert statement might fail here.

Assuming it does return a DiagonalGaussianDistribution, should we append .mode() or .sample() to the encode call during training in these specific branches?

Thank you in advance for your time and clarification!

[RuijieZhu94]

Hi @maoxindong123, thanks for pointing it out. I used .sample() in my original code, but now I feel that .mode() would be more stable here. I'm not entirely sure.

[maoxindong123]

Hi @RuijieZhu94, thanks for the quick reply and clarification!

While looking further into this, I found another potential issue right in the UnifyAutoencoderKL branch. Currently, the code passes three separate tensors to geometry_motion_vae.encode:

elif self.config.vae_type == "UnifyAutoencoderKL":
    latent = self.geometry_motion_vae.encode(
        pmap_[i : i + chunk_size],
        sflow_[i : i + chunk_size],
        vmask_[i : i + chunk_size],
    ).mode()

This throws an error because geometry_motion_vae.encode expects a single tensor input, but here it receives three (pmap_, sflow_, and vmask_).

Thanks again for your help!

[RuijieZhu94]

Hi @maoxindong123, thank you for pointing it out. I use copilot to merge my code, it may lead to these errors. Sorry for that.