HartOMat/render-worker/scripts/_blender_gpu.py

"""GPU activation and engine configuration helpers for Blender headless renders.

activate_gpu() must be called BEFORE open_mainfile / Cycles engine initialisation
so that the CUDA/OptiX kernel is compiled with the correct compute_device_type.
"""
from __future__ import annotations

import os
import sys


def activate_gpu(cycles_device: str = "auto") -> str | None:
    """Probe for GPU compute devices and activate them.

    Args:
        cycles_device: "auto" | "gpu" | "cpu"

    Returns:
        Device type string (e.g. "OPTIX", "CUDA") if GPU was activated,
        or None if CPU-only.
    """
    if cycles_device == "cpu":
        return None
    import bpy  # type: ignore[import]
    try:
        cprefs = bpy.context.preferences.addons['cycles'].preferences
        for dt in ('OPTIX', 'CUDA', 'HIP', 'ONEAPI'):
            try:
                cprefs.compute_device_type = dt
                cprefs.get_devices()
                gpu = [d for d in cprefs.devices if d.type != 'CPU']
                if gpu:
                    for d in cprefs.devices:
                        d.use = (d.type != 'CPU')
                    print(f"[blender_render] early GPU activation: {dt}, "
                          f"devices={[(d.name, d.type) for d in gpu]}", flush=True)
                    return dt
            except Exception as e:
                print(f"[blender_render] {dt} not available: {e}", flush=True)
    except Exception as e:
        print(f"[blender_render] early GPU probe failed: {e}", flush=True)
    return None


def configure_engine(
    scene,
    engine: str,
    samples: int,
    cycles_device: str,
    early_gpu_type: str | None,
    noise_threshold_arg: str = "",
    denoiser_arg: str = "",
    denoising_input_passes_arg: str = "",
    denoising_prefilter_arg: str = "",
    denoising_quality_arg: str = "",
    denoising_use_gpu_arg: str = "",
) -> str:
    """Configure the Blender render engine (EEVEE or Cycles) on *scene*.

    Returns the effective engine name ("eevee" or "cycles").
    Exits with code 2 if GPU required but unavailable (CYCLES_DEVICE=gpu env var).
    """
    if engine == "eevee":
        set_ok = False
        for eevee_id in ('BLENDER_EEVEE', 'BLENDER_EEVEE_NEXT'):
            try:
                scene.render.engine = eevee_id
                set_ok = True
                print(f"[blender_render] EEVEE engine id: {eevee_id}")
                break
            except TypeError:
                continue
        if not set_ok:
            print("[blender_render] WARNING: could not set EEVEE engine – falling back to Cycles")
            engine = "cycles"
        if engine == "eevee":
            for attr in ('taa_render_samples', 'samples'):
                try:
                    import bpy  # type: ignore[import]
                    setattr(scene.eevee, attr, samples)
                    print(f"[blender_render] EEVEE samples: scene.eevee.{attr}={samples}")
                    break
                except AttributeError:
                    continue

    if engine != "eevee":
        gpu_type_found = activate_gpu(cycles_device) or early_gpu_type
        scene.render.engine = 'CYCLES'
        if gpu_type_found:
            scene.cycles.device = 'GPU'
            activate_gpu(cycles_device)
            print(f"[blender_render] Cycles GPU ({gpu_type_found}), samples={samples}", flush=True)
            print(f"RENDER_DEVICE_USED: engine=CYCLES device=GPU compute_type={gpu_type_found}", flush=True)
        else:
            scene.cycles.device = 'CPU'
            print(f"[blender_render] WARNING: GPU not found — falling back to CPU, samples={samples}", flush=True)
            print("RENDER_DEVICE_USED: engine=CYCLES device=CPU compute_type=NONE (fallback)", flush=True)
            if os.environ.get("CYCLES_DEVICE", "auto").lower() == "gpu":
                print("GPU_REQUIRED_BUT_CPU_USED: strict mode active (CYCLES_DEVICE=gpu)", flush=True)
                sys.exit(2)

        scene.cycles.samples       = samples
        scene.cycles.use_denoising = True
        scene.cycles.denoiser = denoiser_arg if denoiser_arg else 'OPENIMAGEDENOISE'
        if denoising_input_passes_arg:
            try: scene.cycles.denoising_input_passes = denoising_input_passes_arg
            except Exception: pass
        if denoising_prefilter_arg:
            try: scene.cycles.denoising_prefilter = denoising_prefilter_arg
            except Exception: pass
        if denoising_quality_arg:
            try: scene.cycles.denoising_quality = denoising_quality_arg
            except Exception: pass
        if denoising_use_gpu_arg:
            try: scene.cycles.denoising_use_gpu = (denoising_use_gpu_arg == "1")
            except AttributeError: pass
        if noise_threshold_arg:
            scene.cycles.use_adaptive_sampling = True
            scene.cycles.adaptive_threshold = float(noise_threshold_arg)

    return engine