284 lines
14 KiB
Python
284 lines
14 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont, ImageColor
|
|
import os
|
|
|
|
class TwoPassTypingColorASCIIAnimationNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
default_charset = " .'`^\":;Il!i~_?[{1(|/fjxnvzXYJCLQ0Zmwqdkh*#M&8%B@$"
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"char_width": ("INT", {"default": 80, "min": 10, "max": 1000, "step": 10, "display": "slider"}),
|
|
"font_path": ("STRING", {"default": "cour.ttf", "multiline": False}),
|
|
"font_size": ("INT", {"default": 15, "min": 5, "max": 100, "step": 1, "display": "slider"}),
|
|
"ascii_charset": ("STRING", {"default": default_charset, "multiline": True}),
|
|
"background_color": ("STRING", {"default": "#000000", "multiline": False}),
|
|
"initial_text_color": ("STRING", {"default": "#FFFFFF", "multiline": False}), # Color for first pass
|
|
"invert_brightness_mapping": ("BOOLEAN", {"default": False}),
|
|
"chars_per_frame": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1, "display":"slider"}),
|
|
"cursor_char": ("STRING", {"default": "_", "multiline": False}),
|
|
"cursor_blink_frames": ("INT", {"default": 5, "min": 1, "max": 60, "step": 1, "display":"slider"}),
|
|
"cursor_color": ("STRING", {"default": "#FFFFFF", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("two_pass_animated_frames",)
|
|
FUNCTION = "generate_two_pass_animation"
|
|
CATEGORY = "image/animation"
|
|
|
|
def tensor_to_pil(self, tensor_image: torch.Tensor) -> Image.Image:
|
|
img_np = tensor_image.cpu().numpy()
|
|
img_np = (img_np * 255).astype(np.uint8)
|
|
if img_np.shape[-1] == 1:
|
|
return Image.fromarray(img_np.squeeze(-1), 'L').convert('RGB')
|
|
else:
|
|
return Image.fromarray(img_np, 'RGB' if img_np.shape[-1] == 3 else 'RGBA').convert('RGB')
|
|
|
|
def pil_to_tensor_frame(self, pil_image: Image.Image) -> torch.Tensor:
|
|
img_np = np.array(pil_image.convert("RGB")).astype(np.float32) / 255.0
|
|
if img_np.ndim == 2:
|
|
img_np = np.expand_dims(img_np, axis=2)
|
|
img_np = np.repeat(img_np, 3, axis=2)
|
|
return torch.from_numpy(img_np)
|
|
|
|
def get_char_pixel_width(self, draw, char, font):
|
|
try:
|
|
bbox = draw.textbbox((0,0), char, font=font)
|
|
return bbox[2] - bbox[0]
|
|
except AttributeError:
|
|
width, _ = draw.textsize(char, font=font)
|
|
return width
|
|
|
|
def generate_two_pass_animation(self, image: torch.Tensor, char_width: int, font_path: str,
|
|
font_size: int, ascii_charset: str, background_color: str,
|
|
initial_text_color: str, invert_brightness_mapping: bool,
|
|
chars_per_frame: int, cursor_char: str,
|
|
cursor_blink_frames: int, cursor_color: str):
|
|
|
|
if not ascii_charset: raise ValueError("ASCII charset cannot be empty.")
|
|
if chars_per_frame < 1: chars_per_frame = 1
|
|
if cursor_blink_frames < 1: cursor_blink_frames = 1
|
|
cursor_char = cursor_char[0] if cursor_char else "_"
|
|
|
|
if image.ndim == 3: image = image.unsqueeze(0)
|
|
|
|
input_pil_image_rgb = self.tensor_to_pil(image[0])
|
|
original_width, original_height = input_pil_image_rgb.size
|
|
|
|
aspect_ratio = original_height / original_width
|
|
char_height_ascii_grid = max(1, int(char_width * aspect_ratio * 0.5))
|
|
|
|
# --- Generate the full ASCII map (char_str, image_color, char_width) ---
|
|
# This map is static for the whole animation.
|
|
resized_image_for_sampling = input_pil_image_rgb.resize((char_width, char_height_ascii_grid), Image.Resampling.LANCZOS)
|
|
grayscale_image = resized_image_for_sampling.convert("L")
|
|
color_sample_image = resized_image_for_sampling # Already RGB
|
|
|
|
try:
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
except IOError:
|
|
print(f"Warning: Font '{font_path}' not found. Falling back to default PIL font.")
|
|
try: font = ImageFont.load_default(font_size=font_size)
|
|
except AttributeError: font = ImageFont.load_default()
|
|
except TypeError: font = ImageFont.load_default()
|
|
|
|
temp_draw = ImageDraw.Draw(Image.new("RGB", (1,1)))
|
|
try:
|
|
metrics_bbox = temp_draw.textbbox((0,0), "My", font=font)
|
|
line_pixel_height = metrics_bbox[3] - metrics_bbox[1]
|
|
except AttributeError:
|
|
_, line_pixel_height = temp_draw.textsize("My", font=font)
|
|
if line_pixel_height == 0: line_pixel_height = font_size
|
|
|
|
flat_full_ascii_map = [] # Stores (char_str, image_sampled_color_tuple, char_pixel_width, x_pos, y_pos)
|
|
max_text_pixel_width = 0
|
|
|
|
gray_pixels = grayscale_image.load()
|
|
color_pixels = color_sample_image.load()
|
|
|
|
for y_idx in range(char_height_ascii_grid):
|
|
current_line_width_pixels = 0
|
|
y_pos_pixels = y_idx * line_pixel_height
|
|
for x_idx in range(char_width):
|
|
brightness = gray_pixels[x_idx, y_idx]
|
|
if invert_brightness_mapping: brightness = 255 - brightness
|
|
|
|
char_index = int((brightness / 255) * (len(ascii_charset) - 1))
|
|
char_to_draw_str = ascii_charset[char_index]
|
|
sampled_color_tuple = color_pixels[x_idx, y_idx]
|
|
|
|
char_pixel_w = self.get_char_pixel_width(temp_draw, char_to_draw_str, font)
|
|
if char_pixel_w == 0 and char_to_draw_str != ' ':
|
|
char_pixel_w = self.get_char_pixel_width(temp_draw, "M", font) // 2
|
|
|
|
flat_full_ascii_map.append(
|
|
(char_to_draw_str, sampled_color_tuple, char_pixel_w, current_line_width_pixels, y_pos_pixels)
|
|
)
|
|
current_line_width_pixels += char_pixel_w
|
|
max_text_pixel_width = max(max_text_pixel_width, current_line_width_pixels)
|
|
|
|
if not flat_full_ascii_map:
|
|
default_char_w = self.get_char_pixel_width(temp_draw, " ", font)
|
|
flat_full_ascii_map = [(' ', (128,128,128), default_char_w, 0, 0)]
|
|
max_text_pixel_width = default_char_w
|
|
|
|
total_chars_in_map = len(flat_full_ascii_map)
|
|
text_render_width = max(1, max_text_pixel_width)
|
|
text_render_height = max(1, char_height_ascii_grid * line_pixel_height)
|
|
|
|
try: bg_color_rgb = ImageColor.getrgb(background_color)
|
|
except ValueError: bg_color_rgb = (0,0,0)
|
|
try: initial_text_color_rgb = ImageColor.getrgb(initial_text_color)
|
|
except ValueError: initial_text_color_rgb = (255,255,255)
|
|
try: cursor_color_rgb = ImageColor.getrgb(cursor_color)
|
|
except ValueError: cursor_color_rgb = (255,255,255)
|
|
|
|
output_frame_tensors = []
|
|
main_state_progress = 0 # How many characters have been fully processed (i.e., colored with image color)
|
|
animation_frame_index = 0
|
|
|
|
while main_state_progress <= total_chars_in_map:
|
|
current_frame_pil = Image.new("RGB", (text_render_width, text_render_height), bg_color_rgb)
|
|
draw_frame = ImageDraw.Draw(current_frame_pil)
|
|
|
|
target_main_state_for_frame = min(total_chars_in_map, main_state_progress + chars_per_frame)
|
|
|
|
cursor_x_pixel = 0
|
|
cursor_y_pixel = 0
|
|
cursor_char_w = 0 # Width of the character the cursor is after
|
|
|
|
for char_idx in range(total_chars_in_map):
|
|
char_str, img_color, char_w, x_pos, y_pos = flat_full_ascii_map[char_idx]
|
|
|
|
current_char_color = None
|
|
if char_idx < main_state_progress: # Fully processed, use image color
|
|
current_char_color = img_color
|
|
elif char_idx < target_main_state_for_frame: # Being typed in initial color this frame
|
|
current_char_color = initial_text_color_rgb
|
|
|
|
if current_char_color:
|
|
draw_frame.text((x_pos, y_pos), char_str, font=font, fill=current_char_color)
|
|
|
|
if char_idx == target_main_state_for_frame - 1 and target_main_state_for_frame > 0 : # Cursor position
|
|
cursor_x_pixel = x_pos + char_w
|
|
cursor_y_pixel = y_pos
|
|
cursor_char_w = char_w # Store for potential adjustment if it's a space
|
|
|
|
# If target_main_state_for_frame is 0, cursor is at the very beginning (0,0)
|
|
if target_main_state_for_frame == 0:
|
|
cursor_x_pixel = 0
|
|
cursor_y_pixel = 0
|
|
# Get width of the first char for cursor if it exists, otherwise default cursor char width
|
|
if total_chars_in_map > 0:
|
|
_, _, first_char_w, _, _ = flat_full_ascii_map[0]
|
|
cursor_char_w = first_char_w
|
|
else:
|
|
cursor_char_w = self.get_char_pixel_width(draw_frame, cursor_char, font)
|
|
|
|
|
|
# --- Draw flashing cursor ---
|
|
# Cursor is drawn if not all characters are fully processed (colored)
|
|
# OR if all are processed but it's the very last step of revealing them
|
|
if target_main_state_for_frame < total_chars_in_map or \
|
|
(target_main_state_for_frame == total_chars_in_map and main_state_progress < total_chars_in_map):
|
|
is_cursor_blink_on = (animation_frame_index // cursor_blink_frames) % 2 == 0
|
|
if is_cursor_blink_on:
|
|
draw_frame.text((cursor_x_pixel, cursor_y_pixel),
|
|
cursor_char, font=font, fill=cursor_color_rgb)
|
|
|
|
final_frame_pil = current_frame_pil.resize((original_width, original_height), Image.Resampling.LANCZOS)
|
|
output_frame_tensors.append(self.pil_to_tensor_frame(final_frame_pil))
|
|
|
|
if main_state_progress >= total_chars_in_map: # All chars are in their final image color
|
|
# Add one last frame if cursor was on, to show text without cursor
|
|
if is_cursor_blink_on and (target_main_state_for_frame == total_chars_in_map):
|
|
current_frame_pil_no_cursor = Image.new("RGB", (text_render_width, text_render_height), bg_color_rgb)
|
|
draw_final = ImageDraw.Draw(current_frame_pil_no_cursor)
|
|
for char_str_f, img_color_f, _, x_pos_f, y_pos_f in flat_full_ascii_map:
|
|
draw_final.text((x_pos_f, y_pos_f), char_str_f, font=font, fill=img_color_f)
|
|
final_pil_no_cursor = current_frame_pil_no_cursor.resize((original_width, original_height), Image.Resampling.LANCZOS)
|
|
output_frame_tensors.append(self.pil_to_tensor_frame(final_pil_no_cursor))
|
|
break
|
|
|
|
main_state_progress = target_main_state_for_frame
|
|
animation_frame_index += 1
|
|
|
|
if not output_frame_tensors:
|
|
fallback_pil = Image.new("RGB", (original_width, original_height), bg_color_rgb)
|
|
output_frame_tensors.append(self.pil_to_tensor_frame(fallback_pil))
|
|
|
|
batched_output_tensor = torch.stack(output_frame_tensors, dim=0)
|
|
return (batched_output_tensor,)
|
|
|
|
# For testing the node independently (optional)
|
|
if __name__ == '__main__':
|
|
print("Testing TwoPassTypingColorASCIIAnimationNode locally...")
|
|
node = TwoPassTypingColorASCIIAnimationNode()
|
|
|
|
test_w, test_h = 100, 70
|
|
gradient_img = Image.new("RGB", (test_w, test_h))
|
|
gradient_draw = ImageDraw.Draw(gradient_img)
|
|
for i in range(test_w):
|
|
r = int((i / test_w) * 255)
|
|
for j in range(test_h):
|
|
g = int((j / test_h) * 255)
|
|
b = int(((i+j)/(test_w+test_h)) * 150) + 50 # Ensure some color
|
|
gradient_draw.point((i,j), fill=(r,g,b))
|
|
|
|
test_input_np = np.array(gradient_img).astype(np.float32) / 255.0
|
|
dummy_tensor_bhwc = torch.from_numpy(test_input_np).unsqueeze(0)
|
|
|
|
print(f"Input tensor shape for test: {dummy_tensor_bhwc.shape}")
|
|
|
|
try:
|
|
output_batched_tensor_tuple = node.generate_two_pass_animation(
|
|
image=dummy_tensor_bhwc,
|
|
char_width=35,
|
|
font_path="cour.ttf",
|
|
font_size=12,
|
|
ascii_charset=" ABCDEFGHIJKLMNOPQRSTUVWXYZ.:-=+*#%@",
|
|
background_color="#111111",
|
|
initial_text_color="#DDDDDD",
|
|
invert_brightness_mapping=False,
|
|
chars_per_frame=2, # How many new chars get the "initial_color" treatment per frame
|
|
cursor_char="|",
|
|
cursor_blink_frames=4,
|
|
cursor_color="#33FF33"
|
|
)
|
|
output_batched_tensor = output_batched_tensor_tuple[0]
|
|
print(f"Output batched tensor shape: {output_batched_tensor.shape}")
|
|
|
|
num_frames = output_batched_tensor.shape[0]
|
|
print(f"Generated {num_frames} frames.")
|
|
|
|
save_dir = "test_two_pass_frames"
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
print(f"Saving frames to ./{save_dir}/")
|
|
|
|
indices_to_save = []
|
|
if num_frames > 0:
|
|
indices_to_save.extend(range(min(10, num_frames)))
|
|
if num_frames > 20:
|
|
indices_to_save.extend(range(num_frames // 2 - 2, num_frames // 2 + 3))
|
|
indices_to_save.extend(range(max(0, num_frames - 10), num_frames))
|
|
indices_to_save = sorted(list(set(i for i in indices_to_save if 0 <= i < num_frames)))
|
|
|
|
for frame_idx in indices_to_save:
|
|
frame_tensor_hwc = output_batched_tensor[frame_idx]
|
|
frame_np = (frame_tensor_hwc.cpu().numpy() * 255).astype(np.uint8)
|
|
frame_pil = Image.fromarray(frame_np, 'RGB')
|
|
frame_pil.save(os.path.join(save_dir, f"two_pass_frame_{frame_idx:03d}.png"))
|
|
if indices_to_save:
|
|
print(f"Saved selected frames to '{save_dir}' directory.")
|
|
|
|
except Exception as e:
|
|
print(f"Error during local test: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|