360 lines
11 KiB
Python
Executable File
360 lines
11 KiB
Python
Executable File
#!/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 = """
|
|
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 noise and 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())
|