mirror of
https://github.com/hanzoai/live.git
synced 2026-10-30 18:29:41 +00:00
zeekayandHanzo Dev
12b64ce59f
Completes the interrupted kill/shadcn work (81ffa84):
- vite: GUI_TARGET/TAMAGUI_TARGET=web defines + react-native/react-native-svg
web aliases (@hanzogui/react-native-svg) — fixes the unresolved
react-native-svg import from @hanzogui/lucide-icons-2
- deps: @hanzo/gui 8.0.2, @hanzo/ui 8.0.48, @hanzogui/lucide-icons-2 8.0.2;
react/react-dom aligned at 19.2.8
- docs/: drop cdn.tailwindcss.com from all 3 static pages — rewritten on
plain-CSS tokens (docs/docs.css), same visual design
- kill last residue: vscode tailwind extension rec, stale LLM.md stack notes
- gui-env.d.ts: eslint disable for the intentional declaration-merge pattern
typescript stays 5.8.3: TS7.0.2 builds green here but typescript-eslint
(latest 8.66.0, peer <6.1.0) crashes against it — lint is part of this
repo's gate, so 5.x is correct until typescript-eslint supports TS7.
Verified: tsc -b && vite build green; playwright render of the built app —
full UI styled (cards/selects/switches/sliders), zero console errors.
Co-authored-by: Hanzo Dev <[email protected]>
402 lines
17 KiB
HTML
402 lines
17 KiB
HTML
<!DOCTYPE html>
|
|
<html lang="en">
|
|
<head>
|
|
<meta charset="UTF-8">
|
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
|
<title>MLX-Native WAN Implementation - Hanzo Live</title>
|
|
<link rel="stylesheet" href="docs.css">
|
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css">
|
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github-dark.min.css">
|
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>
|
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/languages/python.min.js"></script>
|
|
</head>
|
|
<body>
|
|
<div class="container mid">
|
|
<!-- Back Navigation -->
|
|
<div class="mb-8">
|
|
<a href="index.html" class="link"><i class="fas fa-arrow-left"></i> Back to Documentation</a>
|
|
</div>
|
|
|
|
<!-- Header -->
|
|
<header class="mb-12">
|
|
<div class="row mb-6">
|
|
<i class="fab fa-apple blue" style="font-size:3.75rem"></i>
|
|
<div>
|
|
<h1 class="gradient-text-blue" style="font-size:3rem">MLX-Native WAN Implementation</h1>
|
|
<p class="muted mb-0" style="font-size:1.25rem">Complete Apple Silicon optimization with 2-3x speedup</p>
|
|
</div>
|
|
</div>
|
|
<div class="callout green">
|
|
<p class="mb-0">
|
|
<i class="fas fa-check-circle"></i>
|
|
<strong>Production Ready!</strong> All 9 tests passing. Full feature parity with PyTorch.
|
|
</p>
|
|
</div>
|
|
</header>
|
|
|
|
<!-- Overview -->
|
|
<section class="mb-12">
|
|
<h2 class="purple">Overview</h2>
|
|
<div class="card solid">
|
|
<p class="muted">
|
|
The MLX implementation provides significant performance improvements over the PyTorch MPS (Metal Performance Shaders)
|
|
backend by using Apple's native MLX framework. MLX is specifically designed for Apple Silicon and offers:
|
|
</p>
|
|
<ul>
|
|
<li><strong style="color:var(--text)">2-3x faster inference</strong> compared to PyTorch MPS</li>
|
|
<li>Native Metal optimization without CPU-GPU synchronization overhead</li>
|
|
<li>Better memory efficiency through unified memory architecture</li>
|
|
<li>Seamless integration with Apple's Neural Engine</li>
|
|
</ul>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Performance Comparison -->
|
|
<section class="mb-12">
|
|
<h2 class="pink">Performance Comparison</h2>
|
|
<div class="card solid">
|
|
<h3 class="mb-6">PyTorch MPS vs MLX</h3>
|
|
<div class="table-wrap">
|
|
<table>
|
|
<thead>
|
|
<tr>
|
|
<th class="purple">Operation</th>
|
|
<th class="yellow">PyTorch MPS</th>
|
|
<th class="green">MLX Native</th>
|
|
<th class="blue">Speedup</th>
|
|
</tr>
|
|
</thead>
|
|
<tbody>
|
|
<tr>
|
|
<td>Attention (256x256)</td>
|
|
<td>~15ms</td>
|
|
<td>~6ms</td>
|
|
<td class="bold green">2.5x</td>
|
|
</tr>
|
|
<tr>
|
|
<td>RoPE Apply</td>
|
|
<td>~8ms</td>
|
|
<td>~3ms</td>
|
|
<td class="bold green">2.7x</td>
|
|
</tr>
|
|
<tr>
|
|
<td>RMS Norm</td>
|
|
<td>~2ms</td>
|
|
<td>~0.8ms</td>
|
|
<td class="bold green">2.5x</td>
|
|
</tr>
|
|
<tr>
|
|
<td class="bold">Full Forward Pass</td>
|
|
<td class="bold">~150ms</td>
|
|
<td class="bold">~60ms</td>
|
|
<td class="bold green">2.5x</td>
|
|
</tr>
|
|
</tbody>
|
|
</table>
|
|
</div>
|
|
<p class="small dim mt-4 mb-0">
|
|
<i class="fas fa-info-circle"></i>
|
|
Benchmarks on M2 Max with 128x128 resolution, 16 frames
|
|
</p>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Why MLX is Faster -->
|
|
<section class="mb-12">
|
|
<h2 class="blue">Why MLX is Faster</h2>
|
|
<div class="grid cols-2">
|
|
<div class="card solid">
|
|
<div class="icon purple" style="font-size:1.875rem"><i class="fas fa-memory"></i></div>
|
|
<h3>No CPU-GPU Sync</h3>
|
|
<p class="muted mb-0">PyTorch MPS requires synchronization between CPU and GPU, while MLX uses unified memory</p>
|
|
</div>
|
|
|
|
<div class="card solid">
|
|
<div class="icon pink" style="font-size:1.875rem"><i class="fas fa-bolt"></i></div>
|
|
<h3>Metal-Native</h3>
|
|
<p class="muted mb-0">MLX compiles directly to Metal shaders, avoiding PyTorch's abstraction overhead</p>
|
|
</div>
|
|
|
|
<div class="card solid">
|
|
<div class="icon blue" style="font-size:1.875rem"><i class="fab fa-apple"></i></div>
|
|
<h3>Apple Silicon First</h3>
|
|
<p class="muted mb-0">Designed specifically for M-series chips, not a general cross-platform solution</p>
|
|
</div>
|
|
|
|
<div class="card solid">
|
|
<div class="icon green" style="font-size:1.875rem"><i class="fas fa-database"></i></div>
|
|
<h3>Better Memory</h3>
|
|
<p class="muted mb-0">Unified memory architecture reduces data copying between CPU and GPU</p>
|
|
</div>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Usage -->
|
|
<section class="mb-12">
|
|
<h2 class="green">Quick Start</h2>
|
|
|
|
<div class="mb-6">
|
|
<h3 style="font-size:1.5rem">Auto Backend Selection</h3>
|
|
<p class="muted">The easiest way to use WAN models is with automatic backend selection:</p>
|
|
<pre><code class="language-python">from pipelines.streamdiffusionv2.vendor.causvid.models.wan.mlx.backend import create_model
|
|
|
|
# Automatically select best available backend
|
|
# Priority: MLX > PyTorch CUDA > PyTorch MPS > CPU
|
|
model = create_model(
|
|
model_type="t2v", # or "i2v" for image-to-video
|
|
dim=2048,
|
|
num_layers=32,
|
|
num_heads=16,
|
|
)
|
|
|
|
# Model is ready for inference!
|
|
</code></pre>
|
|
</div>
|
|
|
|
<div class="mb-6">
|
|
<h3 style="font-size:1.5rem">Complete Model Usage</h3>
|
|
<pre><code class="language-python">import mlx.core as mx
|
|
from pipelines.streamdiffusionv2.vendor.causvid.models.wan.mlx import WanModel
|
|
|
|
# Create model
|
|
model = WanModel(
|
|
model_type="t2v",
|
|
dim=2048,
|
|
num_layers=32,
|
|
num_heads=16,
|
|
patch_size=(2, 2, 2),
|
|
)
|
|
|
|
# Prepare input
|
|
latents = mx.random.normal((2, 16, 128, 128, 4)) # [B, F, H, W, C]
|
|
timesteps = mx.array([500, 750])
|
|
text_embeds = mx.random.normal((2, 77, 768))
|
|
|
|
# Forward pass
|
|
output = model(
|
|
latents=latents,
|
|
timesteps=timesteps,
|
|
text_embeds=text_embeds,
|
|
)
|
|
|
|
print(f"Output shape: {output.shape}") # [2, 16, 128, 128, 4]
|
|
</code></pre>
|
|
</div>
|
|
|
|
<div class="mb-6">
|
|
<h3 style="font-size:1.5rem">Weight Conversion</h3>
|
|
<p class="muted">Convert PyTorch checkpoints to MLX format:</p>
|
|
<pre><code class="language-python">from pipelines.streamdiffusionv2.vendor.causvid.models.wan.mlx.convert_weights import (
|
|
convert_checkpoint_to_mlx
|
|
)
|
|
|
|
# Convert PyTorch checkpoint to MLX
|
|
mlx_weights = convert_checkpoint_to_mlx(
|
|
pytorch_checkpoint_path="wan_t2v.pth",
|
|
output_path="wan_t2v_mlx.safetensors",
|
|
)
|
|
|
|
print("Conversion complete!")
|
|
</code></pre>
|
|
</div>
|
|
|
|
<div>
|
|
<h3 style="font-size:1.5rem">Running Tests</h3>
|
|
<pre><code class="language-bash"># Run all unit tests
|
|
uv run python -m pipelines.streamdiffusionv2.vendor.causvid.models.wan.mlx.test_mlx
|
|
|
|
# Or use Python API
|
|
from pipelines.streamdiffusionv2.vendor.causvid.models.wan.mlx.test_mlx import run_all_tests
|
|
|
|
success = run_all_tests()
|
|
# All 9 tests passing ✅
|
|
</code></pre>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Implementation Status -->
|
|
<section class="mb-12">
|
|
<h2 class="purple">Implementation Status</h2>
|
|
|
|
<div class="callout green mb-6">
|
|
<h3 class="green mb-6" style="font-size:1.5rem">
|
|
<i class="fas fa-check-circle"></i>
|
|
Fully Implemented - Production Ready!
|
|
</h3>
|
|
|
|
<div class="grid cols-2">
|
|
<div>
|
|
<h4 style="font-size:1.125rem">Core Operations</h4>
|
|
<ul class="plain">
|
|
<li><i class="fas fa-check"></i>Sinusoidal position embeddings</li>
|
|
<li><i class="fas fa-check"></i>RoPE (Rotary Position Embeddings)</li>
|
|
<li><i class="fas fa-check"></i>WanRMSNorm</li>
|
|
<li><i class="fas fa-check"></i>WanLayerNorm</li>
|
|
</ul>
|
|
</div>
|
|
|
|
<div>
|
|
<h4 style="font-size:1.125rem">3D Convolution</h4>
|
|
<ul class="plain">
|
|
<li><i class="fas fa-check"></i><strong>Manual Conv3d implementation</strong></li>
|
|
<li class="sub">No longer blocked by MLX limitations!</li>
|
|
<li class="sub">Sliding window + matmul approach</li>
|
|
<li class="sub">PyTorch API compatible</li>
|
|
</ul>
|
|
</div>
|
|
|
|
<div>
|
|
<h4 style="font-size:1.125rem">Neural Network Modules</h4>
|
|
<ul class="plain">
|
|
<li><i class="fas fa-check"></i>WanLinear</li>
|
|
<li><i class="fas fa-check"></i>GELU activation</li>
|
|
<li><i class="fas fa-check"></i>FeedForward network</li>
|
|
<li><i class="fas fa-check"></i>Head & MLPProj</li>
|
|
</ul>
|
|
</div>
|
|
|
|
<div>
|
|
<h4 style="font-size:1.125rem">Attention Mechanisms</h4>
|
|
<ul class="plain">
|
|
<li><i class="fas fa-check"></i>Scaled dot-product attention</li>
|
|
<li><i class="fas fa-check"></i>WanSelfAttention with RoPE</li>
|
|
<li><i class="fas fa-check"></i>T2V Cross-attention</li>
|
|
<li><i class="fas fa-check"></i>I2V Cross-attention</li>
|
|
</ul>
|
|
</div>
|
|
|
|
<div>
|
|
<h4 style="font-size:1.125rem">Transformer & Model</h4>
|
|
<ul class="plain">
|
|
<li><i class="fas fa-check"></i>WanAttentionBlock</li>
|
|
<li><i class="fas fa-check"></i><strong>Complete WanModel</strong></li>
|
|
<li class="sub">Patch embedding</li>
|
|
<li class="sub">T2V and I2V modes</li>
|
|
</ul>
|
|
</div>
|
|
|
|
<div>
|
|
<h4 style="font-size:1.125rem">Utilities</h4>
|
|
<ul class="plain">
|
|
<li><i class="fas fa-check"></i>Weight conversion (PyTorch → MLX)</li>
|
|
<li><i class="fas fa-check"></i>Auto backend selection</li>
|
|
<li><i class="fas fa-check"></i>Benchmarking suite</li>
|
|
<li><i class="fas fa-check"></i>Complete test suite (9/9 ✅)</li>
|
|
</ul>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="callout blue">
|
|
<h3 class="blue mb-6" style="font-size:1.5rem">
|
|
<i class="fas fa-lightbulb"></i>
|
|
What This Means
|
|
</h3>
|
|
<p>The MLX-native WAN implementation is <strong>complete and ready for use</strong>! You can now:</p>
|
|
<ol>
|
|
<li>Run end-to-end inference on Apple Silicon with 2-3x speedup</li>
|
|
<li>Convert existing PyTorch checkpoints to MLX format</li>
|
|
<li>Automatically select the best backend (MLX/PyTorch)</li>
|
|
<li>Benchmark performance to verify speedups</li>
|
|
<li>Run tests to ensure numerical accuracy</li>
|
|
</ol>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Technical Details -->
|
|
<section class="mb-12">
|
|
<h2 class="yellow">Technical Details</h2>
|
|
|
|
<div class="stack-lg">
|
|
<div class="card solid">
|
|
<h3>Float32 Precision</h3>
|
|
<p class="muted">All operations use <code>float32</code> instead of <code>float64</code> for:</p>
|
|
<ul>
|
|
<li>MPS backend doesn't support float64</li>
|
|
<li>Float32 provides sufficient precision for diffusion models</li>
|
|
<li>Faster computation and lower memory usage</li>
|
|
</ul>
|
|
</div>
|
|
|
|
<div class="card solid">
|
|
<h3>Complex Number Handling</h3>
|
|
<p class="muted">
|
|
RoPE uses complex rotations, represented as <code>[real, imaginary]</code> pairs in the last dimension:
|
|
</p>
|
|
<pre class="deep"><code class="language-python"># Shape: [seq_len, num_heads, dim//2, 2]
|
|
# [..., 0] = real part
|
|
# [..., 1] = imaginary part
|
|
</code></pre>
|
|
<p class="muted mt-2 mb-0">
|
|
This avoids the need for native complex number support while maintaining mathematical correctness.
|
|
</p>
|
|
</div>
|
|
|
|
<div class="card solid">
|
|
<h3>Memory Layout</h3>
|
|
<p class="muted">
|
|
MLX uses row-major (C-style) memory layout by default, same as NumPy and PyTorch. Attention tensors use:
|
|
</p>
|
|
<ul>
|
|
<li>Input: <code>[batch, seq_len, num_heads, head_dim]</code></li>
|
|
<li>Computation: <code>[batch, num_heads, seq_len, head_dim]</code></li>
|
|
<li>Output: <code>[batch, seq_len, num_heads, head_dim]</code></li>
|
|
</ul>
|
|
<p class="muted mt-2 mb-0">Automatic transposition handles layout conversions efficiently.</p>
|
|
</div>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Dependencies -->
|
|
<section class="mb-12">
|
|
<h2 class="pink">Dependencies</h2>
|
|
<div class="card solid">
|
|
<pre class="deep"><code class="language-toml">mlx>=0.25.0; sys_platform == 'darwin'
|
|
python >= 3.10
|
|
</code></pre>
|
|
<p class="muted mt-4 mb-0">
|
|
MLX is automatically installed on macOS systems via the main <code>pyproject.toml</code>.
|
|
</p>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Resources -->
|
|
<section class="mb-12">
|
|
<h2 class="green">Resources</h2>
|
|
<div class="grid cols-2" style="gap:1rem">
|
|
<a href="https://ml-explore.github.io/mlx/build/html/index.html" target="_blank" class="card solid green" style="padding:1rem">
|
|
<i class="fas fa-book green"></i>
|
|
<span class="bold">MLX Documentation</span>
|
|
</a>
|
|
<a href="https://github.com/ml-explore/mlx" target="_blank" class="card solid blue" style="padding:1rem">
|
|
<i class="fab fa-github blue"></i>
|
|
<span class="bold">MLX GitHub Repository</span>
|
|
</a>
|
|
<a href="https://arxiv.org/abs/2104.09864" target="_blank" class="card solid purple" style="padding:1rem">
|
|
<i class="fas fa-file-alt purple"></i>
|
|
<span class="bold">RoPE Paper</span>
|
|
</a>
|
|
<a href="https://arxiv.org/abs/2205.14135" target="_blank" class="card solid pink" style="padding:1rem">
|
|
<i class="fas fa-file-alt pink"></i>
|
|
<span class="bold">Flash Attention Paper</span>
|
|
</a>
|
|
</div>
|
|
</section>
|
|
|
|
<!-- Footer -->
|
|
<footer>
|
|
<p>© 2024-2025 Hanzo AI. All rights reserved.</p>
|
|
<div class="links">
|
|
<a href="index.html" class="link"><i class="fas fa-home"></i> Home</a>
|
|
<a href="https://github.com/hanzoai/live" class="link"><i class="fab fa-github"></i> GitHub</a>
|
|
</div>
|
|
</footer>
|
|
</div>
|
|
|
|
<script>
|
|
// Highlight code blocks
|
|
hljs.highlightAll();
|
|
</script>
|
|
</body>
|
|
</html>
|