diff --git a/wsinfer/cli/convert_csv_to_sbubmi.py b/wsinfer/cli/convert_csv_to_sbubmi.py
index 594822c..005fa9a 100644
--- a/wsinfer/cli/convert_csv_to_sbubmi.py
+++ b/wsinfer/cli/convert_csv_to_sbubmi.py
@@ -36,7 +36,7 @@
 import pandas as pd
 import tqdm
 
-from ..wsi import WSI
+from ..wsi import get_wsi_cls
 from ..wsi import CanReadRegion
 
 
@@ -353,7 +353,8 @@ def tosbu(
             click.secho(f"WSI file not found: {wsi_file}", bg="red")
             click.secho("Skipping...", bg="red")
             continue
-        slide = WSI(wsi_file)
+        wsi_reader = get_wsi_cls()
+        slide = wsi_reader(wsi_file)
 
         slide_width, slide_height = slide.level_dimensions[0]
 
diff --git a/wsinfer/modellib/data.py b/wsinfer/modellib/data.py
index b0d9ed7..be4e838 100644
--- a/wsinfer/modellib/data.py
+++ b/wsinfer/modellib/data.py
@@ -10,7 +10,7 @@
 import torch
 from PIL import Image
 
-from wsinfer.wsi import WSI
+from wsinfer.wsi import get_wsi_cls
 
 
 def _read_patch_coords(path: str | Path) -> npt.NDArray[np.int_]:
@@ -87,7 +87,8 @@ def __init__(
 
     def worker_init(self, worker_id: int | None = None) -> None:
         del worker_id
-        self.slide = WSI(self.wsi_path)
+        wsi_reader = get_wsi_cls()
+        self.slide = wsi_reader(self.wsi_path)
 
     def __len__(self) -> int:
         return self.patches.shape[0]
diff --git a/wsinfer/patchlib/__init__.py b/wsinfer/patchlib/__init__.py
index ee15598..b42b465 100644
--- a/wsinfer/patchlib/__init__.py
+++ b/wsinfer/patchlib/__init__.py
@@ -10,7 +10,7 @@
 import numpy.typing as npt
 from PIL import Image
 
-from ..wsi import WSI
+from ..wsi import get_wsi_cls
 from ..wsi import _validate_wsi_directory
 from ..wsi import get_avg_mpp
 from .patch import get_multipolygon_from_binary_arr
@@ -103,7 +103,7 @@ def segment_and_patch_one_slide(
         logger.info(f"mask_path={mask_path}")
         return None
 
-    slide = WSI(slide_path)
+    slide = get_wsi_cls()(slide_path)
     mpp = get_avg_mpp(slide_path)
     logger.info(f"Slide has WxH {slide.dimensions} and MPP={mpp}")
 
diff --git a/wsinfer/wsi.py b/wsinfer/wsi.py
index d283376..79b3666 100644
--- a/wsinfer/wsi.py
+++ b/wsinfer/wsi.py
@@ -17,6 +17,9 @@
 
 logger = logging.getLogger(__name__)
 
+_BACKEND: str = "tiffslide"
+
+_allowed_backends = {"openslide", "tiffslide"}
 
 try:
     import openslide
@@ -46,48 +49,46 @@
     )
 
 
-@overload
-def set_backend(name: Literal["openslide"]) -> type[openslide.OpenSlide]:
-    ...
-
-
-@overload
-def set_backend(name: Literal["tiffslide"]) -> type[tiffslide.TiffSlide]:
-    ...
-
-
-def set_backend(
-    name: Literal["openslide"] | Literal["tiffslide"],
-) -> type[tiffslide.TiffSlide] | type[openslide.OpenSlide]:
-    global WSI
-    if name not in ["openslide", "tiffslide"]:
-        raise ValueError(f"Unknown backend: {name}")
-    logger.info(f"Setting backend to {name}")
-    if name == "openslide":
-        if not HAS_OPENSLIDE:
-            raise BackendNotAvailable(
-                "OpenSlide is not available. Please install the OpenSlide compiled"
-                " library and the Python package 'openslide-python'."
-                " See https://openslide.org/ for more information."
-            )
-        WSI = openslide.OpenSlide
+def set_backend(name: str) -> None:
+    global _BACKEND
+    if name not in _allowed_backends:
+        raise ValueError(f"Unknown backend: '{name}'")
+    if name == "openslide" and not HAS_OPENSLIDE:
+        raise BackendNotAvailable(
+            "OpenSlide is not available. Please install the OpenSlide compiled"
+            " library and the Python package 'openslide-python'."
+            " See https://openslide.org/ for more information."
+        )
     elif name == "tiffslide":
         if not HAS_TIFFSLIDE:
             raise BackendNotAvailable(
                 "TiffSlide is not available. Please install 'tiffslide'."
             )
-        WSI = tiffslide.TiffSlide
+
+    logger.debug(f"Set backend to {name}")
+
+    _BACKEND = name
+
+
+def get_wsi_cls() -> type[openslide.OpenSlide] | type[tiffslide.TiffSlide]:
+    if _BACKEND not in _allowed_backends:
+        raise ValueError(
+            f"Unknown backend: '{_BACKEND}'. Please contact the developer!"
+        )
+    if _BACKEND == "openslide":
+        return openslide.OpenSlide  # type: ignore
+    elif _BACKEND == "tiffslide":
+        return tiffslide.TiffSlide
     else:
-        raise ValueError(f"Unknown backend: {name}")
-    return WSI
+        raise ValueError("Contact the developer, slide backend not known")
 
 
 # Set the slide backend based on the environment.
-WSI: type[openslide.OpenSlide] | type[tiffslide.TiffSlide]
-if HAS_OPENSLIDE:
-    WSI = set_backend("openslide")
-elif HAS_TIFFSLIDE:
-    WSI = set_backend("tiffslide")
+# Prioritize TiffSlide if the user has it installed.
+if HAS_TIFFSLIDE:
+    set_backend("tiffslide")
+elif HAS_OPENSLIDE:
+    set_backend("openslide")
 else:
     raise NoBackendException("No backend found! Please install openslide or tiffslide")