render.py

#!/usr/bin/env python3
"""Render generated Blender Python scripts at 4 reference camera angles.

Camera path matches the Infinigen 40-frame turntable used to produce
the reference renders in `output_factories/<cat>/<F>/<F>_NNN/Image_*.png`:

    cam(t) = (cam_r·cos θ, cam_r·sin θ, cam_h),   θ = 9° · frame_idx
    cam_r  = 1.8 · extent      cam_h = 0.6 · extent

We render frames {5, 15, 25, 35} which correspond to azimuths
{45°, 135°, 225°, 315°} — these are the four reference views the eval
data was selected at.

Two-mode script:

    1. Orchestrator (normal Python): discover instances under
       `results/<model>/<inst>/` and spawn a Blender subprocess per
       generated `.py` file.
    2. Blender mode (with `--blender-render`): runs inside Blender,
       executes the generated script, sets up its own camera + lights,
       renders 4 views and writes `render_log.json`.

Per-instance log fields (written into `<inst>/renders/render_log.json`):
    status:    OK | ERR_EXEC | ERR_NO_MESH | ERR_RENDER | ERR_TIMEOUT | ERR_NOLOG
    n_meshes:  count of mesh objects after script execution
    extent:    bounding-box extent of the scene
    n_views_rendered: 0..4
    latency_s, error, ...

Usage:
    # Render all 212 instances of one model
    python eval/utils/render.py --model gemini-3-flash-preview

    # A subset for testing
    python eval/utils/render.py --model gemini-3-flash-preview \\
        --instances ArmChair_seed0 BeverageFridge_seed0
"""

import sys
import os

_BL_ARGV = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []


# ╔══════════════════════════════════════════════════════════════════╗
# ║  Blender-side  (active when --blender-render is in -- argv)      ║
# ╚══════════════════════════════════════════════════════════════════╝
if "--blender-render" in _BL_ARGV:
    import argparse
    import json
    import math
    import runpy
    import time
    import traceback

    import bpy
    from mathutils import Vector

    # frame_idx → output filename (azimuth_deg = 9° × frame_idx)
    REF_VIEWS = [(5,  "Image_005.png"),
                 (15, "Image_015.png"),
                 (25, "Image_025.png"),
                 (35, "Image_035.png")]

    def _bl_parse():
        p = argparse.ArgumentParser()
        p.add_argument("--blender-render", action="store_true")
        p.add_argument("--script",      required=True)
        p.add_argument("--output-dir",  required=True)
        p.add_argument("--samples",     type=int, default=64)
        p.add_argument("--resolution",  type=int, default=512)
        p.add_argument("--engine",      default="CYCLES",
                       choices=["CYCLES", "BLENDER_EEVEE"])
        return p.parse_args(_BL_ARGV)

    def _clear_scene():
        bpy.ops.object.select_all(action="SELECT")
        bpy.ops.object.delete()
        for col in (bpy.data.meshes, bpy.data.curves, bpy.data.cameras,
                    bpy.data.lights, bpy.data.materials, bpy.data.node_groups):
            for item in list(col):
                col.remove(item)
        bpy.context.scene.cursor.location = (0, 0, 0)

    def _run_user_script(path):
        """Execute the generated script, then strip any camera/light it created."""
        runpy.run_path(path, run_name="__main__")
        for o in list(bpy.context.scene.objects):
            if o.type in ("CAMERA", "LIGHT"):
                bpy.data.objects.remove(o, do_unlink=True)
        return [o for o in bpy.context.scene.objects if o.type == "MESH"]

    def _setup_scene(meshes, resolution, samples, engine):
        bpy.context.view_layer.update()
        corners = [o.matrix_world @ Vector(c) for o in meshes for c in o.bound_box]
        xs, ys, zs = zip(*[(v.x, v.y, v.z) for v in corners])
        cx, cy, cz = (min(xs)+max(xs))/2, (min(ys)+max(ys))/2, (min(zs)+max(zs))/2
        extent = max(max(xs)-min(xs), max(ys)-min(ys), max(zs)-min(zs)) or 1.0
        for o in meshes:
            if o.parent is None:
                o.location = o.location - Vector((cx, cy, cz))
        center = Vector((0, 0, 0))
        cam_r, cam_h = extent * 1.8, extent * 0.6

        # World — neutral dark background
        w = bpy.context.scene.world or bpy.data.worlds.new("World")
        bpy.context.scene.world = w
        w.use_nodes = True
        for n in list(w.node_tree.nodes):
            w.node_tree.nodes.remove(n)
        bg = w.node_tree.nodes.new("ShaderNodeBackground")
        out = w.node_tree.nodes.new("ShaderNodeOutputWorld")
        bg.inputs[0].default_value = (0.012, 0.013, 0.021, 1)
        bg.inputs[1].default_value = 1.0
        w.node_tree.links.new(bg.outputs[0], out.inputs[0])

        # 3-point area lights
        def _add_light(name, loc, energy, sz=1.0):
            bpy.ops.object.light_add(type="AREA", location=loc)
            lt = bpy.context.object
            lt.name, lt.data.energy, lt.data.size = name, energy, extent * sz
            lt.rotation_euler = (center - Vector(loc)).to_track_quat("-Z", "Y").to_euler()
        _add_light("Key",  (cam_r*0.9, -cam_r*0.7, cam_h*1.7), 1200, 0.9)
        _add_light("Fill", (-cam_r*0.6, -cam_r*0.4, cam_h),     400, 1.2)
        _add_light("Rim",  (0, cam_r*0.9, cam_h*1.3),           600, 0.7)

        # Default neutral material (only for meshes that have none)
        mat = bpy.data.materials.new("eval_default")
        mat.use_nodes = True
        bsdf = mat.node_tree.nodes.get("Principled BSDF")
        if bsdf:
            bsdf.inputs["Base Color"].default_value = (0.72, 0.72, 0.75, 1)
            bsdf.inputs["Roughness"].default_value = 0.65
        for o in meshes:
            if not o.data.materials:
                o.data.materials.append(mat)

        # Camera
        bpy.ops.object.camera_add(location=(cam_r, 0, cam_h))
        cam = bpy.context.object
        cam.name = "EvalCam"
        bpy.context.scene.camera = cam
        cam.data.lens = 50
        cam.data.clip_end = extent * 25

        # Render settings
        sc = bpy.context.scene
        sc.render.resolution_x = sc.render.resolution_y = resolution
        sc.render.image_settings.file_format = "PNG"
        sc.render.image_settings.color_mode = "RGBA"
        sc.render.film_transparent = True
        sc.render.engine = engine

        if engine == "CYCLES":
            gpu_ok = False
            try:
                pr = bpy.context.preferences.addons["cycles"].preferences
                for dev_type in ("OPTIX", "CUDA"):
                    try:
                        pr.compute_device_type = dev_type
                    except TypeError:
                        continue  # device type not in this Blender build
                    # Blender 5.x: refresh_devices(); older Blender: get_devices()
                    if hasattr(pr, "refresh_devices"):
                        pr.refresh_devices()
                    elif hasattr(pr, "get_devices"):
                        pr.get_devices()
                    gpus = [d for d in pr.devices if d.type == dev_type]
                    if gpus:
                        for d in pr.devices:
                            d.use = (d.type == dev_type)
                        sc.cycles.device = "GPU"
                        print(f"[render] GPU: {dev_type} × {len(gpus)}")
                        gpu_ok = True
                        break
                if not gpu_ok:
                    print("[render] No GPU available, falling back to CPU")
                    sc.cycles.device = "CPU"
            except Exception as e:
                print(f"[render] GPU init failed, using CPU: {type(e).__name__}: {e}")
                sc.cycles.device = "CPU"
            sc.cycles.samples = samples
            sc.cycles.use_denoising = True
        else:  # EEVEE (Blender 5.0 may use BLENDER_EEVEE_NEXT internally)
            for attr in ("taa_render_samples", "samples"):
                try:
                    setattr(sc.eevee, attr, samples)
                    break
                except AttributeError:
                    pass

        return cam, center, cam_r, cam_h, float(extent)

    def _render_views(cam, center, cam_r, cam_h, output_dir):
        sc = bpy.context.scene
        n_done = 0
        for frame_idx, fname in REF_VIEWS:
            angle = 2 * math.pi * frame_idx / 40
            cam.location = (cam_r * math.cos(angle),
                            cam_r * math.sin(angle),
                            cam_h)
            cam.rotation_euler = (center - cam.location).to_track_quat("-Z", "Y").to_euler()
            sc.render.filepath = os.path.join(output_dir, fname)
            bpy.ops.render.render(write_still=True)
            n_done += 1
        return n_done

    # ── main ──
    args = _bl_parse()
    os.makedirs(args.output_dir, exist_ok=True)
    log_path = os.path.join(args.output_dir, "render_log.json")
    record = {
        "script":           args.script,
        "samples":          args.samples,
        "resolution":       args.resolution,
        "engine":           args.engine,
        "status":           None,
        "n_meshes":         0,
        "extent":           None,
        "n_views_rendered": 0,
        "latency_s":        None,
        "error":            None,
    }
    t0 = time.time()
    try:
        _clear_scene()
        try:
            meshes = _run_user_script(args.script)
        except Exception as e:
            record["status"] = "ERR_EXEC"
            record["error"] = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
        else:
            record["n_meshes"] = len(meshes)
            if not meshes:
                record["status"] = "ERR_NO_MESH"
                record["error"] = "Script ran but produced no mesh objects"
            else:
                try:
                    cam, center, cam_r, cam_h, extent = _setup_scene(
                        meshes, args.resolution, args.samples, args.engine)
                    record["extent"] = extent
                    n = _render_views(cam, center, cam_r, cam_h, args.output_dir)
                    record["n_views_rendered"] = n
                    record["status"] = "OK" if n == 4 else "ERR_RENDER"
                except Exception as e:
                    record["status"] = "ERR_RENDER"
                    record["error"] = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
    finally:
        record["latency_s"] = round(time.time() - t0, 2)
        with open(log_path, "w") as f:
            json.dump(record, f, indent=2)

    sys.exit(0 if record["status"] == "OK" else 0)  # never bubble error to subprocess


# ╔══════════════════════════════════════════════════════════════════╗
# ║  Orchestrator (normal Python)                                    ║
# ╚══════════════════════════════════════════════════════════════════╝
import argparse
import json
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

EVAL_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_RESULTS_ROOT = EVAL_ROOT / "results"
DEFAULT_BLENDER = os.environ.get("BLENDER", "blender")


def parse_cli():
    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=__doc__.split("\n\n", 1)[0],
    )
    p.add_argument("--model",         required=True,
                   help="Sub-folder name under `results/`.")
    p.add_argument("--results-root",  type=Path, default=DEFAULT_RESULTS_ROOT)
    p.add_argument("--blender",       default=DEFAULT_BLENDER)
    p.add_argument("--samples",       type=int, default=64)
    p.add_argument("--resolution",    type=int, default=512)
    p.add_argument("--engine",        default="CYCLES",
                   choices=["CYCLES", "BLENDER_EEVEE"])
    p.add_argument("--instances",     nargs="*", default=None,
                   help="Limit to these instance folder names.")
    p.add_argument("--workers",       type=int, default=1,
                   help="Parallel Blender subprocesses (1 is safest on single GPU).")
    p.add_argument("--timeout",       type=int, default=240,
                   help="Per-instance Blender timeout in seconds.")
    p.add_argument("--overwrite",     action="store_true")
    return p.parse_args()


def render_one(args, this_script: Path, instance_dir: Path):
    name = instance_dir.name
    gen_script = instance_dir / f"{name}.py"
    renders_dir = instance_dir / "renders"
    log_path = renders_dir / "render_log.json"

    if not gen_script.exists():
        return name, "MISSING", None, 0.0

    if log_path.exists() and not args.overwrite:
        try:
            existing = json.loads(log_path.read_text())
            return name, existing.get("status", "?"), "skip-existing", 0.0
        except Exception:
            pass  # fall through and re-render

    renders_dir.mkdir(parents=True, exist_ok=True)
    cmd = [args.blender, "--background", "--python", str(this_script), "--",
           "--blender-render",
           "--script",     str(gen_script),
           "--output-dir", str(renders_dir),
           "--samples",    str(args.samples),
           "--resolution", str(args.resolution),
           "--engine",     args.engine]
    t0 = time.time()
    try:
        subprocess.run(cmd, capture_output=True, text=True, timeout=args.timeout)
    except subprocess.TimeoutExpired:
        log_path.write_text(json.dumps({
            "script":   str(gen_script),
            "status":   "ERR_TIMEOUT",
            "error":    f"Blender subprocess exceeded {args.timeout}s",
            "n_meshes": 0,
            "n_views_rendered": 0,
            "latency_s": round(time.time() - t0, 2),
        }, indent=2))
        return name, "ERR_TIMEOUT", None, time.time() - t0

    dt = time.time() - t0
    if log_path.exists():
        rec = json.loads(log_path.read_text())
        return name, rec.get("status", "?"), rec.get("error"), dt
    log_path.write_text(json.dumps({
        "script": str(gen_script),
        "status": "ERR_NOLOG",
        "error":  "Blender exited without writing render_log.json (likely segfault)",
        "n_meshes": 0,
        "n_views_rendered": 0,
        "latency_s": round(dt, 2),
    }, indent=2))
    return name, "ERR_NOLOG", None, dt


def main():
    args = parse_cli()
    model_dir = args.results_root / args.model
    if not model_dir.exists():
        raise SystemExit(f"No such model dir: {model_dir}")

    instance_dirs = sorted(d for d in model_dir.iterdir() if d.is_dir())
    if args.instances:
        keep = set(args.instances)
        instance_dirs = [d for d in instance_dirs if d.name in keep]
        if not instance_dirs:
            raise SystemExit(f"No matches for {args.instances}")

    this_script = Path(__file__).resolve()

    print(f"Model:       {args.model}")
    print(f"Instances:   {len(instance_dirs)}")
    print(f"Engine:      {args.engine}, samples={args.samples}, res={args.resolution}")
    print(f"Workers:     {args.workers}")
    print(f"Timeout:     {args.timeout}s per instance")
    print(f"Overwrite:   {args.overwrite}")
    print()

    counts = {}
    t_run = time.time()
    if args.workers <= 1:
        for d in instance_dirs:
            name, status, err, dt = render_one(args, this_script, d)
            counts[status] = counts.get(status, 0) + 1
            tail = f" -- {err.splitlines()[0]}" if err else ""
            print(f"  [{status:<13}] {name}  ({dt:.1f}s){tail}")
    else:
        with ThreadPoolExecutor(max_workers=args.workers) as ex:
            futs = {ex.submit(render_one, args, this_script, d): d
                    for d in instance_dirs}
            for fut in as_completed(futs):
                name, status, err, dt = fut.result()
                counts[status] = counts.get(status, 0) + 1
                tail = f" -- {err.splitlines()[0]}" if err else ""
                print(f"  [{status:<13}] {name}  ({dt:.1f}s){tail}")

    print()
    print(f"Done in {time.time()-t_run:.1f}s")
    for k in sorted(counts):
        print(f"  {k:<13} {counts[k]:>4}")


if __name__ == "__main__":
    main()