203 lines
9.4 KiB
Python
203 lines
9.4 KiB
Python
import folder_paths
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import server # Required for preview generation
|
|
import threading
|
|
import shutil # For potential file operations
|
|
from datetime import datetime # For timestamps
|
|
|
|
# --- Constants ---
|
|
# Determine the base path of the custom node
|
|
EXTENSION_PATH = os.path.normpath(os.path.dirname(__file__))
|
|
DEFAULT_HISTORY_FOLDER = os.path.join(EXTENSION_PATH, "history_folder")
|
|
|
|
# Tensor to PIL
|
|
def tensor2pil(image):
|
|
if image.dim() > 3:
|
|
image = image[0]
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
# Placeholder generation function (remains unused in the optimized path, keep for robustness?)
|
|
# Note: Placeholder generation is removed from the main optimized logic below.
|
|
# If an error occurs during copy, the image is simply skipped in the preview.
|
|
def create_placeholder(size=(128, 128), text="?"):
|
|
"""Creates a simple placeholder PIL image."""
|
|
img = Image.new('RGB', size, color = (40, 40, 40))
|
|
d = ImageDraw.Draw(img)
|
|
font = ImageFont.load_default() # Keep it simple for placeholders
|
|
try:
|
|
# Basic centering for default font
|
|
tw, th = d.textbbox((0,0), text, font=font)[2:] if hasattr(d, 'textbbox') else (10, 10) # Use textbbox if available
|
|
d.text(((size[0]-tw)/2, (size[1]-th)/2), text, font=font, fill=(180, 180, 180))
|
|
except Exception as e:
|
|
print(f"[PreviewHistory] Error drawing placeholder text: {e}")
|
|
d.text((10, 10), text, fill=(180, 180, 180)) # Fallback position
|
|
return img
|
|
|
|
|
|
class PreviewHistory:
|
|
# Lock for file system operations in the target directory
|
|
_dir_lock = threading.Lock()
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_temp_directory() # For UI previews
|
|
self.type = "temp"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"history_size": ("INT", {
|
|
"default": 9,
|
|
"min": 1,
|
|
"max": 1000000,
|
|
"step": 1,
|
|
"tooltip": "How many images to keep in the history folder."
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "execute"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Raffle/Previews"
|
|
|
|
# --- Main Execution ---
|
|
|
|
def execute(self, image, history_size):
|
|
|
|
history_folder = DEFAULT_HISTORY_FOLDER
|
|
|
|
# Ensure target directory exists
|
|
if not os.path.exists(history_folder):
|
|
try:
|
|
os.makedirs(history_folder, exist_ok=True)
|
|
print(f"[PreviewHistory] Created history directory: {history_folder}")
|
|
except OSError as e:
|
|
print(f"[PreviewHistory] Error creating directory {history_folder}: {e}. Cannot proceed.")
|
|
# Return empty previews if directory creation fails
|
|
return {"ui": {"images": []}} # Return empty list
|
|
|
|
# --- Save New Image (if provided) ---
|
|
new_image_saved_path = None
|
|
if image is not None and image.nelement() > 0:
|
|
new_pil_image = tensor2pil(image)
|
|
with PreviewHistory._dir_lock: # Lock specifically for saving
|
|
try:
|
|
now = datetime.now()
|
|
timestamp_final = now.strftime("%d-%m-%Y_%H-%M-%S")
|
|
# Include milliseconds for higher uniqueness, preventing rare collisions
|
|
timestamp_final += f"_{now.microsecond // 1000:03d}"
|
|
new_filename = f"history_{timestamp_final}.png"
|
|
new_path = os.path.join(history_folder, new_filename)
|
|
new_pil_image.save(new_path, "PNG", compress_level=1) # Faster compression
|
|
new_image_saved_path = new_path # Keep track of the path if saved
|
|
# print(f"[PreviewHistory] Saved new file: {new_path}") # Debug
|
|
except Exception as e:
|
|
print(f"[PreviewHistory] Error saving new image to '{history_folder}': {e}")
|
|
# Continue even if saving fails, try to show existing history
|
|
|
|
# --- Cleanup Old Files & Get List for Preview ---
|
|
sorted_history_files = [] # List to hold full paths of files to preview
|
|
with PreviewHistory._dir_lock: # Lock for listing, cleanup
|
|
try:
|
|
# Get all .png files with modification times for cleanup
|
|
all_files = []
|
|
for filename in os.listdir(history_folder):
|
|
if filename.lower().endswith(".png"):
|
|
full_path = os.path.join(history_folder, filename)
|
|
try:
|
|
if os.path.isfile(full_path):
|
|
# Use creation time if available and potentially more stable, else modification time
|
|
# Note: ctime might be platform-dependent (inode change on Unix, creation on Win)
|
|
# Stick to mtime for broader consistency unless ctime is specifically desired
|
|
mod_time = os.path.getmtime(full_path)
|
|
all_files.append((mod_time, full_path))
|
|
except OSError:
|
|
print(f"[PreviewHistory] Warning: Could not access file {full_path} during scan.")
|
|
continue # Skip file if cannot get mod_time
|
|
|
|
# Sort by modification time, newest first
|
|
all_files.sort(key=lambda x: x[0], reverse=True)
|
|
|
|
# Determine which files to keep based on history_size
|
|
files_to_keep_info = all_files[:history_size]
|
|
files_to_remove_info = all_files[history_size:]
|
|
|
|
# Cleanup: Remove files exceeding the current history_size
|
|
if files_to_remove_info:
|
|
# print(f"[PreviewHistory] Found {len(all_files)} files, keeping {history_size}, removing {len(files_to_remove_info)}.") # Debug
|
|
for mod_time, path_to_remove in files_to_remove_info:
|
|
try:
|
|
# print(f"[PreviewHistory] Removing old file: {path_to_remove}") # Debug
|
|
os.remove(path_to_remove)
|
|
except OSError as e:
|
|
print(f"[PreviewHistory] Error removing old file {path_to_remove}: {e}")
|
|
|
|
# Get the final list of full paths for the files we are keeping for the preview
|
|
sorted_history_files = [f[1] for f in files_to_keep_info]
|
|
|
|
except Exception as e:
|
|
print(f"[PreviewHistory] Error listing/cleaning files in {history_folder}: {e}")
|
|
# Fall through with empty list if error during file operations
|
|
|
|
|
|
# --- Generate Preview Data for the UI by Copying Files ---
|
|
previews = []
|
|
preview_server = server.PromptServer.instance # Get server instance
|
|
|
|
# Process the sorted list of history files to generate previews
|
|
for i, history_file_path in enumerate(sorted_history_files):
|
|
if not os.path.exists(history_file_path):
|
|
print(f"[PreviewHistory] Warning: File {history_file_path} not found during preview generation (maybe removed?). Skipping.")
|
|
continue
|
|
|
|
try:
|
|
# 1. Get image dimensions efficiently
|
|
# Use a context manager to ensure the file handle is closed
|
|
# PIL often reads headers without loading the full image data
|
|
with Image.open(history_file_path) as img:
|
|
width, height = img.size
|
|
|
|
# 2. Determine temporary path information using dimensions
|
|
# Use a distinct prefix based on the original filename hash or index to avoid collisions in temp dir
|
|
# Using index 'i' is simple and effective here.
|
|
filename_prefix = f"PreviewHistory_Item_{i:02d}_"
|
|
full_output_folder, fname, count, subfolder, _ = folder_paths.get_save_image_path(
|
|
filename_prefix, self.output_dir, width, height
|
|
)
|
|
temp_filename = f"{fname}_{count:05}_.png" # Previews are always PNG
|
|
temp_file_path = os.path.join(full_output_folder, temp_filename)
|
|
|
|
# 3. Copy the file from history folder to temp folder
|
|
# shutil.copy2 preserves more metadata (like mtime), copy is slightly faster if metadata isn't needed.
|
|
# Using copy should be sufficient here.
|
|
shutil.copy(history_file_path, temp_file_path)
|
|
|
|
# 4. Append preview info for this image
|
|
previews.append({
|
|
"filename": temp_filename,
|
|
"subfolder": subfolder,
|
|
"type": self.type
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"[PreviewHistory] Error processing history image '{history_file_path}' for preview: {e}. Skipping this image.")
|
|
# Optionally, could create and copy a placeholder image here, but skipping is simpler/faster.
|
|
|
|
|
|
# Return the list of preview data dictionaries
|
|
# print(f"[PreviewHistory] Generated {len(previews)} previews.") # Debug
|
|
return {"ui": {"images": previews}}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PreviewHistory": PreviewHistory
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PreviewHistory": "Preview History (Raffle)"
|
|
} |