Update components/CustomNode.tsx

components/CustomNode.tsx

@@ -1,10 +1,11 @@
+
 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,
   ArrowUpCircle, Sliders, RefreshCcw, Brain, Crosshair, Network, Clock, Eye, Workflow,
-  Terminal, MinusCircle
+  Terminal, MinusCircle, Scaling, BoxSelect, Wifi
 } from 'lucide-react';
 import { NodeData, LayerType } from '../types';
 
@@ -13,19 +14,24 @@ const getIcon = (type: LayerType) => {
     case LayerType.INPUT: return <ArrowRight 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.LINEAR: return <Grid 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.MAXPOOL: 
+    case LayerType.MAXPOOL3D:
     case LayerType.AVGPOOL:
     case LayerType.ADAPTIVEAVGPOOL: return <Minimize className="w-4 h-4" />;
+    case LayerType.GLOBAL_AVG_POOL: return <BoxSelect className="w-4 h-4" />;
     case LayerType.UPSAMPLE: return <Maximize className="w-4 h-4" />;
+    case LayerType.PIXEL_SHUFFLE: return <Scaling className="w-4 h-4" />;
     
     // Transformer / GenAI
     case LayerType.ATTENTION: 
@@ -35,6 +41,7 @@ const getIcon = (type: LayerType) => {
     case LayerType.TRANSFORMER_DECODER: return <Brain className="w-4 h-4" />;
     case LayerType.MOE_BLOCK: return <Network className="w-4 h-4" />;
     case LayerType.ACTION_HEAD: return <Crosshair className="w-4 h-4" />;
+    case LayerType.SE_BLOCK: return <Wifi className="w-4 h-4" />;
     case LayerType.PATCH_EMBED:
     case LayerType.SAM_PROMPT_ENCODER: return <Eye className="w-4 h-4" />;
     case LayerType.SAM_MASK_DECODER: return <Workflow className="w-4 h-4" />;
@@ -46,6 +53,7 @@ const getIcon = (type: LayerType) => {
     case LayerType.EMBEDDING: return <Database className="w-4 h-4" />;
     case LayerType.LAYERNORM:
     case LayerType.BATCHNORM: 
+    case LayerType.GROUPNORM:
     case LayerType.INSTANCENORM: 
     case LayerType.RMSNORM: return <Sliders className="w-4 h-4" />;
     case LayerType.FLATTEN: return <Type className="w-4 h-4" />;
@@ -53,7 +61,9 @@ const getIcon = (type: LayerType) => {
     case LayerType.ADD: return <Combine className="w-4 h-4" />;
     
     case LayerType.CUSTOM: return <Terminal className="w-4 h-4" />;
-    case LayerType.IDENTITY: return <MinusCircle className="w-4 h-4" />;
+    case LayerType.IDENTITY: 
+    case LayerType.DROPOUT:
+    case LayerType.DROPPATH: return <MinusCircle className="w-4 h-4" />;
     
     default: return <Box className="w-4 h-4" />;
   }
@@ -66,11 +76,15 @@ const getColor = (type: LayerType) => {
     
     case LayerType.CONV2D: 
     case LayerType.CONV1D:
+    case LayerType.CONV3D:
     case LayerType.CONV_TRANSPOSE2D:
     case LayerType.MAXPOOL: 
+    case LayerType.MAXPOOL3D:
     case LayerType.AVGPOOL:
+    case LayerType.GLOBAL_AVG_POOL:
     case LayerType.ADAPTIVEAVGPOOL:
-    case LayerType.UPSAMPLE: return 'border-blue-500 shadow-blue-500/20';
+    case LayerType.UPSAMPLE:
+    case LayerType.PIXEL_SHUFFLE: return 'border-blue-500 shadow-blue-500/20';
     
     case LayerType.LINEAR: 
     case LayerType.ACTION_HEAD: return 'border-violet-500 shadow-violet-500/20';
@@ -82,6 +96,7 @@ const getColor = (type: LayerType) => {
     case LayerType.TRANSFORMER_ENCODER:
     case LayerType.TRANSFORMER_DECODER:
     case LayerType.MOE_BLOCK:
+    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';
@@ -94,6 +109,7 @@ const getColor = (type: LayerType) => {
     
     case LayerType.BATCHNORM:
     case LayerType.LAYERNORM: 
+    case LayerType.GROUPNORM:
     case LayerType.INSTANCENORM: 
     case LayerType.RMSNORM: return 'border-cyan-500 shadow-cyan-500/20';
     
@@ -180,4 +196,4 @@ const CustomNode = ({ data, selected }: NodeProps<NodeData>) => {
   );
 };
 
-export default memo(CustomNode);
+export default memo(CustomNode);