LOCALHOST1978/Modules/SSK_Demucs_MLX/ARCHITECTURE.md

SSK_Demucs_MLX Architecture

Version Control and State Tracking Rule: This repository exclusively uses lore for version control, state management, and documentation tracking. Do NOT use git.

---

1. Overview and Purpose

SSK_Demucs_MLX is a standalone Swift Package module within SovereignSurvivalKit providing hardware-accelerated music source separation using the HTDemucs (Hybrid Transformer Demucs) neural architecture.

Built on Apple's MLX array framework via mlx-swift, the module decomposes stereo mix tracks (44.1 kHz) into 4 isolated audio stems:

• Vocals
• Drums
• Bass
• Other

---

2. Package Specifications and Dependencies

Defined in Package.swift:

• Target Name: SSK_Demucs_MLX
• Swift Tools Version: 6.4
• Supported Platforms: macOS 27.0+, iOS 27.0+
• External Package Dependencies:
  • mlx-swift (0.10.0+) — Apple MLX tensor computing framework:
    • MLX: Tensor operations, array allocations, and execution evaluation (eval).
    • MLXNN: Neural network modules, activations, and normalizations.
    • MLXRandom: Random state generators.
    • MLXOptimizers: Optimization algorithms.

---

3. Component Architecture and Layers

                 +---------------------------------+
                 |      Input Mix (Batch, 2, T)    |
                 +---------------------------------+
                                  │
                          ┌───────┴───────┐
                          │               │
                          ▼               ▼
                 +----------------+  +-------------------+
                 | RealSTFT (Time |  | Raw Time Branch   |
                 | -> Spectrogram)|  +-------------------+
                 +----------------+            │
                          │                    │
                          ▼                    │
                 +----------------+            │
                 |  HEncLayer x 4 |            │
                 | (Encoder Stack)|            │
                 +----------------+            │
                          │                    │
                          ┌───────┐   ┌────────┘
                                  │   │
                                  ▼   ▼
                 +---------------------------------+
                 |        CrossTransformer         |
                 |  (Bottleneck Cross-Attention)   |
                 +---------------------------------+
                                  │
                          ┌───────┘
                          │
                          ▼
                 +----------------+
                 |  HDecLayer x 4 | (With Residual Skip Connections)
                 | (Decoder Stack)|
                 +----------------+
                          │
                          ▼
                 +----------------+
                 |   RealISTFT    |
                 |  (Inverse FFT) |
                 +----------------+
                          │
                          ▼
                 +---------------------------------+
                 |   Output Stems (Batch, 4, 2, T) |
                 +---------------------------------+

Core Source Files

Component / File	Description
DemucsModel.swift	Primary Module orchestrating STFT, encoder pyramid, cross-transformer bottleneck, decoder stack, and ISTFT synthesis.
WeightLoader.swift	Downloads htdemucs.safetensors from Hugging Face via URLSession and maps flat weight keys into model parameters.
Layers/Transforms.swift	Implements RealSTFT and RealISTFT using Hann windowing, reflect padding, and rfft/irfft operations.
Layers/HEncLayer.swift	Hierarchical Encoder Layer featuring 1D convolutions, GroupNorm, GELU activation, DConv, and GLU rewrite blocks.
Layers/HDecLayer.swift	Hierarchical Decoder Layer featuring transposed 1D convolutions, residual skip connections from encoder levels, and GLU blocks.
Layers/DConv.swift	Deep residual dilated convolution block incorporating multi-rate dilations, LayerScale, and LocalState attention.
Layers/CrossTransformer.swift	Bottleneck transformer block cross-attending time-domain and frequency-domain representations with sinusoidal positional encodings.

---

4. Technical Specifications

• FFT Window: 4,096 samples (nFft = 4096), Hop Length: 1,024 samples (hopLength = 1024).
• Segment Length: 10 seconds at 44.1 kHz (441,000 samples).
• Channel Scaling: Encoder channel depth increases across [48, 96, 192, 384].
• Skip Connections: Features from each encoder level are cached and concatenated into matching decoder levels before transposed convolutions.

---

5. Architectural Position

                        SSK_Hardware
                               ▲
                               │
                        SSK_Demucs_MLX
                               ▲
                               │
                           mlx-swift