Change weight loading method

app.py

@@ -562,9 +562,7 @@ if __name__ == '__main__':
 
     print("Initializing and loading Pi3 model...")
 
-    model = Pi3()
-    _URL = "https://huggingface.co/yyfz233/Pi3/resolve/main/model.safetensors"
-    model.load_state_dict(torch.hub.load_state_dict_from_url(_URL))
+    model = Pi3.from_pretrained("yyfz233/Pi3")
     # model.load_state_dict(torch.load('ckpts/pi3.pt', weights_only=False, map_location=device))
 
     model.eval()

example.py

@@ -41,10 +41,6 @@ if __name__ == '__main__':
         model.load_state_dict(weight)
     else:
         model = Pi3.from_pretrained("yyfz233/Pi3").to(device).eval()
-        # or
-        # model = Pi3().to(device).eval()
-        # _URL = "https://huggingface.co/yyfz233/Pi3/resolve/main/model.safetensors"
-        # model.load_state_dict(torch.hub.load_state_dict_from_url(_URL))
 
     # 2. Prepare input data
     # The load_images_as_tensor function will print the loading path