diff --git a/migrations/env.py b/migrations/env.py
index b9ac505..db84064 100644
--- a/migrations/env.py
+++ b/migrations/env.py
@@ -7,7 +7,7 @@ from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import async_engine_from_config
 from sqlmodel import SQLModel
 
-# from app.models import Bookmark
+from src.digimarks.main import Bookmark, PublicTag, User
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -62,10 +62,7 @@ def do_run_migrations(connection: Connection) -> None:
 
 
 async def run_async_migrations() -> None:
-    """In this scenario we need to create an Engine
-    and associate a connection with the context.
-
-    """
+    """In this scenario we need to create an Engine and associate a connection with the context."""
     connectable = async_engine_from_config(
         config.get_section(config.config_ini_section, {}),
         prefix='sqlalchemy.',
diff --git a/migrations/script.py.mako b/migrations/script.py.mako
index 1101630..697cf67 100644
--- a/migrations/script.py.mako
+++ b/migrations/script.py.mako
@@ -9,6 +9,7 @@ from typing import Sequence, Union
 
 from alembic import op
 import sqlalchemy as sa
+import sqlmodel
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.