diff --git a/animatediff/model_injection.py b/animatediff/model_injection.py
index 6332471..71fdeef 100644
--- a/animatediff/model_injection.py
+++ b/animatediff/model_injection.py
@@ -13,7 +13,8 @@ import comfy.conds
 import comfy.lora
 import comfy.model_management
 import comfy.utils
-from comfy.model_patcher import ModelPatcher, PatcherInjection, WrappersMP
+from comfy.model_patcher import ModelPatcher
+from comfy.patcher_extension import WrappersMP, PatcherInjection
 from comfy.model_base import BaseModel
 from comfy.sd import CLIP, VAE
 
diff --git a/animatediff/sampling.py b/animatediff/sampling.py
index c7d28af..7dcc4a6 100644
--- a/animatediff/sampling.py
+++ b/animatediff/sampling.py
@@ -13,7 +13,8 @@ import comfy.sampler_helpers
 import comfy.utils
 from comfy.controlnet import ControlBase
 from comfy.model_base import BaseModel
-from comfy.model_patcher import ModelPatcher, WrapperExecutor
+from comfy.model_patcher import ModelPatcher
+from comfy.patcher_extension import WrapperExecutor
 import comfy.conds
 import comfy.ops