296 lines
9.9 KiB
Python
296 lines
9.9 KiB
Python
import folder_paths
|
||
import math
|
||
|
||
class LoraKeyframes:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"clip": ("CLIP",),
|
||
"lora_name": (folder_paths.get_filename_list("loras"),),
|
||
"strength_multiplier": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.01,
|
||
}),
|
||
"preset": ([
|
||
"custom",
|
||
"constant_full",
|
||
"fade_in",
|
||
"fade_out",
|
||
"fade_in_out",
|
||
"early_only",
|
||
"late_only",
|
||
"middle_only",
|
||
"two_peaks"
|
||
], {"default": "custom"}),
|
||
"keyframes": ("STRING", {
|
||
"default": "0@0.0, 1@0.5, 0@1.0",
|
||
}),
|
||
"interpolation": ([
|
||
"step",
|
||
"linear",
|
||
"smoothstep",
|
||
"ease_in",
|
||
"ease_out",
|
||
"ease_in_out"
|
||
], {"default": "linear"}),
|
||
"resolution": ("INT", {
|
||
"default": 20,
|
||
"min": 2,
|
||
"max": 100,
|
||
"step": 1
|
||
})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("CLIP", "STRING")
|
||
RETURN_NAMES = ("CLIP", "schedule_graph")
|
||
FUNCTION = "apply_lora_keyframes"
|
||
CATEGORY = "loaders"
|
||
DESCRIPTION = """
|
||
LoRA Keyframes - Schedule LoRA strength during sampling
|
||
|
||
FORMAT: strength@percent, strength@percent, ...
|
||
- strength: LoRA strength multiplier at this point (0 = off, 1 = normal, 2 = double)
|
||
- percent: when to apply (0.0 = start, 0.5 = middle, 1.0 = end)
|
||
|
||
EXAMPLES:
|
||
• "0@0.0, 1@0.5, 0@1.0" → Bell-shaped curve (fade in/out)
|
||
• "0@0.0, 1@0.3" → Fade in and stay on
|
||
• "1@0.0, 0@0.7" → Start strong, then fade out
|
||
• "0@0.2, 1@0.5, 0@0.8" → Active only in middle
|
||
|
||
INTERPOLATION MODES:
|
||
• step → Rectangular jumps (default ComfyUI behavior)
|
||
• linear → Linear interpolation between keyframes
|
||
• smoothstep → Smooth fade
|
||
• ease_in → Starts slow, accelerates
|
||
• ease_out → Starts fast, slows down
|
||
• ease_in_out → Smooth ease in/out
|
||
|
||
RESOLUTION:
|
||
- Controls how many intermediate keyframes are generated between your defined keyframes
|
||
- Number of points = resolution + 1
|
||
- Example: resolution = 20 → 21 keyframes at 0%, 5%, 10%, ..., 100%
|
||
- Higher resolution → smoother curve, more granular strength changes
|
||
- Lower resolution → coarser, fewer keyframes, faster processing
|
||
|
||
TIP: Use fewer keyframes for smoother sampling (each keyframe adds a tiny delay)
|
||
"""
|
||
|
||
|
||
PRESETS = {
|
||
"custom": None,
|
||
"constant_full": "1@0.0",
|
||
"fade_in": "0@0.0, 1@0.3",
|
||
"fade_out": "1@0.0, 0@0.7",
|
||
"fade_in_out": "0@0.0, 1@0.5, 0@1.0",
|
||
"early_only": "1@0.0, 0@0.3",
|
||
"late_only": "0@0.7, 1@0.8",
|
||
"middle_only": "0@0.0, 1@0.3, 1@0.7, 0@1.0",
|
||
"two_peaks": "0@0.0, 1@0.25, 0@0.5, 1@0.75, 0@1.0"
|
||
}
|
||
|
||
# ------------------------------------------------------------------
|
||
# Keyframe parsing / normalization
|
||
# ------------------------------------------------------------------
|
||
|
||
def parse_keyframes(self, text):
|
||
result = []
|
||
for part in text.split(","):
|
||
part = part.strip()
|
||
if "@" not in part:
|
||
continue
|
||
try:
|
||
s, p = part.split("@")
|
||
result.append((float(s), float(p)))
|
||
except ValueError:
|
||
print(f"⚠️ Invalid keyframe ignored: {part}")
|
||
return result
|
||
|
||
def normalize_keyframes(self, keyframes):
|
||
kf = {}
|
||
for s, p in keyframes:
|
||
p = max(0.0, min(1.0, p))
|
||
kf[p] = s # last wins
|
||
|
||
items = sorted([(s, p) for p, s in kf.items()], key=lambda x: x[1])
|
||
|
||
if not items:
|
||
return []
|
||
|
||
if items[0][1] > 0.0:
|
||
items.insert(0, (0.0, 0.0))
|
||
if items[-1][1] < 1.0:
|
||
items.append((0.0, 1.0))
|
||
|
||
return items
|
||
|
||
# ------------------------------------------------------------------
|
||
# Interpolation curves
|
||
# ------------------------------------------------------------------
|
||
|
||
def apply_curve(self, t, mode):
|
||
if mode == "linear":
|
||
return t
|
||
if mode == "smoothstep":
|
||
return t * t * (3 - 2 * t)
|
||
if mode == "ease_in":
|
||
return t * t
|
||
if mode == "ease_out":
|
||
return 1 - (1 - t) * (1 - t)
|
||
if mode == "ease_in_out":
|
||
return 0.5 * (1 - math.cos(math.pi * t))
|
||
return 0.0 # step handled separately
|
||
|
||
def interpolate(self, keyframes, percent, multiplier, mode):
|
||
before = None
|
||
after = None
|
||
|
||
for s, p in keyframes:
|
||
if p <= percent:
|
||
before = (s, p)
|
||
if p >= percent and after is None:
|
||
after = (s, p)
|
||
|
||
if before is None:
|
||
return keyframes[0][0] * multiplier
|
||
if after is None:
|
||
return keyframes[-1][0] * multiplier
|
||
|
||
if mode == "step" or before[1] == after[1]:
|
||
return before[0] * multiplier
|
||
|
||
t = (percent - before[1]) / (after[1] - before[1])
|
||
t = self.apply_curve(t, mode)
|
||
|
||
return (before[0] + t * (after[0] - before[0])) * multiplier
|
||
|
||
# ------------------------------------------------------------------
|
||
# Densify keyframes
|
||
# ------------------------------------------------------------------
|
||
|
||
def densify_keyframes(self, keyframes, multiplier, mode, resolution):
|
||
dense = []
|
||
for i in range(resolution + 1):
|
||
percent = i / resolution
|
||
strength = self.interpolate(keyframes, percent, multiplier=1.0, mode=mode)
|
||
dense.append((strength, percent))
|
||
return dense
|
||
|
||
# ------------------------------------------------------------------
|
||
# ASCII graph
|
||
# ------------------------------------------------------------------
|
||
|
||
def generate_graph(self, keyframes, multiplier, mode):
|
||
width = 60
|
||
height = 15
|
||
|
||
values = [s * multiplier for s, _ in keyframes]
|
||
min_v = min(0.0, min(values))
|
||
max_v = max(values)
|
||
span = max(max_v - min_v, 1e-6)
|
||
|
||
rows = []
|
||
for r in range(height):
|
||
row_val = max_v - (r / (height - 1)) * span
|
||
line = []
|
||
for c in range(width):
|
||
p = c / (width - 1)
|
||
v = self.interpolate(keyframes, p, multiplier, mode)
|
||
line.append("█" if abs(v - row_val) < span / height else " ")
|
||
rows.append("".join(line))
|
||
|
||
out = []
|
||
out.append("=" * 70)
|
||
out.append(f"LoRA Strength Schedule | {mode} ×{multiplier}")
|
||
out.append("=" * 70)
|
||
out.append(f"{max_v:>6.2f} ┤{rows[0]}")
|
||
for r in rows[1:-1]:
|
||
out.append(" │" + r)
|
||
out.append(f"{min_v:>6.2f} ┤{rows[-1]}")
|
||
out.append(" └" + "─" * width)
|
||
out.append(" 0%" + " " * (width - 7) + "100%")
|
||
out.append("")
|
||
out.append("Keyframes:")
|
||
for s, p in keyframes:
|
||
out.append(f" {p*100:>5.1f}% → {s*multiplier:.2f}")
|
||
out.append("=" * 70)
|
||
return "\n".join(out)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Main
|
||
# ------------------------------------------------------------------
|
||
|
||
def apply_lora_keyframes(self, clip, lora_name, strength_multiplier,
|
||
preset, keyframes, interpolation, resolution):
|
||
|
||
# 1️⃣ Select preset
|
||
if preset != "custom":
|
||
keyframes = self.PRESETS[preset]
|
||
|
||
# 2️⃣ Parse + normalize
|
||
parsed = self.normalize_keyframes(
|
||
self.parse_keyframes(keyframes)
|
||
)
|
||
|
||
# 3️⃣ Densify keyframes to match interpolation
|
||
parsed = self.densify_keyframes(parsed, strength_multiplier, interpolation, resolution)
|
||
|
||
# 4️⃣ Generate ASCII graph
|
||
graph = self.generate_graph(parsed, strength_multiplier, interpolation)
|
||
print("\n" + graph)
|
||
|
||
if not parsed:
|
||
return (clip, graph)
|
||
|
||
# 5️⃣ Call native ComfyUI nodes safely
|
||
from nodes import NODE_CLASS_MAPPINGS
|
||
|
||
CreateHookLora = NODE_CLASS_MAPPINGS["CreateHookLora"]()
|
||
lora_fn = getattr(CreateHookLora, CreateHookLora.FUNCTION)
|
||
|
||
CreateHookKeyframe = NODE_CLASS_MAPPINGS["CreateHookKeyframe"]()
|
||
kf_fn = getattr(CreateHookKeyframe, CreateHookKeyframe.FUNCTION)
|
||
|
||
SetHookKeyframes = NODE_CLASS_MAPPINGS["SetHookKeyframes"]()
|
||
set_kf_fn = getattr(SetHookKeyframes, SetHookKeyframes.FUNCTION)
|
||
|
||
SetClipHooks = NODE_CLASS_MAPPINGS["SetClipHooks"]()
|
||
set_clip_fn = getattr(SetClipHooks, SetClipHooks.FUNCTION)
|
||
|
||
# 6️⃣ Create LoRA hooks
|
||
hooks = lora_fn(
|
||
lora_name=lora_name,
|
||
strength_model=strength_multiplier,
|
||
strength_clip=strength_multiplier,
|
||
prev_hooks=None
|
||
)[0]
|
||
|
||
# 7️⃣ Create keyframes
|
||
hook_kf = None
|
||
for s, p in parsed:
|
||
hook_kf = kf_fn(
|
||
strength_mult=s,
|
||
start_percent=p,
|
||
prev_hook_kf=hook_kf
|
||
)[0]
|
||
|
||
# 8️⃣ Apply keyframes to hooks
|
||
hooks = set_kf_fn(
|
||
hooks=hooks,
|
||
hook_kf=hook_kf
|
||
)[0]
|
||
|
||
# 9️⃣ Apply hooks to CLIP
|
||
clip = set_clip_fn(
|
||
clip=clip,
|
||
hooks=hooks,
|
||
schedule_clip=False,
|
||
apply_to_conds=True
|
||
)[0]
|
||
|
||
return (clip, graph)
|