Made AdvancedPromptEmbeds that is compatable with previous PromptEmbeds functionality, but is more streamlined and can accomidate more model embedding paradigms.

toolkit/prompt_utils.py

@@ -8,6 +8,8 @@ import random
 
 from toolkit.train_tools import get_torch_dtype
 import itertools
+from safetensors import safe_open
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
 
 if TYPE_CHECKING:
     from toolkit.config_modules import SliderTargetConfig
@@ -145,6 +147,12 @@ class PromptEmbeds:
         :param path: The path to load the prompt embeds from.
         :return: An instance of PromptEmbeds.
         """
+        # first check if it is advanced prompt embed file
+        f = safe_open(path, framework='pt')
+        metadata = f.metadata()
+        if metadata is not None and metadata.get("class_name", "") == "AdvancedPromptEmbeds":
+            return AdvancedPromptEmbeds.load(path=path)
+        
         state_dict = load_file(path, device='cpu')
         text_embeds = []
         pooled_embeds = None
@@ -245,6 +253,9 @@ class EncodedPromptPair:
 
 
 def concat_prompt_embeds(prompt_embeds: list["PromptEmbeds"], padding_side: str = "right") -> PromptEmbeds:
+    # check if first item has a classmethod of concat_prompt_embeds
+    if hasattr(prompt_embeds[0].__class__, "concat_prompt_embeds"):
+        return prompt_embeds[0].__class__.concat_prompt_embeds(prompt_embeds, padding_side=padding_side)
     # --- pad text_embeds ---
     if isinstance(prompt_embeds[0].text_embeds, (list, tuple)):
         text_embeds = []
@@ -335,6 +346,8 @@ def concat_prompt_pairs(prompt_pairs: list[EncodedPromptPair]):
 
 
 def split_prompt_embeds(concatenated: PromptEmbeds, num_parts=None) -> List[PromptEmbeds]:
+    if hasattr(concatenated.__class__, "split_prompt_embeds"):
+        return concatenated.__class__.split_prompt_embeds(concatenated, num_parts=num_parts)
     if num_parts is None:
         # use batch size
         num_parts = concatenated.text_embeds.shape[0]