mirror of
https://github.com/ollama/ollama.git
synced 2026-04-17 19:54:03 +02:00
30fdd229a4
* create: Clean up experimental paths
This cleans up the experimental features, and adds both unit and integration test coverage to verify no regressions.
* create: preserve config and layer names when creating from safetensors models
When creating a model FROM an existing safetensors model, ModelFormat,
Capabilities, and layer Name fields were lost. ModelFormat stayed empty
because it's only set from GGML layers (which safetensors models lack),
and layer names weren't copied in parseFromModel. This caused derived
models to fail loading ("config.json not found in manifest").
* review comments
108 lines
3.5 KiB
Go
108 lines
3.5 KiB
Go
//go:build integration && imagegen
|
|
|
|
package integration
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
)
|
|
|
|
func TestCreateImageGen(t *testing.T) {
|
|
skipIfRemote(t)
|
|
skipUnderMinVRAM(t, 13)
|
|
|
|
// Allow overriding the model directory via env var for local testing,
|
|
// since the model is ~33GB and may already be downloaded elsewhere.
|
|
modelDir := os.Getenv("OLLAMA_TEST_IMAGEGEN_MODEL_DIR")
|
|
if modelDir == "" {
|
|
modelDir = filepath.Join(testdataModelsDir, "Z-Image-Turbo")
|
|
downloadHFModel(t, "Tongyi-MAI/Z-Image-Turbo", modelDir)
|
|
} else {
|
|
t.Logf("Using existing imagegen model at %s", modelDir)
|
|
}
|
|
|
|
// Verify it looks like a valid imagegen model directory
|
|
if _, err := os.Stat(filepath.Join(modelDir, "model_index.json")); err != nil {
|
|
t.Fatalf("model_index.json not found in %s — not a valid imagegen model directory", modelDir)
|
|
}
|
|
|
|
ensureMLXLibraryPath(t)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
|
|
defer cancel()
|
|
|
|
client, _, cleanup := InitServerConnection(ctx, t)
|
|
defer cleanup()
|
|
|
|
modelName := "test-z-image-turbo-create"
|
|
|
|
absModelDir, err := filepath.Abs(modelDir)
|
|
if err != nil {
|
|
t.Fatalf("Failed to get absolute path: %v", err)
|
|
}
|
|
|
|
// Create a Modelfile pointing to the diffusers model directory
|
|
tmpModelfile := filepath.Join(t.TempDir(), "Modelfile")
|
|
if err := os.WriteFile(tmpModelfile, []byte("FROM "+absModelDir+"\n"), 0o644); err != nil {
|
|
t.Fatalf("Failed to write Modelfile: %v", err)
|
|
}
|
|
|
|
t.Logf("Creating imagegen model from %s", absModelDir)
|
|
runOllamaCreate(ctx, t, modelName, "--experimental", "-f", tmpModelfile)
|
|
|
|
// Verify model exists via show
|
|
showReq := &api.ShowRequest{Name: modelName}
|
|
showResp, err := client.Show(ctx, showReq)
|
|
if err != nil {
|
|
t.Fatalf("Model show failed after create: %v", err)
|
|
}
|
|
t.Logf("Created model details: %+v", showResp.Details)
|
|
|
|
// Generate an image to verify the model isn't corrupted
|
|
t.Log("Generating test image...")
|
|
imageBase64, err := generateImage(ctx, client, modelName, "A red circle on a white background")
|
|
if err != nil {
|
|
if strings.Contains(err.Error(), "image generation not available") {
|
|
t.Skip("Target system does not support image generation")
|
|
} else if strings.Contains(err.Error(), "insufficient memory for image generation") {
|
|
t.Skip("insufficient memory for image generation")
|
|
} else if strings.Contains(err.Error(), "ollama-mlx: no such file or directory") {
|
|
t.Skip("unsupported architecture")
|
|
}
|
|
t.Fatalf("Image generation failed: %v", err)
|
|
}
|
|
|
|
// Verify we got valid image data
|
|
imageData, err := base64.StdEncoding.DecodeString(imageBase64)
|
|
if err != nil {
|
|
t.Fatalf("Failed to decode base64 image: %v", err)
|
|
}
|
|
|
|
t.Logf("Generated image: %d bytes", len(imageData))
|
|
|
|
if len(imageData) < 1000 {
|
|
t.Fatalf("Generated image suspiciously small (%d bytes), likely corrupted", len(imageData))
|
|
}
|
|
|
|
// Check for PNG or JPEG magic bytes
|
|
isPNG := len(imageData) >= 4 && imageData[0] == 0x89 && imageData[1] == 'P' && imageData[2] == 'N' && imageData[3] == 'G'
|
|
isJPEG := len(imageData) >= 2 && imageData[0] == 0xFF && imageData[1] == 0xD8
|
|
if !isPNG && !isJPEG {
|
|
t.Fatalf("Generated image is neither PNG nor JPEG (first bytes: %x)", imageData[:min(8, len(imageData))])
|
|
}
|
|
t.Logf("Image format validated (PNG=%v, JPEG=%v)", isPNG, isJPEG)
|
|
|
|
// Cleanup: delete the model
|
|
deleteReq := &api.DeleteRequest{Model: modelName}
|
|
if err := client.Delete(ctx, deleteReq); err != nil {
|
|
t.Logf("Warning: failed to delete test model: %v", err)
|
|
}
|
|
}
|