Update components/CustomNode.tsx

components/CustomNode.tsx

@@ -1,4 +1,5 @@
 
+
 import React, { memo } from 'react';
 import { Handle, Position, NodeProps } from 'reactflow';
 import { 
@@ -7,7 +8,9 @@ import {
   ArrowUpCircle, Sliders, RefreshCcw, Brain, Crosshair, Network, Clock, Eye, Workflow,
   Terminal, MinusCircle, Scaling, BoxSelect, Wifi,
   Mic, Speaker, Radio, Cuboid, Target, Scan,
-  AudioLines, FileAudio, Hexagon, Component
+  AudioLines, FileAudio, Hexagon, Component,
+  ScanText, Mountain, Move, Radar, Map, Orbit, Wind,
+  Film, Video, FastForward, Timer, Clapperboard, Merge
 } from 'lucide-react';
 import { NodeData, LayerType } from '../types';
 
@@ -78,6 +81,29 @@ const getIcon = (type: LayerType) => {
     case LayerType.POINTNET_BLOCK:
     case LayerType.POINT_TRANSFORMER: return <Scan className="w-4 h-4" />;
 
+    // OCR
+    case LayerType.TPS_TRANSFORM: 
+    case LayerType.CRNN_BLOCK:
+    case LayerType.CTC_DECODER: return <ScanText className="w-4 h-4" />;
+
+    // Robotics / Motion
+    case LayerType.DEPTH_DECODER:
+    case LayerType.DISPARITY_HEAD: return <Mountain className="w-4 h-4" />;
+    case LayerType.OPTICAL_FLOW: 
+    case LayerType.VELOCITY_HEAD: return <Wind className="w-4 h-4" />;
+    case LayerType.KALMAN_FILTER: return <Orbit className="w-4 h-4" />;
+    case LayerType.BEV_TRANSFORM: return <Map className="w-4 h-4" />;
+    case LayerType.RADAR_ENCODER: return <Radar className="w-4 h-4" />;
+
+    // Video / Generation
+    case LayerType.VIDEO_DIFFUSION_BLOCK: return <Film className="w-4 h-4" />;
+    case LayerType.SPATIO_TEMPORAL_ATTN: return <Timer className="w-4 h-4" />;
+    case LayerType.VIDEO_TOKENIZER: return <Video className="w-4 h-4" />;
+    case LayerType.FRAME_INTERPOLATOR: return <FastForward className="w-4 h-4" />;
+    case LayerType.TEMPORAL_SHIFT: return <Move className="w-4 h-4" />;
+    case LayerType.NON_LOCAL_BLOCK: return <Clapperboard className="w-4 h-4" />;
+    case LayerType.MULTIMODAL_FUSION: return <Merge className="w-4 h-4" />;
+
     case LayerType.LSTM:
     case LayerType.GRU: return <GitBranch className="w-4 h-4" />;
     case LayerType.EMBEDDING: return <Database className="w-4 h-4" />;
@@ -171,6 +197,29 @@ const getColor = (type: LayerType) => {
     case LayerType.TRIPLANE_ENC:
     case LayerType.GAUSSIAN_SPLAT: return 'border-teal-500 shadow-teal-500/20';
 
+    // OCR - Yellow/Lime
+    case LayerType.TPS_TRANSFORM: 
+    case LayerType.CRNN_BLOCK:
+    case LayerType.CTC_DECODER: return 'border-yellow-500 shadow-yellow-500/20';
+
+    // Robotics/Motion - Cyan/Sky
+    case LayerType.DEPTH_DECODER:
+    case LayerType.DISPARITY_HEAD:
+    case LayerType.OPTICAL_FLOW: 
+    case LayerType.VELOCITY_HEAD:
+    case LayerType.KALMAN_FILTER:
+    case LayerType.BEV_TRANSFORM:
+    case LayerType.RADAR_ENCODER: return 'border-sky-500 shadow-sky-500/20';
+
+    // Video - Fuchsia
+    case LayerType.VIDEO_DIFFUSION_BLOCK:
+    case LayerType.SPATIO_TEMPORAL_ATTN:
+    case LayerType.VIDEO_TOKENIZER:
+    case LayerType.FRAME_INTERPOLATOR:
+    case LayerType.TEMPORAL_SHIFT:
+    case LayerType.NON_LOCAL_BLOCK:
+    case LayerType.MULTIMODAL_FUSION: return 'border-fuchsia-500 shadow-fuchsia-500/20';
+
     case LayerType.CONCAT:
     case LayerType.ADD: return 'border-pink-500 shadow-pink-500/20';
     
@@ -258,4 +307,4 @@ const CustomNode = ({ data, selected }: NodeProps<NodeData>) => {
   );
 };
 
-export default memo(CustomNode);
+export default memo(CustomNode);