diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py
index 8ece0a569e..dbedeeca00 100644
--- a/studio/backend/models/__init__.py
+++ b/studio/backend/models/__init__.py
@@ -14,9 +14,6 @@ from .training import (
     TrainingRunDetailResponse,
     TrainingRunDeleteResponse,
     TrainingRunUpdateRequest,
-    DiffusionTrainingStartRequest,
-    DiffusionTrainingStartResponse,
-    DiffusionTrainingStatusResponse,
 )
 from .models import (
     CheckpointInfo,
@@ -75,9 +72,6 @@ from .data_recipe import (
 __all__ = [
     # Training schemas
     "TrainingStartRequest",
-    "DiffusionTrainingStartRequest",
-    "DiffusionTrainingStartResponse",
-    "DiffusionTrainingStatusResponse",
     "TrainingJobResponse",
     "TrainingStatus",
     "TrainingProgress",
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index b2ea6dfaa4..0c311a6e57 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -56,6 +56,8 @@ from models import (
     TrainingJobResponse,
     TrainingStatus,
     TrainingProgress,
+)
+from models.training import (
     DiffusionTrainingStartRequest,
     DiffusionTrainingStartResponse,
     DiffusionTrainingStatusResponse,