Update z-prompt-fusion-extension/scripts/promptlang.py

z-prompt-fusion-extension/scripts/promptlang.py

@@ -134,7 +134,8 @@ def _adapt_flattened_schedules(result: Any, total_steps: int):
                 if isinstance(src, dict):
                     out = dict(dst)
                     for k, v in src.items():
-                        out[k] = out.get(k, 0) + v * w
+                        prev = out.get(k)
+                        out[k] = (prev + v * w) if prev is not None else (v * w)
                     return out
                 return dst + src * w
 
@@ -214,41 +215,54 @@ def _build_tensor_for_prompt(
         ]
         wrapped_conds.append(wrapped_sched)
 
-    # 4) Собрать интерполяционный тензор и вычислить расписания по шагам
-    tensor = tensor_builder.build(wrapped_conds, empty_cond.get())
-
-    slerp_scale = global_state.get_slerp_scale()
-    slerp_epsilon = global_state.get_slerp_epsilon()
-
-    # Сформировать расписание, склеивая одинаковые соседние сегменты
-    schedules: List[prompt_parser.ScheduledPromptConditioning] = []
-    prev_wrapper = None
-    for step in range(total_steps):
-        params = interpolation_tensor.InterpolationParams(
-            t=step / max(1, total_steps - 1),
-            step=step,
-            total_steps=total_steps,
-            slerp_scale=slerp_scale,
-            slerp_epsilon=slerp_epsilon,
-        )
-        origin = global_state.get_origin_cond_at(step, is_hires=is_hires)
-        cond_wrapper = tensor.interpolate(params, origin, empty_cond.get())
+        # 4) Собрать интерполяционный тензор и вычислить расписания по шагам
+        tensor = tensor_builder.build(wrapped_conds, empty_cond.get())
 
-        # Склейка соседних, если равны (по wrapper.__eq__)
-        if prev_wrapper is not None and cond_wrapper == prev_wrapper:
-            schedules[-1].end_at_step = step
-        else:
-            raw = cond_wrapper.original_cond if hasattr(cond_wrapper, "original_cond") else cond_wrapper
-            schedules.append(
-                prompt_parser.ScheduledPromptConditioning(end_at_step=step, cond=raw)
+        slerp_scale = global_state.get_slerp_scale()
+        slerp_epsilon = global_state.get_slerp_epsilon()
+
+        # Вспомогалка: безопасно «продлить» последний сегмент до step
+        def _extend_last_segment_to(step_val: int):
+            last = schedules[-1]
+            if getattr(last, "end_at_step", None) != step_val:
+                schedules[-1] = prompt_parser.ScheduledPromptConditioning(
+                    end_at_step=step_val,
+                    cond=last.cond,
+                )
+
+        # Сформировать расписание, склеивая одинаковые соседние сегменты
+        schedules: List[prompt_parser.ScheduledPromptConditioning] = []
+        prev_wrapper = None
+        for step in range(total_steps):
+            params = interpolation_tensor.InterpolationParams(
+                t=step / max(1, total_steps - 1),
+                step=step,
+                total_steps=total_steps,
+                slerp_scale=slerp_scale,
+                slerp_epsilon=slerp_epsilon,
             )
-            prev_wrapper = cond_wrapper
+            origin = global_state.get_origin_cond_at(step, is_hires=is_hires)
+            cond_wrapper = tensor.interpolate(params, origin, empty_cond.get())
+
+            if prev_wrapper is not None and cond_wrapper == prev_wrapper:
+                # Было: schedules[-1].end_at_step = step  (иммутабельно -> нельзя)
+                _extend_last_segment_to(step)
+            else:
+                raw = getattr(cond_wrapper, "original_cond", cond_wrapper)
+                schedules.append(
+                    prompt_parser.ScheduledPromptConditioning(
+                        end_at_step=step,
+                        cond=raw,
+                    )
+                )
+                prev_wrapper = cond_wrapper
+
+        # Финализировать последний сегмент на последний шаг
+        if schedules:
+            _extend_last_segment_to(total_steps - 1)
 
-    # Финализировать последний сегмент на последний шаг
-    if schedules:
-        schedules[-1].end_at_step = total_steps - 1
+        return schedules
 
-    return schedules
 
 
 # -----------------------------------------------------------------------------