feat: 🚧 add my CLIs
just convenience tools, not exposed but usable from CI etc... still need some work.
This commit is contained in:
melMass
@@ -0,0 +1,142 @@
|
||||
import glob
|
||||
from pathlib import Path
|
||||
import uuid
|
||||
import sys
|
||||
from typing import List
|
||||
|
||||
sys.path.append((Path(__file__).parent / "extern").as_posix())
|
||||
|
||||
|
||||
import argparse
|
||||
from rich_argparse import RichHelpFormatter
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
|
||||
import numpy as np
|
||||
import subprocess
|
||||
|
||||
|
||||
def write_prores_444_video(output_file, frames: List[np.ndarray], fps):
|
||||
# Convert float images to the range of 0-65535 (12-bit color depth)
|
||||
frames = [(frame * 65535).clip(0, 65535).astype(np.uint16) for frame in frames]
|
||||
|
||||
height, width, _ = frames[0].shape
|
||||
|
||||
# Prepare the FFmpeg command
|
||||
command = [
|
||||
"ffmpeg",
|
||||
"-y", # Overwrite output file if it already exists
|
||||
"-f",
|
||||
"rawvideo",
|
||||
"-vcodec",
|
||||
"rawvideo",
|
||||
"-s",
|
||||
f"{width}x{height}",
|
||||
"-pix_fmt",
|
||||
"rgb48le",
|
||||
"-r",
|
||||
str(fps),
|
||||
"-i",
|
||||
"-",
|
||||
"-c:v",
|
||||
"prores_ks",
|
||||
"-profile:v",
|
||||
"4",
|
||||
"-pix_fmt",
|
||||
"yuva444p10le",
|
||||
"-r",
|
||||
str(fps),
|
||||
"-y", # Overwrite output file if it already exists
|
||||
output_file,
|
||||
]
|
||||
|
||||
process = subprocess.Popen(command, stdin=subprocess.PIPE)
|
||||
|
||||
for frame in frames:
|
||||
process.stdin.write(frame.tobytes())
|
||||
|
||||
process.stdin.close()
|
||||
process.wait()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
default_output = f"./output_{uuid.uuid4()}.mov"
|
||||
parser = argparse.ArgumentParser(
|
||||
description="FILM frame interpolation", formatter_class=RichHelpFormatter
|
||||
)
|
||||
parser.add_argument("inputs", nargs="*", help="Input image files")
|
||||
parser.add_argument("--output", help="Output JSON file", default=default_output)
|
||||
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose mode")
|
||||
parser.add_argument(
|
||||
"--glob", help="Enable glob pattern matching", metavar="PATTERN"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--interpolate", type=int, default=4, help="Time for interpolated frames"
|
||||
)
|
||||
parser.add_argument("--fps", type=int, default=30, help="Out FPS")
|
||||
align = 64
|
||||
block_width = 2
|
||||
block_height = 2
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# - checks
|
||||
if not args.glob and not args.inputs:
|
||||
parser.error("Either --glob flag or inputs must be provided.")
|
||||
if args.glob:
|
||||
glob_pattern = args.glob
|
||||
try:
|
||||
pattern_path = str(Path(glob_pattern).expanduser().resolve())
|
||||
|
||||
if not any(glob.glob(pattern_path)):
|
||||
raise ValueError(f"No files found for glob pattern: {glob_pattern}")
|
||||
except Exception as e:
|
||||
console = Console()
|
||||
console.print(
|
||||
f"[bold red]Error: Invalid glob pattern '{glob_pattern}': {e}[/bold red]"
|
||||
)
|
||||
|
||||
exit(1)
|
||||
else:
|
||||
glob_pattern = None
|
||||
|
||||
input_files: List[Path] = []
|
||||
|
||||
if glob_pattern:
|
||||
input_files = [
|
||||
Path(p)
|
||||
for p in list(glob.glob(str(Path(glob_pattern).expanduser().resolve())))
|
||||
]
|
||||
else:
|
||||
input_files = [Path(p) for p in args.inputs]
|
||||
|
||||
console = Console()
|
||||
console.print("Input Files:", style="bold", end=" ")
|
||||
console.print(f"{len(input_files):03d} files", style="cyan")
|
||||
# for input_file in args.inputs:
|
||||
# console.print(f"- {input_file}", style="cyan")
|
||||
console.print("\nOutput File:", style="bold", end=" ")
|
||||
console.print(f"{Path(args.output).resolve().absolute()}", style="cyan")
|
||||
|
||||
with Progress(console=console, auto_refresh=True) as progress:
|
||||
from frame_interpolation.eval import util
|
||||
from frame_interpolation.eval import util, interpolator
|
||||
|
||||
# files = Path(pth).rglob("*.png")
|
||||
|
||||
model = interpolator.Interpolator(
|
||||
"G:/MODELS/FILM/pretrained_models/film_net/Style", None
|
||||
) # [2,2]
|
||||
|
||||
task = progress.add_task("[cyan]Interpolating frames...", total=1)
|
||||
|
||||
frames = list(
|
||||
util.interpolate_recursively_from_files(
|
||||
[x.as_posix() for x in input_files], args.interpolate, model
|
||||
)
|
||||
)
|
||||
|
||||
# mediapy.write_video(args.output, frames, fps=args.fps)
|
||||
write_prores_444_video(args.output, frames, fps=args.fps)
|
||||
progress.update(task, advance=1)
|
||||
progress.refresh()
|
||||
@@ -0,0 +1,112 @@
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngImageFile, PngInfo
|
||||
import json
|
||||
from pprint import pprint
|
||||
import argparse
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
from rich_argparse import RichHelpFormatter
|
||||
|
||||
|
||||
def parse_a111(params, verbose=False):
|
||||
# params = [p.split(": ") for p in params.split("\n")]
|
||||
params = params.split("\n")
|
||||
|
||||
prompt = params[0].strip()
|
||||
neg = params[1].split(":")[1].strip()
|
||||
|
||||
settings = {}
|
||||
try:
|
||||
settings = {
|
||||
s.split(":")[0].strip(): s.split(":")[1].strip()
|
||||
for s in params[2].split(",")
|
||||
}
|
||||
|
||||
except IndexError:
|
||||
settings = {"raw": params[2].strip()}
|
||||
|
||||
if verbose:
|
||||
print(f"PROMPT: {prompt}")
|
||||
print(f"NEG: {neg}")
|
||||
print("SETTINGS:")
|
||||
pprint(settings, indent=4)
|
||||
|
||||
return {"prompt": prompt, "negative": neg, "settings": settings}
|
||||
|
||||
|
||||
import glob
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Crude metadata extractor from A111 pngs",
|
||||
formatter_class=RichHelpFormatter
|
||||
)
|
||||
parser.add_argument("inputs", nargs="*", help="Input image files")
|
||||
parser.add_argument("--output", help="Output JSON file")
|
||||
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose mode")
|
||||
parser.add_argument(
|
||||
"--glob", help="Enable glob pattern matching", metavar="PATTERN"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# - checks
|
||||
if not args.glob and not args.inputs:
|
||||
parser.error("Either --glob flag or inputs must be provided.")
|
||||
if args.glob:
|
||||
glob_pattern = args.glob
|
||||
try:
|
||||
pattern_path = str(Path(glob_pattern).expanduser().resolve())
|
||||
|
||||
if not any(glob.glob(pattern_path)):
|
||||
raise ValueError(f"No files found for glob pattern: {glob_pattern}")
|
||||
except Exception as e:
|
||||
console = Console()
|
||||
console.print(
|
||||
f"[bold red]Error: Invalid glob pattern '{glob_pattern}': {e}[/bold red]"
|
||||
)
|
||||
|
||||
exit(1)
|
||||
else:
|
||||
glob_pattern = None
|
||||
|
||||
input_files = []
|
||||
|
||||
if glob_pattern:
|
||||
input_files = list(glob.glob(str(Path(glob_pattern).expanduser().resolve())))
|
||||
else:
|
||||
input_files = [Path(p) for p in args.inputs]
|
||||
|
||||
console = Console()
|
||||
console.print("Input Files:", style="bold", end=" ")
|
||||
console.print(f"{len(input_files):03d} files", style="cyan")
|
||||
# for input_file in args.inputs:
|
||||
# console.print(f"- {input_file}", style="cyan")
|
||||
console.print("\nOutput File:", style="bold", end=" ")
|
||||
console.print(f"{Path(args.output).resolve().absolute()}", style="cyan")
|
||||
|
||||
with Progress(console=console, auto_refresh=True) as progress:
|
||||
# files = Path(pth).rglob("*.png")
|
||||
unique_info = {}
|
||||
last = None
|
||||
|
||||
task = progress.add_task("[cyan]Extracting meta...", total=len(input_files) + 1)
|
||||
for p in input_files:
|
||||
im = Image.open(p)
|
||||
parsed = parse_a111(im.info["parameters"], args.verbose)
|
||||
|
||||
if parsed != last:
|
||||
unique_info[Path(p).stem] = parsed
|
||||
|
||||
last = parsed
|
||||
progress.update(task, advance=1)
|
||||
progress.refresh()
|
||||
|
||||
unique_info = json.dumps(unique_info, indent=4)
|
||||
with open(args.output, "w") as f:
|
||||
f.write(unique_info)
|
||||
progress.update(task, advance=1)
|
||||
progress.refresh()
|
||||
|
||||
console.print("\nProcessing completed!", style="bold green")
|
||||
@@ -0,0 +1,213 @@
|
||||
import argparse
|
||||
import json
|
||||
from PIL import Image, PngImagePlugin
|
||||
from rich.console import Console
|
||||
from rich import print
|
||||
from rich_argparse import RichHelpFormatter
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
console = Console()
|
||||
|
||||
# BNK_CutoffSetRegions
|
||||
# BNK_CutoffRegionsToConditioning
|
||||
# BNK_CutoffBasePrompt
|
||||
|
||||
|
||||
# Extracts metadata from a PNG image and returns it as a dictionary
|
||||
def extract_metadata(image_path):
|
||||
image = Image.open(image_path)
|
||||
prompt = image.info.get("prompt", "")
|
||||
workflow = image.info.get("workflow", "")
|
||||
|
||||
if workflow:
|
||||
workflow = json.loads(workflow)
|
||||
|
||||
if prompt:
|
||||
prompt = json.loads(prompt)
|
||||
|
||||
console.print(f"Metadata extracted from [cyan]{image_path}[/cyan].")
|
||||
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"workflow": workflow,
|
||||
}
|
||||
|
||||
|
||||
# Embeds metadata into a PNG image
|
||||
def embed_metadata(image_path, metadata):
|
||||
image = Image.open(image_path)
|
||||
o_metadata = image.info
|
||||
|
||||
pnginfo = PngImagePlugin.PngInfo()
|
||||
if prompt := metadata.get("prompt"):
|
||||
pnginfo.add_text("prompt", json.dumps(prompt))
|
||||
elif "prompt" in o_metadata:
|
||||
pnginfo.add_text("prompt", o_metadata["prompt"])
|
||||
|
||||
if workflow := metadata.get("workflow"):
|
||||
pnginfo.add_text("workflow", json.dumps(workflow))
|
||||
elif "workflow" in o_metadata:
|
||||
pnginfo.add_text("workflow", o_metadata["workflow"])
|
||||
|
||||
imgp = Path(image_path)
|
||||
output = imgp.with_stem(f"{imgp.stem}_comfy_embed")
|
||||
index = 1
|
||||
while output.exists():
|
||||
output = imgp.with_stem(f"{imgp.stem}_{index}_comfy_embed").with_suffix(".png")
|
||||
index += 1
|
||||
|
||||
image.save(output, pnginfo=pnginfo)
|
||||
console.print(f"Metadata embedded into [cyan]{output}[/cyan].")
|
||||
|
||||
|
||||
# CLI subcommand: extract
|
||||
def extract(args):
|
||||
input_files = []
|
||||
for input_path in args.input:
|
||||
if os.path.isdir(input_path):
|
||||
folder_path = input_path
|
||||
input_files.extend(
|
||||
[
|
||||
os.path.join(folder_path, file_name)
|
||||
for file_name in os.listdir(folder_path)
|
||||
if file_name.lower().endswith((".png", ".jpg", ".jpeg"))
|
||||
]
|
||||
)
|
||||
else:
|
||||
input_files.append(input_path)
|
||||
|
||||
if len(input_files) == 1:
|
||||
metadata = extract_metadata(input_files[0])
|
||||
if args.print_output:
|
||||
print(json.dumps(metadata, indent=4))
|
||||
else:
|
||||
if not args.output:
|
||||
output = Path(input_files[0]).with_suffix(".json")
|
||||
index = 1
|
||||
while output.exists():
|
||||
output = (
|
||||
Path(input_files[0])
|
||||
.with_stem(f"{Path(input_files[0]).stem}_{index}")
|
||||
.with_suffix(".json")
|
||||
)
|
||||
index += 1
|
||||
else:
|
||||
output = args.output
|
||||
with open(output, "w") as file:
|
||||
json.dump(metadata, file, indent=4)
|
||||
console.print(f"Metadata extracted and saved to [cyan]{output}[/cyan].")
|
||||
else:
|
||||
metadata_dict = {}
|
||||
for input_file in input_files:
|
||||
metadata = extract_metadata(input_file)
|
||||
filename = os.path.basename(input_file)
|
||||
output = (
|
||||
Path(args.output) / f"{filename}.json"
|
||||
if args.output
|
||||
else Path(input_file).with_suffix(".json")
|
||||
)
|
||||
index = 1
|
||||
while output.exists():
|
||||
output = Path(args.output).parent / f"{filename}_{index}.json"
|
||||
index += 1
|
||||
with open(output, "w") as file:
|
||||
json.dump(metadata, file, indent=4)
|
||||
metadata_dict[filename] = metadata
|
||||
if args.output:
|
||||
with open(args.output, "w") as file:
|
||||
json.dump(metadata_dict, file, indent=4)
|
||||
console.print(
|
||||
f"Metadata extracted and saved to [cyan]{args.output}[/cyan]."
|
||||
)
|
||||
else:
|
||||
console.print("Multiple metadata files created.")
|
||||
|
||||
|
||||
# CLI subcommand: embed
|
||||
def embed(args):
|
||||
input_files = []
|
||||
for input_path in args.input:
|
||||
if os.path.isdir(input_path):
|
||||
folder_path = input_path
|
||||
input_files.extend(
|
||||
[
|
||||
os.path.join(folder_path, file_name)
|
||||
for file_name in os.listdir(folder_path)
|
||||
if file_name.lower().endswith(".json")
|
||||
]
|
||||
)
|
||||
else:
|
||||
input_files.append(input_path)
|
||||
|
||||
for input_file in input_files:
|
||||
with open(input_file) as file:
|
||||
metadata = json.load(file)
|
||||
image_path = input_file.replace(".json", ".png")
|
||||
if args.output:
|
||||
output_dir = args.output
|
||||
if os.path.isdir(output_dir):
|
||||
output_path = os.path.join(output_dir, os.path.basename(image_path))
|
||||
index = 1
|
||||
while os.path.exists(output_path):
|
||||
output_path = os.path.join(
|
||||
output_dir,
|
||||
f"{os.path.basename(image_path)}_{index}.png",
|
||||
)
|
||||
index += 1
|
||||
else:
|
||||
output_path = output_dir
|
||||
else:
|
||||
output_path = image_path.replace(".png", "_comfy_embed.png")
|
||||
|
||||
embed_metadata(image_path, metadata)
|
||||
# os.rename(image_path, output_path)
|
||||
console.print(f"Metadata embedded into [cyan]{output_path}[/cyan].")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Create the main CLI parser
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="image-metadata-cli", formatter_class=RichHelpFormatter
|
||||
)
|
||||
subparsers = parser.add_subparsers(title="subcommands")
|
||||
|
||||
# Parser for the "extract" subcommand
|
||||
extract_parser = subparsers.add_parser(
|
||||
"extract",
|
||||
help="Extract metadata from PNG image(s) or folder",
|
||||
formatter_class=RichHelpFormatter,
|
||||
)
|
||||
extract_parser.add_argument(
|
||||
"input", nargs="+", help="Input PNG image file(s) or folder path"
|
||||
)
|
||||
extract_parser.add_argument(
|
||||
"--print",
|
||||
dest="print_output",
|
||||
action="store_true",
|
||||
help="Print the output to stdout",
|
||||
)
|
||||
extract_parser.add_argument("--output", help="Output JSON file(s) or directory")
|
||||
extract_parser.set_defaults(func=extract)
|
||||
|
||||
# Parser for the "embed" subcommand
|
||||
embed_parser = subparsers.add_parser(
|
||||
"embed",
|
||||
help="Embed metadata into PNG image(s) or folder",
|
||||
formatter_class=RichHelpFormatter,
|
||||
)
|
||||
embed_parser.add_argument(
|
||||
"input", nargs="+", help="Input JSON file(s) or folder path"
|
||||
)
|
||||
embed_parser.add_argument("--output", help="Output PNG image file(s) or directory")
|
||||
embed_parser.set_defaults(func=embed)
|
||||
|
||||
# Parse the command-line arguments and execute the appropriate subcommand
|
||||
args = parser.parse_args()
|
||||
if hasattr(args, "func"):
|
||||
try:
|
||||
args.func(args)
|
||||
except ValueError as e:
|
||||
console.print(f"[bold red]Error:[/bold red] {str(e)}")
|
||||
else:
|
||||
parser.print_help()
|
||||
Reference in New Issue
Block a user