diff --git a/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/FeatureGenerationAddon.java b/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/FeatureGenerationAddon.java
index b59f1774b..10c57d320 100644
--- a/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/FeatureGenerationAddon.java
+++ b/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/FeatureGenerationAddon.java
@@ -20,12 +20,13 @@ import com.dfsek.terra.api.event.functional.FunctionalEventHandler;
 import com.dfsek.terra.api.inject.annotations.Inject;
 import com.dfsek.terra.api.util.reflection.TypeKey;
 import com.dfsek.terra.api.world.biome.Biome;
+import com.dfsek.terra.api.world.chunk.generation.stage.GenerationStage;
 
 import java.util.function.Supplier;
 
 
 public class FeatureGenerationAddon implements AddonInitializer {
-    public static final TypeKey<Supplier<ObjectTemplate<FeatureGenerationStage>>> STAGE_TYPE_KEY = new TypeKey<>() {};
+    public static final TypeKey<Supplier<ObjectTemplate<GenerationStage>>> STAGE_TYPE_KEY = new TypeKey<>() {};
     @Inject
     private Platform platform;
     
diff --git a/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/config/FeatureStageTemplate.java b/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/config/FeatureStageTemplate.java
index d280dd8ad..b2652aed8 100644
--- a/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/config/FeatureStageTemplate.java
+++ b/common/addons/generation-stage-feature/src/main/java/com/dfsek/terra/addons/generation/feature/config/FeatureStageTemplate.java
@@ -5,9 +5,10 @@ import com.dfsek.tectonic.loading.object.ObjectTemplate;
 
 import com.dfsek.terra.addons.generation.feature.FeatureGenerationStage;
 import com.dfsek.terra.api.Platform;
+import com.dfsek.terra.api.world.chunk.generation.stage.GenerationStage;
 
 
-public class FeatureStageTemplate implements ObjectTemplate<FeatureGenerationStage> {
+public class FeatureStageTemplate implements ObjectTemplate<GenerationStage> {
     @Value("id")
     private String id;