--- a/net/minecraft/client/renderer/model/ModelBakery.java
+++ b/net/minecraft/client/renderer/model/ModelBakery.java
@@ -146,8 +146,17 @@
    });
 
    public ModelBakery(IResourceManager p_i226056_1_, BlockColors p_i226056_2_, IProfiler p_i226056_3_, int p_i226056_4_) {
+      this(p_i226056_1_, p_i226056_2_, true);
+      processLoading(p_i226056_3_, p_i226056_4_);
+   }
+
+   protected ModelBakery(IResourceManager p_i226056_1_, BlockColors p_i226056_2_, boolean vanillaBakery) {
       this.field_177598_f = p_i226056_1_;
       this.field_225365_D = p_i226056_2_;
+   }
+
+   protected void processLoading(IProfiler p_i226056_3_, int p_i226056_4_) {
+      net.minecraftforge.client.model.ModelLoaderRegistry.onModelLoadingStart();
       p_i226056_3_.func_76320_a("missing_model");
 
       try {
@@ -180,12 +189,16 @@
 
       p_i226056_3_.func_219895_b("special");
       this.func_217843_a(new ModelResourceLocation("minecraft:trident_in_hand#inventory"));
+      for (ResourceLocation rl : getSpecialModels()) {
+         addModelToCache(rl);
+      }
       p_i226056_3_.func_219895_b("textures");
       Set<Pair<String, String>> set = Sets.newLinkedHashSet();
       Set<RenderMaterial> set1 = this.field_217851_H.values().stream().flatMap((p_229342_2_) -> {
          return p_229342_2_.func_225614_a_(this::func_209597_a, set).stream();
       }).collect(Collectors.toSet());
       set1.addAll(field_177602_b);
+      net.minecraftforge.client.ForgeHooksClient.gatherFluidTextures(set1);
       set.stream().filter((p_229346_0_) -> {
          return !p_229346_0_.getSecond().equals(field_229321_r_);
       }).forEach((p_229330_0_) -> {
@@ -224,6 +237,7 @@
          try {
             ibakedmodel = this.func_217845_a(p_229350_1_, ModelRotation.X0_Y0);
          } catch (Exception exception) {
+            exception.printStackTrace();
             field_177603_c.warn("Unable to bake model: '{}': {}", p_229350_1_, exception);
          }
 
@@ -442,6 +456,13 @@
       this.field_217848_D.addAll(p_209593_2_.func_187965_e());
    }
 
+   // Same as loadTopModel but without restricting to MRL's
+   private void addModelToCache(ResourceLocation p_217843_1_) {
+      IUnbakedModel iunbakedmodel = this.func_209597_a(p_217843_1_);
+      this.field_217849_F.put(p_217843_1_, iunbakedmodel);
+      this.field_217851_H.put(p_217843_1_, iunbakedmodel);
+   }
+
    private void func_217843_a(ModelResourceLocation p_217843_1_) {
       IUnbakedModel iunbakedmodel = this.func_209597_a(p_217843_1_);
       this.field_217849_F.put(p_217843_1_, iunbakedmodel);
@@ -456,7 +477,13 @@
    }
 
    @Nullable
+   @Deprecated
    public IBakedModel func_217845_a(ResourceLocation p_217845_1_, IModelTransform p_217845_2_) {
+      return getBakedModel(p_217845_1_, p_217845_2_, this.field_229322_z_::func_229151_a_);
+   }
+
+   @Nullable
+   public IBakedModel getBakedModel(ResourceLocation p_217845_1_, IModelTransform p_217845_2_, java.util.function.Function<RenderMaterial, net.minecraft.client.renderer.texture.TextureAtlasSprite> textureGetter) {
       Triple<ResourceLocation, TransformationMatrix, Boolean> triple = Triple.of(p_217845_1_, p_217845_2_.func_225615_b_(), p_217845_2_.func_188049_c());
       if (this.field_217850_G.containsKey(triple)) {
          return this.field_217850_G.get(triple);
@@ -467,11 +494,11 @@
          if (iunbakedmodel instanceof BlockModel) {
             BlockModel blockmodel = (BlockModel)iunbakedmodel;
             if (blockmodel.func_178310_f() == field_177606_o) {
-               return field_217854_z.func_209579_a(this.field_229322_z_::func_229151_a_, blockmodel).func_228813_a_(this, blockmodel, this.field_229322_z_::func_229151_a_, p_217845_2_, p_217845_1_, false);
+               return field_217854_z.func_209579_a(textureGetter, blockmodel).func_228813_a_(this, blockmodel, this.field_229322_z_::func_229151_a_, p_217845_2_, p_217845_1_, false);
             }
          }
 
-         IBakedModel ibakedmodel = iunbakedmodel.func_225613_a_(this, this.field_229322_z_::func_229151_a_, p_217845_2_, p_217845_1_);
+         IBakedModel ibakedmodel = iunbakedmodel.func_225613_a_(this, textureGetter, p_217845_2_, p_217845_1_);
          this.field_217850_G.put(triple, ibakedmodel);
          return ibakedmodel;
       }
@@ -524,6 +551,10 @@
       return this.field_225367_M;
    }
 
+   public Set<ResourceLocation> getSpecialModels() {
+      return java.util.Collections.emptySet();
+   }
+
    @OnlyIn(Dist.CLIENT)
    static class BlockStateDefinitionException extends RuntimeException {
       public BlockStateDefinitionException(String p_i49526_1_) {
@@ -531,6 +562,10 @@
       }
    }
 
+   public SpriteMap getSpriteMap() {
+      return this.field_229322_z_;
+   }
+
    @OnlyIn(Dist.CLIENT)
    static class ModelListWrapper {
       private final List<IUnbakedModel> field_225339_a;
