gpu_render.py

"""
Rendering entry points.

- ``render_blend_gpu``: Cycles CUDA/OptiX render inside a ZeroGPU worker. The Blender
  subprocess must be spawned *inside* the ``@spaces.GPU`` function so it inherits the
  worker's CUDA devices.
- ``render_blend_local``: CPU Cycles (reliable) or EEVEE via Mesa software GL
  (experimental, slow) in the main container.
"""

from __future__ import annotations

import os
import subprocess

import spaces

from blender_session import BLENDER_BIN, USER_RESOURCES, wrap_with_memory_cap

RENDER_SCRIPT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "render_script.py")

_MAX_CPU_SAMPLES = 256


DEFAULT_CAMERA = {"mode": "scene", "azimuth": 315.0, "elevation": 30.0, "zoom": 1.0, "focal": 50.0}


def _run_blender_render(
    blend_path: str,
    out_path: str,
    engine: str,
    samples: int,
    res_x: int,
    res_y: int,
    camera: dict | None = None,
    extra_env: dict[str, str] | None = None,
    timeout: float = 600.0,
    anim: dict | None = None,
) -> str:
    env = {**os.environ, "BLENDER_USER_RESOURCES": USER_RESOURCES}
    if extra_env:
        env.update(extra_env)
    cam = {**DEFAULT_CAMERA, **(camera or {})}
    cmd = [
        BLENDER_BIN, "-b", blend_path, "--python", RENDER_SCRIPT, "--",
        out_path, engine, str(samples), str(res_x), str(res_y),
        str(cam["mode"]), str(cam["azimuth"]), str(cam["elevation"]),
        str(cam["zoom"]), str(cam["focal"]),
    ]
    if anim is not None:
        cmd += [
            "anim", str(anim.get("fps", 0)),
            str(anim.get("frame_start", 0)), str(anim.get("frame_end", 0)),
        ]
    proc = subprocess.run(
        wrap_with_memory_cap(cmd),
        env=env, capture_output=True, text=True, timeout=timeout, check=False,
    )
    # Blender writes "file.png" possibly with frame suffix; render_script sets exact path.
    if proc.returncode != 0 or not os.path.exists(out_path):
        tail = "\n".join((proc.stdout + "\n" + proc.stderr).splitlines()[-30:])
        raise RuntimeError(f"Render failed (exit {proc.returncode}):\n{tail}")
    for line in proc.stdout.splitlines():
        if line.startswith("RENDER_DONE"):
            print(line)
    return out_path


def estimate_render_duration(
    blend_path: str, out_path: str, samples: int, res_x: int, res_y: int,
    camera: dict | None = None,
) -> int:
    """ZeroGPU duration estimate: Blender startup + scene load + render time."""
    base = 30.0
    pixel_factor = (res_x * res_y) / (1920 * 1080)
    render = 60.0 * (samples / 128.0) * max(pixel_factor, 0.1)
    return int(min(max(base + render, 45.0), 300.0))


@spaces.GPU(duration=estimate_render_duration)
def render_blend_gpu(
    blend_path: str, out_path: str, samples: int, res_x: int, res_y: int,
    camera: dict | None = None,
) -> str:
    """Render a .blend with Cycles on the ZeroGPU worker's GPU (OptiX/CUDA)."""
    return _run_blender_render(blend_path, out_path, "CYCLES_GPU", samples, res_x, res_y, camera)


def render_blend_local(
    blend_path: str, out_path: str, engine: str, samples: int, res_x: int, res_y: int,
    camera: dict | None = None,
) -> str:
    """CPU Cycles or experimental software-GL EEVEE render in the main container."""
    if engine == "CYCLES_CPU":
        samples = min(samples, _MAX_CPU_SAMPLES)
        return _run_blender_render(blend_path, out_path, "CYCLES_CPU", samples, res_x, res_y, camera)
    if engine == "EEVEE":
        env = {"LIBGL_ALWAYS_SOFTWARE": "1", "GALLIUM_DRIVER": "llvmpipe"}
        return _run_blender_render(
            blend_path, out_path, "EEVEE", samples, res_x, res_y, camera, extra_env=env
        )
    raise ValueError(f"Unknown engine: {engine}")


# --- animation rendering -----------------------------------------------------


def estimate_animation_duration(
    blend_path: str, out_path: str, samples: int, res_x: int, res_y: int,
    fps: int = 0, frame_start: int = 0, frame_end: int = 0,
    camera: dict | None = None,
) -> int:
    """
    ZeroGPU duration estimate for animations: startup + per-frame render cost.
    Kept as tight as realistic - over-declaring blocks low-quota visitors.
    """
    frames = max((frame_end - frame_start + 1) if frame_end > frame_start else 72, 1)
    pixel_factor = max((res_x * res_y) / (1280 * 720), 0.1)
    per_frame = 1.0 * (samples / 64.0) * pixel_factor
    return int(min(max(35.0 + frames * per_frame, 60.0), 300.0))


@spaces.GPU(duration=estimate_animation_duration)
def render_animation_gpu(
    blend_path: str, out_path: str, samples: int, res_x: int, res_y: int,
    fps: int = 0, frame_start: int = 0, frame_end: int = 0,
    camera: dict | None = None,
) -> str:
    """Render the scene's animation to MP4 with Cycles on the ZeroGPU worker."""
    return _run_blender_render(
        blend_path, out_path, "CYCLES_GPU", samples, res_x, res_y, camera,
        timeout=340.0,
        anim={"fps": fps, "frame_start": frame_start, "frame_end": frame_end},
    )


def render_animation_local(
    blend_path: str, out_path: str, engine: str, samples: int, res_x: int, res_y: int,
    fps: int = 0, frame_start: int = 0, frame_end: int = 0,
    camera: dict | None = None,
) -> str:
    """Workbench/EEVEE preview animation via software GL in the main container."""
    if engine not in ("WORKBENCH", "EEVEE"):
        raise ValueError(f"Unknown animation preview engine: {engine}")
    env = {"LIBGL_ALWAYS_SOFTWARE": "1", "GALLIUM_DRIVER": "llvmpipe"}
    return _run_blender_render(
        blend_path, out_path, engine, samples, res_x, res_y, camera,
        extra_env=env, timeout=900.0,
        anim={"fps": fps, "frame_start": frame_start, "frame_end": frame_end},
    )