diff --git a/internal/schema/schema.go b/internal/schema/schema.go
index 071aaa50..467e32ea 100644
--- a/internal/schema/schema.go
+++ b/internal/schema/schema.go
@@ -28,19 +28,19 @@ type GeneratorSchema struct {
 func (g GeneratorSchema) ImportsString() (string, error) {
 	imports := NewImports()
 
-	// Both context and terraform-plugin-framework/diag packages are required if
-	// model object helpers are generated. Refer to the logic in
-	// ModelsObjectHelpersBytes() method.
 	for _, a := range g.Attributes {
 		if a == nil {
 			continue
 		}
 
+		imports.Add([]code.Import{
+			{
+				Path: ContextImport,
+			},
+		}...)
+
 		if _, ok := a.(Attributes); ok {
 			imports.Add([]code.Import{
-				{
-					Path: ContextImport,
-				},
 				{
 					Path: FmtImport,
 				},
@@ -56,23 +56,26 @@ func (g GeneratorSchema) ImportsString() (string, error) {
 				{
 					Path: TfTypesImport,
 				},
+				{
+					Path: BaseTypesImport,
+				},
 			}...)
 		}
 	}
 
-	// Both context and terraform-plugin-framework/diag packages are required if
-	// model object helpers are generated. Refer to the logic in
-	// ModelsObjectHelpersBytes() method.
 	for _, b := range g.Blocks {
 		if b == nil {
 			continue
 		}
 
+		imports.Add([]code.Import{
+			{
+				Path: ContextImport,
+			},
+		}...)
+
 		if _, ok := b.(Blocks); ok {
 			imports.Add([]code.Import{
-				{
-					Path: ContextImport,
-				},
 				{
 					Path: FmtImport,
 				},
@@ -88,6 +91,9 @@ func (g GeneratorSchema) ImportsString() (string, error) {
 				{
 					Path: TfTypesImport,
 				},
+				{
+					Path: BaseTypesImport,
+				},
 			}...)
 		}
 	}