music-video-gen/flow-state/src/engine/Compositor.js

import * as THREE from 'three';
import { makePassMaterial } from './Renderer.js';
import {
    BLEND_FRAG, FEEDBACK_FRAG, BRIGHT_FRAG, BLUR_FRAG, COMPOSITE_FRAG, COPY_FRAG,
    BLEND_MODE_IDS,
} from './passes.js';

const DEFAULT_POST = {
    bloom: 0.35,
    bloomThreshold: 0.6,
    bloomKnee: 0.3,
    chroma: 0.15,
    // Grain defaults to absent. What it looks like when a look does ask for it
    // is described by the four fields below — see look/grain.js.
    grain: 0,
    grainScale: 1,
    grainRate: 1,
    grainMask: 0,
    grainChroma: 0,
    vignette: 0.35,
    contrast: 1.05,
    saturation: 1.1,
    lift: 0.0,
    exposure: 1.0,
};

const DEFAULT_FEEDBACK = {
    amount: 0.0,
    decay: 0.9,
    zoom: 0.995,
    rotate: 0.0,
};

/**
 * The layer stack. Layers render into their own target, then blend into an
 * accumulator; the result goes through feedback and the post chain.
 *
 * Every target is explicitly cleared on allocation and on reset, because
 * inheriting stale GPU memory is exactly the kind of thing that makes an export
 * differ from a preview.
 */
export class Compositor {
    constructor(renderer, { width, height } = {}) {
        this.renderer = renderer;
        this.width = width || renderer.width;
        this.height = height || renderer.height;

        this.layers = [];
        this.post = { ...DEFAULT_POST };
        this.feedback = { ...DEFAULT_FEEDBACK };
        this.fade = 1;
        this.soloIndex = -1;    // debug: render one layer alone
        this.postEnabled = true;
        this._primed = new WeakSet();

        this._buildTargets();
        this._buildMaterials();
    }

    _buildTargets() {
        const r = this.renderer;
        const w = this.width, h = this.height;
        const bw = Math.max(1, Math.floor(w / 2));
        const bh = Math.max(1, Math.floor(h / 2));

        this.layerTarget = r.createTarget(w, h, { depth: true });
        this.accumA = r.createTarget(w, h);
        this.accumB = r.createTarget(w, h);
        this.historyA = r.createTarget(w, h, { float: true });
        this.historyB = r.createTarget(w, h, { float: true });
        this.bloomA = r.createTarget(bw, bh);
        this.bloomB = r.createTarget(bw, bh);
        this.outputTarget = r.createTarget(w, h);
    }

    _buildMaterials() {
        this.blendMaterial = makePassMaterial(BLEND_FRAG, {
            u_base: { value: null },
            u_src: { value: null },
            u_amount: { value: 1 },
            u_mode: { value: 0 },
        });
        this.feedbackMaterial = makePassMaterial(FEEDBACK_FRAG, {
            u_current: { value: null },
            u_history: { value: null },
            u_decay: { value: 0.9 },
            u_amount: { value: 0 },
            u_zoom: { value: 0.995 },
            u_rotate: { value: 0 },
            u_aspect: { value: 1 },
        });
        this.brightMaterial = makePassMaterial(BRIGHT_FRAG, {
            u_tex: { value: null },
            u_threshold: { value: 0.6 },
            u_knee: { value: 0.3 },
        });
        this.blurMaterial = makePassMaterial(BLUR_FRAG, {
            u_tex: { value: null },
            u_direction: { value: new THREE.Vector2(0, 0) },
        });
        this.compositeMaterial = makePassMaterial(COMPOSITE_FRAG, {
            u_tex: { value: null },
            u_bloom: { value: null },
            u_bloomAmount: { value: 0 },
            u_chroma: { value: 0 },
            u_grain: { value: 0 },
            u_grainScale: { value: 1 },
            u_grainRate: { value: 1 },
            u_grainMask: { value: 0 },
            u_grainChroma: { value: 0 },
            u_vignette: { value: 0 },
            u_contrast: { value: 1 },
            u_saturation: { value: 1 },
            u_lift: { value: 0 },
            u_exposure: { value: 1 },
            u_fade: { value: 1 },
            u_frame: { value: 0 },
            u_resolution: { value: new THREE.Vector2(1, 1) },
        });
        this.copyMaterial = makePassMaterial(COPY_FRAG, { u_tex: { value: null } });
    }

    setSize(width, height) {
        if (width === this.width && height === this.height) return;
        this.width = width;
        this.height = height;
        this.disposeTargets();
        this._buildTargets();
    }

    /**
     * The compositor does NOT own its layers and never disposes them — the arc
     * driver caches Layer instances across sections and swaps them in and out
     * every crossfade, and disposing on removal would destroy shaders that are
     * about to be reused (and recompile them on the way back).
     */
    setLayers(layers) {
        for (const layer of layers) {
            if (!this._primed.has(layer)) {
                this._primeLayer(layer);
                this._primed.add(layer);
            }
        }
        this.layers = layers;
        return this;
    }

    /**
     * Compile layers that are not on screen yet, so the frame they first appear
     * on does not pay for the link. Used by the arc driver's prewarm pass.
     */
    primeLayers(layers) {
        for (const layer of layers) {
            if (this._primed.has(layer)) continue;
            this._primeLayer(layer);
            this._primed.add(layer);
        }
        return this;
    }

    /**
     * Force a layer's shader program to finish linking before it is used for real.
     *
     * three.js links programs through KHR_parallel_shader_compile, so the first
     * draws after a material is created can run against a program that is not
     * ready and produce wrong output. Measured on the heaviest scene in the
     * library, the first TEN frames rendered differently from every later render
     * of the same frames. Preview hides this — the frames go by and the next pass
     * is correct — but an export renders each frame exactly once, so those frames
     * would ship broken.
     *
     * Rendering a throwaway frame and reading it back is NOT sufficient: measured,
     * it left 3-5 frames still wrong. WebGLRenderer.compile() is the API that
     * actually waits for the link, and it clears the problem completely.
     */
    _primeLayer(layer) {
        try {
            if (layer.material) this.renderer.compileMaterial(layer.material);
            else if (layer.scene && layer.camera) this.renderer.compileScene(layer.scene, layer.camera);
        } catch (err) {
            console.warn('[compositor] priming failed for', layer.module && layer.module.name, err);
        }
    }

    /**
     * Bring the whole chain to a state where the next rendered frame is correct.
     *
     * Shader programs link asynchronously (KHR_parallel_shader_compile), and until
     * they are ready a draw produces wrong output. compile() covers the programs;
     * the discarded frame covers everything else that is lazily created on first
     * use. Cheap, and it converts "the first few frames may be wrong" into "the
     * first few frames were thrown away".
     *
     * The exporter calls this before encoding anything, because an export renders
     * each frame exactly once and has no second chance to get frame 0 right.
     */
    prime(ctx = null) {
        const materials = [
            this.blendMaterial, this.feedbackMaterial, this.brightMaterial,
            this.blurMaterial, this.compositeMaterial, this.copyMaterial,
        ];
        for (const material of materials) {
            try { this.renderer.compileMaterial(material); } catch { /* non-fatal */ }
        }
        for (const layer of this.layers) this._primeLayer(layer);

        if (ctx) {
            try { this.render(ctx); } catch { /* non-fatal */ }
        }
        this.reset();
    }

    setPost(post) {
        this.post = { ...this.post, ...post };
        return this;
    }

    setFeedback(feedback) {
        this.feedback = { ...this.feedback, ...feedback };
        return this;
    }

    /**
     * Wipe all history. Called on seek and before an export run so a render never
     * depends on what was on screen beforehand.
     */
    reset() {
        const r = this.renderer;
        [this.accumA, this.accumB, this.historyA, this.historyB,
         this.bloomA, this.bloomB, this.layerTarget, this.outputTarget]
            .forEach((t) => r.clear(t));
    }

    /**
     * Render one frame. Returns the target holding the finished image, so the
     * caller decides whether it goes to the canvas or to the encoder.
     */
    render(ctx) {
        const r = this.renderer;
        const { timeline, features } = ctx;

        r.clear(this.accumA);
        let accum = this.accumA;
        let spare = this.accumB;

        const active = this.soloIndex >= 0
            ? this.layers.slice(this.soloIndex, this.soloIndex + 1)
            : this.layers;

        for (const layer of active) {
            if (layer.opacity <= 0.001) continue;

            r.clear(this.layerTarget);
            layer.render(r, this.layerTarget, {
                timeline,
                features,
                prevTexture: this.historyA.texture,
            });

            const bu = this.blendMaterial.uniforms;
            bu.u_base.value = accum.texture;
            bu.u_src.value = this.layerTarget.texture;
            bu.u_amount.value = 1.0;   // layer opacity already applied in-shader
            bu.u_mode.value = BLEND_MODE_IDS[layer.blend] ?? 0;
            r.blit(this.blendMaterial, spare);

            const t = accum; accum = spare; spare = t;
        }

        // --- feedback -------------------------------------------------------
        let composited = accum;
        if (this.feedback.amount > 0.001) {
            const fu = this.feedbackMaterial.uniforms;
            fu.u_current.value = accum.texture;
            fu.u_history.value = this.historyA.texture;
            fu.u_decay.value = Math.min(0.99, this.feedback.decay);
            fu.u_amount.value = this.feedback.amount;
            fu.u_zoom.value = this.feedback.zoom;
            fu.u_rotate.value = this.feedback.rotate;
            fu.u_aspect.value = this.width / this.height;
            r.blit(this.feedbackMaterial, this.historyB);
            composited = this.historyB;
            const t = this.historyA; this.historyA = this.historyB; this.historyB = t;
        } else {
            // Keep history tracking the image even when feedback is off, so
            // enabling it mid-track doesn't pop from black.
            this.copyMaterial.uniforms.u_tex.value = accum.texture;
            r.blit(this.copyMaterial, this.historyA);
        }

        if (!this.postEnabled) {
            this.copyMaterial.uniforms.u_tex.value = composited.texture;
            r.blit(this.copyMaterial, this.outputTarget);
            return this.outputTarget;
        }

        // --- bloom ----------------------------------------------------------
        const p = this.post;
        if (p.bloom > 0.001) {
            this.brightMaterial.uniforms.u_tex.value = composited.texture;
            this.brightMaterial.uniforms.u_threshold.value = p.bloomThreshold;
            this.brightMaterial.uniforms.u_knee.value = p.bloomKnee;
            r.blit(this.brightMaterial, this.bloomA);

            const bw = this.bloomA.width, bh = this.bloomA.height;
            for (let i = 0; i < 2; i++) {
                this.blurMaterial.uniforms.u_tex.value = this.bloomA.texture;
                this.blurMaterial.uniforms.u_direction.value.set((1 + i) / bw, 0);
                r.blit(this.blurMaterial, this.bloomB);

                this.blurMaterial.uniforms.u_tex.value = this.bloomB.texture;
                this.blurMaterial.uniforms.u_direction.value.set(0, (1 + i) / bh);
                r.blit(this.blurMaterial, this.bloomA);
            }
        } else {
            r.clear(this.bloomA);
        }

        // --- final grade ----------------------------------------------------
        const cu = this.compositeMaterial.uniforms;
        cu.u_tex.value = composited.texture;
        cu.u_bloom.value = this.bloomA.texture;
        cu.u_bloomAmount.value = p.bloom;
        cu.u_chroma.value = p.chroma;
        cu.u_grain.value = p.grain;
        cu.u_grainScale.value = p.grainScale;
        cu.u_grainRate.value = p.grainRate;
        cu.u_grainMask.value = p.grainMask;
        cu.u_grainChroma.value = p.grainChroma;
        cu.u_vignette.value = p.vignette;
        cu.u_contrast.value = p.contrast;
        cu.u_saturation.value = p.saturation;
        cu.u_lift.value = p.lift;
        cu.u_exposure.value = p.exposure;
        cu.u_fade.value = this.fade;
        cu.u_frame.value = timeline.frame;
        cu.u_resolution.value.set(this.width, this.height);
        r.blit(this.compositeMaterial, this.outputTarget);

        return this.outputTarget;
    }

    /** Present a finished target to the canvas. */
    present(target) {
        this.copyMaterial.uniforms.u_tex.value = target.texture;
        this.renderer.blit(this.copyMaterial, null);
    }

    disposeTargets() {
        [this.layerTarget, this.accumA, this.accumB, this.historyA, this.historyB,
         this.bloomA, this.bloomB, this.outputTarget].forEach((t) => t && t.dispose());
    }

    dispose() {
        this.layers = [];       // owned elsewhere; see setLayers
        this.disposeTargets();
    }
}