diff --git a/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlGenerateTask.kt b/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlGenerateTask.kt
index 09df1af..5f1a61a 100644
--- a/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlGenerateTask.kt
+++ b/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlGenerateTask.kt
@@ -3,13 +3,20 @@ package dev.ktml.gradle
 import dev.ktml.KtmlProcessor
 import org.gradle.api.DefaultTask
 import org.gradle.api.file.FileTree
+import org.gradle.api.file.ProjectLayout
+import org.gradle.api.provider.Provider
 import org.gradle.api.tasks.*
 import java.io.File
 import java.io.Serializable
+import javax.inject.Inject
 
 data class SourceInfo(val ktmlDir: String, val outputDir: String) : Serializable
 
-open class KtmlGenerateTask : DefaultTask() {
+open class KtmlGenerateTask @Inject constructor(private var projectLayout: ProjectLayout) : DefaultTask() {
+
+    @Internal
+    lateinit var errorCollectorProvider: Provider<KtmlCompilationErrorCollector>
+
     @Input
     lateinit var moduleName: String
 
@@ -23,25 +30,25 @@ open class KtmlGenerateTask : DefaultTask() {
     @SkipWhenEmpty
     @PathSensitive(PathSensitivity.RELATIVE)
     fun getInputFiles(): FileTree {
-        return project.files(dirSets.map { dirSet -> project.fileTree(dirSet.ktmlDir) { it.include("**/*.ktml") } }).asFileTree
+        return projectLayout.files(dirSets.map { it.ktmlDir }).asFileTree.matching { it.include("**/*.ktml") }
     }
 
     @OutputDirectories
     fun getOutputDirectories(): List<File> {
-        return dirSets.map { project.file(it.outputDir) }
+        return dirSets.flatMap { projectLayout.files(it.outputDir).files }
     }
 
     @TaskAction
     open fun generate() {
         dirSets.forEach { (ktmlDir, outputDir) ->
-            project.file(outputDir).deleteRecursively()
+            projectLayout.files(outputDir).forEach { it.deleteRecursively() }
             KtmlProcessor(
                 moduleName = moduleName,
                 templatePackage = templatePackage,
                 outputDirectory = outputDir,
                 removeContentComments = false,
             ).apply {
-                project.getErrorCollector().get().addProjectInfo(ProjectInfo(templatePackage, ktmlDir, outputDir, this))
+                errorCollectorProvider.get().addProjectInfo(ProjectInfo(templatePackage, ktmlDir, outputDir, this))
                 processRootDirectory(ktmlDir)
                 generateTemplateCode()
             }
diff --git a/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlPlugin.kt b/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlPlugin.kt
index d6c817b..998cbc6 100644
--- a/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlPlugin.kt
+++ b/integrations/gradle/src/main/kotlin/dev/ktml/gradle/KtmlPlugin.kt
@@ -83,14 +83,15 @@ open class KtmlPlugin : Plugin<Project> {
             }
         }
 
+        val errorCollectorProvider = project.getErrorCollector()
+
         val generateTask = project.tasks.register("generateKtml", KtmlGenerateTask::class.java) {
             it.moduleName = moduleName
             it.templatePackage = templatePackage
             it.dirSets = sets.map { it.second }
+            it.errorCollectorProvider = errorCollectorProvider
         }
 
-        val errorCollectorProvider = project.getErrorCollector()
-
         project.afterEvaluate {
             sets.forEach { it.first.kotlin.srcDir(it.second.outputDir) }