diff --git a/extensions-builtin/sd_forge_stylealign/scripts/forge_stylealign.py b/extensions-builtin/sd_forge_stylealign/scripts/forge_stylealign.py
index 541e2be5..901cc0e7 100644
--- a/extensions-builtin/sd_forge_stylealign/scripts/forge_stylealign.py
+++ b/extensions-builtin/sd_forge_stylealign/scripts/forge_stylealign.py
@@ -2,11 +2,11 @@ import torch
 import gradio as gr
 
 from modules import scripts
-import ldm_patched.ldm.modules.attention as attention
+from backend import attention
 
 
 def sdp(q, k, v, transformer_options):
-    return attention.optimized_attention(q, k, v, heads=transformer_options["n_heads"], mask=None)
+    return attention.attention_function(q, k, v, heads=transformer_options["n_heads"], mask=None)
 
 
 class StyleAlignForForge(scripts.Script):