Updated HF node + Created PDF Caption Layout
This commit is contained in:
+4
-1
@@ -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",
|
||||
|
||||
}
|
||||
|
||||
|
||||
@@ -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")
|
||||
@@ -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
@@ -7,4 +7,5 @@ scipy >=1.13.1
|
||||
requests
|
||||
aiohttp
|
||||
moviepy
|
||||
matplotlib
|
||||
matplotlib
|
||||
reportlab
|
||||
Reference in New Issue
Block a user