Update app.py

app.py

@@ -26,6 +26,8 @@ try_cuda = not args.disable_cuda
 torch.inference_mode()
 torch.no_grad()
 
+device = get_device(try_cuda=try_cuda)
+
 # Load segmentation models
 def load_segmentation_models(model_name: str = 'facebook/detr-resnet-50-panoptic'):
     feature_extractor = DetrFeatureExtractor.from_pretrained(model_name)
@@ -67,7 +69,6 @@ def clean_mask(mask, max_kernel: int = 23, min_kernel: int = 5):
 feature_extractor, segmentation_model, segmentation_cfg = load_segmentation_models()
 pipe = load_diffusion_pipeline()
 
-device = get_device(try_cuda=try_cuda)
 segmentation_model = segmentation_model.to(device)
 pipe = pipe.to(device)
 if args.attention_slicing: