HartOMat/render-worker/scripts/blender_render.py

"""
Blender Python script for rendering a GLB file to PNG.
Targets Blender 5.0+ (EEVEE / Cycles).

Called by Blender:
  blender --background --python blender_render.py -- \
      <glb_path> <output_path> <width> <height> [engine] [samples]

engine: "cycles" (default) | "eevee"

Features:
- OCC-generated GLB: one mesh per STEP part, already in metres.
- Bounding-box-aware camera: object fills ~85 % of the frame.
- Isometric-style angle (elevation 28°, azimuth 40°).
- Dynamic clip planes.
- Standard (non-Filmic) colour management → no grey tint.
"""
import sys
import os

# Force unbuffered stdout so render log lines appear immediately
os.environ["PYTHONUNBUFFERED"] = "1"
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(line_buffering=True)

# Add script directory to sys.path so Blender Python finds our submodules
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import bpy  # type: ignore[import]

from _blender_gpu       import activate_gpu, configure_engine
from _blender_import    import import_glb, apply_rotation
from _blender_materials import (
    FAILED_MATERIAL_NAME, assign_failed_material,
    build_mat_map_lower, apply_material_library,
)
from _blender_camera    import setup_auto_camera, setup_auto_lights
from _blender_scene     import (
    ensure_collection, apply_smooth_batch,
    apply_sharp_edges_from_occ, setup_shadow_catcher,
)

# ── Parse arguments ────────────────────────────────────────────────────────────
import json as _json

def _arg(n, default="", transform=str):
    return transform(argv[n]) if len(argv) > n and argv[n] else default

argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
if len(argv) < 4:
    print("Usage: blender --background --python blender_render.py -- "
          "<glb_path> <output_path> <width> <height> ...")
    sys.exit(1)

glb_path              = argv[0]
output_path           = argv[1]
width                 = int(argv[2])
height                = int(argv[3])
engine                = _arg(4, "cycles", str.lower)
samples               = _arg(5, None, int)
smooth_angle          = _arg(6, 30, int)
cycles_device         = _arg(7, "auto", str.lower)
transparent_bg        = argv[8] == "1" if len(argv) > 8 else False
template_path         = _arg(9, "")
target_collection     = _arg(10, "Product")
material_library_path = _arg(11, "")
material_map          = _json.loads(_arg(12, "{}")) if _arg(12, "{}") else {}
part_names_ordered    = _json.loads(_arg(13, "[]")) if _arg(13, "[]") else []
lighting_only         = argv[14] == "1" if len(argv) > 14 else False
shadow_catcher        = argv[15] == "1" if len(argv) > 15 else False
rotation_x            = _arg(16, 0.0, float)
rotation_y            = _arg(17, 0.0, float)
rotation_z            = _arg(18, 0.0, float)
noise_threshold_arg        = _arg(19, "")
denoiser_arg               = _arg(20, "")
denoising_input_passes_arg = _arg(21, "")
denoising_prefilter_arg    = _arg(22, "")
denoising_quality_arg      = _arg(23, "")
denoising_use_gpu_arg      = _arg(24, "")

if samples is None:
    samples = 64 if engine == "eevee" else 256

# Named argument: --mesh-attributes <json>
_mesh_attrs: dict = {}
if "--mesh-attributes" in sys.argv:
    _idx = sys.argv.index("--mesh-attributes")
    try:
        _mesh_attrs = _json.loads(sys.argv[_idx + 1])
    except Exception:
        pass

if template_path and not os.path.isfile(template_path):
    print(f"[blender_render] ERROR: template not found: {template_path}")
    sys.exit(1)

use_template = bool(template_path)
print(f"[blender_render] engine={engine}, samples={samples}, size={width}x{height}, smooth_angle={smooth_angle}°, device={cycles_device}, transparent={transparent_bg}")
print(f"[blender_render] part_names_ordered: {len(part_names_ordered)} entries")
print(f"[blender_render] {'template='+template_path+', collection='+target_collection+', lighting_only='+str(lighting_only) if use_template else 'no template — Mode A'}")
if material_library_path:
    print(f"[blender_render] material_library={material_library_path}, material_map keys={list(material_map.keys())}")

# ── Early GPU activation (must happen BEFORE open_mainfile / Cycles init) ─────
_early_gpu_type = activate_gpu(cycles_device)

# ── Timing harness ─────────────────────────────────────────────────────────────
import time as _time
_t0 = _time.monotonic()
_timings: dict = {}


def _lap(label: str) -> None:
    now = _time.monotonic()
    if not hasattr(_lap, '_last'):
        _lap._last = _t0
    delta = now - _lap._last
    total = now - _t0
    _timings[label] = round(delta, 3)
    print(f"[blender_render] TIMING {label}={delta:.2f}s (total={total:.2f}s)", flush=True)
    _lap._last = now


# ── SCENE SETUP ───────────────────────────────────────────────────────────────

if use_template:
    # ── MODE B: Template-based render ─────────────────────────────────────────
    print(f"[blender_render] Opening template: {template_path}")
    bpy.ops.wm.open_mainfile(filepath=template_path)
    _lap("template_load")

    target_col = ensure_collection(target_collection)
    parts = import_glb(glb_path)
    _lap("glb_import")
    apply_rotation(parts, rotation_x, rotation_y, rotation_z)
    _lap("rotation")

    for part in parts:
        for col in list(part.users_collection):
            col.objects.unlink(part)
        target_col.objects.link(part)

    apply_smooth_batch(parts, smooth_angle)
    _occ_pairs = _mesh_attrs.get("sharp_edge_pairs") or []
    if _occ_pairs:
        apply_sharp_edges_from_occ(parts, _occ_pairs)
    _lap("smooth_shading")

    if material_library_path and material_map:
        apply_material_library(parts, material_library_path, build_mat_map_lower(material_map), part_names_ordered)
    else:
        for part in parts:
            assign_failed_material(part)
    _lap("material_assign")

    if shadow_catcher:
        setup_shadow_catcher(parts)

    needs_auto_camera = (lighting_only and not shadow_catcher) or not bpy.context.scene.camera
    if lighting_only and not shadow_catcher:
        print("[blender_render] lighting_only mode: using template World/HDRI, forcing auto-camera")
    elif needs_auto_camera:
        print("[blender_render] WARNING: template has no camera — will create auto-camera")

    if not needs_auto_camera and bpy.context.scene.camera:
        bpy.context.scene.camera.data.clip_start = 0.001

    print(f"[blender_render] template mode: {len(parts)} parts imported into collection '{target_collection}'")

else:
    # ── MODE A: Factory settings ───────────────────────────────────────────────
    needs_auto_camera = True
    bpy.ops.wm.read_factory_settings(use_empty=True)
    parts = import_glb(glb_path)
    apply_rotation(parts, rotation_x, rotation_y, rotation_z)

    _t_smooth_a = _time.time()
    apply_smooth_batch(parts, smooth_angle)
    _occ_pairs = _mesh_attrs.get("sharp_edge_pairs") or []
    if _occ_pairs:
        apply_sharp_edges_from_occ(parts, _occ_pairs)
    for part in parts:
        assign_failed_material(part)
    print(f"[blender_render] smooth+fallback-material: {len(parts)} parts ({_time.time()-_t_smooth_a:.2f}s)", flush=True)

    if material_library_path and material_map:
        apply_material_library(parts, material_library_path, build_mat_map_lower(material_map), part_names_ordered)

if needs_auto_camera:
    bbox_center, bsphere_radius = setup_auto_camera(parts, width, height)
    if not use_template:
        setup_auto_lights(bbox_center, bsphere_radius)
        # Mode A world background
        world = bpy.data.worlds.new("World")
        bpy.context.scene.world = world
        world.use_nodes = True
        bg = world.node_tree.nodes["Background"]
        bg.inputs["Color"].default_value    = (0.96, 0.96, 0.97, 1.0)
        bg.inputs["Strength"].default_value = 0.15

# ── Render engine ──────────────────────────────────────────────────────────────
scene = bpy.context.scene
engine = configure_engine(
    scene, engine, samples, cycles_device, _early_gpu_type,
    noise_threshold_arg, denoiser_arg,
    denoising_input_passes_arg, denoising_prefilter_arg,
    denoising_quality_arg, denoising_use_gpu_arg,
)

# ── Colour management ──────────────────────────────────────────────────────────
if not use_template:
    scene.view_settings.view_transform = 'Standard'
    scene.view_settings.exposure        = 0.0
    scene.view_settings.gamma           = 1.0
    try:
        scene.view_settings.look = 'None'
    except Exception:
        pass

# ── Render settings ────────────────────────────────────────────────────────────
scene.render.resolution_x               = width
scene.render.resolution_y               = height
scene.render.resolution_percentage      = 100
scene.render.image_settings.file_format = 'PNG'
scene.render.filepath                   = output_path
scene.render.film_transparent           = transparent_bg

# ── Final verification + render ────────────────────────────────────────────────
if scene.render.engine == 'CYCLES':
    cprefs = bpy.context.preferences.addons['cycles'].preferences
    print(f"[blender_render] VERIFY: engine={scene.render.engine}, "
          f"cycles.device={scene.cycles.device}, "
          f"compute_device_type={cprefs.compute_device_type}, "
          f"gpu_devices={[(d.name, d.type, d.use) for d in cprefs.devices if d.type != 'CPU']}",
          flush=True)

_lap("pre_render_setup")
print(f"[blender_render] Rendering → {output_path}  (Blender {bpy.app.version_string})", flush=True)
sys.stdout.flush()
bpy.ops.render.render(write_still=True)
print("[blender_render] render done.", flush=True)
_lap("gpu_render")

# ── Final timing summary ───────────────────────────────────────────────────────
_total = _time.monotonic() - _t0
print(f"[blender_render] TIMING_SUMMARY total={_total:.2f}s | " +
      " | ".join(f"{k}={v:.2f}s" for k, v in _timings.items()), flush=True)
print("[blender_render] Done.")