Update app.py

app.py

@@ -347,7 +347,7 @@ def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[s
         # GLB 변환
         with torch.inference_mode():
             try:
-                # 모든 텐서를 CUDA로 이동
+                # 모든 텐서를 CUDA로 이동하고 gradient 활성화
                 device = torch.device('cuda:0')
                 
                 # Gaussian 텐서들을 변환
@@ -355,12 +355,13 @@ def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[s
                     if hasattr(gs, attr_name):
                         tensor = getattr(gs, attr_name)
                         if torch.is_tensor(tensor):
-                            new_tensor = tensor.detach().clone().float().to(device)
+                            # gradient 계산이 필요한 텐서로 변환
+                            new_tensor = tensor.detach().clone().float().to(device).requires_grad_(True)
                             setattr(gs, attr_name, new_tensor)
                 
                 # Mesh 텐서들을 변환
                 if hasattr(mesh, 'vertices') and torch.is_tensor(mesh.vertices):
-                    mesh.vertices = mesh.vertices.detach().clone().float().to(device)
+                    mesh.vertices = mesh.vertices.detach().clone().float().to(device).requires_grad_(True)
                 if hasattr(mesh, 'faces') and torch.is_tensor(mesh.faces):
                     mesh.faces = mesh.faces.detach().clone().long().to(device)
                 
@@ -370,11 +371,14 @@ def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[s
                         continue
                     attr = getattr(mesh, attr_name)
                     if torch.is_tensor(attr):
-                        setattr(mesh, attr_name, attr.to(device))
+                        if attr.dtype in [torch.float32, torch.float64]:
+                            setattr(mesh, attr_name, attr.to(device).requires_grad_(True))
+                        else:
+                            setattr(mesh, attr_name, attr.to(device))
                 
-                print("Device check before GLB conversion:")
-                print(f"Gaussian xyz device: {gs._xyz.device}")
-                print(f"Mesh vertices device: {mesh.vertices.device}")
+                print("Device and gradient check before GLB conversion:")
+                print(f"Gaussian xyz device: {gs._xyz.device}, requires_grad: {gs._xyz.requires_grad}")
+                print(f"Mesh vertices device: {mesh.vertices.device}, requires_grad: {mesh.vertices.requires_grad}")
                 
                 # GLB 변환
                 glb = postprocessing_utils.to_glb(
@@ -387,11 +391,11 @@ def extract_glb(state: dict, mesh_simplify: float, texture_size: int) -> Tuple[s
                 
             except Exception as e:
                 print(f"Error during GLB conversion: {str(e)}")
-                # 디바이스 정보 출력
+                # 디바이스와 gradient 정보 출력
                 if hasattr(gs, '_xyz'):
-                    print(f"Gaussian xyz device: {gs._xyz.device}")
+                    print(f"Gaussian xyz device: {gs._xyz.device}, requires_grad: {gs._xyz.requires_grad}")
                 if hasattr(mesh, 'vertices'):
-                    print(f"Mesh vertices device: {mesh.vertices.device}")
+                    print(f"Mesh vertices device: {mesh.vertices.device}, requires_grad: {mesh.vertices.requires_grad}")
                 return None, None
             
         if glb is None: