[WIP] adopt model replacement from Forge vesion

+ bunch of renames to bring the 2 version closer + at-least-not-crashing levels of Optifine support
2021-05-13 21:13:10 +02:00
parent 9dacdde761
commit 40fd46b278
51 changed files with 1082 additions and 749 deletions
--- a/src/main/java/mods/betterfoliage/MixinConfigPlugin.java
+++ b/src/main/java/mods/betterfoliage/MixinConfigPlugin.java
@@ -0,0 +1,49 @@
+package mods.betterfoliage;
+
+import net.fabricmc.loader.api.FabricLoader;
+import org.apache.logging.log4j.Level;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.objectweb.asm.tree.ClassNode;
+import org.spongepowered.asm.mixin.extensibility.IMixinConfigPlugin;
+import org.spongepowered.asm.mixin.extensibility.IMixinInfo;
+
+import java.util.List;
+import java.util.Set;
+
+public class MixinConfigPlugin implements IMixinConfigPlugin {
+
+    Logger logger = LogManager.getLogger(this);
+
+    Boolean hasOptifine = null;
+
+    @Override
+    public boolean shouldApplyMixin(String targetClassName, String mixinClassName) {
+        if (hasOptifine == null) {
+            hasOptifine = FabricLoader.getInstance().isModLoaded("optifabric");
+            if (hasOptifine) logger.log(Level.INFO, "[BetterFoliage] Optifabric detected, applying Optifine mixins");
+            else logger.log(Level.INFO, "[BetterFoliage] Optifabric not detected, applying Vanilla mixins");
+        }
+        if (mixinClassName.endsWith("Vanilla") && hasOptifine) return false;
+        if (mixinClassName.endsWith("Optifine") && !hasOptifine) return false;
+        return true;
+    }
+
+    @Override
+    public void onLoad(String mixinPackage) { }
+
+    @Override
+    public String getRefMapperConfig() { return null; }
+
+    @Override
+    public void acceptTargets(Set<String> myTargets, Set<String> otherTargets) { }
+
+    @Override
+    public List<String> getMixins() { return null; }
+
+    @Override
+    public void preApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) { }
+
+    @Override
+    public void postApply(String targetClassName, ClassNode targetClass, String mixinClassName, IMixinInfo mixinInfo) { }
+}
--- a/src/main/java/mods/betterfoliage/mixin/MixinModelLoader.java
+++ b/src/main/java/mods/betterfoliage/mixin/MixinModelLoader.java
@@ -1,27 +1,36 @@
 package mods.betterfoliage.mixin;

 import mods.betterfoliage.ModelLoadingCallback;
+import mods.betterfoliage.resource.discovery.BakeWrapperManager;
+import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.ModelBakeSettings;
 import net.minecraft.client.render.model.ModelLoader;
+import net.minecraft.client.render.model.UnbakedModel;
+import net.minecraft.client.texture.Sprite;
 import net.minecraft.client.util.ModelIdentifier;
+import net.minecraft.client.util.SpriteIdentifier;
 import net.minecraft.resource.ResourceManager;
+import net.minecraft.util.Identifier;
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;

-import static net.minecraft.client.render.model.ModelLoader.MISSING;
+import java.util.function.Function;

@Mixin(ModelLoader.class)
 public class MixinModelLoader {

    @Shadow @Final private ResourceManager resourceManager;

+    // use the same trick fabric-api does to get around the no-mixins-in-constructors policy
    @Inject(at = @At("HEAD"), method = "addModel")
    private void addModelHook(ModelIdentifier id, CallbackInfo info) {
-        // use the same trick fabric-api does to get around the no-mixins-in-constructors policy
-        if (id == MISSING) {
+        if (id.getPath().equals("trident_in_hand")) {
+            // last step before stitching
            ModelLoadingCallback.EVENT.invoker().beginLoadModels((ModelLoader) (Object) this, resourceManager);
        }
    }
--- a/src/main/java/mods/betterfoliage/mixin/MixinModelLoaderOptifine.java
+++ b/src/main/java/mods/betterfoliage/mixin/MixinModelLoaderOptifine.java
@@ -0,0 +1,34 @@
+package mods.betterfoliage.mixin;
+
+import mods.betterfoliage.resource.discovery.BakeWrapperManager;
+import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.ModelBakeSettings;
+import net.minecraft.client.render.model.ModelLoader;
+import net.minecraft.client.render.model.UnbakedModel;
+import net.minecraft.client.texture.Sprite;
+import net.minecraft.client.util.SpriteIdentifier;
+import net.minecraft.util.Identifier;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Redirect;
+
+import java.util.function.Function;
+
+@Mixin(ModelLoader.class)
+public class MixinModelLoaderOptifine {
+
+    private static final String loaderBake = "Lnet/minecraft/class_1088;getBakedModel(Lnet/minecraft/class_2960;Lnet/minecraft/class_3665;Ljava/util/function/Function;)Lnet/minecraft/class_1087;";
+    private static final String modelBake = "Lnet/minecraft/class_1100;method_4753(Lnet/minecraft/class_1088;Ljava/util/function/Function;Lnet/minecraft/class_3665;Lnet/minecraft/class_2960;)Lnet/minecraft/class_1087;";
+
+    @SuppressWarnings("UnresolvedMixinReference")
+    @Redirect(method = loaderBake, at = @At(value = "INVOKE", target = modelBake), remap = false)
+    BakedModel onBakeModel(
+            UnbakedModel unbaked,
+            ModelLoader loader,
+            Function<SpriteIdentifier, Sprite> textureGetter,
+            ModelBakeSettings rotationContainer,
+            Identifier modelId
+    ) {
+        return BakeWrapperManager.INSTANCE.onBake(unbaked, loader, textureGetter, rotationContainer, modelId);
+    }
+}
--- a/src/main/java/mods/betterfoliage/mixin/MixinModelLoaderVanilla.java
+++ b/src/main/java/mods/betterfoliage/mixin/MixinModelLoaderVanilla.java
@@ -0,0 +1,34 @@
+package mods.betterfoliage.mixin;
+
+import mods.betterfoliage.resource.discovery.BakeWrapperManager;
+import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.ModelBakeSettings;
+import net.minecraft.client.render.model.ModelLoader;
+import net.minecraft.client.render.model.UnbakedModel;
+import net.minecraft.client.texture.Sprite;
+import net.minecraft.client.util.SpriteIdentifier;
+import net.minecraft.util.Identifier;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Redirect;
+
+import java.util.function.Function;
+
+@Mixin(ModelLoader.class)
+public class MixinModelLoaderVanilla {
+
+    private static final String loaderBake = "Lnet/minecraft/client/render/model/ModelLoader;bake(Lnet/minecraft/util/Identifier;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;";
+    private static final String modelBake = "Lnet/minecraft/client/render/model/UnbakedModel;bake(Lnet/minecraft/client/render/model/ModelLoader;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Lnet/minecraft/util/Identifier;)Lnet/minecraft/client/render/model/BakedModel;";
+
+    @Redirect(method = loaderBake, at = @At(value = "INVOKE", target = modelBake))
+    BakedModel onBakeModel(
+            UnbakedModel unbaked,
+            ModelLoader loader,
+            Function<SpriteIdentifier, Sprite> textureGetter,
+            ModelBakeSettings rotationContainer,
+            Identifier modelId
+    ) {
+        return BakeWrapperManager.INSTANCE.onBake(unbaked, loader, textureGetter, rotationContainer, modelId);
+    }
+}
+
--- a/src/main/kotlin/mods/betterfoliage/BetterFoliage.kt
+++ b/src/main/kotlin/mods/betterfoliage/BetterFoliage.kt
@@ -8,11 +8,13 @@ import mods.betterfoliage.render.ShadersModIntegration
 import mods.betterfoliage.render.block.vanilla.*
 import mods.betterfoliage.render.particle.LeafParticleRegistry
 import mods.betterfoliage.render.particle.RisingSoulParticle
-import mods.betterfoliage.resource.discovery.BakedModelReplacer
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.BlockTypeCache
 import mods.betterfoliage.resource.generated.GeneratedBlockTexturePack
 import net.fabricmc.api.ClientModInitializer
 import net.fabricmc.fabric.api.resource.ResourceManagerHelper
 import net.fabricmc.loader.api.FabricLoader
+import net.minecraft.block.BlockState
 import net.minecraft.client.MinecraftClient
 import net.minecraft.resource.ResourceType
 import net.minecraft.util.Identifier
@@ -28,18 +30,14 @@ import java.util.*
 object BetterFoliage : ClientModInitializer {
    const val MOD_ID = "betterfoliage"

-    var logger = LogManager.getLogger()
-    var logDetail = SimpleLogger(
-        "BetterFoliage",
-        Level.DEBUG,
-        false, false, true, false,
-        "yyyy-MM-dd HH:mm:ss",
-        null,
-        PropertiesUtil(Properties()),
-        PrintStream(File(FabricLoader.getInstance().gameDirectory, "logs/betterfoliage.log").apply {
-            parentFile.mkdirs()
-            if (!exists()) createNewFile()
-        })
+    val detailLogStream = PrintStream(File("logs/betterfoliage.log").apply {
+        parentFile.mkdirs()
+        if (!exists()) createNewFile()
+    })
+
+    fun logger(obj: Any) = LogManager.getLogger(obj)
+    fun detailLogger(obj: Any) = SimpleLogger(
+        obj::class.java.simpleName, Level.DEBUG, false, true, true, false, "yyyy-MM-dd HH:mm:ss", null, PropertiesUtil(Properties()), detailLogStream
    )

    val configFile get() = File(FabricLoader.getInstance().configDirectory, "BetterFoliage.json")
@@ -50,39 +48,43 @@ object BetterFoliage : ClientModInitializer {
    }

    val blockConfig = BlockConfig()
-    val generatedPack = GeneratedBlockTexturePack(Identifier(MOD_ID, "generated"), "betterfoliage-generated", "Better Foliage", "Generated leaf textures", logDetail)
-    val modelReplacer = BakedModelReplacer()
+    val generatedPack = GeneratedBlockTexturePack(Identifier(MOD_ID, "generated"), "betterfoliage-generated", "Better Foliage", "Generated leaf textures")
+
+    /** List of recognized [BlockState]s */
+    var blockTypes = BlockTypeCache()

    override fun onInitializeClient() {
        // Register generated resource pack
        ResourceManagerHelper.get(ResourceType.CLIENT_RESOURCES).registerReloadListener(generatedPack.reloader)
        MinecraftClient.getInstance().resourcePackManager.registerProvider(generatedPack.finder)

+        ResourceManagerHelper.get(ResourceType.CLIENT_RESOURCES).registerReloadListener(blockConfig)
+
        // Add standard block support
-        modelReplacer.discoverers.add(StandardLeafDiscovery)
-        modelReplacer.discoverers.add(StandardGrassDiscovery)
-        modelReplacer.discoverers.add(StandardLogDiscovery)
-        modelReplacer.discoverers.add(StandardCactusDiscovery)
-        modelReplacer.discoverers.add(LilyPadDiscovery)
-        modelReplacer.discoverers.add(DirtDiscovery)
-        modelReplacer.discoverers.add(SandDiscovery)
-        modelReplacer.discoverers.add(MyceliumDiscovery)
-        modelReplacer.discoverers.add(NetherrackDiscovery)
+        BakeWrapperManager.discoverers.add(StandardCactusDiscovery)
+        BakeWrapperManager.discoverers.add(StandardDirtDiscovery)
+        BakeWrapperManager.discoverers.add(StandardGrassDiscovery)
+        BakeWrapperManager.discoverers.add(StandardLeafDiscovery)
+        BakeWrapperManager.discoverers.add(StandardLilypadDiscovery)
+        BakeWrapperManager.discoverers.add(StandardMyceliumDiscovery)
+        BakeWrapperManager.discoverers.add(StandardNetherrackDiscovery)
+        BakeWrapperManager.discoverers.add(StandardRoundLogDiscovery)
+        BakeWrapperManager.discoverers.add(StandardSandDiscovery)

        // Init overlay layers
        ChunkOverlayManager.layers.add(RoundLogOverlayLayer)

        // Init singletons
        LeafParticleRegistry
-        NormalLeavesModel.Companion
-        GrassBlockModel.Companion
-        RoundLogModel.Companion
-        CactusModel.Companion
-        LilypadModel.Companion
+        StandardLeafModel.Companion
+        StandardGrassModel.Companion
+        StandardRoundLogModel.Companion
+        StandardCactusModel.Companion
+        StandardLilypadModel.Companion
        DirtModel.Companion
-        SandModel.Companion
-        MyceliumModel.Companion
-        NetherrackModel.Companion
+        StandardSandModel.Companion
+        StandardMyceliumModel.Companion
+        StandardNetherrackModel.Companion
        RisingSoulParticle.Companion
        ShadersModIntegration
    }
--- a/src/main/kotlin/mods/betterfoliage/Hooks.kt
+++ b/src/main/kotlin/mods/betterfoliage/Hooks.kt
@@ -2,15 +2,17 @@
 package mods.betterfoliage

 import mods.betterfoliage.chunk.ChunkOverlayManager
-import mods.betterfoliage.render.block.vanilla.LeafKey
+import mods.betterfoliage.render.block.vanilla.LeafParticleKey
 import mods.betterfoliage.render.block.vanilla.RoundLogKey
 import mods.betterfoliage.render.particle.FallingLeafParticle
 import mods.betterfoliage.render.particle.RisingSoulParticle
 import mods.betterfoliage.util.offset
 import mods.betterfoliage.util.plus
+import mods.betterfoliage.util.random
 import mods.betterfoliage.util.randomD
 import net.minecraft.block.BlockState
 import net.minecraft.block.Blocks
+import net.minecraft.client.MinecraftClient
 import net.minecraft.client.world.ClientWorld
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction
@@ -21,7 +23,7 @@ import net.minecraft.world.BlockView
 fun getAmbientOcclusionLightValueOverride(original: Float, state: BlockState): Float {
    if (BetterFoliage.config.enabled &&
        BetterFoliage.config.roundLogs.enabled &&
-        BetterFoliage.modelReplacer.getTyped<RoundLogKey>(state) != null
+        BetterFoliage.blockTypes.hasTyped<RoundLogKey>(state)
    ) return BetterFoliage.config.roundLogs.dimming.toFloat()
    return original
 }
@@ -49,14 +51,15 @@ fun onRandomDisplayTick(world: ClientWorld, pos: BlockPos) {
        BetterFoliage.config.fallingLeaves.enabled &&
        world.isAir(pos + Direction.DOWN.offset) &&
        randomD() < BetterFoliage.config.fallingLeaves.chance) {
-            BetterFoliage.modelReplacer.getTyped<LeafKey>(state)?.let { key ->
-                FallingLeafParticle(world, pos, key).addIfValid()
+            BetterFoliage.blockTypes.getTyped<LeafParticleKey>(state)?.let { key ->
+                val blockColor = MinecraftClient.getInstance().blockColorMap.getColor(state, world, pos, 0)
+                FallingLeafParticle(world, pos, key, blockColor, random).addIfValid()
            }
    }
 }

 fun getVoxelShapeOverride(state: BlockState, reader: BlockView, pos: BlockPos, dir: Direction): VoxelShape {
-    if (BetterFoliage.modelReplacer[state] is RoundLogKey) {
+    if (BetterFoliage.blockTypes.hasTyped<RoundLogKey>(state)) {
        return VoxelShapes.empty()
    }
    // TODO ?
--- a/src/main/kotlin/mods/betterfoliage/config/BlockConfig.kt
+++ b/src/main/kotlin/mods/betterfoliage/config/BlockConfig.kt
@@ -1,13 +1,13 @@
 package mods.betterfoliage.config

 import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.util.Invalidator
+import mods.betterfoliage.resource.VeryEarlyReloadListener
 import mods.betterfoliage.resource.discovery.ConfigurableBlockMatcher
 import mods.betterfoliage.resource.discovery.ModelTextureListConfiguration
 import net.minecraft.resource.ResourceManager
 import net.minecraft.util.Identifier

-class BlockConfig  {
+class BlockConfig : VeryEarlyReloadListener {
    private val list = mutableListOf<Any>()

    val leafBlocks = blocks("leaves_blocks_default.cfg")
@@ -24,10 +24,13 @@ class BlockConfig  {
 //    val cactus = blocks("cactus_default.cfg")
 //    val netherrack = blocks("netherrack_blocks_default.cfg")

-    private fun blocks(cfgName: String) = ConfigurableBlockMatcher(BetterFoliage.logDetail, Identifier(BetterFoliage.MOD_ID, cfgName)).apply { list.add(this) }
-    private fun models(cfgName: String) = ModelTextureListConfiguration(BetterFoliage.logDetail, Identifier(BetterFoliage.MOD_ID, cfgName)).apply { list.add(this) }
+    private fun blocks(cfgName: String) = ConfigurableBlockMatcher(Identifier(BetterFoliage.MOD_ID, cfgName)).apply { list.add(this) }
+    private fun models(cfgName: String) = ModelTextureListConfiguration(Identifier(BetterFoliage.MOD_ID, cfgName)).apply { list.add(this) }

-    fun reloadConfig(manager: ResourceManager) {
+
+    override fun getFabricId() = Identifier(BetterFoliage.MOD_ID, "block-config")
+
+    override fun onReloadStarted(manager: ResourceManager) {
        list.forEach { when(it) {
            is ConfigurableBlockMatcher -> it.readDefaults(manager)
            is ModelTextureListConfiguration -> it.readDefaults(manager)
--- a/src/main/kotlin/mods/betterfoliage/config/MainConfig.kt
+++ b/src/main/kotlin/mods/betterfoliage/config/MainConfig.kt
@@ -40,6 +40,7 @@ class LeavesConfig(node: ConfigNode) : DelegatingConfigGroup(node) {
    val vOffset by double(0.1, min = 0.0, max = 0.4, langKey = recurring)
    val size by double(1.4, min = 0.75, max = 2.5, langKey = recurring)
    val shaderWind by boolean(true, langKey = recurring)
+    val saturationThreshold by double(0.1, min = 0.0, max = 1.0)
 }

 class ShortGrassConfig(node: ConfigNode) : DelegatingConfigGroup(node), PopulationConfigData {
@@ -71,14 +72,14 @@ class RoundLogConfig(node: ConfigNode) : DelegatingConfigGroup(node) {
    val connectGrass by boolean(true)

    val radiusSmall by double(0.25, min = 0.0, max = 0.5)
-    val radiusLarge by double(0.25, min = 0.0, max = 0.5) { it.coerceAtLeast(radiusSmall) }
+    val radiusLarge by double(0.44, min = 0.0, max = 0.5) { it.coerceAtLeast(radiusSmall) }
    val dimming by double(0.7, min = 0.0, max = 1.0)
    val zProtection by double(0.99, min = 0.9, max = 1.0)
 }

 class CactusConfig(node: ConfigNode) : DelegatingConfigGroup(node) {
    val enabled by boolean(true, langKey = recurring)
-    val size by double(1.3, min = 0.5, max = 1.5, langKey = recurring)
+    val size by double(1.3, min = 1.0, max = 2.0, langKey = recurring)
    val sizeVariation by double(0.1, min = 0.0, max = 0.5)
    val hOffset by double(0.1, min = 0.0, max = 0.5, langKey = recurring)
 }
--- a/src/main/kotlin/mods/betterfoliage/config/MiscDefaults.kt
+++ b/src/main/kotlin/mods/betterfoliage/config/MiscDefaults.kt
@@ -1,12 +1,13 @@
-package mods.betterfoliage.render
+package mods.betterfoliage.config

 import net.minecraft.block.Blocks
 import net.minecraft.block.Material
 import net.minecraft.world.biome.Biome

+val CACTUS_BLOCKS = listOf(Blocks.CACTUS)
 val DIRT_BLOCKS = listOf(Blocks.DIRT, Blocks.COARSE_DIRT, Blocks.PODZOL)
 val SAND_BLOCKS = listOf(Blocks.SAND, Blocks.RED_SAND)
+val NETHERRACK_BLOCKS = listOf(Blocks.NETHERRACK)

 val SALTWATER_BIOMES = listOf(Biome.Category.BEACH, Biome.Category.OCEAN)
-
 val SNOW_MATERIALS = listOf(Material.SNOW_BLOCK)
--- a/src/main/kotlin/mods/betterfoliage/integration/ModMenu.kt
+++ b/src/main/kotlin/mods/betterfoliage/integration/ModMenu.kt
@@ -4,6 +4,7 @@ import io.github.prospector.modmenu.api.ModMenuApi
 import me.shedaniel.clothconfig2.api.ConfigBuilder
 import me.zeroeightsix.fiber.JanksonSettings
 import mods.betterfoliage.BetterFoliage
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
 import net.minecraft.client.MinecraftClient
 import net.minecraft.client.gui.screen.Screen
 import net.minecraft.client.resource.language.I18n
@@ -21,7 +22,7 @@ object ModMenu : ModMenuApi {
        }
        builder.savingRunnable = Runnable {
            JanksonSettings().serialize(BetterFoliage.config.fiberNode, BetterFoliage.configFile.outputStream(), false)
-            BetterFoliage.modelReplacer.invalidate()
+            BakeWrapperManager.invalidate()
            MinecraftClient.getInstance().worldRenderer.reload()
        }
        builder.build()
--- a/src/main/kotlin/mods/betterfoliage/model/Meshify.kt
+++ b/src/main/kotlin/mods/betterfoliage/model/Meshify.kt
@@ -17,12 +17,9 @@ import net.minecraft.client.render.VertexFormatElement.Type.UV
 import net.minecraft.client.render.VertexFormats
 import net.minecraft.client.render.model.BakedModel
 import net.minecraft.client.render.model.BakedQuad
+import net.minecraft.client.render.model.BasicBakedModel
 import net.minecraft.util.math.Direction
-import java.lang.Float
 import java.util.*
-import kotlin.Boolean
-import kotlin.Int
-import kotlin.let

 interface BakedModelConverter {
    /**
@@ -30,12 +27,12 @@ interface BakedModelConverter {
     * @param model Input model
     * @param converter Converter to use for converting nested models.
     */
-    fun convert(model: BakedModel, converter: BakedModelConverter): BakedModel?
+    fun convert(model: BakedModel): BakedModel?
    companion object {
-        fun of(func: (BakedModel, BakedModelConverter)->BakedModel?) = object : BakedModelConverter {
-            override fun convert(model: BakedModel, converter: BakedModelConverter) = func(model, converter)
+        fun of(func: (BakedModel)->BakedModel?) = object : BakedModelConverter {
+            override fun convert(model: BakedModel) = func(model)
        }
-        val identity = of { model, _ -> model }
+        val identity = of { model -> model }
    }
 }

@@ -45,29 +42,29 @@ interface BakedModelConverter {
 */
 fun List<BakedModelConverter>.convert(model: BakedModel) = object : BakedModelConverter {
    val converters = this@convert + BakedModelConverter.identity
-    override fun convert(model: BakedModel, converter: BakedModelConverter) = converters.findFirst { it.convert(model, converter) }
+    override fun convert(model: BakedModel) = converters.findFirst { it.convert(model) }
 }.let { converterStack ->
    // we are guaranteed a result here because of the identity converter
-    converterStack.convert(model, converterStack)!!
+    converterStack.convert(model)!!
 }

-/** List of converters without meaningful configuration that should always be used */
-val COMMON_MESH_CONVERTERS = listOf(WrappedWeightedModel.converter)
-
 /**
- * Convert [BakedModel] into one using fabric-rendering-api [Mesh] instead of the vanilla pipeline.
- * @param blendModeOverride Use the given [BlockRenderLayer] for the [Mesh]
+ * Convert [BasicBakedModel] into one using fabric-rendering-api [Mesh] instead of the vanilla pipeline.
+ * @param blendMode Use the given [BlockRenderLayer] for the [Mesh]
 * instead of the one declared by the corresponding [Block]
 */
-fun meshifyStandard(model: BakedModel, state: BlockState, blendModeOverride: BlendMode? = null) =
-    (COMMON_MESH_CONVERTERS + WrappedMeshModel.converter(state, blendModeOverride = blendModeOverride)).convert(model)
+fun meshifyStandard(model: BasicBakedModel, state: BlockState? = null, blendMode: BlendMode? = null) =
+    WrappedMeshModel.converter(state, blendModeOverride = blendMode).convert(model)!!
+
+fun meshifySolid(model: BasicBakedModel) = meshifyStandard(model, null, BlendMode.SOLID)
+fun meshifyCutoutMipped(model: BasicBakedModel) = meshifyStandard(model, null, BlendMode.CUTOUT_MIPPED)

 /**
 * Convert a vanilla [BakedModel] into intermediate [Quad]s
 * Vertex normals not supported (yet)
 * Vertex data elements not aligned to 32 bit boundaries not supported
 */
-fun unbakeQuads(model: BakedModel, state: BlockState, random: Random, unshade: Boolean): List<Quad> {
+fun unbakeQuads(model: BakedModel, state: BlockState?, random: Random, unshade: Boolean): List<Quad> {
    return (allDirections.toList() + null as Direction?).flatMap { face ->
        model.getQuads(state, face, random).mapIndexed { qIdx, bakedQuad ->
            var quad = Quad(Vertex(), Vertex(), Vertex(), Vertex(), face = face, colorIndex = bakedQuad.colorIndex, sprite = bakedQuad[BakedQuad_sprite])
@@ -76,9 +73,9 @@ fun unbakeQuads(model: BakedModel, state: BlockState, random: Random, unshade: B
            val stride = format.vertexSizeInteger
            format.getIntOffset(POSITION, FLOAT, 3)?.let { posOffset ->
                quad = quad.transformVI { vertex, vIdx -> vertex.copy(xyz = Double3(
-                    x = Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + posOffset + 0]).toDouble(),
-                    y = Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + posOffset + 1]).toDouble(),
-                    z = Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + posOffset + 2]).toDouble()
+                    x = java.lang.Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + posOffset + 0]).toDouble(),
+                    y = java.lang.Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + posOffset + 1]).toDouble(),
+                    z = java.lang.Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + posOffset + 2]).toDouble()
                )) }
            }
            format.getIntOffset(COLOR, UBYTE, 4)?.let { colorOffset ->
@@ -88,8 +85,8 @@ fun unbakeQuads(model: BakedModel, state: BlockState, random: Random, unshade: B
            }
            format.getIntOffset(UV, FLOAT, 2, 0)?.let { uvOffset ->
                quad = quad.transformVI { vertex, vIdx -> vertex.copy(uv = UV(
-                    u = Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + uvOffset + 0]).toDouble(),
-                    v = Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + uvOffset + 1]).toDouble()
+                    u = java.lang.Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + uvOffset + 0]).toDouble(),
+                    v = java.lang.Float.intBitsToFloat(bakedQuad.vertexData[vIdx * stride + uvOffset + 1]).toDouble()
                )) }
            }

--- a/src/main/kotlin/mods/betterfoliage/model/Quads.kt
+++ b/src/main/kotlin/mods/betterfoliage/model/Quads.kt
@@ -12,9 +12,9 @@ import net.minecraft.util.math.Direction
 import net.minecraft.util.math.Direction.*
 import java.lang.Math.max
 import java.lang.Math.min
+import java.util.Random
 import kotlin.math.cos
 import kotlin.math.sin
-import kotlin.random.Random

 /**
 * Vertex UV coordinates
@@ -172,7 +172,7 @@ fun List<Quad>.build(blendMode: BlendMode, noDiffuse: Boolean = false, flatLight
    val builder = renderer.meshBuilder()
    builder.emitter.apply {
        forEach { quad ->
-            val sprite = quad.sprite ?: Atlas.BLOCKS.atlas[MissingSprite.getMissingSpriteId()]!!
+            val sprite = quad.sprite ?: Atlas.BLOCKS[MissingSprite.getMissingSpriteId()]!!
            quad.verts.forEachIndexed { idx, vertex ->
                pos(idx, (vertex.xyz + Double3(0.5, 0.5, 0.5)).asVec3f)
                sprite(idx, 0,
--- a/src/main/kotlin/mods/betterfoliage/model/SpriteSets.kt
+++ b/src/main/kotlin/mods/betterfoliage/model/SpriteSets.kt
@@ -21,8 +21,8 @@ class FixedSpriteSet(val sprites: List<Sprite>) : SpriteSet {
    override fun get(idx: Int) = sprites[idx % num]

    constructor(atlas: Atlas, ids: List<Identifier>) : this(
-        ids.mapNotNull { atlas.atlas[it] }.let { sprites ->
-            if (sprites.isNotEmpty()) sprites else listOf(atlas.atlas[MissingSprite.getMissingSpriteId()]!!)
+        ids.mapNotNull { atlas[it] }.let { sprites ->
+            if (sprites.isNotEmpty()) sprites else listOf(atlas[MissingSprite.getMissingSpriteId()]!!)
        }
    )
 }
@@ -42,7 +42,7 @@ class SpriteDelegate(val atlas: Atlas, val idFunc: ()->Identifier) : ReadOnlyPro
        value?.let { return it }
        synchronized(this) {
            value?.let { return it }
-            atlas.atlas[id!!]!!.let { value = it; return it }
+            atlas[id!!]!!.let { value = it; return it }
        }
    }
 }
@@ -60,7 +60,7 @@ class SpriteSetDelegate(
    override fun registerSprites(atlasTexture: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
        spriteSet = null
        val manager = MinecraftClient.getInstance().resourceManager
-        idList = (0 until 16).map(idFunc).filter { manager.containsResource(atlas.wrap(it)) }.map(idRegister)
+        idList = (0 until 16).map(idFunc).filter { manager.containsResource(atlas.file(it)) }.map(idRegister)
        idList.forEach { registry.register(it) }
    }

--- a/src/main/kotlin/mods/betterfoliage/model/TuftMeshes.kt
+++ b/src/main/kotlin/mods/betterfoliage/model/TuftMeshes.kt
@@ -31,21 +31,21 @@ fun tuftQuadSingle(size: Double, height: Double, flipU: Boolean) =
    verticalRectangle(x1 = -0.5 * size, z1 = 0.5 * size, x2 = 0.5 * size, z2 = -0.5 * size, yBottom = 0.5, yTop = 0.5 + height)
        .mirrorUV(flipU, false)

-fun tuftModelSet(shapes: Array<TuftShapeKey>, overrideColor: Int?, spriteGetter: (Int)->Sprite) = shapes.mapIndexed { idx, shape ->
+fun tuftModelSet(shapes: Array<TuftShapeKey>, tintIndex: Int, spriteGetter: (Int)->Sprite) = shapes.mapIndexed { idx, shape ->
    listOf(
        tuftQuadSingle(shape.size, shape.height, shape.flipU1),
        tuftQuadSingle(shape.size, shape.height, shape.flipU2).rotate(rot(UP))
    ).map { it.move(shape.offset) }
-        .map { it.colorAndIndex(overrideColor) }
+        .map { it.colorIndex(tintIndex) }
        .map { it.sprite(spriteGetter(idx)) }
 }.toTypedArray()

-fun fullCubeTextured(spriteId: Identifier, overrideColor: Int?, scrambleUV: Boolean = true): Mesh {
-    val sprite = Atlas.BLOCKS.atlas[spriteId]!!
+fun fullCubeTextured(spriteId: Identifier, tintIndex: Int, scrambleUV: Boolean = true): Mesh {
+    val sprite = Atlas.BLOCKS[spriteId]!!
    return allDirections.map { faceQuad(it) }
        .map { if (!scrambleUV) it else it.rotateUV(randomI(max = 4)) }
        .map { it.sprite(sprite) }
-        .map { it.colorAndIndex(overrideColor) }
+        .map { it.colorIndex(tintIndex) }
        .build(BlendMode.SOLID)
 }

@@ -61,11 +61,20 @@ fun crossModelsRaw(num: Int, size: Double, hOffset: Double, vOffset: Double): Ar
    }
 }

-fun crossModelsTextured(leafBase: Array<List<Quad>>, overrideColor: Int?, scrambleUV: Boolean, spriteGetter: (Int)->Sprite) = leafBase.map { leaf ->
-    leaf.map { if (scrambleUV) it.scrambleUV(random, canFlipU = true, canFlipV = true, canRotate = true) else it }
-        .map { it.colorAndIndex(overrideColor) }
-        .mapIndexed { idx, quad -> quad.sprite(spriteGetter(idx)) }
-        .withOpposites().build(BlendMode.CUTOUT_MIPPED)
+fun crossModelSingle(base: List<Quad>, sprite: Sprite, tintIndex: Int,scrambleUV: Boolean) =
+    base.map { if (scrambleUV) it.scrambleUV(random, canFlipU = true, canFlipV = true, canRotate = true) else it }
+        .map { it.colorIndex(tintIndex) }
+        .mapIndexed { idx, quad -> quad.sprite(sprite) }
+        .withOpposites()
+        .build(BlendMode.CUTOUT_MIPPED)
+
+fun crossModelsTextured(
+    leafBase: Array<List<Quad>>,
+    tintIndex: Int,
+    scrambleUV: Boolean,
+    spriteGetter: (Int) -> Identifier
+) = leafBase.mapIndexed { idx, leaf ->
+    crossModelSingle(leaf, Atlas.BLOCKS[spriteGetter(idx)], tintIndex, scrambleUV)
 }.toTypedArray()

-fun Array<List<Quad>>.buildTufts() = withOpposites().build(BlendMode.CUTOUT_MIPPED)
+    fun Array<List<Quad>>.buildTufts() = withOpposites().build(BlendMode.CUTOUT_MIPPED)
--- a/src/main/kotlin/mods/betterfoliage/model/VanillaWrappers.kt
+++ b/src/main/kotlin/mods/betterfoliage/model/VanillaWrappers.kt
@@ -1,10 +1,8 @@
 package mods.betterfoliage.model

-import mods.betterfoliage.WeightedBakedModelEntry_model
-import mods.betterfoliage.WeightedBakedModel_models
-import mods.betterfoliage.WeightedBakedModel_totalWeight
-import mods.betterfoliage.WeightedPickerEntry_weight
-import mods.betterfoliage.util.get
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelBakingKey
+import mods.betterfoliage.util.HasLogger
 import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
 import net.fabricmc.fabric.api.renderer.v1.mesh.Mesh
 import net.fabricmc.fabric.api.renderer.v1.model.FabricBakedModel
@@ -13,7 +11,6 @@ import net.minecraft.block.BlockState
 import net.minecraft.client.render.RenderLayers
 import net.minecraft.client.render.model.BakedModel
 import net.minecraft.client.render.model.BasicBakedModel
-import net.minecraft.client.render.model.WeightedBakedModel
 import net.minecraft.item.ItemStack
 import net.minecraft.util.WeightedPicker
 import net.minecraft.util.math.BlockPos
@@ -21,6 +18,18 @@ import net.minecraft.world.BlockRenderView
 import java.util.*
 import java.util.function.Supplier

+abstract class ModelWrapKey : ModelBakingKey, HasLogger() {
+    override fun bake(ctx: ModelBakingContext): BakedModel? {
+        val baseModel = super.bake(ctx)
+        if (baseModel is BasicBakedModel)
+            return bake(ctx, baseModel)
+        else
+            return baseModel
+    }
+
+    abstract fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel): BakedModel
+}
+
 abstract class WrappedBakedModel(val wrapped: BakedModel) : BakedModel by wrapped, FabricBakedModel {
    override fun isVanillaAdapter() = false

@@ -46,7 +55,7 @@ class WrappedMeshModel(wrapped: BasicBakedModel, val mesh: Mesh) : WrappedBakedM
         * @param noDiffuse disable diffuse lighting when baking the [Mesh]
         * @param blendModeOverride [BlockRenderLayer] to use instead of the one declared by the corresponding [Block]
         */
-        fun converter(state: BlockState, unshade: Boolean = false, noDiffuse: Boolean = true, blendModeOverride: BlendMode? = null) = BakedModelConverter.of { model, _ ->
+        fun converter(state: BlockState?, unshade: Boolean = false, noDiffuse: Boolean = true, blendModeOverride: BlendMode? = null) = BakedModelConverter.of { model ->
            if (model is BasicBakedModel) {
                val mesh = unbakeQuads(model, state, Random(42L), unshade).build(
                    blendMode = blendModeOverride ?: BlendMode.fromRenderLayer(RenderLayers.getBlockLayer(state)),
@@ -59,23 +68,20 @@ class WrappedMeshModel(wrapped: BasicBakedModel, val mesh: Mesh) : WrappedBakedM
    }
 }

-class WrappedWeightedModel(wrapped: WeightedBakedModel, transformer: BakedModelConverter) : WrappedBakedModel(wrapped) {
-    val totalWeight = wrapped[WeightedBakedModel_totalWeight] as Int
-    val models = wrapped[WeightedBakedModel_models]!!.map { entry ->
-        Entry(transformer.convert(entry[WeightedBakedModelEntry_model]!!, transformer)!!, entry[WeightedPickerEntry_weight]!!)
-    }
+class WeightedModelWrapper(
+    val models: List<WeightedModel>, baseModel: BakedModel
+): WrappedBakedModel(baseModel), FabricBakedModel {
+
+    class WeightedModel(val model: BakedModel, val weight: Int) : WeightedPicker.Entry(weight)
+    fun getModel(random: Random) = WeightedPicker.getRandom(random, models).model

    override fun emitBlockQuads(blockView: BlockRenderView, state: BlockState, pos: BlockPos, randomSupplier: Supplier<Random>, context: RenderContext) {
-        (WeightedPicker.getRandom(randomSupplier.get(), models, totalWeight).model as FabricBakedModel).emitBlockQuads(blockView, state, pos, randomSupplier, context)
-    }
-
-    class Entry(val model: BakedModel, weight: Int) : WeightedPicker.Entry(weight)
-
-    companion object {
-        val converter = object : BakedModelConverter {
-            override fun convert(model: BakedModel, converter: BakedModelConverter) =
-                (model as? WeightedBakedModel)?.let { WrappedWeightedModel(it, converter) }
-        }
+        (getModel(randomSupplier.get()) as FabricBakedModel).emitBlockQuads(blockView, state, pos, randomSupplier, context)
    }
 }

+fun getUnderlyingModel(model: BakedModel, random: Random): BakedModel = when(model) {
+    is WeightedModelWrapper -> getUnderlyingModel(model.getModel(random), random)
+    is WrappedBakedModel -> model.wrapped
+    else -> model
+}
--- a/src/main/kotlin/mods/betterfoliage/render/OptifineCustomColors.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/OptifineCustomColors.kt
@@ -1,48 +0,0 @@
-package mods.betterfoliage.render
-
-import mods.betterfoliage.*
-import mods.betterfoliage.util.ThreadLocalDelegate
-import net.minecraft.block.BlockState
-import net.minecraft.client.MinecraftClient
-import net.minecraft.client.render.model.BakedQuad
-import net.minecraft.util.math.BlockPos
-import net.minecraft.util.math.Direction.UP
-import org.apache.logging.log4j.Level
-
-/**
- * Integration for OptiFine custom block colors.
- */
-/*
-@Suppress("UNCHECKED_CAST")
-object OptifineCustomColors {
-
-    val isColorAvailable = allAvailable(CustomColors, CustomColors.getColorMultiplier)
-
-    init {
-        BetterFoliage.log(Level.INFO, "Optifine custom color support is ${if (isColorAvailable) "enabled" else "disabled" }")
-    }
-
-    val renderEnv by ThreadLocalDelegate { OptifineRenderEnv() }
-    val fakeQuad = BakedQuad(IntArray(0), 1, UP, null)
-
-    fun getBlockColor(ctx: CombinedContext): Int {
-        val ofColor = if (isColorAvailable && MinecraftClient.getInstance().options.reflectDeclaredField<Boolean>("ofCustomColors") == true) {
-            renderEnv.reset(ctx.state, ctx.pos)
-            CustomColors.getColorMultiplier.invokeStatic(fakeQuad, ctx.state, ctx.world, ctx.pos, renderEnv.wrapped) as? Int
-        } else null
-        return if (ofColor == null || ofColor == -1) ctx.lightingCtx.color else ofColor
-    }
-}
-
-class OptifineRenderEnv {
-    val wrapped: Any = RenderEnv.element!!.getDeclaredConstructor(BlockState.element, BlockPos.element).let {
-        it.isAccessible = true
-        it.newInstance(null, null)
-    }
-
-    fun reset(state: BlockState, pos: BlockPos) {
-        RenderEnv.reset.invoke(wrapped, state, pos)
-    }
-}
-
- */
--- a/src/main/kotlin/mods/betterfoliage/render/ShadersModIntegration.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/ShadersModIntegration.kt
@@ -2,6 +2,7 @@ package mods.betterfoliage.render

 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.render.lighting.getBufferBuilder
+import mods.betterfoliage.util.HasLogger
 import mods.betterfoliage.util.getAllMethods
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
@@ -12,7 +13,7 @@ import net.minecraft.client.render.RenderLayer
 /**
 * Integration for ShadersMod.
 */
-object ShadersModIntegration {
+object ShadersModIntegration : HasLogger() {

    val BufferBuilder_SVertexBuilder = BufferBuilder::class.java.fields.find { it.name == "sVertexBuilder" }
    val SVertexBuilder_pushState = getAllMethods("net.optifine.shaders.SVertexBuilder", "pushEntity").find { it.parameterCount == 1 }
@@ -27,7 +28,7 @@ object ShadersModIntegration {
    val defaultGrass = Blocks.TALL_GRASS.defaultState

    init {
-        BetterFoliage.logger.info("[BetterFoliage] ShadersMod integration is ${if (isAvailable) "enabled" else "disabled" }")
+        logger.info("[BetterFoliage] ShadersMod integration is ${if (isAvailable) "enabled" else "disabled" }")
    }

    /** Quads rendered inside this block will use the given block entity data in shader programs. */
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Cactus.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Cactus.kt
@@ -1,26 +1,29 @@
 package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.render.lighting.grassTuftLighting
-import mods.betterfoliage.render.lighting.roundLeafLighting
-import mods.betterfoliage.render.lighting.withLighting
-import mods.betterfoliage.resource.discovery.BlockRenderKey
-import mods.betterfoliage.resource.discovery.ConfigurableModelDiscovery
-import mods.betterfoliage.resource.discovery.ModelTextureList
-import mods.betterfoliage.resource.discovery.SimpleBlockMatcher
+import mods.betterfoliage.config.CACTUS_BLOCKS
 import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
 import mods.betterfoliage.model.SpriteDelegate
 import mods.betterfoliage.model.SpriteSetDelegate
 import mods.betterfoliage.model.WrappedBakedModel
 import mods.betterfoliage.model.buildTufts
 import mods.betterfoliage.model.crossModelsRaw
 import mods.betterfoliage.model.crossModelsTextured
+import mods.betterfoliage.model.meshifyCutoutMipped
 import mods.betterfoliage.model.meshifyStandard
 import mods.betterfoliage.model.transform
 import mods.betterfoliage.model.tuftModelSet
 import mods.betterfoliage.model.tuftShapeSet
+import mods.betterfoliage.render.lighting.grassTuftLighting
+import mods.betterfoliage.render.lighting.roundLeafLighting
+import mods.betterfoliage.render.lighting.withLighting
+import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.util.Atlas
-import mods.betterfoliage.util.LazyMap
+import mods.betterfoliage.util.LazyInvalidatable
 import mods.betterfoliage.util.Rotation
 import mods.betterfoliage.util.get
 import mods.betterfoliage.util.horizontalDirections
@@ -29,37 +32,36 @@ import mods.betterfoliage.util.randomI
 import net.fabricmc.fabric.api.renderer.v1.model.FabricBakedModel
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
-import net.minecraft.block.CactusBlock
+import net.minecraft.block.Blocks
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.JsonUnbakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.DOWN
 import net.minecraft.world.BlockRenderView
 import java.util.Random
-import java.util.function.Consumer
 import java.util.function.Supplier


-interface CactusKey : BlockRenderKey {
-    val cactusTop: Identifier
-    val cactusBottom: Identifier
-    val cactusSide: Identifier
-}
-
-object StandardCactusDiscovery : ConfigurableModelDiscovery() {
-    override val logger = BetterFoliage.logDetail
-    override val matchClasses = SimpleBlockMatcher(CactusBlock::class.java)
-    override val modelTextures = listOf(ModelTextureList("block/cactus", "top", "bottom", "side"))
-
-    override fun processModel(state: BlockState, textures: List<Identifier>, atlas: Consumer<Identifier>): BlockRenderKey? {
-        return CactusModel.Key(textures[0], textures[1], textures[2])
+object StandardCactusDiscovery : AbstractModelDiscovery() {
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        val model = ctx.getUnbaked()
+        if (model is JsonUnbakedModel && ctx.blockState.block in CACTUS_BLOCKS) {
+            BetterFoliage.blockTypes.dirt.add(ctx.blockState)
+            ctx.addReplacement(StandardCactusKey)
+            ctx.sprites.add(StandardCactusModel.cactusCrossSprite)
+        }
+        super.processModel(ctx)
    }
 }

-class CactusModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wrapped), FabricBakedModel {
+object StandardCactusKey : ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = StandardCactusModel(meshifyCutoutMipped(wrapped))
+}
+
+class StandardCactusModel(wrapped: BakedModel) : WrappedBakedModel(wrapped), FabricBakedModel {

-    val crossModels by cactusCrossModels.delegate(key)
-    val armModels by cactusArmModels.delegate(key)
    val armLighting = horizontalDirections.map { grassTuftLighting(it) }
    val crossLighting = roundLeafLighting()

@@ -71,36 +73,26 @@ class CactusModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wrapped
        val armSide = random.nextInt() and 3

        context.withLighting(armLighting[armSide]) {
-            it.accept(armModels[armSide][random])
+            it.accept(cactusArmModels[armSide][random])
        }
        context.withLighting(crossLighting) {
-            it.accept(crossModels[random])
+            it.accept(cactusCrossModels[random])
        }
    }

-    data class Key(
-        override val cactusTop: Identifier,
-        override val cactusBottom: Identifier,
-        override val cactusSide: Identifier
-    ) : CactusKey {
-        override fun replace(model: BakedModel, state: BlockState) = CactusModel(this, meshifyStandard(model, state))
-    }
-
    companion object {
-        val cactusCrossSprite by SpriteDelegate(Atlas.BLOCKS) {
-            Identifier(BetterFoliage.MOD_ID, "blocks/better_cactus")
-        }
+        val cactusCrossSprite = Identifier(BetterFoliage.MOD_ID, "blocks/better_cactus")
        val cactusArmSprites by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_cactus_arm_$idx")
        }
-        val cactusArmModels = LazyMap(BetterFoliage.modelReplacer) { key: CactusKey ->
+        val cactusArmModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = BetterFoliage.config.cactus.let { tuftShapeSet(0.8, 0.8, 0.8, 0.2) }
            val models = tuftModelSet(shapes, Color.white.asInt) { cactusArmSprites[randomI()] }
            horizontalDirections.map { side ->
                models.transform { move(0.0625 to DOWN).rotate(Rotation.fromUp[side.ordinal]) }.buildTufts()
            }.toTypedArray()
        }
-        val cactusCrossModels = LazyMap(BetterFoliage.modelReplacer) { key: CactusKey ->
+        val cactusCrossModels by LazyInvalidatable(BakeWrapperManager) {
            val models = BetterFoliage.config.cactus.let { config ->
                crossModelsRaw(64, config.size, 0.0, 0.0)
                    .transform { rotateZ(randomD(-config.sizeVariation, config.sizeVariation)) }
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Dirt.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Dirt.kt
@@ -2,42 +2,72 @@ package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.chunk.BasicBlockCtx
-import mods.betterfoliage.render.DIRT_BLOCKS
-import mods.betterfoliage.render.SALTWATER_BIOMES
+import mods.betterfoliage.config.DIRT_BLOCKS
+import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
+import mods.betterfoliage.model.SpriteSetDelegate
+import mods.betterfoliage.model.WrappedBakedModel
+import mods.betterfoliage.model.build
+import mods.betterfoliage.model.meshifyStandard
+import mods.betterfoliage.model.tuftModelSet
+import mods.betterfoliage.model.tuftShapeSet
+import mods.betterfoliage.model.withOpposites
+import mods.betterfoliage.config.SALTWATER_BIOMES
+import mods.betterfoliage.model.getUnderlyingModel
+import mods.betterfoliage.model.meshifySolid
 import mods.betterfoliage.render.ShadersModIntegration
-import mods.betterfoliage.render.lighting.withLighting
 import mods.betterfoliage.render.lighting.grassTuftLighting
 import mods.betterfoliage.render.lighting.reedLighting
 import mods.betterfoliage.render.lighting.renderMasquerade
-import mods.betterfoliage.util.Atlas
-import mods.betterfoliage.resource.discovery.BlockRenderKey
-import mods.betterfoliage.resource.discovery.ModelDiscoveryBase
+import mods.betterfoliage.render.lighting.withLighting
+import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ModelBakingContext
 import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.resource.generated.CenteredSprite
-import mods.betterfoliage.model.*
-import mods.betterfoliage.util.*
+import mods.betterfoliage.util.Atlas
+import mods.betterfoliage.util.Int3
+import mods.betterfoliage.util.LazyInvalidatable
+import mods.betterfoliage.util.get
+import mods.betterfoliage.util.randomI
 import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
+import net.minecraft.block.Blocks
 import net.minecraft.block.Material
+import net.minecraft.client.render.RenderLayer
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.JsonUnbakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.UP
 import net.minecraft.world.BlockRenderView
-import java.util.*
-import java.util.function.Consumer
+import java.util.Random
 import java.util.function.Supplier

-object DirtKey : BlockRenderKey {
-    override fun replace(model: BakedModel, state: BlockState) = DirtModel(meshifyStandard(model, state))
+object StandardDirtDiscovery : AbstractModelDiscovery() {
+    fun canRenderInLayer(layer: RenderLayer) = when {
+        !BetterFoliage.config.enabled -> layer == RenderLayer.getSolid()
+        (!BetterFoliage.config.connectedGrass.enabled &&
+                !BetterFoliage.config.algae.enabled &&
+                !BetterFoliage.config.reed.enabled
+                ) -> layer == RenderLayer.getSolid()
+        else -> layer == RenderLayer.getCutoutMipped()
+    }
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        if (ctx.getUnbaked() is JsonUnbakedModel && ctx.blockState.block in DIRT_BLOCKS) {
+            BetterFoliage.blockTypes.dirt.add(ctx.blockState)
+            ctx.addReplacement(StandardDirtKey)
+//            RenderTypeLookup.setRenderLayer(ctx.blockState.block, ::canRenderInLayer)
+        }
+        super.processModel(ctx)
+    }
 }

-object DirtDiscovery : ModelDiscoveryBase() {
-    override val logger = BetterFoliage.logDetail
-
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>) =
-        if (ctx.state.block in DIRT_BLOCKS) DirtKey else null
+object StandardDirtKey : ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = DirtModel(meshifySolid(wrapped))
 }

 class DirtModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
@@ -45,26 +75,32 @@ class DirtModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
    val algaeLighting = grassTuftLighting(UP)
    val reedLighting = reedLighting()

-    override fun emitBlockQuads(blockView: BlockRenderView, state: BlockState, pos: BlockPos, randomSupplier: Supplier<Random>, context: RenderContext) {
+    override fun emitBlockQuads(
+        blockView: BlockRenderView,
+        state: BlockState,
+        pos: BlockPos,
+        randomSupplier: Supplier<Random>,
+        context: RenderContext
+    ) {
        if (!BetterFoliage.config.enabled) return super.emitBlockQuads(blockView, state, pos, randomSupplier, context)

        val ctx = BasicBlockCtx(blockView, pos)
        val stateUp = ctx.offset(UP).state
-        val keyUp = BetterFoliage.modelReplacer[stateUp]
+        val isGrassUp = stateUp in BetterFoliage.blockTypes.grass

        val isWater = stateUp.material == Material.WATER
        val isDeepWater = isWater && ctx.offset(Int3(2 to UP)).state.material == Material.WATER
        val isShallowWater = isWater && ctx.offset(Int3(2 to UP)).state.isAir
        val isSaltWater = isWater && ctx.biome?.category in SALTWATER_BIOMES

-        if (BetterFoliage.config.connectedGrass.enabled && keyUp is GrassKey) {
-            val grassBaseModel = (ctx.model(UP) as WrappedBakedModel).wrapped
+        val random = randomSupplier.get()
+        if (BetterFoliage.config.connectedGrass.enabled && isGrassUp) {
+            val grassBaseModel = getUnderlyingModel(ctx.model(UP), random)
            context.renderMasquerade(grassBaseModel, blockView, stateUp, pos, randomSupplier, context)
        } else {
            super.emitBlockQuads(blockView, state, pos, randomSupplier, context)
        }

-        val random = randomSupplier.get()
        if (BetterFoliage.config.algae.enabled(random) && isDeepWater) {
            ShadersModIntegration.grass(context, BetterFoliage.config.algae.shaderWind) {
                context.withLighting(algaeLighting) {
@@ -88,14 +124,15 @@ class DirtModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
            idFunc = { idx -> Identifier(BetterFoliage.MOD_ID, "blocks/better_reed_$idx") },
            idRegister = { id -> CenteredSprite(id, aspectHeight = 2).register(BetterFoliage.generatedPack) }
        )
-        val algaeModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
-            val shapes = BetterFoliage.config.algae.let { tuftShapeSet(it.size, it.heightMin, it.heightMax, it.hOffset) }
+        val algaeModels by LazyInvalidatable(BakeWrapperManager) {
+            val shapes =
+                BetterFoliage.config.algae.let { tuftShapeSet(it.size, it.heightMin, it.heightMax, it.hOffset) }
            tuftModelSet(shapes, Color.white.asInt) { algaeSprites[randomI()] }
                .withOpposites()
                .build(BlendMode.CUTOUT_MIPPED, flatLighting = false)

        }
-        val reedModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val reedModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = BetterFoliage.config.reed.let { tuftShapeSet(2.0, it.heightMin, it.heightMax, it.hOffset) }
            tuftModelSet(shapes, Color.white.asInt) { reedSprites[randomI()] }
                .withOpposites()
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Grass.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Grass.kt
@@ -2,7 +2,8 @@ package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.chunk.BasicBlockCtx
-import mods.betterfoliage.render.SNOW_MATERIALS
+import mods.betterfoliage.config.BlockConfig
+import mods.betterfoliage.config.SNOW_MATERIALS
 import mods.betterfoliage.render.ShadersModIntegration
 import mods.betterfoliage.render.lighting.grassTuftLighting
 import mods.betterfoliage.render.lighting.withLighting
@@ -14,6 +15,7 @@ import net.fabricmc.fabric.api.renderer.v1.model.FabricBakedModel
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.*
@@ -22,32 +24,32 @@ import java.util.*
 import java.util.function.Consumer
 import java.util.function.Supplier

-interface GrassKey : BlockRenderKey {
-    val grassTopTexture: Identifier
-
-    /**
-     * Color to use for Short Grass rendering instead of the biome color.
-     *
-     * Value is null if the texture is mostly grey (the saturation of its average color is under a configurable limit),
-     * the average color of the texture otherwise.
-     */
-    val overrideColor: Int?
-}
-
 object StandardGrassDiscovery : ConfigurableModelDiscovery() {
-    override val logger = BetterFoliage.logDetail
    override val matchClasses: ConfigurableBlockMatcher get() = BetterFoliage.blockConfig.grassBlocks
    override val modelTextures: List<ModelTextureList> get() = BetterFoliage.blockConfig.grassModels.modelList

-    override fun processModel(state: BlockState, textures: List<Identifier>, atlas: Consumer<Identifier>): BlockRenderKey? {
-        val grassId = textures[0]
-        log(" block state $state")
-        log("       texture $grassId")
-        return GrassBlockModel.Key(grassId, getAndLogColorOverride(grassId, Atlas.BLOCKS, BetterFoliage.config.shortGrass.saturationThreshold))
+    override fun processModel(ctx: ModelDiscoveryContext, textureMatch: List<Identifier>) {
+        ctx.addReplacement(StandardGrassKey(textureMatch[0], null))
+        BetterFoliage.blockTypes.grass.add(ctx.blockState)
    }
 }

-class GrassBlockModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wrapped), FabricBakedModel {
+data class StandardGrassKey(
+    val grassLocation: Identifier,
+    val overrideColor: Color?
+) : ModelWrapKey() {
+    val tintIndex: Int get() = if (overrideColor == null) 0 else -1
+
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel): BakedModel {
+        val grassSpriteColor = Atlas.BLOCKS[grassLocation].averageColor.let { hsb ->
+            logColorOverride(detailLogger, BetterFoliage.config.shortGrass.saturationThreshold, hsb)
+            hsb.colorOverride(BetterFoliage.config.shortGrass.saturationThreshold)
+        }
+        return StandardGrassModel(meshifyCutoutMipped(wrapped), this.copy(overrideColor = grassSpriteColor))
+    }
+}
+
+class StandardGrassModel(wrapped: BakedModel, val key: StandardGrassKey) : WrappedBakedModel(wrapped) {

    val tuftNormal by grassTuftMeshesNormal.delegate(key)
    val tuftSnowed by grassTuftMeshesSnowed.delegate(key)
@@ -64,9 +66,8 @@ class GrassBlockModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wra

        val isSnowed = stateAbove.material in SNOW_MATERIALS
        val connected = BetterFoliage.config.connectedGrass.enabled &&
-            (!isSnowed || BetterFoliage.config.connectedGrass.snowEnabled) && (
-                BetterFoliage.modelReplacer[stateBelow].let { it is DirtKey || it is GrassKey }
-            )
+            (!isSnowed || BetterFoliage.config.connectedGrass.snowEnabled) &&
+            (stateBelow in BetterFoliage.blockTypes.dirt || stateBelow in BetterFoliage.blockTypes.grass)

        val random = randomSupplier.get()
        if (connected) {
@@ -84,13 +85,6 @@ class GrassBlockModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wra
        }
    }

-    data class Key(
-        override val grassTopTexture: Identifier,
-        override val overrideColor: Int?
-    ) : GrassKey {
-        override fun replace(model: BakedModel, state: BlockState) = GrassBlockModel(this, meshifyStandard(model, state))
-    }
-
    companion object {
        val grassTuftSpritesNormal by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_grass_long_$idx")
@@ -98,23 +92,23 @@ class GrassBlockModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wra
        val grassTuftSpritesSnowed by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_grass_long_$idx")
        }
-        val grassTuftShapes = LazyMap(BetterFoliage.modelReplacer) { key: GrassKey ->
+        val grassTuftShapes = LazyMap(BakeWrapperManager) { key: StandardGrassKey ->
            BetterFoliage.config.shortGrass.let { tuftShapeSet(it.size, it.heightMin, it.heightMax, it.hOffset) }
        }
-        val grassTuftMeshesNormal = LazyMap(BetterFoliage.modelReplacer) { key: GrassKey ->
-            tuftModelSet(grassTuftShapes[key], key.overrideColor) { idx -> grassTuftSpritesNormal[randomI()] }
+        val grassTuftMeshesNormal = LazyMap(BakeWrapperManager) { key: StandardGrassKey ->
+            tuftModelSet(grassTuftShapes[key], key.tintIndex) { idx -> grassTuftSpritesNormal[randomI()] }
                .withOpposites()
                .build(BlendMode.CUTOUT_MIPPED, flatLighting = false)
        }
-        val grassTuftMeshesSnowed = LazyMap(BetterFoliage.modelReplacer) { key: GrassKey ->
+        val grassTuftMeshesSnowed = LazyMap(BakeWrapperManager) { key: StandardGrassKey ->
            tuftModelSet(grassTuftShapes[key], Color.white.asInt) { idx -> grassTuftSpritesSnowed[randomI()] }
                .withOpposites()
                .build(BlendMode.CUTOUT_MIPPED, flatLighting = false)
        }
-        val grassFullBlockMeshes = LazyMap(BetterFoliage.modelReplacer) { key: GrassKey ->
-            Array(64) { fullCubeTextured(key.grassTopTexture, key.overrideColor) }
+        val grassFullBlockMeshes = LazyMap(BakeWrapperManager) { key: StandardGrassKey ->
+            Array(64) { fullCubeTextured(key.grassLocation, key.tintIndex) }
        }
-        val snowFullBlockMeshes by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val snowFullBlockMeshes by LazyInvalidatable(BakeWrapperManager) {
            Array(64) { fullCubeTextured(Identifier("block/snow"), Color.white.asInt) }
        }
    }
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Leaf.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Leaf.kt
@@ -2,76 +2,87 @@ package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.chunk.BasicBlockCtx
-import mods.betterfoliage.render.SNOW_MATERIALS
+import mods.betterfoliage.config.SNOW_MATERIALS
+import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
+import mods.betterfoliage.model.SpriteSetDelegate
+import mods.betterfoliage.model.WrappedBakedModel
+import mods.betterfoliage.model.crossModelsRaw
+import mods.betterfoliage.model.crossModelsTextured
+import mods.betterfoliage.model.meshifyCutoutMipped
 import mods.betterfoliage.render.ShadersModIntegration
-import mods.betterfoliage.render.lighting.withLighting
 import mods.betterfoliage.render.lighting.roundLeafLighting
+import mods.betterfoliage.render.lighting.withLighting
 import mods.betterfoliage.render.particle.LeafParticleRegistry
-import mods.betterfoliage.util.Atlas
-import mods.betterfoliage.resource.discovery.*
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ConfigurableBlockMatcher
+import mods.betterfoliage.resource.discovery.ConfigurableModelDiscovery
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
+import mods.betterfoliage.resource.discovery.ModelTextureList
 import mods.betterfoliage.resource.generated.GeneratedLeafSprite
-import mods.betterfoliage.model.*
-import mods.betterfoliage.util.*
-import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
-import net.fabricmc.fabric.api.renderer.v1.model.FabricBakedModel
+import mods.betterfoliage.util.Atlas
+import mods.betterfoliage.util.LazyMap
+import mods.betterfoliage.util.averageColor
+import mods.betterfoliage.util.colorOverride
+import mods.betterfoliage.util.get
+import mods.betterfoliage.util.logColorOverride
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
-import net.minecraft.util.math.Direction.*
+import net.minecraft.util.math.Direction.UP
 import net.minecraft.world.BlockRenderView
-import java.util.*
-import java.util.function.Consumer
+import org.apache.logging.log4j.Level
+import java.util.Random
 import java.util.function.Supplier

-interface LeafKey : BlockRenderKey {
-    val roundLeafTexture: Identifier
+interface LeafBlockModel {
+    val key: LeafParticleKey
+}

-    /** Type of the leaf block (configurable by user). */
+interface LeafParticleKey {
    val leafType: String
-
-    /** Average color of the round leaf texture. */
-    val overrideColor: Int?
+    val overrideColor: Color?
 }

 object StandardLeafDiscovery : ConfigurableModelDiscovery() {
-    override val logger = BetterFoliage.logDetail
    override val matchClasses: ConfigurableBlockMatcher get() = BetterFoliage.blockConfig.leafBlocks
    override val modelTextures: List<ModelTextureList> get() = BetterFoliage.blockConfig.leafModels.modelList

-    override fun processModel(state: BlockState, textures: List<Identifier>, atlas: Consumer<Identifier>) =
-        defaultRegisterLeaf(textures[0], atlas)
+    override fun processModel(ctx: ModelDiscoveryContext, textureMatch: List<Identifier>) {
+        val leafType = LeafParticleRegistry.typeMappings.getType(textureMatch[0]) ?: "default"
+        val generated = GeneratedLeafSprite(textureMatch[0], leafType)
+            .register(BetterFoliage.generatedPack)
+            .apply { ctx.sprites.add(this) }

-}
-
-fun HasLogger.defaultRegisterLeaf(sprite: Identifier, atlas: Consumer<Identifier>): BlockRenderKey? {
-    val leafType = LeafParticleRegistry.typeMappings.getType(sprite) ?: "default"
-    val leafId = GeneratedLeafSprite(sprite, leafType).register(BetterFoliage.generatedPack)
-    atlas.accept(leafId)
-
-    log(" leaf texture $sprite")
-    log("     particle $leafType")
-
-    return NormalLeavesModel.Key(
-        leafId, leafType,
-        getAndLogColorOverride(leafId, Atlas.BLOCKS, BetterFoliage.config.shortGrass.saturationThreshold)
-    )
-}
-
-fun HasLogger.getAndLogColorOverride(sprite: Identifier, atlas: Atlas, threshold: Double): Int? {
-    val hsb = resourceManager.averageImageColorHSB(sprite, atlas)
-    return if (hsb.saturation >= threshold) {
-        log("         brightness ${hsb.brightness}")
-        log("         saturation ${hsb.saturation} >= ${threshold}, using texture color")
-        hsb.copy(brightness = 0.9f.coerceAtMost(hsb.brightness * 2.0f)).asColor
-    } else {
-        log("         saturation ${hsb.saturation} < ${threshold}, using block color")
-        null
+        detailLogger.log(Level.INFO, "     particle $leafType")
+        ctx.addReplacement(StandardLeafKey(generated, leafType, null))
    }
 }

-class NormalLeavesModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(wrapped), FabricBakedModel {
+data class StandardLeafKey(
+    val roundLeafTexture: Identifier,
+    override val leafType: String,
+    override val overrideColor: Color?
+) : ModelWrapKey(), LeafParticleKey {
+    val tintIndex: Int get() = if (overrideColor == null) 0 else -1
+
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel): BakedModel {
+        val leafSpriteColor = Atlas.BLOCKS[roundLeafTexture].averageColor.let { hsb ->
+            logColorOverride(detailLogger, BetterFoliage.config.leaves.saturationThreshold, hsb)
+            hsb.colorOverride(BetterFoliage.config.leaves.saturationThreshold)
+        }
+        return StandardLeafModel(meshifyCutoutMipped(wrapped), this.copy(overrideColor = leafSpriteColor))
+    }
+}
+
+class StandardLeafModel(
+    wrapped: BakedModel,
+    override val key: StandardLeafKey
+) : WrappedBakedModel(wrapped), LeafBlockModel {

    val leafNormal by leafModelsNormal.delegate(key)
    val leafSnowed by leafModelsSnowed.delegate(key)
@@ -94,26 +105,18 @@ class NormalLeavesModel(val key: Key, wrapped: BakedModel) : WrappedBakedModel(w
        }
    }

-    data class Key(
-        override val roundLeafTexture: Identifier,
-        override val leafType: String,
-        override val overrideColor: Int?
-    ) : LeafKey {
-        override fun replace(model: BakedModel, state: BlockState) = NormalLeavesModel(this, meshifyStandard(model, state, blendModeOverride = BlendMode.CUTOUT_MIPPED))
-    }
-
    companion object {
        val leafSpritesSnowed by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_leaves_snowed_$idx")
        }
-        val leafModelsBase = LazyMap(BetterFoliage.modelReplacer) { key: LeafKey ->
+        val leafModelsBase = LazyMap(BakeWrapperManager) { key: StandardLeafKey ->
            BetterFoliage.config.leaves.let { crossModelsRaw(64, it.size, it.hOffset, it.vOffset) }
        }
-        val leafModelsNormal = LazyMap(BetterFoliage.modelReplacer) { key: LeafKey ->
-            crossModelsTextured(leafModelsBase[key], key.overrideColor, true) { Atlas.BLOCKS.atlas[key.roundLeafTexture]!! }
+        val leafModelsNormal = LazyMap(BakeWrapperManager) { key: StandardLeafKey ->
+            crossModelsTextured(leafModelsBase[key], key.tintIndex, true) { key.roundLeafTexture }
        }
-        val leafModelsSnowed = LazyMap(BetterFoliage.modelReplacer) { key: LeafKey ->
-            crossModelsTextured(leafModelsBase[key], Color.white.asInt, false) { leafSpritesSnowed[it] }
+        val leafModelsSnowed = LazyMap(BakeWrapperManager) { key: StandardLeafKey ->
+            crossModelsTextured(leafModelsBase[key], Color.white.asInt, false) { leafSpritesSnowed[it].id }
        }
    }
 }
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Lilypad.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Lilypad.kt
@@ -1,18 +1,21 @@
 package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.render.ShadersModIntegration
-import mods.betterfoliage.resource.discovery.BlockRenderKey
-import mods.betterfoliage.resource.discovery.ModelDiscoveryBase
-import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
 import mods.betterfoliage.model.SpriteSetDelegate
 import mods.betterfoliage.model.WrappedBakedModel
 import mods.betterfoliage.model.buildTufts
+import mods.betterfoliage.model.meshifyCutoutMipped
 import mods.betterfoliage.model.meshifyStandard
 import mods.betterfoliage.model.transform
 import mods.betterfoliage.model.tuftModelSet
 import mods.betterfoliage.model.tuftShapeSet
+import mods.betterfoliage.render.ShadersModIntegration
+import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.util.Atlas
 import mods.betterfoliage.util.LazyInvalidatable
 import mods.betterfoliage.util.get
@@ -20,25 +23,31 @@ import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
 import net.minecraft.block.Blocks
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.JsonUnbakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.DOWN
 import net.minecraft.world.BlockRenderView
 import java.util.Random
-import java.util.function.Consumer
 import java.util.function.Supplier

-object LilypadKey : BlockRenderKey {
-    override fun replace(model: BakedModel, state: BlockState) = LilypadModel(meshifyStandard(model, state))
+object StandardLilypadDiscovery : AbstractModelDiscovery() {
+    val LILYPAD_BLOCKS = listOf(Blocks.LILY_PAD)
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        if (ctx.getUnbaked() is JsonUnbakedModel && ctx.blockState.block in LILYPAD_BLOCKS) {
+            ctx.addReplacement(StandardLilypadKey)
+        }
+        super.processModel(ctx)
+    }
 }

-object LilyPadDiscovery : ModelDiscoveryBase() {
-    override val logger = BetterFoliage.logDetail
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>) =
-        if (ctx.state.block == Blocks.LILY_PAD) LilypadKey else null
+object StandardLilypadKey : ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = StandardLilypadModel(meshifyCutoutMipped(wrapped))
 }

-class LilypadModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
+class StandardLilypadModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
    override fun emitBlockQuads(blockView: BlockRenderView, state: BlockState, pos: BlockPos, randomSupplier: Supplier<Random>, context: RenderContext) {
        super.emitBlockQuads(blockView, state, pos, randomSupplier, context)
        if (!BetterFoliage.config.enabled || !BetterFoliage.config.lilypad.enabled) return
@@ -59,13 +68,13 @@ class LilypadModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
        val lilypadFlowerSprites by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_lilypad_flower_$idx")
        }
-        val lilypadRootModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val lilypadRootModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = tuftShapeSet(1.0, 1.0, 1.0, BetterFoliage.config.lilypad.hOffset)
            tuftModelSet(shapes, Color.white.asInt) { lilypadRootSprites[it] }
                .transform { move(2.0 to DOWN) }
                .buildTufts()
        }
-        val lilypadFlowerModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val lilypadFlowerModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = tuftShapeSet(0.5, 0.5, 0.5, BetterFoliage.config.lilypad.hOffset)
            tuftModelSet(shapes, Color.white.asInt) { lilypadFlowerSprites[it] }
                .transform { move(1.0 to DOWN) }
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Mycelium.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Mycelium.kt
@@ -1,19 +1,22 @@
 package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.render.ShadersModIntegration
-import mods.betterfoliage.render.lighting.grassTuftLighting
-import mods.betterfoliage.render.lighting.withLighting
-import mods.betterfoliage.resource.discovery.BlockRenderKey
-import mods.betterfoliage.resource.discovery.ModelDiscoveryBase
-import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
 import mods.betterfoliage.model.SpriteSetDelegate
 import mods.betterfoliage.model.WrappedBakedModel
 import mods.betterfoliage.model.buildTufts
+import mods.betterfoliage.model.meshifyCutoutMipped
 import mods.betterfoliage.model.meshifyStandard
 import mods.betterfoliage.model.tuftModelSet
 import mods.betterfoliage.model.tuftShapeSet
+import mods.betterfoliage.render.ShadersModIntegration
+import mods.betterfoliage.render.lighting.grassTuftLighting
+import mods.betterfoliage.render.lighting.withLighting
+import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.util.Atlas
 import mods.betterfoliage.util.LazyInvalidatable
 import mods.betterfoliage.util.get
@@ -24,26 +27,32 @@ import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
 import net.minecraft.block.Blocks
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.JsonUnbakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.UP
 import net.minecraft.world.BlockRenderView
 import java.util.Random
-import java.util.function.Consumer
 import java.util.function.Supplier

-object MyceliumKey : BlockRenderKey {
-    override fun replace(model: BakedModel, state: BlockState) = MyceliumModel(meshifyStandard(model, state))
+object StandardMyceliumDiscovery : AbstractModelDiscovery() {
+    val MYCELIUM_BLOCKS = listOf(Blocks.MYCELIUM)
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        if (ctx.getUnbaked() is JsonUnbakedModel && ctx.blockState.block in MYCELIUM_BLOCKS) {
+            ctx.addReplacement(StandardMyceliumKey)
+//            RenderTypeLookup.setRenderLayer(ctx.blockState.block, RenderType.getCutout())
+        }
+        super.processModel(ctx)
+    }
 }

-object MyceliumDiscovery : ModelDiscoveryBase() {
-    override val logger = BetterFoliage.logDetail
-    val myceliumBlocks = listOf(Blocks.MYCELIUM)
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>) =
-        if (ctx.state.block in myceliumBlocks) MyceliumKey else null
+object StandardMyceliumKey : ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = StandardMyceliumModel(meshifyCutoutMipped(wrapped))
 }

-class MyceliumModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
+class StandardMyceliumModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {

    val tuftLighting = grassTuftLighting(UP)

@@ -67,7 +76,7 @@ class MyceliumModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
        val myceliumTuftSprites by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_mycel_$idx")
        }
-        val myceliumTuftModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val myceliumTuftModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = BetterFoliage.config.shortGrass.let { tuftShapeSet(it.size, it.heightMin, it.heightMax, it.hOffset) }
            tuftModelSet(shapes, Color.white.asInt) { idx -> myceliumTuftSprites[randomI()] }.buildTufts()
        }
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Netherrack.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Netherrack.kt
@@ -1,38 +1,68 @@
 package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
+import mods.betterfoliage.config.NETHERRACK_BLOCKS
+import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
+import mods.betterfoliage.model.SpriteSetDelegate
+import mods.betterfoliage.model.WrappedBakedModel
+import mods.betterfoliage.model.build
+import mods.betterfoliage.model.meshifyCutoutMipped
+import mods.betterfoliage.model.meshifyStandard
+import mods.betterfoliage.model.transform
+import mods.betterfoliage.model.tuftModelSet
+import mods.betterfoliage.model.tuftShapeSet
+import mods.betterfoliage.model.withOpposites
 import mods.betterfoliage.render.lighting.grassTuftLighting
 import mods.betterfoliage.render.lighting.withLighting
-import mods.betterfoliage.resource.discovery.BlockRenderKey
-import mods.betterfoliage.resource.discovery.ModelDiscoveryBase
+import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ModelBakingContext
 import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
-import mods.betterfoliage.model.*
-import mods.betterfoliage.util.*
+import mods.betterfoliage.util.Atlas
+import mods.betterfoliage.util.LazyInvalidatable
+import mods.betterfoliage.util.Rotation
+import mods.betterfoliage.util.get
+import mods.betterfoliage.util.offset
+import mods.betterfoliage.util.plus
+import mods.betterfoliage.util.randomI
 import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
-import net.minecraft.block.Blocks
+import net.minecraft.client.render.RenderLayer
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.JsonUnbakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.DOWN
 import net.minecraft.world.BlockRenderView
-import java.util.*
-import java.util.function.Consumer
+import java.util.Random
 import java.util.function.Supplier

-object NetherrackKey : BlockRenderKey {
-    override fun replace(model: BakedModel, state: BlockState) = NetherrackModel(meshifyStandard(model, state))
+object StandardNetherrackDiscovery : AbstractModelDiscovery() {
+
+    fun canRenderInLayer(layer: RenderLayer) = when {
+        !BetterFoliage.config.enabled -> layer == RenderLayer.getSolid()
+        !BetterFoliage.config.netherrack.enabled -> layer == RenderLayer.getSolid()
+        else -> layer == RenderLayer.getCutoutMipped()
+    }
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        if (ctx.getUnbaked() is JsonUnbakedModel && ctx.blockState.block in NETHERRACK_BLOCKS) {
+            BetterFoliage.blockTypes.dirt.add(ctx.blockState)
+            ctx.addReplacement(StandardNetherrackKey)
+//            RenderTypeLookup.setRenderLayer(ctx.blockState.block, ::canRenderInLayer)
+        }
+        super.processModel(ctx)
+    }
 }

-object NetherrackDiscovery : ModelDiscoveryBase() {
-    override val logger = BetterFoliage.logDetail
-    val netherrackBlocks = listOf(Blocks.NETHERRACK)
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>) =
-        if (ctx.state.block in netherrackBlocks) NetherrackKey else null
+object StandardNetherrackKey : ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = StandardNetherrackModel(meshifyCutoutMipped(wrapped))
 }

-class NetherrackModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
+class StandardNetherrackModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {

    val tuftLighting = grassTuftLighting(DOWN)

@@ -53,7 +83,7 @@ class NetherrackModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
        val netherrackTuftSprites by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_netherrack_$idx")
        }
-        val netherrackTuftModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val netherrackTuftModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = BetterFoliage.config.netherrack.let { tuftShapeSet(it.size, it.heightMin, it.heightMax, it.hOffset) }
            tuftModelSet(shapes, Color.white.asInt) { netherrackTuftSprites[randomI()] }
                .transform { rotate(Rotation.fromUp[DOWN.ordinal]).rotateUV(2) }
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/RoundLog.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/RoundLog.kt
@@ -1,36 +1,51 @@
 package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.render.column.*
-import mods.betterfoliage.util.Atlas
-import mods.betterfoliage.resource.discovery.*
+import mods.betterfoliage.model.ModelWrapKey
+import mods.betterfoliage.model.meshifySolid
 import mods.betterfoliage.model.meshifyStandard
+import mods.betterfoliage.render.column.ColumnBlockKey
+import mods.betterfoliage.render.column.ColumnMeshSet
+import mods.betterfoliage.render.column.ColumnModelBase
+import mods.betterfoliage.render.column.ColumnRenderLayer
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ConfigurableBlockMatcher
+import mods.betterfoliage.resource.discovery.ConfigurableModelDiscovery
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelBakingKey
+import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
+import mods.betterfoliage.resource.discovery.ModelTextureList
+import mods.betterfoliage.util.Atlas
 import mods.betterfoliage.util.LazyMap
-import mods.betterfoliage.util.get
 import mods.betterfoliage.util.tryDefault
 import net.minecraft.block.BlockState
 import net.minecraft.block.LogBlock
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.Direction.Axis
-import java.util.function.Consumer
+import org.apache.logging.log4j.Level
+
+interface RoundLogKey : ColumnBlockKey, ModelBakingKey {
+    val barkSprite: Identifier
+    val endSprite: Identifier
+}

 object RoundLogOverlayLayer : ColumnRenderLayer() {
-    override fun getColumnKey(state: BlockState) = BetterFoliage.modelReplacer.getTyped<ColumnBlockKey>(state)
+    override fun getColumnKey(state: BlockState) = BetterFoliage.blockTypes.getTyped<ColumnBlockKey>(state)
    override val connectSolids: Boolean get() = BetterFoliage.config.roundLogs.connectSolids
    override val lenientConnect: Boolean get() = BetterFoliage.config.roundLogs.lenientConnect
    override val defaultToY: Boolean get() = BetterFoliage.config.roundLogs.defaultY
 }

-object StandardLogDiscovery : ConfigurableModelDiscovery() {
-    override val logger = BetterFoliage.logDetail
+object StandardRoundLogDiscovery : ConfigurableModelDiscovery() {
    override val matchClasses: ConfigurableBlockMatcher get() = BetterFoliage.blockConfig.logBlocks
    override val modelTextures: List<ModelTextureList> get() = BetterFoliage.blockConfig.logModels.modelList

-    override fun processModel(state: BlockState, textures: List<Identifier>, atlas: Consumer<Identifier>): BlockRenderKey? {
-        val axis = getAxis(state)
-        log("       axis $axis")
-        return RoundLogModel.Key(axis, textures[0], textures[1])
+    override fun processModel(ctx: ModelDiscoveryContext, textureMatch: List<Identifier>) {
+        val axis = getAxis(ctx.blockState)
+        detailLogger.log(Level.INFO, "       axis $axis")
+        ctx.addReplacement(StandardRoundLogKey(axis, textureMatch[0], textureMatch[1]))
    }

    fun getAxis(state: BlockState): Axis? {
@@ -45,12 +60,15 @@ object StandardLogDiscovery : ConfigurableModelDiscovery() {
    }
 }

-interface RoundLogKey : ColumnBlockKey, BlockRenderKey {
-    val barkSprite: Identifier
-    val endSprite: Identifier
+data class StandardRoundLogKey(
+    override val axis: Axis?,
+    override val barkSprite: Identifier,
+    override val endSprite: Identifier
+) : RoundLogKey, ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = StandardRoundLogModel(meshifySolid(wrapped), this)
 }

-class RoundLogModel(val key: Key, wrapped: BakedModel) : ColumnModelBase(wrapped) {
+class StandardRoundLogModel(wrapped: BakedModel, val key: StandardRoundLogKey) : ColumnModelBase(wrapped) {
    override val enabled: Boolean get() = BetterFoliage.config.enabled && BetterFoliage.config.roundLogs.enabled
    override val overlayLayer: ColumnRenderLayer get() = RoundLogOverlayLayer
    override val connectPerpendicular: Boolean get() = BetterFoliage.config.roundLogs.connectPerpendicular
@@ -58,18 +76,10 @@ class RoundLogModel(val key: Key, wrapped: BakedModel) : ColumnModelBase(wrapped
    val modelSet by modelSets.delegate(key)
    override fun getMeshSet(axis: Axis, quadrant: Int) = modelSet

-    data class Key(
-        override val axis: Axis?,
-        override val barkSprite: Identifier,
-        override val endSprite: Identifier
-    ) : RoundLogKey {
-        override fun replace(model: BakedModel, state: BlockState) = RoundLogModel(this, meshifyStandard(model, state))
-    }
-
    companion object {
-        val modelSets = LazyMap(BetterFoliage.modelReplacer) { key: Key ->
-            val barkSprite = Atlas.BLOCKS.atlas[key.barkSprite]!!
-            val endSprite = Atlas.BLOCKS.atlas[key.endSprite]!!
+        val modelSets = LazyMap(BakeWrapperManager) { key: StandardRoundLogKey ->
+            val barkSprite = Atlas.BLOCKS[key.barkSprite]!!
+            val endSprite = Atlas.BLOCKS[key.endSprite]!!
            BetterFoliage.config.roundLogs.let { config ->
                ColumnMeshSet(
                    config.radiusSmall, config.radiusLarge, config.zProtection,
--- a/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Sand.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/block/vanilla/Sand.kt
@@ -2,23 +2,26 @@ package mods.betterfoliage.render.block.vanilla

 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.chunk.CachedBlockCtx
-import mods.betterfoliage.render.SALTWATER_BIOMES
-import mods.betterfoliage.render.SAND_BLOCKS
-import mods.betterfoliage.render.lighting.grassTuftLighting
-import mods.betterfoliage.render.lighting.withLighting
-import mods.betterfoliage.resource.discovery.BlockRenderKey
-import mods.betterfoliage.resource.discovery.ModelDiscoveryBase
-import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
+import mods.betterfoliage.config.SALTWATER_BIOMES
+import mods.betterfoliage.config.SAND_BLOCKS
 import mods.betterfoliage.model.Color
+import mods.betterfoliage.model.ModelWrapKey
 import mods.betterfoliage.model.SpriteSetDelegate
 import mods.betterfoliage.model.WrappedBakedModel
 import mods.betterfoliage.model.build
 import mods.betterfoliage.model.horizontalRectangle
+import mods.betterfoliage.model.meshifySolid
 import mods.betterfoliage.model.meshifyStandard
 import mods.betterfoliage.model.transform
 import mods.betterfoliage.model.tuftModelSet
 import mods.betterfoliage.model.tuftShapeSet
 import mods.betterfoliage.model.withOpposites
+import mods.betterfoliage.render.lighting.grassTuftLighting
+import mods.betterfoliage.render.lighting.withLighting
+import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
+import mods.betterfoliage.resource.discovery.BakeWrapperManager
+import mods.betterfoliage.resource.discovery.ModelBakingContext
+import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
 import mods.betterfoliage.util.Atlas
 import mods.betterfoliage.util.LazyInvalidatable
 import mods.betterfoliage.util.Rotation
@@ -30,28 +33,36 @@ import mods.betterfoliage.util.randomI
 import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
 import net.fabricmc.fabric.api.renderer.v1.render.RenderContext
 import net.minecraft.block.BlockState
+import net.minecraft.block.Blocks
 import net.minecraft.block.Material
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.JsonUnbakedModel
 import net.minecraft.util.Identifier
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.Direction.UP
 import net.minecraft.world.BlockRenderView
 import java.util.Random
-import java.util.function.Consumer
 import java.util.function.Supplier

-object SandKey : BlockRenderKey {
-    override fun replace(model: BakedModel, state: BlockState) = SandModel(meshifyStandard(model, state))
+object StandardSandDiscovery : AbstractModelDiscovery() {
+
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        if (ctx.getUnbaked() is JsonUnbakedModel && ctx.blockState.block in SAND_BLOCKS) {
+            BetterFoliage.blockTypes.dirt.add(ctx.blockState)
+            ctx.addReplacement(StandardSandKey)
+//            RenderTypeLookup.setRenderLayer(ctx.blockState.block, RenderType.getCutoutMipped())
+        }
+        super.processModel(ctx)
+    }
 }

-object SandDiscovery : ModelDiscoveryBase() {
-    override val logger = BetterFoliage.logDetail
-
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>) =
-        if (ctx.state.block in SAND_BLOCKS) SandKey else null
+object StandardSandKey : ModelWrapKey() {
+    override fun bake(ctx: ModelBakingContext, wrapped: BasicBakedModel) = StandardSandModel(meshifySolid(wrapped))
 }

-class SandModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
+class StandardSandModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {

    val coralLighting = allDirections.map { grassTuftLighting(it) }.toTypedArray()

@@ -85,7 +96,7 @@ class SandModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
        val coralCrustSprites by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            Identifier(BetterFoliage.MOD_ID, "blocks/better_crust_$idx")
        }
-        val coralTuftModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val coralTuftModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = BetterFoliage.config.coral.let { tuftShapeSet(it.size, 1.0, 1.0, it.hOffset) }
            allDirections.map { face ->
                tuftModelSet(shapes, Color.white.asInt) { coralTuftSprites[randomI()] }
@@ -94,7 +105,7 @@ class SandModel(wrapped: BakedModel) : WrappedBakedModel(wrapped) {
                    .build(BlendMode.CUTOUT_MIPPED)
            }.toTypedArray()
        }
-        val coralCrustModels by LazyInvalidatable(BetterFoliage.modelReplacer) {
+        val coralCrustModels by LazyInvalidatable(BakeWrapperManager) {
            allDirections.map { face ->
                Array(64) { idx ->
                    listOf(horizontalRectangle(x1 = -0.5, x2 = 0.5, z1 = -0.5, z2 = 0.5, y = 0.0)
--- a/src/main/kotlin/mods/betterfoliage/render/column/ColumnModel.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/column/ColumnModel.kt
@@ -22,6 +22,7 @@ abstract class ColumnModelBase(wrapped: BakedModel) : WrappedBakedModel(wrapped)
    abstract fun getMeshSet(axis: Axis, quadrant: Int): ColumnMeshSet

    override fun emitBlockQuads(blockView: BlockRenderView, state: BlockState, pos: BlockPos, randomSupplier: Supplier<Random>, context: RenderContext) {
+        if (!enabled) return super.emitBlockQuads(blockView, state, pos, randomSupplier, context)
        val ctx = CachedBlockCtx(blockView, pos)
        val roundLog = ChunkOverlayManager.get(overlayLayer, ctx)

--- a/src/main/kotlin/mods/betterfoliage/render/column/ColumnOverlayLayer.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/column/ColumnOverlayLayer.kt
@@ -19,6 +19,7 @@ import mods.betterfoliage.util.Int3
 import mods.betterfoliage.util.Rotation
 import mods.betterfoliage.util.allDirections
 import mods.betterfoliage.util.face
+import mods.betterfoliage.util.get
 import mods.betterfoliage.util.plus
 import net.minecraft.block.BlockState
 import net.minecraft.util.math.BlockPos
@@ -89,8 +90,9 @@ abstract class ColumnRenderLayer : ChunkOverlayLayer<ColumnLayerData> {
    }

    override fun calculate(ctx: BlockCtx): ColumnLayerData {
-        if (allDirections.all { dir -> ctx.offset(dir).let { it.isNormalCube && BetterFoliage.modelReplacer[it.state] !is RoundLogKey } }) return ColumnLayerData.SkipRender
-//        val columnTextures = registry[ctx] ?: return ColumnLayerData.ResolveError
+        if (allDirections.all { dir ->
+                ctx.offset(dir).let { it.isNormalCube && !BetterFoliage.blockTypes.hasTyped<RoundLogKey>(it.state) }
+        }) return ColumnLayerData.SkipRender
        val columnTextures = getColumnKey(ctx.state) ?: return ColumnLayerData.ResolveError

        // if log axis is not defined and "Default to vertical" config option is not set, render normally
--- a/src/main/kotlin/mods/betterfoliage/render/particle/FallingLeaves.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/particle/FallingLeaves.kt
@@ -2,8 +2,14 @@ package mods.betterfoliage.render.particle

 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.ClientWorldLoadCallback
-import mods.betterfoliage.render.block.vanilla.LeafKey
-import mods.betterfoliage.util.*
+import mods.betterfoliage.render.block.vanilla.LeafParticleKey
+import mods.betterfoliage.util.Double3
+import mods.betterfoliage.util.PI2
+import mods.betterfoliage.util.minmax
+import mods.betterfoliage.util.randomB
+import mods.betterfoliage.util.randomD
+import mods.betterfoliage.util.randomF
+import mods.betterfoliage.util.randomI
 import net.fabricmc.fabric.api.event.world.WorldTickCallback
 import net.minecraft.client.MinecraftClient
 import net.minecraft.client.particle.ParticleTextureSheet
@@ -11,13 +17,13 @@ import net.minecraft.client.world.ClientWorld
 import net.minecraft.util.math.BlockPos
 import net.minecraft.util.math.MathHelper
 import net.minecraft.world.World
-import java.util.*
+import java.util.Random
 import kotlin.math.abs
 import kotlin.math.cos
 import kotlin.math.sin

 class FallingLeafParticle(
-    world: World, pos: BlockPos, leafKey: LeafKey
+    world: World, pos: BlockPos, leaf: LeafParticleKey, blockColor: Int, random: Random
 ) : AbstractParticle(
    world, pos.x.toDouble() + 0.5, pos.y.toDouble(), pos.z.toDouble() + 0.5
 ) {
@@ -26,12 +32,12 @@ class FallingLeafParticle(
        @JvmStatic val biomeBrightnessMultiplier = 0.5f
    }

-    var rotationSpeed = randomF(min = PI2 / 80.0, max = PI2 / 50.0)
+    var rotationSpeed = random.randomF(min = PI2 / 80.0, max = PI2 / 50.0)
    val isMirrored = randomB()
    var wasCollided = false

    init {
-        angle = randomF(max = PI2)
+        angle = random.randomF(max = PI2)
        prevAngle = angle - rotationSpeed

        maxAge = MathHelper.floor(randomD(0.6, 1.0) * BetterFoliage.config.fallingLeaves.lifetime * 20.0)
@@ -40,9 +46,9 @@ class FallingLeafParticle(
        scale = BetterFoliage.config.fallingLeaves.size.toFloat() * 0.1f

        val state = world.getBlockState(pos)
-        val blockColor = MinecraftClient.getInstance().blockColorMap.getColor(state, world, pos, 0)
-        sprite = LeafParticleRegistry[leafKey.leafType][randomI(max = 1024)]
-        setParticleColor(leafKey.overrideColor, blockColor)
+
+        setColor(leaf.overrideColor?.asInt ?: blockColor)
+        sprite = LeafParticleRegistry[leaf.leafType][randomI(max = 1024)]
    }

    override val isValid: Boolean get() = (sprite != null)
--- a/src/main/kotlin/mods/betterfoliage/render/particle/LeafParticleRegistry.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/particle/LeafParticleRegistry.kt
@@ -3,36 +3,70 @@ package mods.betterfoliage.render.particle
 import mods.betterfoliage.BetterFoliage
 import mods.betterfoliage.model.FixedSpriteSet
 import mods.betterfoliage.model.SpriteSet
-import mods.betterfoliage.util.*
+import mods.betterfoliage.resource.VeryEarlyReloadListener
+import mods.betterfoliage.util.Atlas
+import mods.betterfoliage.util.HasLogger
+import mods.betterfoliage.util.get
+import mods.betterfoliage.util.getLines
+import mods.betterfoliage.util.resourceManager
+import mods.betterfoliage.util.stripStart
 import net.fabricmc.fabric.api.event.client.ClientSpriteRegistryCallback
+import net.minecraft.client.texture.MissingSprite
 import net.minecraft.client.texture.SpriteAtlasTexture
+import net.minecraft.resource.ResourceManager
 import net.minecraft.util.Identifier
+import org.apache.logging.log4j.Level
+import kotlin.collections.MutableList
+import kotlin.collections.distinct
+import kotlin.collections.filter
+import kotlin.collections.firstOrNull
+import kotlin.collections.forEach
+import kotlin.collections.isNotEmpty
+import kotlin.collections.joinToString
+import kotlin.collections.listOf
+import kotlin.collections.map
+import kotlin.collections.mutableListOf
+import kotlin.collections.mutableMapOf
+import kotlin.collections.plus
+import kotlin.collections.set

-object LeafParticleRegistry : ClientSpriteRegistryCallback {
+object LeafParticleRegistry : HasLogger(), ClientSpriteRegistryCallback, VeryEarlyReloadListener {
    val typeMappings = TextureMatcher()
+    val allTypes get() = (typeMappings.mappings.map { it.type } + "default").distinct()

-    val ids = mutableMapOf<String, List<Identifier>>()
    val spriteSets = mutableMapOf<String, SpriteSet>()

-    override fun registerSprites(atlasTexture: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
-        ids.clear()
-        spriteSets.clear()
+    override fun getFabricId() = Identifier(BetterFoliage.MOD_ID, "leaf-particles")
+
+    override fun onReloadStarted(resourceManager: ResourceManager) {
        typeMappings.loadMappings(Identifier(BetterFoliage.MOD_ID, "leaf_texture_mappings.cfg"))
-        (typeMappings.mappings.map { it.type } + "default").distinct().forEach { leafType ->
+        detailLogger.log(Level.INFO, "Loaded leaf particle mappings, types = [${allTypes.joinToString(", ")}]")
+    }
+
+    override fun registerSprites(atlasTexture: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
+        spriteSets.clear()
+        allTypes.forEach { leafType ->
            val validIds = (0 until 16).map { idx -> Identifier(BetterFoliage.MOD_ID, "particle/falling_leaf_${leafType}_$idx") }
-                .filter { resourceManager.containsResource(Atlas.PARTICLES.wrap(it)) }
-            ids[leafType] = validIds
+                .filter { resourceManager.containsResource(Atlas.PARTICLES.file(it)) }
            validIds.forEach { registry.register(it) }
        }
    }

-    operator fun get(type: String): SpriteSet {
-        spriteSets[type]?.let { return it }
-        ids[type]?.let {
-            return FixedSpriteSet(Atlas.PARTICLES, it).apply { spriteSets[type] = this }
-        }
-        return if (type == "default") FixedSpriteSet(Atlas.PARTICLES, emptyList()).apply { spriteSets[type] = this }
-            else get("default")
+    operator fun get(leafType: String): SpriteSet {
+        spriteSets[leafType]?.let { return it }
+
+        val sprites = (0 until 16)
+            .map { idx -> Identifier(BetterFoliage.MOD_ID, "particle/falling_leaf_${leafType}_$idx") }
+            .map { Atlas.PARTICLES[it] }
+            .filter { it !is MissingSprite }
+
+        detailLogger.log(Level.INFO, "Leaf particle type [$leafType], ${sprites.size} sprites in atlas")
+        if (sprites.isNotEmpty()) return FixedSpriteSet(sprites).apply { spriteSets[leafType] = this }
+
+        return if (leafType == "default")
+            FixedSpriteSet(listOf(Atlas.PARTICLES[MissingSprite.getMissingSpriteId()])).apply { spriteSets[leafType] = this }
+        else
+            get("default").apply { spriteSets[leafType] = this }
    }

    init {
--- a/src/main/kotlin/mods/betterfoliage/render/particle/RisingSouls.kt
+++ b/src/main/kotlin/mods/betterfoliage/render/particle/RisingSouls.kt
@@ -21,7 +21,7 @@ class RisingSoulParticle(
    world, pos.x.toDouble() + 0.5, pos.y.toDouble() + 1.0, pos.z.toDouble() + 0.5
 ) {

-    val particleTrail: Deque<Double3> = LinkedList<Double3>()
+    val particleTrail: Deque<Double3> = LinkedList()
    val initialPhase = randomD(max = PI2)

    init {
--- a/src/main/kotlin/mods/betterfoliage/resource/VeryEarlyReloadListener.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/VeryEarlyReloadListener.kt
@@ -0,0 +1,28 @@
+package mods.betterfoliage.resource
+
+import net.fabricmc.fabric.api.resource.IdentifiableResourceReloadListener
+import net.minecraft.resource.ResourceManager
+import net.minecraft.resource.ResourceReloadListener
+import net.minecraft.util.profiler.Profiler
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Executor
+
+/**
+ * Catch resource reload extremely early, before any of the reloaders
+ * have started working.
+ */
+interface VeryEarlyReloadListener : ResourceReloadListener, IdentifiableResourceReloadListener {
+    override fun reload(
+        synchronizer: ResourceReloadListener.Synchronizer,
+        resourceManager: ResourceManager,
+        preparationsProfiler: Profiler,
+        reloadProfiler: Profiler,
+        backgroundExecutor: Executor,
+        gameExecutor: Executor
+    ): CompletableFuture<Void> {
+        onReloadStarted(resourceManager)
+        return synchronizer.whenPrepared(null)
+    }
+
+    fun onReloadStarted(resourceManager: ResourceManager) {}
+}
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/BakedModelReplacer.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/BakedModelReplacer.kt
@@ -1,102 +0,0 @@
-package mods.betterfoliage.resource.discovery
-
-import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.BlockModelsReloadCallback
-import mods.betterfoliage.ModelLoadingCallback
-import mods.betterfoliage.util.HasLogger
-import mods.betterfoliage.util.Invalidator
-import mods.betterfoliage.util.YarnHelper
-import mods.betterfoliage.util.get
-import net.fabricmc.fabric.api.event.client.ClientSpriteRegistryCallback
-import net.minecraft.block.BlockState
-import net.minecraft.client.MinecraftClient
-import net.minecraft.client.render.block.BlockModels
-import net.minecraft.client.render.model.BakedModel
-import net.minecraft.client.render.model.ModelLoader
-import net.minecraft.client.texture.SpriteAtlasTexture
-import net.minecraft.resource.ResourceManager
-import net.minecraft.util.Identifier
-import org.apache.logging.log4j.Level
-import org.apache.logging.log4j.Level.DEBUG
-import org.apache.logging.log4j.Level.WARN
-import java.lang.ref.WeakReference
-import java.util.*
-import java.util.concurrent.CompletableFuture
-import java.util.function.Consumer
-import java.util.function.Supplier
-
-// net.minecraft.client.render.block.BlockModels.models
-val BlockModels_models = YarnHelper.requiredField<Map<BlockState, BakedModel>>("net.minecraft.class_773", "field_4162", "Ljava/util/Map;")
-
-class BakedModelReplacer : ModelLoadingCallback, ClientSpriteRegistryCallback, BlockModelsReloadCallback, Invalidator, HasLogger {
-    override val logger get() = BetterFoliage.logDetail
-
-    val discoverers = mutableListOf<ModelDiscovery>()
-    override val callbacks = mutableListOf<WeakReference<()->Unit>>()
-
-    protected var keys = emptyMap<BlockState, BlockRenderKey>()
-
-    operator fun get(state: BlockState) = keys[state]
-    inline fun <reified T> getTyped(state: BlockState) = get(state) as? T
-
-    var currentLoader: ModelLoader? = null
-
-    override fun beginLoadModels(loader: ModelLoader, manager: ResourceManager) {
-        // Step 1: get a hold of the ModelLoader instance when model reloading starts
-        currentLoader = loader
-        log("reloading block discovery configuration")
-        BetterFoliage.blockConfig.reloadConfig(manager)
-        invalidate()
-    }
-
-    override fun registerSprites(atlasTexture: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
-        // Step 2: ModelLoader is finished with the unbaked models by now, we can inspect them
-        log("discovering blocks")
-        val idSet = Collections.synchronizedSet(mutableSetOf<Identifier>())
-        val allKeys = discoverers.map {
-            // run model discoverers in parallel
-            CompletableFuture.supplyAsync(Supplier {
-                it.discover(currentLoader!!, Consumer { idSet.add(it) })
-            }, MinecraftClient.getInstance())
-        }.map { it.join() }
-        idSet.forEach { registry.register(it) }
-
-        val result = mutableMapOf<BlockState, BlockRenderKey>()
-        allKeys.forEach { keys ->
-            keys.entries.forEach { (state, key) ->
-                val oldKey = result[state]
-                if (oldKey != null) log("Replacing $oldKey with $key for state $state")
-                else log(DEBUG, "Adding replacement $key for state $state")
-                result[state] = key
-            }
-        }
-
-        keys = result
-    }
-
-    override fun reloadBlockModels(blockModels: BlockModels) {
-        // Step 3: replace the baked models with our own
-        log("block model baking finished")
-        val modelMap = blockModels[BlockModels_models] as MutableMap<BlockState, BakedModel>
-        keys.forEach { (state, key) ->
-            val oldModel = modelMap[state]
-            if (oldModel == null) log(WARN, "Cannot find model for state $state, ignoring")
-            else {
-                try {
-                    val newModel = key.replace(oldModel, state)
-                    modelMap[state] = newModel
-                    log(DEBUG, "Replaced model for state $state with $key")
-                } catch (e: Exception) {
-                    log(WARN, "Error creating model for state $state with $key", e)
-                }
-            }
-        }
-    }
-
-    init {
-        ModelLoadingCallback.EVENT.register(this)
-        ClientSpriteRegistryCallback.event(SpriteAtlasTexture.BLOCK_ATLAS_TEX).register(this)
-        BlockModelsReloadCallback.EVENT.register(this)
-    }
-}
-
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/BakingLifecycle.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/BakingLifecycle.kt
@@ -0,0 +1,124 @@
+package mods.betterfoliage.resource.discovery
+
+import mods.betterfoliage.BetterFoliage
+import mods.betterfoliage.BlockModelsReloadCallback
+import mods.betterfoliage.ModelLoadingCallback
+import mods.betterfoliage.util.HasLogger
+import mods.betterfoliage.util.Invalidator
+import net.fabricmc.fabric.api.event.client.ClientSpriteRegistryCallback
+import net.minecraft.block.BlockState
+import net.minecraft.client.render.block.BlockModels
+import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.ModelBakeSettings
+import net.minecraft.client.render.model.ModelLoader
+import net.minecraft.client.render.model.UnbakedModel
+import net.minecraft.client.texture.Sprite
+import net.minecraft.client.texture.SpriteAtlasTexture
+import net.minecraft.client.util.SpriteIdentifier
+import net.minecraft.resource.ResourceManager
+import net.minecraft.util.Identifier
+import org.apache.logging.log4j.Level.INFO
+import org.apache.logging.log4j.Level.WARN
+import org.apache.logging.log4j.Logger
+import java.lang.ref.WeakReference
+import java.util.function.Function
+
+data class ModelDiscoveryContext(
+    val bakery: ModelLoader,
+    val blockState: BlockState,
+    val modelLocation: Identifier,
+    val sprites: MutableSet<Identifier>,
+    val replacements: MutableMap<Identifier, ModelBakingKey>,
+    val logger: Logger
+) {
+    fun getUnbaked(location: Identifier = modelLocation) = bakery.getOrLoadModel(location)
+    fun addReplacement(key: ModelBakingKey, addToStateKeys: Boolean = true) {
+        replacements[modelLocation] = key
+        if (addToStateKeys) BetterFoliage.blockTypes.stateKeys[blockState] = key
+        logger.log(INFO, "Adding model replacement $modelLocation -> $key")
+    }
+}
+
+interface ModelDiscovery {
+    fun onModelsLoaded(
+        bakery: ModelLoader,
+        sprites: MutableSet<Identifier>,
+        replacements: MutableMap<Identifier, ModelBakingKey>
+    )
+}
+
+data class ModelBakingContext(
+    val bakery: ModelLoader,
+    val spriteGetter: Function<SpriteIdentifier, Sprite>,
+    val location: Identifier,
+    val transform: ModelBakeSettings,
+    val logger: Logger
+) {
+    fun getUnbaked() = bakery.getOrLoadModel(location)
+    fun getBaked() = bakery.bake(location, transform)
+}
+
+interface ModelBakingKey {
+    fun bake(ctx: ModelBakingContext): BakedModel? =
+        ctx.getUnbaked().bake(ctx.bakery, ctx.spriteGetter, ctx.transform, ctx.location)
+}
+
+object BakeWrapperManager : HasLogger(), Invalidator, ModelLoadingCallback, ClientSpriteRegistryCallback, BlockModelsReloadCallback {
+    init {
+        ModelLoadingCallback.EVENT.register(this)
+        ClientSpriteRegistryCallback.event(SpriteAtlasTexture.BLOCK_ATLAS_TEX).register(this)
+    }
+    val discoverers = mutableListOf<ModelDiscovery>()
+    override val callbacks = mutableListOf<WeakReference<()->Unit>>()
+
+    private val replacements = mutableMapOf<Identifier, ModelBakingKey>()
+    private val sprites = mutableSetOf<Identifier>()
+
+    override fun beginLoadModels(loader: ModelLoader, manager: ResourceManager) {
+        val startTime = System.currentTimeMillis()
+        replacements.clear()
+        sprites.clear()
+        invalidate()
+        BetterFoliage.blockTypes = BlockTypeCache()
+
+        logger.log(INFO, "starting model discovery (${discoverers.size} listeners)")
+        discoverers.forEach { listener ->
+            val replacementsLocal = mutableMapOf<Identifier, ModelBakingKey>()
+            listener.onModelsLoaded(loader, sprites, replacements)
+        }
+
+        val elapsed = System.currentTimeMillis() - startTime
+        logger.log(INFO, "finished model discovery in $elapsed ms, ${replacements.size} top-level replacements")
+    }
+
+    override fun registerSprites(atlas: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
+        logger.log(INFO, "Adding ${sprites.size} sprites to block atlas")
+        sprites.forEach { registry.register(it) }
+        sprites.clear()
+    }
+
+    override fun reloadBlockModels(blockModels: BlockModels) {
+        replacements.clear()
+    }
+
+    fun onBake(
+        unbaked: UnbakedModel,
+        bakery: ModelLoader,
+        spriteGetter: Function<SpriteIdentifier, Sprite>,
+        transform: ModelBakeSettings,
+        location: Identifier
+    ): BakedModel? {
+        val ctx = ModelBakingContext(bakery, spriteGetter, location, transform, detailLogger)
+        // bake replacement if available
+        replacements[location]?.let { replacement ->
+            detailLogger.log(INFO, "Baking replacement for [${unbaked::class.java.simpleName}] $location -> $replacement")
+            try {
+                return replacement.bake(ctx)
+            } catch (e: Exception) {
+                detailLogger.log(WARN, "Error while baking $replacement", e)
+                logger.log(WARN, "Error while baking $replacement", e)
+            }
+        }
+        return unbaked.bake(bakery, spriteGetter, transform, location)
+    }
+}
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/BlockTypeCache.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/BlockTypeCache.kt
@@ -0,0 +1,14 @@
+package mods.betterfoliage.resource.discovery
+
+import net.minecraft.block.BlockState
+
+class BlockTypeCache {
+    val leaf = mutableSetOf<BlockState>()
+    val grass = mutableSetOf<BlockState>()
+    val dirt = mutableSetOf<BlockState>()
+
+    val stateKeys = mutableMapOf<BlockState, ModelBakingKey>()
+
+    inline fun <reified T> getTyped(state: BlockState) = stateKeys[state] as? T
+    inline fun <reified T> hasTyped(state: BlockState) = stateKeys[state] is T
+}
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/ConfigurableModelDiscovery.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/ConfigurableModelDiscovery.kt
@@ -1,69 +0,0 @@
-package mods.betterfoliage.resource.discovery
-
-import com.google.common.base.Joiner
-import mods.betterfoliage.util.YarnHelper
-import mods.betterfoliage.util.get
-import mods.betterfoliage.util.stripStart
-import net.minecraft.block.BlockState
-import net.minecraft.client.render.model.ModelLoader
-import net.minecraft.client.render.model.json.JsonUnbakedModel
-import net.minecraft.client.texture.MissingSprite
-import net.minecraft.util.Identifier
-import org.apache.logging.log4j.Level
-import org.apache.logging.log4j.Level.DEBUG
-import org.apache.logging.log4j.Level.INFO
-import java.util.function.Consumer
-
-// net.minecraft.client.render.model.json.JsonUnbakedModel.parent
-val JsonUnbakedModel_parent = YarnHelper.requiredField<JsonUnbakedModel>("net.minecraft.class_793", "field_4253", "Lnet/minecraft/class_793;")
-// net.minecraft.client.render.model.json.JsonUnbakedModel.parentId
-val JsonUnbakedModel_parentId = YarnHelper.requiredField<Identifier>("net.minecraft.class_793", "field_4247", "Lnet/minecraft/class_2960;")
-
-fun Pair<JsonUnbakedModel, Identifier>.derivesFrom(targetLocation: Identifier): Boolean {
-    if (second.stripStart("models/") == targetLocation) return true
-    if (first[JsonUnbakedModel_parent] != null && first[JsonUnbakedModel_parentId] != null)
-        return Pair(first[JsonUnbakedModel_parent]!!, first[JsonUnbakedModel_parentId]!!).derivesFrom(targetLocation)
-    return false
-}
-
-abstract class ConfigurableModelDiscovery : ModelDiscoveryBase() {
-
-    abstract val matchClasses: IBlockMatcher
-    abstract val modelTextures: List<ModelTextureList>
-
-    override fun discover(loader: ModelLoader, atlas: Consumer<Identifier>): Map<BlockState, BlockRenderKey> {
-        log(INFO, "Starting model discovery: ${this::class.java.canonicalName}")
-        matchClasses.describe(this)
-        modelTextures.forEach { modelTex ->
-            log(DEBUG, "    model: ${modelTex.modelLocation} textures: ${modelTex.textureNames.joinToString(", ")}")
-        }
-        return super.discover(loader, atlas)
-    }
-
-    abstract fun processModel(state: BlockState, textures: List<Identifier>, atlas: Consumer<Identifier>): BlockRenderKey?
-
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>): BlockRenderKey? {
-        val matchClass = matchClasses.matchingClass(ctx.state.block) ?: return null
-        log(DEBUG, "block state ${ctx.state.toString()}")
-        log(DEBUG, "      class ${ctx.state.block.javaClass.name} matches ${matchClass.name}")
-
-        (ctx.models.filter { it.first is JsonUnbakedModel } as List<Pair<JsonUnbakedModel, Identifier>>).forEach { (model, location) ->
-            val modelMatch = modelTextures.firstOrNull { (model to location).derivesFrom(it.modelLocation) }
-            if (modelMatch != null) {
-                log(DEBUG, "      model ${model} matches ${modelMatch.modelLocation}")
-
-                val textures = modelMatch.textureNames.map { it to model.resolveSprite(it).textureId }
-                val texMapString = Joiner.on(", ").join(textures.map { "${it.first}=${it.second}" })
-                log(DEBUG, "    sprites [$texMapString]")
-
-                if (textures.all { it.second != MissingSprite.getMissingSpriteId() }) {
-                    // found a valid model (all required textures exist)
-                    return processModel(ctx.state, textures.map { it.second }, atlas).also {
-                        log(DEBUG, "    valid model discovered: $it")
-                    }
-                }
-            }
-        }
-        return null
-    }
-}
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/Matchers.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/Matchers.kt
@@ -9,13 +9,12 @@ import net.minecraft.block.Block
 import net.minecraft.resource.ResourceManager
 import net.minecraft.util.Identifier
 import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.Level.INFO
 import org.apache.logging.log4j.Logger

 interface IBlockMatcher {
    fun matchesClass(block: Block): Boolean
    fun matchingClass(block: Block): Class<*>?
-
-    fun describe(logger: HasLogger)
 }

 class SimpleBlockMatcher(vararg val classes: Class<*>) : IBlockMatcher {
@@ -26,15 +25,9 @@ class SimpleBlockMatcher(vararg val classes: Class<*>) : IBlockMatcher {
        classes.forEach { if (it.isAssignableFrom(blockClass)) return it }
        return null
    }
-
-    override fun describe(logger: HasLogger) {
-        classes.forEach { klass ->
-            logger.log(Level.DEBUG, "    class whitelist: ${klass.name}")
-        }
-    }
 }

-class ConfigurableBlockMatcher(val logger: Logger, val location: Identifier) : IBlockMatcher {
+class ConfigurableBlockMatcher(val location: Identifier) : HasLogger(), IBlockMatcher {

    val blackList = mutableListOf<Class<*>>()
    val whiteList = mutableListOf<Class<*>>()
@@ -57,7 +50,7 @@ class ConfigurableBlockMatcher(val logger: Logger, val location: Identifier) : I
        blackList.clear()
        whiteList.clear()
        manager.getAllResources(location).forEach { resource ->
-            logger.info("Reading class list $location from pack ${resource.resourcePackName}")
+            detailLogger.log(INFO, "Reading class list $location from pack ${resource.resourcePackName}")
            resource.getLines().map{ it.trim() }.filter { !it.startsWith("//") && it.isNotEmpty() }.forEach { line ->
                val name = if (line.startsWith("-")) line.substring(1) else line
                val mappedName = FabricLoader.getInstance().mappingResolver.mapClassName(INTERMEDIARY, name)
@@ -75,26 +68,17 @@ class ConfigurableBlockMatcher(val logger: Logger, val location: Identifier) : I
            }
        }
    }
-
-    override fun describe(logger: HasLogger) {
-        whiteList.forEach { klass ->
-            logger.log(Level.DEBUG, "    class whitelist: ${klass.name}")
-        }
-        blackList.forEach { klass ->
-            logger.log(Level.DEBUG, "    class blacklist: ${klass.name}")
-        }
-    }
 }

 data class ModelTextureList(val modelLocation: Identifier, val textureNames: List<String>) {
    constructor(vararg args: String) : this(Identifier(args[0]), listOf(*args).drop(1))
 }

-class ModelTextureListConfiguration(val logger: Logger, val location: Identifier) {
+class ModelTextureListConfiguration(val location: Identifier) : HasLogger() {
    val modelList = mutableListOf<ModelTextureList>()
    fun readDefaults(manager: ResourceManager) {
        manager.getAllResources(location).forEach { resource ->
-            logger.info("Reading model configuration $location from pack ${resource.resourcePackName}")
+            detailLogger.log(INFO, "Reading model configuration $location from pack ${resource.resourcePackName}")
            resource.getLines().map{ it.trim() }.filter { !it.startsWith("//") && it.isNotEmpty() }.forEach { line ->
                val elements = line.split(",")
                modelList.add(ModelTextureList(Identifier(elements.first()), elements.drop(1)))
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/ModelDiscovery.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/ModelDiscovery.kt
@@ -0,0 +1,98 @@
+package mods.betterfoliage.resource.discovery
+
+import com.google.common.base.Joiner
+import mods.betterfoliage.util.HasLogger
+import mods.betterfoliage.util.YarnHelper
+import mods.betterfoliage.util.get
+import net.minecraft.client.render.block.BlockModels
+import net.minecraft.client.render.model.ModelLoader
+import net.minecraft.client.render.model.json.JsonUnbakedModel
+import net.minecraft.client.render.model.json.WeightedUnbakedModel
+import net.minecraft.client.texture.MissingSprite
+import net.minecraft.util.Identifier
+import net.minecraft.util.registry.Registry
+import org.apache.logging.log4j.Level
+
+abstract class AbstractModelDiscovery : HasLogger(), ModelDiscovery {
+    override fun onModelsLoaded(
+        bakery: ModelLoader,
+        sprites: MutableSet<Identifier>,
+        replacements: MutableMap<Identifier, ModelBakingKey>
+    ) {
+        Registry.BLOCK
+            .flatMap { block -> block.stateManager.states }
+            .forEach { state ->
+                val location = BlockModels.getModelId(state)
+                val ctx = ModelDiscoveryContext(bakery, state, location, sprites, replacements, detailLogger)
+                try {
+                    processModel(ctx)
+                } catch (e: Exception) {
+                    logger.log(Level.WARN, "Discovery error in $location", e)
+                }
+            }
+    }
+
+    open fun processModel(ctx: ModelDiscoveryContext) {
+        val model = ctx.getUnbaked()
+
+        // built-in support for container models
+        if (model is WeightedUnbakedModel) {
+            // per-location replacements need to be scoped to the variant list, as replacement models
+            // may need information from the BlockState which is not available at baking time
+            val scopedReplacements = mutableMapOf<Identifier, ModelBakingKey>()
+            model.variants.forEach { variant ->
+                processModel(ctx.copy(modelLocation = variant.location, replacements = scopedReplacements))
+            }
+            if (scopedReplacements.isNotEmpty()) {
+                ctx.addReplacement(WeightedUnbakedKey(scopedReplacements), addToStateKeys = false)
+            }
+        }
+    }
+}
+
+abstract class ConfigurableModelDiscovery : AbstractModelDiscovery() {
+    abstract val matchClasses: IBlockMatcher
+    abstract val modelTextures: List<ModelTextureList>
+
+    abstract fun processModel(
+        ctx: ModelDiscoveryContext,
+        textureMatch: List<Identifier>
+    )
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        val model = ctx.getUnbaked()
+        if (model is JsonUnbakedModel) {
+            val matchClass = matchClasses.matchingClass(ctx.blockState.block) ?: return
+
+            detailLogger.log(Level.INFO, "block state ${ctx.blockState}")
+            detailLogger.log(Level.INFO, "      model ${ctx.modelLocation}")
+            detailLogger.log(Level.INFO, "      class ${ctx.blockState.block.javaClass.name} matches ${matchClass.name}")
+
+            modelTextures
+                .filter { matcher -> ctx.bakery.modelDerivesFrom(model, ctx.modelLocation, matcher.modelLocation) }
+                .forEach { match ->
+                    detailLogger.log(Level.INFO, "      model $model matches ${match.modelLocation}")
+
+                    val materials = match.textureNames.map { it to model.resolveSprite(it) }
+                    val texMapString = Joiner.on(", ").join(materials.map { "${it.first}=${it.second.textureId}" })
+                    detailLogger.log(Level.INFO, "    sprites [$texMapString]")
+
+                    if (materials.all { it.second.textureId != MissingSprite.getMissingSpriteId() }) {
+                        // found a valid model (all required textures exist)
+                        processModel(ctx, materials.map { it.second.textureId })
+                    }
+                }
+        }
+        return super.processModel(ctx)
+    }
+}
+
+// net.minecraft.client.render.model.json.JsonUnbakedModel.parentId
+val JsonUnbakedModel_parentId = YarnHelper.requiredField<Identifier>("net.minecraft.class_793", "field_4247", "Lnet/minecraft/class_2960;")
+
+fun ModelLoader.modelDerivesFrom(model: JsonUnbakedModel, location: Identifier, target: Identifier): Boolean =
+    if (location == target) true
+    else model[JsonUnbakedModel_parentId]
+        ?.let { getOrLoadModel(it) as? JsonUnbakedModel }
+        ?.let { parent -> modelDerivesFrom(parent, model[JsonUnbakedModel_parentId]!!, target) }
+        ?: false
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/ModelDiscoveryBase.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/ModelDiscoveryBase.kt
@@ -1,74 +0,0 @@
-package mods.betterfoliage.resource.discovery
-
-import mods.betterfoliage.util.HasLogger
-import net.minecraft.block.BlockState
-import net.minecraft.client.render.block.BlockModels
-import net.minecraft.client.render.model.BakedModel
-import net.minecraft.client.render.model.ModelLoader
-import net.minecraft.client.render.model.UnbakedModel
-import net.minecraft.client.render.model.json.JsonUnbakedModel
-import net.minecraft.client.render.model.json.ModelVariant
-import net.minecraft.client.render.model.json.WeightedUnbakedModel
-import net.minecraft.client.texture.SpriteAtlasTexture
-import net.minecraft.client.util.ModelIdentifier
-import net.minecraft.util.Identifier
-import net.minecraft.util.registry.Registry
-import java.util.function.Consumer
-
-interface BlockRenderKey {
-    fun replace(model: BakedModel, state: BlockState): BakedModel = model
-}
-
-fun ModelLoader.iterateModels(func: (ModelDiscoveryContext)->Unit) {
-    Registry.BLOCK.flatMap { block ->
-        block.stateManager.states.map { state -> state to BlockModels.getModelId(state) }
-    }.forEach { (state, stateModelResource) ->
-        func(ModelDiscoveryContext(this, state, stateModelResource))
-    }
-}
-
-/**
- * Information about a single [BlockState] and all the [UnbakedModel]s it could render as.
- */
-class ModelDiscoveryContext(
-    loader: ModelLoader,
-    val state: BlockState,
-    val modelId: ModelIdentifier
-) {
-    val models = loader.unwrapVariants(loader.getOrLoadModel(modelId) to modelId)
-        .filter { it.second != loader.getOrLoadModel(ModelLoader.MISSING) }
-
-    fun ModelLoader.unwrapVariants(modelAndLoc: Pair<UnbakedModel, Identifier>): List<Pair<UnbakedModel, Identifier>> = when(val model = modelAndLoc.first) {
-        is WeightedUnbakedModel -> (model.variants as List<ModelVariant>).flatMap {
-            variant -> unwrapVariants(getOrLoadModel(variant.location) to variant.location)
-        }
-        is JsonUnbakedModel -> listOf(modelAndLoc)
-        else -> emptyList()
-    }
-}
-
-interface ModelDiscovery {
-    fun discover(loader: ModelLoader, atlas: Consumer<Identifier>): Map<BlockState, BlockRenderKey>
-}
-
-abstract class ModelDiscoveryBase : ModelDiscovery, HasLogger {
-    override fun discover(loader: ModelLoader, atlas: Consumer<Identifier>): Map<BlockState, BlockRenderKey> {
-        val keys = mutableMapOf<BlockState, BlockRenderKey>()
-        var errors = 0
-
-        loader.iterateModels { ctx ->
-            try {
-                val result = processModel(ctx, atlas)
-                result?.let { keys[ctx.state] = it }
-            } catch (e: Exception) {
-                errors++
-            }
-        }
-        log("${keys.size} BlockStates discovered, $errors errors")
-        return keys
-    }
-
-    abstract fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>): BlockRenderKey?
-}
-
-
--- a/src/main/kotlin/mods/betterfoliage/resource/discovery/VanillaWrappers.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/discovery/VanillaWrappers.kt
@@ -0,0 +1,61 @@
+package mods.betterfoliage.resource.discovery
+
+import mods.betterfoliage.model.WeightedModelWrapper
+import mods.betterfoliage.model.WrappedBakedModel
+import mods.betterfoliage.model.WrappedMeshModel
+import mods.betterfoliage.util.HasLogger
+import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
+import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.WeightedUnbakedModel
+import net.minecraft.util.Identifier
+import org.apache.logging.log4j.Level.INFO
+import org.apache.logging.log4j.Level.WARN
+
+class WeightedUnbakedKey(
+    val replacements: Map<Identifier, ModelBakingKey>
+) : ModelBakingKey {
+
+    override fun bake(ctx: ModelBakingContext): BakedModel? {
+        val unbaked = ctx.getUnbaked()
+        if (unbaked !is WeightedUnbakedModel) return super.bake(ctx)
+
+        // bake all variants, replace as needed
+        val bakedModels = unbaked.variants.mapNotNull {
+            val variantCtx = ctx.copy(location = it.location, transform = it)
+            val replacement = replacements[it.location]
+            val baked = replacement?.let { replacement ->
+                ctx.logger.log(INFO, "Baking replacement for variant [${variantCtx.getUnbaked()::class.java.simpleName}] ${variantCtx.location} -> $replacement")
+                replacement.bake(variantCtx)
+            } ?: variantCtx.getBaked()
+            when(baked) {
+                is WrappedBakedModel -> it to baked
+                // just in case we replaced some variants in the list, but not others
+                // this should not realistically happen, this is just a best-effort fallback
+                is BasicBakedModel -> it to WrappedMeshModel.converter(
+                    state = null, unshade = false, noDiffuse = true, blendModeOverride = BlendMode.CUTOUT_MIPPED
+                ).convert(baked)!!
+                else -> null
+            }
+        }
+
+        // something fishy going on, possibly unknown model type
+        // let it through unchanged
+        if (bakedModels.isEmpty()) return super.bake(ctx)
+
+        if (bakedModels.size < unbaked.variants.size) {
+            detailLogger.log(
+                WARN,
+                "Dropped ${unbaked.variants.size - bakedModels.size} variants from model ${ctx.location}"
+            )
+        }
+        val weightedSpecials = bakedModels.map { (variant, model) ->
+            WeightedModelWrapper.WeightedModel(model, variant.weight)
+        }
+        return WeightedModelWrapper(weightedSpecials, weightedSpecials[0].model)
+    }
+
+    override fun toString() = "[WeightedUnbakedKey, ${replacements.size} replacements]"
+
+    companion object : HasLogger()
+}
--- a/src/main/kotlin/mods/betterfoliage/resource/generated/CenteredSprite.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/generated/CenteredSprite.kt
@@ -13,7 +13,7 @@ data class CenteredSprite(val sprite: Identifier, val atlas: Atlas = Atlas.BLOCK
    fun register(pack: GeneratedBlockTexturePack) = pack.register(this, this::draw)

    fun draw(resourceManager: ResourceManager): ByteArray {
-        val baseTexture = resourceManager.loadSprite(atlas.wrap(sprite))
+        val baseTexture = resourceManager.loadSprite(atlas.file(sprite))

        val frameWidth = baseTexture.width
        val frameHeight = baseTexture.width * aspectHeight / aspectWidth
--- a/src/main/kotlin/mods/betterfoliage/resource/generated/GeneratedBlockTexturePack.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/generated/GeneratedBlockTexturePack.kt
@@ -10,6 +10,8 @@ import net.minecraft.resource.metadata.ResourceMetadataReader
 import net.minecraft.text.LiteralText
 import net.minecraft.util.Identifier
 import net.minecraft.util.profiler.Profiler
+import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.Level.INFO
 import org.apache.logging.log4j.Logger
 import java.io.IOException
 import java.lang.IllegalStateException
@@ -29,7 +31,9 @@ import java.util.function.Supplier
 * @param[packDesc] Description of pack
 * @param[logger] Logger to log to when generating resources
 */
-class GeneratedBlockTexturePack(val reloadId: Identifier, val nameSpace: String, val packName: String, val packDesc: String, override val logger: Logger) : HasLogger, ResourcePack {
+class GeneratedBlockTexturePack(
+    val reloadId: Identifier, val nameSpace: String, val packName: String, val packDesc: String
+) : HasLogger(), ResourcePack {

    override fun getName() = reloadId.toString()
    override fun getNamespaces(type: ResourceType) = setOf(nameSpace)
@@ -51,8 +55,8 @@ class GeneratedBlockTexturePack(val reloadId: Identifier, val nameSpace: String,
        val resource = func(manager!!)

        identifiers[key] = id
-        resources[Atlas.BLOCKS.wrap(id)] = resource
-        log("generated resource $key -> $id")
+        resources[Atlas.BLOCKS.file(id)] = resource
+        detailLogger.log(INFO, "generated resource $key -> $id")
        return id
    }

--- a/src/main/kotlin/mods/betterfoliage/resource/generated/GeneratedGrassSprite.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/generated/GeneratedGrassSprite.kt
@@ -17,7 +17,7 @@ data class GeneratedGrassSprite(val sprite: Identifier, val isSnowed: Boolean, v
    fun register(pack: GeneratedBlockTexturePack) = pack.register(this, this::draw)

    fun draw(resourceManager: ResourceManager): ByteArray {
-        val baseTexture = resourceManager.loadSprite(atlas.wrap(sprite))
+        val baseTexture = resourceManager.loadSprite(atlas.file(sprite))

        val result = BufferedImage(baseTexture.width, baseTexture.height, BufferedImage.TYPE_4BYTE_ABGR)
        val graphics = result.createGraphics()
--- a/src/main/kotlin/mods/betterfoliage/resource/generated/GeneratedLeafSprite.kt
+++ b/src/main/kotlin/mods/betterfoliage/resource/generated/GeneratedLeafSprite.kt
@@ -20,7 +20,7 @@ data class GeneratedLeafSprite(val sprite: Identifier, val leafType: String, val
    fun register(pack: GeneratedBlockTexturePack) = pack.register(this, this::draw)

    fun draw(resourceManager: ResourceManager): ByteArray {
-        val baseTexture = resourceManager.loadSprite(atlas.wrap(sprite))
+        val baseTexture = resourceManager.loadSprite(atlas.file(sprite))

        val size = baseTexture.width
        val frames = baseTexture.height / size
@@ -67,7 +67,7 @@ data class GeneratedLeafSprite(val sprite: Identifier, val leafType: String, val
     * @param[maxSize] Preferred mask size.
     */
    fun getLeafMask(type: String, maxSize: Int) = getMultisizeTexture(maxSize) { size ->
-        Atlas.BLOCKS.wrap(Identifier(BetterFoliage.MOD_ID, "blocks/leafmask_${size}_${type}"))
+        Atlas.BLOCKS.file(Identifier(BetterFoliage.MOD_ID, "blocks/leafmask_${size}_${type}"))
    }

    /**
--- a/src/main/kotlin/mods/betterfoliage/util/Misc.kt
+++ b/src/main/kotlin/mods/betterfoliage/util/Misc.kt
@@ -1,6 +1,7 @@
@file:Suppress("NOTHING_TO_INLINE")
 package mods.betterfoliage.util

+import mods.betterfoliage.BetterFoliage
 import net.minecraft.text.LiteralText
 import net.minecraft.text.Style
 import net.minecraft.util.Formatting
@@ -63,13 +64,10 @@ fun nextPowerOf2(x: Int): Int {
 //    else -> false
 //}

-interface HasLogger {
-    val logger: Logger
-    val logName: String get() = this::class.simpleName!!
-    fun log(msg: String) = log(Level.INFO, msg)
-    fun log(level: Level, msg: String) = logger.log(level, "[$logName] $msg")
-    fun log(msg: String, e: Throwable) = log(Level.WARN, msg, e)
-    fun log(level: Level, msg: String, e: Throwable) = logger.log(level, "[$logName] $msg", e)
+@Suppress("LeakingThis")
+abstract class HasLogger {
+    val logger = BetterFoliage.logger(this)
+    val detailLogger = BetterFoliage.detailLogger(this)
 }

 fun textComponent(msg: String, color: Formatting = Formatting.GRAY): LiteralText {
--- a/src/main/kotlin/mods/betterfoliage/util/Random.kt
+++ b/src/main/kotlin/mods/betterfoliage/util/Random.kt
@@ -1,15 +1,18 @@
 package mods.betterfoliage.util

 import net.minecraft.util.math.BlockPos
-import kotlin.random.Random
+import java.util.Random

 val random = Random(System.nanoTime())

 fun randomB() = random.nextBoolean()
-fun randomI(min: Int = 0, max: Int = Int.MAX_VALUE) = random.nextInt(min, max)
-fun randomL(min: Long = 0, max: Long = Long.MAX_VALUE) = random.nextLong(min, max)
-fun randomF(min: Double = 0.0, max: Double = 1.0) = random.nextDouble(min, max).toFloat()
-fun randomD(min: Double = 0.0, max: Double = 1.0) = if (min == max) min else random.nextDouble(min, max)
+fun randomI(min: Int = 0, max: Int = Int.MAX_VALUE) = min + random.nextInt(max - min)
+fun randomF(min: Float = 0.0f, max: Float = 1.0f) = random.randomF(min, max)
+fun randomD(min: Double = 0.0, max: Double = 1.0) = random.randomD(min, max)
+
+fun Random.randomF(min: Float = 0.0f, max: Float = 1.0f) = nextFloat() * (max - min) + min
+fun Random.randomF(min: Double = 0.0, max: Double = 1.0) = randomF(min.toFloat(), max.toFloat())
+fun Random.randomD(min: Double = 0.0, max: Double = 1.0) = nextDouble() * (max - min) + min

 fun semiRandom(x: Int, y: Int, z: Int, seed: Int): Int {
    var value = (x * x + y * y + z * z + x * y + y * z + z * x + (seed * seed))
--- a/src/main/kotlin/mods/betterfoliage/util/Reflection.kt
+++ b/src/main/kotlin/mods/betterfoliage/util/Reflection.kt
@@ -3,6 +3,7 @@ package mods.betterfoliage.util
 import java.lang.reflect.Field
 import java.lang.reflect.Method
 import net.fabricmc.loader.api.FabricLoader
+import net.fabricmc.mappings.EntryTriple
 import org.apache.logging.log4j.Level
 import org.apache.logging.log4j.LogManager
 import java.lang.Exception
@@ -22,7 +23,7 @@ fun <T> Any.reflectField(name: String) = getFieldRecursive(this::class.java, nam
 fun getFieldRecursive(cls: Class<*>, name: String): Field = try {
    cls.getDeclaredField(name)
 } catch (e: NoSuchFieldException) {
-    cls.superclass?.let { getFieldRecursive(it, name) } ?: throw e
+    cls.superclass?.let { getFieldRecursive(it, name) } ?: throw IllegalArgumentException(e)
 }

 /** Get the method on the class with the given name.
@@ -31,7 +32,7 @@ fun getFieldRecursive(cls: Class<*>, name: String): Field = try {
 fun getMethodRecursive(cls: Class<*>, name: String): Method = try {
    cls.declaredMethods.find { it.name == name } ?: throw NoSuchMethodException()
 } catch (e: NoSuchMethodException) {
-    cls.superclass?.let { getMethodRecursive(it, name) } ?: throw e
+    cls.superclass?.let { getMethodRecursive(it, name) } ?: throw IllegalArgumentException(e)
 }

 fun getAllMethods(className: String, methodName: String): List<Method> =
@@ -68,7 +69,7 @@ object YarnHelper {
    val resolver = FabricLoader.getInstance().mappingResolver

    fun <T> requiredField(className: String, fieldName: String, descriptor: String) = Field<T>(false, className, fieldName, descriptor)
-    fun <T> requiredMethod(className: String, methodName: String, descriptor: String, vararg params: String) = Method<T>(false, className, methodName, descriptor)
+    fun <T> requiredMethod(className: String, methodName: String, descriptor: String) = Method<T>(false, className, methodName, descriptor)

    class Field<T>(val optional: Boolean, val className: String, val fieldName: String, descriptor: String) : FieldRef<T> {
        override val field = FabricLoader.getInstance().mappingResolver.let { resolver ->
@@ -103,8 +104,6 @@ object YarnHelper {
    }
 }

-//fun Any.isInstance(cls: ClassRefOld<*>) = cls.isInstance(this)
-
 interface ReflectionCallable<T> {
    operator fun invoke(vararg args: Any): T
 }
--- a/src/main/kotlin/mods/betterfoliage/util/Sprites.kt
+++ b/src/main/kotlin/mods/betterfoliage/util/Sprites.kt
@@ -1,30 +1,34 @@
 package mods.betterfoliage.util

+import mods.betterfoliage.model.Color
 import mods.betterfoliage.model.HSB
 import net.minecraft.client.MinecraftClient
+import net.minecraft.client.texture.NativeImage
 import net.minecraft.client.texture.Sprite
 import net.minecraft.client.texture.SpriteAtlasTexture
 import net.minecraft.resource.Resource
 import net.minecraft.resource.ResourceManager
 import net.minecraft.util.Identifier
+import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.Logger
 import java.awt.image.BufferedImage
 import java.io.ByteArrayOutputStream
 import java.io.IOException
 import javax.imageio.ImageIO
 import kotlin.math.atan2

-enum class Atlas(val basePath: String, val resourceId: Identifier) {
-    BLOCKS("textures", SpriteAtlasTexture.BLOCK_ATLAS_TEX),
-    PARTICLES("textures", SpriteAtlasTexture.PARTICLE_ATLAS_TEX);
+enum class Atlas(val resourceId: Identifier) {
+    BLOCKS(SpriteAtlasTexture.BLOCK_ATLAS_TEX),
+    PARTICLES(SpriteAtlasTexture.PARTICLE_ATLAS_TEX);

-    /** Get the fully-qualified resource name for sprites belonging to this atlas*/
-    fun wrap(resource: Identifier) = Identifier(resource.namespace, "$basePath/${resource.path}.png")
-
-    /** Get the short resource name for sprites belonging to this atlas*/
-    fun unwrap(resource: Identifier) = resource.stripStart("$basePath/").stripEnd(".png")
+    /** Get the fully-qualified resource name for sprites belonging to this atlas */
+    fun file(resource: Identifier) = Identifier(resource.namespace, "textures/${resource.path}.png")

    /** Reference to the atlas itself */
-    val atlas: SpriteAtlasTexture get() = MinecraftClient.getInstance().textureManager.getTexture(resourceId) as SpriteAtlasTexture
+    private val atlas: SpriteAtlasTexture get() = MinecraftClient.getInstance().textureManager.getTexture(resourceId) as SpriteAtlasTexture
+
+    /** Get a sprite from this atlas */
+    operator fun get(location: Identifier) = atlas.getSprite(location)
 }

 operator fun SpriteAtlasTexture.get(res: Identifier): Sprite? = getSprite(res)
@@ -48,15 +52,17 @@ val BufferedImage.bytes: ByteArray get() =
 * Only non-transparent pixels are considered. Averages are taken in the HSB color space (note: Hue is a circular average),
 * and the result transformed back to the RGB color space.
 */
-fun ResourceManager.averageImageColorHSB(id: Identifier, atlas: Atlas) = loadSprite(atlas.wrap(id)).let { image ->
+val Sprite_images = YarnHelper.requiredField<Array<NativeImage>>("net.minecraft.class_1058", "field_5262", "[Lnet/minecraft/class_1011;")
+
+val Sprite.averageColor: HSB get() {
    var numOpaque = 0
    var sumHueX = 0.0
    var sumHueY = 0.0
    var sumSaturation = 0.0f
    var sumBrightness = 0.0f
-    for (x in 0 until image.width)
-        for (y in 0 until image.height) {
-            val pixel = image.get(x, y)
+    for (x in 0 until width)
+        for (y in 0 until height) {
+            val pixel = this[Sprite_images]!![0].getPixelRgba(x, y)
            val alpha = (pixel shr 24) and 255
            val hsb = HSB.fromColor(pixel)
            if (alpha == 255) {
@@ -70,7 +76,7 @@ fun ResourceManager.averageImageColorHSB(id: Identifier, atlas: Atlas) = loadSpr

    // circular average - transform sum vector to polar angle
    val avgHue = (atan2(sumHueY, sumHueX) / PI2 + 0.5).toFloat()
-    HSB(avgHue, sumSaturation / numOpaque.toFloat(), sumBrightness / numOpaque.toFloat())
+    return HSB(avgHue, sumSaturation / numOpaque.toFloat(), sumBrightness / numOpaque.toFloat())
 }

 /** Weighted blend of 2 packed RGB colors */
@@ -82,3 +88,15 @@ fun blendRGB(rgb1: Int, rgb2: Int, weight1: Int, weight2: Int): Int {
    val result = ((a shl 24) or (r shl 16) or (g shl 8) or b)
    return result
 }
+
+fun logColorOverride(logger: Logger, threshold: Double, hsb: HSB) {
+    return if (hsb.saturation >= threshold) {
+        logger.log(Level.INFO, "         brightness ${hsb.brightness}")
+        logger.log(Level.INFO, "         saturation ${hsb.saturation} >= ${threshold}, will use texture color")
+    } else {
+        logger.log(Level.INFO, "         saturation ${hsb.saturation} < ${threshold}, will use block color")
+    }
+}
+
+fun HSB.colorOverride(threshold: Double) =
+    if (saturation < threshold) null else copy(brightness = (brightness * 2.0f).coerceAtMost(0.9f)).asColor.let { Color(it) }
--- a/src/main/resources/META-INF/MANIFEST.MF
+++ b/src/main/resources/META-INF/MANIFEST.MF
@@ -1,2 +1,6 @@
 Manifest-Version: 1.0
-Implementation-Title: betterfoliage
+Specification-Title: BetterFoliage
+Implementation-Title: BetterFoliage
+Specification-Vendor: octarine-noise
+Implementation-Vendor: octarine-noise
+MixinConnector: mods.betterfoliage.MixinConnector
--- a/src/main/resources/betterfoliage.mixins.json
+++ b/src/main/resources/betterfoliage.mixins.json
@@ -4,6 +4,7 @@
  "refmap": "betterfoliage-refmap.json",
  "compatibilityLevel": "JAVA_8",
  "minVersion": "0.8-SNAPSHOT",
+  "plugin": "mods.betterfoliage.MixinConfigPlugin",
  "mixins": [
  ],
  "client": [
@@ -13,7 +14,9 @@
    "MixinClientWorld",
    "MixinClientChunkManager",
    "MixinClientChunkManagerChunkMap",
-    "MixinModelLoader"
+    "MixinModelLoader",
+    "MixinModelLoaderVanilla",
+    "MixinModelLoaderOptifine"
  ],
  "server": [
  ],
--- a/src/main/resources/bf_generated_pack.png
+++ b/src/main/resources/bf_generated_pack.png