add configurable device support

app.py

@@ -19,6 +19,7 @@ from super_gradients.training import models
 
 class Kosmos2:
     def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.colors = [
             (0, 255, 0),
             (0, 0, 255),
@@ -43,7 +44,7 @@ class Kosmos2:
         }
 
         self.ckpt = "ydshieh/kosmos-2-patch14-224"
-        self.model = AutoModelForVision2Seq.from_pretrained(self.ckpt, trust_remote_code=True).to("cuda")
+        self.model = AutoModelForVision2Seq.from_pretrained(self.ckpt, trust_remote_code=True).to(self.device)
         self.processor = AutoProcessor.from_pretrained(self.ckpt, trust_remote_code=True)
 
     def is_overlapping(self, rect1, rect2):
@@ -191,11 +192,11 @@ class Kosmos2:
         inputs = self.processor(text=text_input, images=image_input, return_tensors="pt")
 
         generated_ids = self.model.generate(
-            pixel_values=inputs["pixel_values"].to("cuda"),
-            input_ids=inputs["input_ids"][:, :-1].to("cuda"),
-            attention_mask=inputs["attention_mask"][:, :-1].to("cuda"),
+            pixel_values=inputs["pixel_values"].to(self.device),
+            input_ids=inputs["input_ids"][:, :-1].to(self.device),
+            attention_mask=inputs["attention_mask"][:, :-1].to(self.device),
             img_features=None,
-            img_attn_mask=inputs["img_attn_mask"][:, :-1].to("cuda"),
+            img_attn_mask=inputs["img_attn_mask"][:, :-1].to(self.device),
             use_cache=True,
             max_new_tokens=128,
         )