diff --git a/src/components/nodes/Generate3DNode.tsx b/src/components/nodes/Generate3DNode.tsx new file mode 100644 index 00000000..403877f4 --- /dev/null +++ b/src/components/nodes/Generate3DNode.tsx @@ -0,0 +1,443 @@ +"use client"; + +import React, { useCallback, useState, useEffect, useMemo, useRef } from "react"; +import { Handle, Position, NodeProps, Node, useReactFlow } from "@xyflow/react"; +import { BaseNode } from "./BaseNode"; +import { useCommentNavigation } from "@/hooks/useCommentNavigation"; +import { ModelParameters } from "./ModelParameters"; +import { useWorkflowStore, useProviderApiKeys } from "@/store/workflowStore"; +import { Generate3DNodeData, ProviderType, SelectedModel, ModelInputDef } from "@/types"; +import { ProviderModel, ModelCapability } from "@/lib/providers/types"; +import { ModelSearchDialog } from "@/components/modals/ModelSearchDialog"; +import { useToast } from "@/components/Toast"; + +// Provider badge component +function ProviderBadge({ provider }: { provider: ProviderType }) { + const providerName = provider === "gemini" ? "Gemini" : provider === "replicate" ? "Replicate" : provider === "kie" ? "Kie.ai" : provider === "wavespeed" ? "WaveSpeed" : "fal.ai"; + + return ( + + {provider === "replicate" ? ( + + + + + + ) : provider === "wavespeed" ? ( + + + + + + ) : ( + + + + )} + + ); +} + +// 3D generation capabilities +const THREE_D_CAPABILITIES: ModelCapability[] = ["text-to-3d", "image-to-3d"]; + +type Generate3DNodeType = Node; + +export function Generate3DNode({ id, data, selected }: NodeProps) { + const nodeData = data; + const commentNavigation = useCommentNavigation(id); + const updateNodeData = useWorkflowStore((state) => state.updateNodeData); + const { replicateApiKey, falApiKey, kieApiKey } = useProviderApiKeys(); + const [isBrowseDialogOpen, setIsBrowseDialogOpen] = useState(false); + + // Get the current selected provider (default to fal since most 3D models are there) + const currentProvider: ProviderType = nodeData.selectedModel?.provider || "fal"; + + const handleParametersChange = useCallback( + (parameters: Record) => { + updateNodeData(id, { parameters }); + }, + [id, updateNodeData] + ); + + // Handle inputs loaded from schema + const handleInputsLoaded = useCallback( + (inputs: ModelInputDef[]) => { + updateNodeData(id, { inputSchema: inputs }); + }, + [id, updateNodeData] + ); + + // Handle parameters expand/collapse - resize node height + const { setNodes } = useReactFlow(); + const handleParametersExpandChange = useCallback( + (expanded: boolean, parameterCount: number) => { + const parameterHeight = expanded ? Math.max(parameterCount * 28 + 16, 60) : 0; + const baseHeight = 300; + const newHeight = baseHeight + parameterHeight; + + setNodes((nodes) => + nodes.map((node) => + node.id === id + ? { ...node, style: { ...node.style, height: newHeight } } + : node + ) + ); + }, + [id, setNodes] + ); + + const regenerateNode = useWorkflowStore((state) => state.regenerateNode); + const isRunning = useWorkflowStore((state) => state.isRunning); + + const handleRegenerate = useCallback(() => { + regenerateNode(id); + }, [id, regenerateNode]); + + // Handle model selection from browse dialog + const handleBrowseModelSelect = useCallback((model: ProviderModel) => { + const newSelectedModel: SelectedModel = { + provider: model.provider, + modelId: model.id, + displayName: model.name, + }; + updateNodeData(id, { selectedModel: newSelectedModel, parameters: {} }); + setIsBrowseDialogOpen(false); + }, [id, updateNodeData]); + + // Dynamic title based on selected model + const displayTitle = useMemo(() => { + if (nodeData.selectedModel?.displayName && nodeData.selectedModel.modelId) { + return nodeData.selectedModel.displayName; + } + return "Select 3D model..."; + }, [nodeData.selectedModel?.displayName, nodeData.selectedModel?.modelId]); + + // Provider badge as title prefix + const titlePrefix = useMemo(() => ( + + ), [currentProvider]); + + // Header action element - browse button + const headerAction = useMemo(() => ( + + ), []); + + // Track previous status to detect error transitions + const prevStatusRef = useRef(nodeData.status); + + // Show toast when error occurs + useEffect(() => { + if (nodeData.status === "error" && prevStatusRef.current !== "error" && nodeData.error) { + useToast.getState().show("3D generation failed", "error", true, nodeData.error); + } + prevStatusRef.current = nodeData.status; + }, [nodeData.status, nodeData.error]); + + const handleClear3D = useCallback(() => { + updateNodeData(id, { output3dUrl: null, status: "idle", error: null }); + }, [id, updateNodeData]); + + return ( + <> + updateNodeData(id, { customTitle: title || undefined })} + onCommentChange={(comment) => updateNodeData(id, { comment: comment || undefined })} + onRun={handleRegenerate} + selected={selected} + isExecuting={isRunning} + hasError={nodeData.status === "error"} + headerAction={headerAction} + titlePrefix={titlePrefix} + commentNavigation={commentNavigation ?? undefined} + > + {/* Dynamic input handles based on model schema */} + {nodeData.inputSchema && nodeData.inputSchema.length > 0 ? ( + (() => { + const imageInputs = nodeData.inputSchema!.filter(i => i.type === "image"); + const textInputs = nodeData.inputSchema!.filter(i => i.type === "text"); + + const hasImageInput = imageInputs.length > 0; + const hasTextInput = textInputs.length > 0; + + const handles: Array<{ + id: string; + type: "image" | "text"; + label: string; + schemaName: string | null; + description: string | null; + isPlaceholder: boolean; + }> = []; + + if (hasImageInput) { + imageInputs.forEach((input, index) => { + handles.push({ + id: `image-${index}`, + type: "image", + label: input.label, + schemaName: input.name, + description: input.description || null, + isPlaceholder: false, + }); + }); + } else { + handles.push({ + id: "image", + type: "image", + label: "Image", + schemaName: null, + description: "Not used by this model", + isPlaceholder: true, + }); + } + + if (hasTextInput) { + textInputs.forEach((input, index) => { + handles.push({ + id: `text-${index}`, + type: "text", + label: input.label, + schemaName: input.name, + description: input.description || null, + isPlaceholder: false, + }); + }); + } else { + handles.push({ + id: "text", + type: "text", + label: "Prompt", + schemaName: null, + description: "Not used by this model", + isPlaceholder: true, + }); + } + + const imageHandles = handles.filter(h => h.type === "image"); + const textHandles = handles.filter(h => h.type === "text"); + const totalSlots = imageHandles.length + textHandles.length + 1; + + const renderedHandles = handles.map((handle) => { + const isImage = handle.type === "image"; + const typeIndex = isImage + ? imageHandles.findIndex(h => h.id === handle.id) + : textHandles.findIndex(h => h.id === handle.id); + const adjustedIndex = isImage ? typeIndex : imageHandles.length + 1 + typeIndex; + const topPercent = ((adjustedIndex + 1) / (totalSlots + 1)) * 100; + + return ( + + +
+ {handle.label} +
+ + ); + }); + + return ( + <> + {renderedHandles} + {hasImageInput && ( + + )} + {hasTextInput && ( + + )} + + ); + })() + ) : ( + // Default handles when no schema + <> + +
+ Image +
+ +
+ Prompt +
+ + )} + + {/* 3D output handle */} + + {/* Output label */} +
+ 3D +
+ +
+ {/* Preview area */} + {nodeData.output3dUrl ? ( +
+ + + + 3D Model Generated + Connect to 3D Viewer + {/* Loading overlay for re-generation */} + {nodeData.status === "loading" && ( +
+ + + + +
+ )} + {/* Error overlay */} + {nodeData.status === "error" && ( +
+ + + + Generation failed + See toast for details +
+ )} +
+ +
+
+ ) : ( +
+ {nodeData.status === "loading" ? ( + + + + + ) : nodeData.status === "error" ? ( + + {nodeData.error || "Failed"} + + ) : ( + + Run to generate + + )} +
+ )} + + {/* Model-specific parameters */} + {nodeData.selectedModel?.modelId && ( + + )} +
+ + + {/* Model browser dialog */} + {isBrowseDialogOpen && ( + setIsBrowseDialogOpen(false)} + onModelSelected={handleBrowseModelSelect} + initialCapabilityFilter="3d" + /> + )} + + ); +}