From 8e9c2ed56b3329b37e9b891f4ee2ad163f65c422 Mon Sep 17 00:00:00 2001
From: Matthew Muckley <matt.muckley@gmail.com>
Date: Wed, 18 Sep 2024 11:28:38 -0400
Subject: [PATCH] Add very low MS-ILLM rates (#241)

---
 hubconf.py                        | 46 +++++++++++++++++++++++++++++++
 neuralcompression/zoo/__init__.py |  2 ++
 neuralcompression/zoo/_msillm.py  | 20 ++++++++++++++
 3 files changed, 68 insertions(+)

diff --git a/hubconf.py b/hubconf.py
index 04a1482..7e9fe81 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -10,6 +10,8 @@
 from neuralcompression.zoo import msillm_quality_4 as _msillm_quality_4
 from neuralcompression.zoo import msillm_quality_5 as _msillm_quality_5
 from neuralcompression.zoo import msillm_quality_6 as _msillm_quality_6
+from neuralcompression.zoo import msillm_quality_vlo1 as _msillm_quality_vlo1
+from neuralcompression.zoo import msillm_quality_vlo2 as _msillm_quality_vlo2
 from neuralcompression.zoo import noganms_quality_1 as _noganms_quality_1
 from neuralcompression.zoo import noganms_quality_2 as _noganms_quality_2
 from neuralcompression.zoo import noganms_quality_3 as _noganms_quality_3
@@ -23,6 +25,50 @@
 dependencies = ["torch"]
 
 
+def msillm_quality_vlo1(pretrained=True, **kwargs):
+    """
+    Pretrained MS-ILLM model
+
+    This model was trained for the very low rates in the paper:
+
+    M Careil, MJ Muckley, J Verbeek, S Lathuliere.
+    Towards image compression with perfect realism at ultra-low bitrates.
+    In *ICLR*, 2024.
+
+    The target bitrate is 0.00218 bits per pixel
+
+    The pretrained weights are released under the CC-BY-NC 4.0 license
+    available at
+    https://github.com/facebookresearch/NeuralCompression/blob/main/WEIGHTS_LICENSE
+
+    pretrained (bool): kwargs, load pretrained weights into the model
+    """
+
+    return _msillm_quality_vlo1(pretrained=pretrained, **kwargs)
+
+
+def msillm_quality_vlo2(pretrained=True, **kwargs):
+    """
+    Pretrained MS-ILLM model
+
+    This model was trained for the very low rates in the paper:
+
+    M Careil, MJ Muckley, J Verbeek, S Lathuliere.
+    Towards image compression with perfect realism at ultra-low bitrates.
+    In *ICLR*, 2024.
+
+    The target bitrate is 0.00438 bits per pixel
+
+    The pretrained weights are released under the CC-BY-NC 4.0 license
+    available at
+    https://github.com/facebookresearch/NeuralCompression/blob/main/WEIGHTS_LICENSE
+
+    pretrained (bool): kwargs, load pretrained weights into the model
+    """
+
+    return _msillm_quality_vlo2(pretrained=pretrained, **kwargs)
+
+
 def msillm_quality_1(pretrained=True, **kwargs):
     """
     Pretrained MS-ILLM model
diff --git a/neuralcompression/zoo/__init__.py b/neuralcompression/zoo/__init__.py
index fe3725e..7e7888c 100644
--- a/neuralcompression/zoo/__init__.py
+++ b/neuralcompression/zoo/__init__.py
@@ -10,6 +10,8 @@
     msillm_quality_4,
     msillm_quality_5,
     msillm_quality_6,
+    msillm_quality_vlo1,
+    msillm_quality_vlo2,
 )
 from ._msillm_vqvae import vqvae_xcit_p8_ch64_cb1024_h8
 from ._noganms import (
diff --git a/neuralcompression/zoo/_msillm.py b/neuralcompression/zoo/_msillm.py
index d88c9ab..f259b04 100644
--- a/neuralcompression/zoo/_msillm.py
+++ b/neuralcompression/zoo/_msillm.py
@@ -13,6 +13,8 @@
 LOGGER = logging.getLogger(__file__)
 
 VALID_WEIGHTS = [
+    "target_0.00218bpp",
+    "target_0.00438bpp",
     "target_0.035bpp",
     "target_0.07bpp",
     "target_0.14bpp",
@@ -48,6 +50,24 @@ def _build_msillm(weights: Optional[str] = None):
     return model
 
 
+def msillm_quality_vlo1(pretrained=False, **kwargs):
+    if pretrained is True:
+        weights = "target_0.00218bpp"
+    else:
+        weights = None
+
+    return _build_msillm(weights=weights)
+
+
+def msillm_quality_vlo2(pretrained=False, **kwargs):
+    if pretrained is True:
+        weights = "target_0.00438bpp"
+    else:
+        weights = None
+
+    return _build_msillm(weights=weights)
+
+
 def msillm_quality_1(pretrained=False, **kwargs):
     if pretrained is True:
         weights = "target_0.035bpp"