Integrate More Models

[Yiozolm]

Feature

Code Link: https://github.com/Yiozolm/CompressAI

Add support for recent learned image compression models beyond the current CompressAI model zoo, preferably through staged PRs rather than one large
merge.

The local integration currently includes candidate model entries for:

• GLIC, CMIC
• MLIC++, WACNN/STF, LIC-TCM/TCM, DCAE, SAAF, WeConvene
• TIC, TinyLIC, ShiftLIC small/middle/large
• Informer, FTIC, InvCompress, CCA
• HPCM base/large/phi
• MambaIC, MambaVC

The implementation also extracts shared LIC components instead of duplicating model-specific code, including reusable attention layers, wavelet wrappers, graph/context blocks, SSM/selective-scan wrappers, invertible helpers, and latent codecs such as ChannelSliceLatentCodec and MultistageCheckerboardLatentCodec.

Motivation

CompressAI is widely used as infrastructure for learned image compression research, but the official model coverage mostly stops around ELIC-era architectures. Many newer LIC papers rely on similar building blocks but are currently spread across independent forks with inconsistent APIs, checkpoint layouts, entropy-model implementations, and dependency handling.

Adding these models, or a selected subset of them, would make it easier to:

• compare recent LIC architectures under a common CompressAI API;
• reuse shared layers and latent-codec implementations across models;
• load and convert upstream checkpoints more consistently;
• provide tested reference implementations for post-ELIC learned image compression research.

Additional context

All source code is taken from the implementation on GitHub. The intergration was developed in collaboration with Opus 4.7.

This is not intended as a single large PR. I would like to first check whether maintainers are open to this direction and how the work should be split.

A possible PR sequence would be:

1. Shared layers, entropy-model utilities, and latent-codec infrastructure.
2. Lower-risk licensed models with tests and checkpoint conversion scripts.
3. Optional-dependency models guarded by lazy imports and clear error messages.
4. Remaining models after agreeing on license and pretrained-weight policy.

The local branch currently includes smoke tests in tests/test_models.py and tests/test_zoo.py, plus examples/convert_*_checkpoint.py scripts for several upstream checkpoint formats. Optional dependencies such as pytorch_wavelets, FrEIA, mamba_ssm, triton, and natten are gated so that importing compressai should still work without them.

Pretrained weights are not mirrored to CompressAI S3 yet. For candidate entries, pretrained=True can remain disabled until the project decides how to handle official weight hosting. Some upstream candidates also have incomplete license metadata, so those models can be excluded or deferred if that is required for mainline inclusion.

[Yiozolm]

I am about to complete my master's thesis, and I will not be working on image compression during my doctoral studies.

I am very grateful for the convenienve the compressai library has provided me over the past two years.

So I’d also like to do my part in the end.

[fracape]

Hi Yiozolm,

Thank you for offering to add support for these models. We’re happy to include this work in the library.

As you mentioned, we need to carefully review the license of the original source code you used to implement them. I see that you’ve reported (at least some of) them in your code. Regarding pretrained models, are they all the output of your training? The weights in the existing S3 bucket are outdated, but they all correspond to our training using similar parameters and datasets for the most accurate comparisons. It would be safer to host them at the same location once we’ve cleared the license and training configurations, rather than relying on external google drive links. We might want to move both the existing and new models to a new location.

Your PR sequence sounds good. PR #1 and #2 can be quickly agreed upon.

Thanks,

[Yiozolm]

Thank you for the positive response. That sounds good to me.

1. Regarding pretrained models: no, they are not all the output of my own training. For several models, I used publicly released upstream checkpoints only to validate checkpoint conversion and numerical compatibility with the original implementation. I did not intend to rely on external Google Drive links for mainline pretrained support.

   My suggestion is to first integrate the model definitions, reusable layers, latent codecs, zoo factory entries, tests, and checkpoint conversion utilities, while leaving pretrained=True disabled for these new models. This would already be useful for users who want to study, train, or reproduce these architectures within CompressAI, without making any claim about official pretrained weights or benchmark comparability.

2. PR sequence: since the full set of models is quite large, I can submit them in smaller batches to reduce review burden and avoid introducing too much instability at once. Since I’ve been pretty busy lately, this might take a few weeks to finish.

[fracape]

Great, thanks. Regarding the first point, yes. At the beginning of the repo, the objective was to verify whether the performance metrics reported by authors could be achieved. The lib’s popularity was actually attributed to its implementation in pytorch of modules like the entropy layers. For now, pretrained models can be skipped.