Skip to content

Conversation

@NielsRogge
Copy link

@NielsRogge NielsRogge commented May 12, 2025

Hi @wenbowen123, congrats on this work!

This PR ensures the FoundationStereo model can be pushed and loaded to and from the 🤗 hub.

It equips the main model class with:

  • from_pretrained, save_pretrainedandpush_to_hub capabilities
  • safetensors for weights serialization (if you do model.save_pretrained("...") or model.push_to_hub("..."))
  • track download numbers for models pushed to the hub, similar to models in the Transformers library
  • push an automated model card.

It leverages the PyTorchModelHubMixin class which allows to inherits these methods.

Usage is as follows:

from core.foundation_stereo import FoundationStereo

model = FoundationStereo(...)

# save it locally like so
model.save_pretrained("my-foundation-stereo-model")

# or push a trained model to the hub
model.push_to_hub("nvidia/foundation-stereo")

# now anyone can use it like so
model = FoundationStereo.from_pretrained("nvidia/foundation-stereo")

@NielsRogge
Copy link
Author

Disclaimer: please don't merge this PR before trying to push and reload to and from the hub to ensure it works as expected :) thank you

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.

1 participant