diff --git a/src/main/java/me/xginko/villageroptimizer/models/WrappedVillager.java b/src/main/java/me/xginko/villageroptimizer/models/WrappedVillager.java
index c3aa678..7a0a704 100644
--- a/src/main/java/me/xginko/villageroptimizer/models/WrappedVillager.java
+++ b/src/main/java/me/xginko/villageroptimizer/models/WrappedVillager.java
@@ -1,5 +1,6 @@
 package me.xginko.villageroptimizer.models;
 
+import me.xginko.villageroptimizer.VillagerOptimizer;
 import me.xginko.villageroptimizer.enums.Keys;
 import me.xginko.villageroptimizer.enums.OptimizationType;
 import org.bukkit.entity.Villager;
@@ -45,11 +46,15 @@ public final class WrappedVillager {
     public void setOptimization(OptimizationType type) {
         if (type.equals(OptimizationType.OFF) && isOptimized()) {
             dataContainer.remove(Keys.OPTIMIZATION_TYPE.key());
-            villager.setAware(true);
-            villager.setAI(true);
+            villager.getScheduler().run(VillagerOptimizer.getInstance(), enableAI -> {
+                villager.setAware(true);
+                villager.setAI(true);
+            }, null);
         } else {
             dataContainer.set(Keys.OPTIMIZATION_TYPE.key(), PersistentDataType.STRING, type.name());
-            villager.setAware(false);
+            villager.getScheduler().run(VillagerOptimizer.getInstance(), disableAI -> {
+                villager.setAware(false);
+            }, null);
         }
     }
 
diff --git a/src/main/java/me/xginko/villageroptimizer/modules/LevelVillagers.java b/src/main/java/me/xginko/villageroptimizer/modules/LevelVillagers.java
index 2d7071d..8d45ede 100644
--- a/src/main/java/me/xginko/villageroptimizer/modules/LevelVillagers.java
+++ b/src/main/java/me/xginko/villageroptimizer/modules/LevelVillagers.java
@@ -68,8 +68,8 @@ public class LevelVillagers implements VillagerOptimizerModule, Listener {
                     && wVillager.calculateLevel() > villager.getVillagerLevel()
             ) {
                 villager.addPotionEffect(new PotionEffect(PotionEffectType.SLOW, (int) (20 + (cooldown / 50L)), 120, false, false));
-                villager.setAware(true);
-                villager.getScheduler().runDelayed(plugin, reOptimize -> {
+                villager.getScheduler().run(VillagerOptimizer.getInstance(), enableAI -> villager.setAware(true), null);
+                villager.getScheduler().runDelayed(plugin, disableAI -> {
                     villager.setAware(false);
                     wVillager.saveLastLevelUp();
                 }, null, 100L);