Skip to content

fix: DINOv3 layer path and hardcoded .cuda() device hardening - #178

Open
adwaitm1301 wants to merge 1 commit into
microsoft:mainfrom
adwaitm1301:fix/dinov3-cuda-device-hardening
Open

fix: DINOv3 layer path and hardcoded .cuda() device hardening#178
adwaitm1301 wants to merge 1 commit into
microsoft:mainfrom
adwaitm1301:fix/dinov3-cuda-device-hardening

Conversation

@adwaitm1301

Copy link
Copy Markdown

Three fixes dug up from reading issue #161 and cross-referencing against existing PRs.

1. DINOv3 layer path (both copies) — Closes #161

The extract_features method hardcodes self.model.layer, but DINOv3ViTModel from HuggingFace transformers moved the encoder to self.model.model.layer (v5.0+) and later to self.model.encoder.layer (v5.8+). The fix uses a fallback chain that tries all three paths and raises a clear AttributeError if none match.

Two copies of this code exist — modules/image_feature_extractor.py AND rainers/.../image_conditioned.py. All existing PRs (#144, #148, #156) only fixed the modules copy. This PR fixes both.

2. Hardcoded .cuda() calls (both copies)

Every orch.stack(image).cuda() and self.transform(image).cuda() was hardcoded to CUDA, meaning the feature extractor crashes on CPU or multi-GPU setups. Now uses .to(self.device) with a device property added to both DinoV2FeatureExtractor and DinoV3FeatureExtractor.

Also removed the unconditional .cuda() call in ImageConditionedMixin._init_image_cond_model(), which destroyed device placement for the training code path.

3. Device type consistency + renderer hardening

  • _device = 'cpu' — stored as a string, breaking callers that check isinstance(device, torch.device). Fixed to orch.device('cpu') in both pipeline files.
  • latlong_to_cubemap() in pbr_mesh_renderer.py — hardcoded device='cuda' for all intermediate tensors. Changed to latlong_map.device so it follows the input.

Three interconnected bugs fixed across 5 files.

1. DINOv3 layer path (modules + trainers image_feature_extractor)
   - self.model.layer -> fallback chain to support all transformers versions
   - The trainers copy was missed by every existing DINOv3 PR

2. Hardcoded .cuda() calls in both image_feature_extractor files
   - Replaced .cuda() with .to(self.device) for CPU and multi-GPU support
   - Added device property to both extractor classes

3. Device type consistency + renderer hardening
   - _device = 'cpu' -> torch.device('cpu') in both pipeline files
   - latlong_to_cubemap: device='cuda' -> latlong_map.device
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Total Program Failure: Trellis2 DINOv3 extractor cannot load the HuggingFace DINOv3 model

1 participant