support cpu

README.md

@@ -22,12 +22,12 @@ conda activate medversa
 ## Inference
 ``` python
 from utils import *
+from torch import cuda
 
 # ---  Launch Model ---
-device = 'cuda:0'
+device = 'cuda' if cuda.is_available() else 'cpu'
 model_cls = registry.get_model_class('medomni') # medomni is the architecture name :)
-model = model_cls.from_pretrained('hyzhou/MedVersa').to(device)
-model.eval()
+model = model_cls.from_pretrained('hyzhou/MedVersa').to(device).eval()
 
 # --- Define examples ---
 examples = [

demo.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import argparse
 import torch
+from torch import cuda
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 from torchvision import transforms
@@ -32,15 +33,14 @@ def parse_args():
     args = parser.parse_args()
     return args
 
-device = 'cuda:0'
+device = 'cuda' if cuda.is_available() else 'cpu'
 # Launch model
 args = parse_args()
 cfg = Config(args)
 
 model_config = cfg.model_cfg
 model_cls = registry.get_model_class(model_config.arch)
-model = model_cls.from_pretrained('hyzhou/MedVersa').to(device)
-model.eval()
+model = model_cls.from_pretrained('hyzhou/MedVersa').to(device).eval()
 global global_images
 global_images = None
 
@@ -146,7 +146,7 @@ def task_seg_2d(model, preds, hidden_states, image):
         seg_feats = model.model_seg_2d.decoder(*feats)
         seg_preds = model.model_seg_2d.segmentation_head(seg_feats)
         seg_probs = F.sigmoid(seg_preds)
-        seg_mask = seg_probs.cpu().squeeze().numpy() >= 0.5
+        seg_mask = seg_probs.to(torch.float32).cpu().squeeze().numpy() >= 0.5
         return seg_mask
     else:
         return None
@@ -165,7 +165,7 @@ def task_seg_3d(model, preds, hidden_states, img_embeds_list):
         new_img_embeds_list[-1] = last_feats
         seg_preds = model.visual_encoder_3d(encoder_only=False, x_=new_img_embeds_list)
         seg_probs = F.sigmoid(seg_preds)
-        seg_mask = seg_probs.cpu().squeeze().numpy() >= 0.5
+        seg_mask = seg_probs.to(torch.float32).cpu().squeeze().numpy() >= 0.5
         return seg_mask
 
 def task_det_2d(model, preds, hidden_states):
@@ -175,7 +175,7 @@ def task_det_2d(model, preds, hidden_states):
     if target_states:
         target_states = torch.cat(target_states).squeeze()
         det_states = model.text_det(target_states).detach().cpu()
-        return det_states.numpy()
+        return det_states.to(torch.float32).numpy()
     return torch.zeros_like(indices)
 
 class StoppingCriteriaSub(StoppingCriteria):
@@ -240,7 +240,7 @@ def load_and_preprocess_image(image):
         transforms.ToTensor(),
         transforms.Normalize(mean, std)
     ])
-    image = transform(image).type(torch.bfloat16).cuda().unsqueeze(0)
+    image = transform(image).type(torch.bfloat16).unsqueeze(0)
     return image
 
 def load_and_preprocess_volume(image):
@@ -249,7 +249,7 @@ def load_and_preprocess_volume(image):
     transform = tio.Compose([
         tio.ZNormalization(masking_method=tio.ZNormalization.mean),
     ])
-    image = transform(image.unsqueeze(0)).type(torch.bfloat16).cuda()
+    image = transform(image.unsqueeze(0)).type(torch.bfloat16)
     return image
 
 def read_image(image_path):
@@ -328,14 +328,14 @@ def generate(image_path, image, context, modal, num_imgs, prompt, num_beams, do_
 def generate_predictions(images, context, prompt, modality, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature):
     num_imgs = len(images)
     modal = modality.lower()
-    image_tensors = [read_image(img) for img in images]
+    image_tensors = [read_image(img).to(device) for img in images]
     if modality == 'ct':
         time.sleep(2)
     else:
         time.sleep(1)
     image_tensor = torch.cat(image_tensors)
     
-    with torch.autocast("cuda"):
+    with torch.autocast(device):
         with torch.no_grad():
             generated_image, seg_mask_2d, seg_mask_3d, output_text = generate(images, image_tensor, context, modal, num_imgs, prompt, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature)
     
@@ -388,38 +388,6 @@ def gradio_interface(chatbot, images, context, prompt, modality, num_beams, do_s
 
     return chatbot, snapshot, gr.update(maximum=len(output_images)-1)
 
-# my_dict = {}
-# def gradio_interface(images, task, context, prompt, modality, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature):
-#     if not images:
-#         return None, "Error: At least one image is required to proceed."
-#     if not prompt or not task or not modality:
-#         return None, "Error: Please provide prompt, select task and modality to proceed."
-    
-#     generated_images, seg_mask_2d, seg_mask_3d, output_text = generate_predictions(images, task, context, prompt, modality, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature)
-#     output_images = []
-
-#     input_images = [np.asarray(Image.open(img.name).convert('RGB')).astype(np.uint8) if img.name.endswith(('.jpg', '.jpeg', '.png')) else f"{img.name} (3D Volume)" for img in images]
-#     if generated_images is not None:
-#         for generated_image in generated_images:
-#             output_images.append(np.asarray(generated_image).astype(np.uint8)) 
-#         snapshot = (output_images[0], [])
-#         if seg_mask_2d is not None:
-#             snapshot = (output_images[0], [(seg_mask_2d[0], "Mask")])
-#         if seg_mask_3d is not None:
-#             snapshot = (output_images[0], [(seg_mask_3d[0], "Mask")])
-#     else:
-#         output_images = input_images.copy()
-#         snapshot = (output_images[0], [])
-    
-#     my_dict['image'] = output_images
-#     my_dict['mask'] = None
-#     if seg_mask_2d is not None:
-#         my_dict['mask'] = seg_mask_2d
-#     if seg_mask_3d is not None:
-#         my_dict['mask'] = seg_mask_3d
-    
-#     return output_text, snapshot, gr.update(maximum=len(output_images)-1)
-
 def render(x):
     if x > len(my_dict['image'])-1:
         x = len(my_dict['image'])-1

inference.py

@@ -1,10 +1,10 @@
 from utils import *
+from torch import cuda
 
 # ---  Launch Model ---
-device = 'cuda:0'
+device = 'cuda' if cuda.is_available() else 'cpu'
 model_cls = registry.get_model_class('medomni') # medomni is the architecture name :)
-model = model_cls.from_pretrained('hyzhou/MedVersa').to(device)
-model.eval()
+model = model_cls.from_pretrained('hyzhou/MedVersa').to(device).eval()
 
 # --- Define examples ---
 examples = [
@@ -85,14 +85,14 @@ temperature = 0.1
 index = 0
 demo_ex = examples[index]
 images, context, prompt, modality, task = demo_ex[0], demo_ex[1], demo_ex[2], demo_ex[3], demo_ex[4]
-seg_mask_2d, seg_mask_3d, output_text = generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature)
+seg_mask_2d, seg_mask_3d, output_text = generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature, device)
 print(output_text)
 
 # --- Segment the lesion in the dermatology image ---
 index = 6
 demo_ex = examples[index]
 images, context, prompt, modality, task = demo_ex[0], demo_ex[1], demo_ex[2], demo_ex[3], demo_ex[4]
-seg_mask_2d, seg_mask_3d, output_text = generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature)
+seg_mask_2d, seg_mask_3d, output_text = generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature, device)
 print(output_text)
 print(seg_mask_2d[0].shape) # H, W
 
@@ -100,7 +100,7 @@ print(seg_mask_2d[0].shape) # H, W
 index = -2
 demo_ex = examples[index]
 images, context, prompt, modality, task = demo_ex[0], demo_ex[1], demo_ex[2], demo_ex[3], demo_ex[4]
-seg_mask_2d, seg_mask_3d, output_text = generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature)
+seg_mask_2d, seg_mask_3d, output_text = generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature, device)
 print(output_text)
 print(len(seg_mask_3d)) # Number of slices
 print(seg_mask_3d[0].shape) # H, W

utils.py

@@ -133,7 +133,7 @@ def task_seg_2d(model, preds, hidden_states, image):
         seg_feats = model.model_seg_2d.decoder(*feats)
         seg_preds = model.model_seg_2d.segmentation_head(seg_feats)
         seg_probs = F.sigmoid(seg_preds)
-        seg_mask = seg_probs.cpu().squeeze().numpy() >= 0.5
+        seg_mask = seg_probs.to(dtype=torch.float32).cpu().squeeze().numpy() >= 0.5
         return seg_mask
     else:
         return None
@@ -152,7 +152,7 @@ def task_seg_3d(model, preds, hidden_states, img_embeds_list):
         new_img_embeds_list[-1] = last_feats
         seg_preds = model.visual_encoder_3d(encoder_only=False, x_=new_img_embeds_list)
         seg_probs = F.sigmoid(seg_preds)
-        seg_mask = seg_probs.cpu().squeeze().numpy() >= 0.5
+        seg_mask = seg_probs.to(dtype=torch.float32).cpu().squeeze().numpy() >= 0.5
         return seg_mask
 
 def task_det_2d(model, preds, hidden_states):
@@ -227,7 +227,7 @@ def load_and_preprocess_image(image):
         transforms.ToTensor(),
         transforms.Normalize(mean, std)
     ])
-    image = transform(image).type(torch.bfloat16).cuda().unsqueeze(0)
+    image = transform(image).type(torch.bfloat16).unsqueeze(0)
     return image
 
 def load_and_preprocess_volume(image):
@@ -236,7 +236,7 @@ def load_and_preprocess_volume(image):
     transform = tio.Compose([
         tio.ZNormalization(masking_method=tio.ZNormalization.mean),
     ])
-    image = transform(image.unsqueeze(0)).type(torch.bfloat16).cuda()
+    image = transform(image.unsqueeze(0)).type(torch.bfloat16)
     return image
 
 def read_image(image_path):
@@ -285,7 +285,6 @@ def generate(model, image_path, image, context, modal, task, num_imgs, prompt, n
         seg_mask = task_seg_2d(model, preds, hidden_states, image)
         output_image, seg_mask_2d = seg_2d_process(image_path, seg_mask)
     if sum(preds == model.seg_token_idx_3d):
-        ipdb.set_trace()
         seg_mask = task_seg_3d(model, preds, hidden_states, img_embeds_list)
         output_image, seg_mask_3d = seg_3d_process(image_path, seg_mask)
     if sum(preds == model.det_token_idx):
@@ -304,17 +303,17 @@ def generate(model, image_path, image, context, modal, task, num_imgs, prompt, n
         output_text = 'The main diagnosis is melanoma.'
     return output_image, seg_mask_2d, seg_mask_3d, output_text
 
-def generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature):
+def generate_predictions(model, images, context, prompt, modality, task, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature, device):
     num_imgs = len(images)
     modal = modality.lower()
-    image_tensors = [read_image(img) for img in images]
+    image_tensors = [read_image(img).to(device) for img in images]
     if modality == 'ct':
         time.sleep(2)
     else:
         time.sleep(1)
     image_tensor = torch.cat(image_tensors)
     
-    with torch.autocast("cuda"):
+    with torch.autocast(device):
         with torch.no_grad():
             generated_image, seg_mask_2d, seg_mask_3d, output_text = generate(model, images, image_tensor, context, modal, task, num_imgs, prompt, num_beams, do_sample, min_length, top_p, repetition_penalty, length_penalty, temperature)