Update components/CustomNode.tsx

components/CustomNode.tsx

@@ -1,37 +1,44 @@
-
-
 import React, { memo } from 'react';
 import { Handle, Position, NodeProps } from 'reactflow';
 import { 
   Layers, Box, Activity, Zap, ArrowRight, Grid, Minimize, 
-  Database, GitBranch, AlignJustify, Type, Combine, Maximize,
+  Database, GitBranch, Type, Combine, Maximize,
   ArrowUpCircle, Sliders, RefreshCcw, Brain, Crosshair, Network, Clock, Eye, Workflow,
   Terminal, MinusCircle, Scaling, BoxSelect, Wifi,
   Mic, Speaker, Radio, Cuboid, Target, Scan,
   AudioLines, FileAudio, Hexagon, Component,
   ScanText, Mountain, Move, Radar, Map, Orbit, Wind,
   Film, Video, FastForward, Timer, Clapperboard, Merge,
-  Share2, Atom, Dna, Gamepad2, Sparkles, FlipVertical, RefreshCw
+  Share2, Atom, Dna, Gamepad2, Sparkles, FlipVertical, RefreshCw,
+  Scissors, Hash, Sigma, Calculator, BarChart3, Binary, X, Circle
 } from 'lucide-react';
 import { NodeData, LayerType } from '../types';
 
 const getIcon = (type: LayerType) => {
   switch (type) {
     case LayerType.INPUT: return <ArrowRight className="w-4 h-4" />;
+    case LayerType.OUTPUT: return <Circle className="w-4 h-4" />;
+    
     case LayerType.CONV1D: return <Activity className="w-4 h-4" />;
     case LayerType.CONV2D: return <Layers className="w-4 h-4" />;
     case LayerType.CONV3D: return <Box className="w-4 h-4" />;
     case LayerType.CONV_TRANSPOSE2D: return <ArrowUpCircle className="w-4 h-4" />;
     case LayerType.DEFORMABLE_CONV: return <Hexagon className="w-4 h-4" />;
+    case LayerType.SEPARABLE_CONV2D: return <Layers className="w-4 h-4" />;
+    case LayerType.DEPTHWISE_CONV2D: return <Layers className="w-4 h-4" />;
     case LayerType.LINEAR: return <Grid className="w-4 h-4" />;
+    
+    case LayerType.SIGMOID: return <Sigma className="w-4 h-4" />;
     case LayerType.RELU: 
     case LayerType.LEAKYRELU: 
     case LayerType.PRELU:
     case LayerType.GELU: 
     case LayerType.SILU: 
     case LayerType.SWIGLU:
-    case LayerType.SIGMOID: 
-    case LayerType.TANH: return <Zap className="w-4 h-4" />;
+    case LayerType.TANH:
+    case LayerType.SOFTPLUS:
+    case LayerType.SOFTSIGN: return <Zap className="w-4 h-4" />;
+    
     case LayerType.MAXPOOL: 
     case LayerType.MAXPOOL3D:
     case LayerType.AVGPOOL:
@@ -40,6 +47,19 @@ const getIcon = (type: LayerType) => {
     case LayerType.UPSAMPLE: return <Maximize className="w-4 h-4" />;
     case LayerType.PIXEL_SHUFFLE: return <Scaling className="w-4 h-4" />;
     
+    // TF / Preprocessing
+    case LayerType.RESCALING: 
+    case LayerType.RESIZING: return <Scaling className="w-4 h-4" />;
+    case LayerType.CENTER_CROP: return <Scissors className="w-4 h-4" />;
+    case LayerType.RANDOM_FLIP:
+    case LayerType.RANDOM_ROTATION:
+    case LayerType.RANDOM_ZOOM:
+    case LayerType.RANDOM_CONTRAST: return <RefreshCw className="w-4 h-4" />;
+    case LayerType.TEXT_VECTORIZATION: return <Type className="w-4 h-4" />;
+    case LayerType.NORMALIZATION_LAYER: return <BarChart3 className="w-4 h-4" />;
+    case LayerType.DISCRETIZATION: return <Binary className="w-4 h-4" />;
+    case LayerType.CATEGORY_ENCODING: return <Hash className="w-4 h-4" />;
+
     // Transformer / GenAI
     case LayerType.ATTENTION: 
     case LayerType.CROSS_ATTENTION: 
@@ -152,12 +172,22 @@ const getIcon = (type: LayerType) => {
     case LayerType.UNFLATTEN: return <Type className="w-4 h-4" />;
     case LayerType.RESHAPE: return <FlipVertical className="w-4 h-4" />;
     case LayerType.PERMUTE: return <RefreshCw className="w-4 h-4" />;
+    
+    // Merge
     case LayerType.CONCAT:
-    case LayerType.ADD: return <Combine className="w-4 h-4" />;
+    case LayerType.ADD: 
+    case LayerType.SUBTRACT:
+    case LayerType.MULTIPLY:
+    case LayerType.AVERAGE:
+    case LayerType.MAXIMUM:
+    case LayerType.MINIMUM: return <Combine className="w-4 h-4" />;
+    case LayerType.DOT: return <X className="w-4 h-4" />;
     
+    case LayerType.LAMBDA: return <Calculator className="w-4 h-4" />;
     case LayerType.CUSTOM: return <Terminal className="w-4 h-4" />;
     case LayerType.IDENTITY: 
     case LayerType.DROPOUT:
+    case LayerType.SPATIAL_DROPOUT:
     case LayerType.DROPPATH: return <MinusCircle className="w-4 h-4" />;
     
     default: return <Box className="w-4 h-4" />;
@@ -174,6 +204,8 @@ const getColor = (type: LayerType) => {
     case LayerType.CONV3D:
     case LayerType.CONV_TRANSPOSE2D:
     case LayerType.DEFORMABLE_CONV:
+    case LayerType.SEPARABLE_CONV2D:
+    case LayerType.DEPTHWISE_CONV2D:
     case LayerType.MAXPOOL: 
     case LayerType.MAXPOOL3D:
     case LayerType.AVGPOOL:
@@ -185,7 +217,20 @@ const getColor = (type: LayerType) => {
     case LayerType.LINEAR: 
     case LayerType.ACTION_HEAD: return 'border-violet-500 shadow-violet-500/20';
     
-    // GenAI / Transformer - Gold/Amber theme
+    // TF / Preprocessing (Orange/Amber)
+    case LayerType.RESCALING: 
+    case LayerType.RESIZING:
+    case LayerType.CENTER_CROP:
+    case LayerType.RANDOM_FLIP:
+    case LayerType.RANDOM_ROTATION:
+    case LayerType.RANDOM_ZOOM:
+    case LayerType.RANDOM_CONTRAST:
+    case LayerType.TEXT_VECTORIZATION:
+    case LayerType.NORMALIZATION_LAYER:
+    case LayerType.DISCRETIZATION:
+    case LayerType.CATEGORY_ENCODING: return 'border-amber-600 shadow-amber-600/20';
+
+    // GenAI / Transformer - Gold
     case LayerType.ATTENTION: 
     case LayerType.CROSS_ATTENTION:
     case LayerType.WINDOW_ATTENTION:
@@ -196,7 +241,7 @@ const getColor = (type: LayerType) => {
     case LayerType.SE_BLOCK:
     case LayerType.PATCH_EMBED:
     case LayerType.SAM_PROMPT_ENCODER:
-    case LayerType.SAM_MASK_DECODER: return 'border-amber-500 shadow-amber-500/20';
+    case LayerType.SAM_MASK_DECODER: return 'border-yellow-500 shadow-yellow-500/20';
     
     case LayerType.LSTM:
     case LayerType.GRU:
@@ -294,9 +339,17 @@ const getColor = (type: LayerType) => {
     case LayerType.ECHO_STATE: return 'border-white shadow-white/30';
 
     case LayerType.CONCAT:
-    case LayerType.ADD: return 'border-pink-500 shadow-pink-500/20';
+    case LayerType.ADD: 
+    case LayerType.SUBTRACT:
+    case LayerType.MULTIPLY:
+    case LayerType.AVERAGE:
+    case LayerType.MAXIMUM:
+    case LayerType.MINIMUM:
+    case LayerType.DOT: return 'border-pink-500 shadow-pink-500/20';
     
+    case LayerType.LAMBDA: return 'border-lime-500 shadow-lime-500/20';
     case LayerType.CUSTOM: return 'border-lime-500 shadow-lime-500/20';
+    
     case LayerType.IDENTITY: 
     case LayerType.FLATTEN:
     case LayerType.RESHAPE: