bugfix: model path

app.py

@@ -41,17 +41,20 @@ def main():
     """
     Main function to set up and run the Gradio demo.
     """
+    
     args = parse_arguments()
     
     args.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     # Set up model
     die_token = os.getenv("DIE_TOKEN")
-    model_path = hf_hub_download(
+    
+    args.die_model_path = hf_hub_download(
         repo_id="gabar92/die",
         filename=args.die_model_path,
         use_auth_token=die_token
     )
+    
     die_model = UNetDIEModel(args=args)
 
     # Prepare example images
@@ -105,9 +108,10 @@ def parse_arguments():
     :return: argument namespace
     """
     parser = argparse.ArgumentParser()
-    parser.add_argument("--die_model_path", default="2024_08_09_model_epoch_89.pt", help="Path to the DIE model checkpoint")
-    parser.add_argument("--example_image_path", default="example_images", help="Path to directory with example images")
-    return parser.parse_args()
+
+    parser.add_argument("--die_model_path", default="2024_08_09_model_epoch_89.pt")
+
+    parser.add_argument("--example_image_path", default="example_images")
     
 
 if __name__ == "__main__":