228 lines
11 KiB
Python
228 lines
11 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont, ImageColor
|
|
import os
|
|
|
|
class ASCIIAnimationGeneratorNode:
|
|
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"}), # Max width reduced for performance
|
|
"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}),
|
|
"text_color": ("STRING", {"default": "#FFFFFF", "multiline": False}),
|
|
"invert_brightness_mapping": ("BOOLEAN", {"default": False}),
|
|
"chars_per_frame": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1, "display":"slider"}), # New: control typing speed
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",) # Output is a batch of images
|
|
RETURN_NAMES = ("animated_ascii_frames",)
|
|
FUNCTION = "generate_ascii_animation"
|
|
CATEGORY = "image/animation" # Or image/art
|
|
|
|
def tensor_to_pil(self, tensor_image: torch.Tensor) -> Image.Image:
|
|
"""Converts a HWC PyTorch tensor (0-1 float) from the input batch to a PIL Image."""
|
|
img_np = tensor_image.cpu().numpy()
|
|
# Input tensor_image is expected to be HWC (from image[0])
|
|
img_np = (img_np * 255).astype(np.uint8)
|
|
|
|
if img_np.shape[-1] == 1: # Grayscale
|
|
return Image.fromarray(img_np.squeeze(-1), 'L')
|
|
else: # RGB or RGBA
|
|
return Image.fromarray(img_np, 'RGB' if img_np.shape[-1] == 3 else 'RGBA')
|
|
|
|
def pil_to_tensor_frame(self, pil_image: Image.Image) -> torch.Tensor:
|
|
"""Converts a PIL Image to a HWC PyTorch tensor (0-1 float) for a single animation frame."""
|
|
img_np = np.array(pil_image).astype(np.float32) / 255.0
|
|
if img_np.ndim == 2: # Grayscale image, add channel dimension
|
|
img_np = np.expand_dims(img_np, axis=2)
|
|
|
|
tensor_image = torch.from_numpy(img_np) # HWC
|
|
return tensor_image
|
|
|
|
def generate_ascii_animation(self, image: torch.Tensor, char_width: int, font_path: str,
|
|
font_size: int, ascii_charset: str, background_color: str,
|
|
text_color: str, invert_brightness_mapping: bool,
|
|
chars_per_frame: int):
|
|
|
|
# 0. Basic Input Validation
|
|
if not ascii_charset:
|
|
raise ValueError("ASCII charset cannot be empty.")
|
|
if chars_per_frame < 1:
|
|
chars_per_frame = 1
|
|
|
|
# 1. Convert Input Tensor to PIL Image (use the first image in the batch)
|
|
if image.ndim == 3: # Should be (B,H,W,C), if (H,W,C) add batch
|
|
image = image.unsqueeze(0)
|
|
|
|
input_pil_image = self.tensor_to_pil(image[0]) # image[0] is HWC
|
|
original_width, original_height = input_pil_image.size
|
|
|
|
# 2. Convert PIL Image to Full ASCII Representation (list of strings)
|
|
aspect_ratio = original_height / original_width
|
|
# Character cell aspect ratio adjustment (0.5 means chars are ~twice as tall as wide)
|
|
char_height = max(1, int(char_width * aspect_ratio * 0.5))
|
|
|
|
resized_image_for_ascii = input_pil_image.resize((char_width, char_height), Image.Resampling.LANCZOS)
|
|
grayscale_image = resized_image_for_ascii.convert("L")
|
|
|
|
full_ascii_lines = []
|
|
pixels = grayscale_image.load()
|
|
for y_idx in range(char_height):
|
|
line_of_text = ""
|
|
for x_idx in range(char_width):
|
|
brightness = pixels[x_idx, y_idx]
|
|
if invert_brightness_mapping:
|
|
brightness = 255 - brightness
|
|
char_index = int((brightness / 255) * (len(ascii_charset) - 1))
|
|
line_of_text += ascii_charset[char_index]
|
|
full_ascii_lines.append(line_of_text)
|
|
|
|
if not any(full_ascii_lines): # If all lines are empty
|
|
full_ascii_lines = [" "] # Ensure at least one space to avoid errors later
|
|
|
|
# 3. Prepare for Rendering (Font, Colors, Base Text Dimensions)
|
|
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()
|
|
|
|
try:
|
|
bg_color_rgb = ImageColor.getrgb(background_color)
|
|
text_color_rgb = ImageColor.getrgb(text_color)
|
|
except ValueError:
|
|
print(f"Warning: Invalid color string. Using black background/white text.")
|
|
bg_color_rgb = (0,0,0)
|
|
text_color_rgb = (255,255,255)
|
|
|
|
# Calculate dimensions of the text block if all ASCII art was rendered
|
|
temp_draw = ImageDraw.Draw(Image.new("RGB", (1,1)))
|
|
max_text_pixel_width = 0
|
|
for line in full_ascii_lines:
|
|
try: line_bbox = temp_draw.textbbox((0,0), line, font=font); width = line_bbox[2] - line_bbox[0]
|
|
except AttributeError: width, _ = temp_draw.textsize(line, font=font) # Older Pillow
|
|
max_text_pixel_width = max(max_text_pixel_width, width)
|
|
|
|
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) # Older Pillow
|
|
|
|
if line_pixel_height == 0 and font_size > 0: line_pixel_height = int(font_size * 1.2)
|
|
if max_text_pixel_width == 0: max_text_pixel_width = char_width * font_size // 2
|
|
if line_pixel_height == 0: line_pixel_height = font_size
|
|
|
|
text_render_width = max(1, max_text_pixel_width)
|
|
text_render_height = max(1, len(full_ascii_lines) * line_pixel_height)
|
|
|
|
# 4. Generate Animation Frames
|
|
output_frame_tensors = []
|
|
total_chars_to_type = sum(len(line) for line in full_ascii_lines)
|
|
chars_typed_so_far = 0
|
|
|
|
# Always generate at least one frame, even if it's just the background
|
|
# or the first few characters if total_chars_to_type is small.
|
|
|
|
while chars_typed_so_far <= total_chars_to_type:
|
|
current_frame_pil = Image.new("RGB", (text_render_width, text_render_height), bg_color_rgb)
|
|
draw_frame = ImageDraw.Draw(current_frame_pil)
|
|
|
|
chars_drawn_this_frame_total = 0
|
|
temp_chars_typed_count = 0 # Relative to start of full_ascii_lines
|
|
|
|
for line_idx, line_content in enumerate(full_ascii_lines):
|
|
chars_to_draw_on_this_line = 0
|
|
if temp_chars_typed_count < chars_typed_so_far:
|
|
remaining_to_type_on_line = chars_typed_so_far - temp_chars_typed_count
|
|
chars_to_draw_on_this_line = min(len(line_content), remaining_to_type_on_line)
|
|
|
|
if chars_to_draw_on_this_line > 0:
|
|
draw_frame.text((0, line_idx * line_pixel_height),
|
|
line_content[:chars_to_draw_on_this_line],
|
|
font=font, fill=text_color_rgb)
|
|
|
|
temp_chars_typed_count += len(line_content)
|
|
if temp_chars_typed_count >= chars_typed_so_far:
|
|
break # Stop processing lines if all typed characters for this frame are drawn
|
|
|
|
# Resize the rendered text frame to original input dimensions
|
|
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 chars_typed_so_far >= total_chars_to_type:
|
|
break # Animation complete
|
|
|
|
chars_typed_so_far += chars_per_frame
|
|
if chars_typed_so_far > total_chars_to_type and (chars_typed_so_far - chars_per_frame) < total_chars_to_type :
|
|
chars_typed_so_far = total_chars_to_type # Ensure the last frame shows everything
|
|
|
|
# 5. Batch and Return
|
|
if not output_frame_tensors: # Should not happen with the new loop logic, but as a fallback
|
|
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) # (num_frames, H, W, C)
|
|
return (batched_output_tensor,)
|
|
|
|
# For testing the node independently (optional)
|
|
if __name__ == '__main__':
|
|
print("Testing ASCIIAnimationGeneratorNode locally...")
|
|
node = ASCIIAnimationGeneratorNode()
|
|
|
|
# Create a dummy input PIL image for testing
|
|
test_input_pil = Image.new('RGB', (100, 75), color = 'darkcyan') # W, H
|
|
# Convert to ComfyUI-like tensor (B, H, W, C)
|
|
test_input_np = np.array(test_input_pil).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_ascii_animation(
|
|
image=dummy_tensor_bhwc,
|
|
char_width=40, # Smaller for quicker test
|
|
font_path="cour.ttf",
|
|
font_size=10,
|
|
ascii_charset=" .:oO0@",
|
|
background_color="#202020",
|
|
text_color="#33FF33",
|
|
invert_brightness_mapping=False,
|
|
chars_per_frame=5 # Type 5 characters per frame
|
|
)
|
|
output_batched_tensor = output_batched_tensor_tuple[0]
|
|
print(f"Output batched tensor shape: {output_batched_tensor.shape}") # (num_frames, H, W, C)
|
|
|
|
num_frames = output_batched_tensor.shape[0]
|
|
print(f"Generated {num_frames} frames.")
|
|
|
|
# To save a few frames for visual inspection (e.g., first, middle, last):
|
|
if num_frames > 0:
|
|
indices_to_save = [0]
|
|
if num_frames > 1: indices_to_save.append(num_frames // 2)
|
|
if num_frames > 2: indices_to_save.append(num_frames - 1)
|
|
indices_to_save = sorted(list(set(indices_to_save))) # Unique sorted
|
|
|
|
for i, frame_idx in enumerate(indices_to_save):
|
|
if frame_idx < num_frames:
|
|
frame_tensor_hwc = output_batched_tensor[frame_idx] # H, W, C
|
|
# Need a tensor_to_pil that handles HWC directly for saving
|
|
frame_pil = node.tensor_to_pil(frame_tensor_hwc) # Re-use existing one
|
|
frame_pil.save(f"test_animation_frame_{i+1}_idx{frame_idx}.png")
|
|
print(f"Saved test_animation_frame_{i+1}_idx{frame_idx}.png")
|
|
|
|
except Exception as e:
|
|
print(f"Error during local test: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|