#!/usr/bin/env python3 import base64 import json import os import subprocess import sys import tempfile import time from pathlib import Path from typing import Any import replicate import requests from PIL import Image MODEL_NAME = "tencent/hunyuan-3d-3.1" TIMEOUT = 900 PREDICTION_TIMEOUT = 10 * 60 POLL_INTERVAL = 2.0 MAX_INPUT_BYTES = 6 * 1024 * 1024 MAX_DATA_URI_BYTES = 1024 * 1024 MAX_INPUT_SIDE = 2048 SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"} def notify(title: str, message: str) -> None: script = f"display notification {json.dumps(message)} with title {json.dumps(title)}" try: subprocess.run(["osascript", "-e", script], check=False) except OSError: pass def _has_alpha(img: Image.Image) -> bool: return img.mode in {"RGBA", "LA"} or ( img.mode == "P" and "transparency" in img.info ) def _save_compact_image(img: Image.Image, output_path: str) -> None: if max(img.size) > MAX_INPUT_SIDE: img = img.copy() img.thumbnail((MAX_INPUT_SIDE, MAX_INPUT_SIDE), Image.Resampling.LANCZOS) if _has_alpha(img): img.save(output_path, "WEBP", quality=95, method=6) return if img.mode != "RGB": img = img.convert("RGB") img.save(output_path, "JPEG", quality=92, optimize=True) def prepare_input_image(src_path: str, temp_paths: list[str]) -> str: ext = Path(src_path).suffix.lower() if ext in SUPPORTED_EXTENSIONS and os.path.getsize(src_path) <= MAX_DATA_URI_BYTES: return src_path img = Image.open(src_path) suffix = ".webp" if _has_alpha(img) else ".jpg" fd, temp_path = tempfile.mkstemp(suffix=suffix) os.close(fd) _save_compact_image(img, temp_path) temp_paths.append(temp_path) if os.path.getsize(temp_path) > MAX_INPUT_BYTES: raise RuntimeError("Prepared image is still larger than Replicate's 6MB limit.") return temp_path def run_replicate(image_path: str, api_token: str) -> Any: client = replicate.Client(api_token=api_token) client.poll_interval = POLL_INTERVAL with open(image_path, "rb") as image_file: prediction = client.models.predictions.create( model=MODEL_NAME, input={ "image": image_file, "generate_type": "Normal", "face_count": 500000, "enable_pbr": False, }, wait=False, file_encoding_strategy="base64", ) print(f"Replicate prediction started: {prediction.id} ({prediction.status})") deadline = time.monotonic() + PREDICTION_TIMEOUT last_status = prediction.status while prediction.status not in {"succeeded", "failed", "canceled"}: if time.monotonic() >= deadline: raise TimeoutError( f"Timed out waiting for Replicate prediction {prediction.id} " f"after {PREDICTION_TIMEOUT // 60} minutes. It may still be running." ) time.sleep(client.poll_interval) prediction.reload() if prediction.status != last_status: print(f"Replicate prediction {prediction.id}: {prediction.status}") last_status = prediction.status if prediction.status != "succeeded": detail = prediction.error or prediction.logs or f"status={prediction.status}" raise RuntimeError(f"Replicate prediction {prediction.id} failed: {detail}") return prediction.output def _extract_output_file(output: Any) -> Any: if isinstance(output, (list, tuple)): if not output: raise RuntimeError("Replicate returned an empty output.") return output[0] if isinstance(output, dict): for key in ("output", "model", "mesh", "glb", "url"): if output.get(key): return _extract_output_file(output[key]) raise RuntimeError(f"Replicate returned an unsupported output shape: {output}") return output def _bytes_from_url(url: str) -> bytes: if url.startswith("data:"): _, encoded = url.split(",", 1) return base64.b64decode(encoded) response = requests.get(url, timeout=TIMEOUT) response.raise_for_status() return response.content def write_output(output: Any, output_path: str) -> None: output_file = _extract_output_file(output) if hasattr(output_file, "read"): data = output_file.read() elif hasattr(output_file, "url"): data = _bytes_from_url(str(output_file.url)) elif isinstance(output_file, str): data = _bytes_from_url(output_file) else: raise RuntimeError(f"Replicate returned an unsupported output type: {type(output_file)!r}") with open(output_path, "wb") as f: f.write(data) def process_image(img_path: str, api_token: str | None = None) -> str | None: api_token = (api_token or os.environ.get("REPLICATE_API_TOKEN", "")).strip() if not api_token: msg = "Missing Replicate API token. Add it in app settings or set REPLICATE_API_TOKEN." print(msg) notify("3D conversion error", msg) return None img_path = os.path.abspath(img_path) if not os.path.isfile(img_path): msg = f"Image not found: {img_path}" print(msg) notify("3D conversion error", msg) return None temp_paths: list[str] = [] try: input_path = prepare_input_image(img_path, temp_paths) base_name = os.path.splitext(os.path.basename(img_path))[0] output_path = os.path.join(os.path.dirname(img_path), f"{base_name}.glb") print(f"Running {MODEL_NAME} on Replicate...") output = run_replicate(input_path, api_token) write_output(output, output_path) msg = f"3D model saved: {output_path}" print(msg) notify("3D conversion complete", msg) return output_path except Exception as e: msg = f"3D conversion failed for {img_path}: {e}" print(msg) notify("3D conversion error", msg) return None finally: for temp_path in temp_paths: try: os.remove(temp_path) except OSError: pass def main() -> None: if len(sys.argv) < 2: msg = "Usage: python image_to_3d.py [image2 ...]" notify("3D conversion error", "No image files provided.") print(msg) sys.exit(1) outputs = [] for img_path in sys.argv[1:]: print(f"\nProcessing: {img_path}") result = process_image(img_path) if result: outputs.append(result) if not outputs: sys.exit(1) notify("3D conversion", f"Finished {len(outputs)} model(s).") if __name__ == "__main__": main()