diff --git a/comfy/ldm/common_dit.py b/comfy/ldm/common_dit.py
index 3f6f38917..a31c9f571 100644
--- a/comfy/ldm/common_dit.py
+++ b/comfy/ldm/common_dit.py
@@ -1,5 +1,5 @@
 import torch
-from einops import rearrange
+from einops import rearrange, repeat
 import comfy.rmsnorm
 
 
diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py
index 1333b45e7..fb38d1a73 100644
--- a/comfy/ldm/flux/model.py
+++ b/comfy/ldm/flux/model.py
@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 from torch import Tensor, nn
-from einops import rearrange, repeat
+from einops import rearrange
 import comfy.ldm.common_dit
 import comfy.patcher_extension