Update components/CustomNode.tsx

components/CustomNode.tsx

@@ -6,7 +6,8 @@ import {
   Database, GitBranch, AlignJustify, Type, Combine, Maximize,
   ArrowUpCircle, Sliders, RefreshCcw, Brain, Crosshair, Network, Clock, Eye, Workflow,
   Terminal, MinusCircle, Scaling, BoxSelect, Wifi,
-  Mic, Speaker, Radio, Cuboid, Target, Scan
+  Mic, Speaker, Radio, Cuboid, Target, Scan,
+  AudioLines, FileAudio, Hexagon, Component
 } from 'lucide-react';
 import { NodeData, LayerType } from '../types';
 
@@ -17,6 +18,7 @@ const getIcon = (type: LayerType) => {
     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.LINEAR: return <Grid className="w-4 h-4" />;
     case LayerType.RELU: 
     case LayerType.LEAKYRELU: 
@@ -51,22 +53,30 @@ const getIcon = (type: LayerType) => {
     case LayerType.ROPE: return <Clock className="w-4 h-4" />;
     
     // Audio
-    case LayerType.MEL_SPECTROGRAM: return <Radio className="w-4 h-4" />;
+    case LayerType.MEL_SPECTROGRAM:
+    case LayerType.STFT: return <Radio className="w-4 h-4" />;
+    case LayerType.SPEC_AUGMENT: return <AudioLines className="w-4 h-4" />;
     case LayerType.CONFORMER_BLOCK:
     case LayerType.WAVENET_BLOCK: return <Activity className="w-4 h-4" />;
+    case LayerType.RVC_ENCODER:
+    case LayerType.WAV2VEC2_ENC: return <FileAudio className="w-4 h-4" />;
     case LayerType.VOCODER: return <Speaker className="w-4 h-4" />;
     case LayerType.AUDIO_EMBEDDING: return <Mic className="w-4 h-4" />;
 
     // Detection / YOLO
     case LayerType.C2F_BLOCK:
-    case LayerType.SPPF_BLOCK: return <Layers className="w-4 h-4" />;
+    case LayerType.SPPF_BLOCK:
+    case LayerType.DARKNET_BLOCK: return <Layers className="w-4 h-4" />;
     case LayerType.DETECT_HEAD:
     case LayerType.ANCHOR_BOX: return <Target className="w-4 h-4" />;
+    case LayerType.NMS: return <Component className="w-4 h-4" />;
 
     // 3D
     case LayerType.NERF_BLOCK:
-    case LayerType.GAUSSIAN_SPLAT: return <Cuboid className="w-4 h-4" />;
-    case LayerType.POINTNET_BLOCK: return <Scan className="w-4 h-4" />;
+    case LayerType.GAUSSIAN_SPLAT:
+    case LayerType.TRIPLANE_ENC: return <Cuboid className="w-4 h-4" />;
+    case LayerType.POINTNET_BLOCK:
+    case LayerType.POINT_TRANSFORMER: return <Scan className="w-4 h-4" />;
 
     case LayerType.LSTM:
     case LayerType.GRU: return <GitBranch className="w-4 h-4" />;
@@ -98,6 +108,7 @@ const getColor = (type: LayerType) => {
     case LayerType.CONV1D:
     case LayerType.CONV3D:
     case LayerType.CONV_TRANSPOSE2D:
+    case LayerType.DEFORMABLE_CONV:
     case LayerType.MAXPOOL: 
     case LayerType.MAXPOOL3D:
     case LayerType.AVGPOOL:
@@ -137,19 +148,27 @@ const getColor = (type: LayerType) => {
     // Detection - Rose/Red theme
     case LayerType.C2F_BLOCK:
     case LayerType.SPPF_BLOCK:
+    case LayerType.DARKNET_BLOCK:
     case LayerType.DETECT_HEAD:
-    case LayerType.ANCHOR_BOX: return 'border-rose-500 shadow-rose-500/20';
+    case LayerType.ANCHOR_BOX:
+    case LayerType.NMS: return 'border-rose-500 shadow-rose-500/20';
 
     // Audio - Indigo theme
     case LayerType.MEL_SPECTROGRAM:
+    case LayerType.STFT:
+    case LayerType.SPEC_AUGMENT:
     case LayerType.CONFORMER_BLOCK:
     case LayerType.WAVENET_BLOCK:
     case LayerType.VOCODER:
+    case LayerType.RVC_ENCODER:
+    case LayerType.WAV2VEC2_ENC:
     case LayerType.AUDIO_EMBEDDING: return 'border-indigo-500 shadow-indigo-500/20';
 
     // 3D - Teal theme
     case LayerType.NERF_BLOCK:
     case LayerType.POINTNET_BLOCK:
+    case LayerType.POINT_TRANSFORMER:
+    case LayerType.TRIPLANE_ENC:
     case LayerType.GAUSSIAN_SPLAT: return 'border-teal-500 shadow-teal-500/20';
 
     case LayerType.CONCAT: