From fe474d6b0541cc53dad90dc558d07885b49db02c Mon Sep 17 00:00:00 2001
From: PepperCode1 <44146161+PepperCode1@users.noreply.github.com>
Date: Sat, 8 Jun 2024 06:02:54 -0700
Subject: [PATCH] Port Model Loading API v1 to 1.21 (#3824)

* Port Model Loading API v1

* Fix checkstyle and address review comments
---
 fabric-model-loading-api-v1/build.gradle      |   1 +
 .../model/loading/v1/BlockStateResolver.java  |   6 +-
 .../loading/v1/DelegatingUnbakedModel.java    |   8 +-
 .../loading/v1/FabricBakedModelManager.java   |  14 +-
 .../model/loading/v1/ModelLoadingPlugin.java  |   8 +-
 .../model/loading/v1/ModelModifier.java       |  64 ++++--
 .../model/loading/v1/ModelResolver.java       |   6 +-
 ...riantProvider.java => BakerImplHooks.java} |  14 +-
 .../model/loading/BlockStatesLoaderHooks.java |  30 +++
 .../model/loading/ModelLoaderHooks.java       |  10 +-
 .../model/loading/ModelLoadingConstants.java  |  42 ++++
 .../loading/ModelLoadingEventDispatcher.java  | 200 ++++++-----------
 ...ava => ModelLoadingPluginContextImpl.java} |  82 +------
 .../model/loading/BakedModelManagerMixin.java |  20 +-
 .../model/loading/BlockStatesLoaderMixin.java |  51 +++++
 .../loading/ModelLoaderBakerImplMixin.java    |  32 ++-
 .../model/loading/ModelLoaderMixin.java       | 204 +++++++++---------
 .../fabric-model-loading-api-v1.mixins.json   |   1 +
 .../src/client/resources/fabric.mod.json      |   3 -
 .../model/loading/ModelTestModClient.java     |  28 ++-
 .../model/loading/NestedModelLoadingTest.java |  31 ++-
 .../model/loading/PreparablePluginTest.java   |  15 +-
 .../loading/SpecificModelReloadListener.java  |  20 +-
 fabric-renderer-api-v1/build.gradle           |   2 +-
 .../renderer/client/ModelResolverImpl.java    |  11 +-
 .../renderer/client/RendererClientTest.java   |   8 +-
 settings.gradle                               |   2 +-
 27 files changed, 469 insertions(+), 444 deletions(-)
 rename fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/{LegacyModelVariantProvider.java => BakerImplHooks.java} (66%)
 create mode 100644 fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BlockStatesLoaderHooks.java
 create mode 100644 fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingConstants.java
 rename fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/{ModelLoaderPluginContextImpl.java => ModelLoadingPluginContextImpl.java} (66%)
 create mode 100644 fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BlockStatesLoaderMixin.java

diff --git a/fabric-model-loading-api-v1/build.gradle b/fabric-model-loading-api-v1/build.gradle
index ca3eb5eb1..3d2819264 100644
--- a/fabric-model-loading-api-v1/build.gradle
+++ b/fabric-model-loading-api-v1/build.gradle
@@ -4,6 +4,7 @@ moduleDependencies(project, ['fabric-api-base'])
 
 testDependencies(project, [
 	':fabric-renderer-api-v1',
+	':fabric-renderer-indigo',
 	':fabric-rendering-v1',
 	':fabric-resource-loader-v0'
 ])
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/BlockStateResolver.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/BlockStateResolver.java
index dacd2f1a4..4f350e43d 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/BlockStateResolver.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/BlockStateResolver.java
@@ -22,7 +22,6 @@ import net.minecraft.block.Block;
 import net.minecraft.block.BlockState;
 import net.minecraft.client.render.model.ModelLoader;
 import net.minecraft.client.render.model.UnbakedModel;
-import net.minecraft.client.util.ModelIdentifier;
 import net.minecraft.util.Identifier;
 
 /**
@@ -72,7 +71,7 @@ public interface BlockStateResolver {
 		void setModel(BlockState state, UnbakedModel model);
 
 		/**
-		 * Loads a model using an {@link Identifier} or {@link ModelIdentifier}, or gets it if it was already loaded.
+		 * Loads a model using an {@link Identifier}, or gets it if it was already loaded.
 		 *
 		 * @param id the model identifier
 		 * @return the unbaked model, or a missing model if it is not present
@@ -81,9 +80,6 @@ public interface BlockStateResolver {
 
 		/**
 		 * The current model loader instance, which changes between resource reloads.
-		 *
-		 * <p>Do <b>not</b> call {@link ModelLoader#getOrLoadModel} as it does not supported nested model resolution;
-		 * use {@link #getOrLoadModel} from the context instead.
 		 */
 		ModelLoader loader();
 	}
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/DelegatingUnbakedModel.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/DelegatingUnbakedModel.java
index ee37263fd..081f5f6a0 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/DelegatingUnbakedModel.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/DelegatingUnbakedModel.java
@@ -27,7 +27,6 @@ import net.minecraft.client.render.model.Baker;
 import net.minecraft.client.render.model.ModelBakeSettings;
 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.util.Identifier;
 
@@ -43,7 +42,7 @@ public final class DelegatingUnbakedModel implements UnbakedModel {
 	/**
 	 * Constructs a new delegating model.
 	 *
-	 * @param delegate The identifier (can be a {@link ModelIdentifier}) of the underlying baked model.
+	 * @param delegate The identifier of the underlying baked model.
 	 */
 	public DelegatingUnbakedModel(Identifier delegate) {
 		this.delegate = delegate;
@@ -57,11 +56,12 @@ public final class DelegatingUnbakedModel implements UnbakedModel {
 
 	@Override
 	public void setParents(Function<Identifier, UnbakedModel> modelLoader) {
+		modelLoader.apply(delegate).setParents(modelLoader);
 	}
 
-	@Nullable
 	@Override
-	public BakedModel bake(Baker baker, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings rotationContainer, Identifier modelId) {
+	@Nullable
+	public BakedModel bake(Baker baker, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings rotationContainer) {
 		return baker.bake(delegate, rotationContainer);
 	}
 }
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/FabricBakedModelManager.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/FabricBakedModelManager.java
index dbc5e0c4c..e0b8caba0 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/FabricBakedModelManager.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/FabricBakedModelManager.java
@@ -16,8 +16,6 @@
 
 package net.fabricmc.fabric.api.client.model.loading.v1;
 
-import org.jetbrains.annotations.Nullable;
-
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.BakedModelManager;
 import net.minecraft.client.util.ModelIdentifier;
@@ -30,20 +28,18 @@ import net.minecraft.util.Identifier;
  */
 public interface FabricBakedModelManager {
 	/**
-	 * An alternative to {@link BakedModelManager#getModel(ModelIdentifier)} that accepts an
-	 * {@link Identifier} instead. Models loaded using {@link ModelLoadingPlugin.Context#addModels}
-	 * do not have a corresponding {@link ModelIdentifier}, so the vanilla method cannot be used to
-	 * retrieve them. The {@link Identifier} that was used to load them can be used in this method
-	 * to retrieve them.
+	 * Similar to {@link BakedModelManager#getModel(ModelIdentifier)}, but accepts an {@link Identifier} instead of a
+	 * {@link ModelIdentifier}. Use this method to retrieve models loaded using
+	 * {@link ModelLoadingPlugin.Context#addModels}, since those models do not have corresponding
+	 * {@link ModelIdentifier}s.
 	 *
 	 * <p><b>This method, as well as its vanilla counterpart, should only be used after the
 	 * {@link BakedModelManager} has completed reloading.</b> Otherwise, the result will be
-	 * outdated or null.
+	 * outdated or an exception will be thrown.
 	 *
 	 * @param id the id of the model
 	 * @return the model
 	 */
-	@Nullable
 	default BakedModel getModel(Identifier id) {
 		throw new UnsupportedOperationException("Implemented via mixin.");
 	}
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelLoadingPlugin.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelLoadingPlugin.java
index 002deb2fa..5d15c8428 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelLoadingPlugin.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelLoadingPlugin.java
@@ -22,7 +22,6 @@ import org.jetbrains.annotations.ApiStatus;
 
 import net.minecraft.block.Block;
 import net.minecraft.client.render.model.json.JsonUnbakedModel;
-import net.minecraft.client.util.ModelIdentifier;
 import net.minecraft.resource.ResourceManager;
 import net.minecraft.util.Identifier;
 
@@ -53,13 +52,14 @@ public interface ModelLoadingPlugin {
 	@ApiStatus.NonExtendable
 	interface Context {
 		/**
-		 * Adds one or more models (can be {@link ModelIdentifier}s) to the list of models that will be loaded and
-		 * baked.
+		 * Adds one or more models that will be loaded, baked, and made available through
+		 * {@link FabricBakedModelManager#getModel(Identifier)}.
 		 */
 		void addModels(Identifier... ids);
 
 		/**
-		 * Adds multiple models (can be {@link ModelIdentifier}s) to the list of models that will be loaded and baked.
+		 * Adds multiple models that will be loaded, baked, and made available through
+		 * {@link FabricBakedModelManager#getModel(Identifier)}.
 		 */
 		void addModels(Collection<? extends Identifier> ids);
 
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelModifier.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelModifier.java
index da6777996..0d8596b6d 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelModifier.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelModifier.java
@@ -20,6 +20,7 @@ import java.util.function.Function;
 
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.UnknownNullability;
 
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.Baker;
@@ -39,7 +40,7 @@ import net.fabricmc.fabric.api.event.Event;
  *
  * <p>Example use cases:
  * <ul>
- *     <li>Overriding a model for a particular block state - check if the given identifier is a {@link ModelIdentifier},
+ *     <li>Overriding a model for a particular block state - check if the given top-level identifier is not null,
  *     and then check if it has the appropriate variant for that block state. If so, return your desired model,
  *     otherwise return the given model.</li>
  *     <li>Wrapping a model to override certain behaviors - simply return a new model instance and delegate calls
@@ -91,16 +92,26 @@ public final class ModelModifier {
 		@ApiStatus.NonExtendable
 		interface Context {
 			/**
-			 * The identifier of this model (may be a {@link ModelIdentifier}).
+			 * Models with a resource ID are loaded directly from JSON or a {@link ModelModifier}.
 			 *
-			 * <p>For item models, only the {@link ModelIdentifier} with the {@code inventory} variant is passed, and
-			 * not the corresponding plain identifier.
+			 * @return the identifier of the given model as an {@link Identifier}, or null if {@link #topLevelId()} is
+			 * not null
 			 */
-			Identifier id();
+			@UnknownNullability("#topLevelId() != null")
+			Identifier resourceId();
 
 			/**
-			 * Loads a model using an {@link Identifier} or {@link ModelIdentifier}, or gets it if it was already
-			 * loaded.
+			 * Models with a top-level ID are loaded from blockstate files, {@link BlockStateResolver}s, or by copying
+			 * a previously loaded model.
+			 *
+			 * @return the identifier of the given model as a {@link ModelIdentifier}, or null if {@link #resourceId()}
+			 * is not null
+			 */
+			@UnknownNullability("#resourceId() != null")
+			ModelIdentifier topLevelId();
+
+			/**
+			 * Loads a model using an {@link Identifier}, or gets it if it was already loaded.
 			 *
 			 * @param id the model identifier
 			 * @return the unbaked model, or a missing model if it is not present
@@ -109,9 +120,6 @@ public final class ModelModifier {
 
 			/**
 			 * The current model loader instance, which changes between resource reloads.
-			 *
-			 * <p>Do <b>not</b> call {@link ModelLoader#getOrLoadModel} as it does not supported nested model
-			 * resolution; use {@link #getOrLoadModel} from the context instead.
 			 */
 			ModelLoader loader();
 		}
@@ -135,9 +143,23 @@ public final class ModelModifier {
 		@ApiStatus.NonExtendable
 		interface Context {
 			/**
-			 * The identifier of this model (may be a {@link ModelIdentifier}).
+			 * Models with a resource ID are loaded directly from JSON or a {@link ModelModifier}.
+			 *
+			 * @return the identifier of the given model as an {@link Identifier}, or null if {@link #topLevelId()} is
+			 * not null
 			 */
-			Identifier id();
+			@UnknownNullability("#topLevelId() != null")
+			Identifier resourceId();
+
+			/**
+			 * Models with a top-level ID are loaded from blockstate files, {@link BlockStateResolver}s, or by copying
+			 * a previously loaded model.
+			 *
+			 * @return the identifier of the given model as a {@link ModelIdentifier}, or null if {@link #resourceId()}
+			 * is not null
+			 */
+			@UnknownNullability("#resourceId() != null")
+			ModelIdentifier topLevelId();
 
 			/**
 			 * The function that can be used to retrieve sprites.
@@ -189,9 +211,23 @@ public final class ModelModifier {
 		@ApiStatus.NonExtendable
 		interface Context {
 			/**
-			 * The identifier of this model (may be a {@link ModelIdentifier}).
+			 * Models with a resource ID are loaded directly from JSON or a {@link ModelModifier}.
+			 *
+			 * @return the identifier of the given model as an {@link Identifier}, or null if {@link #topLevelId()} is
+			 * not null
 			 */
-			Identifier id();
+			@UnknownNullability("#topLevelId() != null")
+			Identifier resourceId();
+
+			/**
+			 * Models with a top-level ID are loaded from blockstate files, {@link BlockStateResolver}s, or by copying
+			 * a previously loaded model.
+			 *
+			 * @return the identifier of the given model as a {@link ModelIdentifier}, or null if {@link #resourceId()}
+			 * is not null
+			 */
+			@UnknownNullability("#resourceId() != null")
+			ModelIdentifier topLevelId();
 
 			/**
 			 * The unbaked model that is being baked.
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelResolver.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelResolver.java
index 8236af031..7f0843e2c 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelResolver.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/api/client/model/loading/v1/ModelResolver.java
@@ -21,7 +21,6 @@ import org.jetbrains.annotations.Nullable;
 
 import net.minecraft.client.render.model.ModelLoader;
 import net.minecraft.client.render.model.UnbakedModel;
-import net.minecraft.client.util.ModelIdentifier;
 import net.minecraft.util.Identifier;
 
 /**
@@ -64,7 +63,7 @@ public interface ModelResolver {
 		Identifier id();
 
 		/**
-		 * Loads a model using an {@link Identifier} or {@link ModelIdentifier}, or gets it if it was already loaded.
+		 * Loads a model using an {@link Identifier}, or gets it if it was already loaded.
 		 *
 		 * @param id the model identifier
 		 * @return the unbaked model, or a missing model if it is not present
@@ -73,9 +72,6 @@ public interface ModelResolver {
 
 		/**
 		 * The current model loader instance, which changes between resource reloads.
-		 *
-		 * <p>Do <b>not</b> call {@link ModelLoader#getOrLoadModel} as it does not supported nested model resolution;
-		 * use {@link #getOrLoadModel} from the context instead.
 		 */
 		ModelLoader loader();
 	}
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/LegacyModelVariantProvider.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BakerImplHooks.java
similarity index 66%
rename from fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/LegacyModelVariantProvider.java
rename to fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BakerImplHooks.java
index ecc5ddc0b..8d088dee1 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/LegacyModelVariantProvider.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BakerImplHooks.java
@@ -16,15 +16,11 @@
 
 package net.fabricmc.fabric.impl.client.model.loading;
 
-import org.jetbrains.annotations.Nullable;
+import java.util.function.Function;
 
-import net.minecraft.client.render.model.UnbakedModel;
-import net.minecraft.client.util.ModelIdentifier;
+import net.minecraft.client.texture.Sprite;
+import net.minecraft.client.util.SpriteIdentifier;
 
-/**
- * Legacy v0 bridge - remove if the legacy v0 module is removed.
- */
-public interface LegacyModelVariantProvider {
-	@Nullable
-	UnbakedModel loadModelVariant(ModelIdentifier modelId);
+public interface BakerImplHooks {
+	Function<SpriteIdentifier, Sprite> fabric_getTextureGetter();
 }
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BlockStatesLoaderHooks.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BlockStatesLoaderHooks.java
new file mode 100644
index 000000000..eda910323
--- /dev/null
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/BlockStatesLoaderHooks.java
@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.client.model.loading;
+
+import net.minecraft.block.Block;
+import net.minecraft.block.BlockState;
+import net.minecraft.state.StateManager;
+import net.minecraft.util.Identifier;
+
+public interface BlockStatesLoaderHooks {
+	void fabric_setLoadingOverride(LoadingOverride override);
+
+	interface LoadingOverride {
+		boolean loadBlockStates(Identifier id, StateManager<Block, BlockState> stateManager);
+	}
+}
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderHooks.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderHooks.java
index 67e1dfd26..434988699 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderHooks.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderHooks.java
@@ -17,7 +17,7 @@
 package net.fabricmc.fabric.impl.client.model.loading;
 
 import net.minecraft.client.render.model.UnbakedModel;
-import net.minecraft.client.render.model.json.JsonUnbakedModel;
+import net.minecraft.client.util.ModelIdentifier;
 import net.minecraft.util.Identifier;
 
 public interface ModelLoaderHooks {
@@ -27,11 +27,5 @@ public interface ModelLoaderHooks {
 
 	UnbakedModel fabric_getOrLoadModel(Identifier id);
 
-	void fabric_putModel(Identifier id, UnbakedModel model);
-
-	void fabric_putModelDirectly(Identifier id, UnbakedModel model);
-
-	void fabric_queueModelDependencies(UnbakedModel model);
-
-	JsonUnbakedModel fabric_loadModelFromJson(Identifier id);
+	void fabric_add(ModelIdentifier id, UnbakedModel model);
 }
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingConstants.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingConstants.java
new file mode 100644
index 000000000..7d27e12e0
--- /dev/null
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingConstants.java
@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.client.model.loading;
+
+import net.minecraft.client.util.ModelIdentifier;
+import net.minecraft.util.Identifier;
+
+public final class ModelLoadingConstants {
+	/**
+	 * This variant is used to convert user-provided Identifiers for extra models to ModelIdentifiers, since top-level
+	 * models that will be baked must have a ModelIdentifier. Models corresponding to the Identifiers will go through
+	 * ModelModifier.OnLoad, but models corresponding to the ModelIdentifiers will not.
+	 *
+	 * <p>This variant must be non-empty, must not contain "=", and must not be equal to "inventory" or "missingno".
+	 */
+	public static final String RESOURCE_SPECIAL_VARIANT = "fabric_resource";
+
+	private ModelLoadingConstants() {
+	}
+
+	public static ModelIdentifier toResourceModelId(Identifier id) {
+		return new ModelIdentifier(id, RESOURCE_SPECIAL_VARIANT);
+	}
+
+	public static boolean isResourceModelId(ModelIdentifier id) {
+		return id.variant().equals(RESOURCE_SPECIAL_VARIANT);
+	}
+}
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingEventDispatcher.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingEventDispatcher.java
index c0a62ce67..7d6a3dcc5 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingEventDispatcher.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingEventDispatcher.java
@@ -25,9 +25,8 @@ import com.google.common.collect.ImmutableList;
 import it.unimi.dsi.fastutil.objects.ObjectArrayList;
 import it.unimi.dsi.fastutil.objects.Reference2ReferenceMap;
 import it.unimi.dsi.fastutil.objects.Reference2ReferenceOpenHashMap;
-import it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet;
-import it.unimi.dsi.fastutil.objects.ReferenceSet;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.UnknownNullability;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -42,6 +41,7 @@ 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.state.StateManager;
 import net.minecraft.util.Identifier;
 
 import net.fabricmc.fabric.api.client.model.loading.v1.BlockStateResolver;
@@ -53,12 +53,10 @@ public class ModelLoadingEventDispatcher {
 	private static final Logger LOGGER = LoggerFactory.getLogger(ModelLoadingEventDispatcher.class);
 
 	private final ModelLoader loader;
-	private final ModelLoaderPluginContextImpl pluginContext;
+	private final ModelLoadingPluginContextImpl pluginContext;
 
 	private final ObjectArrayList<ModelResolverContext> modelResolverContextStack = new ObjectArrayList<>();
-
-	private final ObjectArrayList<BlockStateResolverContext> blockStateResolverContextStack = new ObjectArrayList<>();
-	private final ReferenceSet<Block> resolvingBlocks = new ReferenceOpenHashSet<>();
+	private final BlockStateResolverContext blockStateResolverContext = new BlockStateResolverContext();
 
 	private final ObjectArrayList<OnLoadModifierContext> onLoadModifierContextStack = new ObjectArrayList<>();
 	private final ObjectArrayList<BeforeBakeModifierContext> beforeBakeModifierContextStack = new ObjectArrayList<>();
@@ -66,7 +64,7 @@ public class ModelLoadingEventDispatcher {
 
 	public ModelLoadingEventDispatcher(ModelLoader loader, List<ModelLoadingPlugin> plugins) {
 		this.loader = loader;
-		this.pluginContext = new ModelLoaderPluginContextImpl(((ModelLoaderHooks) loader)::fabric_getOrLoadModel);
+		this.pluginContext = new ModelLoadingPluginContextImpl();
 
 		for (ModelLoadingPlugin plugin : plugins) {
 			try {
@@ -83,107 +81,19 @@ public class ModelLoadingEventDispatcher {
 		}
 	}
 
-	/**
-	 * @return {@code true} to cancel the vanilla method
-	 */
-	public boolean loadModel(Identifier id) {
-		if (id instanceof ModelIdentifier modelId) {
-			if ("inventory".equals(modelId.getVariant())) {
-				// We ALWAYS override the vanilla inventory model code path entirely, even for vanilla item models.
-				// See loadItemModel for an explanation.
-				loadItemModel(modelId);
-				return true;
-			} else {
-				// Prioritize block state resolver over legacy variant provider
-				BlockStateResolverHolder resolver = pluginContext.getBlockStateResolver(modelId);
+	public boolean loadBlockStateModels(Identifier id, StateManager<Block, BlockState> stateManager) {
+		BlockStateResolver resolver = pluginContext.blockStateResolvers.get(id);
 
-				if (resolver != null) {
-					loadBlockStateModels(resolver.resolver(), resolver.block(), resolver.blockId());
-					return true;
-				}
-
-				UnbakedModel legacyModel = legacyLoadModelVariant(modelId);
-
-				if (legacyModel != null) {
-					((ModelLoaderHooks) loader).fabric_putModel(id, legacyModel);
-					return true;
-				}
-
-				return false;
-			}
+		if (resolver != null) {
+			resolveBlockStates(resolver, stateManager.getOwner(), id);
+			return true;
 		} else {
-			UnbakedModel model = resolveModel(id);
-
-			if (model != null) {
-				((ModelLoaderHooks) loader).fabric_putModel(id, model);
-				return true;
-			}
-
 			return false;
 		}
 	}
 
-	@Nullable
-	private UnbakedModel legacyLoadModelVariant(ModelIdentifier modelId) {
-		return pluginContext.legacyVariantProviders().invoker().loadModelVariant(modelId);
-	}
-
-	/**
-	 * This function handles both modded item models and vanilla item models.
-	 * The vanilla code path for item models is never used.
-	 * See the long comment in the function for an explanation.
-	 */
-	private void loadItemModel(ModelIdentifier modelId) {
-		ModelLoaderHooks loaderHooks = (ModelLoaderHooks) loader;
-
-		Identifier id = modelId.withPrefixedPath("item/");
-
-		// Legacy variant provider
-		UnbakedModel model = legacyLoadModelVariant(modelId);
-
-		// Model resolver
-		if (model == null) {
-			model = resolveModel(id);
-		}
-
-		// Load from the vanilla code path otherwise.
-		if (model == null) {
-			model = loaderHooks.fabric_loadModelFromJson(id);
-		}
-
-		// This is a bit tricky:
-		// We have a single UnbakedModel now, but there are two identifiers:
-		// the ModelIdentifier (...#inventory) and the Identifier (...:item/...).
-		// So we call the on load modifier now and then directly add the model to the ModelLoader,
-		// reimplementing the behavior of ModelLoader#put.
-		// Calling ModelLoader#put is not an option as the model for the Identifier would not be replaced by an on load modifier.
-		// This is why we override the vanilla code path entirely.
-		model = modifyModelOnLoad(modelId, model);
-
-		loaderHooks.fabric_putModelDirectly(modelId, model);
-		loaderHooks.fabric_putModelDirectly(id, model);
-		loaderHooks.fabric_queueModelDependencies(model);
-	}
-
-	private void loadBlockStateModels(BlockStateResolver resolver, Block block, Identifier blockId) {
-		if (!resolvingBlocks.add(block)) {
-			throw new IllegalStateException("Circular reference while resolving models for block " + block);
-		}
-
-		try {
-			resolveBlockStates(resolver, block, blockId);
-		} finally {
-			resolvingBlocks.remove(block);
-		}
-	}
-
 	private void resolveBlockStates(BlockStateResolver resolver, Block block, Identifier blockId) {
-		// Get and prepare context
-		if (blockStateResolverContextStack.isEmpty()) {
-			blockStateResolverContextStack.add(new BlockStateResolverContext());
-		}
-
-		BlockStateResolverContext context = blockStateResolverContextStack.pop();
+		BlockStateResolverContext context = blockStateResolverContext;
 		context.prepare(block);
 
 		Reference2ReferenceMap<BlockState, UnbakedModel> resolvedModels = context.models;
@@ -204,14 +114,14 @@ public class ModelLoadingEventDispatcher {
 
 			for (BlockState state : allStates) {
 				ModelIdentifier modelId = BlockModels.getModelId(blockId, state);
-				((ModelLoaderHooks) loader).fabric_putModelDirectly(modelId, missingModel);
+				((ModelLoaderHooks) loader).fabric_add(modelId, missingModel);
 			}
 		} else if (resolvedModels.size() == allStates.size()) {
 			// If there are as many resolved models as total states, all states have
 			// been resolved and models do not need to be null-checked.
 			resolvedModels.forEach((state, model) -> {
 				ModelIdentifier modelId = BlockModels.getModelId(blockId, state);
-				((ModelLoaderHooks) loader).fabric_putModel(modelId, model);
+				((ModelLoaderHooks) loader).fabric_add(modelId, model);
 			});
 		} else {
 			UnbakedModel missingModel = ((ModelLoaderHooks) loader).fabric_getMissingModel();
@@ -223,21 +133,18 @@ public class ModelLoadingEventDispatcher {
 
 				if (model == null) {
 					LOGGER.error("Block state resolver did not provide a model for state {} in block {}. Using missing model.", state, block);
-					((ModelLoaderHooks) loader).fabric_putModelDirectly(modelId, missingModel);
+					((ModelLoaderHooks) loader).fabric_add(modelId, missingModel);
 				} else {
-					((ModelLoaderHooks) loader).fabric_putModel(modelId, model);
+					((ModelLoaderHooks) loader).fabric_add(modelId, model);
 				}
 			}
 		}
 
 		resolvedModels.clear();
-
-		// Store context for reuse
-		blockStateResolverContextStack.add(context);
 	}
 
 	@Nullable
-	private UnbakedModel resolveModel(Identifier id) {
+	public UnbakedModel resolveModel(Identifier id) {
 		if (modelResolverContextStack.isEmpty()) {
 			modelResolverContextStack.add(new ModelResolverContext());
 		}
@@ -251,13 +158,13 @@ public class ModelLoadingEventDispatcher {
 		return model;
 	}
 
-	public UnbakedModel modifyModelOnLoad(Identifier id, UnbakedModel model) {
+	public UnbakedModel modifyModelOnLoad(UnbakedModel model, @UnknownNullability Identifier resourceId, @UnknownNullability ModelIdentifier topLevelId) {
 		if (onLoadModifierContextStack.isEmpty()) {
 			onLoadModifierContextStack.add(new OnLoadModifierContext());
 		}
 
 		OnLoadModifierContext context = onLoadModifierContextStack.pop();
-		context.prepare(id);
+		context.prepare(resourceId, topLevelId);
 
 		model = pluginContext.modifyModelOnLoad().invoker().modifyModelOnLoad(model, context);
 
@@ -265,13 +172,13 @@ public class ModelLoadingEventDispatcher {
 		return model;
 	}
 
-	public UnbakedModel modifyModelBeforeBake(UnbakedModel model, Identifier id, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
+	public UnbakedModel modifyModelBeforeBake(UnbakedModel model, @UnknownNullability Identifier resourceId, @UnknownNullability ModelIdentifier topLevelId, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
 		if (beforeBakeModifierContextStack.isEmpty()) {
 			beforeBakeModifierContextStack.add(new BeforeBakeModifierContext());
 		}
 
 		BeforeBakeModifierContext context = beforeBakeModifierContextStack.pop();
-		context.prepare(id, textureGetter, settings, baker);
+		context.prepare(resourceId, topLevelId, textureGetter, settings, baker);
 
 		model = pluginContext.modifyModelBeforeBake().invoker().modifyModelBeforeBake(model, context);
 
@@ -280,13 +187,13 @@ public class ModelLoadingEventDispatcher {
 	}
 
 	@Nullable
-	public BakedModel modifyModelAfterBake(@Nullable BakedModel model, Identifier id, UnbakedModel sourceModel, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
+	public BakedModel modifyModelAfterBake(@Nullable BakedModel model, @UnknownNullability Identifier resourceId, @UnknownNullability ModelIdentifier topLevelId, UnbakedModel sourceModel, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
 		if (afterBakeModifierContextStack.isEmpty()) {
 			afterBakeModifierContextStack.add(new AfterBakeModifierContext());
 		}
 
 		AfterBakeModifierContext context = afterBakeModifierContextStack.pop();
-		context.prepare(id, sourceModel, textureGetter, settings, baker);
+		context.prepare(resourceId, topLevelId, sourceModel, textureGetter, settings, baker);
 
 		model = pluginContext.modifyModelAfterBake().invoker().modifyModelAfterBake(model, context);
 
@@ -357,15 +264,26 @@ public class ModelLoadingEventDispatcher {
 	}
 
 	private class OnLoadModifierContext implements ModelModifier.OnLoad.Context {
-		private Identifier id;
+		@UnknownNullability
+		private Identifier resourceId;
+		@UnknownNullability
+		private ModelIdentifier topLevelId;
 
-		private void prepare(Identifier id) {
-			this.id = id;
+		private void prepare(@UnknownNullability Identifier resourceId, @UnknownNullability ModelIdentifier topLevelId) {
+			this.resourceId = resourceId;
+			this.topLevelId = topLevelId;
 		}
 
 		@Override
-		public Identifier id() {
-			return id;
+		@UnknownNullability("#topLevelId() != null")
+		public Identifier resourceId() {
+			return resourceId;
+		}
+
+		@Override
+		@UnknownNullability("#resourceId() != null")
+		public ModelIdentifier topLevelId() {
+			return topLevelId;
 		}
 
 		@Override
@@ -380,21 +298,32 @@ public class ModelLoadingEventDispatcher {
 	}
 
 	private class BeforeBakeModifierContext implements ModelModifier.BeforeBake.Context {
-		private Identifier id;
+		@UnknownNullability
+		private Identifier resourceId;
+		@UnknownNullability
+		private ModelIdentifier topLevelId;
 		private Function<SpriteIdentifier, Sprite> textureGetter;
 		private ModelBakeSettings settings;
 		private Baker baker;
 
-		private void prepare(Identifier id, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
-			this.id = id;
+		private void prepare(@UnknownNullability Identifier resourceId, @UnknownNullability ModelIdentifier topLevelId, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
+			this.resourceId = resourceId;
+			this.topLevelId = topLevelId;
 			this.textureGetter = textureGetter;
 			this.settings = settings;
 			this.baker = baker;
 		}
 
 		@Override
-		public Identifier id() {
-			return id;
+		@UnknownNullability("#topLevelId() != null")
+		public Identifier resourceId() {
+			return resourceId;
+		}
+
+		@Override
+		@UnknownNullability("#resourceId() != null")
+		public ModelIdentifier topLevelId() {
+			return topLevelId;
 		}
 
 		@Override
@@ -419,14 +348,18 @@ public class ModelLoadingEventDispatcher {
 	}
 
 	private class AfterBakeModifierContext implements ModelModifier.AfterBake.Context {
-		private Identifier id;
+		@UnknownNullability
+		private Identifier resourceId;
+		@UnknownNullability
+		private ModelIdentifier topLevelId;
 		private UnbakedModel sourceModel;
 		private Function<SpriteIdentifier, Sprite> textureGetter;
 		private ModelBakeSettings settings;
 		private Baker baker;
 
-		private void prepare(Identifier id, UnbakedModel sourceModel, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
-			this.id = id;
+		private void prepare(@UnknownNullability Identifier resourceId, @UnknownNullability ModelIdentifier topLevelId, UnbakedModel sourceModel, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Baker baker) {
+			this.resourceId = resourceId;
+			this.topLevelId = topLevelId;
 			this.sourceModel = sourceModel;
 			this.textureGetter = textureGetter;
 			this.settings = settings;
@@ -434,8 +367,15 @@ public class ModelLoadingEventDispatcher {
 		}
 
 		@Override
-		public Identifier id() {
-			return id;
+		@UnknownNullability("#topLevelId() != null")
+		public Identifier resourceId() {
+			return resourceId;
+		}
+
+		@Override
+		@UnknownNullability("#resourceId() != null")
+		public ModelIdentifier topLevelId() {
+			return topLevelId;
 		}
 
 		@Override
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderPluginContextImpl.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingPluginContextImpl.java
similarity index 66%
rename from fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderPluginContextImpl.java
rename to fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingPluginContextImpl.java
index ba4ceec64..3231be50b 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoaderPluginContextImpl.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/impl/client/model/loading/ModelLoadingPluginContextImpl.java
@@ -23,15 +23,12 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.Function;
 
-import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import net.minecraft.block.Block;
 import net.minecraft.client.render.model.UnbakedModel;
-import net.minecraft.client.util.ModelIdentifier;
 import net.minecraft.registry.Registries;
 import net.minecraft.registry.RegistryKey;
 import net.minecraft.util.Identifier;
@@ -43,13 +40,11 @@ import net.fabricmc.fabric.api.client.model.loading.v1.ModelResolver;
 import net.fabricmc.fabric.api.event.Event;
 import net.fabricmc.fabric.api.event.EventFactory;
 
-public class ModelLoaderPluginContextImpl implements ModelLoadingPlugin.Context {
-	private static final Logger LOGGER = LoggerFactory.getLogger(ModelLoaderPluginContextImpl.class);
+public class ModelLoadingPluginContextImpl implements ModelLoadingPlugin.Context {
+	private static final Logger LOGGER = LoggerFactory.getLogger(ModelLoadingPluginContextImpl.class);
 
 	final Set<Identifier> extraModels = new LinkedHashSet<>();
-
-	private final Map<BlockKey, BlockStateResolverHolder> blockStateResolvers = new HashMap<>();
-	private final BlockKey lookupKey = new BlockKey();
+	final Map<Identifier, BlockStateResolver> blockStateResolvers = new HashMap<>();
 
 	private final Event<ModelResolver> modelResolvers = EventFactory.createArrayBacked(ModelResolver.class, resolvers -> context -> {
 		for (ModelResolver resolver : resolvers) {
@@ -103,13 +98,7 @@ public class ModelLoaderPluginContextImpl implements ModelLoadingPlugin.Context
 		return model;
 	}, MODEL_MODIFIER_PHASES);
 
-	/**
-	 * This field is used by the v0 wrapper to avoid constantly wrapping the context in hot code.
-	 */
-	public final Function<Identifier, UnbakedModel> modelGetter;
-
-	public ModelLoaderPluginContextImpl(Function<Identifier, UnbakedModel> modelGetter) {
-		this.modelGetter = modelGetter;
+	public ModelLoadingPluginContextImpl() {
 	}
 
 	@Override
@@ -136,23 +125,12 @@ public class ModelLoaderPluginContextImpl implements ModelLoadingPlugin.Context
 		}
 
 		Identifier blockId = optionalKey.get().getValue();
-		BlockKey key = new BlockKey(blockId.getNamespace(), blockId.getPath());
-		BlockStateResolverHolder holder = new BlockStateResolverHolder(resolver, block, blockId);
 
-		if (blockStateResolvers.put(key, holder) != null) {
-			throw new IllegalArgumentException("Duplicate block state resolver for block " + blockId);
+		if (blockStateResolvers.put(blockId, resolver) != null) {
+			throw new IllegalArgumentException("Duplicate block state resolver for " + block);
 		}
 	}
 
-	@Nullable
-	BlockStateResolverHolder getBlockStateResolver(ModelIdentifier modelId) {
-		BlockKey key = lookupKey;
-		key.namespace = modelId.getNamespace();
-		key.path = modelId.getPath();
-
-		return blockStateResolvers.get(key);
-	}
-
 	@Override
 	public Event<ModelResolver> resolveModel() {
 		return modelResolvers;
@@ -172,52 +150,4 @@ public class ModelLoaderPluginContextImpl implements ModelLoadingPlugin.Context
 	public Event<ModelModifier.AfterBake> modifyModelAfterBake() {
 		return afterBakeModifiers;
 	}
-
-	private static class BlockKey {
-		private String namespace;
-		private String path;
-
-		private BlockKey() {
-		}
-
-		private BlockKey(String namespace, String path) {
-			this.namespace = namespace;
-			this.path = path;
-		}
-
-		@Override
-		public boolean equals(Object o) {
-			if (this == o) return true;
-			if (o == null || getClass() != o.getClass()) return false;
-			BlockKey blockKey = (BlockKey) o;
-			return namespace.equals(blockKey.namespace) && path.equals(blockKey.path);
-		}
-
-		@Override
-		public int hashCode() {
-			return 31 * namespace.hashCode() + path.hashCode();
-		}
-	}
-
-	// Legacy v0 bridge - remove if the legacy v0 module is removed.
-
-	private final Event<LegacyModelVariantProvider> legacyVariantProviders = EventFactory.createArrayBacked(LegacyModelVariantProvider.class, providers -> modelId -> {
-		for (LegacyModelVariantProvider provider : providers) {
-			try {
-				UnbakedModel model = provider.loadModelVariant(modelId);
-
-				if (model != null) {
-					return model;
-				}
-			} catch (Exception exception) {
-				LOGGER.error("Failed to run legacy model variant provider", exception);
-			}
-		}
-
-		return null;
-	});
-
-	public Event<LegacyModelVariantProvider> legacyVariantProviders() {
-		return legacyVariantProviders;
-	}
 }
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BakedModelManagerMixin.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BakedModelManagerMixin.java
index f5c9b6488..27e7ce596 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BakedModelManagerMixin.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BakedModelManagerMixin.java
@@ -28,11 +28,12 @@ import org.spongepowered.asm.mixin.Shadow;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Redirect;
 
-import net.minecraft.class_9824;
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.BakedModelManager;
+import net.minecraft.client.render.model.BlockStatesLoader;
 import net.minecraft.client.render.model.ModelLoader;
 import net.minecraft.client.render.model.json.JsonUnbakedModel;
+import net.minecraft.client.util.ModelIdentifier;
 import net.minecraft.resource.ResourceManager;
 import net.minecraft.resource.ResourceReloader;
 import net.minecraft.util.Identifier;
@@ -41,16 +42,17 @@ import net.minecraft.util.profiler.Profiler;
 
 import net.fabricmc.fabric.api.client.model.loading.v1.FabricBakedModelManager;
 import net.fabricmc.fabric.api.client.model.loading.v1.ModelLoadingPlugin;
+import net.fabricmc.fabric.impl.client.model.loading.ModelLoadingConstants;
 import net.fabricmc.fabric.impl.client.model.loading.ModelLoadingPluginManager;
 
 @Mixin(BakedModelManager.class)
-public class BakedModelManagerMixin implements FabricBakedModelManager {
+abstract class BakedModelManagerMixin implements FabricBakedModelManager {
 	@Shadow
-	private Map<Identifier, BakedModel> models;
+	private Map<ModelIdentifier, BakedModel> models;
 
 	@Override
 	public BakedModel getModel(Identifier id) {
-		return models.get(id);
+		return models.get(ModelLoadingConstants.toResourceModelId(id));
 	}
 
 	@Redirect(
@@ -63,8 +65,8 @@ public class BakedModelManagerMixin implements FabricBakedModelManager {
 			allow = 1)
 	private CompletableFuture<ModelLoader> loadModelPluginData(
 			CompletableFuture<Map<Identifier, JsonUnbakedModel>> self,
-			CompletionStage<Map<Identifier, List<class_9824.SourceTrackedData>>> otherFuture,
-			BiFunction<Map<Identifier, JsonUnbakedModel>, Map<Identifier, List<class_9824.SourceTrackedData>>, ModelLoader> modelLoaderConstructor,
+			CompletionStage<Map<Identifier, List<BlockStatesLoader.SourceTrackedData>>> otherFuture,
+			BiFunction<Map<Identifier, JsonUnbakedModel>, Map<Identifier, List<BlockStatesLoader.SourceTrackedData>>, ModelLoader> modelLoaderConstructor,
 			Executor executor,
 			// reload args
 			ResourceReloader.Synchronizer synchronizer,
@@ -74,10 +76,12 @@ public class BakedModelManagerMixin implements FabricBakedModelManager {
 			Executor prepareExecutor,
 			Executor applyExecutor) {
 		CompletableFuture<List<ModelLoadingPlugin>> pluginsFuture = ModelLoadingPluginManager.preparePlugins(manager, prepareExecutor);
-		CompletableFuture<Pair<Map<Identifier, JsonUnbakedModel>, Map<Identifier, List<class_9824.SourceTrackedData>>>> pairFuture = self.thenCombine(otherFuture, Pair::new);
+		CompletableFuture<Pair<Map<Identifier, JsonUnbakedModel>, Map<Identifier, List<BlockStatesLoader.SourceTrackedData>>>> pairFuture = self.thenCombine(otherFuture, Pair::new);
 		return pairFuture.thenCombineAsync(pluginsFuture, (pair, plugins) -> {
 			ModelLoadingPluginManager.CURRENT_PLUGINS.set(plugins);
-			return modelLoaderConstructor.apply(pair.getLeft(), pair.getRight());
+			ModelLoader modelLoader = modelLoaderConstructor.apply(pair.getLeft(), pair.getRight());
+			ModelLoadingPluginManager.CURRENT_PLUGINS.remove();
+			return modelLoader;
 		}, executor);
 	}
 }
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BlockStatesLoaderMixin.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BlockStatesLoaderMixin.java
new file mode 100644
index 000000000..426e25378
--- /dev/null
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/BlockStatesLoaderMixin.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.client.model.loading;
+
+import org.jetbrains.annotations.Nullable;
+import org.spongepowered.asm.mixin.Mixin;
+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;
+
+import net.minecraft.block.Block;
+import net.minecraft.block.BlockState;
+import net.minecraft.client.render.model.BlockStatesLoader;
+import net.minecraft.state.StateManager;
+import net.minecraft.util.Identifier;
+
+import net.fabricmc.fabric.impl.client.model.loading.BlockStatesLoaderHooks;
+
+@Mixin(BlockStatesLoader.class)
+abstract class BlockStatesLoaderMixin implements BlockStatesLoaderHooks {
+	@Unique
+	@Nullable
+	private LoadingOverride loadingOverride;
+
+	@Inject(method = "loadBlockStates(Lnet/minecraft/util/Identifier;Lnet/minecraft/state/StateManager;)V", at = @At("HEAD"), cancellable = true)
+	private void onHeadLoadBlockStates(Identifier id, StateManager<Block, BlockState> stateManager, CallbackInfo ci) {
+		if (loadingOverride != null && loadingOverride.loadBlockStates(id, stateManager)) {
+			ci.cancel();
+		}
+	}
+
+	@Override
+	public void fabric_setLoadingOverride(LoadingOverride override) {
+		loadingOverride = override;
+	}
+}
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderBakerImplMixin.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderBakerImplMixin.java
index ef55ae900..9e207c323 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderBakerImplMixin.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderBakerImplMixin.java
@@ -18,28 +18,29 @@ package net.fabricmc.fabric.mixin.client.model.loading;
 
 import java.util.function.Function;
 
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
 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.ModifyVariable;
-import org.spongepowered.asm.mixin.injection.Redirect;
+import org.spongepowered.asm.mixin.injection.Coerce;
 
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.client.render.model.Baker;
 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.render.model.json.JsonUnbakedModel;
 import net.minecraft.client.texture.Sprite;
 import net.minecraft.client.util.SpriteIdentifier;
 import net.minecraft.util.Identifier;
 
+import net.fabricmc.fabric.impl.client.model.loading.BakerImplHooks;
 import net.fabricmc.fabric.impl.client.model.loading.ModelLoaderHooks;
 import net.fabricmc.fabric.impl.client.model.loading.ModelLoadingEventDispatcher;
 
 @Mixin(targets = "net/minecraft/client/render/model/ModelLoader$BakerImpl")
-public class ModelLoaderBakerImplMixin {
+abstract class ModelLoaderBakerImplMixin implements BakerImplHooks {
 	@Shadow
 	@Final
 	private ModelLoader field_40571;
@@ -47,23 +48,16 @@ public class ModelLoaderBakerImplMixin {
 	@Final
 	private Function<SpriteIdentifier, Sprite> textureGetter;
 
-	@ModifyVariable(method = "bake", at = @At(value = "INVOKE_ASSIGN", target = "Lnet/minecraft/client/render/model/ModelLoader$BakerImpl;getOrLoadModel(Lnet/minecraft/util/Identifier;)Lnet/minecraft/client/render/model/UnbakedModel;"))
-	private UnbakedModel invokeModifyBeforeBake(UnbakedModel model, Identifier id, ModelBakeSettings settings) {
+	@WrapOperation(method = "bake(Lnet/minecraft/util/Identifier;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/ModelLoader$BakerImpl;bake(Lnet/minecraft/client/render/model/UnbakedModel;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;"))
+	private BakedModel wrapInnerBake(@Coerce Baker self, UnbakedModel unbakedModel, ModelBakeSettings settings, Operation<BakedModel> operation, Identifier id) {
 		ModelLoadingEventDispatcher dispatcher = ((ModelLoaderHooks) this.field_40571).fabric_getDispatcher();
-		return dispatcher.modifyModelBeforeBake(model, id, textureGetter, settings, (Baker) this);
+		unbakedModel = dispatcher.modifyModelBeforeBake(unbakedModel, id, null, textureGetter, settings, self);
+		BakedModel model = operation.call(self, unbakedModel, settings);
+		return dispatcher.modifyModelAfterBake(model, id, null, unbakedModel, textureGetter, settings, self);
 	}
 
-	@Redirect(method = "bake", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/UnbakedModel;bake(Lnet/minecraft/client/render/model/Baker;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Lnet/minecraft/util/Identifier;)Lnet/minecraft/client/render/model/BakedModel;"))
-	private BakedModel invokeModifyAfterBake(UnbakedModel unbakedModel, Baker baker, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Identifier id) {
-		BakedModel model = unbakedModel.bake(baker, textureGetter, settings, id);
-		ModelLoadingEventDispatcher dispatcher = ((ModelLoaderHooks) this.field_40571).fabric_getDispatcher();
-		return dispatcher.modifyModelAfterBake(model, id, unbakedModel, textureGetter, settings, baker);
-	}
-
-	@Redirect(method = "bake", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/json/JsonUnbakedModel;bake(Lnet/minecraft/client/render/model/Baker;Lnet/minecraft/client/render/model/json/JsonUnbakedModel;Ljava/util/function/Function;Lnet/minecraft/client/render/model/ModelBakeSettings;Lnet/minecraft/util/Identifier;Z)Lnet/minecraft/client/render/model/BakedModel;"))
-	private BakedModel invokeModifyAfterBake(JsonUnbakedModel unbakedModel, Baker baker, JsonUnbakedModel parent, Function<SpriteIdentifier, Sprite> textureGetter, ModelBakeSettings settings, Identifier id, boolean hasDepth) {
-		BakedModel model = unbakedModel.bake(baker, parent, textureGetter, settings, id, hasDepth);
-		ModelLoadingEventDispatcher dispatcher = ((ModelLoaderHooks) this.field_40571).fabric_getDispatcher();
-		return dispatcher.modifyModelAfterBake(model, id, unbakedModel, textureGetter, settings, baker);
+	@Override
+	public Function<SpriteIdentifier, Sprite> fabric_getTextureGetter() {
+		return textureGetter;
 	}
 }
diff --git a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderMixin.java b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderMixin.java
index 5c2231825..fad4c0961 100644
--- a/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderMixin.java
+++ b/fabric-model-loading-api-v1/src/client/java/net/fabricmc/fabric/mixin/client/model/loading/ModelLoaderMixin.java
@@ -19,36 +19,48 @@ package net.fabricmc.fabric.mixin.client.model.loading;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
+import com.llamalad7.mixinextras.sugar.Local;
+import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
 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.Coerce;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.ModifyVariable;
+import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
-import net.minecraft.class_9824;
 import net.minecraft.client.color.block.BlockColors;
+import net.minecraft.client.render.model.BakedModel;
+import net.minecraft.client.render.model.Baker;
+import net.minecraft.client.render.model.BlockStatesLoader;
+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.render.model.json.JsonUnbakedModel;
+import net.minecraft.client.texture.Sprite;
 import net.minecraft.client.util.ModelIdentifier;
+import net.minecraft.client.util.SpriteIdentifier;
 import net.minecraft.util.Identifier;
 import net.minecraft.util.profiler.Profiler;
 
+import net.fabricmc.fabric.impl.client.model.loading.BakerImplHooks;
+import net.fabricmc.fabric.impl.client.model.loading.BlockStatesLoaderHooks;
 import net.fabricmc.fabric.impl.client.model.loading.ModelLoaderHooks;
+import net.fabricmc.fabric.impl.client.model.loading.ModelLoadingConstants;
 import net.fabricmc.fabric.impl.client.model.loading.ModelLoadingEventDispatcher;
 import net.fabricmc.fabric.impl.client.model.loading.ModelLoadingPluginManager;
 
 @Mixin(ModelLoader.class)
-public abstract class ModelLoaderMixin implements ModelLoaderHooks {
-	// The missing model is always loaded and added first.
-	@Final
-	@Shadow
-	public static ModelIdentifier MISSING_ID;
+abstract class ModelLoaderMixin implements ModelLoaderHooks {
 	@Final
 	@Shadow
 	private Set<Identifier> modelsToLoad;
@@ -57,94 +69,119 @@ public abstract class ModelLoaderMixin implements ModelLoaderHooks {
 	private Map<Identifier, UnbakedModel> unbakedModels;
 	@Shadow
 	@Final
-	private Map<Identifier, UnbakedModel> modelsToBake;
+	private Map<ModelIdentifier, UnbakedModel> modelsToBake;
+	@Shadow
+	@Final
+	private UnbakedModel missingModel;
 
 	@Unique
 	private ModelLoadingEventDispatcher fabric_eventDispatcher;
-	// Explicitly not @Unique to allow mods that heavily rework model loading to reimplement the guard.
-	// Note that this is an implementation detail; it can change at any time.
-	private int fabric_guardGetOrLoadModel = 0;
-	private boolean fabric_enableGetOrLoadModelGuard = true;
+	@Unique
+	private final ObjectLinkedOpenHashSet<Identifier> modelLoadingStack = new ObjectLinkedOpenHashSet<>();
 
 	@Shadow
-	private void addModel(ModelIdentifier id) {
-	}
+	abstract UnbakedModel getOrLoadModel(Identifier id);
 
 	@Shadow
-	public abstract UnbakedModel getOrLoadModel(Identifier id);
+	abstract void add(ModelIdentifier id, UnbakedModel unbakedModel);
 
 	@Shadow
-	private void loadModel(Identifier id) {
-	}
+	abstract JsonUnbakedModel loadModelFromJson(Identifier id);
 
-	@Shadow
-	private void method_61076(ModelIdentifier id, UnbakedModel unbakedModel) {
-	}
-
-	@Shadow
-	public abstract JsonUnbakedModel loadModelFromJson(Identifier id);
-
-	@Inject(method = "<init>", at = @At(value = "INVOKE", target = "net/minecraft/util/profiler/Profiler.swap(Ljava/lang/String;)V", ordinal = 0))
-	private void afterMissingModelInit(BlockColors blockColors, Profiler profiler, Map<Identifier, JsonUnbakedModel> jsonUnbakedModels, Map<Identifier, List<class_9824.SourceTrackedData>> blockStates, CallbackInfo info) {
+	@Inject(method = "<init>", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/BlockStatesLoader;load()V"))
+	private void afterMissingModelInit(BlockColors blockColors, Profiler profiler, Map<Identifier, JsonUnbakedModel> jsonUnbakedModels, Map<Identifier, List<BlockStatesLoader.SourceTrackedData>> blockStates, CallbackInfo info, @Local BlockStatesLoader blockStatesLoader) {
 		// Sanity check
-		if (!unbakedModels.containsKey(MISSING_ID)) {
+		if (missingModel == null || !modelsToBake.containsKey(ModelLoader.MISSING_MODEL_ID)) {
 			throw new AssertionError("Missing model not initialized. This is likely a Fabric API porting bug.");
 		}
 
+		// Add the missing model to the cache since vanilla doesn't. Mods may load/bake the missing model directly.
+		unbakedModels.put(ModelLoader.MISSING_ID, missingModel);
 		profiler.swap("fabric_plugins_init");
 
 		fabric_eventDispatcher = new ModelLoadingEventDispatcher((ModelLoader) (Object) this, ModelLoadingPluginManager.CURRENT_PLUGINS.get());
-		ModelLoadingPluginManager.CURRENT_PLUGINS.remove();
-		fabric_eventDispatcher.addExtraModels(this::addModel);
+		fabric_eventDispatcher.addExtraModels(this::addExtraModel);
+		((BlockStatesLoaderHooks) blockStatesLoader).fabric_setLoadingOverride(fabric_eventDispatcher::loadBlockStateModels);
 	}
 
 	@Unique
-	private void addModel(Identifier id) {
-		if (id instanceof ModelIdentifier) {
-			addModel((ModelIdentifier) id);
-		} else {
-			// The vanilla addModel method is arbitrarily limited to ModelIdentifiers,
-			// but it's useful to tell the game to just load and bake a direct model path as well.
-			// Replicate the vanilla logic of addModel here.
-			UnbakedModel unbakedModel = getOrLoadModel(id);
-			this.unbakedModels.put(id, unbakedModel);
-			this.modelsToBake.put(id, unbakedModel);
-		}
+	private void addExtraModel(Identifier id) {
+		ModelIdentifier modelId = ModelLoadingConstants.toResourceModelId(id);
+		UnbakedModel unbakedModel = getOrLoadModel(id);
+		add(modelId, unbakedModel);
 	}
 
-	@Inject(method = "getOrLoadModel", at = @At("HEAD"))
-	private void fabric_preventNestedGetOrLoadModel(Identifier id, CallbackInfoReturnable<UnbakedModel> cir) {
-		if (fabric_enableGetOrLoadModelGuard && fabric_guardGetOrLoadModel > 0) {
-			throw new IllegalStateException("ModelLoader#getOrLoadModel called from a ModelResolver or ModelModifier.OnBake instance. This is not allowed to prevent errors during model loading. Use getOrLoadModel from the context instead.");
-		}
-	}
-
-	@Inject(method = "loadModel", at = @At("HEAD"), cancellable = true)
-	private void onLoadModel(Identifier id, CallbackInfo ci) {
-		// Prevent calls to getOrLoadModel from loadModel as it will cause problems.
-		// Mods should call getOrLoadModel on the ModelResolver.Context instead.
-		fabric_guardGetOrLoadModel++;
-
-		try {
-			if (fabric_eventDispatcher.loadModel(id)) {
-				ci.cancel();
+	@Inject(method = "getOrLoadModel", at = @At("HEAD"), cancellable = true)
+	private void allowRecursiveLoading(Identifier id, CallbackInfoReturnable<UnbakedModel> cir) {
+		// If the stack is empty, this is the top-level call, so proceed as normal.
+		if (!modelLoadingStack.isEmpty()) {
+			if (unbakedModels.containsKey(id)) {
+				cir.setReturnValue(unbakedModels.get(id));
+			} else if (modelLoadingStack.contains(id)) {
+				throw new IllegalStateException("Circular reference while loading model '" + id + "' (" + modelLoadingStack.stream().map(i -> i + "->").collect(Collectors.joining()) + id + ")");
+			} else {
+				UnbakedModel model = loadModel(id);
+				unbakedModels.put(id, model);
+				// These will be loaded at the top-level call.
+				modelsToLoad.addAll(model.getModelDependencies());
+				cir.setReturnValue(model);
 			}
-		} finally {
-			fabric_guardGetOrLoadModel--;
 		}
 	}
 
-	@ModifyVariable(method = "method_61076", at = @At("HEAD"), argsOnly = true)
-	private UnbakedModel onPutModel(UnbakedModel model, ModelIdentifier id) {
-		fabric_guardGetOrLoadModel++;
+	// This is the call that needs to be redirected to support ModelResolvers, but it returns a JsonUnbakedModel.
+	// Redirect it to always return null and handle the logic in a ModifyVariable right after the call.
+	@Redirect(method = "getOrLoadModel", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/ModelLoader;loadModelFromJson(Lnet/minecraft/util/Identifier;)Lnet/minecraft/client/render/model/json/JsonUnbakedModel;"))
+	private JsonUnbakedModel cancelLoadModelFromJson(ModelLoader self, Identifier id) {
+		return null;
+	}
+
+	@ModifyVariable(method = "getOrLoadModel", at = @At(value = "INVOKE_ASSIGN", target = "Lnet/minecraft/client/render/model/ModelLoader;loadModelFromJson(Lnet/minecraft/util/Identifier;)Lnet/minecraft/client/render/model/json/JsonUnbakedModel;"))
+	private UnbakedModel doLoadModel(UnbakedModel model, @Local(ordinal = 1) Identifier id) {
+		return loadModel(id);
+	}
+
+	@Unique
+	private UnbakedModel loadModel(Identifier id) {
+		modelLoadingStack.add(id);
 
 		try {
-			return fabric_eventDispatcher.modifyModelOnLoad(id, model);
+			UnbakedModel model = fabric_eventDispatcher.resolveModel(id);
+
+			if (model == null) {
+				model = loadModelFromJson(id);
+			}
+
+			return fabric_eventDispatcher.modifyModelOnLoad(model, id, null);
 		} finally {
-			fabric_guardGetOrLoadModel--;
+			modelLoadingStack.removeLast();
 		}
 	}
 
+	@ModifyVariable(method = "add", at = @At("HEAD"), argsOnly = true)
+	private UnbakedModel onAdd(UnbakedModel model, ModelIdentifier id) {
+		if (ModelLoadingConstants.isResourceModelId(id)) {
+			return model;
+		}
+
+		return fabric_eventDispatcher.modifyModelOnLoad(model, null, id);
+	}
+
+	@WrapOperation(method = "method_61072(Lnet/minecraft/client/render/model/ModelLoader$SpriteGetter;Lnet/minecraft/client/util/ModelIdentifier;Lnet/minecraft/client/render/model/UnbakedModel;)V", at = @At(value = "INVOKE", target = "Lnet/minecraft/client/render/model/ModelLoader$BakerImpl;bake(Lnet/minecraft/client/render/model/UnbakedModel;Lnet/minecraft/client/render/model/ModelBakeSettings;)Lnet/minecraft/client/render/model/BakedModel;"))
+	private BakedModel wrapSingleOuterBake(@Coerce Baker baker, UnbakedModel unbakedModel, ModelBakeSettings settings, Operation<BakedModel> operation, ModelLoader.SpriteGetter spriteGetter, ModelIdentifier id) {
+		if (ModelLoadingConstants.isResourceModelId(id) || id.equals(ModelLoader.MISSING_MODEL_ID)) {
+			// Call the baker instead of the operation to ensure the baked model is cached and doesn't end up going
+			// through events twice.
+			// This ignores the UnbakedModel in modelsToBake but it should be the same as the one in unbakedModels.
+			return baker.bake(id.id(), settings);
+		}
+
+		Function<SpriteIdentifier, Sprite> textureGetter = ((BakerImplHooks) baker).fabric_getTextureGetter();
+		unbakedModel = fabric_eventDispatcher.modifyModelBeforeBake(unbakedModel, null, id, textureGetter, settings, baker);
+		BakedModel model = operation.call(baker, unbakedModel, settings);
+		return fabric_eventDispatcher.modifyModelAfterBake(model, null, id, unbakedModel, textureGetter, settings, baker);
+	}
+
 	@Override
 	public ModelLoadingEventDispatcher fabric_getDispatcher() {
 		return fabric_eventDispatcher;
@@ -152,51 +189,16 @@ public abstract class ModelLoaderMixin implements ModelLoaderHooks {
 
 	@Override
 	public UnbakedModel fabric_getMissingModel() {
-		return unbakedModels.get(MISSING_ID);
+		return missingModel;
 	}
 
-	/**
-	 * Unlike getOrLoadModel, this method supports nested model loading.
-	 *
-	 * <p>Vanilla does not due to the iteration over modelsToLoad which causes models to be resolved multiple times,
-	 * possibly leading to crashes.
-	 */
 	@Override
 	public UnbakedModel fabric_getOrLoadModel(Identifier id) {
-		if (this.unbakedModels.containsKey(id)) {
-			return this.unbakedModels.get(id);
-		}
-
-		if (!modelsToLoad.add(id)) {
-			throw new IllegalStateException("Circular reference while loading " + id);
-		}
-
-		try {
-			loadModel(id);
-		} finally {
-			modelsToLoad.remove(id);
-		}
-
-		return unbakedModels.get(id);
+		return getOrLoadModel(id);
 	}
 
 	@Override
-	public void fabric_putModel(Identifier id, UnbakedModel model) {
-		putModel(id, model);
-	}
-
-	@Override
-	public void fabric_putModelDirectly(Identifier id, UnbakedModel model) {
-		unbakedModels.put(id, model);
-	}
-
-	@Override
-	public void fabric_queueModelDependencies(UnbakedModel model) {
-		modelsToLoad.addAll(model.getModelDependencies());
-	}
-
-	@Override
-	public JsonUnbakedModel fabric_loadModelFromJson(Identifier id) {
-		return loadModelFromJson(id);
+	public void fabric_add(ModelIdentifier id, UnbakedModel model) {
+		add(id, model);
 	}
 }
diff --git a/fabric-model-loading-api-v1/src/client/resources/fabric-model-loading-api-v1.mixins.json b/fabric-model-loading-api-v1/src/client/resources/fabric-model-loading-api-v1.mixins.json
index 133a6b442..e59c298da 100644
--- a/fabric-model-loading-api-v1/src/client/resources/fabric-model-loading-api-v1.mixins.json
+++ b/fabric-model-loading-api-v1/src/client/resources/fabric-model-loading-api-v1.mixins.json
@@ -4,6 +4,7 @@
   "compatibilityLevel": "JAVA_17",
   "client": [
     "BakedModelManagerMixin",
+    "BlockStatesLoaderMixin",
     "ModelLoaderMixin",
     "ModelLoaderBakerImplMixin"
   ],
diff --git a/fabric-model-loading-api-v1/src/client/resources/fabric.mod.json b/fabric-model-loading-api-v1/src/client/resources/fabric.mod.json
index dfb3e1cab..01169141c 100644
--- a/fabric-model-loading-api-v1/src/client/resources/fabric.mod.json
+++ b/fabric-model-loading-api-v1/src/client/resources/fabric.mod.json
@@ -19,9 +19,6 @@
     "fabricloader": ">=0.15.11",
     "fabric-api-base": "*"
   },
-  "breaks": {
-    "fabric-models-v0": "<0.4.0"
-  },
   "description": "Provides hooks for model loading.",
   "mixins": [
     {
diff --git a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/ModelTestModClient.java b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/ModelTestModClient.java
index 11687a9f7..ea1649b4d 100644
--- a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/ModelTestModClient.java
+++ b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/ModelTestModClient.java
@@ -47,7 +47,9 @@ import net.fabricmc.fabric.api.resource.ResourceManagerHelper;
 public class ModelTestModClient implements ClientModInitializer {
 	public static final String ID = "fabric-model-loading-api-v1-testmod";
 
-	public static final Identifier MODEL_ID = Identifier.of(ID, "half_red_sand");
+	public static final Identifier HALF_RED_SAND_MODEL_ID = id("half_red_sand");
+	public static final Identifier GOLD_BLOCK_MODEL_ID = Identifier.ofVanilla("block/gold_block");
+	public static final Identifier BROWN_GLAZED_TERRACOTTA_MODEL_ID = Identifier.ofVanilla("block/brown_glazed_terracotta");
 
 	static class DownQuadRemovingModel extends ForwardingBakedModel {
 		DownQuadRemovingModel(BakedModel model) {
@@ -65,19 +67,23 @@ public class ModelTestModClient implements ClientModInitializer {
 	@Override
 	public void onInitializeClient() {
 		ModelLoadingPlugin.register(pluginContext -> {
-			pluginContext.addModels(MODEL_ID);
+			pluginContext.addModels(HALF_RED_SAND_MODEL_ID);
 			// remove bottom face of gold blocks
 			pluginContext.modifyModelAfterBake().register(ModelModifier.WRAP_PHASE, (model, context) -> {
-				if (context.id().getPath().equals("block/gold_block")) {
+				Identifier id = context.resourceId();
+
+				if (id != null && id.equals(GOLD_BLOCK_MODEL_ID)) {
 					return new DownQuadRemovingModel(model);
-				} else {
-					return model;
 				}
+
+				return model;
 			});
 			// make fences with west: true and everything else false appear to be a missing model visually
 			ModelIdentifier fenceId = BlockModels.getModelId(Blocks.OAK_FENCE.getDefaultState().with(HorizontalConnectingBlock.WEST, true));
 			pluginContext.modifyModelOnLoad().register(ModelModifier.OVERRIDE_PHASE, (model, context) -> {
-				if (fenceId.equals(context.id())) {
+				ModelIdentifier id = context.topLevelId();
+
+				if (id != null && id.equals(fenceId)) {
 					return context.getOrLoadModel(ModelLoader.MISSING_ID);
 				}
 
@@ -86,8 +92,10 @@ public class ModelTestModClient implements ClientModInitializer {
 			// make brown glazed terracotta appear to be a missing model visually, but without affecting the item, by using pre-bake
 			// using load here would make the item also appear missing
 			pluginContext.modifyModelBeforeBake().register(ModelModifier.OVERRIDE_PHASE, (model, context) -> {
-				if (context.id().getPath().equals("block/brown_glazed_terracotta")) {
-					return context.loader().getOrLoadModel(ModelLoader.MISSING_ID);
+				Identifier id = context.resourceId();
+
+				if (id != null && id.equals(BROWN_GLAZED_TERRACOTTA_MODEL_ID)) {
+					return context.baker().getOrLoadModel(ModelLoader.MISSING_ID);
 				}
 
 				return model;
@@ -119,4 +127,8 @@ public class ModelTestModClient implements ClientModInitializer {
 			}
 		});
 	}
+
+	public static Identifier id(String path) {
+		return Identifier.of(ID, path);
+	}
 }
diff --git a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/NestedModelLoadingTest.java b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/NestedModelLoadingTest.java
index 862a116aa..8ae600a2e 100644
--- a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/NestedModelLoadingTest.java
+++ b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/NestedModelLoadingTest.java
@@ -16,6 +16,8 @@
 
 package net.fabricmc.fabric.test.model.loading;
 
+import static net.fabricmc.fabric.test.model.loading.ModelTestModClient.id;
+
 import com.mojang.logging.LogUtils;
 import org.slf4j.Logger;
 
@@ -32,10 +34,6 @@ import net.fabricmc.fabric.api.client.model.loading.v1.ModelLoadingPlugin;
 public class NestedModelLoadingTest implements ClientModInitializer {
 	private static final Logger LOGGER = LogUtils.getLogger();
 
-	private static Identifier id(String path) {
-		return Identifier.of("fabric-model-loading-api-v1-testmod", path);
-	}
-
 	private static final Identifier BASE_MODEL = id("nested_base");
 	private static final Identifier NESTED_MODEL_1 = id("nested_1");
 	private static final Identifier NESTED_MODEL_2 = id("nested_2");
@@ -84,21 +82,22 @@ public class NestedModelLoadingTest implements ClientModInitializer {
 
 			pluginContext.modifyModelOnLoad().register((model, context) -> {
 				UnbakedModel ret = model;
+				Identifier id = context.resourceId();
 
-				if (context.id().equals(NESTED_MODEL_3)) {
-					Identifier id = context.id();
+				if (id != null) {
+					if (id.equals(NESTED_MODEL_3)) {
+						LOGGER.info("   Nested model 4 started loading");
+						ret = context.getOrLoadModel(NESTED_MODEL_4);
+						LOGGER.info("   Nested model 4 finished loading");
 
-					LOGGER.info("   Nested model 4 started loading");
-					ret = context.getOrLoadModel(NESTED_MODEL_4);
-					LOGGER.info("   Nested model 4 finished loading");
-
-					if (!id.equals(context.id())) {
-						throw new AssertionError("Context object should not have changed.");
+						if (!id.equals(context.resourceId())) {
+							throw new AssertionError("Context object should not have changed.");
+						}
+					} else if (id.equals(NESTED_MODEL_4)) {
+						LOGGER.info("    Nested model 5 started loading");
+						ret = context.getOrLoadModel(NESTED_MODEL_5);
+						LOGGER.info("    Nested model 5 finished loading");
 					}
-				} else if (context.id().equals(NESTED_MODEL_4)) {
-					LOGGER.info("    Nested model 5 started loading");
-					ret = context.getOrLoadModel(NESTED_MODEL_5);
-					LOGGER.info("    Nested model 5 finished loading");
 				}
 
 				return ret;
diff --git a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/PreparablePluginTest.java b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/PreparablePluginTest.java
index eb4ec3331..f2af763a6 100644
--- a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/PreparablePluginTest.java
+++ b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/PreparablePluginTest.java
@@ -55,9 +55,18 @@ public class PreparablePluginTest implements ClientModInitializer {
 	public void onInitializeClient() {
 		PreparableModelLoadingPlugin.register(PreparablePluginTest::loadModelReplacements, (replacementModels, pluginContext) -> {
 			pluginContext.modifyModelOnLoad().register((model, ctx) -> {
-				@Nullable
-				UnbakedModel replacementModel = replacementModels.get(ctx.id());
-				return replacementModel == null ? model : replacementModel;
+				Identifier id = ctx.resourceId();
+
+				if (id != null) {
+					@Nullable
+					UnbakedModel replacementModel = replacementModels.get(id);
+
+					if (replacementModel != null) {
+						return replacementModel;
+					}
+				}
+
+				return model;
 			});
 		});
 	}
diff --git a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/SpecificModelReloadListener.java b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/SpecificModelReloadListener.java
index 38d402a95..0e2d19afe 100644
--- a/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/SpecificModelReloadListener.java
+++ b/fabric-model-loading-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/model/loading/SpecificModelReloadListener.java
@@ -16,21 +16,18 @@
 
 package net.fabricmc.fabric.test.model.loading;
 
-import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.render.model.BakedModel;
 import net.minecraft.resource.ResourceManager;
-import net.minecraft.resource.SinglePreparationResourceReloader;
 import net.minecraft.util.Identifier;
-import net.minecraft.util.Unit;
-import net.minecraft.util.profiler.Profiler;
 
-import net.fabricmc.fabric.api.resource.IdentifiableResourceReloadListener;
 import net.fabricmc.fabric.api.resource.ResourceReloadListenerKeys;
+import net.fabricmc.fabric.api.resource.SimpleSynchronousResourceReloadListener;
 
-public class SpecificModelReloadListener extends SinglePreparationResourceReloader<Unit> implements IdentifiableResourceReloadListener {
+public class SpecificModelReloadListener implements SimpleSynchronousResourceReloadListener {
 	public static final SpecificModelReloadListener INSTANCE = new SpecificModelReloadListener();
 	public static final Identifier ID = Identifier.of(ModelTestModClient.ID, "specific_model");
 
@@ -41,13 +38,8 @@ public class SpecificModelReloadListener extends SinglePreparationResourceReload
 	}
 
 	@Override
-	protected Unit prepare(ResourceManager manager, Profiler profiler) {
-		return Unit.INSTANCE;
-	}
-
-	@Override
-	protected void apply(Unit loader, ResourceManager manager, Profiler profiler) {
-		specificModel = MinecraftClient.getInstance().getBakedModelManager().getModel(ModelTestModClient.MODEL_ID);
+	public void reload(ResourceManager manager) {
+		specificModel = MinecraftClient.getInstance().getBakedModelManager().getModel(ModelTestModClient.HALF_RED_SAND_MODEL_ID);
 	}
 
 	@Override
@@ -57,6 +49,6 @@ public class SpecificModelReloadListener extends SinglePreparationResourceReload
 
 	@Override
 	public Collection<Identifier> getFabricDependencies() {
-		return Arrays.asList(ResourceReloadListenerKeys.MODELS);
+		return List.of(ResourceReloadListenerKeys.MODELS);
 	}
 }
diff --git a/fabric-renderer-api-v1/build.gradle b/fabric-renderer-api-v1/build.gradle
index 9bc0edaa6..4c5a4747f 100644
--- a/fabric-renderer-api-v1/build.gradle
+++ b/fabric-renderer-api-v1/build.gradle
@@ -6,7 +6,7 @@ testDependencies(project, [
 	':fabric-block-api-v1',
 	':fabric-block-view-api-v2',
 	':fabric-blockrenderlayer-v1',
-//	':fabric-model-loading-api-v1', TODO 1.21
+	':fabric-model-loading-api-v1',
 	':fabric-object-builder-api-v1',
 	':fabric-renderer-indigo',
 	':fabric-resource-loader-v0'
diff --git a/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/ModelResolverImpl.java b/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/ModelResolverImpl.java
index d04af7215..6cb8950f2 100644
--- a/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/ModelResolverImpl.java
+++ b/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/ModelResolverImpl.java
@@ -16,7 +16,15 @@
 
 package net.fabricmc.fabric.test.renderer.client;
 
-/* TODO 1.21
+import java.util.Set;
+
+import org.jetbrains.annotations.Nullable;
+
+import net.minecraft.client.render.model.UnbakedModel;
+import net.minecraft.util.Identifier;
+
+import net.fabricmc.fabric.api.client.model.loading.v1.ModelResolver;
+import net.fabricmc.fabric.test.renderer.RendererTest;
 
 public class ModelResolverImpl implements ModelResolver {
 	private static final Set<Identifier> FRAME_MODEL_LOCATIONS = Set.of(
@@ -65,4 +73,3 @@ public class ModelResolverImpl implements ModelResolver {
 		return null;
 	}
 }
-*/
diff --git a/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/RendererClientTest.java b/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/RendererClientTest.java
index 8578427cb..79645dffe 100644
--- a/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/RendererClientTest.java
+++ b/fabric-renderer-api-v1/src/testmodClient/java/net/fabricmc/fabric/test/renderer/client/RendererClientTest.java
@@ -20,16 +20,16 @@ import net.minecraft.client.render.RenderLayer;
 
 import net.fabricmc.api.ClientModInitializer;
 import net.fabricmc.fabric.api.blockrenderlayer.v1.BlockRenderLayerMap;
+import net.fabricmc.fabric.api.client.model.loading.v1.ModelLoadingPlugin;
 import net.fabricmc.fabric.test.renderer.FrameBlock;
 import net.fabricmc.fabric.test.renderer.Registration;
 
 public final class RendererClientTest implements ClientModInitializer {
 	@Override
 	public void onInitializeClient() {
-		// TODO 1.21
-		// ModelLoadingPlugin.register(pluginContext -> {
-		// 	pluginContext.resolveModel().register(new ModelResolverImpl());
-		// });
+		ModelLoadingPlugin.register(pluginContext -> {
+			pluginContext.resolveModel().register(new ModelResolverImpl());
+		});
 
 		for (FrameBlock frameBlock : Registration.FRAME_BLOCKS) {
 			// We don't specify a material for the frame mesh,
diff --git a/settings.gradle b/settings.gradle
index f31c6d62b..cef611adc 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -39,7 +39,7 @@ include 'fabric-key-binding-api-v1'
 include 'fabric-lifecycle-events-v1'
 include 'fabric-loot-api-v2'
 include 'fabric-message-api-v1'
-//include 'fabric-model-loading-api-v1' TODO 1.21
+include 'fabric-model-loading-api-v1'
 include 'fabric-networking-api-v1'
 include 'fabric-object-builder-api-v1'
 include 'fabric-particles-v1'