2 files changed, 30 insertions, 11 deletions
diff --git a/src/main/kotlin/features/texturepack/BakedModelExtra.kt b/src/main/kotlin/features/texturepack/BakedModelExtra.kt
index 32f419a..6305748 100644
--- a/src/main/kotlin/features/texturepack/BakedModelExtra.kt
+++ b/src/main/kotlin/features/texturepack/BakedModelExtra.kt
@@ -1,11 +1,30 @@
-
 package moe.nea.firmament.features.texturepack
 
+import net.fabricmc.fabric.api.renderer.v1.model.WrapperBakedModel as WrapperBakedModelFabric
 import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.WrapperBakedModel
+import moe.nea.firmament.util.ErrorUtil
 
 interface BakedModelExtra {
+	companion object {
+		@JvmStatic
+		fun cast(originalModel: BakedModel): BakedModelExtra? {
+			var p = originalModel
+			for (i in 0..256) {
+				p = when (p) {
+					is BakedModelExtra -> return p
+					is WrapperBakedModel -> p.wrapped
+					is WrapperBakedModelFabric -> WrapperBakedModelFabric.unwrap(p)
+					else -> break
+				}
+			}
+			ErrorUtil.softError("Could not find a baked model for $originalModel")
+			return null
+		}
+	}
+
 	var tintOverrides_firmament: TintOverrides?
 
 	fun getHeadModel_firmament(): BakedModel?
-    fun setHeadModel_firmament(headModel: BakedModel?)
+	fun setHeadModel_firmament(headModel: BakedModel?)
 }
diff --git a/src/main/kotlin/features/texturepack/TintOverrides.kt b/src/main/kotlin/features/texturepack/TintOverrides.kt
index 8006db8..85fcae4 100644
--- a/src/main/kotlin/features/texturepack/TintOverrides.kt
+++ b/src/main/kotlin/features/texturepack/TintOverrides.kt
@@ -3,7 +3,6 @@ package moe.nea.firmament.features.texturepack
 import com.google.gson.JsonObject
 import com.google.gson.JsonPrimitive
 import moe.nea.firmament.util.ErrorUtil
-import moe.nea.firmament.util.assertNotNullOr
 
 data class TintOverrides(
 	val layerMap: Map<Int, TintOverride> = mapOf()
@@ -14,21 +13,22 @@ data class TintOverrides(
 		val EMPTY = TintOverrides()
 		private val threadLocal = object : ThreadLocal<TintOverrides>() {}
 		fun enter(overrides: TintOverrides?) {
-			ErrorUtil.softCheck("Double entered tintOverrides") {
-				threadLocal.get() == null
-			}
+			ErrorUtil.softCheck("Double entered tintOverrides",
+			                    threadLocal.get() == null)
 			threadLocal.set(overrides ?: EMPTY)
 		}
 
 		fun exit(overrides: TintOverrides?) {
-			ErrorUtil.softCheck("Exited with non matching enter tintOverrides") {
-				threadLocal.get() == (overrides ?: EMPTY)
-			}
+			ErrorUtil.softCheck("Exited with non matching enter tintOverrides",
+			                    threadLocal.get() == (overrides ?: EMPTY))
 			threadLocal.remove()
 		}
 
-		fun getCurrentOverrides() =
-			assertNotNullOr(threadLocal.get(), "Got current tintOverrides without entering") { EMPTY }
+		fun getCurrentOverrides(): TintOverrides {
+			return ErrorUtil.notNullOr(threadLocal.get(), "Got current tintOverrides without entering") {
+				EMPTY
+			}
+		}
 
 		fun parse(jsonObject: JsonObject): TintOverrides {
 			val map = mutableMapOf<Int, TintOverride>()