popiart-node-canvas/src/utils/highDefinitionVideoNodes.ts

import type { Connection, XYPosition } from "@xyflow/react";
import type { ProviderModel } from "@/lib/providers/types";
import type { SelectedModel, SmartVideoNodeData } from "@/types";
import { POPI_VIDEO_SUBTYPE_ENHANCE } from "@/constants/generationTask";
import { SINGLE_INPUT_HANDLE_ID, SINGLE_OUTPUT_HANDLE_ID } from "@/utils/nodeHandles";
import { toSelectedModel } from "@/utils/selectedModel";

export const HIGH_DEFINITION_VIDEO_BATCH_SIZE = 1;

export const HIGH_DEFINITION_VIDEO_FALLBACK_MODEL: SelectedModel = {
  provider: "popiserver",
  modelId: "",
  displayName: "画质增强",
  capabilities: ["image-to-video"],
  isSupportVideos: true,
  type: 2,
  subType: POPI_VIDEO_SUBTYPE_ENHANCE,
  subTypes: [POPI_VIDEO_SUBTYPE_ENHANCE],
};

function modelSubTypes(model: ProviderModel): number[] {
  const subTypes = model.metadata?.subTypes;
  if (Array.isArray(subTypes)) {
    return subTypes.map((item) => Number(item)).filter((item) => Number.isFinite(item));
  }
  const subType = Number(model.metadata?.subType);
  return Number.isFinite(subType) ? [subType] : [];
}

export function isHighDefinitionVideoModel(model: ProviderModel): boolean {
  return modelSubTypes(model).includes(POPI_VIDEO_SUBTYPE_ENHANCE);
}

export function chooseHighDefinitionVideoModel(models: ProviderModel[]): SelectedModel {
  const model = models.find(isHighDefinitionVideoModel);
  return model ? toSelectedModel(model) : HIGH_DEFINITION_VIDEO_FALLBACK_MODEL;
}

export function isHighDefinitionDerivedVideoNode(data: unknown): boolean {
  const record = data && typeof data === "object" ? data as { derivedVideo?: unknown } : {};
  const derivedVideo = record.derivedVideo && typeof record.derivedVideo === "object"
    ? record.derivedVideo as { operation?: unknown }
    : null;
  return derivedVideo?.operation === "highDefinition";
}

export interface CreateHighDefinitionVideoNodeOptions {
  sourceNodeId: string;
  position: XYPosition;
  dimensions?: { width: number; height: number } | null;
  selectedModel: SelectedModel;
  addNode: (type: "smartVideo", position: XYPosition, initialData?: Partial<SmartVideoNodeData>) => string;
  onConnect: (connection: Connection) => void;
  selectSingleNode?: (nodeId: string) => void;
}

export function createHighDefinitionVideoNode({
  sourceNodeId,
  position,
  dimensions = null,
  selectedModel,
  addNode,
  onConnect,
  selectSingleNode,
}: CreateHighDefinitionVideoNodeOptions): string {
  const nodeId = addNode("smartVideo", position, {
    video: null,
    videoRef: undefined,
    assetId: undefined,
    previewVideoPoster: undefined,
    filename: `high-definition-${Date.now()}.mp4`,
    duration: null,
    dimensions,
    format: null,
    inputImages: [],
    inputPrompt: "",
    outputVideo: null,
    outputVideoCover: undefined,
    outputVideoRef: undefined,
    selectedModel,
    batchSize: HIGH_DEFINITION_VIDEO_BATCH_SIZE,
    referenceSubjectList: [],
    parameters: {},
    extraTaskParams: {},
    subType: POPI_VIDEO_SUBTYPE_ENHANCE,
    customTitle: "高清",
    derivedVideo: {
      sourceNodeId,
      operation: "highDefinition",
    },
    status: "idle",
    error: null,
    videoHistory: [],
    selectedVideoHistoryIndex: 0,
  });

  onConnect({
    source: sourceNodeId,
    sourceHandle: SINGLE_OUTPUT_HANDLE_ID,
    target: nodeId,
    targetHandle: SINGLE_INPUT_HANDLE_ID,
  });

  selectSingleNode?.(nodeId);
  return nodeId;
}