Added feature to switch between cpu and gpu based on availability

internvl_utils.py

@@ -82,7 +82,9 @@ def load_image(image, input_size=448, max_num=12):
     images = dynamic_preprocess(pil_image, image_size=input_size, use_thumbnail=True, max_num=max_num)
     pixel_values = [transform(image) for image in images]
     pixel_values = torch.stack(pixel_values)
-    pixel_values = pixel_values.to(torch.bfloat16).to("cuda")
+    pixel_values = pixel_values.to(torch.bfloat16)
+    if torch.cuda.is_available():
+        pixel_values = pixel_values.to("cuda")
     return pixel_values
 
 def split_model(model_name):

models/InternVL3/intervl3.py

@@ -25,7 +25,7 @@ class InternVL3(BaseModel):
             low_cpu_mem_usage=True,
             use_flash_attn=True,
             trust_remote_code=True,
-            device_map="auto",
+            device_map="cuda" if torch.cuda.is_available() else "cpu",
         ).eval()
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_name,