252 lines
8.5 KiB
Python
252 lines
8.5 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import posixpath
|
|
import re
|
|
import shutil
|
|
import tempfile
|
|
import zipfile
|
|
from pathlib import PurePosixPath
|
|
|
|
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
|
_HIDDEN_PARTS = {"__MACOSX"}
|
|
_INDEX_PART_PATTERN = re.compile(r"^\d+(?:-\d+)?$")
|
|
|
|
|
|
def ensure_temp_batch_dir(temp_root, prefix="batchimages_"):
|
|
if not temp_root:
|
|
raise ValueError("Temp directory is not configured.")
|
|
os.makedirs(temp_root, exist_ok=True)
|
|
return tempfile.mkdtemp(prefix=prefix, dir=temp_root)
|
|
|
|
|
|
def is_supported_image_file(path):
|
|
suffix = PurePosixPath(path).suffix.lower()
|
|
return suffix in IMAGE_EXTENSIONS
|
|
|
|
|
|
def should_skip_zip_member(member_name):
|
|
if not member_name or member_name.endswith("/"):
|
|
return True
|
|
|
|
path = PurePosixPath(member_name)
|
|
parts = path.parts
|
|
if not parts:
|
|
return True
|
|
|
|
for part in parts:
|
|
if part in ("", "."):
|
|
continue
|
|
if part.startswith(".") or part in _HIDDEN_PARTS:
|
|
return True
|
|
|
|
return not is_supported_image_file(member_name)
|
|
|
|
|
|
def validate_zip_member_path(member_name):
|
|
normalized = member_name.replace("\\", "/")
|
|
path = PurePosixPath(normalized)
|
|
|
|
if path.is_absolute():
|
|
raise ValueError(f"Zip member uses absolute path: {member_name}")
|
|
|
|
parts = path.parts
|
|
if not parts:
|
|
raise ValueError("Zip member path is empty.")
|
|
|
|
for part in parts:
|
|
if part in ("", "."):
|
|
continue
|
|
if part == "..":
|
|
raise ValueError(f"Zip member escapes extraction root: {member_name}")
|
|
|
|
return normalized
|
|
|
|
|
|
def sanitize_relative_path(relative_path):
|
|
normalized = posixpath.normpath(relative_path.replace("\\", "/")).lstrip("/")
|
|
if normalized in ("", "."):
|
|
raise ValueError("Relative path is empty.")
|
|
|
|
parts = []
|
|
for part in normalized.split("/"):
|
|
if part in ("", "."):
|
|
continue
|
|
if part == "..":
|
|
raise ValueError(f"Relative path escapes root: {relative_path}")
|
|
if part.startswith("."):
|
|
raise ValueError(f"Hidden path segment is not allowed: {relative_path}")
|
|
safe_part = re.sub(r"[^A-Za-z0-9._ -]", "_", part).strip()
|
|
safe_part = safe_part.rstrip(". ")
|
|
if not safe_part:
|
|
safe_part = "item"
|
|
parts.append(safe_part)
|
|
|
|
if not parts:
|
|
raise ValueError(f"Relative path has no usable segments: {relative_path}")
|
|
|
|
return "/".join(parts)
|
|
|
|
|
|
def build_manifest_entry(extracted_root, absolute_path, relative_path, index, total):
|
|
safe_relative_path = sanitize_relative_path(relative_path)
|
|
relative_parent = posixpath.dirname(relative_path)
|
|
safe_relative_parent = posixpath.dirname(safe_relative_path)
|
|
stem, suffix = os.path.splitext(posixpath.basename(relative_path))
|
|
|
|
return {
|
|
"source_path": os.path.abspath(absolute_path),
|
|
"relative_path": relative_path,
|
|
"safe_relative_path": safe_relative_path,
|
|
"relative_parent": "" if relative_parent == "." else relative_parent,
|
|
"safe_relative_parent": "" if safe_relative_parent == "." else safe_relative_parent,
|
|
"stem": stem,
|
|
"suffix": suffix.lower(),
|
|
"index": index,
|
|
"total": total,
|
|
"batch_root": os.path.abspath(extracted_root),
|
|
}
|
|
|
|
|
|
def extract_images_from_zip(zip_path, temp_root):
|
|
if not zipfile.is_zipfile(zip_path):
|
|
raise ValueError(f"Invalid zip file: {zip_path}")
|
|
|
|
extract_root = ensure_temp_batch_dir(temp_root)
|
|
extracted_items = []
|
|
|
|
try:
|
|
with zipfile.ZipFile(zip_path) as archive:
|
|
for member in archive.infolist():
|
|
validated_name = validate_zip_member_path(member.filename)
|
|
if should_skip_zip_member(validated_name):
|
|
continue
|
|
|
|
destination = os.path.abspath(os.path.join(extract_root, validated_name.replace("/", os.sep)))
|
|
if os.path.commonpath((extract_root, destination)) != extract_root:
|
|
raise ValueError(f"Zip member escapes extraction root: {member.filename}")
|
|
|
|
os.makedirs(os.path.dirname(destination), exist_ok=True)
|
|
with archive.open(member, "r") as source, open(destination, "wb") as target:
|
|
shutil.copyfileobj(source, target)
|
|
extracted_items.append((destination, validated_name))
|
|
except Exception:
|
|
shutil.rmtree(extract_root, ignore_errors=True)
|
|
raise
|
|
|
|
if not extracted_items:
|
|
shutil.rmtree(extract_root, ignore_errors=True)
|
|
raise ValueError("No supported images found in zip archive.")
|
|
|
|
manifest = []
|
|
total = len(extracted_items)
|
|
for index, (absolute_path, relative_path) in enumerate(sorted(extracted_items, key=lambda item: item[1].lower())):
|
|
manifest.append(build_manifest_entry(extract_root, absolute_path, relative_path, index, total))
|
|
|
|
return extract_root, manifest
|
|
|
|
|
|
def parse_image_indices(image_indices, total):
|
|
if total <= 0:
|
|
raise ValueError("No images available for selection.")
|
|
|
|
raw = str(image_indices).strip()
|
|
if not raw:
|
|
raise ValueError("Image indices cannot be empty.")
|
|
|
|
selected = []
|
|
seen = set()
|
|
|
|
for part in raw.split(","):
|
|
token = part.strip()
|
|
if not token or not _INDEX_PART_PATTERN.match(token):
|
|
raise ValueError(f"Invalid image index expression: {image_indices}")
|
|
|
|
if "-" in token:
|
|
start_text, end_text = token.split("-", 1)
|
|
start = int(start_text)
|
|
end = int(end_text)
|
|
if end < start:
|
|
raise ValueError(f"Invalid range in image indices: {token}")
|
|
values = range(start, end + 1)
|
|
else:
|
|
values = [int(token)]
|
|
|
|
for value in values:
|
|
if value < 0 or value >= total:
|
|
continue
|
|
if value not in seen:
|
|
seen.add(value)
|
|
selected.append(value)
|
|
|
|
if not selected:
|
|
raise ValueError(f"Image indices produced no valid results: {image_indices}")
|
|
|
|
return selected
|
|
|
|
|
|
def manifest_to_json(manifest):
|
|
return json.dumps(manifest, ensure_ascii=False)
|
|
|
|
|
|
def manifest_from_json(manifest_json):
|
|
if isinstance(manifest_json, str):
|
|
return json.loads(manifest_json)
|
|
|
|
if isinstance(manifest_json, dict):
|
|
return [manifest_json]
|
|
|
|
if isinstance(manifest_json, tuple):
|
|
manifest_json = list(manifest_json)
|
|
|
|
if isinstance(manifest_json, list):
|
|
if not manifest_json:
|
|
return []
|
|
if all(isinstance(item, dict) for item in manifest_json):
|
|
return manifest_json
|
|
if len(manifest_json) == 1:
|
|
return manifest_from_json(manifest_json[0])
|
|
raise ValueError("Manifest list must contain either one JSON payload or manifest items.")
|
|
|
|
raise TypeError(f"Unsupported manifest payload type: {type(manifest_json).__name__}")
|
|
|
|
|
|
def build_unique_output_path(output_root, filename_prefix, manifest_item):
|
|
prefix = sanitize_relative_path(filename_prefix) if filename_prefix else "ComfyUI"
|
|
output_root = os.path.abspath(output_root)
|
|
safe_parent = manifest_item.get("safe_relative_parent", "")
|
|
safe_stem = re.sub(r"[^A-Za-z0-9._ -]", "_", manifest_item["stem"]).strip().rstrip(". ") or "image"
|
|
safe_suffix = manifest_item["suffix"] if manifest_item["suffix"] in IMAGE_EXTENSIONS else ".png"
|
|
|
|
relative_dir = prefix if not safe_parent else posixpath.join(prefix, safe_parent)
|
|
output_dir = os.path.abspath(os.path.join(output_root, relative_dir.replace("/", os.sep)))
|
|
if os.path.commonpath((output_root, output_dir)) != output_root:
|
|
raise ValueError("Saving image outside the output folder is not allowed.")
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
counter = 0
|
|
while True:
|
|
suffix = "" if counter == 0 else f"_{counter}"
|
|
filename = f"{safe_stem}{suffix}{safe_suffix}"
|
|
absolute_path = os.path.abspath(os.path.join(output_dir, filename))
|
|
if os.path.commonpath((output_root, absolute_path)) != output_root:
|
|
raise ValueError("Saving image outside the output folder is not allowed.")
|
|
try:
|
|
handle = os.open(absolute_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
|
os.close(handle)
|
|
break
|
|
except FileExistsError:
|
|
counter += 1
|
|
|
|
subfolder = relative_dir.replace("\\", "/")
|
|
return absolute_path, filename, subfolder
|
|
|
|
|
|
def split_manifest_by_indices(manifest, indices):
|
|
return [manifest[index] for index in indices]
|
|
|
|
|
|
def log_skip(message, *args):
|
|
logging.warning(message, *args)
|