From a16522d61fc55a4636e36fac5dd3ebaebdc5965d Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Mon, 8 Jul 2024 23:18:41 +0900
Subject: [PATCH] fix: consider compatibility with A1111

---
 build_dynamic_prompt.py        | 20 ++++++++++++++++----
 superprompter/superprompter.py | 10 +++++++++-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/build_dynamic_prompt.py b/build_dynamic_prompt.py
index fd82e43..9bc5d86 100644
--- a/build_dynamic_prompt.py
+++ b/build_dynamic_prompt.py
@@ -1,9 +1,21 @@
 import random
 import re
-from .csv_reader import *
-from .random_functions import *
-from .one_button_presets import OneButtonPresets
-from .superprompter.superprompter import *
+
+if __package__ is None or __package__ == '':
+    # A1111 style (standalone script or direct module execution)
+    # Use absolute imports for compatibility with A1111 WebUI environment
+    from csv_reader import *
+    from random_functions import *
+    from one_button_presets import OneButtonPresets
+    from superprompter.superprompter import *
+else:
+    # ComfyUI style (imported as a package)
+    # Use relative imports for proper integration with ComfyUI
+    from .csv_reader import *
+    from .random_functions import *
+    from .one_button_presets import OneButtonPresets
+    from .superprompter.superprompter import *
+
 OBPresets = OneButtonPresets()
 
 
diff --git a/superprompter/superprompter.py b/superprompter/superprompter.py
index ceb2b14..1326cfd 100644
--- a/superprompter/superprompter.py
+++ b/superprompter/superprompter.py
@@ -3,7 +3,15 @@ import os
 import random
 import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration
-from .download_models import download_models
+
+if __package__ is None or __package__ == '':
+    # A1111 style (standalone script or direct module execution)
+    # Use absolute imports for compatibility with A1111 WebUI environment
+    from download_models import download_models
+else:
+    # ComfyUI style (imported as a package)
+    # Use relative imports for proper integration with ComfyUI
+    from .download_models import download_models
 
 global tokenizer, model
 script_dir = os.path.dirname(os.path.abspath(__file__))  # Script directory