diff --git a/infer.py b/infer.py
index 88b9f09..c03a452 100644
--- a/infer.py
+++ b/infer.py
@@ -46,7 +46,7 @@ def edit_image(
     final_timestep = torch.ones((1,), dtype=torch.int64, device="cuda") * 999
 
     # Input Image
-    pil_img_cond = Image.open(img_path).resize((512, 512))
+    pil_img_cond = Image.open(img_path).convert("RGB").resize((512, 512))
 
     processed_image = to_tensor(pil_img_cond).unsqueeze(0).to("cuda") * 2 - 1
 
@@ -110,4 +110,4 @@ def edit_image(
     start_time = time.time()
     result = edit_image(img_path, src_p, edit_p, inverse_model, aux_model, ip_sb_model)
     print(f"Edit {src_p}->{edit_p} in {time.time()-start_time}")
-    save_image(result, f"result_{src_p}->{edit_p}.png")
+    save_image(result, f"result_{src_p}->{edit_p}.png")
\ No newline at end of file
diff --git a/models.py b/models.py
index aefe1ee..02f29d8 100644
--- a/models.py
+++ b/models.py
@@ -12,9 +12,9 @@
 )
 
 from src.mask_ip_controller import *
-from src.ip_adapter.attention_processor import AttnProcessor2_0 as AttnProcessor
-from src.ip_adapter.attention_processor import IPAttnProcessor2_0 as IPAttnProcessor
-from src.ip_adapter.mask_attention_processor import IPAttnProcessor2_0WithIPMaskController
+from src.attention_processor import AttnProcessor2_0 as AttnProcessor
+from src.attention_processor import IPAttnProcessor2_0 as IPAttnProcessor
+from src.mask_attention_processor import IPAttnProcessor2_0WithIPMaskController
 
 def tokenize_captions(tokenizer, captions):
     inputs = tokenizer(