check for gpu if none exists

app.py

@@ -1,24 +1,27 @@
 import torch
 import gradio as gr
 
-model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3").cuda()  # or "resnet50"
+model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3")
+
+if torch.cuda.is_available():
+    model = model.cuda()
 
 convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")
 
 
 def inference(video):
     convert_video(
-        model,                           # The loaded model, can be on any device (cpu or cuda).
-        input_source=video,        # A video file or an image sequence directory.
-        downsample_ratio=0.25,           # [Optional] If None, make downsampled max size be 512px.
-        output_type='video',             # Choose "video" or "png_sequence"
-        output_composition='com.mp4',    # File path if video; directory path if png sequence.
-        output_alpha=None,          # [Optional] Output the raw alpha prediction.
-        output_foreground=None,     # [Optional] Output the raw foreground prediction.
-        output_video_mbps=4,             # Output video mbps. Not needed for png sequence.
-        seq_chunk=12,                    # Process n frames at once for better parallelism.
-        num_workers=1,                   # Only for image sequence input. Reader threads.
-        progress=True                    # Print conversion progress.
+        model,  # The loaded model, can be on any device (cpu or cuda).
+        input_source=video,  # A video file or an image sequence directory.
+        downsample_ratio=0.25,  # [Optional] If None, make downsampled max size be 512px.
+        output_type="video",  # Choose "video" or "png_sequence"
+        output_composition="com.mp4",  # File path if video; directory path if png sequence.
+        output_alpha=None,  # [Optional] Output the raw alpha prediction.
+        output_foreground=None,  # [Optional] Output the raw foreground prediction.
+        output_video_mbps=4,  # Output video mbps. Not needed for png sequence.
+        seq_chunk=12,  # Process n frames at once for better parallelism.
+        num_workers=1,  # Only for image sequence input. Reader threads.
+        progress=True,  # Print conversion progress.
     )
     return "com.mp4"