ComfyUI_frontend/src/extensions/core/load3d/Load3dAnimation.ts

import * as THREE from 'three'

import { LGraphNode } from '@/lib/litegraph/src/litegraph'

import { AnimationManager } from './AnimationManager'
import Load3d from './Load3d'
import { type Load3DOptions } from './interfaces'

class Load3dAnimation extends Load3d {
  private animationManager: AnimationManager

  constructor(
    container: Element | HTMLElement,
    options: Load3DOptions = {
      node: {} as LGraphNode
    }
  ) {
    super(container, options)

    this.animationManager = new AnimationManager(
      this.eventManager,
      this.getCurrentModel.bind(this)
    )

    this.animationManager.init()

    this.overrideAnimationLoop()
  }

  private overrideAnimationLoop(): void {
    if (this.animationFrameId !== null) {
      cancelAnimationFrame(this.animationFrameId)
    }

    const animate = () => {
      this.animationFrameId = requestAnimationFrame(animate)

      if (!this.isActive()) {
        return
      }

      const delta = this.clock.getDelta()

      this.animationManager.update(delta)

      this.viewHelperManager.update(delta)

      this.controlsManager.update()

      this.renderMainScene()

      if (this.previewManager.showPreview) {
        this.previewManager.renderPreview()
      }

      this.resetViewport()

      if (this.viewHelperManager.viewHelper.render) {
        this.viewHelperManager.viewHelper.render(this.renderer)
      }
    }

    animate()
  }

  override async loadModel(
    url: string,
    originalFileName?: string
  ): Promise<void> {
    await super.loadModel(url, originalFileName)

    if (this.modelManager.currentModel) {
      this.animationManager.setupModelAnimations(
        this.modelManager.currentModel,
        this.modelManager.originalModel
      )
    }
  }

  override clearModel(): void {
    this.animationManager.dispose()
    super.clearModel()
  }

  updateAnimationList(): void {
    this.animationManager.updateAnimationList()
  }

  setAnimationSpeed(speed: number): void {
    this.animationManager.setAnimationSpeed(speed)
  }

  updateSelectedAnimation(index: number): void {
    this.animationManager.updateSelectedAnimation(index)
  }

  toggleAnimation(play?: boolean): void {
    this.animationManager.toggleAnimation(play)
  }

  get isAnimationPlaying(): boolean {
    return this.animationManager.isAnimationPlaying
  }

  get animationSpeed(): number {
    return this.animationManager.animationSpeed
  }

  get selectedAnimationIndex(): number {
    return this.animationManager.selectedAnimationIndex
  }

  get animationClips(): THREE.AnimationClip[] {
    return this.animationManager.animationClips
  }

  get animationActions(): THREE.AnimationAction[] {
    return this.animationManager.animationActions
  }

  get currentAnimation(): THREE.AnimationMixer | null {
    return this.animationManager.currentAnimation
  }

  override remove(): void {
    this.animationManager.dispose()
    super.remove()
  }
}

export default Load3dAnimation