valknar
4cd88ba477
kling-o1-pro and kling-o1-std silently fail when aspect_ratio is included in the payload — they derive it from the input image. Added VIDEO_ASPECT_RATIO_MODELS whitelist so only kling-elements and minimax-hailuo receive the parameter. Also switched image_to_base64 to use Pillow for format detection instead of trusting the file extension, which correctly handles files saved with the wrong extension. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
"""File I/O utilities: base64 encoding, output path generation, download."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
DownloadColumn,
|
|
Progress,
|
|
SpinnerColumn,
|
|
TimeElapsedColumn,
|
|
TransferSpeedColumn,
|
|
)
|
|
|
|
|
|
def image_to_base64(path: Path) -> str:
|
|
"""Read an image file and return a base64-encoded string.
|
|
|
|
Uses Pillow to detect the actual format rather than trusting the file
|
|
extension — mismatched extensions (e.g. a JPEG saved as .png) would
|
|
produce an incorrect MIME type that causes silent failures with some models.
|
|
"""
|
|
from PIL import Image
|
|
|
|
_pillow_to_mime = {
|
|
"JPEG": "image/jpeg",
|
|
"PNG": "image/png",
|
|
"GIF": "image/gif",
|
|
"WEBP": "image/webp",
|
|
}
|
|
with Image.open(path) as img:
|
|
fmt = img.format or "JPEG"
|
|
mime = _pillow_to_mime.get(fmt, "image/jpeg")
|
|
|
|
with open(path, "rb") as f:
|
|
encoded = base64.b64encode(f.read()).decode()
|
|
return f"data:{mime};base64,{encoded}"
|
|
|
|
|
|
def video_to_base64(path: Path) -> str:
|
|
"""Read a video file and return a base64-encoded string."""
|
|
suffix = path.suffix.lower().lstrip(".")
|
|
mime_map = {
|
|
"mp4": "video/mp4",
|
|
"mov": "video/quicktime",
|
|
"avi": "video/x-msvideo",
|
|
"webm": "video/webm",
|
|
}
|
|
mime = mime_map.get(suffix, "video/mp4")
|
|
with open(path, "rb") as f:
|
|
encoded = base64.b64encode(f.read()).decode()
|
|
return f"data:{mime};base64,{encoded}"
|
|
|
|
|
|
def auto_output_path(task_type: str, model: str, ext: str = "jpg", output_dir: str = ".") -> Path:
|
|
"""Generate a timestamped output filename."""
|
|
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
safe_model = model.replace("/", "-").replace(" ", "-")
|
|
filename = f"freepik_{task_type}_{safe_model}_{ts}.{ext}"
|
|
return Path(output_dir) / filename
|
|
|
|
|
|
def save_from_url(url: str, dest: Path, console: Console) -> None:
|
|
"""Stream-download a URL to a file, displaying a Rich progress bar."""
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with Progress(
|
|
SpinnerColumn(style="bold magenta"),
|
|
"[progress.description]{task.description}",
|
|
BarColumn(bar_width=30, style="magenta", complete_style="green"),
|
|
"[progress.percentage]{task.percentage:>3.0f}%",
|
|
DownloadColumn(),
|
|
TransferSpeedColumn(),
|
|
TimeElapsedColumn(),
|
|
console=console,
|
|
transient=True,
|
|
) as progress:
|
|
with httpx.stream("GET", url, follow_redirects=True, timeout=120) as r:
|
|
r.raise_for_status()
|
|
total = int(r.headers.get("content-length", 0)) or None
|
|
task = progress.add_task(
|
|
f"[dim]Saving[/dim] [bold]{dest.name}[/bold]",
|
|
total=total,
|
|
)
|
|
with open(dest, "wb") as f:
|
|
for chunk in r.iter_bytes(chunk_size=65536):
|
|
f.write(chunk)
|
|
progress.advance(task, len(chunk))
|
|
|
|
|
|
def get_image_dimensions(path: Path) -> tuple[Optional[int], Optional[int]]:
|
|
"""Return (width, height) of an image using Pillow."""
|
|
try:
|
|
from PIL import Image
|
|
|
|
with Image.open(path) as img:
|
|
return img.width, img.height
|
|
except Exception:
|
|
return None, None
|