diff --git a/src/prefect/main.py b/src/prefect/main.py
index 42b662c4e352..9b25e827f954 100644
--- a/src/prefect/main.py
+++ b/src/prefect/main.py
@@ -24,12 +24,13 @@
 import prefect.context
 
 # Perform any forward-ref updates needed for Pydantic models
-import prefect.client.schemas
+from prefect.client.schemas.objects import State
+from prefect.client.schemas.actions import StateCreate
 
 prefect.context.FlowRunContext.model_rebuild()
 prefect.context.TaskRunContext.model_rebuild()
-prefect.client.schemas.State.model_rebuild()
-prefect.client.schemas.StateCreate.model_rebuild()
+State.model_rebuild()
+StateCreate.model_rebuild()
 Transaction.model_rebuild()
 
 # Configure logging