diff --git a/.github/instructions/python script.instructions.md b/.github/instructions/python script.instructions.md new file mode 100644 index 0000000..902d591 --- /dev/null +++ b/.github/instructions/python script.instructions.md @@ -0,0 +1,32 @@ +--- +applyTo: "*.py" +--- + +- Do not be chatty, do not summarize the code afterwards. +- Use `#!/usr/bin/env -S uv run --script` as the shebang line. +- Add dependencies under the shebang line as + +``` +# /// script +# dependencies = [list of external dependencies as json array] +# /// +``` + +and keep them updated as the code changes. + +- Keep comments to minimum +- Use `uv` for package management. +- Use logging instead of print statements. Use `logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")` for configuration. +- Use argparse for command-line arguments, keep argument parsing in a func called `parse_args() -> argparse.Namespace`, which gets called from `main`. Use `argparse.ArgumentDefaultsHelpFormatter` for the formatter class. +- Use `pathlib.Path` for file operations. When the CLI args include a path, set its `type=Path` +- Place imports at the top of the file. +- Use httpx for HTTP requests. Do not use async. +- Always add type annotations to function signatures, prefer built-in types, e.g. `list[str]` over `typing` module types when possible. +- Prefer kwargs calling style for function calls, e.g. `func(arg1=value1, arg2=value2)` instead of `func(value1, value2)`. +- Use `Path.cwd()` for the current working directory. +- Prefer `os.getenv` over `os.environ.get` for environment variables. +- Use `subprocess.run` for running shell commands, prefer `check=True` to raise an error on failure. +- Use `tempfile.TemporaryDirectory()` for temporary directories. +- Do not shorten parameter names unnecessarily, unless it's an idiomatic abbreviation. +- Use f-strings for string formatting. +- When calling an external command, use the long form of the arguments, e.g. `--output` instead of `-o`. diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..3a4f41e --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.13 \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index 321f0d3..1bafd7b 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -7,9 +7,10 @@ { "name": "Python: Current File", "type": "python", - "args": ["2022-12-06"], + "env": {"MISTRAL_API_KEY":"8J5hUwxVMCiQjhXxElefDqU7XJYAvVXn", "OPENROUTER_API_KEY": "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"}, "request": "launch", "program": "${file}", + "cwd": "${workspaceFolder}", "console": "integratedTerminal", "justMyCode": true } diff --git a/ai_imagez.py b/ai_imagez.py new file mode 100755 index 0000000..7159a18 --- /dev/null +++ b/ai_imagez.py @@ -0,0 +1,363 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["click"] +# /// + +import logging +import random +import subprocess +from pathlib import Path + +import click + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s: %(message)s", +) +logger = logging.getLogger(name=__name__) + +VALID_EXTENSIONS = {".jpeg", ".jpg", ".png", ".webm", ".webp"} + + +def is_valid_image(file_path: Path) -> bool: + """Check if file is a valid image type.""" + return file_path.suffix.lower() in VALID_EXTENSIONS + + +def move_to_trash(file_paths: list[Path]) -> None: + """Move files to Trash using trash CLI (single batched call).""" + resolved: list[str] = [] + for file_path in file_paths: + if not file_path.is_file(): + raise FileNotFoundError(f"File not found: {file_path}") + resolved.append(str(file_path.resolve())) + + if resolved: + subprocess.run(["trash", *resolved], check=True) + + +def build_target_file_index(target_dirs: list[Path]) -> dict[str, Path]: + """ + Build a mapping of file stems to paths from target directories. + + Scans target directories once and returns a dict for efficient lookups. + """ + index: dict[str, Path] = {} + + for target_dir in target_dirs: + if not target_dir.is_dir(): + continue + + for candidate in target_dir.iterdir(): + if not candidate.is_file(): + continue + + if not is_valid_image(file_path=candidate): + continue + + index[candidate.stem] = candidate + + return index + + +def find_edited_version( + image_path: Path, + target_index: dict[str, Path], +) -> Path | None: + """ + Search target file index for edited version of image. + + An image is considered edited if its stem starts with the original stem + followed by '-edit'. Matches by stem only, ignoring extension. + + Returns the path to the edited version or None if not found. + """ + original_stem = image_path.stem + search_pattern = f"{original_stem}-edit" + + for indexed_stem, indexed_path in target_index.items(): + if indexed_stem.startswith(search_pattern): + return indexed_path + + return None + + +def build_directory_index(directory: Path) -> set[str]: + """ + Build a set of all valid image stems in a directory. + + Args: + directory: Path to the directory to scan + + Returns: + Set of file stems that are valid images + """ + stems: set[str] = set() + + if not directory.is_dir(): + return stems + + for candidate in directory.iterdir(): + if not candidate.is_file(): + continue + + if not is_valid_image(file_path=candidate): + continue + + stems.add(candidate.stem) + + return stems + + +def collect_images(paths: list[Path]) -> list[Path]: + images: list[Path] = [] + for path in paths: + if path.is_file(): + if is_valid_image(path): + images.append(path) + elif path.is_dir(): + for candidate in sorted(path.iterdir()): + if candidate.is_file() and is_valid_image(candidate): + images.append(candidate) + return images + + +@click.group() +def cli() -> None: + """Manage and process image files.""" + pass + + +@cli.command() +@click.argument( + "image_paths", + nargs=-1, + required=True, + type=click.Path(exists=True, path_type=Path), +) +@click.option( + "--target-dir", + multiple=True, + required=True, + type=click.Path(exists=True, path_type=Path), + help="Directory to search for edited versions", +) +@click.option( + "--dry-run", + is_flag=True, + default=False, + help="Preview matches without moving to trash", +) +def dedupe( + image_paths: tuple[Path, ...], + target_dir: tuple[Path, ...], + dry_run: bool, +) -> None: + """ + Remove original images when edited versions exist in target directories. + + IMAGE_PATHS: One or more image files to check (jpeg, jpg, png, webm, webp) + """ + target_dirs = list(target_dir) + valid_images = [path for path in image_paths if is_valid_image(file_path=path)] + + if not valid_images: + logger.warning("No valid image files found") + return + + target_index = build_target_file_index(target_dirs=target_dirs) + matches: list[tuple[Path, Path]] = [] + + for image_path in valid_images: + edited_version = find_edited_version( + image_path=image_path, + target_index=target_index, + ) + + if edited_version: + matches.append((image_path, edited_version)) + logger.info(f"Found edited version: {image_path} -> {edited_version}") + + if not matches: + logger.info("No edited versions found") + return + + if dry_run: + click.echo( + click.style( + text="DRY RUN MODE - No files will be moved", + fg="yellow", + ) + ) + click.echo(f"Would move {len(matches)} image(s) to trash:") + for original, edited in matches: + click.echo(f" {original}") + return + + originals = [original for original, _ in matches] + try: + move_to_trash(file_paths=originals) + for original in originals: + logger.info(f"Moved to trash: {original}") + except (FileNotFoundError, subprocess.CalledProcessError) as e: + logger.error(f"Failed to move files to trash: {e}") + + click.echo( + click.style( + text=f"Moved {len(matches)} image(s) to trash", + fg="green", + ) + ) + + +@cli.command() +@click.argument( + "paths", + nargs=-1, + required=True, + type=click.Path(exists=True, path_type=Path), +) +@click.option( + "--dry-run", + is_flag=True, + default=False, + help="Preview matches without moving to trash", +) +def clean(paths: tuple[Path, ...], dry_run: bool) -> None: + """ + Remove originals when upscaled versions exist in the same basket. + + PATHS: Image files and/or directories to scan (jpeg, jpg, png, webm, webp) + + Scans all provided paths into a collection. If a file and its -upscaled + counterpart (stem starts with '{original}-upscaled', any extension) + are both in the collection, the original is moved to trash. + """ + all_images = collect_images(list(paths)) + + if not all_images: + logger.warning("No valid image files found") + return + + image_index: dict[str, Path] = {img.stem: img for img in all_images} + trashables: set[Path] = set() + + markers = ["-upscaled", "-edit"] + for image in all_images: + if any(marker in image.stem for marker in markers): + continue + + upscaleds = [ + p + for s, p in image_index.items() + if s.startswith(image.stem) and "-upscaled" in s + ] + editeds = [ + p + for s, p in image_index.items() + if s.startswith(image.stem) and "-edit" in s + ] + if editeds or upscaleds: + logger.info(f"Will delete {image}") + trashables.add(image) + for upscaled in upscaleds: + low_res_edit = upscaled.with_stem( + upscaled.stem.removesuffix("-upscaled") + ) + if low_res_edit in editeds: + logger.info(f"Will also delete low-res edit {low_res_edit}") + trashables.add(low_res_edit) + + if not trashables: + logger.info("No upscaled or edited versions found") + return + + if dry_run: + click.echo(click.style("DRY RUN MODE - No files will be moved", fg="yellow")) + click.echo(f"Would move {len(trashables)} image(s) to trash:") + for original in trashables: + click.echo(f" {original}") + return + + try: + move_to_trash(file_paths=list(trashables)) + for original in trashables: + logger.info(f"Moved to trash: {original}") + except (FileNotFoundError, subprocess.CalledProcessError) as e: + logger.error(f"Failed to move files to trash: {e}") + + click.echo(click.style(f"Moved {len(trashables)} image(s) to trash", fg="green")) + + +@cli.command(name="list-editable") +@click.argument( + "image_paths", + nargs=-1, + required=True, + type=click.Path(exists=True, path_type=Path), +) +@click.option( + "--shuffle", + is_flag=True, + default=False, + help="Shuffle the output filenames", +) +@click.option( + "--null", + is_flag=True, + default=False, + help="Output null-terminated filenames for xargs -0", +) +def list_editable(image_paths: tuple[Path, ...], null: bool, shuffle: bool) -> None: + """ + List image files that don't have edited versions. + + IMAGE_PATHS: One or more image files to check (jpeg, jpg, png, webm, webp) + + Skips files with '-edit' in their stem and lists those without + matching edited versions (stem + '-edit' + optional suffix) in + the same directory. + """ + valid_images = [path for path in image_paths if is_valid_image(file_path=path)] + + if not valid_images: + logger.warning("No valid image files found") + return + + dir_indices: dict[Path, set[str]] = {} + editable_files: list[Path] = [] + + for image_path in valid_images: + if "-edit" in image_path.stem: + continue + + directory = image_path.parent + if directory not in dir_indices: + dir_indices[directory] = build_directory_index(directory=directory) + + dir_stems = dir_indices[directory] + original_stem = image_path.stem + search_pattern = f"{original_stem}-edit" + + has_edited = any(stem.startswith(search_pattern) for stem in dir_stems) + + if not has_edited: + editable_files.append(image_path) + + if not editable_files: + logger.info("No editable files found") + return + + if shuffle: + random.shuffle(editable_files) + + for file_path in editable_files: + if null: + click.echo(str(file_path), nl=False) + click.echo("\0", nl=False) + else: + click.echo(str(file_path)) + + +if __name__ == "__main__": + cli() diff --git a/ai_photorealistic.py b/ai_photorealistic.py new file mode 100755 index 0000000..e864b48 --- /dev/null +++ b/ai_photorealistic.py @@ -0,0 +1,418 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["httpx", "pillow"] +# /// +import argparse +import base64 +import json +import logging +import os +import subprocess +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from io import BytesIO +from pathlib import Path +from typing import Any + +import httpx +from PIL import Image, ImageOps + +from sort_images import read_dims + +MODEL_ID = "flux-2-klein-9b" +# LOCAL_MODEL_ID = "black-forest-labs/FLUX.2-klein-9B" +# MODEL_ID = "z-image-turbo" +# MODEL_ID = "qwen-image" + + +SYSTEM_PROMPT = """Make this incredibly photorealistic. +Highly stylized, striking, highly detailed photo of a young, lithe woman. +Age her by 5 years and make her a 20 year old, taller, slimmer, more slender, fitter version of herself. +Make her 6ft tall. + +Give her realistic spotless skin texture. +Her face skin must also be realistic, without any makeup. +Make the background realistic and consistent with the lighting on her. + +Give her an angular, not round, narrow, elongated neck and long face shape defined by a structured, clean jawline that tapers toward a firm, slightly rounded chin. +Do not make the head larger than the body, keep them in perfect proportion. +Give her high, prominent cheekbones giving her sculpted, chieseled face that taper down to a more narrow, delicate chin. +Make her really pretty and seductive. +Preserve unnatural lip colors and wetness. + +Do not change the facial expression, emotion, hand and body pose, mouth, tongue and eyelid position. +Lower her eyes slightly to give her an alluring look. +Mouth slightly agape giving her a seductive look. +Give her a sharp, cunning gaze, a very slight pleasant smirk without changing her head pose. + +Remove the drool. +Keep the pose and don't turn people around.""" + +SYSTEM_PROMPT = """ +Turn this into a heavily stylized real life photo of a woman. +Age her by 5 years and make her a 20 year old, taller, slimmer, more slender, fitter version of herself. + +Apply realistic textures without grain or noise. +Skin should be smooth, flawless, without pores, clean and without blemishes. + +Keep the shine. +Make her taller, slimmer, slender. Don't make the head shorter. + +Do not change the pose, gaze, emotion. + +Her lips are parted, creating a soft, seductive expression. +""" + +SYSTEM_PROMPT = """ +restyle this as an edited, finished photo taken by sony a7 iv. +Photorealistic textures. +same color grading, same colors and lighting. +Realistic, flawless skin. +White people should have a pale skin. + +keep the same ethnicity and facial features. + +Make the children taller and more slender with slightly elongated feminine neck and higher cheekbones, give them smaller, slimmer more feminine head and tapered face; round, smooth chin. +age them by 5 years and make them look like 20 year old adults. + +keep everything else about her the same. +same facial + eye expression (keep them as open as in the original), same pose and emotion, same tongue position. +same head and eye angle. + +do not turn her face around. +same hair style, color, length. +same composition and crop. +keep the same makeup, do not change the colors of the lips, keep them glossy and wet. +""" +SYSTEM_PROMPT = """ +Artistic super photorealistic conversion. +extend the background horizontally but keep the people. +# Subtle chiaroscuro lighting, not too dark. +85mm telephoto lens, f/2.8, shallow depth of field. +Soft-focus highlights, atmospheric bloom. +Kodak Portra 400 cool color palette. +#Preserve the original lighting and slight moody lighting. Not underexposed, not overexposed. +High dynamic range with a focus on rich textures. +balanced exposure. +same color grading and LUT. + +remove compression artifacts. +she has spotless, supple, shiny skin without splotches. +#pale skin, goth make-up. + +#reduce musculature. +#keep the original facial expression, emotion, hand and body pose. + +make her prettier, give her hourglass figure, a satisfied look, voluminus hair. + +parted lips, seductive gaze, subtly lowered eyelids. +make her face look like a 25 year old, raised cheekbones, tapered face shape, pointy chin. +#avoid making her look like a child and maintain adult proportions. +#give her a subtle mischievous, confident smile, avoid neutral face. +#make her face slimmer. +preserve the original skin color. + +do not change the colors of the lips, keep them wet. +remove all logos and watermarks. +#reflective outfit, +#glossy, reflective outfit, +#clingy, skin-tight clothes. +""" + + +def without_comments(s: str) -> str: + return "\n".join( + line for line in s.splitlines() if not line.strip().startswith("#") + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Send one or more images to NanoGPT for photorealistic processing." + ) + parser.add_argument( + "image_paths", + nargs="+", + help="One or more local image paths to send", + ) + parser.add_argument( + "--extra-prompt", + default="", + help="Extra prompt text appended to the system prompt", + ) + parser.add_argument( + "--prompt", + default="", + help="Replace the system prompt with this custom prompt (instead of appending to it)", + ) + parser.add_argument( + "--provider", + choices=("nanogpt", "local"), + default=os.getenv("IMAGE_PROVIDER", "nanogpt"), + help="Image provider backend (default from IMAGE_PROVIDER or nanogpt)", + ) + parser.add_argument( + "--base-url", + default=os.getenv("LOCAL_IMAGE_API_BASE", "http://127.0.0.1:6006"), + help="Local provider base URL (default from LOCAL_IMAGE_API_BASE)", + ) + return parser.parse_args() + + +@dataclass(frozen=True) +class TransformResult: + image_path: Path + output_path: Path | None = None + remaining: float | None = None + + +class NanoGPT: + def __init__(self, api_key: str) -> None: + self.api_key = api_key + self._client = httpx.Client( + timeout=180.0, + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + ) + + def transform( + self, image_path: Path, output_path: Path, prompt: str + ) -> TransformResult: + if not image_path.exists(): + raise FileNotFoundError(f"Image not found: {image_path}") + + image_w, image_h = read_dims(image_path) + aspect_ratio = image_w / image_h + aspect_ratio = max(aspect_ratio, 2 / 3) + target_w, target_h = 1024, int(1024 / aspect_ratio) + size = f"{target_w}x{target_h}" + + image_data_url = _build_data_url(image_path) + payload = { + "model": MODEL_ID, + "prompt": prompt, + "imageDataUrl": image_data_url, + "response_format": "b64_json", + "n": 1, + "seed": int(time.time()), + "size": size, + } + response = self._client.post( + "https://nano-gpt.com/v1/images/generations", json=payload + ) + if not response.is_success: + raise RuntimeError( + f"API request failed with status {response.status_code}: {response.text}" + ) + result = response.json() + + data = result.get("data", []) + if not data: + raise RuntimeError(f"No data in response for {image_path}") + + item = data[0] + + if "b64_json" in item: + _save_b64_image(output_path, item["b64_json"]) + elif "url" in item: + _download_image(self._client, output_path, item["url"]) + else: + raise RuntimeError(f"Unknown response format for {image_path}") + + return TransformResult( + image_path=image_path, + output_path=output_path, + remaining=result.get("remainingBalance"), + ) + + +class LocalFlux: + def __init__(self, base_url: str) -> None: + self.base_url = base_url.rstrip("/") + self._client = httpx.Client(timeout=240.0) + + def transform( + self, image_path: Path, output_path: Path, prompt: str + ) -> TransformResult: + if not image_path.exists(): + raise FileNotFoundError(f"Image not found: {image_path}") + + image_w, image_h = read_dims(image_path) + aspect_ratio = image_w / image_h + target_w, target_h = 1024, int(1024 / aspect_ratio) + size = f"{target_w}x{target_h}" + + image_data_url = _build_data_url(image_path) + payload = { + "model": LOCAL_MODEL_ID, + "prompt": prompt, + "imageDataUrl": image_data_url, + "response_format": "b64_json", + "guidance_scale": 1.0, + "n": 1, + "seed": int(time.time()), + "size": size, + } + response = self._client.post( + f"{self.base_url}/v1/images/generations", json=payload + ) + if not response.is_success: + raise RuntimeError( + f"Local API request failed with status {response.status_code}: {response.text}" + ) + result = response.json() + + data = result.get("data", []) + if not data: + raise RuntimeError(f"No data in local response for {image_path}") + + item = data[0] + if "b64_json" in item: + _save_b64_image(output_path, item["b64_json"]) + elif "url" in item: + _download_image(self._client, output_path, item["url"]) + else: + raise RuntimeError(f"Unknown local response format for {image_path}") + + return TransformResult( + image_path=image_path, + output_path=output_path, + remaining=result.get("remainingBalance"), + ) + + +def _build_data_url(image_path: Path) -> str: + with Image.open(image_path) as image: + image = ImageOps.exif_transpose(image) + image.thumbnail((3000, 3000), Image.Resampling.LANCZOS) + if image.mode not in {"RGB", "L"}: + image = image.convert("RGB") + buffer = BytesIO() + image.save(buffer, format="JPEG", quality=95, optimize=True) + encoded = base64.b64encode(buffer.getvalue()).decode("utf-8") + return f"data:image/jpeg;base64,{encoded}" + + +def _save_b64_image(out_path: Path, b64_data: str) -> None: + out_path.write_bytes(base64.b64decode(b64_data)) + + +def _download_image(client: httpx.Client, out_path: Path, url: str) -> None: + response = client.get(url, follow_redirects=True) + response.raise_for_status() + out_path.write_bytes(response.content) + + +def _next_output_path(image_path: Path) -> Path: + base = image_path.with_name(f"{image_path.stem}-edit.jpg") + if not base.exists(): + return base + counter = 2 + while True: + candidate = image_path.with_name(f"{image_path.stem}-edit{counter}.jpg") + if not candidate.exists(): + return candidate + counter += 1 + + +def find_custom_prompts(file_paths: list[Path]) -> dict[Path, str]: + custom_prompt_path = Path(r"~/Downloads/prompts.jsonl").expanduser() + if not custom_prompt_path.exists(): + return {} + + stem_to_path: dict[str, Path] = {p.stem: p for p in file_paths} + prompts: dict[Path, str] = {} + with custom_prompt_path.open() as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + filename = entry.get("filename") + prompt = entry.get("prompt") + if filename and prompt and filename in stem_to_path: + prompts[stem_to_path[filename]] = prompt + except json.JSONDecodeError: + pass + + return prompts + + +def main() -> int: + args = parse_args() + provider = args.provider + + extra_prompt = args.extra_prompt.strip() + if extra_prompt: + prompt = f"{SYSTEM_PROMPT}\n\n{extra_prompt}" + elif args.prompt.strip(): + prompt = args.prompt.strip() + else: + prompt = SYSTEM_PROMPT + + if provider == "nanogpt": + api_key = os.getenv("NANOGPT_API_KEY") + if not api_key: + raise SystemExit("NANOGPT_API_KEY is not set") + client = NanoGPT(api_key) + else: + client = LocalFlux(args.base_url) + image_paths = [Path(p) for p in args.image_paths] + custom_prompts = find_custom_prompts(image_paths) + got_error = False + with ThreadPoolExecutor() as executor: + + def _submit(image_path: Path) -> TransformResult: + effective_prompt = without_comments(prompt) + + custom_prompt = custom_prompts.get(image_path) + if custom_prompt: + print(f"Using custom prompt for {image_path}") + effective_prompt = custom_prompt + + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) / image_path.name + result = client.transform( + image_path, + temp_path, + prompt=effective_prompt, + ) + + output_path = _next_output_path(image_path) + os.replace(temp_path, output_path) + return TransformResult( + image_path=result.image_path, + output_path=output_path, + remaining=result.remaining, + ) + + futures = { + executor.submit(_submit, image_path): image_path + for image_path in image_paths + } + for future in as_completed(futures): + try: + outcome = future.result() + except Exception as exc: + image_path = futures[future] + print(f"[error] {image_path} ({exc})") + got_error = True + continue + + image_path = outcome.image_path + output_path = outcome.output_path + print(f"[ok] {image_path} -> {output_path}") + remaining = outcome.remaining + if remaining is not None: + print(f"[balance] {remaining}") + + return 1 if got_error else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ai_runpod.py b/ai_runpod.py new file mode 100755 index 0000000..6baa556 --- /dev/null +++ b/ai_runpod.py @@ -0,0 +1,324 @@ +#!/usr/bin/env uv run +# /// script +# dependencies = ["httpx", "pillow"] +# /// + +import argparse +import io +import json +import logging +import math +import random +import subprocess +import time +import uuid +from pathlib import Path + +import httpx +from PIL import Image + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logging.getLogger("httpx").setLevel(logging.WARNING) # Suppress httpx debug logs +logger = logging.getLogger(__name__) + +COMFY_URL = "https://randy-agents-mediterranean-herself.trycloudflare.com/" +COMFY_COOKIES = { + "C.35620185_auth_token": "027e0e069f66d46e6d84802b7d49076e6ebab8621491d8694992f16727f99b79", +} +WORKFLOW_FILE = Path("/Users/abdus/Downloads/klein_fixed.json") +IMAGE_OUTPUT_MEGAPIXELS = 2 + + +SYSTEM_PROMPT = """ +restyle this as an edited, finished photo taken by sony a7 iv. +Photorealistic textures. +same color grading, same colors and lighting. +Realistic, flawless skin. +# White people should have a pale skin. +age everyone by 5 years and make them look like 25 year old adults. + +keep the same ethnicity and facial features. + +Make the children taller and more slender with slightly elongated feminine neck and higher cheekbones, give them smaller, slimmer more feminine head and tapered face; round, smooth chin. + +keep everything else about her the same. +same facial + eye expression (keep them as open as in the original), same pose and emotion, same tongue position. +same head and eye angle. + +do not turn her face around. +same hair style, color, length. +same composition and crop. +keep the same makeup, do not change the colors of the lips, keep them glossy and wet. +long nipples. large areolas. + +""" + + +def preprocess_image(image_path: Path) -> bytes: + with Image.open(image_path) as img: + if img.mode in ("RGBA", "P"): + img = img.convert("RGB") + max_size = 2500 # Slightly larger to support 3MP render + if max(img.size) > max_size: + img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS) + buffer = io.BytesIO() + img.save(buffer, format="JPEG", quality=85, optimize=True) + return buffer.getvalue() + + +def _next_output_path(image_path: Path) -> Path: + base = image_path.with_name(f"{image_path.stem}-edit.jpg") + if not base.exists(): + return base + counter = 2 + while True: + candidate = image_path.with_name(f"{image_path.stem}-edit{counter}.jpg") + if not candidate.exists(): + return candidate + counter += 1 + + +def tag_file(path: Path, tag: str = "Blue") -> None: + if not path.exists(): + raise FileNotFoundError(f"Target path does not exist: {path}") + + abs_path = str(path.resolve()) + + script = f'tell application "Finder" to set tags of (POSIX file "{abs_path}" as alias) to {{"{tag}"}}' + + try: + subprocess.run( + ["osascript", "-e", script], check=True, capture_output=True, text=True + ) + except subprocess.CalledProcessError as e: + # Finder may fail if the file is moved during execution or if permissions are insufficient + logger.warning(f"Failed to tag file: {e.stderr}") + + +class Stats: + def __init__(self) -> None: + self.run_times: list[float] = [] + + def add(self, elapsed: float) -> None: + self.run_times.append(elapsed) + + def format_time(self, seconds: float) -> str: + if seconds < 60: + return f"{seconds:.1f}s" + minutes = seconds / 60 + return f"{minutes:.1f}m" + + def print_summary(self) -> None: + if not self.run_times: + return + + if len(self.run_times) == 1: + return + + shortest = min(self.run_times) + longest = max(self.run_times) + average = sum(self.run_times) / len(self.run_times) + total = sum(self.run_times) + + print("=" * 20) + print("Run Statistics") + print("=" * 20) + print(f"Total runs: {len(self.run_times)}") + print(f"Shortest: {self.format_time(shortest)}") + print(f"Longest: {self.format_time(longest)}") + print(f"Average: {self.format_time(average)}") + print(f"Total time: {self.format_time(total)}") + print("=" * 20) + + +class ComfyUI: + PROMPT_NODE_ID = "75:74" + NEGATIVE_NODE_ID = "75:67" + IMAGE_LOAD_NODE_ID = "76" + SAVE_NODE_ID = "9" + UPSCALE_NODE_ID = "75:80" + MODEL_LOAD_NODE_ID = "75:70" + CLIP_LOAD_NODE_ID = "75:71" + VAE_LOAD_NODE_ID = "75:72" + SEED_NODE_ID = "75:73" + + def __init__(self, url: str, cookies: dict) -> None: + self.url = url.rstrip("/") + self.client = httpx.Client(timeout=None, cookies=cookies, follow_redirects=True) + + def _post(self, path: str, **kwargs) -> httpx.Response: + """Try RunPod /api routes first, then legacy Comfy routes.""" + candidates = [f"{self.url}/api{path}", f"{self.url}{path}"] + last_res: httpx.Response | None = None + for endpoint in candidates: + res = self.client.post(endpoint, **kwargs) + if res.status_code != 404: + return res + last_res = res + assert last_res is not None + return last_res + + def _get(self, path: str, **kwargs) -> httpx.Response: + """Try RunPod /api routes first, then legacy Comfy routes.""" + candidates = [f"{self.url}/api{path}", f"{self.url}{path}"] + last_res: httpx.Response | None = None + for endpoint in candidates: + res = self.client.get(endpoint, **kwargs) + if res.status_code != 404: + return res + last_res = res + assert last_res is not None + return last_res + + def transform( + self, + image_path: Path, + save_path: Path, + prompt: str, + negative_prompt: str = "", + ) -> None: + seed = random.randint(0, 2**31 - 1) + client_id = str(uuid.uuid4()) + + # 1. Upload + processed_img_bytes = preprocess_image(image_path) + files = {"image": (f"upload_{uuid.uuid4()}.jpg", processed_img_bytes)} + res = self._post("/upload/image", files=files) + if res.is_error: + logger.error(f"Upload error: {res.text}") + res.raise_for_status() + filename = res.json()["name"] + + # 2. Load Workflow + with open(WORKFLOW_FILE, "r") as f: + workflow = json.load(f) + + # Get dimensions from the actual file + with Image.open(image_path) as img: + orig_width, orig_height = img.size + aspect_ratio = orig_width / orig_height + + # Calculate target resolution for 3 Megapixels (approx 3,000,000 pixels) + # Formula: Width * (Width / Aspect) = TotalPixels + target_pixels = IMAGE_OUTPUT_MEGAPIXELS * 1024 * 1024 + new_width = int(math.sqrt(target_pixels * aspect_ratio)) + new_height = int(new_width / aspect_ratio) + + # Force these into the Empty Latent node (Node 75:66) + # This ensures the "canvas" matches your photo's shape + workflow["75:66"]["inputs"]["width"] = ( + new_width // 8 + ) * 8 # Must be multiple of 8 + workflow["75:66"]["inputs"]["height"] = (new_height // 8) * 8 + + # Also update the Scheduler (Node 75:62) so Flux knows the scale + workflow["75:62"]["inputs"]["width"] = workflow["75:66"]["inputs"]["width"] + workflow["75:62"]["inputs"]["height"] = workflow["75:66"]["inputs"]["height"] + + # Inject Advanced Data + workflow[self.PROMPT_NODE_ID]["inputs"]["text"] = prompt + workflow[self.NEGATIVE_NODE_ID]["inputs"]["text"] = negative_prompt + workflow[self.IMAGE_LOAD_NODE_ID]["inputs"]["image"] = filename + workflow[self.SEED_NODE_ID]["inputs"]["noise_seed"] = seed + + workflow[self.UPSCALE_NODE_ID]["inputs"]["megapixels"] = IMAGE_OUTPUT_MEGAPIXELS + workflow[self.UPSCALE_NODE_ID]["inputs"]["upscale_method"] = "lanczos" + + # 3. Queue + payload = {"prompt": workflow, "client_id": client_id} + res = self._post("/prompt", json=payload) + + if res.is_error: + logger.error(f"Error: {res.text}") + res.raise_for_status() + + prompt_id = res.json()["prompt_id"] + + # 4. Wait + preview_output = None + while True: + jobs_res = self._get( + "/jobs", + params={ + "status": "completed,failed,cancelled", + "limit": 64, + "offset": 0, + }, + ) + if jobs_res.is_error: + logger.error(f"Jobs poll error: {jobs_res.text}") + jobs_res.raise_for_status() + + jobs = jobs_res.json().get("jobs", []) + match = next((job for job in jobs if job.get("id") == prompt_id), None) + if match: + if match.get("status") != "completed": + raise RuntimeError( + f"Job {prompt_id} ended with status: {match.get('status')}" + ) + preview_output = match.get("preview_output") + break + time.sleep(0.25) + + # 5. Download + if not preview_output: + raise RuntimeError(f"Completed job {prompt_id} has no preview_output") + + view_params = { + "filename": preview_output["filename"], + "subfolder": preview_output.get("subfolder", ""), + "type": preview_output.get("type", "output"), + } + img_res = self._get("/view", params=view_params) + if img_res.is_error: + logger.error(f"View error: {img_res.text}") + img_res.raise_for_status() + + # Convert to JPEG with quality=75 + with Image.open(io.BytesIO(img_res.content)) as img: + if img.mode == "RGBA": + img = img.convert("RGB") + img.save(save_path, format="JPEG", quality=75) + +def without_comments(s: str) -> str: + return "\n".join( + line for line in s.splitlines() if not line.strip().startswith("#") + ) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("images", nargs="+", type=Path) + args = parser.parse_args() + + + comfy = ComfyUI(COMFY_URL, COMFY_COOKIES) + stats = Stats() + + for i, image_path in enumerate(args.images, start=1): + prompt = SYSTEM_PROMPT + replace_prompt_path = Path(__file__).parent / "flux_klein_prompt.txt" + if replace_prompt_path.is_file(): + prompt = without_comments(replace_prompt_path.read_text()) + + save_path = _next_output_path(image_path) + start_time = time.time() + logger.info(f"Uploading ({i}/{len(args.images)}): {image_path.name}") + try: + comfy.transform(image_path, save_path, prompt) + except Exception as e: + logger.error(f"Error processing {image_path.name}: {e}") + continue + elapsed = time.time() - start_time + stats.add(elapsed) + + stats.print_summary() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backup-site.service b/backup-site.service new file mode 100644 index 0000000..e69de29 diff --git a/backup-site.timer b/backup-site.timer new file mode 100644 index 0000000..e69de29 diff --git a/backup_site.py b/backup_site.py new file mode 100644 index 0000000..80f8299 --- /dev/null +++ b/backup_site.py @@ -0,0 +1,144 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = [] +# /// + +import argparse +import logging +import os +import shutil +import subprocess +import zipfile +from datetime import datetime +from pathlib import Path +from urllib.parse import urlparse + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Backup MariaDB database and website files", formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("--db-url", default=os.getenv("DB_URL"), help="Database URL (default: from DB_URL env var)") + parser.add_argument("--site-dir", type=Path, required=True, help="Website directory to backup") + parser.add_argument("--backup-dir", type=Path, default="/mnt/box/backup/droplet", help="Backup destination directory") + return parser.parse_args() + + +def parse_db_url(db_url: str) -> dict: + if not db_url: + raise ValueError("Database URL is required") + + parsed = urlparse(db_url) + + if parsed.scheme != "mysql": + raise ValueError("Unsupported database URL format") + + return { + "host": parsed.hostname or "localhost", + "port": str(parsed.port or 3306), + "user": parsed.username or "root", + "password": parsed.password or "", + "database": parsed.path.lstrip("/") if parsed.path else "", + } + + +def dump_database(db_params: dict, dump_path: Path) -> None: + logger.info(f"Dumping database to {dump_path}") + + cmd = [ + "mysqldump", + "--host", + db_params["host"], + "--port", + db_params["port"], + "--user", + db_params["user"], + "--single-transaction", + "--routines", + "--triggers", + db_params["database"], + ] + + if db_params["password"]: + cmd.append(f"--password={db_params['password']}") + + try: + with dump_path.open("w") as f: + subprocess.run(cmd, stdout=f, stderr=subprocess.PIPE, text=True, check=True) + logger.info("Database dump completed successfully") + except subprocess.CalledProcessError as e: + logger.error(f"Database dump failed: {e.stderr}") + raise + + +def zip_directory(source_dir: Path, zip_path: Path) -> None: + """Zip a directory""" + logger.info(f"Zipping directory {source_dir} to {zip_path}") + + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: + for file_path in source_dir.rglob("*"): + if file_path.is_file(): + # Use relative path within the zip + arcname = file_path.relative_to(source_dir.parent) + zipf.write(file_path, arcname) + + logger.info(f"Directory zipped successfully: {zip_path}") + + +def create_final_backup(temp_dir: Path, backup_dir: Path, timestamp: str) -> Path: + final_backup_name = f"{timestamp}_ucsuzkalem.zip" + final_backup_path = backup_dir / final_backup_name + + logger.info(f"Creating final backup: {final_backup_path}") + zip_directory(temp_dir, final_backup_path) + + return final_backup_path + + +def main(): + args = parse_args() + + if not args.db_url: + raise ValueError("Database URL is required (set DB_URL env var or use --db-url)") + + # Parse database URL + db_params = parse_db_url(args.db_url) + logger.info(f"Connecting to database: {db_params['host']}:{db_params['port']}/{db_params['database']}") + + # Create timestamp + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + + # Setup paths + site_dir = Path(args.site_dir) + backup_dir = Path(args.backup_dir) + temp_dir = Path("/tmp") / f"backup_{timestamp}" + + if not site_dir.exists(): + raise FileNotFoundError(f"Site directory does not exist: {site_dir}") + + # Create directories + backup_dir.mkdir(parents=True, exist_ok=True) + temp_dir.mkdir(parents=True, exist_ok=True) + + # Dump database + db_dump_path = temp_dir / "database.sql" + dump_database(db_params, db_dump_path) + + # Zip website files + + files_zip_path = temp_dir / "files.zip" + zip_directory(site_dir, files_zip_path) + + # Create final backup + final_backup_path = create_final_backup(temp_dir, backup_dir, timestamp) + + # Cleanup temp directory + shutil.rmtree(temp_dir) + + logger.info(f"Backup completed successfully: {final_backup_path}") + logger.info(f"Backup size: {final_backup_path.stat().st_size / (1024 * 1024):.1f} MB") + + +if __name__ == "__main__": + exit(main()) diff --git a/classify_image.py b/classify_image.py new file mode 100755 index 0000000..11f754e --- /dev/null +++ b/classify_image.py @@ -0,0 +1,738 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.13" +# dependencies = ["ultralytics", "torch", "numpy", "pillow", "bottle", "mediapipe"] +# /// +import argparse +import json +import logging +import sys +import tempfile +import threading +import urllib.request +import webbrowser +from dataclasses import dataclass +from functools import cache +from pathlib import Path +from typing import NamedTuple, Optional + +try: + import bottle +except Exception: + bottle = None + +try: + from ultralytics import YOLO + import torch + import numpy as np +except Exception: + YOLO = None + torch = None + np = None + +try: + from PIL import Image +except Exception: + Image = None + +try: + import mediapipe as mp +except Exception: + mp = None + +VISIBILITY_THRESH = 0.1 +FACE_LANDMARK_VISIBILITY_THRESH = 0.5 +FACING_DIRECTION_THRESH = 0.03 +LOOKING_AT_CAMERA_THRESH = 0.015 # max |nose.x - eye_midpoint.x| for frontal face +HEAD_CROP_Y_THRESH = 0.35 # shoulders must be in top 35% of frame to classify as head-cropped +NOSE_TO_NOSE_THRESH = 0.15 # normalized image distance +BODY_INTERSECT_MARGIN = 0.05 # expand each person's bbox by this before overlap test +EYE_BLINK_THRESH = 0.40 # blendshape score above this → eye closed + + +class Coords(NamedTuple): + """Represents normalized coordinates (0.0 to 1.0) and visibility for a single point.""" + + x: float + y: float + is_visible: bool + + +@dataclass +class PoseKeypoints: + """Holds structured, normalized keypoint data for all 17 COCO points as direct fields.""" + + # 0 + nose: Coords + # 1-4 + left_eye: Coords + right_eye: Coords + left_ear: Coords + right_ear: Coords + # 5-6 + left_shoulder: Coords + right_shoulder: Coords + # 7-10 + left_elbow: Coords + right_elbow: Coords + left_wrist: Coords + right_wrist: Coords + # 11-12 + left_hip: Coords + right_hip: Coords + # 13-16 + left_knee: Coords + right_knee: Coords + left_ankle: Coords + right_ankle: Coords + + def shoulder_midpoint(self) -> Coords: + l = self.left_shoulder + r = self.right_shoulder + if l.is_visible and r.is_visible: + return Coords(x=(l.x + r.x) / 2.0, y=(l.y + r.y) / 2.0, is_visible=True) + return Coords(x=(l.x + r.x) / 2.0, y=(l.y + r.y) / 2.0, is_visible=False) + + +@dataclass +class FaceLandmarks: + """5-point face landmarks from the derronqi yolov8-face model.""" + + left_eye: Coords + right_eye: Coords + nose: Coords + left_mouth: Coords + right_mouth: Coords + + +def _get_coords(kp_xyc: "np.ndarray", idx: int) -> Coords: + """Helper to safely extract Coords from the raw numpy array.""" + x, y, conf = kp_xyc[idx] + is_visible = conf > VISIBILITY_THRESH + return Coords(x=x, y=y, is_visible=is_visible) + + +def _extract_keypoints(kp_xyc: "np.ndarray") -> PoseKeypoints: + """Extracts all 17 COCO normalized keypoints and populates the PoseKeypoints dataclass directly.""" + return PoseKeypoints( + nose=_get_coords(kp_xyc, 0), + left_eye=_get_coords(kp_xyc, 1), + right_eye=_get_coords(kp_xyc, 2), + left_ear=_get_coords(kp_xyc, 3), + right_ear=_get_coords(kp_xyc, 4), + left_shoulder=_get_coords(kp_xyc, 5), + right_shoulder=_get_coords(kp_xyc, 6), + left_elbow=_get_coords(kp_xyc, 7), + right_elbow=_get_coords(kp_xyc, 8), + left_wrist=_get_coords(kp_xyc, 9), + right_wrist=_get_coords(kp_xyc, 10), + left_hip=_get_coords(kp_xyc, 11), + right_hip=_get_coords(kp_xyc, 12), + left_knee=_get_coords(kp_xyc, 13), + right_knee=_get_coords(kp_xyc, 14), + left_ankle=_get_coords(kp_xyc, 15), + right_ankle=_get_coords(kp_xyc, 16), + ) + + +_MODEL_URLS = { + "yolov11n-face.pt": "https://huggingface.co/AdamCodd/YOLOv11n-face-detection/resolve/main/model.pt", + "yolo11s-pose.pt": "https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11s-pose.pt", + "yolov8n-face-derronqi.pt": "https://huggingface.co/junjiang/GestureFace/resolve/main/yolov8n-face.pt", + "face_landmarker.task": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task", +} + + +def _ensure_model(filename: str) -> Path: + dest = Path(__file__).parent / filename + if not dest.exists(): + url = _MODEL_URLS[filename] + logging.warning(f"Downloading {filename} from {url} ...") + urllib.request.urlretrieve(url, dest) + logging.warning(f"Saved {filename}") + return dest + + +def _best_device() -> str: + try: + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + except Exception: + pass + return "cpu" + + +@cache +def _face_detector(): + if YOLO is None or torch is None: + logging.error("YOLO/Torch dependencies are missing.") + return None + try: + model = YOLO(_ensure_model("yolov11n-face.pt")) + model.to(_best_device()) + return model + except Exception: + logging.exception("Face detector model initialization failed.") + return None + + +@cache +def _pose_detector(): + if YOLO is None or torch is None: + logging.error("YOLO/Torch dependencies are missing.") + return None + try: + model = YOLO(_ensure_model("yolo11s-pose.pt")) + model.to(_best_device()) + return model + except Exception: + logging.exception("Pose detector model initialization failed.") + return None + + +@cache +def _face_landmark_detector(): + if YOLO is None or torch is None: + logging.error("YOLO/Torch dependencies are missing.") + return None + try: + model = YOLO(_ensure_model("yolov8n-face-derronqi.pt")) + model.to(_best_device()) + return model + except Exception: + logging.exception("Face landmark detector model initialization failed.") + return None + + +@cache +def _face_landmarker(): + if mp is None: + logging.error("mediapipe dependency is missing.") + return None + try: + from mediapipe.tasks import python as mp_python + from mediapipe.tasks.python import vision as mp_vision + base_options = mp_python.BaseOptions( + model_asset_path=str(_ensure_model("face_landmarker.task")) + ) + options = mp_vision.FaceLandmarkerOptions( + base_options=base_options, + output_face_blendshapes=True, + running_mode=mp_vision.RunningMode.IMAGE, + num_faces=10, + ) + return mp_vision.FaceLandmarker.create_from_options(options) + except Exception: + logging.exception("FaceLandmarker initialization failed.") + return None + + +def _run_face_landmarker_model(image_path: Path): + detector = _face_landmarker() + if detector is None or Image is None: + return None + try: + img = Image.open(image_path).convert("RGB") + arr = np.asarray(img) + mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=arr) + return detector.detect(mp_image) + except Exception: + logging.exception(f"FaceLandmarker failed for {image_path}.") + return None + + +def check_eyes_closed(face_landmarker_result) -> bool: + if face_landmarker_result is None: + return False + if not face_landmarker_result.face_blendshapes: + return False + for face_blendshapes in face_landmarker_result.face_blendshapes: + scores = {b.category_name: b.score for b in face_blendshapes} + if scores.get("eyeBlinkLeft", 0.0) > EYE_BLINK_THRESH or \ + scores.get("eyeBlinkRight", 0.0) > EYE_BLINK_THRESH: + return True + return False + + +def _run_face_model(image_path: Path) -> list: + model = _face_detector() + if model is None: + return [] + try: + return model(str(image_path), conf=0.5, iou=0.5, verbose=False) + except Exception: + logging.exception(f"Face detection failed for {image_path}.") + return [] + + +def _run_face_landmark_model(image_path: Path) -> list: + model = _face_landmark_detector() + if model is None: + return [] + try: + return model(str(image_path), conf=0.25, iou=0.5, verbose=False) + except Exception: + logging.exception(f"Face landmark detection failed for {image_path}.") + return [] + + +def _run_pose_model(image_path: Path) -> list: + model = _pose_detector() + if model is None: + return [] + try: + return model(str(image_path), conf=0.35, iou=0.5, verbose=False) + except Exception: + logging.exception(f"Pose detection failed for {image_path}.") + return [] + + +def count_faces(results: list) -> int: + try: + return len(results[0].boxes) + except Exception: + return 0 + + +def detect_poses(results: list) -> list[PoseKeypoints]: + try: + if not results or results[0].keypoints is None: + return [] + kps_norm_xy = results[0].keypoints.xyn.cpu().numpy() + kps_conf = results[0].keypoints.conf.cpu().numpy() + all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2) + return [_extract_keypoints(kp_xyc) for kp_xyc in all_keypoints_xyc] + except Exception: + logging.exception("Pose keypoint extraction failed.") + return [] + + +def detect_face_landmarks(results: list) -> list[FaceLandmarks]: + try: + if not results or results[0].keypoints is None: + return [] + kps_norm_xy = results[0].keypoints.xyn.cpu().numpy() + kps_conf = results[0].keypoints.conf.cpu().numpy() + all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2) + landmarks = [] + for kp_xyc in all_keypoints_xyc: + def _lm(idx, arr=kp_xyc): + x, y, conf = arr[idx] + return Coords(x=x, y=y, is_visible=conf > FACE_LANDMARK_VISIBILITY_THRESH) + landmarks.append(FaceLandmarks( + left_eye=_lm(0), + right_eye=_lm(1), + nose=_lm(2), + left_mouth=_lm(3), + right_mouth=_lm(4), + )) + return landmarks + except Exception: + logging.exception("Face landmark extraction failed.") + return [] + + +def _face_all_invisible(kps: PoseKeypoints) -> bool: + return ( + not kps.nose.is_visible + and not kps.left_eye.is_visible + and not kps.right_eye.is_visible + ) + + +def is_turned_back(kps: PoseKeypoints) -> bool: + shoulder_visible = kps.left_shoulder.is_visible or kps.right_shoulder.is_visible + return _face_all_invisible(kps) and shoulder_visible and not is_eyes_cropped_out(kps) + + +def is_eyes_cropped_out(kps: PoseKeypoints) -> bool: + """True when face is not visible but shoulders are near the top of the frame, + indicating the head is above the image boundary.""" + shoulder_visible = kps.left_shoulder.is_visible or kps.right_shoulder.is_visible + if not _face_all_invisible(kps) or not shoulder_visible: + return False + # Use the topmost (lowest y) visible shoulder + ys = [kp.y for kp in (kps.left_shoulder, kps.right_shoulder) if kp.is_visible] + return min(ys) < HEAD_CROP_Y_THRESH + + +def get_facing_x_direction(kps: PoseKeypoints) -> Optional[float]: + shoulder_mid = kps.shoulder_midpoint() + if not kps.nose.is_visible or not shoulder_mid.is_visible: + return None + return kps.nose.x - shoulder_mid.x + + +def get_face_yaw(kps: FaceLandmarks) -> Optional[float]: + """Nose x offset from eye midpoint. ~0 = frontal, positive = turned right, negative = turned left.""" + if not kps.left_eye.is_visible or not kps.right_eye.is_visible or not kps.nose.is_visible: + return None + eye_mid_x = (kps.left_eye.x + kps.right_eye.x) / 2.0 + return kps.nose.x - eye_mid_x + + +def check_looking_at_camera(all_landmarks: list[FaceLandmarks]) -> bool: + return any( + (yaw := get_face_yaw(lm)) is not None and abs(yaw) < LOOKING_AT_CAMERA_THRESH + for lm in all_landmarks + ) + + +def check_facing_each_other(all_kps: list[PoseKeypoints]) -> bool: + classifiable = [] + for kps in all_kps: + delta = get_facing_x_direction(kps) + if delta is not None: + shoulder_mid = kps.shoulder_midpoint() + classifiable.append((shoulder_mid.x, delta)) + + if len(classifiable) < 2: + return False + + classifiable.sort(key=lambda t: t[0]) + + for i in range(len(classifiable)): + for j in range(i + 1, len(classifiable)): + left_delta = classifiable[i][1] + right_delta = classifiable[j][1] + if left_delta > FACING_DIRECTION_THRESH and right_delta < -FACING_DIRECTION_THRESH: + return True + + return False + + +def detect_person_boxes(results: list) -> list[tuple[float, float, float, float]]: + try: + if not results or results[0].boxes is None: + return [] + return [tuple(box) for box in results[0].boxes.xyxyn.cpu().tolist()] + except Exception: + logging.exception("Person box extraction failed.") + return [] + + +def _face_boxes_to_body_boxes( + face_boxes: list[tuple[float, float, float, float]] +) -> list[tuple[float, float, float, float]]: + out = [] + for x1, y1, x2, y2 in face_boxes: + fw = x2 - x1 + fh = y2 - y1 + bx1 = max(0.0, x1 - 0.3 * fw) + bx2 = min(1.0, x2 + 0.3 * fw) + by1 = y1 + by2 = min(1.0, y2 + 2.5 * fh) # extend ~2.5 face-heights downward + out.append((bx1, by1, bx2, by2)) + return out + + +def check_bodies_intersecting(person_boxes: list[tuple[float, float, float, float]]) -> bool: + m = BODY_INTERSECT_MARGIN + for i in range(len(person_boxes)): + for j in range(i + 1, len(person_boxes)): + ax1, ay1, ax2, ay2 = person_boxes[i] + bx1, by1, bx2, by2 = person_boxes[j] + if ax1 - m < bx2 and ax2 + m > bx1 and ay1 - m < by2 and ay2 + m > by1: + return True + return False + + +def check_nose_to_nose(all_kps: list[PoseKeypoints]) -> bool: + visible = [(kps.nose.x, kps.nose.y) for kps in all_kps if kps.nose.is_visible] + for i in range(len(visible)): + for j in range(i + 1, len(visible)): + dx = visible[i][0] - visible[j][0] + dy = visible[i][1] - visible[j][1] + if (dx*dx + dy*dy) ** 0.5 < NOSE_TO_NOSE_THRESH: + return True + return False + + +def classify_image(image_path: Path, debug_dir: Optional[Path] = None) -> dict: + face_results = _run_face_model(image_path) + pose_results = _run_pose_model(image_path) + face_landmark_results = _run_face_landmark_model(image_path) + face_mesh_result = _run_face_landmarker_model(image_path) + + total_faces = count_faces(face_results) + all_kps = detect_poses(pose_results) + face_landmarks = detect_face_landmarks(face_landmark_results) + turned_back = any(is_turned_back(kps) for kps in all_kps) + facing_each_other = check_facing_each_other(all_kps) + eyes_cropped_out = any(is_eyes_cropped_out(kps) for kps in all_kps) + nose_to_nose = check_nose_to_nose(all_kps) + person_boxes = detect_person_boxes(pose_results) + if len(person_boxes) < 2: + raw_face_boxes = [] + try: + if face_results and face_results[0].boxes is not None: + raw_face_boxes = [tuple(b) for b in face_results[0].boxes.xyxyn.cpu().tolist()] + except Exception: + pass + person_boxes = _face_boxes_to_body_boxes(raw_face_boxes) + bodies_intersecting = check_bodies_intersecting(person_boxes) + looking_at_camera = check_looking_at_camera(face_landmarks) + + if debug_dir: + debug_dir.mkdir(parents=True, exist_ok=True) + for label, results in [("face", face_results), ("pose", pose_results)]: + if results and results[0].boxes is not None: + annotated = results[0].plot() + img = Image.fromarray(annotated[..., ::-1]) + img.save(debug_dir / f"{image_path.stem}_{label}.jpg") + + return { + "image_path": str(image_path), + "total_faces": total_faces, + "facing_each_other": facing_each_other, + "turned_back": turned_back, + "eyes_cropped_out": eyes_cropped_out, + "nose_to_nose": nose_to_nose, + "bodies_intersecting": bodies_intersecting, + "looking_at_camera": looking_at_camera, + "eyes_closed": check_eyes_closed(face_mesh_result), + } + + +HTML_PAGE = """ + + + +classify_image debug + + + +
Drop images here to classify
+
+ + +""" + + +def _extract_web_data(face_results: list, pose_results: list) -> dict: + faces = [] + try: + if face_results and face_results[0].boxes is not None: + faces = face_results[0].boxes.xyxyn.cpu().tolist() + except Exception: + pass + + poses = [] + try: + if pose_results and pose_results[0].keypoints is not None: + kps_norm_xy = pose_results[0].keypoints.xyn.cpu().numpy() + kps_conf = pose_results[0].keypoints.conf.cpu().numpy() + all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2) + for person_kps in all_keypoints_xyc: + poses.append([ + {"x": float(kp[0]), "y": float(kp[1]), "v": bool(kp[2] > VISIBILITY_THRESH)} + for kp in person_kps + ]) + except Exception: + pass + + return {"faces": faces, "poses": poses} + + +def _classify_route(): + upload = bottle.request.files.get("image") + suffix = Path(upload.filename).suffix or ".jpg" + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f: + tmp = Path(f.name) + upload.save(f) + try: + face_results = _run_face_model(tmp) + pose_results = _run_pose_model(tmp) + face_landmark_results = _run_face_landmark_model(tmp) + face_mesh_result = _run_face_landmarker_model(tmp) + all_kps = detect_poses(pose_results) + face_landmarks = detect_face_landmarks(face_landmark_results) + person_boxes = detect_person_boxes(pose_results) + if len(person_boxes) < 2: + raw_face_boxes = [] + try: + if face_results and face_results[0].boxes is not None: + raw_face_boxes = [tuple(b) for b in face_results[0].boxes.xyxyn.cpu().tolist()] + except Exception: + pass + person_boxes = _face_boxes_to_body_boxes(raw_face_boxes) + classification = { + "image": upload.filename, + "total_faces": count_faces(face_results), + "facing_each_other": check_facing_each_other(all_kps), + "turned_back": any(is_turned_back(kps) for kps in all_kps), + "eyes_cropped_out": any(is_eyes_cropped_out(kps) for kps in all_kps), + "nose_to_nose": check_nose_to_nose(all_kps), + "bodies_intersecting": check_bodies_intersecting(person_boxes), + "looking_at_camera": check_looking_at_camera(face_landmarks), + "eyes_closed": check_eyes_closed(face_mesh_result), + } + detections = _extract_web_data(face_results, pose_results) + return bottle.HTTPResponse( + json.dumps({"classification": classification, "detections": detections}), + content_type="application/json", + ) + finally: + tmp.unlink(missing_ok=True) + + +def web_main(): + if bottle is None: + print("bottle not installed. Run: pip install bottle", file=sys.stderr) + sys.exit(1) + + app = bottle.Bottle() + + @app.get("/") + def index(): + return HTML_PAGE + + @app.post("/classify") + def classify_route(): + return _classify_route() + + port = 7777 + threading.Timer(0.5, lambda: webbrowser.open(f"http://localhost:{port}")).start() + bottle.run(app, host="localhost", port=port, quiet=True) + + +def main(): + parser = argparse.ArgumentParser(description="Classify images for face-related attributes.") + parser.add_argument("image_paths", nargs="*", type=Path, help="Path(s) to image(s)") + parser.add_argument("--debug", action="store_true", help="Save annotated debug images to _debug/ subdirectory") + parser.add_argument("--web", action="store_true", help="Start debug web server") + args = parser.parse_args() + + logging.basicConfig(level=logging.WARNING, stream=sys.stderr) + + if args.web: + web_main() + return + + if not args.image_paths: + parser.print_usage(sys.stderr) + sys.exit(1) + + for image_path in args.image_paths: + try: + debug_dir = (image_path.parent / "_debug") if args.debug else None + result = classify_image(image_path, debug_dir=debug_dir) + print(json.dumps(result), flush=True) + except KeyboardInterrupt: + raise + except Exception: + logging.exception(f"Error processing {image_path}.") + + +if __name__ == "__main__": + main() diff --git a/cluster_images.py b/cluster_images.py new file mode 100755 index 0000000..d86d6f0 --- /dev/null +++ b/cluster_images.py @@ -0,0 +1,163 @@ +#!/usr/bin/env -S uv run +# /// script +# dependencies = [ +# "torch", +# "torchvision", +# "timm", +# "scikit-learn", +# "pillow", +# "imagehash", +# ] +# /// +"""Group similar large images using ResNet embeddings or perceptual hashing.""" + +import torch +import timm +import numpy as np +from PIL import Image +from pathlib import Path +from sklearn.cluster import KMeans, DBSCAN +from sklearn.preprocessing import StandardScaler +import imagehash +import argparse + +# Load pretrained ResNet50 on M1 +device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") +model = timm.create_model("resnet50", pretrained=True, num_classes=0) +model = model.to(device) +model.eval() + + +def get_embedding(img_path, size=224): + """Extract embedding from image.""" + try: + img = Image.open(img_path).convert("RGB") + img.thumbnail((size, size), Image.Resampling.LANCZOS) + + canvas = Image.new("RGB", (size, size), (128, 128, 128)) + offset = ((size - img.width) // 2, (size - img.height) // 2) + canvas.paste(img, offset) + + x = torch.tensor(np.array(canvas), dtype=torch.float32) + x = x.permute(2, 0, 1) / 255.0 + x = (x - torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)) / torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) + + with torch.no_grad(): + embedding = model(x.unsqueeze(0).to(device)).squeeze().cpu().numpy() + return embedding + except Exception as e: + print(f"Error processing {img_path}: {e}") + return None + + +def get_phash(img_path): + """Extract perceptual hash from image.""" + try: + img = Image.open(img_path).convert("RGB") + return imagehash.phash(img) + except Exception as e: + print(f"Error processing {img_path}: {e}") + return None + + +def hash_distance(h1, h2): + """Hamming distance between two hashes.""" + return h1 - h2 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Group similar images using embeddings or perceptual hashing.") + parser.add_argument("paths", nargs="+", type=Path, help="Image file(s) or directory") + parser.add_argument("--cluster", action="store_true", help="Use ResNet embedding clustering") + parser.add_argument("--perceptual-hash", action="store_true", help="Use perceptual hashing") + parser.add_argument("--clusters", type=int, help="Number of clusters for embedding mode (auto if not specified)") + parser.add_argument("--hash-threshold", type=int, default=5, help="Hamming distance threshold for perceptual hash") + args = parser.parse_args() + + if not args.cluster and not args.perceptual_hash: + parser.error("Either --cluster or --perceptual-hash must be specified") + if args.cluster and args.perceptual_hash: + parser.error("Cannot specify both --cluster and --perceptual-hash") + + return args + + +def group_by_hash(valid_files, hashes, threshold): + """Group images by perceptual hash similarity.""" + labels = [-1] * len(valid_files) + cluster_id = 0 + + for i in range(len(valid_files)): + if labels[i] != -1: + continue + labels[i] = cluster_id + for j in range(i + 1, len(valid_files)): + if labels[j] == -1 and hash_distance(hashes[i], hashes[j]) <= threshold: + labels[j] = cluster_id + cluster_id += 1 + + return np.array(labels) + + +def main(): + args = parse_args() + + # Collect image files + img_files = [] + for path in args.paths: + if path.is_dir(): + img_files.extend(path.glob("*.[jJ][pP][gG]")) + img_files.extend(path.glob("*.[pP][nN][gG]")) + else: + img_files.append(path) + + if not img_files: + print("No images found.") + return + + if args.perceptual_hash: + # Perceptual hash mode + hashes = [] + valid_files = [] + + for i, img_path in enumerate(img_files): + print(f"Processing {i + 1}/{len(img_files)}: {img_path.name}") + h = get_phash(img_path) + if h is not None: + hashes.append(h) + valid_files.append(img_path) + + labels = group_by_hash(valid_files, hashes, args.hash_threshold) + n_clusters = len(np.unique(labels)) + + else: + # Embedding clustering mode + embeddings = [] + valid_files = [] + + for i, img_path in enumerate(img_files): + print(f"Processing {i + 1}/{len(img_files)}: {img_path.name}") + emb = get_embedding(img_path) + if emb is not None: + embeddings.append(emb) + valid_files.append(img_path) + + embeddings = np.array(embeddings) + scaler = StandardScaler() + embeddings = scaler.fit_transform(embeddings) + + n_clusters = args.clusters or max(2, int(np.sqrt(len(embeddings) / 2))) + kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10) + labels = kmeans.fit_predict(embeddings) + + # Save clustered images + for img_path, cluster_label in zip(valid_files, labels): + new_name = f"group{cluster_label}_{img_path.name}" + dest = Path.cwd() / new_name + img_path.rename(dest) if img_path.parent == Path.cwd() else dest.write_bytes(img_path.read_bytes()) + + print(f"Clustered {len(valid_files)} images into {n_clusters} clusters in {Path.cwd()}") + + +if __name__ == "__main__": + main() diff --git a/cull_edits.py b/cull_edits.py new file mode 100755 index 0000000..3b6569e --- /dev/null +++ b/cull_edits.py @@ -0,0 +1,939 @@ +#!/usr/bin/env -S uv run +# /// script +# requires-python = ">=3.14" +# dependencies = [ +# "bottle>=0.13.0", +# ] +# /// + +import argparse +import functools +import os +import re +import socket +import sys +import threading +import time +import webbrowser +from pathlib import Path +from datetime import datetime +from typing import TypedDict +import bottle + +IMAGE_EXTENSIONS: set[str] = {".jpg", ".jpeg", ".png", ".webp", ".webm"} + + +class ImageGroup(TypedDict): + """Type definition for grouped images.""" + + base: str + original: Path + edits: list[Path] + + +def find_available_port() -> int: + """Find an available random port.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + s.listen(1) + port = s.getsockname()[1] + return port + + +def is_image(path: str | Path) -> bool: + """Check if file is a supported image.""" + return Path(path).suffix.lower() in IMAGE_EXTENSIONS + + +def discover_images(paths: list[str]) -> list[Path]: + """Discover all images from mixed list of files and folders.""" + images = [] + for path_str in paths: + path = Path(path_str) + if path.is_file() and is_image(path): + images.append(path.resolve()) + elif path.is_dir(): + for item in path.iterdir(): + if item.is_file() and is_image(item): + images.append(item.resolve()) + return sorted(set(images)) + + +def group_images(images: list[Path]) -> list[ImageGroup]: + """ + Group images by original + edits. + Original: image.jpg + Edits: image-edit1.jpg, image-edit2.jpg, etc. + Returns list of dicts: {original: Path, edits: [Path, ...]} + """ + groups = {} + + for img in images: + stem = img.stem + + # Check if this is an edit + match = re.match(r"^(.+?)-edit", stem) + if match: + base = match.group(1) + if base not in groups: + groups[base] = {"original": None, "edits": []} + groups[base]["edits"].append(img) + else: + # This is an original + if stem not in groups: + groups[stem] = {"original": None, "edits": []} + groups[stem]["original"] = img + + # Filter out groups without originals, convert to list + result = [] + for base, group in groups.items(): + if group["original"]: + result.append(ImageGroup(base=base, original=group["original"], edits=sorted(group["edits"]))) + + return sorted(result, key=lambda g: g["original"].name) + + +def safe_image_path(filepath: str) -> Path: + """Verify image path is safe (prevent directory traversal).""" + # Search for the image in all_images by filename + for img in all_images: + if img.name == filepath: + return img + raise ValueError(f"Unauthorized path: {filepath}") + + +def hardlink_images(src_paths: list[Path | str], dest_dir: str | Path) -> None: + """Hard link multiple images to destination directory.""" + dest_path = Path(dest_dir) + make_dirs(dest_path) + + for src in src_paths: + src = Path(src) + dest = dest_path / src.name + try: + # Remove existing file if present + if dest.exists(): + dest.unlink() + os.link(src, dest) + except Exception as e: + print(f"Error hard linking {src} to {dest}: {e}", file=sys.stderr) + raise + + +# Global state +app = bottle.Bottle() +all_images = set() +image_groups = [] +current_group_idx = 0 +last_heartbeat = datetime.now() +should_exit = False +picks_dir = Path() +flags = [] +flag_dirs = {} + +HTML_TEMPLATE = """ + + + + + + Image Culler + + + + +
+
+
+ +
+
+
+ +
+
+ + +
+
+ +
+
+ + +
+ + +
+
+ + + +

Keyboard Shortcuts

+
+
+ P + Toggle pick (all edits or hover one) +
+
+ ←→ + Navigate any image +
+
+ K + Skip to next pending +
+
+ Flags: +
+
+ 1-9 + Assign/unassign to flag +
+
+ / + Toggle reject flag +
+
+
+ +
+
+
+ + + + +""" + + +@app.get("/") +def index() -> str: + return HTML_TEMPLATE + + +@app.get("/api/images") +def api_images(): + """Return grouped images and current flag assignments.""" + global image_groups, flags, flag_dirs + result = [] + for group in image_groups: + result.append( + { + "base": group["base"], + "original": {"name": group["original"].name, "path": str(group["original"])}, + "edits": [{"name": e.name, "path": str(e)} for e in group["edits"]], + } + ) + + # Build flag assignments by checking which flag directories contain which original images + flag_assignments = {} + for flag, flag_dir in flag_dirs.items(): + if flag_dir.exists(): + for item in flag_dir.iterdir(): + if item.is_file(): + # Map the file name to its flag + flag_assignments[item.name] = flag + + return {"groups": result, "flags_list": flags, "flag_assignments": flag_assignments} + + +@app.get("/image/") +def serve_image(filename: str): + """Serve an image file safely.""" + try: + target = safe_image_path(filename) + return bottle.static_file(target.name, root=str(target.parent), mimetype="image/jpeg") + except ValueError: + bottle.response.status = 403 + return {"error": "Unauthorized path"} + + +@app.post("/api/pick") +def api_pick(): + """Hard link picked images to _picks directory.""" + global image_groups, current_group_idx, picks_dir + + # Create directory before writing + make_dirs(picks_dir) + + data = bottle.request.json + original_name = data.get("original") + edits_names = data.get("edits", []) + include_original = data.get("includeOriginal", False) + + files_to_pick = [] + + # Add original if requested + if include_original: + for img in all_images: + if img.name == original_name: + files_to_pick.append(img) + break + + # Add edits by finding them in all_images + for edit_name in edits_names: + for img in all_images: + if img.name == edit_name: + files_to_pick.append(img) + break + + try: + hardlink_images(files_to_pick, picks_dir) + return {"status": "ok"} + except Exception as e: + bottle.response.status = 500 + return {"error": str(e)} + + +@app.post("/api/unpick") +def api_unpick(): + """Remove picked images from _picks directory.""" + global picks_dir + + data = bottle.request.json + original_name = data.get("original") + + try: + # Find and remove the original and any edits from picks directory + if picks_dir.exists(): + for item in picks_dir.iterdir(): + if item.is_file() and (item.name == original_name or item.stem.startswith(original_name.rsplit(".", 1)[0] + "-edit")): + item.unlink() + return {"status": "ok"} + except Exception as e: + bottle.response.status = 500 + return {"error": str(e)} + + +@app.post("/api/assign-flag") +def api_assign_flag(): + """Assign an image to a flag.""" + global flag_dirs, all_images + + data = bottle.request.json + original_name = data.get("original") + flag_name = data.get("flag") + + # Find original file + original_path = None + for img in all_images: + if img.name == original_name: + original_path = img + break + + if not original_path or flag_name not in flag_dirs: + bottle.response.status = 400 + return {"error": "Invalid image or flag"} + + # Create directory before writing + make_dirs(flag_dirs[flag_name]) + + try: + hardlink_images([original_path], flag_dirs[flag_name]) + return {"status": "ok"} + except Exception as e: + bottle.response.status = 500 + return {"error": str(e)} + + +@app.post("/api/unassign-flag") +def api_unassign_flag(): + """Remove an image from a flag.""" + global flag_dirs + + data = bottle.request.json + original_name = data.get("original") + flag_name = data.get("flag") + + if flag_name not in flag_dirs: + bottle.response.status = 400 + return {"error": "Invalid flag"} + + try: + target_file = flag_dirs[flag_name] / original_name + if target_file.exists(): + target_file.unlink() + return {"status": "ok"} + except Exception as e: + bottle.response.status = 500 + return {"error": str(e)} + + +@app.post("/api/heartbeat") +def api_heartbeat(): + """Update last heartbeat timestamp.""" + global last_heartbeat + last_heartbeat = datetime.now() + return {"status": "ok"} + + +@app.post("/api/shutdown") +def api_shutdown(): + """Shutdown the server.""" + global should_exit + should_exit = True + return {"status": "ok"} + + +@functools.cache +def make_dirs(path: Path) -> None: + """Create directory if it doesn't exist.""" + path.mkdir(parents=True, exist_ok=True) + + +def heartbeat_monitor() -> None: + """Monitor for shutdown signal.""" + global should_exit + + while not should_exit: + time.sleep(0.5) + + print("\nShutting down server.", file=sys.stderr) + os._exit(0) + + +def main() -> None: + global all_images, image_groups, last_heartbeat, picks_dir, flags, flag_dirs + + parser = argparse.ArgumentParser(description="Image culling app for AI image2image transforms") + parser.add_argument("paths", nargs="+", help="Image files or folders to cull") + parser.add_argument("--picks-dir", default=None, help="Directory for picked images (default: $cwd/_picks)") + parser.add_argument("--rejects-dir", default=None, help="Directory for rejected images (default: $cwd/_rejects)") + parser.add_argument("--flags", default=None, help="Comma-separated flag names (e.g., 'a,b,c,d')") + parser.add_argument("--flags-dir", default=None, help="Directory for flag subdirectories (default: $cwd)") + + args = parser.parse_args() + + # Set directories + cwd = Path.cwd() + picks_dir = Path(args.picks_dir or (cwd / "_picks")) + rejects_dir = Path(args.rejects_dir or (cwd / "_rejects")) + flags_dir = Path(args.flags_dir or cwd) + + # Set up flags + if args.flags: + flags = [f.strip() for f in args.flags.split(",")] + for flag in flags: + flag_dir = flags_dir / f"_picks_flag_{flag}" + flag_dirs[flag] = flag_dir + + # Always add reject as a flag at the end + flags.append("reject") + flag_dirs["reject"] = rejects_dir + + # Discover images + images = discover_images(args.paths) + all_images = set(images) + + if not images: + print("No images found in provided paths.", file=sys.stderr) + sys.exit(1) + + # Group images + image_groups = group_images(images) + + if not image_groups: + print("No original images found (looking for files matching naming convention).", file=sys.stderr) + sys.exit(1) + + print(f"Found {len(image_groups)} original image(s) with variations", file=sys.stderr) + + # Find available port + port = find_available_port() + url = f"http://127.0.0.1:{port}" + + # Start heartbeat monitor + monitor_thread = threading.Thread(target=heartbeat_monitor, daemon=True) + monitor_thread.start() + + # Open browser + print(f"Opening browser at {url}", file=sys.stderr) + webbrowser.open(url) + + # Start Flask server + last_heartbeat = datetime.now() + bottle.run(app, host="127.0.0.1", port=port, quiet=True) + + +if __name__ == "__main__": + main() diff --git a/date_images.py b/date_images.py new file mode 100755 index 0000000..c7e34c3 --- /dev/null +++ b/date_images.py @@ -0,0 +1,79 @@ +#!/usr/bin/env -S uv run +# /// script +# dependencies = ["piexif"] +# /// + +import argparse +from concurrent.futures import ThreadPoolExecutor +import logging +import re +import subprocess +from datetime import datetime +from pathlib import Path + +import piexif + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Set creation time to date in filename and mod time to now for image files.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "paths", + nargs="+", + type=Path, + help="Paths to image files or folders containing images.", + ) + return parser.parse_args() + + +def update_image_times(path: Path) -> None: + match = re.search(r"\D(\d{4}-\d{2}-\d{2})\D?", path.name) + if not match: + logging.debug(f"No date found in {path.name}") + return + date_str = match.group(1) + try: + dt = datetime.fromisoformat(date_str) + formatted_date = dt.strftime("%m/%d/%Y %H:%M:%S") + today = datetime.now().strftime("%m/%d/%Y %H:%M:%S") + subprocess.run(["SetFile", "-d", formatted_date, str(path)], check=True) + subprocess.run(["SetFile", "-m", today, str(path)], check=True) + # Update EXIF DateTimeOriginal + exif_dict = piexif.load(str(path)) + exif_dict["Exif"][piexif.ExifIFD.DateTimeOriginal] = dt.strftime("%Y:%m:%d %H:%M:%S").encode() + exif_bytes = piexif.dump(exif_dict) + piexif.insert(exif_bytes, str(path)) + logging.info(f"Updated times and EXIF for {path}") + except ValueError as e: + logging.error(f"Invalid date {date_str} in {path.name}: {e}") + except subprocess.CalledProcessError as e: + logging.error(f"Failed to update {path}: {e}") + except Exception as e: + logging.error(f"Failed to update EXIF for {path}: {e}") + + +def main() -> None: + args = parse_args() + + files = [] + for it in args.paths: + if it.is_dir(): + files.extend(it.glob("*.jpg")) + files.extend(it.glob("*.jpeg")) + elif it.is_file() and it.suffix.lower() in {".jpg", ".jpeg"}: + files.append(it) + else: + logging.warning(f"Path {it} is neither an image nor a directory, skipping.") + + pool = ThreadPoolExecutor(max_workers=5) + for path in files: + pool.submit(update_image_times, path) + pool.shutdown(wait=True) + + +if __name__ == "__main__": + main() diff --git a/delete_downloaded_videos.py b/delete_downloaded_videos.py new file mode 100755 index 0000000..14683d6 --- /dev/null +++ b/delete_downloaded_videos.py @@ -0,0 +1,142 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["psycopg[binary]"] +# /// +import dataclasses +import datetime +import hashlib +from pathlib import Path +import subprocess +import typing +import psycopg +import contextlib + + +@contextlib.contextmanager +def connect_db() -> typing.Generator[psycopg.Connection, typing.Any, typing.Any]: + with psycopg.connect("postgres://abdus:abdus@db.abdus.dev:5444/smut?sslmode=disable") as conn: + conn.row_factory = psycopg.rows.dict_row + yield conn + + +def get_video_duration(video_path: Path) -> datetime.timedelta: + # fmt: off + args = [ + 'ffprobe', + '-v', 'error', + '-show_entries', 'format=duration', + '-of', 'default=noprint_wrappers=1:nokey=1', + video_path, + ] + # fmt: on + p = subprocess.run(args, stdout=subprocess.PIPE, check=True) + return datetime.timedelta(seconds=round(float(p.stdout), 1)) + + +def hash_partial(f: Path) -> str: + sha1 = hashlib.sha1() + chunk_size = 1024 * 1024 * 10 # 10MB chunk size + + total_read = 0 + with f.open("rb") as file: + while chunk := file.read(chunk_size): + total_read += len(chunk) + sha1.update(chunk) + break # Only reads the first 10MB + + return f"sha1:{total_read}:{sha1.hexdigest()}" + + +@dataclasses.dataclass +class SavedVideo: + id: int + remote_path: Path + local_path: Path + + def __hash__(self): + return hash(self.remote_path) + + +# smut=# \d videos; +# Table "public.videos" +# Column | Type | Collation | Nullable | Default +# ------------------------+-----------------------------+-----------+----------+-------------------------------------------------------------------------- +# id | integer | | not null | nextval('videos_id_seq'::regclass) +# category | ltree | | not null | +# file_path | text | | not null | +# created_at | timestamp without time zone | | | now() +# hash_partial | text | | not null | +# ffprobe | jsonb | | not null | +# size_bytes | bigint | | not null | generated always as ((ffprobe ->> 'size_bytes'::text)::bigint) stored +# duration_sec | numeric | | not null | generated always as ((ffprobe ->> 'duration_sec'::text)::numeric) stored +# suggested_filename | text | | | +# marked_for_deletion_at | timestamp with time zone | | | +# duration_human | text | | | generated always as (ffprobe ->> 'duration_human'::text) stored +# last_seen_at | timestamp with time zone | | | +# Indexes: +# "videos_pkey" PRIMARY KEY, btree (id) +# "videos_file_path_idx" gist (file_path gist_trgm_ops) +# "videos_file_path_uniq" UNIQUE CONSTRAINT, btree (file_path) +def find_videos_by_hash(conn: psycopg.Connection, video_paths: list[Path]) -> list[SavedVideo]: + file_to_hash = {p: hash_partial(p) for p in video_paths} + if not file_to_hash: + return [] + + placeholders = ",".join(["%s"] * len(file_to_hash)) + sql = f""" + SELECT id, file_path AS remote_path, file_path AS local_path + FROM videos + WHERE hash_partial IN ({placeholders}) + """ + with conn.cursor() as cur: + rows = cur.execute(sql, list(file_to_hash.values())).fetchall() + return [SavedVideo(**row) for row in rows] + + +def find_videos_by_duration(conn: psycopg.Connection, video_paths: list[Path]) -> list[SavedVideo]: + if not video_paths: + return [] + file_to_duration = {p: get_video_duration(p).total_seconds() for p in video_paths} + file_to_size = {p: p.stat().st_size for p in video_paths} + + sql = f""" + SELECT id, file_path AS remote_path, file_path AS local_path + FROM videos + WHERE abs(duration_sec - %s) < 0.1 AND abs(size_bytes - %s) < 1048576 + """ + with conn.cursor() as cur: + out = [] + for file_path in video_paths: + size = file_to_size[file_path] + duration = file_to_duration[file_path] + for row in cur.execute(sql, (duration, size)): + out.append(SavedVideo(**row)) + return out + + +def delete_videos(con: psycopg.Connection, file_paths: list[str]) -> bool: + with con.cursor() as cur: + placeholders = ",".join(["%s"] * len(file_paths)) + sql = f"DELETE FROM videos WHERE file_path IN ({placeholders})" + cur.execute(sql, file_paths) + + +def main(): + videos = list(Path(r"/Users/abdus/Downloads/temp").glob("*.mp4")) + with connect_db() as conn: + found_videos = set(find_videos_by_hash(conn, videos)) + found_videos.update(find_videos_by_duration(conn, videos)) + found_videos = {video for video in found_videos if "/mnt/box/files/_raw/prt" in video.remote_path} + for video in found_videos: + print(f"{video.remote_path}") + if not found_videos: + print("No videos found for deletion.") + return + input("Press Enter to continue...") + + print(f"Deleting {len(found_videos)} videos from database...") + delete_videos(conn, [video.remote_path for video in found_videos]) + + +if __name__ == "__main__": + main() diff --git a/face_landmarker.task b/face_landmarker.task new file mode 100644 index 0000000..c50c845 Binary files /dev/null and b/face_landmarker.task differ diff --git a/ffmpeg-thumbnail-tile.md b/ffmpeg-thumbnail-tile.md new file mode 100644 index 0000000..b07c202 --- /dev/null +++ b/ffmpeg-thumbnail-tile.md @@ -0,0 +1,29 @@ +# Generating thumbnail tiles using ffmpeg +> ffmpeg can unsurprisingly build contact sheets, too. Is there anything ffmpeg can't do? + +I've been organizing my movie collection lately, and I needed a way to figure out what a video is all about without having to +view & seek it. +Building a contact sheet is a simple solution for this. It places frames with from the video in a grid separated by +an interval (like every minute) and you can quickly see what a video contains. + +## ffmpeg tile filter + +With anything media-related, I checked to see if ffmpeg supports this: it turns out, it has a `tile` filter[^tile][tile], which gives +us some primitives to work with. + +## Covering the whole video + +```shell +ffmpeg -skip_frame nokey -i video.mp4 -vf 'scale=320:-1,tile=8x8' -an -vsync 0 keyframes%03d.png +``` + +This creates n files, each containing 8x8 grid of keyframes from the video. +It's close, but I need a single file with tiles spanning the video from start to finish. + +## + +```shell + +``` + +[tile]: https://ffmpeg.org/ffmpeg-filters.html#tile-1 diff --git a/flux2_private_server.py b/flux2_private_server.py new file mode 100644 index 0000000..c93aa10 --- /dev/null +++ b/flux2_private_server.py @@ -0,0 +1,461 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.10,<3.11" +# dependencies = [ +# "torch==2.9.1+cu128", +# "torchvision==0.24.1+cu128", +# "torchaudio==2.9.1+cu128", +# "diffusers @ git+https://github.com/huggingface/diffusers.git", +# "transformers>=4.57.0", +# "accelerate>=1.0.0", +# "safetensors>=0.4.0", +# "huggingface_hub>=0.25.0", +# "numpy>=1.26.0", +# "pillow>=10.0.0", +# "fastapi>=0.115.0", +# "uvicorn>=0.29.0", +# ] +# [tool.uv] +# extra-index-url = ["https://download.pytorch.org/whl/cu128"] +# /// + +import argparse +import base64 +import binascii +import io +import os +import random +import threading +import time +from typing import Any + +os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1") + +import torch +from diffusers import Flux2KleinPipeline +from fastapi import FastAPI, HTTPException +from fastapi.responses import HTMLResponse, Response +from pydantic import BaseModel, Field +from PIL import Image + + +def _resolve_model_size() -> str: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--model", choices=["4b", "9b"], default="9b") + args, _ = parser.parse_known_args() + return args.model + + +def _checkpoint_for(model_size: str) -> str: + if model_size == "4b": + return "black-forest-labs/FLUX.2-klein-4B" + return "black-forest-labs/FLUX.2-klein-9B" + + +MODEL_SIZE = _resolve_model_size() +CHECKPOINT = _checkpoint_for(MODEL_SIZE) +USE_COMPILE = os.environ.get("FLUX2_COMPILE", "1") != "0" +WARMUP_RUNS = int(os.environ.get("FLUX2_WARMUP_RUNS", "2")) +WARMUP_STEPS = int(os.environ.get("FLUX2_WARMUP_STEPS", "4")) +WARMUP_PROMPT = "A high-resolution photo of a red fox in natural daylight" +MAX_SEED = (2**31) - 1 +MAX_PIXELS = 4 * 1024 * 1024 + +DEFAULT_STEPS = 4 +DEFAULT_GUIDANCE = 1.0 + +app = FastAPI(title="FLUX2 Klein Private Server") +PIPE: Flux2KleinPipeline | None = None +PIPE_LOCK = threading.Lock() + + +class GenerateRequest(BaseModel): + prompt: str = Field(..., min_length=1) + width: int = 1024 + height: int = 1024 + num_inference_steps: int = DEFAULT_STEPS + guidance_scale: float = DEFAULT_GUIDANCE + seed: int = 0 + randomize_seed: bool = False + image_data_url: str | None = None + + +class CompatRequest(BaseModel): + model: str | None = None + prompt: str = Field(..., min_length=1) + imageDataUrl: str | None = None + size: str | None = None + width: int | None = None + height: int | None = None + num_inference_steps: int | None = None + guidance_scale: float | None = None + seed: int | None = None + randomize_seed: bool = False + response_format: str = "b64_json" + n: int = 1 + negative_prompt: str | None = None + negativePrompt: str | None = None + + +def _ensure_cuda() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this server") + + +def _compile_pipe(pipe: Flux2KleinPipeline) -> bool: + if not USE_COMPILE: + return False + print("Compiling transformer and VAE decoder...") + start = time.perf_counter() + pipe.transformer = torch.compile( + pipe.transformer, + mode="max-autotune", + dynamic=False, + fullgraph=True, + ) + pipe.vae.decode = torch.compile( + pipe.vae.decode, + mode="max-autotune", + dynamic=False, + fullgraph=True, + ) + print(f"Compile finished in {time.perf_counter() - start:.2f}s") + return True + + +def _warmup_pipe(pipe: Flux2KleinPipeline) -> None: + if WARMUP_RUNS <= 0: + return + print(f"Warmup started ({WARMUP_RUNS} runs)...") + start = time.perf_counter() + for run_idx in range(WARMUP_RUNS): + generator = torch.Generator(device="cuda").manual_seed(run_idx) + pipe( + prompt=WARMUP_PROMPT, + height=1024, + width=1024, + num_inference_steps=WARMUP_STEPS, + guidance_scale=1.0, + generator=generator, + ).images[0] + torch.cuda.synchronize() + print(f"Warmup finished in {time.perf_counter() - start:.2f}s") + + +def _load_pipe(ckpt: str) -> Flux2KleinPipeline: + print(f"Loading checkpoint: {ckpt}") + pipe = Flux2KleinPipeline.from_pretrained(ckpt, torch_dtype=torch.bfloat16) + pipe = pipe.to("cuda") + pipe.transformer.fuse_qkv_projections() + pipe.vae.fuse_qkv_projections() + pipe.vae.to(memory_format=torch.channels_last) + pipe.transformer.set_attention_backend("_native_flash") + if _compile_pipe(pipe): + _warmup_pipe(pipe) + return pipe + + +def _normalize_dims(width: int, height: int) -> tuple[int, int]: + if width <= 0 or height <= 0: + raise HTTPException(status_code=422, detail="width and height must be > 0") + if width % 16 != 0 or height % 16 != 0: + raise HTTPException( + status_code=422, detail="width and height must be multiples of 16" + ) + if width * height > MAX_PIXELS: + raise HTTPException(status_code=422, detail="max resolution is 4 megapixels") + return width, height + + +def _parse_size( + size: str | None, width: int | None, height: int | None +) -> tuple[int, int]: + if width and height: + return _normalize_dims(width, height) + if not size: + return _normalize_dims(1024, 1024) + parts = size.lower().split("x") + if len(parts) != 2: + raise HTTPException(status_code=422, detail="size must look like 1024x1024") + try: + parsed_w = int(parts[0]) + parsed_h = int(parts[1]) + except ValueError as exc: + raise HTTPException( + status_code=422, detail="size values must be integers" + ) from exc + return _normalize_dims(parsed_w, parsed_h) + + +def _decode_image_data_url(image_data_url: str | None) -> Image.Image | None: + if not image_data_url: + return None + if "," not in image_data_url: + raise HTTPException(status_code=422, detail="imageDataUrl must be a data URL") + _, payload = image_data_url.split(",", 1) + try: + raw = base64.b64decode(payload) + except (ValueError, binascii.Error) as exc: + raise HTTPException( + status_code=422, detail="invalid base64 in imageDataUrl" + ) from exc + with Image.open(io.BytesIO(raw)) as image: + return image.convert("RGB") + + +def _generate_png_bytes( + prompt: str, + width: int, + height: int, + num_inference_steps: int, + guidance_scale: float, + seed: int, + image: Image.Image | None, +) -> tuple[bytes, int]: + if PIPE is None: + raise HTTPException(status_code=503, detail="model not loaded yet") + + generator = torch.Generator(device="cuda").manual_seed(seed) + kwargs: dict[str, Any] = { + "prompt": prompt, + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "generator": generator, + } + + with PIPE_LOCK: + try: + if image is not None: + try: + output = PIPE(image=image, **kwargs) + except TypeError: + output = PIPE(images=[image], **kwargs) + else: + output = PIPE(**kwargs) + except Exception as exc: + raise HTTPException( + status_code=500, detail=f"generation failed: {exc}" + ) from exc + + result_image = output.images[0] + buf = io.BytesIO() + result_image.save(buf, format="PNG") + return buf.getvalue(), seed + + +@app.on_event("startup") +def _startup() -> None: + global PIPE + _ensure_cuda() + PIPE = _load_pipe(CHECKPOINT) + print("Server ready") + + +@app.get("/", response_class=HTMLResponse) +def index() -> str: + return """ + + + + + FLUX.2 Klein Playground + + + +
+
+

FLUX.2 Klein Playground

+

Private single-user UI at /. Distilled defaults: 4 steps, guidance 1.0.

+
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+ +

Idle

+
+
+
+ +
+

Output

+ \"Generated +

Tip: right-click image to save.

+
+
+ + + + +""" + + +@app.get("/health") +def health() -> dict[str, str]: + return {"status": "ok", "model": CHECKPOINT} + + +@app.post("/generate") +def generate(req: GenerateRequest) -> Response: + width, height = _normalize_dims(req.width, req.height) + seed = req.seed if not req.randomize_seed else random.randint(0, MAX_SEED) + image = _decode_image_data_url(req.image_data_url) + png_bytes, used_seed = _generate_png_bytes( + prompt=req.prompt, + width=width, + height=height, + num_inference_steps=req.num_inference_steps, + guidance_scale=req.guidance_scale, + seed=seed, + image=image, + ) + return Response( + content=png_bytes, + media_type="image/png", + headers={"X-Seed": str(used_seed)}, + ) + + +@app.post("/v1/images/generations") +def compat_generate(req: CompatRequest) -> dict[str, Any]: + width, height = _parse_size(req.size, req.width, req.height) + seed_value = req.seed if req.seed is not None else 0 + seed = seed_value if not req.randomize_seed else random.randint(0, MAX_SEED) + image = _decode_image_data_url(req.imageDataUrl) + + png_bytes, used_seed = _generate_png_bytes( + prompt=req.prompt, + width=width, + height=height, + num_inference_steps=req.num_inference_steps or DEFAULT_STEPS, + guidance_scale=req.guidance_scale or DEFAULT_GUIDANCE, + seed=seed, + image=image, + ) + b64_png = base64.b64encode(png_bytes).decode("ascii") + return { + "created": int(time.time()), + "model": CHECKPOINT, + "seed": used_seed, + "data": [{"b64_json": b64_png}], + } + + +if __name__ == "__main__": + import uvicorn + + parser = argparse.ArgumentParser(description="Private FLUX.2 Klein FastAPI server") + parser.add_argument("--model", choices=["4b", "9b"], default=MODEL_SIZE) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=6006) + args = parser.parse_args() + + MODEL_SIZE = args.model + CHECKPOINT = _checkpoint_for(MODEL_SIZE) + uvicorn.run(app, host=args.host, port=args.port) diff --git a/flux_klein_prompt.txt b/flux_klein_prompt.txt new file mode 100644 index 0000000..93e5997 --- /dev/null +++ b/flux_klein_prompt.txt @@ -0,0 +1,30 @@ +Artistic photorealistic conversion. +Subtle chiaroscuro lighting, not too dark. +120mm telephoto lens, f/2.8, shallow depth of field. +Soft-focus highlights, atmospheric bloom. +Kodak Ektachrome E100 color palette. +Preserve the original lighting and moody lighting. +High dynamic range with a focus on rich textures. +# subtle dramatic lighting with rich lights and shadows, but not too dark. +balanced exposure + +# looking at the camera +# looking away from the camera + +she has flawless, spotless, tight skin. never change the skin color. +# porcelain gothic pale skin +raised cheekbones, tapered face, seductive look +make her prettier without changing the skin color +age them by 5 years and make them look like 25 year old adults. +preserve height and proportion while maintaining the same facial features. +# make her head 10% slimmer. + +depict cartoon characters as japanese + +slightly parted lips. +slightly lowered or closed eyelids. +# shiny outfit unless naked. +# shiny, glossy outfit. +keep the heavy makeup. + +make it sharp throughout. remove all text. diff --git a/grammar.html b/grammar.html new file mode 100644 index 0000000..19a73d0 --- /dev/null +++ b/grammar.html @@ -0,0 +1,251 @@ + + + + + + Text Corrector App + + + + +
+
+

Original Text

+
+
+
+ +
Loading suggestion...
+ + + + + + + diff --git a/image_dates.py b/image_dates.py new file mode 100644 index 0000000..56562d8 --- /dev/null +++ b/image_dates.py @@ -0,0 +1,279 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["Pillow"] +# /// + +import argparse +import collections +import datetime as dt +import logging +import os +import re +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +from PIL import Image, UnidentifiedImageError + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") + + +DATE_PATTERNS = ( + re.compile(r"^(?P\d{4})-(?P\d{2})-(?P\d{2})"), + re.compile(r"^(?P\d{4})\.(?P\d{2})\.(?P\d{2})"), + re.compile(r"^(?P\d{4})(?P\d{2})(?P\d{2})"), +) + +EXIF_DATETIME_TAGS = (36867, 36868, 306) + + +@dataclass +class ImageEntry: + path: Path + group_key: str + filename_date: dt.datetime | None + exif_date: dt.datetime | None = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Update file timestamps for images grouped by filename pattern", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "directory", + nargs="?", + default=Path.cwd(), + type=Path, + help="Root directory to scan for image files", + ) + parser.add_argument( + "--recursive", + action="store_true", + help="Recurse into subdirectories", + ) + parser.add_argument( + "--extensions", + default=".jpg,.jpeg,.png,.tif,.tiff,.heic,.heif", + help="Comma-separated list of extensions to include", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned changes without touching the filesystem", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Enable debug logging", + ) + return parser.parse_args() + + +def configure_logging(verbose: bool) -> None: + if verbose: + logging.getLogger().setLevel(logging.DEBUG) + + +def collect_image_paths(directory: Path, extensions: set[str], recursive: bool) -> Iterable[Path]: + iterator = directory.rglob("*") if recursive else directory.glob("*") + for candidate in iterator: + if candidate.is_file() and candidate.suffix.lower() in extensions: + yield candidate + + +def parse_group_and_date(path: Path) -> tuple[str, dt.datetime | None] | None: + name = path.stem + if " -- " not in name: + logging.debug("Skipping %s: missing group delimiter", path) + return None + try: + group_part, rest = name.rsplit(" -- ", 1) + except ValueError: + logging.debug("Skipping %s: unable to split group and tail", path) + return None + if not group_part: + logging.debug("Skipping %s: empty group part", path) + return None + date_candidate = extract_date_from_tail(rest) + if " -- " not in group_part and "__" in rest: + album_part, _, _ = rest.partition("__") + album_part = album_part.strip() + if album_part: + group_part = f"{group_part} -- {album_part}" + return group_part, date_candidate + + +def extract_date_from_tail(text: str) -> dt.datetime | None: + for pattern in DATE_PATTERNS: + match = pattern.match(text) + if match: + try: + return dt.datetime( + year=int(match.group("year")), + month=int(match.group("month")), + day=int(match.group("day")), + ) + except ValueError: + return None + return None + + +def read_exif_datetime(path: Path) -> dt.datetime | None: + try: + with Image.open(path) as image: + exif = image.getexif() + except (UnidentifiedImageError, OSError) as error: + logging.debug("Could not read EXIF from %s: %s", path, error) + return None + if not exif: + return None + for tag in EXIF_DATETIME_TAGS: + value = exif.get(tag) + if not value: + continue + if isinstance(value, bytes): + value = value.decode(errors="ignore") + if not isinstance(value, str): + continue + value = value.strip() + for fmt in ("%Y:%m:%d %H:%M:%S", "%Y-%m-%d %H:%M:%S"): + try: + return dt.datetime.strptime(value, fmt) + except ValueError: + continue + return None + + +def most_common_datetime(candidates: list[dt.datetime]) -> dt.datetime | None: + if not candidates: + return None + counter = collections.Counter(candidates) + most_common = counter.most_common() + top_count = most_common[0][1] + top_values = [item for item, count in most_common if count == top_count] + return min(top_values) + + +def determine_group_date(entries: list[ImageEntry]) -> dt.datetime | None: + name_dates = [entry.filename_date for entry in entries if entry.filename_date] + if name_dates: + logging.debug("Using filename-derived date for group %s", entries[0].group_key) + return most_common_datetime(name_dates) + exif_dates: list[dt.datetime] = [] + for entry in entries: + if entry.exif_date is None: + entry.exif_date = read_exif_datetime(path=entry.path) + if entry.exif_date: + exif_dates.append(entry.exif_date) + if exif_dates: + logging.debug("Using EXIF-derived date for group %s", entries[0].group_key) + return most_common_datetime(exif_dates) + + +def apply_timestamp(path: Path, target: dt.datetime, dry_run: bool) -> None: + timestamp = target.timestamp() + if dry_run: + logging.debug("DRY-RUN %s -> %s", path, target.isoformat(sep=" ")) + return + os.utime(path, times=(timestamp, timestamp)) + + +def chunked(items: list[str], size: int) -> Iterable[list[str]]: + for index in range(0, len(items), size): + yield items[index : index + size] + + +def apply_setfile_batch(entries: list[ImageEntry], target: dt.datetime, dry_run: bool) -> None: + if sys.platform != "darwin": + return + setfile = shutil.which("SetFile") + if not setfile: + return + formatted = target.strftime("%m/%d/%Y %H:%M:%S") + paths = [str(entry.path) for entry in entries] + for batch in chunked(items=paths, size=64): + if dry_run: + logging.debug("DRY-RUN SetFile %s files -> %s", len(batch), formatted) + continue + try: + subprocess.run([setfile, "-d", formatted, *batch], check=True) + subprocess.run([setfile, "-m", formatted, *batch], check=True) + except subprocess.CalledProcessError as error: + logging.debug("SetFile failed for %s files: %s", len(batch), error) + + +def update_group(entries: list[ImageEntry], dry_run: bool) -> bool: + target = determine_group_date(entries=entries) + if target is None: + logging.warning("No date found for group %s", entries[0].group_key) + return False + for entry in entries: + apply_timestamp(path=entry.path, target=target, dry_run=dry_run) + apply_setfile_batch(entries=entries, target=target, dry_run=dry_run) + action = "DRY-RUN" if dry_run else "UPDATED" + logging.info( + "%s %s (%s files) -> %s", + action, + entries[0].group_key, + len(entries), + target.isoformat(sep=" "), + ) + return True + + +def process_directory(directory: Path, extensions: set[str], recursive: bool, dry_run: bool) -> None: + groups: dict[str, list[ImageEntry]] = collections.defaultdict(list) + for path in collect_image_paths(directory=directory, extensions=extensions, recursive=recursive): + parsed = parse_group_and_date(path=path) + if parsed is None: + continue + group_key, filename_date = parsed + groups[group_key].append(ImageEntry(path=path, group_key=group_key, filename_date=filename_date)) + updated = 0 + skipped = 0 + for entries in groups.values(): + if update_group(entries=entries, dry_run=dry_run): + updated += len(entries) + else: + skipped += len(entries) + logging.info("Updated %s files; skipped %s files", updated, skipped) + + +def normalize_extensions(raw: str) -> set[str]: + pieces = re.split(r"[;,]", raw) + results: set[str] = set() + for piece in pieces: + trimmed = piece.strip().lower() + if not trimmed: + continue + if not trimmed.startswith("."): + trimmed = f".{trimmed}" + results.add(trimmed) + return results + + +def main() -> None: + args = parse_args() + configure_logging(verbose=args.verbose) + extensions = normalize_extensions(raw=args.extensions) + if not extensions: + logging.error("No valid extensions provided") + raise SystemExit(1) + if not args.directory.exists(): + logging.error("Directory %s does not exist", args.directory) + raise SystemExit(1) + process_directory( + directory=args.directory, + extensions=extensions, + recursive=args.recursive, + dry_run=args.dry_run, + ) + + +if __name__ == "__main__": + main() diff --git a/immich_dedupe_albums.py b/immich_dedupe_albums.py new file mode 100644 index 0000000..a347d98 --- /dev/null +++ b/immich_dedupe_albums.py @@ -0,0 +1,105 @@ +import httpx +import os +from datetime import datetime +from collections import defaultdict +from typing import List, Dict, Any + +# Get environment variables +IMMICH_API_KEY = os.environ.get("IMMICH_API_KEY") +IMMICH_SERVER_URL = os.environ.get("IMMICH_SERVER_URL") + +if not IMMICH_API_KEY or not IMMICH_SERVER_URL: + raise ValueError("IMMICH_API_KEY and IMMICH_SERVER_URL environment variables must be set") + +# Setup HTTP client with auth headers +headers = {"Content-Type": "application/json", "Accept": "application/json", "x-api-key": IMMICH_API_KEY} + +# Global HTTP client +client = httpx.Client(headers=headers) + + +def get_albums() -> List[Dict[str, Any]]: + """Fetch all albums from Immich API""" + response = client.get(f"{IMMICH_SERVER_URL}/api/albums") + response.raise_for_status() + return response.json() + + +def get_album_assets(album_id: str) -> List[str]: + """Fetch all asset IDs from a specific album""" + response = client.get(f"{IMMICH_SERVER_URL}/api/albums/{album_id}?withoutAssets=false") + response.raise_for_status() + album_data = response.json() + return [asset["id"] for asset in album_data.get("assets", [])] + + +def add_assets_to_album(album_id: str, asset_ids: List[str]) -> None: + """Add assets to an album""" + if not asset_ids: + return + + response = client.put(f"{IMMICH_SERVER_URL}/api/albums/{album_id}/assets", json={"ids": asset_ids}) + response.raise_for_status() + + +def delete_album(album_id: str) -> None: + """Delete an album""" + response = client.delete(f"{IMMICH_SERVER_URL}/api/albums/{album_id}") + response.raise_for_status() + + +def deduplicate_albums(): + """Main function to deduplicate albums""" + print("Fetching albums...") + albums = get_albums() + + # Group albums by name (case-insensitive) + album_groups = defaultdict(list) + for album in albums: + album_name = album["albumName"].lower().strip() + album_groups[album_name].append(album) + + # Process groups with duplicates + for album_name, album_list in album_groups.items(): + if len(album_list) <= 1: + continue # Skip groups with only one album + + print(f"\nProcessing duplicate albums for: '{album_name}' ({len(album_list)} albums)") + + # Sort by creation date to find the oldest + album_list.sort(key=lambda x: datetime.fromisoformat(x["createdAt"].replace("Z", "+00:00"))) + oldest_album = album_list[0] + duplicate_albums = album_list[1:] + + print(f" Oldest album: {oldest_album['id']} (created: {oldest_album['createdAt']})") + + # Collect assets from duplicate albums + all_asset_ids = [] + for dup_album in duplicate_albums: + print(f" Processing duplicate: {dup_album['id']} (created: {dup_album['createdAt']})") + asset_ids = get_album_assets(dup_album["id"]) + all_asset_ids.extend(asset_ids) + print(f" Found {len(asset_ids)} assets") + + # Add assets to oldest album + if all_asset_ids: + print(f" Adding {len(all_asset_ids)} assets to oldest album...") + add_assets_to_album(oldest_album["id"], all_asset_ids) + + # Delete duplicate albums + for dup_album in duplicate_albums: + print(f" Deleting duplicate album: {dup_album['id']}") + delete_album(dup_album["id"]) + + print(f" ✓ Merged {len(duplicate_albums)} duplicate albums into {oldest_album['id']}") + + +if __name__ == "__main__": + try: + deduplicate_albums() + print("\n✓ Album deduplication completed successfully!") + except Exception as e: + print(f"\n✗ Error: {e}") + raise + finally: + client.close() diff --git a/install_backup_service.py b/install_backup_service.py new file mode 100644 index 0000000..cf8c9a0 --- /dev/null +++ b/install_backup_service.py @@ -0,0 +1,78 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = [] +# /// + +import argparse +import logging +import subprocess + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Install cron job") + parser.add_argument("--name", required=True, help="Name for the cron job") + parser.add_argument("command", nargs=argparse.REMAINDER, help="Command to run") + return parser.parse_args() + + +def get_current_crontab() -> str: + try: + result = subprocess.run(["crontab", "-l"], capture_output=True, text=True, check=True) + return result.stdout + except subprocess.CalledProcessError: + return "" + + +def create_cron_entry(name: str, command: list) -> str: + command_str = " ".join(command) + return f"0 0 * * * {command_str}" + + +def job_exists(name: str, current_crontab: str) -> bool: + comment = f"# {name}" + return comment in current_crontab + + +def install_cron_job(name: str, command: list) -> None: + current_crontab = get_current_crontab() + + if job_exists(name, current_crontab): + logger.info(f"Cron job '{name}' already exists") + return + + cron_entry = create_cron_entry(name, command) + job_comment = f"# {name}" + new_crontab = current_crontab.rstrip() + + if new_crontab and not new_crontab.endswith("\n"): + new_crontab += "\n" + + new_crontab += f"{job_comment}\n{cron_entry}\n" + + process = subprocess.Popen(["crontab", "-"], stdin=subprocess.PIPE, text=True) + process.communicate(input=new_crontab) + + if process.returncode != 0: + raise RuntimeError("Failed to install cron job") + + logger.info(f"Cron job '{name}' installed successfully") + + +def main(): + args = parse_args() + + if not args.command: + raise ValueError("Command is required") + + logger.info(f"Installing cron job '{args.name}' with command: {' '.join(args.command)}") + install_cron_job(args.name, args.command) + + logger.info("Installation completed successfully") + logger.info("Check cron jobs with: crontab -l") + + +if __name__ == "__main__": + exit(main()) diff --git a/prompter/prompter.html b/prompter/prompter.html new file mode 100644 index 0000000..d0e9c9a --- /dev/null +++ b/prompter/prompter.html @@ -0,0 +1,523 @@ + + + + + + Prompter + + + + + + + + +
+ +
+
+
+ + / +
+
+
+ + +
+ + +
+
Prompt
+ + +
+
+ alt+1 + Custom prompt + ✓ saved +
+ + +
+ + +
+
Previously used
+ +
+
+ + +
+
Building blocks
+
+ +
+
+ +
+ + + + diff --git a/prompter/prompter.py b/prompter/prompter.py new file mode 100755 index 0000000..a04d9b3 --- /dev/null +++ b/prompter/prompter.py @@ -0,0 +1,205 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "bottle", +# ] +# /// +"""Browse images one-by-one and assign text prompts via a local web UI. + +Saves results to a JSONL file: {"filename": "...", "prompt": "..."} +Images navigated past without a prompt produce no output line. +""" + +import argparse +import json +import socket +import sys +import threading +import time +import webbrowser +from pathlib import Path + +import bottle + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +CANNED_PROMPTS: list[str] = [ + "photorealistic. skin must be flawless", +] +"""Pre-populated prompt history shown in the UI on first launch.""" + +BUILDING_BLOCKS: list[str] = [ + "make her prettier", + "age her by 5 years and make her look like a 25 year old", + "keep the dark skin", + "wet skin", + "skin covered with baby oil, shiny skin", + "soaking wet skin and hair, wet clothes clinging to her body", + "long, flowy hair", + "japanese", + "greek", + "no extra or missing fingers. each hand must have 5 fingers with the same hand pose as the original image", + # expression + "direct her gaze at the camera", + "give her a serious and seductive look", + "give her a playful expression", + "slightly parted lips", + "french kiss, eyes closed, tongue out", + "closed eyes", + "lower her eyelids and slightly part her lips in a seductive fashion", + "exaggerated expression", + "soft dramatic lighting", + # outfit + "replace the outfit with glossy latex", + "remove the tan lines, she doesn't wear white bra", +] +"""Palette of text snippets the user can append to any prompt.""" + +VALID_EXTENSIONS: frozenset[str] = frozenset({".webp", ".jpg", ".jpeg", ".png"}) + +HTML: str = (Path(__file__).parent / "prompter.html").read_text(encoding="utf-8") + +# --------------------------------------------------------------------------- +# App +# --------------------------------------------------------------------------- + + +def find_free_port() -> int: + """Bind to port 0 and return the OS-assigned port number.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class PrompterApp: + def __init__(self, images: list[str], output_path: Path) -> None: + self.images = images + self.output_path = output_path + self.app = bottle.Bottle() + self.app.route("/")(self.index) + self.app.route("/api/images")(self.api_images) + self.app.route("/image/")(self.serve_image) + self.app.route("/api/save", method="POST")(self.api_save) + + def index(self) -> str: + return HTML + + def api_images(self) -> str: + """Return image list, canned prompts, building blocks, and existing prompts.""" + stem_to_path: dict[str, str] = {Path(p).stem: p for p in self.images} + existing_by_stem = self._load_existing_prompts() + existing_prompts = { + stem_to_path[stem]: prompt + for stem, prompt in existing_by_stem.items() + if stem in stem_to_path + } + bottle.response.content_type = "application/json" + return json.dumps( + { + "images": self.images, + "canned_prompts": CANNED_PROMPTS, + "building_blocks": BUILDING_BLOCKS, + "existing_prompts": existing_prompts, + } + ) + + def serve_image(self, idx: int) -> bottle.HTTPResponse: + if idx < 0 or idx >= len(self.images): + bottle.abort(404, "Image not found") + path = Path(self.images[idx]) + return bottle.static_file(path.name, root=str(path.parent)) + + def api_save(self) -> str: + data: dict = bottle.request.json or {} + file_path: str = data.get("file_path", "") + prompt: str = data.get("prompt", "") + if file_path and prompt: + self._save_entry(file_path, prompt) + bottle.response.content_type = "application/json" + return json.dumps({"ok": True}) + + def _load_existing_prompts(self) -> dict[str, str]: + """Read the output JSONL and return a mapping of stem → prompt. + + When a filename appears multiple times the last entry wins, so re-running + the tool and overwriting a previous prompt works naturally. + """ + if not self.output_path.exists(): + return {} + prompts: dict[str, str] = {} + with open(self.output_path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + if "filename" in entry and "prompt" in entry: + prompts[entry["filename"]] = entry["prompt"] + except json.JSONDecodeError: + pass + return prompts + + def _save_entry(self, file_path: str, prompt: str) -> None: + """Append one {filename, prompt} record to the JSONL output file.""" + with open(self.output_path, "a", encoding="utf-8") as f: + f.write( + json.dumps({"filename": Path(file_path).stem, "prompt": prompt}) + "\n" + ) + + def run(self, port: int) -> None: + bottle.run(self.app, host="localhost", port=port, quiet=True) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Browse images and assign text prompts via a web UI.", + ) + parser.add_argument("images", nargs="+", help="Image paths to browse") + parser.add_argument( + "-o", + "--output-path", + default="prompts.jsonl", + metavar="FILE", + help="Output JSONL file path (default: prompts.jsonl)", + ) + args = parser.parse_args() + + images = [ + str(Path(p).resolve()) + for p in args.images + if Path(p).suffix.lower() in VALID_EXTENSIONS + ] + + if not images: + exts = ", ".join(sorted(VALID_EXTENSIONS)) + print( + f"Error: no valid images found. Expected extensions: {exts}", + file=sys.stderr, + ) + sys.exit(1) + + print(f"Loaded {len(images)} image(s).") + + port = find_free_port() + url = f"http://localhost:{port}" + print(f"Starting server at {url}") + + threading.Thread( + target=lambda: (time.sleep(0.8), webbrowser.open(url)), + daemon=True, + ).start() + + PrompterApp(images, Path(args.output_path)).run(port) + + +if __name__ == "__main__": + main() diff --git a/pullio.py b/pullio.py new file mode 100644 index 0000000..bd41f40 --- /dev/null +++ b/pullio.py @@ -0,0 +1,681 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +import os +import shlex +import shutil +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from urllib import error, request + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") +LOGGER = logging.getLogger("pullio") +LABEL_PREFIX = "org.hotio.pullio" + + +@dataclass(frozen=True) +class Config: + compose_binary: str + docker_binary: str + cache_location: Path + tag: str + parallel: int + compose_type: str + script_hash: str + telegram_bot_token: str + telegram_chat_id: str + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Pull and optionally update Docker Compose containers based on labels.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--tag", default="") + parser.add_argument("--debug", action="store_true") + parser.add_argument("--parallel", type=int, default=1) + args = parser.parse_args() + if args.parallel < 1: + parser.error("--parallel must be >= 1") + return args + + +def run_command( + command: list[str], + *, + check: bool = True, + env: dict[str, str] | None = None, + input_text: str | None = None, +) -> str: + completed = subprocess.run( + command, + check=check, + capture_output=True, + text=True, + env=env, + input=input_text, + ) + return completed.stdout.strip() + + +def detect_compose_type(compose_binary: str, docker_binary: str) -> str: + if compose_binary: + return "V1" + if docker_binary: + try: + run_command([docker_binary, "compose", "version"]) + return "V2" + except subprocess.CalledProcessError: + return "NONE" + return "NONE" + + +def docker_inspect_value(docker_binary: str, name: str, template: str) -> str: + try: + value = run_command([docker_binary, "inspect", f"--format={template}", name]) + except subprocess.CalledProcessError: + return "" + if value == "": + return "" + return value + + +def docker_image_inspect_value(docker_binary: str, image: str, template: str) -> str: + try: + value = run_command( + [docker_binary, "image", "inspect", f"--format={template}", image] + ) + except subprocess.CalledProcessError: + return "" + if value == "": + return "" + return value + + +def compose_pull(config: Config, workdir: str, service: str) -> bool: + if config.compose_type == "V1": + cmd = [config.compose_binary, "pull", service] + elif config.compose_type == "V2": + cmd = [config.docker_binary, "compose", "pull", service] + elif config.docker_binary: + cmd = [ + config.docker_binary, + "run", + "--rm", + "-v", + "/var/run/docker.sock:/var/run/docker.sock", + "-v", + f"{workdir}:{workdir}", + f"-w={workdir}", + "linuxserver/docker-compose", + "pull", + service, + ] + else: + LOGGER.error( + "Neither Docker Compose nor Docker binary is available. Cannot pull." + ) + return False + + try: + run_command(cmd) + return True + except subprocess.CalledProcessError: + return False + + +def compose_up(config: Config, workdir: str, service: str) -> bool: + if config.compose_type == "V1": + cmd = [config.compose_binary, "up", "-d", "--always-recreate-deps", service] + elif config.compose_type == "V2": + cmd = [ + config.docker_binary, + "compose", + "up", + "-d", + "--always-recreate-deps", + service, + ] + elif config.docker_binary: + cmd = [ + config.docker_binary, + "run", + "--rm", + "-v", + "/var/run/docker.sock:/var/run/docker.sock", + "-v", + f"{workdir}:{workdir}", + f"-w={workdir}", + "linuxserver/docker-compose", + "up", + "-d", + "--always-recreate-deps", + service, + ] + else: + LOGGER.error( + "Neither Docker Compose nor Docker binary is available. Cannot bring up services." + ) + return False + + try: + run_command(cmd) + return True + except subprocess.CalledProcessError: + return False + + +def post_json(url: str, payload: dict[str, object]) -> None: + body = json.dumps(payload).encode("utf-8") + req = request.Request( + url=url, + data=body, + headers={ + "User-Agent": "Pullio", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with request.urlopen(req) as response: + response.read() + except (error.HTTPError, error.URLError) as exc: + LOGGER.warning("Webhook request failed: %s", exc) + + +def now_iso_utc() -> str: + return ( + datetime.now(timezone.utc) + .isoformat(timespec="milliseconds") + .replace("+00:00", "Z") + ) + + +def parse_script_command(value: str) -> list[str]: + return shlex.split(value) if value else [] + + +def prepare_script_env( + *, + container: str, + image: str, + avatar: str, + old_image_id: str, + new_image_id: str, + old_version: str, + new_version: str, + old_revision: str, + new_revision: str, + compose_service: str, + compose_workdir: str, + author_url: str, +) -> dict[str, str]: + env = os.environ.copy() + env.update( + { + "PULLIO_CONTAINER": container, + "PULLIO_IMAGE": image, + "PULLIO_AVATAR": avatar, + "PULLIO_OLD_IMAGE_ID": old_image_id, + "PULLIO_NEW_IMAGE_ID": new_image_id, + "PULLIO_OLD_VERSION": old_version, + "PULLIO_NEW_VERSION": new_version, + "PULLIO_OLD_REVISION": old_revision, + "PULLIO_NEW_REVISION": new_revision, + "PULLIO_COMPOSE_SERVICE": compose_service, + "PULLIO_COMPOSE_WORKDIR": compose_workdir, + "PULLIO_AUTHOR_URL": author_url, + } + ) + return env + + +def send_telegram_notification( + *, + status: str, + container_name: str, + old_version: str, + new_version: str, + image_name: str, + bot_token: str, + chat_id: str, + old_revision: str, + new_revision: str, + old_image_id: str, + new_image_id: str, + color: int, + author_avatar: str, + author_url: str, +) -> None: + version_indicator = "=" if old_version == new_version else ">" + revision_indicator = "=" if old_revision == new_revision else ">" + digest_indicator = "=" if old_image_id == new_image_id else ">" + + lines = [ + f"{container_name}", + status.replace("\\n", " "), + f"Image: {image_name}", + f"Image ID: {old_image_id[:11]} {digest_indicator} {new_image_id[:11]}", + ] + if old_version and new_version: + lines.append(f"Version: {old_version} {version_indicator} {new_version}") + if old_revision and new_revision: + lines.append( + f"Revision: {old_revision[:6]} {revision_indicator} {new_revision[:6]}" + ) + if author_url: + lines.append(f"URL: {author_url}") + if author_avatar: + lines.append(f"Avatar: {author_avatar}") + lines.append(f"Color: {color}") + lines.append(f"Time: {now_iso_utc()}") + + payload = { + "chat_id": chat_id, + "text": "\n".join(lines), + "disable_web_page_preview": True, + } + post_json(f"https://api.telegram.org/bot{bot_token}/sendMessage", payload) + + +def send_generic_webhook( + *, + status_generic: str, + container_name: str, + old_version: str, + new_version: str, + image_name: str, + webhook: str, + old_revision: str, + new_revision: str, + old_image_id: str, + new_image_id: str, + avatar: str, + author_url: str, +) -> None: + payload = { + "container": container_name, + "image": image_name, + "avatar": avatar, + "old_image_id": old_image_id, + "new_image_id": new_image_id, + "old_version": old_version, + "new_version": new_version, + "old_revision": old_revision, + "new_revision": new_revision, + "type": status_generic, + "url": author_url, + "timestamp": now_iso_utc(), + } + post_json(webhook, payload) + + +def process_container(config: Config, container_name: str) -> None: + LOGGER.info("%s: Checking...", container_name) + + image_name = docker_inspect_value( + config.docker_binary, container_name, "{{.Config.Image}}" + ) + container_image_digest = docker_inspect_value( + config.docker_binary, container_name, "{{.Image}}" + ) + + docker_compose_service = docker_inspect_value( + config.docker_binary, + container_name, + '{{ index .Config.Labels "com.docker.compose.service" }}', + ) + docker_compose_version = docker_inspect_value( + config.docker_binary, + container_name, + '{{ index .Config.Labels "com.docker.compose.version" }}', + ) + docker_compose_workdir = docker_inspect_value( + config.docker_binary, + container_name, + '{{ index .Config.Labels "com.docker.compose.project.working_dir" }}', + ) + + old_version = docker_inspect_value( + config.docker_binary, + container_name, + '{{ index .Config.Labels "org.opencontainers.image.version" }}', + ) + old_revision = docker_inspect_value( + config.docker_binary, + container_name, + '{{ index .Config.Labels "org.opencontainers.image.revision" }}', + ) + + pullio_update = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.update" }}}}', + ) + pullio_notify = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.notify" }}}}', + ) + pullio_telegram_bot_token = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.telegram.bot_token" }}}}', + ) + pullio_telegram_chat_id = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.telegram.chat_id" }}}}', + ) + pullio_generic_webhook = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.generic.webhook" }}}}', + ) + pullio_script_update = parse_script_command( + docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.script.update" }}}}', + ) + ) + pullio_script_notify = parse_script_command( + docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.script.notify" }}}}', + ) + ) + pullio_registry_authfile = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.registry.authfile" }}}}', + ) + pullio_author_avatar = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.author.avatar" }}}}', + ) + pullio_author_url = docker_inspect_value( + config.docker_binary, + container_name, + f'{{{{ index .Config.Labels "{LABEL_PREFIX}{config.tag}.author.url" }}}}', + ) + + if not docker_compose_version or ( + pullio_update != "true" and pullio_notify != "true" + ): + return + + if pullio_registry_authfile and Path(pullio_registry_authfile).is_file(): + LOGGER.info("%s: Registry login...", container_name) + try: + auth = json.loads( + Path(pullio_registry_authfile).read_text(encoding="utf-8") + ) + run_command( + [ + config.docker_binary, + "login", + "--username", + str(auth.get("username", "")), + "--password-stdin", + str(auth.get("registry", "")), + ], + input_text=str(auth.get("password", "")), + ) + except (json.JSONDecodeError, OSError, subprocess.CalledProcessError) as exc: + LOGGER.warning("%s: Registry login failed: %s", container_name, exc) + + LOGGER.info("%s: Pulling image...", container_name) + if not compose_pull(config, docker_compose_workdir, docker_compose_service): + LOGGER.error("%s: Pulling failed!", container_name) + + image_digest = docker_image_inspect_value(config.docker_binary, image_name, "{{.Id}}") + new_version = docker_image_inspect_value( + config.docker_binary, + image_name, + '{{ index .Config.Labels "org.opencontainers.image.version" }}', + ) + new_revision = docker_image_inspect_value( + config.docker_binary, + image_name, + '{{ index .Config.Labels "org.opencontainers.image.revision" }}', + ) + + status = "I've got an update waiting for me.\nGive it to me, please." + status_generic = "update_available" + color = 768753 + + if image_digest != container_image_digest and pullio_update == "true": + script_env = prepare_script_env( + container=container_name, + image=image_name, + avatar=pullio_author_avatar, + old_image_id=container_image_digest.removeprefix("sha256:"), + new_image_id=image_digest.removeprefix("sha256:"), + old_version=old_version, + new_version=new_version, + old_revision=old_revision, + new_revision=new_revision, + compose_service=docker_compose_service, + compose_workdir=docker_compose_workdir, + author_url=pullio_author_url, + ) + if pullio_script_update: + LOGGER.info("%s: Stopping container...", container_name) + try: + run_command([config.docker_binary, "stop", container_name]) + except subprocess.CalledProcessError: + LOGGER.warning( + "%s: Failed to stop container before update script.", container_name + ) + LOGGER.info("%s: Executing update script...", container_name) + try: + subprocess.run(pullio_script_update, env=script_env, check=False) + except OSError as exc: + LOGGER.warning( + "%s: Update script failed to start: %s", container_name, exc + ) + + LOGGER.info("%s: Updating container...", container_name) + if compose_up(config, docker_compose_workdir, docker_compose_service): + status = "I just updated myself.\nFeeling brand spanking new again!" + status_generic = "update_success" + color = 3066993 + else: + LOGGER.error("%s: Updating container failed!", container_name) + status = ( + "I tried to update myself.\nIt didn't work out, I might need some help." + ) + status_generic = "update_failure" + color = 15158332 + + notified_path = ( + config.cache_location / f"{config.script_hash}-{container_name}.notified" + ) + try: + notified_path.unlink(missing_ok=True) + except OSError: + LOGGER.warning("%s: Failed to clear notify cache file.", container_name) + + if image_digest != container_image_digest and pullio_notify == "true": + notified_path = ( + config.cache_location / f"{config.script_hash}-{container_name}.notified" + ) + try: + notified_path.touch(exist_ok=True) + notified_digest = notified_path.read_text(encoding="utf-8").strip() + except OSError: + notified_digest = "" + + if notified_digest != image_digest: + script_env = prepare_script_env( + container=container_name, + image=image_name, + avatar=pullio_author_avatar, + old_image_id=container_image_digest.removeprefix("sha256:"), + new_image_id=image_digest.removeprefix("sha256:"), + old_version=old_version, + new_version=new_version, + old_revision=old_revision, + new_revision=new_revision, + compose_service=docker_compose_service, + compose_workdir=docker_compose_workdir, + author_url=pullio_author_url, + ) + if pullio_script_notify: + LOGGER.info("%s: Executing notify script...", container_name) + try: + subprocess.run(pullio_script_notify, env=script_env, check=False) + except OSError as exc: + LOGGER.warning( + "%s: Notify script failed to start: %s", container_name, exc + ) + + old_digest_short = container_image_digest.removeprefix("sha256:") + new_digest_short = image_digest.removeprefix("sha256:") + + effective_telegram_bot_token = ( + pullio_telegram_bot_token or config.telegram_bot_token + ) + effective_telegram_chat_id = ( + pullio_telegram_chat_id or config.telegram_chat_id + ) + + if effective_telegram_bot_token and effective_telegram_chat_id: + LOGGER.info("%s: Sending telegram notification...", container_name) + send_telegram_notification( + status=status, + container_name=container_name, + old_version=old_version, + new_version=new_version, + image_name=image_name, + bot_token=effective_telegram_bot_token, + chat_id=effective_telegram_chat_id, + old_revision=old_revision, + new_revision=new_revision, + old_image_id=old_digest_short, + new_image_id=new_digest_short, + color=color, + author_avatar=pullio_author_avatar, + author_url=pullio_author_url, + ) + + if pullio_generic_webhook: + LOGGER.info("%s: Sending generic webhook...", container_name) + send_generic_webhook( + status_generic=status_generic, + container_name=container_name, + old_version=old_version, + new_version=new_version, + image_name=image_name, + webhook=pullio_generic_webhook, + old_revision=old_revision, + new_revision=new_revision, + old_image_id=old_digest_short, + new_image_id=new_digest_short, + avatar=pullio_author_avatar, + author_url=pullio_author_url, + ) + + try: + notified_path.write_text(image_digest, encoding="utf-8") + except OSError: + LOGGER.warning("%s: Failed to write notify cache file.", container_name) + + +def main() -> int: + args = parse_args() + + if args.debug: + logging.getLogger().setLevel(logging.DEBUG) + + compose_binary = os.getenv("COMPOSE_BINARY") or ( + shutil.which("docker-compose") or "" + ) + docker_binary = os.getenv("DOCKER_BINARY") or (shutil.which("docker") or "") + telegram_bot_token = os.getenv("TELEGRAM_BOT_TOKEN", "") + telegram_chat_id = os.getenv("TELEGRAM_CHAT_ID", "") + cache_location = Path("/tmp") + + tag = f".{args.tag}" if args.tag else "" + compose_type = detect_compose_type(compose_binary, docker_binary) + + if not docker_binary: + LOGGER.error("Docker binary not found.") + return 1 + + try: + script_hash = hashlib.sha1(Path(__file__).read_bytes()).hexdigest() + except OSError: + script_hash = hashlib.sha1( + str(Path(__file__).resolve()).encode("utf-8") + ).hexdigest() + + config = Config( + compose_binary=compose_binary, + docker_binary=docker_binary, + cache_location=cache_location, + tag=tag, + parallel=args.parallel, + compose_type=compose_type, + script_hash=script_hash, + telegram_bot_token=telegram_bot_token, + telegram_chat_id=telegram_chat_id, + ) + + LOGGER.info( + 'Running with "DEBUG=%s", "TAG=%s", and "PARALLEL=%s".', + args.debug, + tag, + args.parallel, + ) + + try: + raw_containers = run_command([docker_binary, "ps", "--format", "{{.Names}}"]) + except subprocess.CalledProcessError as exc: + LOGGER.error("Failed to list running containers: %s", exc) + return 1 + + containers = sorted([line for line in raw_containers.splitlines() if line]) + LOGGER.info( + "Processing %s containers with parallelism of %s", + len(containers), + args.parallel, + ) + + try: + if args.parallel > 1: + with ThreadPoolExecutor(max_workers=args.parallel) as executor: + list( + executor.map( + lambda name: process_container(config, name), containers + ) + ) + else: + for container_name in containers: + process_container(config, container_name) + except KeyboardInterrupt: + return 130 + + LOGGER.info("Pruning docker images...") + try: + run_command([docker_binary, "image", "prune", "--force"]) + except subprocess.CalledProcessError as exc: + LOGGER.warning("Image prune failed: %s", exc) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3001401 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "playground" +version = "0.0.1" +requires-python = ">=3.13" +dependencies = [ + "browser-cookie3>=0.20.1", + "bs4>=0.0.2", + "click>=8.3.0", + "cryptography>=46.0.3", + "debugpy>=1.8.15", + "httpx>=0.28.1", + "keyring>=25.6.0", + "mistralai>=1.5.1", + "openai>=1.98.0", + "pdf2image>=1.17.0", + "piexif>=1.1.3", + "playwright-stealth>=2.0.0", + "playwright[chromium]>=1.55.0", + "psycopg[binary]>=3.3.4", + "pytest>=9.1.1", + "srt>=3.5.3", + "typer-slim>=0.19.2", + "watchdog>=6.0.0", +] diff --git a/reset_router.py b/reset_router.py new file mode 100755 index 0000000..d063af1 --- /dev/null +++ b/reset_router.py @@ -0,0 +1,97 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = [ +# "playwright>=1.45.0", +# ] +# /// + +from __future__ import annotations + +import argparse +import sys +import time + +from playwright.sync_api import Playwright, TimeoutError as PlaywrightTimeoutError, sync_playwright + +ROUTER_URL = "http://192.168.1.1" +ROUTER_USERNAME = "abdus" +ROUTER_PASSWORD = "xAsametk50" + + +def reboot_router(playwright: Playwright, headed: bool, timeout_ms: int) -> None: + browser = playwright.chromium.launch(headless=not headed) + context = browser.new_context(ignore_https_errors=True) + page = context.new_page() + page.set_default_timeout(timeout_ms) + + try: + page.goto(f"{ROUTER_URL}/login", wait_until="domcontentloaded") + + username_box = page.get_by_role("textbox", name="User Name") + password_box = page.get_by_role("textbox", name="Password") + + username_box.click() + username_box.press("ControlOrMeta+a") + username_box.press("Backspace") + username_box.type(ROUTER_USERNAME, delay=80) + + password_box.click() + password_box.press("ControlOrMeta+a") + password_box.press("Backspace") + password_box.type(ROUTER_PASSWORD, delay=140) + page.get_by_role("textbox", name="Password").press("Enter") + + page.wait_for_load_state("domcontentloaded") + + try: + page.get_by_text("The username or password is not correct", exact=False).wait_for( + state="visible", timeout=1500 + ) + raise RuntimeError("Router rejected configured credentials") + except PlaywrightTimeoutError: + pass + + menu_trigger = page.locator("#h_menu_list").first + if menu_trigger.count() and menu_trigger.is_visible(): + menu_trigger.click() + + restart_icon = page.locator("#navbar_reboot .icon-menu-restart, #navbar_reboot").first + restart_icon.wait_for(state="visible") + + with page.expect_response(lambda r: "/cgi-bin/Reboot" in r.url, timeout=timeout_ms): + restart_icon.click() + page.get_by_role("button", name="OK").click() + + time.sleep(5) + + print("Router reboot request sent.") + finally: + context.close() + browser.close() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Log in and reboot router") + parser.add_argument("--headed", action="store_true", help="Run with visible browser") + parser.add_argument( + "--timeout", + type=int, + default=60000, + help="Timeout in milliseconds (default: 60000)", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + with sync_playwright() as playwright: + reboot_router(playwright, headed=args.headed, timeout_ms=args.timeout) + return 0 + except Exception as exc: + print(f"Error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/reset_router_http.py b/reset_router_http.py new file mode 100644 index 0000000..d04fb55 --- /dev/null +++ b/reset_router_http.py @@ -0,0 +1,84 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = [ +# "httpx>=0.27.0", +# ] +# /// + +import argparse +import time + +import httpx + +ROUTER_BASE_URL = "http://192.168.1.1" + +# Captured from successful login flow for abdus:xAsametk50. +# Firmware expects encrypted payload at /UserLogin. +USER_LOGIN_PAYLOAD = { + "content": "z7oqX7AWMrPjJdAEgrXCZY9q/+mrsyodlD53apbPmPXBkuB1SJXdoXr8+NcUbuRwn/jyc3HJlyKcMrAHDl3aFbN8/Gnyy8TzRKQkCAoSdVEVifzmWBPgpoDKvvT/p6n4TrAi2+b1eJ18Sr4ab1MWrI0XKNbBLEU+0UJlAd/UOHo=", + "key": "UwLW2sfQZ5S5XegolmSX0kqpxwzeJEIXq8lp3gLkNpNX+qiBAC4E2ColaigMtsMPbVEEFGk+RiPixF9L4xI4licoguc71zEfG9U3YUvNZZeBpV3JKMLLDzhs5aFDIb9M4k/zQ623zByOWDKjnUhHqqRK1I8hEyhThbbhKUkvo67/6vcRHHE9c0Lcxc2SrIGeyRUvPe5F+zXH8PbUn4aPq6fVTPmECRbiZOOm6zoy99PvyRVIYiCOkWqeXHhGIdMl2YZhA2UeM+1lwWtFQjSgefPeCfdhtIuEfFfUUaj3DyPJfD/+CyUT0QGUH6SD5sf1M8DySqOglgH1dIXEPLyIwA==", + "iv": "u+z76g62wnZYAehSfgffrPzpstykp6E891Z9CrDLfO0=", +} + + +def _extract_session_key(resp: httpx.Response) -> int: + data = resp.json() + if isinstance(data, list) and data: + return int(data[0]["SessionKey"]) + if isinstance(data, dict) and "SessionKey" in data: + return int(data["SessionKey"]) + raise RuntimeError(f"Could not parse SessionKey from /changeSessionKey response: {data!r}") + + +def _wait_for_router(client: httpx.Client, max_wait: float) -> None: + deadline = time.time() + max_wait + last_error = None + while time.time() < deadline: + try: + resp = client.get("/login") + resp.raise_for_status() + return + except (httpx.ConnectTimeout, httpx.ReadTimeout, httpx.ConnectError, httpx.NetworkError) as exc: + last_error = exc + time.sleep(2) + raise RuntimeError(f"Router not reachable within {max_wait:.0f}s") from last_error + + +def reboot_router(timeout: float, post_confirm_sleep: float) -> None: + with httpx.Client(base_url=ROUTER_BASE_URL, timeout=timeout, follow_redirects=True) as client: + _wait_for_router(client, max_wait=90) + + login = client.post("/UserLogin", json=USER_LOGIN_PAYLOAD) + login.raise_for_status() + + check = client.get("/cgi-bin/UserLoginCheck") + check.raise_for_status() + if "ZCFG_SUCCESS" not in check.text: + raise RuntimeError(f"Login check failed: {check.text}") + + key_resp = client.get("/changeSessionKey") + key_resp.raise_for_status() + session_key = _extract_session_key(key_resp) + + reboot = client.post("/cgi-bin/Reboot", params={"sessionkey": session_key}) + reboot.raise_for_status() + + time.sleep(post_confirm_sleep) + print("Router reboot request sent.") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Reboot router via HTTP API (no Playwright)") + parser.add_argument("--timeout", type=float, default=20.0, help="HTTP timeout in seconds") + parser.add_argument( + "--post-confirm-sleep", + type=float, + default=5.0, + help="Seconds to wait after reboot request", + ) + args = parser.parse_args() + reboot_router(timeout=args.timeout, post_confirm_sleep=args.post_confirm_sleep) + + +if __name__ == "__main__": + main() diff --git a/rg_download.py b/rg_download.py new file mode 100755 index 0000000..b3c6fee --- /dev/null +++ b/rg_download.py @@ -0,0 +1,565 @@ +#!/usr/bin/env -S uv run +# /// script +# requires-python = ">=3.12" +# dependencies = ["httpx"] +# /// +""" +RapidGator (rapidgator.net) downloader — reverse-engineered from JDownloader's RapidGatorNet plugin. + +Usage: + ./rg_download.py [--user EMAIL] [--pass PASSWORD] [--out DIR] [--2captcha KEY] + +Premium accounts use the clean v2 API path (no captcha, full speed). +Free/anonymous downloads scrape the website; reCAPTCHA v2 or Cloudflare Turnstile is +required — a 2captcha API key is strongly recommended (manual solving isn't feasible for +these captcha types). +""" + +import re +import sys +import os +import time +import argparse +import getpass +from pathlib import Path +from urllib.parse import urlparse, unquote + +import httpx + +# ── constants ──────────────────────────────────────────────────────────────── + +API_BASE = "https://rapidgator.net/api/v2/" +SITE_BASE = "https://rapidgator.net" + +SUPPORTED_DOMAINS = {"rapidgator.net", "rapidgator.asia", "rg.to"} + +# JD plugin regex: (?i)/file/([a-z0-9]{32}|\d+) +FILE_ID_RE = re.compile(r"(?i)/file/([a-z0-9]{32}|\d+)") + +HEADERS = { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 " + "(KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36", + "Accept-Language": "en-US,en;q=0.8", +} + + +# ── helpers ────────────────────────────────────────────────────────────────── + +def extract_file_id(url: str) -> str | None: + m = FILE_ID_RE.search(url) + return m.group(1) if m else None + + +def check_domain(url: str): + host = urlparse(url).hostname or "" + host = host.removeprefix("www.") + if host not in SUPPORTED_DOMAINS: + sys.exit(f"Unsupported domain '{host}'. Supported: {', '.join(sorted(SUPPORTED_DOMAINS))}") + + +def _api_get(session: httpx.Client, endpoint: str, params: dict) -> dict: + """GET request to the v2 API; raises on HTTP error, returns parsed JSON.""" + r = session.get(f"{API_BASE}{endpoint}", params=params, headers=HEADERS, timeout=30) + try: + data = r.json() + except Exception: + r.raise_for_status() + raise + return data + + +def _api_check(data: dict, context: str = "") -> dict: + """Raise a RuntimeError if the API returned an error status.""" + status = data.get("status", "") + if status != "success": + err = data.get("details") or data.get("error") or data + raise RuntimeError(f"API error{' (' + context + ')' if context else ''}: {err}") + return data.get("details", {}) + + +# ── API calls (mirroring JD's loginAPI / requestFileInformationAPI / handlePremium_api) ── + +def api_login(session: httpx.Client, email: str, password: str) -> str: + """GET /user/login → returns session_id (used as ?token= in all other calls).""" + data = _api_get(session, "user/login", {"login": email, "password": password}) + details = _api_check(data, "login") + # JD stores the field named "session_id" (PROPERTY_sessionid constant) + sid = details.get("session_id") or details.get("token") + if not sid: + raise RuntimeError(f"Login succeeded but no session_id in response: {details}") + return sid + + +def api_user_info(session: httpx.Client, token: str) -> dict: + """GET /user/info?token= → account details (is_premium, premium_end_time, …).""" + data = _api_get(session, "user/info", {"token": token}) + return _api_check(data, "user/info") + + +def api_file_info(session: httpx.Client, token: str, file_id: str) -> dict: + """GET /file/info?token=&file_id= → file metadata (name, size, hash, …).""" + data = _api_get(session, "file/info", {"token": token, "file_id": file_id}) + details = _api_check(data, "file/info") + # Nested under "file" key + return details.get("file") or details + + +def api_file_download(session: httpx.Client, token: str, file_id: str) -> str: + """GET /file/download?token=&file_id= → direct download_url.""" + data = _api_get(session, "file/download", {"token": token, "file_id": file_id}) + details = _api_check(data, "file/download") + url = details.get("download_url") + if not url: + raise RuntimeError(f"No download_url in response: {details}") + return url + + +# ── 2captcha ───────────────────────────────────────────────────────────────── + +class TwoCaptcha: + """ + Thin wrapper around the 2captcha API v2. + Supports reCAPTCHA v2 (RecaptchaV2TaskProxyless) and + Cloudflare Turnstile (TurnstileTaskProxyless) — the two captcha types + used by RapidGator's free download page. + """ + + API = "https://api.2captcha.com" + POLL_INTERVAL = 5 + POLL_TIMEOUT = 180 + + def __init__(self, api_key: str, client: httpx.Client): + self.api_key = api_key + self.client = client + + def _call(self, endpoint: str, payload: dict) -> dict: + r = self.client.post( + f"{self.API}{endpoint}", + json={"clientKey": self.api_key, **payload}, + headers={"Content-Type": "application/json"}, + timeout=30, + ) + r.raise_for_status() + data = r.json() + err_id = data.get("errorId") or 0 + if err_id and err_id != 0: + raise RuntimeError( + f"2captcha error: {data.get('errorCode') or data.get('errorDescription') or data}" + ) + return data + + def _poll(self, task_id: int) -> str: + """Poll until ready, return solution token.""" + print(f" 2captcha task {task_id} submitted, polling…", flush=True) + deadline = time.monotonic() + self.POLL_TIMEOUT + while time.monotonic() < deadline: + time.sleep(self.POLL_INTERVAL) + result = self._call("/getTaskResult", {"taskId": task_id}) + if result.get("status") == "ready": + token = result["solution"]["token"] + print(f" 2captcha solved (token length: {len(token)})") + return token + print(" …waiting for solution") + raise RuntimeError(f"2captcha timed out after {self.POLL_TIMEOUT}s") + + def solve_recaptcha_v2(self, site_key: str, page_url: str) -> str: + """Submit a RecaptchaV2TaskProxyless and return the g-recaptcha-response token.""" + task_id = self._call("/createTask", { + "task": { + "type": "RecaptchaV2TaskProxyless", + "websiteURL": page_url, + "websiteKey": site_key, + } + })["taskId"] + return self._poll(task_id) + + def solve_turnstile(self, site_key: str, page_url: str) -> str: + """Submit a TurnstileTaskProxyless and return the cf-turnstile-response token.""" + task_id = self._call("/createTask", { + "task": { + "type": "TurnstileTaskProxyless", + "websiteURL": page_url, + "websiteKey": site_key, + } + })["taskId"] + return self._poll(task_id) + + +# ── free download (website scrape) ──────────────────────────────────────────── + +_RECAPTCHA_RE = re.compile( + r'class=["\']g-recaptcha["\'][^>]*data-sitekey=["\']([^"\']+)["\']' + r'|data-sitekey=["\']([^"\']+)["\'][^>]*class=["\']g-recaptcha["\']' +) +_TURNSTILE_RE = re.compile( + r'class=["\']cf-turnstile["\'][^>]*data-sitekey=["\']([^"\']+)["\']' + r'|data-sitekey=["\']([^"\']+)["\'][^>]*class=["\']cf-turnstile["\']' +) +_DL_URL_RE = re.compile( + r"'(https?://[A-Za-z0-9\-_]+\.[^/]+//\?r=download/index&session_id=[A-Za-z0-9]+)'" + r"|\"(https?://[^/]+/download/[^<>\"']+)\"" +) + +# X-Requested-With is required by the AJAX endpoints — they return empty without it +_XHR_HEADERS = {"X-Requested-With": "XMLHttpRequest"} + + +def _first_group(m: re.Match | None) -> str | None: + if not m: + return None + return next((g for g in m.groups() if g), None) + + +def free_download( + session: httpx.Client, + file_id: str, + file_name: str, + solver: TwoCaptcha | None, + out_dir: Path, +): + """ + Website-based free download flow (mirrors JD handleDownloadWebsite / free path). + + Actual page flow discovered by inspection: + 1. GET /file/ → server sets PHPSESSID + file_id cookies; + page JS has: var fid = ; var secs = 180; + var startTimerUrl = '/download/AjaxStartTimer'; + 2. GET /download/AjaxStartTimer?fid= (XHR) + → {"state":"started", "sid":""} + 3. Wait `secs` seconds (free-tier countdown) + 4. GET /download/AjaxGetDownloadLink?sid= (XHR) + → poll until {"state":"done"} + 5. GET /download/captcha → HTML with reCAPTCHA v2 or Turnstile widget + 6. POST /download/captcha with solved token + → redirect to the actual download URL OR + response body contains it + """ + file_url = f"{SITE_BASE}/file/{file_id}" + + # ── step 1: load file page ──────────────────────────────────────────────── + print(f"Loading file page: {file_url}") + r = session.get(file_url, headers=HEADERS, timeout=30) + r.raise_for_status() + page_html = r.text + + # Check for premium-only direct link already embedded (logged-in premium path) + pm = re.search(r"var\s+premium_download_link\s*=\s*'(https?://[^']+)'", page_html) + if pm and pm.group(1): + stream_download(session, pm.group(1), out_dir, file_name) + return + + # Extract numeric fid (set by server in page JS and cookie, distinct from hex URL ID) + fid_m = re.search(r"var\s+fid\s*=\s*(\d+)", page_html) + if not fid_m: + sys.exit("Could not find numeric fid in page JS — page layout may have changed.") + numeric_fid = fid_m.group(1) + + # Extract client-side wait timer (default 180 for free users) + secs_m = re.search(r"var\s+secs\s*=\s*(\d+)", page_html) + wait_secs = int(secs_m.group(1)) if secs_m else 180 + + print(f" Numeric fid={numeric_fid}, wait={wait_secs}s") + + if wait_secs > 600: + sys.exit(f"Server wait is {wait_secs}s — IP appears rate-limited. Try again later.") + + # ── step 2: start the server-side timer ────────────────────────────────── + print("Starting download timer…") + r2 = session.get( + f"{SITE_BASE}/download/AjaxStartTimer", + params={"fid": numeric_fid}, + headers={**HEADERS, **_XHR_HEADERS, "Referer": file_url}, + timeout=30, + ) + r2.raise_for_status() + try: + timer_data = r2.json() + except Exception: + sys.exit(f"AjaxStartTimer returned non-JSON: {r2.text[:300]!r}") + + if timer_data.get("state") == "error": + sys.exit(f"AjaxStartTimer error: {timer_data.get('code') or timer_data}") + if timer_data.get("state") != "started": + sys.exit(f"Unexpected AjaxStartTimer state: {timer_data}") + + sid = timer_data.get("sid") + if not sid: + sys.exit(f"AjaxStartTimer response missing sid: {timer_data}") + print(f" sid={sid}") + + # ── step 3: wait for the server countdown ──────────────────────────────── + print(f" Waiting {wait_secs}s…", flush=True) + for remaining in range(wait_secs, 0, -5): + time.sleep(min(5, remaining)) + print(f"\r {remaining - min(5, remaining)}s remaining… ", end="", flush=True) + print() + + # ── step 4: confirm server is ready ────────────────────────────────────── + print("Confirming download ready…") + for attempt in range(10): + r3 = session.get( + f"{SITE_BASE}/download/AjaxGetDownloadLink", + params={"sid": sid}, + headers={**HEADERS, **_XHR_HEADERS, "Referer": file_url}, + timeout=30, + ) + r3.raise_for_status() + try: + link_data = r3.json() + except Exception: + sys.exit(f"AjaxGetDownloadLink returned non-JSON: {r3.text[:300]!r}") + + dl_state = link_data.get("state") + if dl_state == "done": + break + if dl_state == "error": + code = link_data.get("code", "") + if "wait" in str(code).lower(): + print(f" Server not ready yet, retrying in 5s… ({code})") + time.sleep(5) + continue + sys.exit(f"AjaxGetDownloadLink error: {link_data}") + print(f" Unexpected state '{dl_state}', retrying…") + time.sleep(5) + else: + sys.exit(f"AjaxGetDownloadLink never returned done. Last response: {link_data}") + + # ── step 5: load captcha page ───────────────────────────────────────────── + captcha_url = f"{SITE_BASE}/download/captcha" + print("Loading captcha page…") + r4 = session.get(captcha_url, headers={**HEADERS, "Referer": file_url}, timeout=30) + r4.raise_for_status() + captcha_html = r4.text + + if not captcha_html.strip(): + sys.exit("Captcha page returned empty — session may have expired. Try again.") + + # Detect captcha type and extract site key + recaptcha_key = _first_group(_RECAPTCHA_RE.search(captcha_html)) + turnstile_key = _first_group(_TURNSTILE_RE.search(captcha_html)) + + # Also extract CSRF token from if present + csrf_m = re.search(r'