move to cuda is device is available

app.py

@@ -1,6 +1,7 @@
 import argparse
 import os
 from functools import partial
+import torch
 
 import gradio as gr
 from PIL import Image
@@ -41,6 +42,8 @@ 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")
@@ -49,8 +52,7 @@ def main():
         filename=args.die_model_path,
         use_auth_token=die_token
     )
-    die_model = UNetDIEModel(args=model_path)
-    device = args.device
+    die_model = UNetDIEModel(args=args)
 
     # Prepare example images
     example_image_list = [
@@ -68,7 +70,7 @@ def main():
     """
 
     # Partial function for inference with model and device arguments
-    partial_die_inference = partial(die_inference, die_model=die_model, device=device)
+    partial_die_inference = partial(die_inference, die_model=die_model, device=args.device)
 
     with gr.Blocks() as demo:
         with gr.Row():
@@ -104,7 +106,6 @@ def parse_arguments():
     """
     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("--device", default="cpu", choices=["cpu", "cuda"], help="Device to run the model on")
     parser.add_argument("--example_image_path", default="example_images", help="Path to directory with example images")
     return parser.parse_args()