"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, saveNanoBananaDefaults, useProviderApiKeys } from "@/store/workflowStore"; import { deduplicatedFetch } from "@/utils/deduplicatedFetch"; import { NanoBananaNodeData, AspectRatio, Resolution, ModelType, ProviderType, SelectedModel, ModelInputDef } from "@/types"; import { ProviderModel, ModelCapability } from "@/lib/providers/types"; import { ModelSearchDialog } from "@/components/modals/ModelSearchDialog"; import { useToast } from "@/components/Toast"; import { getImageDimensions, calculateNodeSizePreservingHeight } from "@/utils/nodeDimensions"; // Provider badge component - shows provider icon for all providers function ProviderBadge({ provider }: { provider: ProviderType }) { const providerName = provider === "gemini" ? "Gemini" : provider === "replicate" ? "Replicate" : provider === "kie" ? "Kie.ai" : "fal.ai"; return ( {provider === "gemini" ? ( ) : provider === "replicate" ? ( ) : provider === "kie" ? ( ) : ( )} ); } // All 10 aspect ratios supported by both models const ASPECT_RATIOS: AspectRatio[] = ["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"]; // Resolutions only for Nano Banana Pro (gemini-3-pro-image-preview) const RESOLUTIONS: Resolution[] = ["1K", "2K", "4K"]; // Hardcoded Gemini image models (always available) const GEMINI_IMAGE_MODELS: { value: ModelType; label: string }[] = [ { value: "nano-banana", label: "Nano Banana" }, { value: "nano-banana-pro", label: "Nano Banana Pro" }, ]; // Image generation capabilities const IMAGE_CAPABILITIES: ModelCapability[] = ["text-to-image", "image-to-image"]; type NanoBananaNodeType = Node; export function GenerateImageNode({ id, data, selected }: NodeProps) { const nodeData = data; const commentNavigation = useCommentNavigation(id); const updateNodeData = useWorkflowStore((state) => state.updateNodeData); const generationsPath = useWorkflowStore((state) => state.generationsPath); // Use stable selector for API keys to prevent unnecessary re-fetches const { replicateApiKey, falApiKey, kieApiKey, replicateEnabled, kieEnabled } = useProviderApiKeys(); const [isLoadingCarouselImage, setIsLoadingCarouselImage] = useState(false); const [externalModels, setExternalModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); const [modelsFetchError, setModelsFetchError] = useState(null); const [isBrowseDialogOpen, setIsBrowseDialogOpen] = useState(false); // Get the current selected provider (default to gemini) const currentProvider: ProviderType = nodeData.selectedModel?.provider || "gemini"; // Get enabled providers const enabledProviders = useMemo(() => { const providers: { id: ProviderType; name: string }[] = []; // Gemini is always available providers.push({ id: "gemini", name: "Gemini" }); // fal.ai is always available (works without key but rate limited) providers.push({ id: "fal", name: "fal.ai" }); // Add Replicate if configured if (replicateEnabled && replicateApiKey) { providers.push({ id: "replicate", name: "Replicate" }); } // Add Kie.ai if configured if (kieEnabled && kieApiKey) { providers.push({ id: "kie", name: "Kie.ai" }); } return providers; }, [replicateEnabled, replicateApiKey, kieEnabled, kieApiKey]); // Check if external providers (Replicate/Fal) are enabled // fal.ai is always available (works without key but rate limited) const hasExternalProviders = useMemo(() => { const hasReplicate = replicateEnabled && replicateApiKey; // fal.ai is always available return !!(hasReplicate || true); }, [replicateEnabled, replicateApiKey]); const isGeminiOnly = !hasExternalProviders; // Migrate legacy data: derive selectedModel from model field if missing useEffect(() => { if (nodeData.model && !nodeData.selectedModel) { const displayName = nodeData.model === "nano-banana" ? "Nano Banana" : "Nano Banana Pro"; const newSelectedModel: SelectedModel = { provider: "gemini", modelId: nodeData.model, displayName, }; updateNodeData(id, { selectedModel: newSelectedModel }); } }, [id, nodeData.model, nodeData.selectedModel, updateNodeData]); // Fetch models from external providers when provider changes const fetchModels = useCallback(async () => { if (currentProvider === "gemini") { setExternalModels([]); setModelsFetchError(null); return; } setIsLoadingModels(true); setModelsFetchError(null); try { const capabilities = IMAGE_CAPABILITIES.join(","); const headers: HeadersInit = {}; if (replicateApiKey) { headers["X-Replicate-Key"] = replicateApiKey; } if (falApiKey) { headers["X-Fal-Key"] = falApiKey; } if (kieApiKey) { headers["X-Kie-Key"] = kieApiKey; } const response = await deduplicatedFetch(`/api/models?provider=${currentProvider}&capabilities=${capabilities}`, { headers }); if (response.ok) { const data = await response.json(); setExternalModels(data.models || []); setModelsFetchError(null); } else { const errorData = await response.json().catch(() => ({})); const errorMsg = errorData.error || `Failed to load models (${response.status})`; setExternalModels([]); setModelsFetchError( currentProvider === "replicate" && response.status === 401 ? "Invalid Replicate API key. Check your settings." : errorMsg ); } } catch (error) { console.error("Failed to fetch models:", error); setExternalModels([]); setModelsFetchError("Failed to load models. Check your connection."); } finally { setIsLoadingModels(false); } }, [currentProvider, replicateApiKey, falApiKey, kieApiKey]); useEffect(() => { fetchModels(); }, [fetchModels]); // Handle provider change const handleProviderChange = useCallback( (e: React.ChangeEvent) => { const provider = e.target.value as ProviderType; if (provider === "gemini") { // Reset to Gemini default const newSelectedModel: SelectedModel = { provider: "gemini", modelId: nodeData.model || "nano-banana-pro", displayName: GEMINI_IMAGE_MODELS.find(m => m.value === (nodeData.model || "nano-banana-pro"))?.label || "Nano Banana Pro", }; // Clear parameters when switching providers (different providers have different schemas) updateNodeData(id, { selectedModel: newSelectedModel, parameters: {} }); } else { // Set placeholder for external provider const newSelectedModel: SelectedModel = { provider, modelId: "", displayName: "Select model...", }; // Clear parameters when switching providers updateNodeData(id, { selectedModel: newSelectedModel, parameters: {} }); } }, [id, nodeData.model, updateNodeData] ); // Handle model change for external providers const handleExternalModelChange = useCallback( (e: React.ChangeEvent) => { const modelId = e.target.value; const model = externalModels.find(m => m.id === modelId); if (model) { const newSelectedModel: SelectedModel = { provider: currentProvider, modelId: model.id, displayName: model.name, }; // Clear parameters when changing models (different models have different schemas) updateNodeData(id, { selectedModel: newSelectedModel, parameters: {} }); } }, [id, currentProvider, externalModels, updateNodeData] ); const handleAspectRatioChange = useCallback( (e: React.ChangeEvent) => { const aspectRatio = e.target.value as AspectRatio; updateNodeData(id, { aspectRatio }); saveNanoBananaDefaults({ aspectRatio }); }, [id, updateNodeData] ); const handleResolutionChange = useCallback( (e: React.ChangeEvent) => { const resolution = e.target.value as Resolution; updateNodeData(id, { resolution }); saveNanoBananaDefaults({ resolution }); }, [id, updateNodeData] ); const handleModelChange = useCallback( (e: React.ChangeEvent) => { const model = e.target.value as ModelType; updateNodeData(id, { model }); saveNanoBananaDefaults({ model }); // Also update selectedModel for consistency const newSelectedModel: SelectedModel = { provider: "gemini", modelId: model, displayName: GEMINI_IMAGE_MODELS.find(m => m.value === model)?.label || model, }; updateNodeData(id, { selectedModel: newSelectedModel }); }, [id, updateNodeData] ); const handleGoogleSearchToggle = useCallback( (e: React.ChangeEvent) => { const useGoogleSearch = e.target.checked; updateNodeData(id, { useGoogleSearch }); saveNanoBananaDefaults({ useGoogleSearch }); }, [id, updateNodeData] ); 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) => { // Each parameter row is ~24px, plus some padding const parameterHeight = expanded ? Math.max(parameterCount * 28 + 16, 60) : 0; const baseHeight = 300; // Default node height const newHeight = baseHeight + parameterHeight; setNodes((nodes) => nodes.map((node) => node.id === id ? { ...node, style: { ...node.style, height: newHeight } } : node ) ); }, [id, setNodes] ); const handleClearImage = useCallback(() => { updateNodeData(id, { outputImage: null, status: "idle", error: null }); }, [id, updateNodeData]); const regenerateNode = useWorkflowStore((state) => state.regenerateNode); const isRunning = useWorkflowStore((state) => state.isRunning); const handleRegenerate = useCallback(() => { regenerateNode(id); }, [id, regenerateNode]); const loadImageById = useCallback(async (imageId: string) => { if (!generationsPath) { console.error("Generations path not configured"); return null; } try { const response = await fetch("/api/load-generation", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ directoryPath: generationsPath, imageId, }), }); const result = await response.json(); if (!result.success) { // Missing images are expected when refs point to deleted/moved files console.log(`Image not found: ${imageId}`); return null; } return result.image; } catch (error) { console.warn("Error loading image:", error); return null; } }, [generationsPath]); const handleCarouselPrevious = useCallback(async () => { const history = nodeData.imageHistory || []; if (history.length === 0 || isLoadingCarouselImage) return; const currentIndex = nodeData.selectedHistoryIndex || 0; const newIndex = currentIndex === 0 ? history.length - 1 : currentIndex - 1; const imageItem = history[newIndex]; setIsLoadingCarouselImage(true); const image = await loadImageById(imageItem.id); setIsLoadingCarouselImage(false); if (image) { updateNodeData(id, { outputImage: image, selectedHistoryIndex: newIndex, }); } }, [id, nodeData.imageHistory, nodeData.selectedHistoryIndex, isLoadingCarouselImage, loadImageById, updateNodeData]); const handleCarouselNext = useCallback(async () => { const history = nodeData.imageHistory || []; if (history.length === 0 || isLoadingCarouselImage) return; const currentIndex = nodeData.selectedHistoryIndex || 0; const newIndex = (currentIndex + 1) % history.length; const imageItem = history[newIndex]; setIsLoadingCarouselImage(true); const image = await loadImageById(imageItem.id); setIsLoadingCarouselImage(false); if (image) { updateNodeData(id, { outputImage: image, selectedHistoryIndex: newIndex, }); } }, [id, nodeData.imageHistory, nodeData.selectedHistoryIndex, isLoadingCarouselImage, loadImageById, updateNodeData]); // 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]); const isGeminiProvider = currentProvider === "gemini"; // Dynamic title based on selected model - just the model name const displayTitle = useMemo(() => { if (nodeData.selectedModel?.displayName && nodeData.selectedModel.modelId) { return nodeData.selectedModel.displayName; } // Fallback for legacy data or no model selected if (nodeData.model) { return GEMINI_IMAGE_MODELS.find(m => m.value === nodeData.model)?.label || nodeData.model; } return "Select model..."; }, [nodeData.selectedModel?.displayName, nodeData.selectedModel?.modelId, nodeData.model]); // Provider badge as title prefix const titlePrefix = useMemo(() => ( ), [currentProvider]); // Header action element based on provider mode const headerAction = useMemo(() => { if (!isGeminiOnly) { // External providers: browse button return ( ); } return null; }, [isGeminiOnly]); // Use selectedModel.modelId for Gemini models, fallback to legacy model field const currentModelId = isGeminiProvider ? (nodeData.selectedModel?.modelId || nodeData.model) : null; const isNanoBananaPro = currentModelId === "nano-banana-pro"; const hasCarouselImages = (nodeData.imageHistory || []).length > 1; // 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("Generation failed", "error", true, nodeData.error); } prevStatusRef.current = nodeData.status; }, [nodeData.status, nodeData.error]); // Auto-resize node when output image changes const prevOutputImageRef = useRef(null); useEffect(() => { // Only resize when outputImage transitions from null/different to a new value if (!nodeData.outputImage || nodeData.outputImage === prevOutputImageRef.current) { prevOutputImageRef.current = nodeData.outputImage ?? null; return; } prevOutputImageRef.current = nodeData.outputImage; // Use requestAnimationFrame to avoid React Flow update conflicts requestAnimationFrame(() => { getImageDimensions(nodeData.outputImage!).then((dims) => { if (!dims) return; const aspectRatio = dims.width / dims.height; setNodes((nodes) => nodes.map((node) => { if (node.id !== id) return node; // Preserve user's manually set height if present const currentHeight = typeof node.style?.height === 'number' ? node.style.height : undefined; const newSize = calculateNodeSizePreservingHeight(aspectRatio, currentHeight); return { ...node, style: { ...node.style, width: newSize.width, height: newSize.height } }; }) ); }); }); }, [id, nodeData.outputImage, setNodes]); 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} > {/* Input handles - ALWAYS use same IDs and positions for connection stability */} {/* Image input at 35%, Text input at 65% - never changes regardless of model */} {/* Image label */}
Image
{/* Prompt label */}
Prompt
{/* Image output */} {/* Output label */}
Image
{/* Preview area */} {nodeData.outputImage ? ( <>
Generated {/* Loading overlay for generation */} {nodeData.status === "loading" && (
)} {/* Error overlay when generation failed */} {nodeData.status === "error" && (
Generation failed See toast for details
)} {/* Loading overlay for carousel navigation */} {isLoadingCarouselImage && (
)}
{/* Carousel controls - only show if there are multiple images */} {hasCarouselImages && (
{(nodeData.selectedHistoryIndex || 0) + 1} / {(nodeData.imageHistory || []).length}
)} ) : (
{nodeData.status === "loading" ? ( ) : nodeData.status === "error" ? ( {nodeData.error || "Failed"} ) : ( Run to generate )}
)} {/* Model-specific parameters for external providers */} {!isGeminiOnly && nodeData.selectedModel?.modelId && ( )} {/* Aspect ratio and resolution row - only for Gemini models */} {currentProvider === "gemini" && (
{isNanoBananaPro && ( )}
)} {/* Google Search toggle - only for Nano Banana Pro */} {currentProvider === "gemini" && isNanoBananaPro && ( )}
{/* Model browser dialog for external providers */} {isBrowseDialogOpen && ( setIsBrowseDialogOpen(false)} onModelSelected={handleBrowseModelSelect} initialCapabilityFilter="image" /> )} ); } /** * @deprecated Use `GenerateImageNode` instead. This alias is kept for backward compatibility * with existing workflows but will be removed in a future version. */ export { GenerateImageNode as NanoBananaNode };