Update app.py

app.py

@@ -26,6 +26,10 @@ try_cuda = not args.disable_cuda
 torch.inference_mode()
 torch.no_grad()
 
+# Device helper
+def get_device(try_cuda=True):
+    return torch.device('cuda' if try_cuda and torch.cuda.is_available() else 'cpu')
+    
 device = get_device(try_cuda=try_cuda)
 
 # Load segmentation models
@@ -45,10 +49,6 @@ def load_diffusion_pipeline(model_name: str = 'stabilityai/stable-diffusion-2-in
         use_auth_token=auth_token
     )
 
-# Device helper
-def get_device(try_cuda=True):
-    return torch.device('cuda' if try_cuda and torch.cuda.is_available() else 'cpu')
-
 def min_pool(x: torch.Tensor, kernel_size: int):
     pad_size = (kernel_size - 1) // 2
     return -torch.nn.functional.max_pool2d(-x, kernel_size, (1, 1), padding=pad_size)