Updated HF node + Created PDF Caption Layout

This commit is contained in:
Fillip
2024-08-20 12:55:12 -07:00
parent 4fc3fa2a8f
commit 712c5f22ee
4 changed files with 226 additions and 7 deletions
+4 -1
View File
@@ -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",
}
+36 -5
View File
@@ -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")
+184
View File
@@ -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
+2 -1
View File
@@ -7,4 +7,5 @@ scipy >=1.13.1
requests
aiohttp
moviepy
matplotlib
matplotlib
reportlab