fix model loading bug

app.py

@@ -49,17 +49,23 @@ 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"
         #model = safely_load_state_dict(metauas_model, ckt_path)
-        model = MetaUAS.load_from_checkpoint("csgaobb/MetaUAS/MetaUAS-512.ckpt")
+        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)
-        model = MetaUAS.load_from_checkpoint("csgaobb/MetaUAS/MetaUAS-256.ckpt")
+        model = MetaUAS.from_pretrained("csgaobb/MetaUAS-256")
         img_size = 256
 
     model.to(device)