219 lines
6.6 KiB
Python
219 lines
6.6 KiB
Python
#!/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 <image1> [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()
|