ComfyUI_frontend/src/CurveEditor.ts

import type { Point, Rect } from "./interfaces"
import { clamp, LGraphCanvas } from "./litegraph"
import { distance } from "./measure"

//used by some widgets to render a curve editor

export class CurveEditor {
    points: Point[]
    selected: number
    nearest: number
    size: Rect
    must_update: boolean
    margin: number
    _nearest: number

    constructor(points: Point[]) {
        this.points = points
        this.selected = -1
        this.nearest = -1
        this.size = null //stores last size used
        this.must_update = true
        this.margin = 5
    }

    static sampleCurve(f: number, points: Point[]): number {
        if (!points)
            return
        for (let i = 0; i < points.length - 1; ++i) {
            const p = points[i]
            const pn = points[i + 1]
            if (pn[0] < f)
                continue
            const r = (pn[0] - p[0])
            if (Math.abs(r) < 0.00001)
                return p[1]
            const local_f = (f - p[0]) / r
            return p[1] * (1.0 - local_f) + pn[1] * local_f
        }
        return 0
    }

    draw(ctx: CanvasRenderingContext2D, size: Rect, graphcanvas?: LGraphCanvas, background_color?: string, line_color?: string, inactive = false): void {
        const points = this.points
        if (!points)
            return
        this.size = size
        const w = size[0] - this.margin * 2
        const h = size[1] - this.margin * 2

        line_color = line_color || "#666"

        ctx.save()
        ctx.translate(this.margin, this.margin)

        if (background_color) {
            ctx.fillStyle = "#111"
            ctx.fillRect(0, 0, w, h)
            ctx.fillStyle = "#222"
            ctx.fillRect(w * 0.5, 0, 1, h)
            ctx.strokeStyle = "#333"
            ctx.strokeRect(0, 0, w, h)
        }
        ctx.strokeStyle = line_color
        if (inactive)
            ctx.globalAlpha = 0.5
        ctx.beginPath()
        for (let i = 0; i < points.length; ++i) {
            const p = points[i]
            ctx.lineTo(p[0] * w, (1.0 - p[1]) * h)
        }
        ctx.stroke()
        ctx.globalAlpha = 1
        if (!inactive)
            for (let i = 0; i < points.length; ++i) {
                const p = points[i]
                ctx.fillStyle = this.selected == i ? "#FFF" : (this.nearest == i ? "#DDD" : "#AAA")
                ctx.beginPath()
                ctx.arc(p[0] * w, (1.0 - p[1]) * h, 2, 0, Math.PI * 2)
                ctx.fill()
            }
        ctx.restore()
    }

    //localpos is mouse in curve editor space
    onMouseDown(localpos: Point, graphcanvas: LGraphCanvas): boolean {
        const points = this.points
        if (!points)
            return
        if (localpos[1] < 0)
            return

        //this.captureInput(true);
        const w = this.size[0] - this.margin * 2
        const h = this.size[1] - this.margin * 2
        const x = localpos[0] - this.margin
        const y = localpos[1] - this.margin
        const pos: Point = [x, y]
        const max_dist = 30 / graphcanvas.ds.scale
        //search closer one
        this.selected = this.getCloserPoint(pos, max_dist)
        //create one
        if (this.selected == -1) {
            const point: Point = [x / w, 1 - y / h]
            points.push(point)
            points.sort(function (a, b) { return a[0] - b[0] })
            this.selected = points.indexOf(point)
            this.must_update = true
        }
        if (this.selected != -1)
            return true
    }

    onMouseMove(localpos: Point, graphcanvas: LGraphCanvas): void {
        const points = this.points
        if (!points)
            return
        const s = this.selected
        if (s < 0)
            return
        const x = (localpos[0] - this.margin) / (this.size[0] - this.margin * 2)
        const y = (localpos[1] - this.margin) / (this.size[1] - this.margin * 2)
        const curvepos: Point = [(localpos[0] - this.margin), (localpos[1] - this.margin)]
        const max_dist = 30 / graphcanvas.ds.scale
        this._nearest = this.getCloserPoint(curvepos, max_dist)
        const point = points[s]
        if (point) {
            const is_edge_point = s == 0 || s == points.length - 1
            if (!is_edge_point && (localpos[0] < -10 || localpos[0] > this.size[0] + 10 || localpos[1] < -10 || localpos[1] > this.size[1] + 10)) {
                points.splice(s, 1)
                this.selected = -1
                return
            }
            if (!is_edge_point) //not edges
                point[0] = clamp(x, 0, 1)
            else
                point[0] = s == 0 ? 0 : 1
            point[1] = 1.0 - clamp(y, 0, 1)
            points.sort(function (a, b) { return a[0] - b[0] })
            this.selected = points.indexOf(point)
            this.must_update = true
        }
    }

    // Former params: localpos, graphcanvas
    onMouseUp(): boolean {
        this.selected = -1
        return false
    }

    getCloserPoint(pos: Point, max_dist: number): number {
        const points = this.points
        if (!points)
            return -1
        max_dist = max_dist || 30
        const w = (this.size[0] - this.margin * 2)
        const h = (this.size[1] - this.margin * 2)
        const num = points.length
        const p2: Point = [0, 0]
        let min_dist = 1000000
        let closest = -1
        for (let i = 0; i < num; ++i) {
            const p = points[i]
            p2[0] = p[0] * w
            p2[1] = (1.0 - p[1]) * h
            const dist = distance(pos, p2)
            if (dist > min_dist || dist > max_dist)
                continue
            closest = i
            min_dist = dist
        }
        return closest
    }
}