155 lines
6.0 KiB
Python
155 lines
6.0 KiB
Python
import re
|
|
import torch
|
|
import os
|
|
import folder_paths
|
|
import nodes
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
|
|
class LTXVTimelineEditor:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"script": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Visual Timeline Script"
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE")
|
|
RETURN_NAMES = ("script_string", "guide_images")
|
|
FUNCTION = "execute"
|
|
CATEGORY = "ErosDiffusion/LTXV"
|
|
OUTPUT_NODE = False
|
|
|
|
DESCRIPTION = "Visually edit scene scripts. Outputs the script string and the batch of images referenced in the script (uploaded via the editor)."
|
|
|
|
def execute(self, script):
|
|
# Regex to find image refs: first:filename | end:filename | mid:time:filename
|
|
# We need to extract them, load them, and replace them with $0, $1 etc.
|
|
|
|
# 1. Parse all image references
|
|
# Format in script:
|
|
# first:filename.png
|
|
# end:filename.png
|
|
# mid:00:05.00:filename.png
|
|
|
|
# We'll use a dict to map filename -> index to avoid duplicates?
|
|
# Or just sequential batch?
|
|
# User said "we want $number style variables".
|
|
# If I reuse the same image in multiple places, I should probably reuse the index?
|
|
# But for simplicity, let's just make a list of Unique images found.
|
|
|
|
image_map = {} # val -> index (val could be just filename)
|
|
images_loaded = []
|
|
|
|
# Helper to load and process image
|
|
def get_image_index(filename):
|
|
filename = filename.strip()
|
|
if not filename: return None
|
|
|
|
if filename in image_map:
|
|
return f"${image_map[filename]}"
|
|
|
|
# Load Image
|
|
try:
|
|
# Use standard ComfyUI load logic
|
|
image_path = folder_paths.get_annotated_filepath(filename)
|
|
|
|
i = Image.open(image_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
image = i.convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,] # [1, H, W, 3]
|
|
|
|
images_loaded.append(image)
|
|
idx = len(images_loaded) - 1
|
|
image_map[filename] = idx
|
|
return f"${idx}"
|
|
|
|
except Exception as e:
|
|
print(f"LTXVTimelineEditor: Failed to load image '{filename}': {e}")
|
|
return None
|
|
|
|
# 2. Process Script Line by Line
|
|
new_lines = []
|
|
lines = script.split('\n')
|
|
|
|
for line in lines:
|
|
# We look for our directives
|
|
# Pattern: | first: (.*?) ($|\|)
|
|
# Pattern: | end: (.*?) ($|\|)
|
|
# Pattern: | mid:(\d+:\d+(?:\.\d+)?): (.*?) ($|\|)
|
|
|
|
# Since regex is tricky with multiple pipes, let's split by pipe
|
|
parts = [p.strip() for p in line.split('|')]
|
|
new_parts = []
|
|
new_parts.append(parts[0]) # The [Time] Prompt part
|
|
|
|
for p in parts[1:]:
|
|
if p.startswith("first:"):
|
|
fname = p[6:]
|
|
idx_str = get_image_index(fname)
|
|
if idx_str: new_parts.append(f"first:{idx_str}")
|
|
|
|
elif p.startswith("end:"):
|
|
fname = p[4:]
|
|
idx_str = get_image_index(fname)
|
|
if idx_str: new_parts.append(f"end:{idx_str}")
|
|
|
|
elif p.startswith("mid:"):
|
|
# mid:time:filename
|
|
# Split by colon, max 2 splits?
|
|
# mid:05:00:file.png -> problem if time has colons
|
|
# We expect MM:SS format usually.
|
|
# Let's extract time str and filename
|
|
# Regex for inner part
|
|
m = re.match(r"mid:(\d+:\d+(?:\.\d+)?):(.*)", p)
|
|
if m:
|
|
time_str = m.group(1)
|
|
fname = m.group(2)
|
|
idx_str = get_image_index(fname)
|
|
if idx_str: new_parts.append(f"mid:{time_str}:{idx_str}")
|
|
else:
|
|
new_parts.append(p) # Keep as is if cant parse
|
|
else:
|
|
new_parts.append(p) # Other directives or unrecognized
|
|
|
|
new_lines.append(" | ".join(new_parts))
|
|
|
|
processed_script = "\n".join(new_lines)
|
|
|
|
# 3. Batch Images
|
|
if not images_loaded:
|
|
# Return empty batch? Or None?
|
|
# If no images, return a dummy small tensor to avoid errors if connected?
|
|
# Or just None? If logic handles it.
|
|
# LTXVSceneExtender "guide_images" is optional.
|
|
empty_tensor = torch.zeros((1, 64, 64, 3), dtype=torch.float32) # Dummy
|
|
return (processed_script, None) # None is clearer for optional input
|
|
|
|
# Resize logic for batching
|
|
# We must resize all to match the first image dimensions (or max?)
|
|
# Standard Comfy batching behavior: usually same size required.
|
|
shape = images_loaded[0].shape # [1, H, W, 3]
|
|
target_h, target_w = shape[1], shape[2]
|
|
|
|
final_list = []
|
|
for img in images_loaded:
|
|
if img.shape[1] != target_h or img.shape[2] != target_w:
|
|
# Resize
|
|
# Permute to [1, 3, H, W] for interpolate
|
|
img_p = img.movedim(-1, 1)
|
|
img_p = torch.nn.functional.interpolate(img_p, size=(target_h, target_w), mode="bilinear", align_corners=False)
|
|
img_p = img_p.movedim(1, -1)
|
|
final_list.append(img_p)
|
|
else:
|
|
final_list.append(img)
|
|
|
|
batch = torch.cat(final_list, dim=0)
|
|
|
|
return (processed_script, batch)
|