diff --git a/Dockerfile b/Dockerfile
index ec6c055..a2a9f5e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -21,8 +21,8 @@ RUN cd /app/custom_nodes \
     && git clone --depth 1 https://github.com/XLabs-AI/x-flux-comfyui \
     && pip install --no-cache-dir -r x-flux-comfyui/requirements.txt
 
-RUN sed -i 's/def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple\[Tensor, Tensor\]:/def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, transformer_options={}, **kwargs):/' /app/custom_nodes/x-flux-comfyui/xflux/src/flux/modules/layers.py \
-    && grep "attn_mask" /app/custom_nodes/x-flux-comfyui/xflux/src/flux/modules/layers.py
+COPY fixes.py /tmp/fixes.py
+RUN python3 /tmp/fixes.py
 
 RUN useradd -u 1000 -m comfyui && chown -R 1000:1000 /app