update cpu config

inference.py

@@ -1,5 +1,6 @@
 from utils import *
 from torch import cuda
+import ipdb
 
 # ---  Launch Model ---
 device = 'cuda' if cuda.is_available() else 'cpu'

medomni/models/blip2.py

@@ -32,10 +32,10 @@ class Blip2Base(BaseModel):
         tokenizer.add_special_tokens({"bos_token": "[DEC]"})
         return tokenizer
 
-    def maybe_autocast(self, dtype=torch.float16):
+    def maybe_autocast(self, device, dtype=torch.float16):
         # if on cpu, don't use autocast
         # if on gpu, use autocast with dtype if provided, otherwise use torch.float16
-        enable_autocast = self.device != torch.device("cpu")
+        enable_autocast = device != torch.device("cpu")
 
         if enable_autocast:
             return torch.cuda.amp.autocast(dtype=dtype)

medomni/models/medomni.py

@@ -285,8 +285,8 @@ class MedOmni(Blip2Base, PyTorchModelHubMixin):
         if self.low_resource:
             self.vit_to_cpu()
             image = image.to("cpu")
-
-        with self.maybe_autocast():
+        ipdb.set_trace()
+        with self.maybe_autocast(device):
             if 'ct' in modals:
                 image_embeds_list = self.visual_encoder_3d(image, encoder_only=True)
                 image_embeds_list = [_.to(device) for _ in image_embeds_list]
@@ -414,7 +414,7 @@ class MedOmni(Blip2Base, PyTorchModelHubMixin):
         to_regress_embeds = self.embed_tokens(to_regress_tokens.input_ids)
         inputs_embeds = torch.cat([bos_embeds, img_embeds, to_regress_embeds], dim=1)
         attention_mask = torch.cat([atts_bos, atts_img, to_regress_tokens.attention_mask], dim=1)
-        with self.maybe_autocast():
+        with self.maybe_autocast(image.device):
             outputs = self.llama_model(
                 inputs_embeds=inputs_embeds,
                 attention_mask=attention_mask,
@@ -425,7 +425,7 @@ class MedOmni(Blip2Base, PyTorchModelHubMixin):
         loss = outputs.loss
 
         if 'detection' in samples['task_type']:
-            with self.maybe_autocast():
+            with self.maybe_autocast(image.device):
                 hidden_states = outputs.hidden_states[-1]
                 token_mask = targets == self.det_token_idx
                 target_states = hidden_states[token_mask]
@@ -437,7 +437,7 @@ class MedOmni(Blip2Base, PyTorchModelHubMixin):
                 loss += det_loss * 1e2
 
         if 'keypoint' in samples['task_type']:
-            with self.maybe_autocast():
+            with self.maybe_autocast(image.device):
                 hidden_states = outputs.hidden_states[-1]
                 token_mask = targets == self.point_token_idx_2d
                 target_states = hidden_states[token_mask]
@@ -451,7 +451,7 @@ class MedOmni(Blip2Base, PyTorchModelHubMixin):
         if 'segmentation' in samples['task_type']:
             if 'ct' in samples['modal']:
                 masks = samples['answer_img']
-                with self.maybe_autocast():
+                with self.maybe_autocast(image.device):
                     img_embeds_list = self.visual_encoder_3d(image, encoder_only=True)
                     img_embeds_list = [_.to(targets.device) for _ in img_embeds_list]
                     hidden_states = outputs.hidden_states[-1]
@@ -466,7 +466,7 @@ class MedOmni(Blip2Base, PyTorchModelHubMixin):
                     loss += self.seg_loss(seg_preds, masks.float()) # +
             else:  
                 masks = samples['answer_img']
-                with self.maybe_autocast():
+                with self.maybe_autocast(image.device):
                     feats = self.model_seg_2d.encoder(image[:,0])
                     last_feats = feats[-1]
                     hidden_states = outputs.hidden_states[-1]