fix model loading bug

app.py

@@ -49,23 +49,24 @@ metauas_model = MetaUAS(encoder_name,
                 fusion_policy
                 )   
 
-
-model_256 = safely_load_state_dict(metauas_model, "weights/metauas-256.ckpt")
-model_512 = safely_load_state_dict(metauas_model, "weights/metauas-512.ckpt")
-model_256.push_to_hub("csgaobb/MetaUAS-256")
-model_512.push_to_hub("csgaobb/MetaUAS-512")
-
 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"
+        ckt_path = "weights/metauas-512.ckpt"
         #model = safely_load_state_dict(metauas_model, ckt_path)
-        model = MetaUAS.from_pretrained("csgaobb/MetaUAS-512")
+        
+        ckpt_path = hf_hub_download(repo_id="csgaobb/MetaUAS", filename=ckt_path)
+        metauas_model.load_state_dict(torch.load(ckpt_path))
+  
+        #model = MetaUAS.from_pretrained("csgaobb/MetaUAS-512")
         img_size = 512
     else:
-        #ckt_path = 'weights/metauas-256.ckpt'
+        ckt_path = 'weights/metauas-256.ckpt'
         #model = safely_load_state_dict(metauas_model, ckt_path)
-        model = MetaUAS.from_pretrained("csgaobb/MetaUAS-256")
+
+        ckpt_path = hf_hub_download(repo_id="csgaobb/MetaUAS", filename=ckt_path)
+        metauas_model.load_state_dict(torch.load(ckpt_path))
+        #model = MetaUAS.from_pretrained("csgaobb/MetaUAS-256")
         img_size = 256
 
     model.to(device)

metauas.py

@@ -132,7 +132,7 @@ class AlignmentLayer(nn.Module):
             return aligned_features
 
 
-class MetaUAS(pl.LightningModule, PyTorchModelHubMixin):
+class MetaUAS(pl.LightningModule):
     def __init__(self, encoder_name, decoder_name, encoder_depth, decoder_depth, num_alignment_layers, alignment_type, fusion_policy):
         super().__init__()
         
@@ -267,7 +267,7 @@ class MetaUAS(pl.LightningModule, PyTorchModelHubMixin):
                 stride=1,
                 padding=0,
             )
-           
+        
     def forward(self, batch):
         query_input = self.preprocess(batch["query_image"])
         prompt_input = self.preprocess(batch["prompt_image"])