diff --git a/llama_stack/cli/model/safety_models.py b/llama_stack/cli/model/safety_models.py
index 39c133f73..9464e0a2d 100644
--- a/llama_stack/cli/model/safety_models.py
+++ b/llama_stack/cli/model/safety_models.py
@@ -6,11 +6,12 @@
 
 from typing import Any, Dict, Optional
 
-from pydantic import BaseModel, ConfigDict, Field
-
-from llama_models.datatypes import *  # noqa: F403
+from llama_models.datatypes import CheckpointQuantizationFormat
+from llama_models.llama3.api.datatypes import SamplingParams
 from llama_models.sku_list import LlamaDownloadInfo
 
+from pydantic import BaseModel, ConfigDict, Field
+
 
 class PromptGuardModel(BaseModel):
     """Make a 'fake' Model-like object for Prompt Guard. Eventually this will be removed."""