Update hf

3D_Stage/lrm/models/camera.py

@@ -29,5 +29,7 @@ class LinearCameraEmbedder(BaseModule):
             cond_tensors.append(cond.view(*cond.shape[:2], -1))
         cond_tensor = torch.cat(cond_tensors, dim=-1)
         assert cond_tensor.shape[-1] == self.cfg.in_channels
+        #self.linear = self.linear.to("cuda")
+        #print(self.linear.weight.device, cond_tensor.device)
         embedding = self.linear(cond_tensor)
         return embedding

3D_Stage/lrm/models/exporters/mesh_exporter.py

@@ -101,6 +101,7 @@ class MeshExporter(Exporter):
     def get_texture_maps(
         self, scene_code: torch.Tensor, mesh: Mesh
     ) -> Dict[str, torch.Tensor]:
+        #self.ctx = NVDiffRasterizerContext(self.cfg.context_type, self.device)
         assert mesh.has_uv
         # clip space transform
         uv_clip = mesh.v_tex * 2.0 - 1.0

3D_Stage/lrm/models/renderers/triplane_dmtet.py

@@ -44,6 +44,8 @@ class TriplaneDMTetRenderer(BaseRenderer):
         assert self.cfg.feature_reduction in ["concat", "mean"]
 
         self.ctx = NVDiffRasterizerContext(self.cfg.context_type, self.device)
+        #self.ctx = None
+        
         self.isosurface_helper = MarchingTetrahedraHelper(
             self.cfg.isosurface_resolution,
             os.path.join(self.cfg.tet_dir, f"{self.cfg.isosurface_resolution}_tets.npz"),
@@ -130,6 +132,7 @@ class TriplaneDMTetRenderer(BaseRenderer):
         background_color: Optional[Float[Tensor, "3"]],
         extra_sdf_query: Any = None,
     ) -> Dict[str, Tensor]:
+        #self.ctx = NVDiffRasterizerContext(self.cfg.context_type, self.device)
         Nv = mvp_mtx.shape[0]
 
         out = {}

3D_Stage/lrm/systems/multiview_lrm.py

@@ -62,19 +62,21 @@ class MultiviewLRM(BaseSystem):
 
     def configure(self):
         super().configure()
+        device = torch.device("cuda")
         self.image_tokenizer = lrm.find(self.cfg.image_tokenizer_cls)(
             self.cfg.image_tokenizer
-        )
+        ).to(device)
         if self.cfg.image_tokenizer.modulation:
             self.camera_embedder = lrm.find(self.cfg.camera_embedder_cls)(
                 self.cfg.camera_embedder
-            )
-        self.tokenizer = lrm.find(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
+            ).to(device)
+        #print(device, self.camera_embedder.linear.weight.device)
+        self.tokenizer = lrm.find(self.cfg.tokenizer_cls)(self.cfg.tokenizer).to(device)
         self.backbone = lrm.find(self.cfg.backbone_cls)(self.cfg.backbone)
         self.post_processor = lrm.find(self.cfg.post_processor_cls)(
             self.cfg.post_processor
         )
-        self.decoder = lrm.find(self.cfg.decoder_cls)(self.cfg.decoder)
+        self.decoder = lrm.find(self.cfg.decoder_cls)(self.cfg.decoder).to(device)
         self.material = lrm.find(self.cfg.material_cls)(self.cfg.material)
         self.background = lrm.find(self.cfg.background_cls)(self.cfg.background)
         self.renderer = lrm.find(self.cfg.renderer_cls)(

3D_Stage/lrm/utils/rasterize.py

@@ -7,8 +7,8 @@ from .typing import *
 class NVDiffRasterizerContext:
     def __init__(self, context_type: str, device: torch.device) -> None:
         self.device = device
-        self.ctx = None
-        #self.ctx = self.initialize_context(context_type, device)
+        #self.ctx = None
+        self.ctx = self.initialize_context(context_type, device)
 
     def initialize_context(
         self, context_type: str, device: torch.device

app.py

@@ -309,7 +309,7 @@ class Inference3D_API:
     def __init__(self, device="cuda"):
         self.cfg = load_config("3D_Stage/configs/infer.yaml", makedirs=False)
         print("Loading system")
-        self.device = device
+        self.device = torch.device(device)
         self.cfg.system.weights = self.cfg.system.weights.replace("./", "./3D_Stage/")
         self.cfg.system.image_tokenizer.pretrained_model_name_or_path = \
             self.cfg.system.image_tokenizer.pretrained_model_name_or_path.replace("./", "./3D_Stage/")
@@ -318,8 +318,9 @@ class Inference3D_API:
         self.system = lrm.find(self.cfg.system_cls)(self.cfg.system).to(self.device)
         self.system.eval()
 
-    @spaces.GPU
+    @spaces.GPU(duration=120)
     def process_images(self, img_input0, img_input1, img_input2, img_input3, back_proj, smooth_iter):
+        #self.system = self.system.to("cuda")
         meta = json.load(open("./3D_Stage/material/meta.json"))
         c2w_cond = [np.array(loc["transform_matrix"]) for loc in meta["locations"]]
         c2w_cond = torch.from_numpy(np.stack(c2w_cond, axis=0)).float()[None].to(self.device)