diff --git a/data_juicer/ops/mapper/sdxl_prompt2prompt_mapper.py b/data_juicer/ops/mapper/sdxl_prompt2prompt_mapper.py
index b8d30b6bd..c268d6f54 100644
--- a/data_juicer/ops/mapper/sdxl_prompt2prompt_mapper.py
+++ b/data_juicer/ops/mapper/sdxl_prompt2prompt_mapper.py
@@ -8,6 +8,8 @@
 from data_juicer.ops.op_fusion import LOADED_IMAGES
 from data_juicer.utils.lazy_loader import LazyLoader
 from data_juicer.utils.model_utils import get_model, prepare_model
+from diffusers.models.attention import Attention
+
 
 diffusers = LazyLoader('diffusers', 'diffusers')
 torch = LazyLoader('torch', 'torch')