fix rot

app.py

@@ -15,7 +15,7 @@ import gradio as gr
 
 # download checkpoints
 from huggingface_hub import hf_hub_download
-ckpt_path = hf_hub_download(repo_id="ashawkey/LGM", filename="model_fp16.safetensors")
+ckpt_path = hf_hub_download(repo_id="ashawkey/LGM", filename="model_fp16_fixrot.safetensors")
 
 subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl"))
 

core/models.py

@@ -40,7 +40,7 @@ class LGM(nn.Module):
         self.pos_act = lambda x: x.clamp(-1, 1)
         self.scale_act = lambda x: 0.1 * F.softplus(x)
         self.opacity_act = lambda x: torch.sigmoid(x)
-        self.rot_act = F.normalize
+        self.rot_act = lambda x: F.normalize(x, dim=-1)
         self.rgb_act = lambda x: 0.5 * torch.tanh(x) + 0.5 # NOTE: may use sigmoid if train again
 
         # LPIPS loss