Add very low MS-ILLM rates (#241)

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

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 (

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"