load metauas models from huggingface model repositories

app.py

@@ -52,12 +52,14 @@ metauas_model = MetaUAS(encoder_name,
 def process_image(prompt_img, query_img, options):
     # Load the model based on selected options
     if 'model-512' in options:
-        ckt_path = "weights/metauas-512.ckpt"
-        model = safely_load_state_dict(metauas_model, ckt_path)
+        #ckt_path = "weights/metauas-512.ckpt"
+        #model = safely_load_state_dict(metauas_model, ckt_path)
+        model = MetaUAS.from_pretrained("csgaobb/MetaUAS-512")
         img_size = 512
     else:
-        ckt_path = 'weights/metauas-256.ckpt'
-        model = safely_load_state_dict(metauas_model, ckt_path)
+        #ckt_path = 'weights/metauas-256.ckpt'
+        #model = safely_load_state_dict(metauas_model, ckt_path)
+        model = MetaUAS.from_pretrained("csgaobb/MetaUAS-256")
         img_size = 256
 
     model.to(device)

metauas.py

@@ -31,6 +31,7 @@ from torchvision.transforms.functional import pil_to_tensor
 from segmentation_models_pytorch.unet.model import UnetDecoder
 from segmentation_models_pytorch.fpn.decoder import FPNDecoder
 from segmentation_models_pytorch.encoders import get_encoder, get_preprocessing_params
+from huggingface_hub import PyTorchModelHubMixin
 
 def set_random_seed(seed=233, reproduce=False):
     np.random.seed(seed)
@@ -131,7 +132,7 @@ class AlignmentLayer(nn.Module):
             return aligned_features
 
 
-class MetaUAS(pl.LightningModule):
+class MetaUAS(pl.LightningModule, PyTorchModelHubMixin):
     def __init__(self, encoder_name, decoder_name, encoder_depth, decoder_depth, num_alignment_layers, alignment_type, fusion_policy):
         super().__init__()