diff --git a/src/main/kotlin/mods/betterfoliage/integration/ShadersModIntegration.kt b/src/main/kotlin/mods/betterfoliage/integration/ShadersModIntegration.kt
index 8196387..169d78f 100644
--- a/src/main/kotlin/mods/betterfoliage/integration/ShadersModIntegration.kt
+++ b/src/main/kotlin/mods/betterfoliage/integration/ShadersModIntegration.kt
@@ -1,9 +1,12 @@
 package mods.betterfoliage.integration
 
 import mods.betterfoliage.BetterFoliage
+import mods.betterfoliage.BlockAliases
+import mods.betterfoliage.BufferBuilder_sVertexBuilder
+import mods.betterfoliage.SVertexBuilder
+import mods.betterfoliage.Shaders
 import mods.betterfoliage.render.pipeline.RenderCtxBase
 import mods.betterfoliage.render.pipeline.RenderCtxVanilla
-import mods.betterfoliage.resource.discovery.BakeWrapperManager
 import mods.betterfoliage.util.HasLogger
 import mods.betterfoliage.util.allAvailable
 import mods.betterfoliage.util.get
@@ -41,7 +44,7 @@ object ShadersModIntegration : HasLogger() {
      * @see mods.betterfoliage.loader.BetterFoliageTransformer
      */
     @JvmStatic fun getBlockStateOverride(state: BlockState, world: IBlockDisplayReader, pos: BlockPos): BlockState {
-//        if (LeafRegistry[state, world, pos] != null) return defaultLeaves
+        if (state in BetterFoliage.blockTypes.leaf) return defaultLeaves
 //        if (BlockConfig.crops.matchesClass(state.block)) return defaultGrass
         return state
     }
diff --git a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Leaf.kt b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Leaf.kt
index ed5a859..f1ba234 100644
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Leaf.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Leaf.kt
@@ -39,6 +39,7 @@ object StandardLeafDiscovery : ParametrizedModelDiscovery() {
 
         detailLogger.log(INFO, "     particle $leafType")
         ctx.addReplacement(StandardLeafKey(generated, leafType, null))
+        BetterFoliage.blockTypes.leaf.add(ctx.blockState)
     }
 }