diff --git a/trellis2/modules/image_feature_extractor.py b/trellis2/modules/image_feature_extractor.py index c3cb515a..7140d608 100644 --- a/trellis2/modules/image_feature_extractor.py +++ b/trellis2/modules/image_feature_extractor.py @@ -19,6 +19,10 @@ def __init__(self, model_name: str): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + def to(self, device): self.model.to(device) @@ -46,11 +50,11 @@ def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tenso image = [i.resize((518, 518), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") - image = self.transform(image).cuda() + image = self.transform(image).to(self.device) features = self.model(image, is_training=True)['x_prenorm'] patchtokens = F.layer_norm(features, features.shape[-1:]) return patchtokens @@ -69,6 +73,10 @@ def __init__(self, model_name: str, image_size=512): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) + @property + def device(self) -> torch.device: + return self.model.device + def to(self, device): self.model.to(device) @@ -83,7 +91,14 @@ def extract_features(self, image: torch.Tensor) -> torch.Tensor: hidden_states = self.model.embeddings(image, bool_masked_pos=None) position_embeddings = self.model.rope_embeddings(image) - for i, layer_module in enumerate(self.model.layer): + layers = getattr(self.model, 'layer', None) or getattr(getattr(self.model, 'model', None), 'layer', None) or getattr(getattr(self.model, 'encoder', None), 'layer', None) + if layers is None: + raise AttributeError( + "Cannot locate DINOv3 encoder layers: expected " + "self.model.layer, self.model.model.layer, or " + "self.model.encoder.layer" + ) + for i, layer_module in enumerate(layers): hidden_states = layer_module( hidden_states, position_embeddings=position_embeddings, @@ -109,10 +124,10 @@ def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tenso image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") - image = self.transform(image).cuda() + image = self.transform(image).to(self.device) features = self.extract_features(image) return features diff --git a/trellis2/pipelines/trellis2_image_to_3d.py b/trellis2/pipelines/trellis2_image_to_3d.py index a7b84b98..af28a6d8 100644 --- a/trellis2/pipelines/trellis2_image_to_3d.py +++ b/trellis2/pipelines/trellis2_image_to_3d.py @@ -76,7 +76,7 @@ def __init__( 'roughness': slice(4, 5), 'alpha': slice(5, 6), } - self._device = 'cpu' + self._device = torch.device('cpu') @classmethod def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trellis2ImageTo3DPipeline": @@ -112,7 +112,7 @@ def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trel 'roughness': slice(4, 5), 'alpha': slice(5, 6), } - pipeline._device = 'cpu' + pipeline._device = torch.device('cpu') return pipeline diff --git a/trellis2/pipelines/trellis2_texturing.py b/trellis2/pipelines/trellis2_texturing.py index c184b5e7..a9d52f1d 100755 --- a/trellis2/pipelines/trellis2_texturing.py +++ b/trellis2/pipelines/trellis2_texturing.py @@ -63,7 +63,7 @@ def __init__( 'roughness': slice(4, 5), 'alpha': slice(5, 6), } - self._device = 'cpu' + self._device = torch.device('cpu') @classmethod def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trellis2TexturingPipeline": @@ -92,7 +92,7 @@ def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trel 'roughness': slice(4, 5), 'alpha': slice(5, 6), } - pipeline._device = 'cpu' + pipeline._device = torch.device('cpu') return pipeline def to(self, device: torch.device) -> None: diff --git a/trellis2/renderers/pbr_mesh_renderer.py b/trellis2/renderers/pbr_mesh_renderer.py index b3b46553..0ac8ad72 100644 --- a/trellis2/renderers/pbr_mesh_renderer.py +++ b/trellis2/renderers/pbr_mesh_renderer.py @@ -20,10 +20,11 @@ def cube_to_dir(s, x, y): def latlong_to_cubemap(latlong_map, res): if 'dr' not in globals(): import nvdiffrast.torch as dr - cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda') + device = latlong_map.device + cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device=device) for s in range(6): - gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), - torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'), + gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device=device), + torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device=device), indexing='ij') v = F.normalize(cube_to_dir(s, gx, gy), dim=-1) diff --git a/trellis2/trainers/flow_matching/mixins/image_conditioned.py b/trellis2/trainers/flow_matching/mixins/image_conditioned.py index 13b932eb..f8cab411 100644 --- a/trellis2/trainers/flow_matching/mixins/image_conditioned.py +++ b/trellis2/trainers/flow_matching/mixins/image_conditioned.py @@ -21,6 +21,10 @@ def __init__(self, model_name: str): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + def to(self, device): self.model.to(device) @@ -48,11 +52,11 @@ def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tenso image = [i.resize((518, 518), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") - image = self.transform(image).cuda() + image = self.transform(image).to(self.device) features = self.model(image, is_training=True)['x_prenorm'] patchtokens = F.layer_norm(features, features.shape[-1:]) return patchtokens @@ -71,6 +75,10 @@ def __init__(self, model_name: str, image_size=512): transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) + @property + def device(self) -> torch.device: + return self.model.device + def to(self, device): self.model.to(device) @@ -85,7 +93,14 @@ def extract_features(self, image: torch.Tensor) -> torch.Tensor: hidden_states = self.model.embeddings(image, bool_masked_pos=None) position_embeddings = self.model.rope_embeddings(image) - for i, layer_module in enumerate(self.model.layer): + layers = getattr(self.model, 'layer', None) or getattr(getattr(self.model, 'model', None), 'layer', None) or getattr(getattr(self.model, 'encoder', None), 'layer', None) + if layers is None: + raise AttributeError( + "Cannot locate DINOv3 encoder layers: expected " + "self.model.layer, self.model.model.layer, or " + "self.model.encoder.layer" + ) + for i, layer_module in enumerate(layers): hidden_states = layer_module( hidden_states, position_embeddings=position_embeddings, @@ -111,14 +126,14 @@ def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tenso image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image] image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image] image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image] - image = torch.stack(image).cuda() + image = torch.stack(image).to(self.device) else: raise ValueError(f"Unsupported type of image: {type(image)}") - image = self.transform(image).cuda() + image = self.transform(image).to(self.device) features = self.extract_features(image) return features - + class ImageConditionedMixin: """ @@ -138,7 +153,6 @@ def _init_image_cond_model(self): """ with dist_utils.local_master_first(): self.image_cond_model = globals()[self.image_cond_model_config['name']](**self.image_cond_model_config.get('args', {})) - self.image_cond_model.cuda() @torch.no_grad() def encode_image(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor: