Upload 13 files

.env.local

@@ -0,0 +1 @@
+GEMINI_API_KEY=PLACEHOLDER_API_KEY

.gitignore

@@ -0,0 +1,24 @@
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
+
+node_modules
+dist
+dist-ssr
+*.local
+
+# Editor directories and files
+.vscode/*
+!.vscode/extensions.json
+.idea
+.DS_Store
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?

App.tsx

@@ -0,0 +1,408 @@
+import React, { useCallback, useRef, useState, useMemo, useEffect } from 'react';
+import ReactFlow, {
+  ReactFlowProvider,
+  addEdge,
+  useNodesState,
+  useEdgesState,
+  Controls,
+  Background,
+  Connection,
+  Edge,
+  Node,
+  BackgroundVariant,
+  Panel,
+  ConnectionMode
+} from 'reactflow';
+import { FileText, CheckCircle, AlertTriangle, X, Copy, Check, Lightbulb, Key } from 'lucide-react';
+import Sidebar from './components/Sidebar';
+import PropertiesPanel from './components/PropertiesPanel';
+import CustomNode from './components/CustomNode';
+import CodeViewer from './components/CodeViewer';
+import AIBuilderModal from './components/AIBuilderModal';
+import SuggestionsModal from './components/SuggestionsModal';
+import ApiKeyModal from './components/ApiKeyModal';
+import { INITIAL_NODES, INITIAL_EDGES, LAYER_DEFINITIONS, TEMPLATES } from './constants';
+import { generateRefinedPrompt, validateArchitecture, getArchitectureSuggestions, getUserApiKey, implementArchitectureSuggestions } from './services/geminiService';
+import { NodeData, LayerType } from './types';
+
+let id = 1000;
+const getId = () => `${id++}`;
+
+const AppContent = () => {
+  const reactFlowWrapper = useRef<HTMLDivElement>(null);
+  const [nodes, setNodes, onNodesChange] = useNodesState(INITIAL_NODES);
+  const [edges, setEdges, onEdgesChange] = useEdgesState(INITIAL_EDGES);
+  const [reactFlowInstance, setReactFlowInstance] = useState<any>(null);
+  const [selectedNodeId, setSelectedNodeId] = useState<string | null>(null);
+  
+  const [generatedPrompt, setGeneratedPrompt] = useState('');
+  const [isPromptViewerOpen, setIsPromptViewerOpen] = useState(false);
+  const [isGeneratingPrompt, setIsGeneratingPrompt] = useState(false);
+  
+  const [validationMsg, setValidationMsg] = useState<string | null>(null);
+  const [layoutCopied, setLayoutCopied] = useState(false);
+
+  // Suggestions State
+  const [isSuggestionsOpen, setIsSuggestionsOpen] = useState(false);
+  const [suggestionsText, setSuggestionsText] = useState('');
+  const [isSuggestionsLoading, setIsSuggestionsLoading] = useState(false);
+  const [isImplementingSuggestions, setIsImplementingSuggestions] = useState(false);
+
+  // AI Builder and Template State
+  const [isAIBuilderOpen, setIsAIBuilderOpen] = useState(false);
+  const [isTemplateMenuOpen, setIsTemplateMenuOpen] = useState(false);
+
+  // API Key State
+  const [isApiKeyModalOpen, setIsApiKeyModalOpen] = useState(false);
+  const [isConnected, setIsConnected] = useState(false);
+
+  // Load connection status on mount
+  useEffect(() => {
+    setIsConnected(!!getUserApiKey());
+  }, []);
+
+  // Define custom node types
+  const nodeTypes = useMemo(() => ({ custom: CustomNode }), []);
+
+  const onConnect = useCallback(
+    (params: Connection) => setEdges((eds) => addEdge({ ...params, animated: true, style: { stroke: '#94a3b8' } }, eds)),
+    [setEdges]
+  );
+
+  const onDragOver = useCallback((event: React.DragEvent) => {
+    event.preventDefault();
+    event.dataTransfer.dropEffect = 'move';
+  }, []);
+
+  const onDrop = useCallback(
+    (event: React.DragEvent) => {
+      event.preventDefault();
+
+      if (!reactFlowWrapper.current || !reactFlowInstance) return;
+
+      const type = event.dataTransfer.getData('application/reactflow') as LayerType;
+      
+      if (!type) return;
+
+      const position = reactFlowInstance.screenToFlowPosition({
+        x: event.clientX,
+        y: event.clientY,
+      });
+
+      const newNode: Node<NodeData> = {
+        id: getId(),
+        type: 'custom',
+        position,
+        data: { 
+          label: LAYER_DEFINITIONS[type].label, 
+          type: type, 
+          params: LAYER_DEFINITIONS[type].parameters.reduce((acc, p) => ({...acc, [p.name]: p.default}), {}) 
+        },
+      };
+
+      setNodes((nds) => nds.concat(newNode));
+      setSelectedNodeId(newNode.id);
+    },
+    [reactFlowInstance, setNodes]
+  );
+
+  const onNodeClick = useCallback((_: React.MouseEvent, node: Node) => {
+    setSelectedNodeId(node.id);
+  }, []);
+
+  const onPaneClick = useCallback(() => {
+    setSelectedNodeId(null);
+  }, []);
+
+  // Update node params from Properties Panel
+  const updateNodeParams = (id: string, newParams: Record<string, any>) => {
+    setNodes((nds) =>
+      nds.map((node) => {
+        if (node.id === id) {
+          return {
+            ...node,
+            data: { ...node.data, params: newParams }
+          };
+        }
+        return node;
+      })
+    );
+  };
+
+  const deleteNode = (id: string) => {
+    setNodes((nds) => nds.filter((node) => node.id !== id));
+    setEdges((eds) => eds.filter((edge) => edge.source !== id && edge.target !== id));
+    setSelectedNodeId(null);
+  };
+
+  const handleGeneratePrompt = async () => {
+    if (!isConnected) { setIsApiKeyModalOpen(true); return; }
+    setIsGeneratingPrompt(true);
+    setIsPromptViewerOpen(true);
+    try {
+        const promptText = await generateRefinedPrompt(nodes, edges);
+        setGeneratedPrompt(promptText);
+    } catch (e) {
+        setGeneratedPrompt("Failed to generate prompt. Please try again.");
+    } finally {
+        setIsGeneratingPrompt(false);
+    }
+  };
+
+  const handleValidate = async () => {
+    if (!isConnected) { setIsApiKeyModalOpen(true); return; }
+    setValidationMsg("Validating...");
+    const result = await validateArchitecture(nodes, edges);
+    setValidationMsg(result);
+    // Auto hide after 5 seconds if brief
+    if (result.length < 50) {
+      setTimeout(() => setValidationMsg(null), 5000);
+    }
+  };
+
+  const handleGetSuggestions = async () => {
+    if (!isConnected) { setIsApiKeyModalOpen(true); return; }
+    setIsSuggestionsOpen(true);
+    setIsSuggestionsLoading(true);
+    try {
+      const suggestions = await getArchitectureSuggestions(nodes, edges);
+      setSuggestionsText(suggestions);
+    } catch (error) {
+      setSuggestionsText("Failed to get suggestions.");
+    } finally {
+      setIsSuggestionsLoading(false);
+    }
+  };
+
+  const handleImplementSuggestions = async () => {
+    if (!suggestionsText || isImplementingSuggestions) return;
+    setIsImplementingSuggestions(true);
+    try {
+      const result = await implementArchitectureSuggestions(nodes, edges, suggestionsText);
+      if (result && result.nodes) {
+         // Fix types and defaults
+         const processedNodes = result.nodes.map((n: any) => ({
+            ...n,
+            type: 'custom',
+            data: { ...n.data, label: n.data.label || n.data.type }
+         }));
+         const processedEdges = result.edges.map((e: any) => ({
+            ...e, animated: true, style: { stroke: '#94a3b8' }
+         }));
+         
+         setNodes(processedNodes);
+         setEdges(processedEdges);
+         setIsSuggestionsOpen(false); // Close modal on success
+      }
+    } catch (e) {
+      console.error(e);
+      alert("Failed to implement suggestions automatically.");
+    } finally {
+      setIsImplementingSuggestions(false);
+    }
+  };
+
+  const handleCopyLayout = () => {
+    const layout = {
+      nodes,
+      edges
+    };
+    navigator.clipboard.writeText(JSON.stringify(layout, null, 2));
+    setLayoutCopied(true);
+    setTimeout(() => setLayoutCopied(false), 2000);
+  };
+
+  const handleApplyAIBuild = (newNodes: any[], newEdges: any[]) => {
+    setNodes(newNodes);
+    setEdges(newEdges);
+  };
+
+  const loadTemplate = (templateKey: string) => {
+    const template = TEMPLATES[templateKey];
+    if (template) {
+      // Need to deep clone to avoid reference issues if loaded multiple times
+      const clonedNodes = JSON.parse(JSON.stringify(template.nodes));
+      const clonedEdges = JSON.parse(JSON.stringify(template.edges));
+      
+      // Update styling and functional props that aren't in JSON
+      const hydratedNodes = clonedNodes.map((n: any) => ({
+        ...n,
+        data: {
+          ...n.data,
+          label: LAYER_DEFINITIONS[n.data.type as LayerType]?.label || n.data.label
+        }
+      }));
+
+      const hydratedEdges = clonedEdges.map((e: any) => ({
+        ...e,
+        animated: true,
+        style: { stroke: '#94a3b8' }
+      }));
+
+      setNodes(hydratedNodes);
+      setEdges(hydratedEdges);
+      setIsTemplateMenuOpen(false);
+    }
+  };
+
+  const selectedNode = useMemo(() => 
+    nodes.find((n) => n.id === selectedNodeId) || null, 
+    [nodes, selectedNodeId]
+  );
+
+  return (
+    <div className="flex h-screen w-screen bg-slate-950">
+      <Sidebar 
+        onOpenAIBuilder={() => isConnected ? setIsAIBuilderOpen(true) : setIsApiKeyModalOpen(true)}
+        onSelectTemplate={(id) => id === 'menu' ? setIsTemplateMenuOpen(true) : loadTemplate(id)}
+        isConnected={isConnected}
+      />
+      
+      <div className="flex-1 h-full relative" ref={reactFlowWrapper}>
+        <ReactFlow
+          nodes={nodes}
+          edges={edges}
+          onNodesChange={onNodesChange}
+          onEdgesChange={onEdgesChange}
+          onConnect={onConnect}
+          onInit={setReactFlowInstance}
+          onDrop={onDrop}
+          onDragOver={onDragOver}
+          onNodeClick={onNodeClick}
+          onPaneClick={onPaneClick}
+          nodeTypes={nodeTypes}
+          connectionMode={ConnectionMode.Loose}
+          fitView
+          className="bg-slate-950"
+        >
+          <Background color="#1e293b" variant={BackgroundVariant.Dots} gap={20} size={1} />
+          <Controls className="!bg-slate-800 !border-slate-700 [&>button]:!fill-slate-300 [&>button:hover]:!bg-slate-700" />
+          
+          <Panel position="top-right" className="flex gap-2">
+            <button 
+              onClick={handleCopyLayout}
+              className="flex items-center gap-2 px-4 py-2 bg-slate-800 hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-sm font-medium transition-colors shadow-lg"
+            >
+              {layoutCopied ? <Check size={16} className="text-emerald-500" /> : <Copy size={16} className="text-slate-400" />}
+              {layoutCopied ? 'Copied' : 'Copy Layout'}
+            </button>
+            <button
+               onClick={() => setIsApiKeyModalOpen(true)}
+               className="flex items-center gap-2 px-4 py-2 bg-slate-800 hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-sm font-medium transition-colors shadow-lg"
+            >
+              <Key size={16} className={isConnected ? "text-emerald-500" : "text-slate-400"} />
+              {isConnected ? 'Key Active' : 'Add Key'}
+            </button>
+            <button 
+              onClick={handleGetSuggestions}
+              className="flex items-center gap-2 px-4 py-2 bg-slate-800 hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-sm font-medium transition-colors shadow-lg"
+            >
+              <Lightbulb size={16} className="text-amber-400" />
+              Suggestions
+            </button>
+            <button 
+              onClick={handleValidate}
+              className="flex items-center gap-2 px-4 py-2 bg-slate-800 hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-sm font-medium transition-colors shadow-lg"
+            >
+              <CheckCircle size={16} className="text-emerald-500" />
+              Validate
+            </button>
+            <button 
+              onClick={handleGeneratePrompt}
+              className="flex items-center gap-2 px-4 py-2 bg-blue-600 hover:bg-blue-500 rounded-lg text-white text-sm font-medium transition-colors shadow-lg shadow-blue-900/20"
+            >
+              <FileText size={16} />
+              Generate Prompt
+            </button>
+          </Panel>
+
+          {validationMsg && (
+            <Panel position="bottom-center" className="mb-8">
+              <div className="bg-slate-900/90 backdrop-blur border border-slate-700 text-slate-200 px-6 py-4 rounded-xl shadow-2xl max-w-lg">
+                <div className="flex items-start gap-3">
+                  <AlertTriangle className="text-amber-500 shrink-0 mt-0.5" size={18} />
+                  <p className="text-sm whitespace-pre-wrap">{validationMsg}</p>
+                </div>
+                <button 
+                  onClick={() => setValidationMsg(null)}
+                  className="mt-3 text-xs text-slate-500 hover:text-slate-300 underline"
+                >
+                  Dismiss
+                </button>
+              </div>
+            </Panel>
+          )}
+
+          {/* Template Selection Modal Overlay */}
+          {isTemplateMenuOpen && (
+             <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/50 backdrop-blur-sm">
+                <div className="bg-slate-900 border border-slate-700 rounded-xl shadow-2xl p-6 w-96 max-w-full">
+                  <div className="flex justify-between items-center mb-4">
+                     <h3 className="text-lg font-bold text-white">Select Template</h3>
+                     <button onClick={() => setIsTemplateMenuOpen(false)}><X size={20} className="text-slate-500 hover:text-white"/></button>
+                  </div>
+                  <div className="grid gap-3">
+                    {Object.entries(TEMPLATES).map(([key, t]) => (
+                      <button 
+                        key={key}
+                        onClick={() => loadTemplate(key)}
+                        className="text-left p-3 rounded bg-slate-800 hover:bg-slate-750 border border-slate-700 hover:border-blue-500/50 transition-all group"
+                      >
+                         <div className="font-semibold text-slate-200 group-hover:text-blue-400">{t.name}</div>
+                         <div className="text-xs text-slate-500">{t.description}</div>
+                      </button>
+                    ))}
+                  </div>
+                </div>
+             </div>
+          )}
+
+        </ReactFlow>
+      </div>
+
+      <PropertiesPanel 
+        selectedNode={selectedNode} 
+        onChange={updateNodeParams}
+        onDelete={deleteNode}
+        onClose={() => setSelectedNodeId(null)}
+      />
+
+      <CodeViewer 
+        isOpen={isPromptViewerOpen} 
+        onClose={() => setIsPromptViewerOpen(false)} 
+        code={generatedPrompt}
+        isLoading={isGeneratingPrompt}
+      />
+      
+      <SuggestionsModal
+        isOpen={isSuggestionsOpen}
+        onClose={() => setIsSuggestionsOpen(false)}
+        suggestions={suggestionsText}
+        isLoading={isSuggestionsLoading}
+        onImplement={handleImplementSuggestions}
+        isImplementing={isImplementingSuggestions}
+      />
+
+      <AIBuilderModal 
+        isOpen={isAIBuilderOpen}
+        onClose={() => setIsAIBuilderOpen(false)}
+        onApply={handleApplyAIBuild}
+        currentNodes={nodes}
+      />
+
+      <ApiKeyModal 
+        isOpen={isApiKeyModalOpen}
+        onClose={() => setIsApiKeyModalOpen(false)}
+        onSuccess={() => setIsConnected(true)}
+      />
+    </div>
+  );
+};
+
+const App = () => (
+  <ReactFlowProvider>
+    <AppContent />
+  </ReactFlowProvider>
+);
+
+export default App;

Dockerfile

@@ -0,0 +1,16 @@
+# Stage 1: Build React/Vite app
+FROM node:18-alpine AS builder
+WORKDIR /app
+COPY package*.json ./
+RUN npm install
+COPY . .
+ARG GEMINI_API_KEY
+ENV GEMINI_API_KEY=$GEMINI_API_KEY
+RUN npm run build
+
+# Stage 2: Serve with Python
+FROM python:3.11-slim
+WORKDIR /app
+COPY --from=builder /app/dist .
+EXPOSE 7860
+CMD ["python3", "-m", "http.server", "7860", "--bind", "0.0.0.0"]

README.md

@@ -1,10 +1,8 @@
----
-title: Testarcbuilder
-emoji: 🐢
-colorFrom: blue
-colorTo: yellow
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Dockerized React App
+emoji: 🐳
+colorFrom: purple
+colorTo: gray
+sdk: docker
+app_port: 7860
+---

constants.ts

@@ -0,0 +1,663 @@
+import { LayerDefinition, LayerType, GraphTemplate } from './types';
+
+export const LAYER_DEFINITIONS: Record<LayerType, LayerDefinition> = {
+  // --- CORE ---
+  [LayerType.INPUT]: {
+    type: LayerType.INPUT,
+    label: 'Input Data',
+    description: 'Entry point for data tensors',
+    category: 'Core',
+    parameters: [
+      { name: 'modality', type: 'select', label: 'Modality', default: 'Tensor', options: ['Tensor', 'Image', 'Text', 'Audio', 'Video', 'Latent', 'State'] },
+      { name: 'shape', type: 'string', label: 'Shape (e.g. 3,224,224)', default: '3, 224, 224' },
+      { name: 'batch_size', type: 'number', label: 'Batch Size', default: 32 },
+      { name: 'dtype', type: 'select', label: 'Data Type', default: 'float32', options: ['float32', 'int64', 'bool'] }
+    ]
+  },
+  [LayerType.LINEAR]: {
+    type: LayerType.LINEAR,
+    label: 'Linear (Dense)',
+    description: 'Fully connected layer',
+    category: 'Core',
+    parameters: [
+      { name: 'in_features', type: 'number', label: 'In Features (Opt)', default: 0, description: "0 = Auto-infer" },
+      { name: 'out_features', type: 'number', label: 'Output Features', default: 128 },
+      { name: 'bias', type: 'boolean', label: 'Use Bias', default: true }
+    ]
+  },
+  [LayerType.OUTPUT]: {
+    type: LayerType.OUTPUT,
+    label: 'Output Head',
+    description: 'Final model output',
+    category: 'Core',
+    parameters: [
+      { name: 'num_classes', type: 'number', label: 'Classes', default: 10 },
+      { name: 'activation', type: 'select', label: 'Activation', default: 'Softmax', options: ['None', 'Softmax', 'Sigmoid', 'LogSoftmax'] }
+    ]
+  },
+  [LayerType.EMBEDDING]: {
+    type: LayerType.EMBEDDING,
+    label: 'Embedding',
+    description: 'Lookup table for embeddings',
+    category: 'Core',
+    parameters: [
+      { name: 'num_embeddings', type: 'number', label: 'Vocab Size', default: 10000 },
+      { name: 'embedding_dim', type: 'number', label: 'Embed Dim', default: 256 },
+      { name: 'padding_idx', type: 'number', label: 'Padding Index', default: 0 },
+      { name: 'max_norm', type: 'number', label: 'Max Norm', default: 0, description: "0 = None" }
+    ]
+  },
+  [LayerType.POS_EMBED]: {
+    type: LayerType.POS_EMBED,
+    label: 'Positional Embed',
+    description: 'Learnable positional embeddings',
+    category: 'Core',
+    parameters: [
+      { name: 'num_embeddings', type: 'number', label: 'Max Positions', default: 1024 },
+      { name: 'embedding_dim', type: 'number', label: 'Embed Dim', default: 256 }
+    ]
+  },
+
+  // --- MERGE ---
+  [LayerType.CONCAT]: {
+    type: LayerType.CONCAT,
+    label: 'Concatenate',
+    description: 'Merge inputs along a dim',
+    category: 'Merge',
+    parameters: [
+      { name: 'dim', type: 'number', label: 'Dimension', default: 1 }
+    ]
+  },
+  [LayerType.ADD]: {
+    type: LayerType.ADD,
+    label: 'Add (Sum)',
+    description: 'Element-wise addition (Residual)',
+    category: 'Merge',
+    parameters: []
+  },
+
+  // --- GENAI / ADVANCED ---
+  [LayerType.RMSNORM]: {
+    type: LayerType.RMSNORM,
+    label: 'RMSNorm',
+    description: 'Root Mean Square Norm (LLMs)',
+    category: 'GenAI',
+    parameters: [
+      { name: 'dim', type: 'number', label: 'Dimension', default: 512 },
+      { name: 'eps', type: 'number', label: 'Epsilon', default: 1e-6 }
+    ]
+  },
+  [LayerType.ROPE]: {
+    type: LayerType.ROPE,
+    label: 'RoPE',
+    description: 'Rotary Positional Embedding',
+    category: 'GenAI',
+    parameters: [
+      { name: 'dim', type: 'number', label: 'Head Dim', default: 64 },
+      { name: 'max_position', type: 'number', label: 'Max Pos', default: 2048 }
+    ]
+  },
+  [LayerType.PATCH_EMBED]: {
+    type: LayerType.PATCH_EMBED,
+    label: 'Patch Embed (ViT)',
+    description: 'Image to Sequence Patches',
+    category: 'GenAI',
+    parameters: [
+      { name: 'img_size', type: 'number', label: 'Image Size', default: 224 },
+      { name: 'patch_size', type: 'number', label: 'Patch Size', default: 16 },
+      { name: 'in_chans', type: 'number', label: 'In Channels', default: 3 },
+      { name: 'embed_dim', type: 'number', label: 'Embed Dim', default: 768 }
+    ]
+  },
+  [LayerType.MOE_BLOCK]: {
+    type: LayerType.MOE_BLOCK,
+    label: 'Sparse MoE Block',
+    description: 'Mixture of Experts Layer',
+    category: 'GenAI',
+    parameters: [
+      { name: 'num_experts', type: 'number', label: 'Num Experts', default: 8 },
+      { name: 'top_k', type: 'number', label: 'Top K (Active)', default: 2 },
+      { name: 'hidden_dim', type: 'number', label: 'Hidden Dim', default: 512 },
+      { name: 'expert_dim', type: 'number', label: 'Expert Dim', default: 2048 }
+    ]
+  },
+  [LayerType.ACTION_HEAD]: {
+    type: LayerType.ACTION_HEAD,
+    label: 'Action Head',
+    description: 'Decision output for Agents/LAMs',
+    category: 'GenAI',
+    parameters: [
+      { name: 'input_dim', type: 'number', label: 'Input Dim', default: 512 },
+      { name: 'num_actions', type: 'number', label: 'Num Actions', default: 50 },
+      { name: 'action_type', type: 'select', label: 'Type', default: 'Discrete', options: ['Discrete', 'Continuous'] }
+    ]
+  },
+  [LayerType.TIME_EMBEDDING]: {
+    type: LayerType.TIME_EMBEDDING,
+    label: 'Time Embedding',
+    description: 'Sinusoidal Time Embed (Diffusion)',
+    category: 'GenAI',
+    parameters: [
+      { name: 'dim', type: 'number', label: 'Dimension', default: 256 }
+    ]
+  },
+  [LayerType.SAM_PROMPT_ENCODER]: {
+    type: LayerType.SAM_PROMPT_ENCODER,
+    label: 'SAM Prompt Enc',
+    description: 'Encodes points/boxes (SAM)',
+    category: 'GenAI',
+    parameters: [
+      { name: 'embed_dim', type: 'number', label: 'Embed Dim', default: 256 }
+    ]
+  },
+  [LayerType.SAM_MASK_DECODER]: {
+    type: LayerType.SAM_MASK_DECODER,
+    label: 'SAM Mask Dec',
+    description: 'Decodes segmentation masks',
+    category: 'GenAI',
+    parameters: [
+      { name: 'transformer_dim', type: 'number', label: 'Model Dim', default: 256 },
+      { name: 'num_multimask_outputs', type: 'number', label: 'Num Masks', default: 3 }
+    ]
+  },
+
+  // --- CONVOLUTION ---
+  [LayerType.CONV1D]: {
+    type: LayerType.CONV1D,
+    label: 'Conv1D',
+    description: '1D Convolution (Audio/Text)',
+    category: 'Convolution',
+    parameters: [
+      { name: 'in_channels', type: 'number', label: 'In Channels (Opt)', default: 0 },
+      { name: 'out_channels', type: 'number', label: 'Filters', default: 32 },
+      { name: 'kernel_size', type: 'number', label: 'Kernel Size', default: 3 },
+      { name: 'stride', type: 'number', label: 'Stride', default: 1 },
+      { name: 'padding', type: 'number', label: 'Padding', default: 1 }
+    ]
+  },
+  [LayerType.CONV2D]: {
+    type: LayerType.CONV2D,
+    label: 'Conv2D',
+    description: '2D Convolutional Layer',
+    category: 'Convolution',
+    parameters: [
+      { name: 'in_channels', type: 'number', label: 'In Channels (Opt)', default: 0 },
+      { name: 'out_channels', type: 'number', label: 'Filters', default: 64 },
+      { name: 'kernel_size', type: 'number', label: 'Kernel Size', default: 3 },
+      { name: 'stride', type: 'number', label: 'Stride', default: 1 },
+      { name: 'padding', type: 'number', label: 'Padding', default: 1 },
+      { name: 'dilation', type: 'number', label: 'Dilation', default: 1 },
+      { name: 'groups', type: 'number', label: 'Groups', default: 1, description: "For depthwise separable" },
+      { name: 'bias', type: 'boolean', label: 'Bias', default: true }
+    ]
+  },
+  [LayerType.CONV_TRANSPOSE2D]: {
+    type: LayerType.CONV_TRANSPOSE2D,
+    label: 'ConvTranspose2D',
+    description: 'Deconvolution (Upsampling)',
+    category: 'Convolution',
+    parameters: [
+      { name: 'out_channels', type: 'number', label: 'Filters', default: 64 },
+      { name: 'kernel_size', type: 'number', label: 'Kernel Size', default: 2 },
+      { name: 'stride', type: 'number', label: 'Stride', default: 2 },
+      { name: 'padding', type: 'number', label: 'Padding', default: 0 }
+    ]
+  },
+  [LayerType.MAXPOOL]: {
+    type: LayerType.MAXPOOL,
+    label: 'MaxPool2D',
+    description: 'Max pooling operation',
+    category: 'Convolution',
+    parameters: [
+      { name: 'kernel_size', type: 'number', label: 'Kernel Size', default: 2 },
+      { name: 'stride', type: 'number', label: 'Stride', default: 2 }
+    ]
+  },
+  [LayerType.AVGPOOL]: {
+    type: LayerType.AVGPOOL,
+    label: 'AvgPool2D',
+    description: 'Average pooling operation',
+    category: 'Convolution',
+    parameters: [
+      { name: 'kernel_size', type: 'number', label: 'Kernel Size', default: 2 },
+      { name: 'stride', type: 'number', label: 'Stride', default: 2 }
+    ]
+  },
+  [LayerType.ADAPTIVEAVGPOOL]: {
+    type: LayerType.ADAPTIVEAVGPOOL,
+    label: 'AdaptAvgPool2D',
+    description: 'Pools to specific output size',
+    category: 'Convolution',
+    parameters: [
+      { name: 'output_size', type: 'number', label: 'Output Size (NxN)', default: 1 }
+    ]
+  },
+  [LayerType.UPSAMPLE]: {
+    type: LayerType.UPSAMPLE,
+    label: 'Upsample',
+    description: 'Increases spatial size',
+    category: 'Convolution',
+    parameters: [
+      { name: 'scale_factor', type: 'number', label: 'Scale Factor', default: 2 },
+      { name: 'mode', type: 'select', label: 'Mode', default: 'nearest', options: ['nearest', 'bilinear', 'bicubic'] }
+    ]
+  },
+
+  // --- RECURRENT ---
+  [LayerType.LSTM]: {
+    type: LayerType.LSTM,
+    label: 'LSTM',
+    description: 'Long Short-Term Memory',
+    category: 'Recurrent',
+    parameters: [
+      { name: 'input_size', type: 'number', label: 'Input Size (Opt)', default: 0 },
+      { name: 'hidden_size', type: 'number', label: 'Hidden Size', default: 128 },
+      { name: 'num_layers', type: 'number', label: 'Num Layers', default: 1 },
+      { name: 'bidirectional', type: 'boolean', label: 'Bidirectional', default: false },
+      { name: 'dropout', type: 'number', label: 'Dropout', default: 0.0 },
+      { name: 'batch_first', type: 'boolean', label: 'Batch First', default: true }
+    ]
+  },
+  [LayerType.GRU]: {
+    type: LayerType.GRU,
+    label: 'GRU',
+    description: 'Gated Recurrent Unit',
+    category: 'Recurrent',
+    parameters: [
+      { name: 'input_size', type: 'number', label: 'Input Size (Opt)', default: 0 },
+      { name: 'hidden_size', type: 'number', label: 'Hidden Size', default: 128 },
+      { name: 'num_layers', type: 'number', label: 'Num Layers', default: 1 },
+      { name: 'bidirectional', type: 'boolean', label: 'Bidirectional', default: false },
+      { name: 'dropout', type: 'number', label: 'Dropout', default: 0.0 },
+      { name: 'batch_first', type: 'boolean', label: 'Batch First', default: true }
+    ]
+  },
+
+  // --- UTILITY / ACTIVATION ---
+  [LayerType.RELU]: {
+    type: LayerType.RELU,
+    label: 'ReLU',
+    description: 'Rectified Linear Unit',
+    category: 'Utility',
+    parameters: []
+  },
+  [LayerType.SWIGLU]: {
+    type: LayerType.SWIGLU,
+    label: 'SwiGLU',
+    description: 'Swish-Gated Linear Unit',
+    category: 'Utility',
+    parameters: [
+      { name: 'dim', type: 'number', label: 'Dim (Opt)', default: 0 }
+    ]
+  },
+  [LayerType.LEAKYRELU]: {
+    type: LayerType.LEAKYRELU,
+    label: 'LeakyReLU',
+    description: 'Leaky ReLU Activation',
+    category: 'Utility',
+    parameters: [
+      { name: 'negative_slope', type: 'number', label: 'Negative Slope', default: 0.01 }
+    ]
+  },
+  [LayerType.GELU]: {
+    type: LayerType.GELU,
+    label: 'GELU',
+    description: 'Gaussian Error Linear Unit',
+    category: 'Utility',
+    parameters: []
+  },
+  [LayerType.SILU]: {
+    type: LayerType.SILU,
+    label: 'SiLU (Swish)',
+    description: 'Sigmoid Linear Unit',
+    category: 'Utility',
+    parameters: []
+  },
+  [LayerType.SIGMOID]: {
+    type: LayerType.SIGMOID,
+    label: 'Sigmoid',
+    description: 'Sigmoid Activation',
+    category: 'Utility',
+    parameters: []
+  },
+  [LayerType.TANH]: {
+    type: LayerType.TANH,
+    label: 'Tanh',
+    description: 'Hyperbolic Tangent',
+    category: 'Utility',
+    parameters: []
+  },
+  [LayerType.DROPOUT]: {
+    type: LayerType.DROPOUT,
+    label: 'Dropout',
+    description: 'Random zeroing of elements',
+    category: 'Utility',
+    parameters: [
+      { name: 'p', type: 'number', label: 'Probability', default: 0.5 }
+    ]
+  },
+  [LayerType.FLATTEN]: {
+    type: LayerType.FLATTEN,
+    label: 'Flatten',
+    description: 'Flattens input to 1D',
+    category: 'Utility',
+    parameters: []
+  },
+  [LayerType.CUSTOM]: {
+    type: LayerType.CUSTOM,
+    label: 'Custom / Code',
+    description: 'Any PyTorch/Python Module',
+    category: 'Utility',
+    parameters: [
+        { name: 'class_name', type: 'string', label: 'Class Name', default: 'MyCustomLayer', description: "e.g., 'torch.nn.Identity' or 'MyBlock'" },
+        { name: 'args', type: 'string', label: 'Arguments', default: '', description: "e.g., 'dim=128, p=0.1'" }
+    ]
+  },
+  [LayerType.IDENTITY]: {
+    type: LayerType.IDENTITY,
+    label: 'Identity',
+    description: 'Passthrough layer',
+    category: 'Utility',
+    parameters: []
+  },
+
+  // --- NORMALIZATION ---
+  [LayerType.BATCHNORM]: {
+    type: LayerType.BATCHNORM,
+    label: 'BatchNorm2D',
+    description: 'Batch Normalization',
+    category: 'Normalization',
+    parameters: [
+      { name: 'num_features', type: 'number', label: 'Num Features', default: 64 }
+    ]
+  },
+  [LayerType.LAYERNORM]: {
+    type: LayerType.LAYERNORM,
+    label: 'LayerNorm',
+    description: 'Layer Normalization',
+    category: 'Normalization',
+    parameters: [
+      { name: 'normalized_shape', type: 'string', label: 'Norm Shape', default: '128' }
+    ]
+  },
+  [LayerType.INSTANCENORM]: {
+    type: LayerType.INSTANCENORM,
+    label: 'InstanceNorm2d',
+    description: 'Instance Normalization',
+    category: 'Normalization',
+    parameters: [
+       { name: 'num_features', type: 'number', label: 'Num Features', default: 64 },
+       { name: 'affine', type: 'boolean', label: 'Learnable', default: false }
+    ]
+  },
+
+  // --- TRANSFORMER ---
+  [LayerType.ATTENTION]: {
+    type: LayerType.ATTENTION,
+    label: 'Self Attention',
+    description: 'Multi-Head Self Attention',
+    category: 'Transformer',
+    parameters: [
+      { name: 'embed_dim', type: 'number', label: 'Embed Dim', default: 512 },
+      { name: 'num_heads', type: 'number', label: 'Num Heads', default: 8 },
+      { name: 'dropout', type: 'number', label: 'Dropout', default: 0.1 },
+      { name: 'batch_first', type: 'boolean', label: 'Batch First', default: true }
+    ]
+  },
+  [LayerType.CROSS_ATTENTION]: {
+    type: LayerType.CROSS_ATTENTION,
+    label: 'Cross Attention',
+    description: 'Attention between two sequences',
+    category: 'Transformer',
+    parameters: [
+      { name: 'embed_dim', type: 'number', label: 'Embed Dim', default: 512 },
+      { name: 'num_heads', type: 'number', label: 'Num Heads', default: 8 },
+      { name: 'dropout', type: 'number', label: 'Dropout', default: 0.1 }
+    ]
+  },
+  [LayerType.TRANSFORMER_ENCODER]: {
+    type: LayerType.TRANSFORMER_ENCODER,
+    label: 'Encoder Layer',
+    description: 'Standard Transformer Encoder',
+    category: 'Transformer',
+    parameters: [
+      { name: 'd_model', type: 'number', label: 'D Model', default: 512 },
+      { name: 'nhead', type: 'number', label: 'Heads', default: 8 },
+      { name: 'dim_feedforward', type: 'number', label: 'FF Dim', default: 2048 },
+      { name: 'dropout', type: 'number', label: 'Dropout', default: 0.1 },
+      { name: 'activation', type: 'select', label: 'Activation', default: 'relu', options: ['relu', 'gelu'] }
+    ]
+  },
+  [LayerType.TRANSFORMER_DECODER]: {
+    type: LayerType.TRANSFORMER_DECODER,
+    label: 'Decoder Layer',
+    description: 'Standard Transformer Decoder',
+    category: 'Transformer',
+    parameters: [
+      { name: 'd_model', type: 'number', label: 'D Model', default: 512 },
+      { name: 'nhead', type: 'number', label: 'Heads', default: 8 },
+      { name: 'dim_feedforward', type: 'number', label: 'FF Dim', default: 2048 },
+      { name: 'dropout', type: 'number', label: 'Dropout', default: 0.1 }
+    ]
+  },
+  [LayerType.TRANSFORMER_BLOCK]: {
+    type: LayerType.TRANSFORMER_BLOCK,
+    label: 'Transformer Block',
+    description: 'Generic Block',
+    category: 'Transformer',
+    parameters: [
+      { name: 'd_model', type: 'number', label: 'D Model', default: 512 },
+      { name: 'nhead', type: 'number', label: 'Heads', default: 8 },
+      { name: 'dim_feedforward', type: 'number', label: 'FF Dim', default: 2048 }
+    ]
+  },
+};
+
+export const INITIAL_NODES = [
+  {
+    id: '1',
+    type: 'custom',
+    position: { x: 250, y: 50 },
+    data: { label: 'Input Data', type: LayerType.INPUT, params: LAYER_DEFINITIONS[LayerType.INPUT].parameters.reduce((acc, p) => ({...acc, [p.name]: p.default}), {}) }
+  },
+  {
+    id: '2',
+    type: 'custom',
+    position: { x: 250, y: 200 },
+    data: { label: 'Conv2D', type: LayerType.CONV2D, params: LAYER_DEFINITIONS[LayerType.CONV2D].parameters.reduce((acc, p) => ({...acc, [p.name]: p.default}), {}) }
+  },
+  {
+    id: '3',
+    type: 'custom',
+    position: { x: 250, y: 350 },
+    data: { label: 'ReLU', type: LayerType.RELU, params: {} }
+  }
+];
+
+export const INITIAL_EDGES = [
+  { id: 'e1-2', source: '1', target: '2', animated: true, style: { stroke: '#94a3b8' } },
+  { id: 'e2-3', source: '2', target: '3', animated: true, style: { stroke: '#94a3b8' } }
+];
+
+export const TEMPLATES: Record<string, GraphTemplate> = {
+  'gpt_style': {
+    id: 'gpt_style',
+    name: 'LLM (GPT Style)',
+    description: 'Decoder-only Transformer with RoPE & RMSNorm.',
+    nodes: [
+      { id: 'in', type: 'custom', position: {x: 300, y: 0}, data: {label: 'Token Input', type: LayerType.INPUT, params: {modality: 'Text', shape: '128', dtype: 'int64'}} },
+      { id: 'emb', type: 'custom', position: {x: 300, y: 100}, data: {label: 'Token Embed', type: LayerType.EMBEDDING, params: {num_embeddings: 50257, embedding_dim: 768}} },
+      { id: 'rope', type: 'custom', position: {x: 300, y: 200}, data: {label: 'RoPE', type: LayerType.ROPE, params: {dim: 64}} },
+      { id: 'blk1', type: 'custom', position: {x: 300, y: 300}, data: {label: 'Decoder Block 1', type: LayerType.TRANSFORMER_DECODER, params: {d_model: 768, nhead: 12}} },
+      { id: 'blk2', type: 'custom', position: {x: 300, y: 400}, data: {label: 'Decoder Block 2', type: LayerType.TRANSFORMER_DECODER, params: {d_model: 768, nhead: 12}} },
+      { id: 'ln', type: 'custom', position: {x: 300, y: 500}, data: {label: 'RMSNorm', type: LayerType.RMSNORM, params: {dim: 768}} },
+      { id: 'head', type: 'custom', position: {x: 300, y: 600}, data: {label: 'LM Head', type: LayerType.LINEAR, params: {out_features: 50257, bias: false}} }
+    ],
+    edges: [
+      { id: '1', source: 'in', target: 'emb' }, { id: '2', source: 'emb', target: 'rope' }, 
+      { id: '3', source: 'rope', target: 'blk1' }, { id: '4', source: 'blk1', target: 'blk2' },
+      { id: '5', source: 'blk2', target: 'ln' }, { id: '6', source: 'ln', target: 'head' }
+    ]
+  },
+  'bert_encoder': {
+    id: 'bert_encoder',
+    name: 'BERT (Encoder)',
+    description: 'Bidirectional Transformer Encoder (NLP)',
+    nodes: [
+      { id: 'in', type: 'custom', position: {x: 300, y: 0}, data: {label: 'Token Input', type: LayerType.INPUT, params: {modality: 'Text'}} },
+      { id: 'emb', type: 'custom', position: {x: 300, y: 100}, data: {label: 'Embedding', type: LayerType.EMBEDDING, params: {embedding_dim: 768}} },
+      { id: 'pos', type: 'custom', position: {x: 300, y: 200}, data: {label: 'Pos Embed', type: LayerType.POS_EMBED, params: {embedding_dim: 768}} },
+      { id: 'add', type: 'custom', position: {x: 300, y: 300}, data: {label: 'Combine', type: LayerType.ADD, params: {}} },
+      { id: 'enc1', type: 'custom', position: {x: 300, y: 400}, data: {label: 'Encoder Layer 1', type: LayerType.TRANSFORMER_ENCODER, params: {d_model: 768, nhead: 12}} },
+      { id: 'enc2', type: 'custom', position: {x: 300, y: 500}, data: {label: 'Encoder Layer 2', type: LayerType.TRANSFORMER_ENCODER, params: {d_model: 768, nhead: 12}} },
+      { id: 'pool', type: 'custom', position: {x: 300, y: 600}, data: {label: 'Pooler', type: LayerType.LINEAR, params: {out_features: 768}} },
+    ],
+    edges: [
+      { id: '1', source: 'in', target: 'emb' }, { id: '2', source: 'emb', target: 'add' },
+      { id: '3', source: 'pos', target: 'add' }, { id: '4', source: 'add', target: 'enc1' },
+      { id: '5', source: 'enc1', target: 'enc2' }, { id: '6', source: 'enc2', target: 'pool' }
+    ]
+  },
+  'autoencoder_conv': {
+    id: 'autoencoder_conv',
+    name: 'Autoencoder (Conv)',
+    description: 'Image compression and reconstruction.',
+    nodes: [
+      { id: 'in', type: 'custom', position: {x: 300, y: 0}, data: {label: 'Image', type: LayerType.INPUT, params: {modality: 'Image', shape: '1,28,28'}} },
+      // Encoder
+      { id: 'enc1', type: 'custom', position: {x: 300, y: 100}, data: {label: 'Conv 1', type: LayerType.CONV2D, params: {out_channels: 16, kernel_size: 3, stride: 2, padding: 1}} },
+      { id: 'act1', type: 'custom', position: {x: 300, y: 180}, data: {label: 'ReLU', type: LayerType.RELU, params: {}} },
+      { id: 'enc2', type: 'custom', position: {x: 300, y: 260}, data: {label: 'Conv 2', type: LayerType.CONV2D, params: {out_channels: 32, kernel_size: 3, stride: 2, padding: 1}} },
+      { id: 'act2', type: 'custom', position: {x: 300, y: 340}, data: {label: 'ReLU', type: LayerType.RELU, params: {}} },
+      // Decoder
+      { id: 'dec1', type: 'custom', position: {x: 300, y: 420}, data: {label: 'Deconv 1', type: LayerType.CONV_TRANSPOSE2D, params: {out_channels: 16, kernel_size: 3, stride: 2, padding: 1}} }, // Needs output_padding often in code but usually inferred
+      { id: 'act3', type: 'custom', position: {x: 300, y: 500}, data: {label: 'ReLU', type: LayerType.RELU, params: {}} },
+      { id: 'dec2', type: 'custom', position: {x: 300, y: 580}, data: {label: 'Deconv 2', type: LayerType.CONV_TRANSPOSE2D, params: {out_channels: 1, kernel_size: 3, stride: 2, padding: 1}} },
+      { id: 'sig', type: 'custom', position: {x: 300, y: 660}, data: {label: 'Sigmoid', type: LayerType.SIGMOID, params: {}} }
+    ],
+    edges: [
+       { id: '1', source: 'in', target: 'enc1' }, { id: '2', source: 'enc1', target: 'act1' }, { id: '3', source: 'act1', target: 'enc2' },
+       { id: '4', source: 'enc2', target: 'act2' }, { id: '5', source: 'act2', target: 'dec1' }, { id: '6', source: 'dec1', target: 'act3' },
+       { id: '7', source: 'act3', target: 'dec2' }, { id: '8', source: 'dec2', target: 'sig' }
+    ]
+  },
+  'resnet_mini': {
+    id: 'resnet_mini',
+    name: 'ResNet Block (Mini)',
+    description: 'CNN with Skip Connection.',
+    nodes: [
+       { id: 'in', type: 'custom', position: {x: 200, y: 0}, data: {label: 'Input', type: LayerType.INPUT, params: {modality: 'Image'}} },
+       { id: 'conv1', type: 'custom', position: {x: 200, y: 100}, data: {label: 'Conv1', type: LayerType.CONV2D, params: {out_channels: 64}} },
+       { id: 'relu1', type: 'custom', position: {x: 200, y: 200}, data: {label: 'ReLU', type: LayerType.RELU, params: {}} },
+       { id: 'conv2', type: 'custom', position: {x: 200, y: 300}, data: {label: 'Conv2', type: LayerType.CONV2D, params: {out_channels: 64}} },
+       
+       // Skip connection path handled by edges, but node layout implies it
+       { id: 'add', type: 'custom', position: {x: 200, y: 450}, data: {label: 'Residual Add', type: LayerType.ADD, params: {}} },
+       { id: 'relu2', type: 'custom', position: {x: 200, y: 550}, data: {label: 'Final ReLU', type: LayerType.RELU, params: {}} }
+    ],
+    edges: [
+      { id: '1', source: 'in', target: 'conv1' }, { id: '2', source: 'conv1', target: 'relu1' },
+      { id: '3', source: 'relu1', target: 'conv2' }, { id: '4', source: 'conv2', target: 'add' },
+      { id: '5', source: 'in', target: 'add' }, // Skip connection
+      { id: '6', source: 'add', target: 'relu2' }
+    ]
+  },
+  'moe_transformer': {
+    id: 'moe_transformer',
+    name: 'Mixture of Experts',
+    description: 'Sparse MoE model with routing.',
+    nodes: [
+      { id: 'in', type: 'custom', position: {x: 300, y: 0}, data: {label: 'Input', type: LayerType.INPUT, params: {modality: 'Text'}} },
+      { id: 'emb', type: 'custom', position: {x: 300, y: 100}, data: {label: 'Embedding', type: LayerType.EMBEDDING, params: {}} },
+      { id: 'att', type: 'custom', position: {x: 300, y: 200}, data: {label: 'Self Attention', type: LayerType.ATTENTION, params: {}} },
+      { id: 'moe', type: 'custom', position: {x: 300, y: 300}, data: {label: 'Sparse MoE Block', type: LayerType.MOE_BLOCK, params: {num_experts: 8, top_k: 2}} },
+      { id: 'norm', type: 'custom', position: {x: 300, y: 400}, data: {label: 'RMSNorm', type: LayerType.RMSNORM, params: {}} },
+      { id: 'out', type: 'custom', position: {x: 300, y: 500}, data: {label: 'Output', type: LayerType.LINEAR, params: {}} }
+    ],
+    edges: [
+      { id: '1', source: 'in', target: 'emb' }, { id: '2', source: 'emb', target: 'att' },
+      { id: '3', source: 'att', target: 'moe' }, { id: '4', source: 'moe', target: 'norm' },
+      { id: '5', source: 'norm', target: 'out' }
+    ]
+  },
+  'vlm_llava': {
+    id: 'vlm_llava',
+    name: 'VLM (LlaVA Style)',
+    description: 'Visual Language Model connecting Vision Encoder to LLM.',
+    nodes: [
+      // Vision Branch
+      { id: 'img', type: 'custom', position: {x: 100, y: 0}, data: {label: 'Image Input', type: LayerType.INPUT, params: {modality: 'Image'}} },
+      { id: 'patch', type: 'custom', position: {x: 100, y: 100}, data: {label: 'Patch Embed', type: LayerType.PATCH_EMBED, params: {patch_size: 14}} },
+      { id: 'vit', type: 'custom', position: {x: 100, y: 200}, data: {label: 'ViT Encoder', type: LayerType.TRANSFORMER_ENCODER, params: {}} },
+      { id: 'proj', type: 'custom', position: {x: 100, y: 300}, data: {label: 'Projection', type: LayerType.LINEAR, params: {out_features: 4096}} },
+
+      // Text Branch
+      { id: 'txt', type: 'custom', position: {x: 500, y: 0}, data: {label: 'Text Prompts', type: LayerType.INPUT, params: {modality: 'Text'}} },
+      { id: 'temb', type: 'custom', position: {x: 500, y: 200}, data: {label: 'Text Embed', type: LayerType.EMBEDDING, params: {embedding_dim: 4096}} },
+      
+      // Merge
+      { id: 'cat', type: 'custom', position: {x: 300, y: 400}, data: {label: 'Concat Tokens', type: LayerType.CONCAT, params: {dim: 1}} },
+      { id: 'llm', type: 'custom', position: {x: 300, y: 500}, data: {label: 'LLM Decoder', type: LayerType.TRANSFORMER_DECODER, params: {d_model: 4096}} },
+      { id: 'out', type: 'custom', position: {x: 300, y: 600}, data: {label: 'Response', type: LayerType.OUTPUT, params: {}} }
+    ],
+    edges: [
+      { id: '1', source: 'img', target: 'patch' }, { id: '2', source: 'patch', target: 'vit' }, { id: '3', source: 'vit', target: 'proj' },
+      { id: '4', source: 'txt', target: 'temb' },
+      { id: '5', source: 'proj', target: 'cat' }, { id: '6', source: 'temb', target: 'cat' },
+      { id: '7', source: 'cat', target: 'llm' }, { id: '8', source: 'llm', target: 'out' }
+    ]
+  },
+  'sam_model': {
+    id: 'sam_model',
+    name: 'Segment Anything (SAM)',
+    description: 'Image Encoder + Prompt Encoder + Mask Decoder.',
+    nodes: [
+      { id: 'img', type: 'custom', position: {x: 100, y: 0}, data: {label: 'Image', type: LayerType.INPUT, params: {modality: 'Image'}} },
+      { id: 'enc', type: 'custom', position: {x: 100, y: 150}, data: {label: 'Image Encoder (ViT)', type: LayerType.TRANSFORMER_ENCODER, params: {}} },
+      
+      { id: 'prm', type: 'custom', position: {x: 500, y: 0}, data: {label: 'Points/Boxes', type: LayerType.INPUT, params: {modality: 'Tensor'}} },
+      { id: 'penc', type: 'custom', position: {x: 500, y: 150}, data: {label: 'Prompt Enc', type: LayerType.SAM_PROMPT_ENCODER, params: {}} },
+
+      { id: 'dec', type: 'custom', position: {x: 300, y: 300}, data: {label: 'Mask Decoder', type: LayerType.SAM_MASK_DECODER, params: {}} },
+      { id: 'out', type: 'custom', position: {x: 300, y: 400}, data: {label: 'Masks', type: LayerType.OUTPUT, params: {}} }
+    ],
+    edges: [
+      { id: '1', source: 'img', target: 'enc' }, { id: '2', source: 'enc', target: 'dec' },
+      { id: '3', source: 'prm', target: 'penc' }, { id: '4', source: 'penc', target: 'dec' },
+      { id: '5', source: 'dec', target: 'out' }
+    ]
+  },
+  'lam_agent': {
+    id: 'lam_agent',
+    name: 'Large Action Model (LAM)',
+    description: 'LLM backbone with Action Head for agents.',
+    nodes: [
+      { id: 'state', type: 'custom', position: {x: 300, y: 0}, data: {label: 'Env State', type: LayerType.INPUT, params: {modality: 'State'}} },
+      { id: 'llm', type: 'custom', position: {x: 300, y: 150}, data: {label: 'LLM Backbone', type: LayerType.TRANSFORMER_DECODER, params: {d_model: 1024}} },
+      { id: 'head', type: 'custom', position: {x: 300, y: 300}, data: {label: 'Action Head', type: LayerType.ACTION_HEAD, params: {num_actions: 50, action_type: 'Discrete'}} },
+      { id: 'out', type: 'custom', position: {x: 300, y: 400}, data: {label: 'Action Logits', type: LayerType.OUTPUT, params: {}} }
+    ],
+    edges: [
+      { id: '1', source: 'state', target: 'llm' }, { id: '2', source: 'llm', target: 'head' }, { id: '3', source: 'head', target: 'out' }
+    ]
+  },
+  'lcm_diff': {
+    id: 'lcm_diff',
+    name: 'Latent Consistency (LCM)',
+    description: 'Diffusion backbone with Time Embeddings.',
+    nodes: [
+      { id: 'lat', type: 'custom', position: {x: 200, y: 0}, data: {label: 'Latent Input', type: LayerType.INPUT, params: {modality: 'Latent', shape: '4,64,64'}} },
+      { id: 'time', type: 'custom', position: {x: 500, y: 0}, data: {label: 'Time Step', type: LayerType.INPUT, params: {modality: 'Tensor', shape: '1'}} },
+      { id: 'temb', type: 'custom', position: {x: 500, y: 100}, data: {label: 'Time Embed', type: LayerType.TIME_EMBEDDING, params: {}} },
+      
+      { id: 'cat', type: 'custom', position: {x: 350, y: 250}, data: {label: 'Inject Time', type: LayerType.ADD, params: {}} },
+      { id: 'unet', type: 'custom', position: {x: 350, y: 350}, data: {label: 'UNet Block', type: LayerType.CONV2D, params: {}} },
+      { id: 'out', type: 'custom', position: {x: 350, y: 450}, data: {label: 'Denoised', type: LayerType.OUTPUT, params: {}} }
+    ],
+    edges: [
+      { id: '1', source: 'lat', target: 'cat' },
+      { id: '2', source: 'time', target: 'temb' }, { id: '3', source: 'temb', target: 'cat' },
+      { id: '4', source: 'cat', target: 'unet' }, { id: '5', source: 'unet', target: 'out' }
+    ]
+  }
+};

index.html

@@ -0,0 +1,44 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Architecture Agents</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/reactflow@11.10.1/dist/style.min.css" />
+    <style>
+      /* Custom scrollbar for dark theme */
+      ::-webkit-scrollbar {
+        width: 8px;
+        height: 8px;
+      }
+      ::-webkit-scrollbar-track {
+        background: #0f172a; 
+      }
+      ::-webkit-scrollbar-thumb {
+        background: #334155; 
+        border-radius: 4px;
+      }
+      ::-webkit-scrollbar-thumb:hover {
+        background: #475569; 
+      }
+    </style>
+  <script type="importmap">
+{
+  "imports": {
+    "react-dom/": "https://aistudiocdn.com/react-dom@^19.2.1/",
+    "react/": "https://aistudiocdn.com/react@^19.2.1/",
+    "react": "https://aistudiocdn.com/react@^19.2.1",
+    "@google/genai": "https://aistudiocdn.com/@google/genai@^1.31.0",
+    "lucide-react": "https://aistudiocdn.com/lucide-react@^0.556.0",
+    "reactflow": "https://aistudiocdn.com/reactflow@^11.11.4"
+  }
+}
+</script>
+<link rel="stylesheet" href="/index.css">
+</head>
+  <body class="bg-slate-950 text-slate-100 overflow-hidden">
+    <div id="root"></div>
+  <script type="module" src="/index.tsx"></script>
+</body>
+</html>

index.tsx

@@ -0,0 +1,15 @@
+import React from 'react';
+import ReactDOM from 'react-dom/client';
+import App from './App';
+
+const rootElement = document.getElementById('root');
+if (!rootElement) {
+  throw new Error("Could not find root element to mount to");
+}
+
+const root = ReactDOM.createRoot(rootElement);
+root.render(
+  <React.StrictMode>
+    <App />
+  </React.StrictMode>
+);

metadata.json

@@ -0,0 +1,5 @@
+{
+  "name": "Architecture Agents",
+  "description": "Visual AI Architect for Deep Learning. Design with agents, build with code.",
+  "requestFramePermissions": []
+}

package.json

@@ -0,0 +1,24 @@
+{
+  "name": "architecture-agents",
+  "private": true,
+  "version": "0.0.0",
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "react-dom": "^19.2.1",
+    "react": "^19.2.1",
+    "@google/genai": "^1.31.0",
+    "lucide-react": "^0.556.0",
+    "reactflow": "^11.11.4"
+  },
+  "devDependencies": {
+    "@types/node": "^22.14.0",
+    "@vitejs/plugin-react": "^5.0.0",
+    "typescript": "~5.8.2",
+    "vite": "^6.2.0"
+  }
+}

tsconfig.json

@@ -0,0 +1,29 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "experimentalDecorators": true,
+    "useDefineForClassFields": false,
+    "module": "ESNext",
+    "lib": [
+      "ES2022",
+      "DOM",
+      "DOM.Iterable"
+    ],
+    "skipLibCheck": true,
+    "types": [
+      "node"
+    ],
+    "moduleResolution": "bundler",
+    "isolatedModules": true,
+    "moduleDetection": "force",
+    "allowJs": true,
+    "jsx": "react-jsx",
+    "paths": {
+      "@/*": [
+        "./*"
+      ]
+    },
+    "allowImportingTsExtensions": true,
+    "noEmit": true
+  }
+}

types.ts

@@ -0,0 +1,84 @@
+export enum LayerType {
+  INPUT = 'Input',
+  LINEAR = 'Linear',
+  CONV1D = 'Conv1D',
+  CONV2D = 'Conv2D',
+  CONV_TRANSPOSE2D = 'ConvTranspose2D',
+  RELU = 'ReLU',
+  LEAKYRELU = 'LeakyReLU',
+  GELU = 'GELU',
+  SILU = 'SiLU',
+  SWIGLU = 'SwiGLU',
+  SIGMOID = 'Sigmoid',
+  TANH = 'Tanh',
+  MAXPOOL = 'MaxPool2D',
+  AVGPOOL = 'AvgPool2D',
+  ADAPTIVEAVGPOOL = 'AdaptiveAvgPool2D',
+  UPSAMPLE = 'Upsample',
+  DROPOUT = 'Dropout',
+  BATCHNORM = 'BatchNorm2D',
+  LAYERNORM = 'LayerNorm',
+  RMSNORM = 'RMSNorm',
+  INSTANCENORM = 'InstanceNorm2d',
+  FLATTEN = 'Flatten',
+  EMBEDDING = 'Embedding',
+  PATCH_EMBED = 'PatchEmbed',
+  POS_EMBED = 'PositionalEmbed',
+  ROPE = 'RotaryPosEmbed',
+  TIME_EMBEDDING = 'TimeEmbedding',
+  LSTM = 'LSTM',
+  GRU = 'GRU',
+  ATTENTION = 'SelfAttention',
+  CROSS_ATTENTION = 'CrossAttention',
+  TRANSFORMER_BLOCK = 'TransformerBlock',
+  TRANSFORMER_ENCODER = 'TransformerEncoderLayer',
+  TRANSFORMER_DECODER = 'TransformerDecoderLayer',
+  MOE_BLOCK = 'SparseMoEBlock',
+  ACTION_HEAD = 'ActionHead',
+  SAM_PROMPT_ENCODER = 'SamPromptEncoder',
+  SAM_MASK_DECODER = 'SamMaskDecoder',
+  CONCAT = 'Concat',
+  ADD = 'Add',
+  OUTPUT = 'Output',
+  CUSTOM = 'CustomLayer',
+  IDENTITY = 'Identity'
+}
+
+export interface LayerParameter {
+  name: string;
+  type: 'number' | 'string' | 'select' | 'boolean';
+  label: string;
+  default: any;
+  options?: string[]; // For select types
+  description?: string;
+}
+
+export interface LayerDefinition {
+  type: LayerType;
+  label: string;
+  description: string;
+  category: 'Core' | 'Convolution' | 'Recurrent' | 'Normalization' | 'Transformer' | 'GenAI' | 'Utility' | 'Merge';
+  parameters: LayerParameter[];
+}
+
+export interface NodeData {
+  label: string;
+  type: LayerType;
+  params: Record<string, any>;
+  isValid?: boolean;
+  error?: string;
+  onDelete?: (id: string) => void;
+}
+
+export interface GraphTemplate {
+  id: string;
+  name: string;
+  description: string;
+  nodes: any[];
+  edges: any[];
+}
+
+export interface GeneratedCodeResponse {
+  code: string;
+  explanation: string;
+}

vite.config.ts

@@ -0,0 +1,23 @@
+import path from 'path';
+import { defineConfig, loadEnv } from 'vite';
+import react from '@vitejs/plugin-react';
+
+export default defineConfig(({ mode }) => {
+    const env = loadEnv(mode, '.', '');
+    return {
+      server: {
+        port: 3000,
+        host: '0.0.0.0',
+      },
+      plugins: [react()],
+      define: {
+        'process.env.API_KEY': JSON.stringify(env.GEMINI_API_KEY),
+        'process.env.GEMINI_API_KEY': JSON.stringify(env.GEMINI_API_KEY)
+      },
+      resolve: {
+        alias: {
+          '@': path.resolve(__dirname, '.'),
+        }
+      }
+    };
+});