opencv

Signed-off-by: Sertac Ozercan <[email protected]>
sozercan · Jan 27, 2025 · dfb49e8 · dfb49e8
1 parent aa5581d
commit dfb49e8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/pkg/aikit2llb/inference/convert.go b/pkg/aikit2llb/inference/convert.go
@@ -50,7 +50,7 @@ func Aikit2LLB(c *config.InferenceConfig, platform *specs.Platform) (llb.State,
 		case utils.BackendExllamaV2:
 			merge = installExllama(state, merge)
 		case utils.BackendStableDiffusion:
-			merge = installOpenCV(state, merge)
+			merge = installOpenCV(state, merge, *platform)
 		case utils.BackendMamba:
 			merge = installMamba(state, merge)
 		case utils.BackendDiffusers:

diff --git a/pkg/aikit2llb/inference/stablediffusion.go b/pkg/aikit2llb/inference/stablediffusion.go
@@ -4,12 +4,23 @@ import (
 	"fmt"
 
 	"github.com/moby/buildkit/client/llb"
+	specs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sozercan/aikit/pkg/utils"
 )
 
-func installOpenCV(s llb.State, merge llb.State) llb.State {
+func installOpenCV(s llb.State, merge llb.State, platform specs.Platform) llb.State {
+	libPaths := map[string]string{
+		utils.PlatformAMD64: "/usr/lib/x86_64-linux-gnu",
+		utils.PlatformARM64: "/usr/lib/aarch64-linux-gnu",
+	}
+	libPath, exists := libPaths[platform.Architecture]
+	if !exists {
+		return s
+	}
+
 	savedState := s
-	s = s.Run(utils.Sh("apt-get update && mkdir -p /tmp/generated/images && apt-get install --no-install-recommends -y libopencv-imgcodecs4.5d && apt-get clean"), llb.IgnoreCache).Root()
+
+	s = s.Run(utils.Shf("apt-get update && mkdir -p /tmp/generated/images && apt-get install --no-install-recommends -y libopencv-imgcodecs4.5d && ln -s %[1]s/libopencv_core.so.4.5d %[1]s/libopencv_core.so.406 && ln -s %[1]s/libopencv_imgcodecs.so.4.5d %[1]s/libopencv_imgcodecs.so.406 && ln -s %[1]s/ libopencv_imgproc.so.4.5d %[1]s/ libopencv_imgproc.so.406 && apt-get clean", libPath), llb.IgnoreCache).Root()
 	diff := llb.Diff(savedState, s)
 	merge = llb.Merge([]llb.State{merge, diff})