fix visual registration

Author: maximumpower55

src/main/java/com/simibubi/create/foundation/data/CreateBlockEntityBuilder.java

@@ -2,6 +2,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Objects;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -99,13 +100,13 @@ public CreateBlockEntityBuilder<T, P> visual(
 	}
 
 	protected void registerVisualizer() {
-		var visualFactory = this.visualFactory;
-		if (visualFactory != null) {
+		this.onRegister((entry) -> {
+			Objects.requireNonNull(this.visualFactory);
 			Predicate<T> renderNormally = this.renderNormally;
-			SimpleBlockEntityVisualizer.builder(getEntry())
-					.factory(visualFactory.get())
-					.skipVanillaRender(be -> !renderNormally.test(be))
-					.apply();
-		}
+			SimpleBlockEntityVisualizer.builder(this.getEntry())
+				.factory(this.visualFactory.get())
+				.skipVanillaRender(be -> !renderNormally.test(be))
+				.apply();
+		});
 	}
 }

src/main/java/com/simibubi/create/foundation/data/CreateEntityBuilder.java

@@ -1,5 +1,6 @@
 package com.simibubi.create.foundation.data;
 
+import java.util.Objects;
 import java.util.function.Predicate;
 
 import javax.annotation.Nullable;
@@ -11,6 +12,7 @@
 import com.tterrag.registrate.fabric.EnvExecutor;
 import com.tterrag.registrate.util.nullness.NonNullSupplier;
 
+import dev.engine_room.flywheel.lib.visualization.SimpleBlockEntityVisualizer;
 import dev.engine_room.flywheel.lib.visualization.SimpleEntityVisualizer;
 import net.minecraft.world.entity.Entity;
 import net.minecraft.world.entity.EntityType;
@@ -53,13 +55,13 @@ public CreateEntityBuilder<T, P> visual(NonNullSupplier<SimpleEntityVisualizer.F
 	}
 
 	protected void registerVisualizer() {
-		var visualFactory = this.visualFactory;
-		if (visualFactory != null) {
+		this.onRegister((entry) -> {
+			Objects.requireNonNull(this.visualFactory);
 			Predicate<T> renderNormally = this.renderNormally;
-			SimpleEntityVisualizer.builder(getEntry())
-					.factory(visualFactory.get())
-					.skipVanillaRender(entity -> !renderNormally.test(entity))
-					.apply();
-		}
+			SimpleEntityVisualizer.builder(this.getEntry())
+				.factory(this.visualFactory.get())
+				.skipVanillaRender(entity -> !renderNormally.test(entity))
+				.apply();
+		});
 	}
 }