diff --git a/interpolate_frames.py b/interpolate_frames.py new file mode 100644 index 0000000..6c235f0 --- /dev/null +++ b/interpolate_frames.py @@ -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() diff --git a/scripts/a111_extract.py b/scripts/a111_extract.py new file mode 100644 index 0000000..7e38b0b --- /dev/null +++ b/scripts/a111_extract.py @@ -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") diff --git a/scripts/comfy_meta.py b/scripts/comfy_meta.py new file mode 100644 index 0000000..cb9cb1e --- /dev/null +++ b/scripts/comfy_meta.py @@ -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()