From 712c5f22eea4e51f01816b4966cfe574b2e6bcc4 Mon Sep 17 00:00:00 2001 From: Fillip Date: Tue, 20 Aug 2024 12:55:12 -0700 Subject: [PATCH] Updated HF node + Created PDF Caption Layout --- __init__.py | 5 +- nodes/FL_HF_Character.py | 41 ++++++- nodes/FL_ImageCaptionLayoutPDF.py | 184 ++++++++++++++++++++++++++++++ requirements.txt | 3 +- 4 files changed, 226 insertions(+), 7 deletions(-) create mode 100644 nodes/FL_ImageCaptionLayoutPDF.py diff --git a/__init__.py b/__init__.py index e47c95c..d080817 100644 --- a/__init__.py +++ b/__init__.py @@ -54,6 +54,7 @@ from .nodes.FL_SaveCSV import FL_SaveCSV from. nodes.FL_KSamplerXYZPlot import FL_KSamplerXYZPlot from .nodes.FL_SamplerStrings import FL_SamplerStrings from .nodes.FL_SchedulerStrings import FL_SchedulerStrings +from .nodes.FL_ImageCaptionLayoutPDF import FL_ImageCaptionLayoutPDF @@ -115,6 +116,7 @@ NODE_CLASS_MAPPINGS = { "FL_KSamplerXYZPlot": FL_KSamplerXYZPlot, "FL_SamplerStrings": FL_SamplerStrings, "FL_SchedulerStrings": FL_SchedulerStrings, + "FL_ImageCaptionLayoutPDF": FL_ImageCaptionLayoutPDF, } @@ -174,7 +176,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FL_SaveCSV": "FL Save CSV", "FL_KSamplerXYZPlot": "FL KSampler XYZ Plot", "FL_SamplerStrings": "FL Sampler String XYZ", - "FL_SchedulerStrings": "FL Scheduler String XYZ" + "FL_SchedulerStrings": "FL Scheduler String XYZ", + "FL_ImageCaptionLayoutPDF": "FL Image Caption Layout PDF", } diff --git a/nodes/FL_HF_Character.py b/nodes/FL_HF_Character.py index b91e315..bf7b104 100644 --- a/nodes/FL_HF_Character.py +++ b/nodes/FL_HF_Character.py @@ -17,7 +17,6 @@ from PIL import Image from huggingface_hub import HfApi, create_repo, repo_exists from tqdm import tqdm - class FL_HF_Character: @classmethod def INPUT_TYPES(cls): @@ -36,6 +35,7 @@ class FL_HF_Character: "lora_file": ("STRING", {"default": ""}), "dataset_zip": ("ZIP",), "caption_layout": ("IMAGE",), + "caption_PDF_layout": ("PDF",), "csv_file": ("CSV",), } } @@ -47,7 +47,8 @@ class FL_HF_Character: def upload_to_hub(self, api_key: str, owner: str, repo_name: str, studio_name: str, project_name: str, character_name: str, create_new_repo: str, repo_type: str, lora_file: str = "", dataset_zip: bytes = None, - caption_layout: torch.Tensor = None, csv_file: bytes = None) -> tuple[str]: + caption_layout: torch.Tensor = None, caption_PDF_layout: bytes = None, + csv_file: bytes = None) -> tuple[str]: # Initialize Hugging Face API api = HfApi(token=api_key) @@ -62,8 +63,7 @@ class FL_HF_Character: print(f"Repository created or already exists: {repo_url}") else: if not repo_exists(repo_id=full_repo_id, token=api_key): - return ( - f"Error: Repository {full_repo_id} does not exist. Please create it first or use the 'Create New Repo' option.",) + return (f"Error: Repository {full_repo_id} does not exist. Please create it first or use the 'Create New Repo' option.",) repo_url = f"https://huggingface.co/{full_repo_id}" print(f"Using existing repository: {repo_url}") @@ -77,6 +77,8 @@ class FL_HF_Character: self.upload_zip(api, dataset_zip, f"{base_path}/dataset", full_repo_id, api_key, "Dataset") if caption_layout is not None: self.upload_image(api, caption_layout, base_path, full_repo_id, api_key, "caption_layout") + if caption_PDF_layout is not None: + self.upload_pdf(api, caption_PDF_layout, base_path, full_repo_id, api_key, "caption_PDF_layout") if csv_file is not None: self.upload_csv(api, csv_file, base_path, full_repo_id, api_key) @@ -164,6 +166,35 @@ class FL_HF_Character: ) print(f"{image_type} uploaded successfully") + def upload_pdf(self, api, pdf_data, repo_dir, full_repo_id, api_key, pdf_type): + repo_path = f"{repo_dir}/{pdf_type}.pdf" + + pbar = tqdm(total=100, unit='%', desc=f"Uploading {pdf_type} PDF") + + def update_progress(): + progress = 0 + while progress < 95: + time.sleep(0.5) + increment = min(5, 95 - progress) + progress += increment + pbar.update(increment) + + progress_thread = threading.Thread(target=update_progress) + progress_thread.start() + + api.upload_file( + path_or_fileobj=pdf_data, + path_in_repo=repo_path, + repo_id=full_repo_id, + token=api_key + ) + + progress_thread.join() + pbar.update(100 - pbar.n) + pbar.close() + + print(f"{pdf_type} PDF uploaded successfully to {repo_path}") + def upload_csv(self, api, csv_data, repo_dir, full_repo_id, api_key): repo_path = f"{repo_dir}/metadata.csv" @@ -195,5 +226,5 @@ class FL_HF_Character: @classmethod def IS_CHANGED(cls, api_key, owner, repo_name, studio_name, project_name, character_name, - create_new_repo, repo_type, lora_file, dataset_zip, caption_layout, csv_file): + create_new_repo, repo_type, lora_file, dataset_zip, caption_layout, caption_PDF_layout, csv_file): return float("NaN") \ No newline at end of file diff --git a/nodes/FL_ImageCaptionLayoutPDF.py b/nodes/FL_ImageCaptionLayoutPDF.py new file mode 100644 index 0000000..68a897a --- /dev/null +++ b/nodes/FL_ImageCaptionLayoutPDF.py @@ -0,0 +1,184 @@ +import os +import math +import torch +import numpy as np +from PIL import Image, ImageDraw, ImageFont +import textwrap +from reportlab.lib.pagesizes import letter, portrait, landscape +from reportlab.pdfgen import canvas +from reportlab.lib.units import inch +from reportlab.pdfbase import pdfmetrics +from reportlab.pdfbase.ttfonts import TTFont +from io import BytesIO + + +class FL_ImageCaptionLayoutPDF: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_directory": ("STRING", {"default": ""}), + "images_per_row": ("INT", {"default": 6, "min": 1, "max": 10}), + "display_size": ("INT", {"default": 100, "min": 64, "max": 512}), + "caption_height": ("INT", {"default": 80, "min": 32, "max": 256}), + "font_size": ("INT", {"default": 4, "min": 4, "max": 32}), + "padding": ("INT", {"default": 10, "min": 0, "max": 100}), + "output_directory": ("STRING", {"default": ""}), + "output_filename": ("STRING", {"default": "output"}), + "orientation": (["horizontal", "vertical"], {"default": "horizontal"}), + }, + } + + RETURN_TYPES = ("STRING", "IMAGE", "PDF") + FUNCTION = "create_layout" + CATEGORY = "🏵️Fill Nodes/Captioning" + OUTPUT_NODE = True + + def create_layout(self, image_directory, images_per_row, display_size, caption_height, font_size, padding, + output_directory, output_filename, orientation): + output_path, pdf_bytes = self.create_pdf_layout(image_directory, images_per_row, display_size, caption_height, + font_size, padding, output_directory, output_filename, + orientation) + preview_image = self.create_image_preview(image_directory, images_per_row, display_size, caption_height, + font_size, padding, orientation) + return (output_path, preview_image, pdf_bytes) + + def create_pdf_layout(self, image_directory, images_per_row, display_size, caption_height, font_size, padding, + output_directory, output_filename, orientation): + # Get the path to the fonts directory + current_dir = os.path.dirname(os.path.abspath(__file__)) + fonts_dir = os.path.join(os.path.dirname(current_dir), "fonts") + font_path = os.path.join(fonts_dir, "arial.ttf") + + # Check if the font file exists + if not os.path.exists(font_path): + raise FileNotFoundError(f"Font file not found: {font_path}") + + # Register the font + pdfmetrics.registerFont(TTFont('Arial', font_path)) + + # Get all image files and their corresponding caption files + image_files = [f for f in os.listdir(image_directory) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] + image_files.sort() # Sort files to ensure consistent order + + # PDF setup + page_size = landscape(letter) if orientation == "horizontal" else portrait(letter) + width, height = page_size + output_path = os.path.join(output_directory, f"{output_filename}.pdf") + pdf_buffer = BytesIO() + c = canvas.Canvas(pdf_buffer, pagesize=page_size) + + # Calculate layout dimensions + display_size_pt = display_size + caption_height_pt = caption_height + padding_pt = padding + item_width = display_size_pt + padding_pt + item_height = display_size_pt + caption_height_pt + padding_pt + + # Calculate how many items can fit on a page + items_per_row = min(images_per_row, math.floor((width - padding_pt) / item_width)) + rows_per_page = math.floor((height - padding_pt) / item_height) + items_per_page = items_per_row * rows_per_page + + for i in range(0, len(image_files), items_per_page): + page_images = image_files[i:i + items_per_page] + + for j, image_file in enumerate(page_images): + # Calculate position + row = j // items_per_row + col = j % items_per_row + x = padding_pt + col * item_width + y = height - padding_pt - (row + 1) * item_height + + # Load and draw image + img_path = os.path.join(image_directory, image_file) + img = Image.open(img_path) + aspect_ratio = img.width / img.height + display_height = display_size_pt / aspect_ratio + + c.drawImage(img_path, x, y + caption_height_pt, width=display_size_pt, height=display_height, + preserveAspectRatio=True, anchor='sw') + + # Load caption + caption_file = os.path.splitext(image_file)[0] + '.txt' + caption_path = os.path.join(image_directory, caption_file) + try: + with open(caption_path, 'r') as f: + caption = f.read().strip() + except FileNotFoundError: + caption = "No caption found" + + # Draw caption + c.setFont("Arial", font_size) + text_object = c.beginText(x, y + caption_height_pt - font_size) + wrapped_text = textwrap.fill(caption, width=int(display_size_pt / (font_size * 0.6))) + for line in wrapped_text.split('\n'): + text_object.textLine(line) + c.drawText(text_object) + + c.showPage() # Start a new page after each set of items + + c.save() + pdf_bytes = pdf_buffer.getvalue() + + # Save the PDF to file + with open(output_path, 'wb') as f: + f.write(pdf_bytes) + + print(f"PDF saved as {output_path}") + return output_path, pdf_bytes + + def create_image_preview(self, image_directory, images_per_row, display_size, caption_height, font_size, padding, + orientation): + # Get all image files and their corresponding caption files + image_files = [f for f in os.listdir(image_directory) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] + image_files.sort() # Sort files to ensure consistent order + + # Calculate layout dimensions + total_width = images_per_row * (display_size + padding) + padding + rows = (len(image_files) + images_per_row - 1) // images_per_row + total_height = rows * (display_size + caption_height + padding) + padding + + # Create the layout + layout = Image.new('RGB', (total_width, total_height), color=(255, 255, 255)) + + # Load font + font_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "fonts", "arial.ttf") + try: + font = ImageFont.truetype(font_path, font_size) + except IOError: + font = ImageFont.load_default() + + for i, image_file in enumerate(image_files): + # Load and resize image + img_path = os.path.join(image_directory, image_file) + img = Image.open(img_path).convert('RGB') + img.thumbnail((display_size, display_size), Image.LANCZOS) + + # Calculate position + row = i // images_per_row + col = i % images_per_row + x = padding + col * (display_size + padding) + y = padding + row * (display_size + caption_height + padding) + + # Paste image + layout.paste(img, (x, y)) + + # Load caption + caption_file = os.path.splitext(image_file)[0] + '.txt' + caption_path = os.path.join(image_directory, caption_file) + try: + with open(caption_path, 'r') as f: + caption = f.read().strip() + except FileNotFoundError: + caption = "No caption found" + + # Draw caption + draw = ImageDraw.Draw(layout) + wrapped_text = textwrap.fill(caption, width=int(display_size / (font_size * 0.6))) + draw.text((x, y + display_size + 5), wrapped_text, font=font, fill=(0, 0, 0)) + + # Convert to tensor for preview + preview_tensor = torch.from_numpy(np.array(layout).astype(np.float32) / 255.0).unsqueeze(0) + + return preview_tensor \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 2016299..d3d3dce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ scipy >=1.13.1 requests aiohttp moviepy -matplotlib \ No newline at end of file +matplotlib +reportlab \ No newline at end of file