mirror of
https://github.com/ollama/ollama.git
synced 2026-01-12 00:06:57 +08:00
33ee7168ba
* WIP - MLX backend with gemma3 * MLX: add cmake and go tag build toggles To build the new MLX backend code: cmake --preset MLX cmake --build --preset MLX --parallel cmake --install build --component MLX go build -tags mlx . Note: the main.go entrypoint for the MLX engine will change in a follow up commit. * add experimental image generation runtime * add experimental image generation runtime * MLX: wire up cuda build for linux * MLX: get dependencies correct and dedup This is still too large for a unified github artifact, but is now "correct" for the mlx_cuda_v13 directory. * fix relative link bug in dedup * Add darwin build and readme * add go build tag for mlx dependent code and wire up build_darwin.sh * lint cleanup * macos: build mlx for x86 This will be CPU only. * cuda build instructions and fix drift from mlx bump * stale comment * Delete agent helper doc * Clean up readme.md * Revise README for tokenizer clarity and details Updated README to clarify tokenizer functionality and removed correctness section. --------- Co-authored-by: jmorganca <jmorganca@gmail.com>
90 lines
1.8 KiB
Go
90 lines
1.8 KiB
Go
//go:build mlx
|
|
|
|
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"image"
|
|
"image/png"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
"github.com/ollama/ollama/x/imagegen/mlx"
|
|
)
|
|
|
|
// saveImageArray saves an MLX array as a PNG image.
|
|
// Expected format: [B, C, H, W] with values in [0, 1] range and C=3 (RGB).
|
|
func saveImageArray(arr *mlx.Array, path string) error {
|
|
img, err := arrayToImage(arr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return savePNG(img, path)
|
|
}
|
|
|
|
func savePNG(img *image.RGBA, path string) error {
|
|
if filepath.Ext(path) != ".png" {
|
|
path = path + ".png"
|
|
}
|
|
f, err := os.Create(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer f.Close()
|
|
return png.Encode(f, img)
|
|
}
|
|
|
|
func arrayToImage(arr *mlx.Array) (*image.RGBA, error) {
|
|
shape := arr.Shape()
|
|
if len(shape) != 4 {
|
|
return nil, fmt.Errorf("expected 4D array [B, C, H, W], got %v", shape)
|
|
}
|
|
|
|
// Transform to [H, W, C] for image conversion
|
|
img := mlx.Squeeze(arr, 0)
|
|
arr.Free()
|
|
img = mlx.Transpose(img, 1, 2, 0)
|
|
img = mlx.Contiguous(img)
|
|
mlx.Eval(img)
|
|
|
|
imgShape := img.Shape()
|
|
H := int(imgShape[0])
|
|
W := int(imgShape[1])
|
|
C := int(imgShape[2])
|
|
|
|
if C != 3 {
|
|
img.Free()
|
|
return nil, fmt.Errorf("expected 3 channels (RGB), got %d", C)
|
|
}
|
|
|
|
// Copy to CPU and free GPU memory
|
|
data := img.Data()
|
|
img.Free()
|
|
|
|
// Write directly to Pix slice (faster than SetRGBA)
|
|
goImg := image.NewRGBA(image.Rect(0, 0, W, H))
|
|
pix := goImg.Pix
|
|
for y := 0; y < H; y++ {
|
|
for x := 0; x < W; x++ {
|
|
srcIdx := (y*W + x) * C
|
|
dstIdx := (y*W + x) * 4
|
|
pix[dstIdx+0] = uint8(clampF(data[srcIdx+0]*255+0.5, 0, 255))
|
|
pix[dstIdx+1] = uint8(clampF(data[srcIdx+1]*255+0.5, 0, 255))
|
|
pix[dstIdx+2] = uint8(clampF(data[srcIdx+2]*255+0.5, 0, 255))
|
|
pix[dstIdx+3] = 255
|
|
}
|
|
}
|
|
|
|
return goImg, nil
|
|
}
|
|
|
|
func clampF(v, min, max float32) float32 {
|
|
if v < min {
|
|
return min
|
|
}
|
|
if v > max {
|
|
return max
|
|
}
|
|
return v
|
|
}
|