check model is on device

app.py

@@ -10,6 +10,14 @@ import time
 feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
 model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+print(f"Is CUDA available: {torch.cuda.is_available()} --> {device=}")
+if (torch.cuda.is_available()):
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+    
+model.to(device)
+
 # https://github.com/NielsRogge/Transformers-Tutorials/blob/master/SegFormer/Segformer_inference_notebook.ipynb
 
 def cityscapes_palette():
@@ -53,23 +61,23 @@ def annotation(image:ImageDraw, color_seg:np.array):
 def call(image): #nparray
     start = time.time()
 
-    print(f"Is CUDA available: {torch.cuda.is_available()}")
-    if (torch.cuda.is_available()):
-        print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-
     resized = Image.fromarray(image).resize((1024,1024))
     resized_image = np.array(resized)
     print(f"{np.array(resized_image).shape=}") # 1024, 1024, 3
 
+    print(f"*processing time: {(time.time() - start):.2f} s")
     # resized_image = Image.fromarray(resized_image_np)
     # print(f"{resized_image=}")
 
     inputs = feature_extractor(images=resized_image, return_tensors="pt")
 
+    print(f"**processing time: {(time.time() - start):.2f} s")
+
     outputs  = model(**inputs)
     print(f"{outputs.logits.shape=}") # shape (batch_size, num_labels, height/4, width/4) -> 3, 19, 256 ,256
     # print(f"{logits}")
 
+    print(f"***processing time: {(time.time() - start):.2f} s")
     # First, rescale logits to original image size
     interpolated_logits =  nn.functional.interpolate(
         outputs.logits,
@@ -91,6 +99,8 @@ def call(image): #nparray
     color_seg = color_seg[..., ::-1]
     print(f"{color_seg.shape=}")
 
+    print(f"****processing time: {(time.time() - start):.2f} s")
+
     # Show image + mask
     img = np.array(resized_image) * 0.5 + color_seg * 0.5
     img = img.astype(np.uint8)
@@ -98,8 +108,7 @@ def call(image): #nparray
     out_im_file = Image.fromarray(img)
     annotation(out_im_file, color_seg)
 
-    end = time.time()
-    print(f"processing time: {(end - start):.2f} s")
+    print(f"--> processing time: {(time.time() - start):.2f} s")
 
     return out_im_file