feat(inference): load Stage-2 training components

[Haruko386]

Summary

• add a component loader that overlays a saved Stage 2 U-Net and optional fine-tuned VAE on the original SD2 pipeline
• add --training_checkpoint to folder and dataset inference without changing legacy checkpoint behavior
• require the learned VAE when checkpoint metadata declares it, preventing silent inference with the wrong decoder
• default direct-SD2 component inference to native processing resolution when no explicit resolution is provided

Related Issue

N/A

Type of Change

• Bug fix
• New feature
• Documentation update
• Refactor
• CI / build change
• Other

Test Results

• D:\Anaconda\envs\roll\python.exe -m pytest tests -q - 3 passed
• D:\Anaconda\Scripts\ruff.exe check apdepth\util\checkpoint.py infer.py run.py tests - passed
• git diff --check - passed
• python run.py --help - passed
• python infer.py --help - passed

Summary by CodeRabbit

• New Features

  • Added optional --training_checkpoint support to inference and execution commands.
  • Stage-2 checkpoints can now provide fine-tuned U-Net and VAE components for depth processing.
  • When using a training checkpoint without specifying processing resolution, inputs are processed at their native resolution.

• Bug Fixes

  • Added validation and a clear error when checkpoint metadata expects a fine-tuned VAE that is unavailable.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 518ac326-202f-4495-89af-e0992f423a80

📥 Commits

Reviewing files that changed from the base of the PR and between 52470e2 and 034745f.

📒 Files selected for processing (4)

• apdepth/util/checkpoint.py
• infer.py
• run.py
• tests/test_training_checkpoint.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The change adds training-checkpoint loading for the U-Net and optional VAE. Both infer.py and run.py accept --training_checkpoint and pass loaded components to ApDepthPipeline. run.py uses native input resolution when no processing resolution is provided.

Changes

Training checkpoint support

Layer / File(s)	Summary
Checkpoint component loader and validation apdepth/util/checkpoint.py, tests/test_training_checkpoint.py	Adds load_training_components. It loads the U-Net, conditionally loads the VAE, validates required VAE files, and tests complete, legacy, and invalid checkpoints.
Inference checkpoint wiring infer.py	Adds --training_checkpoint and passes loaded component overrides to ApDepthPipeline.
Batch runner checkpoint wiring run.py	Adds --training_checkpoint, passes loaded components to ApDepthPipeline, and selects native-resolution processing when no resolution is specified.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant CLI as infer.py or run.py
  participant Loader as load_training_components
  participant Pipeline as ApDepthPipeline
  CLI->>Loader: Load components from --training_checkpoint
  Loader-->>CLI: Return U-Net and optional VAE
  CLI->>Pipeline: Pass component overrides to from_pretrained

Loading

Merge Risk: ⚪ Minimal · up to 03474

Training checkpoints load into both inference paths as intended, with native-resolution processing retained where applicable. The change is ready to merge.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 4 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: inference support for loading Stage-2 training components.
Description check	✅ Passed	The description follows the required template. It explains the changes, identifies the change type, and lists test results. The Related Issue section uses N/A, which is acceptable when no issue applie…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}