From da49dd3e865d7db79efc5b7e417cb8091dd455b1 Mon Sep 17 00:00:00 2001
From: FizzleDorf <1fizzledorf@gmail.com>
Date: Tue, 27 Aug 2024 21:31:21 -0400
Subject: [PATCH] start and end frame controls actually work now

---
 BatchFuncs.py    | 17 +++++++++++------
 ScheduleTypes.py | 16 ++++++++--------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/BatchFuncs.py b/BatchFuncs.py
index cca0b4b..c74d12b 100644
--- a/BatchFuncs.py
+++ b/BatchFuncs.py
@@ -313,16 +313,16 @@ def interpolate_prompt_series(animation_prompts, max_frames, start_frame, pre_te
     # if it is an in-between frame and the prompts differ, composable diffusion will be performed.
     return (cur_prompt_series, nxt_prompt_series, weight_series)
 
-def BatchPoolAnimConditioning(cur_prompt_series, nxt_prompt_series, weight_series, clip):
+def BatchPoolAnimConditioning(cur_prompt_series, nxt_prompt_series, weight_series, clip, settings:ScheduleSettings):
     pooled_out = []
     cond_out = []
     max_size = 0
     if max_size == 0:
-        for i in range(len(cur_prompt_series)):
+        for i in range(0, settings.end_frame):
             tokens = clip.tokenize(str(cur_prompt_series[i]))
             cond_to, pooled_to = clip.encode_from_tokens(tokens, return_pooled=True)
             max_size = max(max_size, cond_to.shape[1])
-    for i in range(len(cur_prompt_series)):
+    for i in range(settings.start_frame, settings.end_frame):
         tokens = clip.tokenize(str(cur_prompt_series[i]))
         cond_to, pooled_to = clip.encode_from_tokens(tokens, return_pooled=True)
 
@@ -380,11 +380,16 @@ def BatchGLIGENConditioning(cur_prompt_series, nxt_prompt_series, weight_series,
 
     return cond_out, pooled_out
 
-def BatchPoolAnimConditioningSDXL(cur_prompt_series, nxt_prompt_series, weight_series):
+def BatchPoolAnimConditioningSDXL(cur_prompt_series, nxt_prompt_series, weight_series, settings:ScheduleSettings):
     pooled_out = []
     cond_out = []
-
-    for i in range(len(cur_prompt_series)):
+    max_size = 0
+    if max_size == 0:
+        for i in range(0, settings.end_frame):
+            tokens = clip.tokenize(str(cur_prompt_series[i]))
+            cond_to, pooled_to = clip.encode_from_tokens(tokens, return_pooled=True)
+            max_size = max(max_size, cond_to.shape[1])
+    for i in range(settings.start_frame,settings.end_frame):
         interpolated_conditioning = addWeighted(cur_prompt_series[i],
                                                 nxt_prompt_series[i],
                                                 weight_series[i])
diff --git a/ScheduleTypes.py b/ScheduleTypes.py
index a59e7a0..6f9173b 100644
--- a/ScheduleTypes.py
+++ b/ScheduleTypes.py
@@ -46,8 +46,8 @@ def batch_prompt_schedule(settings:ScheduleSettings,clip):
     neg_cur_prompt, neg_nxt_prompt, weight = interpolate_prompt_seriesA(neg, settings)
 
     # Apply composable diffusion across the batch
-    p = BatchPoolAnimConditioning(pos_cur_prompt, pos_nxt_prompt, weight, clip, )
-    n = BatchPoolAnimConditioning(neg_cur_prompt, neg_nxt_prompt, weight, clip, )
+    p = BatchPoolAnimConditioning(pos_cur_prompt, pos_nxt_prompt, weight, clip, settings)
+    n = BatchPoolAnimConditioning(neg_cur_prompt, neg_nxt_prompt, weight, clip, settings)
 
     # return positive and negative conditioning as well as the current and next prompts for each
     return (p, n,)
@@ -63,13 +63,13 @@ def batch_prompt_schedule_latentInput(settings:ScheduleSettings,clip, latents):
     pos_cur_prompt, pos_nxt_prompt, weight = interpolate_prompt_seriesA(pos, settings)
 
     # Apply composable diffusion across the batch
-    p = BatchPoolAnimConditioning(pos_cur_prompt, pos_nxt_prompt, weight, clip)
+    p = BatchPoolAnimConditioning(pos_cur_prompt, pos_nxt_prompt, weight, clip, settings)
 
     # Interpolate the negative prompt weights over frames
     neg_cur_prompt, neg_nxt_prompt, weight = interpolate_prompt_seriesA(neg, settings)
 
     # Apply composable diffusion across the batch
-    n = BatchPoolAnimConditioning(neg_cur_prompt, neg_nxt_prompt, weight, clip)
+    n = BatchPoolAnimConditioning(neg_cur_prompt, neg_nxt_prompt, weight, clip, settings)
 
     return (p, n, latents,)
 
@@ -144,8 +144,8 @@ def batch_prompt_schedule_SDXL(settings:ScheduleSettings,clip):
     pc, pn, pw = BatchInterpolatePromptsSDXL(posG, posL, clip, settings,)
     nc, nn, nw = BatchInterpolatePromptsSDXL(negG, negL, clip, settings,)
 
-    p = BatchPoolAnimConditioningSDXL(pc, pn, pw)
-    n = BatchPoolAnimConditioningSDXL(nc, nn, nw)
+    p = BatchPoolAnimConditioningSDXL(pc, pn, pw, settings)
+    n = BatchPoolAnimConditioningSDXL(nc, nn, nw, settings)
 
     return (p, n,)
 
@@ -163,8 +163,8 @@ def batch_prompt_schedule_SDXL_latentInput(settings:ScheduleSettings,clip, laten
     pc, pn, pw = BatchInterpolatePromptsSDXL(posG, posL, clip, settings)
     nc, nn, nw = BatchInterpolatePromptsSDXL(negG, negL, clip, settings)
 
-    p = BatchPoolAnimConditioningSDXL(pc, pn, pw)
-    n = BatchPoolAnimConditioningSDXL(nc, nn, nw)
+    p = BatchPoolAnimConditioningSDXL(pc, pn, pw, settings)
+    n = BatchPoolAnimConditioningSDXL(nc, nn, nw, settings)
 
     return (p, n, latents,)