Add Sound Encoder to Cosmos3#13911
Conversation
Signed-off-by: Maciej Bala <[email protected]>
| def _disable_encoder(self): | ||
| self.encoder = None | ||
| self._encoder_available = False | ||
| self.register_to_config(encoder_enabled=False) | ||
|
|
||
| def _fix_state_dict_keys_on_load(self, state_dict: OrderedDict) -> None: | ||
| super()._fix_state_dict_keys_on_load(state_dict) | ||
| if self.encoder is not None and not any(key.startswith("encoder.") for key in state_dict): | ||
| self._disable_encoder() | ||
|
|
There was a problem hiding this comment.
why do we need these two methods?
There was a problem hiding this comment.
It's an extra safety net for checkpoints that do not have the encoder weights. We will update the main checkpoint to have encoder weights, but I think it's still fine to keep this method in case of e.g. cached local checkpoints. We don't want them to break if people don't need the encoder weights.
dg845
left a comment
There was a problem hiding this comment.
Thanks for the PR! Left an initial design review :).
Signed-off-by: Maciej Bala <[email protected]>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
dg845
left a comment
There was a problem hiding this comment.
Thanks for the changes! Left some follow up comments.
Signed-off-by: Maciej Bala <[email protected]>
Signed-off-by: Maciej Bala <[email protected]>
Signed-off-by: Maciej Bala <[email protected]>
|
@bot /style |
|
Style bot fixed some files and pushed the changes. |
| ) -> Cosmos3AudioDecoderOutput | tuple[torch.Tensor]: | ||
| """Encode then decode a waveform; ``sample_posterior=False`` (default) decodes the distribution mode (mean), | ||
| whereas the upstream Cosmos3 AVAE always samples — pass ``sample_posterior=True`` for reference-equivalent | ||
| behavior.""" |
There was a problem hiding this comment.
Can you update the Cosmos3AVAEAudioTokenizer.forward docstring to document the arguments and return values? This should fix the failure in this CI run.
dg845
left a comment
There was a problem hiding this comment.
Thanks for your work on this PR! Left one comment about fixing a CI failure.
Signed-off-by: Maciej Bala <[email protected]>
|
Fixed CI. I think the PR is ready to be merged now. The checkpoints on HF should be updated today or tomorrow with encoder weights, but the current encoder-less checkpoints also work, as long as we don't call the |
* Initial version of sound encoder Signed-off-by: Maciej Bala <[email protected]> * Answered review comments Signed-off-by: Maciej Bala <[email protected]> * Reshaped encoder Snake1d Signed-off-by: Maciej Bala <[email protected]> * Updated test Signed-off-by: Maciej Bala <[email protected]> * Fixed test Signed-off-by: Maciej Bala <[email protected]> * Apply style fixes * Added docstring Signed-off-by: Maciej Bala <[email protected]> --------- Signed-off-by: Maciej Bala <[email protected]> Co-authored-by: dg845 <[email protected]> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
What does this PR do?
Fixes # (issue)
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.