model.py

from segment_anything import sam_model_registry

def create_sam_model(model_type, checkpoint, device: str = "cpu"):
    medsam_model = sam_model_registry[model_type](checkpoint=checkpoint)
    medsam_model = medsam_model.to(device)
    return medsam_model