Auto deploy

stCompressService.py

@@ -1,14 +1,13 @@
 import pathlib
 import torch
 import torch.hub
-from torchvision.transforms.functional import convert_image_dtype, to_tensor
+from torchvision.transforms.functional import convert_image_dtype
 from torchvision.io.image import ImageReadMode, encode_png, decode_image
 
 from mcquic import Config
 from mcquic.modules.compressor import BaseCompressor, Compressor
 from mcquic.datasets.transforms import AlignedCrop
-import mcquic
-from mcquic.utils.specification import File, FileHeader
+from mcquic.utils.specification import File
 from mcquic.utils.vision import DeTransform
 
 try:
@@ -17,15 +16,12 @@ except:
     raise ImportError("To run `mcquic service`, please install Streamlit by `pip install streamlit` firstly.")
 
 
-MODELS_URL = "https://github.com/xiaosu-zhu/McQuic/releases/download/generic/"
+MODELS_URL = "https://github.com/xiaosu-zhu/McQuic/releases/download/generic/qp_3_msssim_fcc58b73.mcquic"
 
 
 @st.experimental_singleton
 def loadModel(qp: int, local: pathlib.Path, device, mse: bool):
-    suffix = "mse" if mse else "msssim"
-    # ckpt = torch.hub.load_state_dict_from_url(MODELS_URL + f"qp_{qp}_{suffix}.mcquic", map_location=device)
-
-    ckpt = torch.load("./qp_3_msssim_fcc58b73.mcquic", map_location=device)
+    ckpt = torch.hub.load_state_dict_from_url(MODELS_URL, map_location=device)
 
     config = Config.deserialize(ckpt["config"])
     model = Compressor(**config.Model.Params).to(device)
@@ -93,6 +89,12 @@ def main(debug: bool, quiet: bool, disable_gpu: bool):
 </a>
 
 > Due to resources limitation, I only provide compression service with model `qp = 3`.
+
+
+<a href="#">
+  <image src="https://img.shields.io/github/stars/xiaosu-zhu/McQuic?style=social" alt="Github"/>
+</a>
+
 """, unsafe_allow_html=True)
 
 
@@ -124,4 +126,4 @@ def main(debug: bool, quiet: bool, disable_gpu: bool):
 
 if __name__ == "__main__":
     with torch.inference_mode():
-        main(False, False, True)
+        main(False, False, False)