From 38da628a01fe6ebc2368653b1d2dd2cfbb324712 Mon Sep 17 00:00:00 2001
From: Brody Childs <maximombro@gmail.com>
Date: Fri, 23 Feb 2024 15:28:10 -0500
Subject: [PATCH] Prompt Sequence Interpolator now returns prompt sequence

---
 mbmInterpPromptSequence.py       | 18 ++++++++++++++++--
 mbmPromptSequence.py             |  2 ++
 mbmPromptSequenceInterpolator.py | 32 ++++++++++++--------------------
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/mbmInterpPromptSequence.py b/mbmInterpPromptSequence.py
index 8f3285c..9ece518 100644
--- a/mbmInterpPromptSequence.py
+++ b/mbmInterpPromptSequence.py
@@ -1,5 +1,5 @@
 # MBM's Music Visualizer: Interpolated Prompt Sequence
-# Represents and facilitates building of a sequence of interpolated prompts.
+# Represents and facilitates building of a sequence of prompts with weighted interpolation.
 
 # Imports
 import torch
@@ -9,7 +9,7 @@ from .mbmPromptSequence import MbmPrompt
 # Classes
 class InterpPromptSequence:
     """
-    Represents and facilitates building of a sequence of interpolated prompts.
+    Represents and facilitates building of a sequence of prompts with weighted interpolation.
     """
     # Constructor
     def __init__(self, start: MbmPrompt, end: MbmPrompt, modifiers: torch.Tensor) -> None:
@@ -93,6 +93,20 @@ class InterpPromptSequence:
         self.positivePools = self.positivePools[:length]
         self.negativePools = self.negativePools[:length]
 
+    def asPromptSequence(self) -> list[MbmPrompt]:
+        """
+        Returns the sequence as a list of MbmPrompt objects.
+        """
+        return [
+            MbmPrompt(
+                self.positives[i],
+                self.negatives[i],
+                positivePool=self.positivePools[i],
+                negativePool=self.negativePools[i]
+            )
+            for i in range(len(self.positives))
+        ]
+
     def _weightedInterpolation(self, start: torch.Tensor, stop: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
         """
         Interpolates between `start` and `stop` based on the given `weights` for each step in the interpolation.
diff --git a/mbmPromptSequence.py b/mbmPromptSequence.py
index bf721f5..4b5d6b5 100644
--- a/mbmPromptSequence.py
+++ b/mbmPromptSequence.py
@@ -28,6 +28,8 @@ class MbmPrompt:
         self.positivePool = positivePool
         self.negativePool = negativePool
 
+        # TODO: Add additional "data" parameter for additional data to be stored with the prompt
+
     # Python Functions
     def __repr__(self) -> str:
         return f"MbmPrompt(positive={self.positive.shape}, negative={self.negative.shape}, positivePool={self.positivePool.shape if self.positivePool is not None else 'None'}, negativePool={self.negativePool.shape if self.negativePool is not None else 'None'})"
diff --git a/mbmPromptSequenceInterpolator.py b/mbmPromptSequenceInterpolator.py
index 80c114c..4592efb 100644
--- a/mbmPromptSequenceInterpolator.py
+++ b/mbmPromptSequenceInterpolator.py
@@ -56,42 +56,34 @@ class PromptSequenceInterpolator:
         # Set intial prompts
         if promptCount > 1:
             # Calculate linear interpolation between prompts
-            promptSeq: Optional[InterpPromptSequence] = None
+            interpPromptSeq: Optional[InterpPromptSequence] = None
             relDesiredFrames = math.ceil(desiredFrames / (promptCount - 1))
             for i in range(promptCount - 1):
                 # Calculate modifiers for this section
                 curModifiers = feat_mods[(relDesiredFrames * i):(relDesiredFrames * (i + 1))]
 
                 # Build prompt interpolation
-                if promptSeq is None:
+                if interpPromptSeq is None:
                     # Start intial prompt sequence
-                    promptSeq = InterpPromptSequence(prompts[i], prompts[i + 1], curModifiers)
+                    interpPromptSeq = InterpPromptSequence(prompts[i], prompts[i + 1], curModifiers)
                 else:
                     # Expand prompt sequence
-                    promptSeq.addToSequence(prompts[i], prompts[i + 1], curModifiers)
+                    interpPromptSeq.addToSequence(prompts[i], prompts[i + 1], curModifiers)
 
             # Trim off any extra frames produced from ceil to int
-            promptSeq.trimToLength(desiredFrames)
+            interpPromptSeq.trimToLength(desiredFrames)
 
-            # Set the initial prompt
-            promptPos = MbmPrompt.buildComfyUiPrompt(
-                promptSeq.positives[0],
-                promptSeq.positivePools[0]
-            )
-            promptNeg = MbmPrompt.buildComfyUiPrompt(
-                promptSeq.negatives[0],
-                promptSeq.negativePools[0]
-            )
+            # Build the prompt sequence
+            promptSeq = interpPromptSeq.asPromptSequence()
         elif promptCount == 1:
-            # Set single prompt
-            promptPos = prompts[0].positivePrompt()
-            promptNeg = prompts[0].negativePrompt()
+            # Send it onward
+            promptSeq = prompts
         else:
-            # No prompts
-            raise ValueError("No prompts were provided to the Music Visualizer node. At least one prompt is required.")
+            # No prompts my guy
+            raise ValueError("At least one prompt is required.")
 
         return (
-            # TODO: prompt sequence
+            promptSeq,
             # TODO: charts
         )