1 files changed, 12 insertions, 3 deletions
diff --git a/src/main/java/moe/nea/firmament/mixins/CustomModelBakerPatch.java b/src/main/java/moe/nea/firmament/mixins/CustomModelBakerPatch.java
index 97f81b0..c1e359d 100644
--- a/src/main/java/moe/nea/firmament/mixins/CustomModelBakerPatch.java
+++ b/src/main/java/moe/nea/firmament/mixins/CustomModelBakerPatch.java
@@ -1,5 +1,3 @@
-
-
 package moe.nea.firmament.mixins;
 
 import moe.nea.firmament.events.BakeExtraModelsEvent;
@@ -10,6 +8,7 @@ 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.Unique;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
@@ -29,9 +28,19 @@ public abstract class CustomModelBakerPatch {
     @Shadow
     abstract UnbakedModel getOrLoadModel(Identifier id);
 
+    @Shadow
+    protected abstract void add(ModelIdentifier id, UnbakedModel model);
+
+    @Unique
+    private void loadNonItemModel(ModelIdentifier identifier) {
+        UnbakedModel unbakedModel = this.getOrLoadModel(identifier.id());
+        this.add(identifier, unbakedModel);
+    }
+
+
     @Inject(method = "bake", at = @At("HEAD"))
     public void onBake(ModelLoader.SpriteGetter spliteGetter, CallbackInfo ci) {
-        BakeExtraModelsEvent.Companion.publish(new BakeExtraModelsEvent(this::loadItemModel));
+        BakeExtraModelsEvent.Companion.publish(new BakeExtraModelsEvent(this::loadItemModel, this::loadNonItemModel));
         modelsToBake.values().forEach(model -> model.setParents(this::getOrLoadModel));
 //        modelsToBake.keySet().stream()
 //            .filter(it -> !it.id().getNamespace().equals("minecraft"))