Initial Commit
This commit is contained in:
218
image_to_3d.py
Normal file
218
image_to_3d.py
Normal file
@@ -0,0 +1,218 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user