Initial commit: img-infinite-canvas AI design workbench MVP

Moteva-style AI design workbench replica. Home prompt creates a project; the project page provides an infinite canvas with node drag, zoom/pan, a design chat panel, history replay, image asset upload, and project save/regenerate. - Backend: go-zero, DDD layering, sqlc/pgx, optional PostgreSQL; memory or Redis cache; asynq job queue; MinIO/S3/R2/OSS object storage; sky-valley/pi agent runtime adapter. - Frontend: Next.js App Router + Vite artifact build, TypeScript, i18n, shadcn/ui components, auth (OTP/Turnstile/Google/WeChat). - Deploy: Docker Compose and k3s manifests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-07 23:15:37 +08:00
commit 44406b72db
349 changed files with 73265 additions and 0 deletions
@@ -0,0 +1,556 @@
+package backgroundremoval
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"image"
+	"image/color"
+	imagedraw "image/draw"
+	"image/png"
+	"math"
+	"os"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"sync"
+	"time"
+
+	"img_infinite_canvas/internal/domain/design"
+
+	"github.com/zeromicro/go-zero/core/logx"
+	"golang.org/x/image/draw"
+
+	ort "github.com/yalue/onnxruntime_go"
+)
+
+const (
+	defaultModelInputSize = 1024
+	minAutoInputSize      = 256
+	maxAutoInputSize      = 2048
+	inputSizeMultiple     = 32
+	defaultInputName      = "input"
+	defaultOutputName     = "output"
+	defaultTimeout        = 5 * time.Minute
+)
+
+var ortEnvironmentMu sync.Mutex
+
+type ONNXRuntimeOptions struct {
+	ModelPath         string
+	SharedLibraryPath string
+	ExecutionProvider string
+	InputSize         int
+	TimeoutSeconds    int
+	InputName         string
+	OutputName        string
+}
+
+type ONNXRuntimeRemover struct {
+	modelPath         string
+	sharedLibraryPath string
+	executionProvider string
+	inputSize         int
+	timeout           time.Duration
+	inputName         string
+	outputName        string
+
+	mu             sync.Mutex
+	session        *ort.DynamicAdvancedSession
+	providerName   string
+	modelInputSize int
+}
+
+func NewONNXRuntimeRemover(opts ONNXRuntimeOptions) *ONNXRuntimeRemover {
+	inputSize := opts.InputSize
+	if inputSize < 0 {
+		inputSize = 0
+	}
+	timeout := time.Duration(opts.TimeoutSeconds) * time.Second
+	if timeout <= 0 {
+		timeout = defaultTimeout
+	}
+	inputName := strings.TrimSpace(opts.InputName)
+	if inputName == "" {
+		inputName = defaultInputName
+	}
+	outputName := strings.TrimSpace(opts.OutputName)
+	if outputName == "" {
+		outputName = defaultOutputName
+	}
+	return &ONNXRuntimeRemover{
+		modelPath:         resolvePath(opts.ModelPath, "model/rmbg1.4.onnx"),
+		sharedLibraryPath: resolveSharedLibraryPath(opts.SharedLibraryPath),
+		executionProvider: normalizeExecutionProvider(opts.ExecutionProvider),
+		inputSize:         inputSize,
+		timeout:           timeout,
+		inputName:         inputName,
+		outputName:        outputName,
+	}
+}
+
+func (r *ONNXRuntimeRemover) RemoveBackground(ctx context.Context, req design.BackgroundRemovalRequest) (design.BackgroundRemoval, error) {
+	if len(req.Image) == 0 {
+		return design.BackgroundRemoval{}, fmt.Errorf("%w: image is required", design.ErrInvalidInput)
+	}
+	if r.timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, r.timeout)
+		defer cancel()
+	}
+	if err := ctx.Err(); err != nil {
+		return design.BackgroundRemoval{}, err
+	}
+
+	source, _, err := image.Decode(bytes.NewReader(req.Image))
+	if err != nil {
+		return design.BackgroundRemoval{}, err
+	}
+	sourceBounds := source.Bounds()
+	r.mu.Lock()
+	if err := r.ensureSession(); err != nil {
+		r.mu.Unlock()
+		return design.BackgroundRemoval{}, err
+	}
+	inputSize := r.effectiveInputSize(req, sourceBounds.Dx(), sourceBounds.Dy())
+	r.mu.Unlock()
+
+	inputData, sourceRGBA := preprocessImage(source, inputSize)
+	inputTensor, err := ort.NewTensor(ort.NewShape(1, 3, int64(inputSize), int64(inputSize)), inputData)
+	if err != nil {
+		return design.BackgroundRemoval{}, err
+	}
+	defer inputTensor.Destroy()
+
+	outputs := []ort.Value{nil}
+	r.mu.Lock()
+	err = r.session.Run([]ort.Value{inputTensor}, outputs)
+	r.mu.Unlock()
+	if err != nil {
+		return design.BackgroundRemoval{}, err
+	}
+	if outputs[0] != nil {
+		defer outputs[0].Destroy()
+	}
+	if err := ctx.Err(); err != nil {
+		return design.BackgroundRemoval{}, err
+	}
+
+	outputTensor, ok := outputs[0].(*ort.Tensor[float32])
+	if !ok {
+		return design.BackgroundRemoval{}, fmt.Errorf("rmbg output tensor has unsupported type %T", outputs[0])
+	}
+	mask := postprocessMask(outputTensor.GetData(), outputTensor.GetShape(), sourceRGBA.Bounds().Dx(), sourceRGBA.Bounds().Dy())
+	result := applyAlphaMask(sourceRGBA, mask)
+
+	var output bytes.Buffer
+	if err := png.Encode(&output, result); err != nil {
+		return design.BackgroundRemoval{}, err
+	}
+	return design.BackgroundRemoval{
+		Image:       output.Bytes(),
+		ContentType: "image/png",
+		Width:       result.Bounds().Dx(),
+		Height:      result.Bounds().Dy(),
+	}, nil
+}
+
+func (r *ONNXRuntimeRemover) ensureSession() error {
+	if r.session != nil {
+		return nil
+	}
+	if err := ensureORTEnvironment(r.sharedLibraryPath); err != nil {
+		return err
+	}
+	if inputSize, err := fixedModelInputSize(r.modelPath, r.inputName); err == nil && inputSize > 0 {
+		r.modelInputSize = inputSize
+	} else if err != nil {
+		logx.Errorf("read background removal model input size failed: %v", err)
+	}
+	session, provider, err := createSessionWithFallback(r.modelPath, r.inputName, r.outputName, r.executionProvider)
+	if err != nil {
+		return err
+	}
+	r.session = session
+	r.providerName = provider
+	logx.Infof("background removal ONNX Runtime provider: %s", provider)
+	return nil
+}
+
+func (r *ONNXRuntimeRemover) effectiveInputSize(req design.BackgroundRemovalRequest, imageWidth int, imageHeight int) int {
+	if r.modelInputSize > 0 {
+		return r.modelInputSize
+	}
+	if req.InputSize > 0 {
+		return normalizeInputSize(req.InputSize)
+	}
+	if r.inputSize > 0 {
+		return normalizeInputSize(r.inputSize)
+	}
+	if req.Width > 0 || req.Height > 0 {
+		return inputSizeFromImageSize(req.Width, req.Height)
+	}
+	return inputSizeFromImageSize(imageWidth, imageHeight)
+}
+
+func fixedModelInputSize(modelPath string, inputName string) (int, error) {
+	inputs, _, err := ort.GetInputOutputInfo(modelPath)
+	if err != nil {
+		return 0, err
+	}
+	if len(inputs) == 0 {
+		return 0, nil
+	}
+	selected := inputs[0]
+	for _, input := range inputs {
+		if input.Name == inputName {
+			selected = input
+			break
+		}
+	}
+	shape := selected.Dimensions
+	if len(shape) < 4 {
+		return 0, nil
+	}
+	height := int(shape[len(shape)-2])
+	width := int(shape[len(shape)-1])
+	if height <= 0 || width <= 0 {
+		return 0, nil
+	}
+	if height > width {
+		return height, nil
+	}
+	return width, nil
+}
+
+func ensureORTEnvironment(sharedLibraryPath string) error {
+	ortEnvironmentMu.Lock()
+	defer ortEnvironmentMu.Unlock()
+	if ort.IsInitialized() {
+		return nil
+	}
+	if sharedLibraryPath != "" {
+		ort.SetSharedLibraryPath(sharedLibraryPath)
+	}
+	if err := ort.InitializeEnvironment(); err != nil {
+		return fmt.Errorf("initialize ONNX Runtime: %w", err)
+	}
+	return nil
+}
+
+func createSessionWithFallback(modelPath string, inputName string, outputName string, provider string) (*ort.DynamicAdvancedSession, string, error) {
+	attempts := executionProviderAttempts(provider)
+	var failures []string
+	for _, attempt := range attempts {
+		session, err := createSession(modelPath, inputName, outputName, attempt)
+		if err == nil {
+			return session, attempt, nil
+		}
+		failures = append(failures, fmt.Sprintf("%s: %v", attempt, err))
+		if provider != "auto" && provider != "gpu" {
+			break
+		}
+	}
+	return nil, "", fmt.Errorf("create ONNX Runtime session failed (%s)", strings.Join(failures, "; "))
+}
+
+func createSession(modelPath string, inputName string, outputName string, provider string) (*ort.DynamicAdvancedSession, error) {
+	options, err := ort.NewSessionOptions()
+	if err != nil {
+		return nil, err
+	}
+	defer options.Destroy()
+	if err := options.SetIntraOpNumThreads(0); err != nil {
+		return nil, err
+	}
+	if err := appendExecutionProvider(options, provider); err != nil {
+		return nil, err
+	}
+	return ort.NewDynamicAdvancedSession(modelPath, []string{inputName}, []string{outputName}, options)
+}
+
+func appendExecutionProvider(options *ort.SessionOptions, provider string) error {
+	switch provider {
+	case "cpu":
+		return nil
+	case "cuda":
+		cudaOptions, err := ort.NewCUDAProviderOptions()
+		if err != nil {
+			return err
+		}
+		defer cudaOptions.Destroy()
+		if err := cudaOptions.Update(map[string]string{"device_id": "0"}); err != nil {
+			return err
+		}
+		return options.AppendExecutionProviderCUDA(cudaOptions)
+	case "coreml":
+		return options.AppendExecutionProviderCoreMLV2(map[string]string{
+			"MLComputeUnits": "ALL",
+		})
+	case "directml":
+		return options.AppendExecutionProviderDirectML(0)
+	default:
+		return fmt.Errorf("%w: unsupported ONNX execution provider %q", design.ErrInvalidInput, provider)
+	}
+}
+
+func executionProviderAttempts(provider string) []string {
+	switch provider {
+	case "cpu":
+		return []string{"cpu"}
+	case "cuda":
+		return []string{"cuda", "cpu"}
+	case "coreml":
+		return []string{"coreml", "cpu"}
+	case "directml":
+		return []string{"directml", "cpu"}
+	case "gpu", "auto":
+		switch runtime.GOOS {
+		case "darwin", "ios":
+			return []string{"coreml", "cpu"}
+		case "windows":
+			return []string{"cuda", "directml", "cpu"}
+		default:
+			return []string{"cuda", "cpu"}
+		}
+	default:
+		return []string{"cpu"}
+	}
+}
+
+func normalizeExecutionProvider(provider string) string {
+	provider = strings.ToLower(strings.TrimSpace(provider))
+	provider = strings.ReplaceAll(provider, "_", "")
+	provider = strings.ReplaceAll(provider, "-", "")
+	switch provider {
+	case "", "auto":
+		return "auto"
+	case "gpu":
+		return "gpu"
+	case "cpu", "cpuexecutionprovider":
+		return "cpu"
+	case "cuda", "cudaexecutionprovider":
+		return "cuda"
+	case "coreml", "coremlexecutionprovider":
+		return "coreml"
+	case "directml", "dml", "directmlexecutionprovider":
+		return "directml"
+	default:
+		return provider
+	}
+}
+
+func preprocessImage(source image.Image, inputSize int) ([]float32, *image.NRGBA) {
+	sourceBounds := source.Bounds()
+	sourceRGBA := image.NewNRGBA(image.Rect(0, 0, sourceBounds.Dx(), sourceBounds.Dy()))
+	imagedraw.Draw(sourceRGBA, sourceRGBA.Bounds(), source, sourceBounds.Min, imagedraw.Src)
+
+	resized := image.NewNRGBA(image.Rect(0, 0, inputSize, inputSize))
+	draw.CatmullRom.Scale(resized, resized.Bounds(), sourceRGBA, sourceRGBA.Bounds(), draw.Over, nil)
+
+	planeSize := inputSize * inputSize
+	data := make([]float32, 3*planeSize)
+	for y := 0; y < inputSize; y++ {
+		for x := 0; x < inputSize; x++ {
+			offset := resized.PixOffset(x, y)
+			index := y*inputSize + x
+			data[index] = float32(resized.Pix[offset])/255 - 0.5
+			data[planeSize+index] = float32(resized.Pix[offset+1])/255 - 0.5
+			data[2*planeSize+index] = float32(resized.Pix[offset+2])/255 - 0.5
+		}
+	}
+	return data, sourceRGBA
+}
+
+func postprocessMask(data []float32, shape ort.Shape, width int, height int) *image.Gray {
+	if len(data) == 0 || width <= 0 || height <= 0 {
+		return image.NewGray(image.Rect(0, 0, maxInt(width, 1), maxInt(height, 1)))
+	}
+	maskWidth, maskHeight := maskDimensions(shape, len(data))
+	mask := image.NewGray(image.Rect(0, 0, maskWidth, maskHeight))
+	minValue, maxValue := data[0], data[0]
+	for _, value := range data {
+		if value < minValue {
+			minValue = value
+		}
+		if value > maxValue {
+			maxValue = value
+		}
+	}
+	denominator := maxValue - minValue
+	for y := 0; y < maskHeight; y++ {
+		for x := 0; x < maskWidth; x++ {
+			index := y*maskWidth + x
+			value := float32(0)
+			if index < len(data) && denominator > 1e-6 {
+				value = (data[index] - minValue) / denominator
+			}
+			value = float32(math.Max(0, math.Min(1, float64(value))))
+			mask.SetGray(x, y, color.Gray{Y: uint8(value*255 + 0.5)})
+		}
+	}
+	if maskWidth == width && maskHeight == height {
+		return mask
+	}
+	resized := image.NewGray(image.Rect(0, 0, width, height))
+	draw.CatmullRom.Scale(resized, resized.Bounds(), mask, mask.Bounds(), draw.Over, nil)
+	return resized
+}
+
+func maskDimensions(shape ort.Shape, dataLength int) (int, int) {
+	if len(shape) >= 2 {
+		width := int(shape[len(shape)-1])
+		height := int(shape[len(shape)-2])
+		if width > 0 && height > 0 && width*height <= dataLength {
+			return width, height
+		}
+	}
+	side := int(math.Sqrt(float64(dataLength)))
+	if side > 0 && side*side == dataLength {
+		return side, side
+	}
+	return dataLength, 1
+}
+
+func inputSizeFromImageSize(width int, height int) int {
+	maxSide := width
+	if height > maxSide {
+		maxSide = height
+	}
+	if maxSide <= 0 {
+		return defaultModelInputSize
+	}
+	return normalizeInputSize(maxSide)
+}
+
+func normalizeInputSize(size int) int {
+	if size <= 0 {
+		return defaultModelInputSize
+	}
+	if size < minAutoInputSize {
+		size = minAutoInputSize
+	}
+	if size > maxAutoInputSize {
+		size = maxAutoInputSize
+	}
+	if remainder := size % inputSizeMultiple; remainder != 0 {
+		size += inputSizeMultiple - remainder
+	}
+	if size > maxAutoInputSize {
+		size = maxAutoInputSize
+	}
+	return size
+}
+
+func applyAlphaMask(source *image.NRGBA, mask *image.Gray) *image.NRGBA {
+	bounds := source.Bounds()
+	result := image.NewNRGBA(bounds)
+	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
+		for x := bounds.Min.X; x < bounds.Max.X; x++ {
+			sourceOffset := source.PixOffset(x, y)
+			maskAlpha := mask.GrayAt(x-bounds.Min.X, y-bounds.Min.Y).Y
+			sourceAlpha := source.Pix[sourceOffset+3]
+			alpha := uint8((uint16(maskAlpha)*uint16(sourceAlpha) + 127) / 255)
+			resultOffset := result.PixOffset(x, y)
+			result.Pix[resultOffset] = source.Pix[sourceOffset]
+			result.Pix[resultOffset+1] = source.Pix[sourceOffset+1]
+			result.Pix[resultOffset+2] = source.Pix[sourceOffset+2]
+			result.Pix[resultOffset+3] = alpha
+		}
+	}
+	return result
+}
+
+func resolveSharedLibraryPath(configured string) string {
+	for _, candidate := range []string{
+		configured,
+		os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH"),
+		os.Getenv("ORT_SHARED_LIBRARY_PATH"),
+		os.Getenv("ORT_DYLIB_PATH"),
+		"model/libonnxruntime.dylib",
+		"model/libonnxruntime.so",
+		"model/onnxruntime.dll",
+		"server/model/libonnxruntime.dylib",
+		"server/model/libonnxruntime.so",
+		"server/model/onnxruntime.dll",
+		"/opt/homebrew/lib/libonnxruntime.dylib",
+		"/usr/local/lib/libonnxruntime.dylib",
+		"/usr/local/lib/libonnxruntime.so",
+	} {
+		path := resolveExistingPath(candidate)
+		if path != "" {
+			return path
+		}
+	}
+	return strings.TrimSpace(configured)
+}
+
+func resolvePath(configured string, fallback string) string {
+	path := resolveExistingPath(configured)
+	if path != "" {
+		return path
+	}
+	path = resolveExistingPath(fallback)
+	if path != "" {
+		return path
+	}
+	return strings.TrimSpace(firstNonEmpty(configured, fallback))
+}
+
+func resolveExistingPath(value string) string {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return ""
+	}
+	candidates := []string{value}
+	if !filepath.IsAbs(value) {
+		for _, base := range ancestorDirs() {
+			candidates = append(candidates, filepath.Join(base, value), filepath.Join(base, "server", value))
+		}
+	}
+	for _, candidate := range candidates {
+		if _, err := os.Stat(candidate); err == nil {
+			if absolute, err := filepath.Abs(candidate); err == nil {
+				return absolute
+			}
+			return candidate
+		}
+	}
+	return ""
+}
+
+func ancestorDirs() []string {
+	cwd, err := os.Getwd()
+	if err != nil {
+		return []string{"."}
+	}
+	var dirs []string
+	for {
+		dirs = append(dirs, cwd)
+		parent := filepath.Dir(cwd)
+		if parent == cwd {
+			break
+		}
+		cwd = parent
+		if len(dirs) >= 8 {
+			break
+		}
+	}
+	return dirs
+}
+
+func firstNonEmpty(values ...string) string {
+	for _, value := range values {
+		if strings.TrimSpace(value) != "" {
+			return value
+		}
+	}
+	return ""
+}
+
+func maxInt(a int, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}
@@ -0,0 +1,58 @@
+package backgroundremoval
+
+import (
+	"bytes"
+	"context"
+	"image"
+	"image/color"
+	"image/png"
+	"os"
+	"testing"
+
+	"img_infinite_canvas/internal/domain/design"
+)
+
+func TestONNXRuntimeRemoverIntegration(t *testing.T) {
+	if os.Getenv("RUN_RMBG_ONNX_INTEGRATION") != "1" {
+		t.Skip("set RUN_RMBG_ONNX_INTEGRATION=1 to run the local rmbg model")
+	}
+	var input bytes.Buffer
+	img := image.NewNRGBA(image.Rect(0, 0, 32, 32))
+	for y := 0; y < 32; y++ {
+		for x := 0; x < 32; x++ {
+			c := color.NRGBA{R: 240, G: 240, B: 240, A: 255}
+			if x >= 8 && x < 24 && y >= 8 && y < 24 {
+				c = color.NRGBA{R: 220, G: 40, B: 40, A: 255}
+			}
+			img.SetNRGBA(x, y, c)
+		}
+	}
+	if err := png.Encode(&input, img); err != nil {
+		t.Fatal(err)
+	}
+
+	remover := NewONNXRuntimeRemover(ONNXRuntimeOptions{
+		ModelPath:         "model/rmbg1.4.onnx",
+		SharedLibraryPath: os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH"),
+		ExecutionProvider: "auto",
+		InputSize:         1024,
+		TimeoutSeconds:    300,
+	})
+	result, err := remover.RemoveBackground(context.Background(), design.BackgroundRemovalRequest{
+		Image:       input.Bytes(),
+		ContentType: "image/png",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if result.ContentType != "image/png" || len(result.Image) == 0 {
+		t.Fatalf("unexpected result metadata: %#v", result)
+	}
+	decoded, _, err := image.Decode(bytes.NewReader(result.Image))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if decoded.Bounds().Dx() != 32 || decoded.Bounds().Dy() != 32 {
+		t.Fatalf("unexpected output size %s", decoded.Bounds())
+	}
+}