Support Cosmos tokenizer TensorRT inference
State: closed | Pull request: yes
Labels: name: common

name: Run CICD

name: skip-docs

# What does this PR do ?

This PR adds TensorRT inference support for the Cosmos Tokenizers

**Collection**: common/video_tokenizers

# Changelog 
- It adds the tokenizer configs and modules from the[ Cosmos Tokenizer repo](https://github.com/NVIDIA/Cosmos-Tokenizer/tree/main/cosmos_tokenizer) which enables the models to be exported to TensorRT.
- It provides an example of how to write a wrapper class that enables the tokenizers to be exported to TensorRT

# Usage
The script `common/video_tokenizers/cosmos_trt_run.py` provides an example of how to export the CosmosTokenizer to TensorRT and run inference.

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


<!-- pylint-output-strict-mode-true -->

beep boop 🤖: 🚨 The following files must be fixed before merge!

---

Your code was analyzed with PyLint. The following annotations have been identified:

```
************* Module nemo.collections.common.video_tokenizers.cosmos_trt_run
nemo/collections/common/video_tokenizers/cosmos_trt_run.py:52:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/cosmos_trt_run.py:55:4: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/cosmos_trt_run.py:60:8: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.modules
nemo/collections/common/video_tokenizers/modules/__init__.py:35:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/__init__.py:39:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/__init__.py:43:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/__init__.py:48:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/__init__.py:53:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/__init__.py:58:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.common.video_tokenizers.modules.distributions
nemo/collections/common/video_tokenizers/modules/distributions.py:22:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/distributions.py:26:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/distributions.py:30:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/distributions.py:36:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/distributions.py:40:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.modules.layers2d
nemo/collections/common/video_tokenizers/modules/layers2d.py:40:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:45:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:50:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:55:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:61:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:85:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:101:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:111:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:141:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:211:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:238:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:307:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers2d.py:29:0: W0611: Unused numpy imported as np (unused-import)
************* Module nemo.collections.common.video_tokenizers.modules.layers3d
nemo/collections/common/video_tokenizers/modules/layers3d.py:50:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:95:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:100:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:105:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:118:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:130:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:138:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:174:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:197:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:233:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:257:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:281:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:296:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:350:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:365:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:375:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:406:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:416:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:450:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:526:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:532:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:572:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:650:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:657:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:688:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:807:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:834:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:946:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/layers3d.py:29:0: W0611: Unused numpy imported as np (unused-import)
************* Module nemo.collections.common.video_tokenizers.modules.patching
nemo/collections/common/video_tokenizers/modules/patching.py:65:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/patching.py:203:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.modules.quantizers
nemo/collections/common/video_tokenizers/modules/quantizers.py:48:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:63:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:264:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:278:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:288:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:351:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:432:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/quantizers.py:510:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.modules.utils
nemo/collections/common/video_tokenizers/modules/utils.py:25:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:30:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:34:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:39:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:43:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:47:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:51:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:55:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:59:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:63:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:67:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:78:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:87:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:91:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:98:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:102:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:112:0: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/modules/utils.py:116:0: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.networks
nemo/collections/common/video_tokenizers/networks/__init__.py:30:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/networks/__init__.py:37:0: C0115: Missing class docstring (missing-class-docstring)
************* Module nemo.collections.common.video_tokenizers.networks.continuous_image
nemo/collections/common/video_tokenizers/networks/continuous_image.py:29:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_image.py:49:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_image.py:60:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_image.py:70:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_image.py:73:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_image.py:78:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_image.py:83:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.networks.continuous_video
nemo/collections/common/video_tokenizers/networks/continuous_video.py:28:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_video.py:53:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_video.py:64:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_video.py:74:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_video.py:77:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_video.py:82:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/continuous_video.py:86:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.networks.discrete_image
nemo/collections/common/video_tokenizers/networks/discrete_image.py:30:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:60:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:64:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:75:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:86:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:89:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:94:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:98:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:103:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_image.py:21:0: W0611: Unused logger imported from loguru as logging (unused-import)
************* Module nemo.collections.common.video_tokenizers.networks.discrete_video
nemo/collections/common/video_tokenizers/networks/discrete_video.py:30:0: C0115: Missing class docstring (missing-class-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:61:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:65:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:76:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:87:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:90:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:95:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:99:4: C0116: Missing function or method docstring (missing-function-docstring)
nemo/collections/common/video_tokenizers/networks/discrete_video.py:104:4: C0116: Missing function or method docstring (missing-function-docstring)
************* Module nemo.collections.common.video_tokenizers.utils
nemo/collections/common/video_tokenizers/utils.py:325:0: C0116: Missing function or method docstring (missing-function-docstring)

-----------------------------------
Your code has been rated at 9.26/10
```

---

Thank you for improving NeMo's documentation!

[🤖]: Hi @meatybobby 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/12320944935) for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g