diff --git a/setup.py b/setup.py
index bdf196b..917a014 100644
--- a/setup.py
+++ b/setup.py
@@ -123,6 +123,7 @@ setup(
     extras_require={
         "tests": ["nose", "parameterized"],
         "docs": ["sphinx", "myst-parser", "sphinx-rtd-theme"],
+        "stubs": ["pybind11-stubgen"],
     },
     entry_points='''
         [console_scripts]
diff --git a/tools/stubgen.py b/tools/stubgen.py
index 32504ec..2519841 100644
--- a/tools/stubgen.py
+++ b/tools/stubgen.py
@@ -1,7 +1,7 @@
 import inspect
 import os.path
 import tempfile
-import mypy.stubgen as sg
+import pybind11_stubgen as sg
 
 package = 'quicktex'
 prefix = '_'
@@ -20,10 +20,15 @@ def find_submodules(pkg):
 
 if __name__ == "__main__":
     find_submodules(__import__(prefix + package))
+
+    pkgdir = os.path.abspath(os.curdir)
+
     with tempfile.TemporaryDirectory() as out:
 
         # generate stubs using mypy Stubgen
-        sg.generate_stubs(sg.parse_options(['-o', out, '-p', prefix + package]))
+        sg.main(['-o', out, '--root-module-suffix', "", prefix + package])
+
+        os.curdir = pkgdir
 
         # walk resulting stubs and move them to their new location
         for root, dirs, files in os.walk(out):
@@ -31,6 +36,9 @@ if __name__ == "__main__":
                 # location of the extension module's stub file
                 ext_module = os.path.relpath(root, out)
 
+                if stub_name.split('.')[-1] != 'pyi':
+                    continue
+
                 if stub_name != '__init__.pyi':
                     ext_module = os.path.join(ext_module, os.path.splitext(stub_name)[0])