Update hf

2D_Stage/tuneavideo/pipelines/pipeline_tuneavideo.py

@@ -360,12 +360,12 @@ class TuneAVideoPipeline(DiffusionPipeline):
             if isinstance(generator, list):
                 shape = (1,) + shape[1:]
                 latents = [
-                    torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype)
+                    torch.randn(shape, device=rand_device, dtype=dtype)
                     for i in range(batch_size)
                 ]
                 latents = torch.cat(latents, dim=0).to(device)
             else:
-                latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
+                latents = torch.randn(shape, device=rand_device, dtype=dtype).to(device)
         else:
             if latents.shape != shape:
                 raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")

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/models/transformers/attention.py

@@ -75,8 +75,8 @@ class MemoryEfficientAttentionMixin:
         # Any children which exposes the set_use_memory_efficient_attention_xformers method
         # gets the message
         def fn_recursive_set_mem_eff(module: torch.nn.Module):
-            if hasattr(module, "set_use_memory_efficient_attention_xformers"):
-                module.set_use_memory_efficient_attention_xformers(valid, attention_op)
+            #if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+            #    module.set_use_memory_efficient_attention_xformers(valid, attention_op)
 
             for child in module.children():
                 fn_recursive_set_mem_eff(child)

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)(