[wip] prototype python model builder using buildkit llb

examples/hello-context/.dockerignore

@@ -0,0 +1,20 @@
+# The .dockerignore file excludes files from the container build process.
+#
+# https://docs.docker.com/engine/reference/builder/#dockerignore-file
+
+# Exclude Git files
+**/.git
+**/.github
+**/.gitignore
+
+# Exclude Python tooling
+.python-version
+
+# Exclude Python cache files
+__pycache__
+.mypy_cache
+.pytest_cache
+.ruff_cache
+
+# Exclude Python virtual environment
+/venv

examples/hello-context/README.md

@@ -0,0 +1,4 @@
+hello-context
+-------------
+
+A simple model that takes no inputs but will echo back any context provided with the prediction as the output.

examples/hello-context/cog.yaml

@@ -0,0 +1,28 @@
+# Configuration for Cog ⚙️
+# Reference: https://cog.run/yaml
+
+build:
+  # set to true if your model requires a GPU
+  gpu: false
+
+  # a list of ubuntu apt packages to install
+  # system_packages:
+  #   - "libgl1-mesa-glx"
+  #   - "libglib2.0-0"
+
+  # python version in the form '3.11' or '3.11.4'
+  python_version: "3.11"
+
+  # path to a Python requirements.txt file
+  python_requirements: requirements.txt
+
+  # enable fast boots
+  fast: true
+
+  # commands run after the environment is setup
+  # run:
+  #   - "echo env is ready!"
+  #   - "echo another command if needed"
+
+# predict.py defines how predictions are run on your model
+predict: "predict.py:run"

examples/hello-context/predict.py

@@ -0,0 +1,5 @@
+from cog import current_scope
+
+
+def run() -> dict[str, str]:
+    return current_scope().context

examples/hello-context/requirements.txt

@@ -0,0 +1,23 @@
+# This is a normal Python requirements.txt file.
+
+# You can add dependencies directly from PyPI:
+# 
+# numpy==1.26.4
+# torch==2.2.1
+# torchvision==0.17.1
+
+
+# You can also add Git repos as dependencies, but you'll need to add git to the system_packages list in cog.yaml:
+# 
+# build:
+#   system_packages:
+#     - "git"
+# 
+# Then you can use a URL like this:
+# 
+# git+https://github.com/huggingface/transformers
+
+
+# You can also pin Git repos to a specific commit:
+# 
+# git+https://github.com/huggingface/transformers@2d1602a

examples/hello-procedure/README.md

@@ -0,0 +1,22 @@
+# hello
+
+A simple pipeline that transforms your text input by converting it to uppercase and prefixing it with "HELLO".
+
+https://replicate.com/pipelines-beta/hello
+
+## Features
+
+- Converts any text input to uppercase
+- Adds a friendly "HELLO" prefix to your text
+- Simple, single-input interface
+
+## Models
+
+Under the hood it uses these models:
+
+- [pipelines-beta/upcase](https://replicate.com/pipelines-beta/upcase): A utility model that converts text to uppercase
+
+## How it works
+
+The pipeline takes a text prompt as input, passes it to the `upcase` model to convert the text to uppercase, and then adds "HELLO" as a prefix to the transformed text. This creates a greeting-style output from any input text.
+Edit model

examples/hello-procedure/cog.yaml

@@ -0,0 +1,4 @@
+predict: "function.py:run"
+build:
+  system_packages:
+    - tini

examples/hello-procedure/function.py

@@ -0,0 +1,12 @@
+from cog import Input
+from cog.ext.pipelines import include
+
+# with run_state("load"):
+upcase = include("pipelines-beta/upcase")
+
+
+def run(
+    prompt: str = Input(),
+) -> str:
+    upcased_prompt = upcase(prompt=prompt)
+    return f"HELLO {upcased_prompt}"

go.mod

@@ -7,18 +7,22 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.36.3
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.67
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.79.3
+	github.com/containerd/containerd/api v1.8.0
 	github.com/creack/pty v1.1.24
+	github.com/distribution/reference v0.6.0
 	github.com/docker/cli v28.1.1+incompatible
 	github.com/docker/docker v28.1.1+incompatible
 	github.com/docker/go-connections v0.5.0
 	github.com/getkin/kin-openapi v0.128.0
 	github.com/google/go-containerregistry v0.20.5
+	github.com/google/uuid v1.6.0
 	github.com/hashicorp/go-version v1.7.0
 	github.com/logrusorgru/aurora v2.0.3+incompatible
 	github.com/mattn/go-isatty v0.0.20
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/moby/buildkit v0.22.0
 	github.com/moby/term v0.5.2
+	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.1
 	github.com/pkg/errors v0.9.1
 	github.com/replicate/go v0.0.0-20250205165008-b772d7cd506b
@@ -28,6 +32,7 @@ require (
 	github.com/stretchr/testify v1.10.0
 	github.com/testcontainers/testcontainers-go v0.37.0
 	github.com/testcontainers/testcontainers-go/modules/registry v0.37.0
+	github.com/tonistiigi/fsutil v0.0.0-20250417144416-3f76f8130144
 	github.com/tonistiigi/go-csvvalue v0.0.0-20240710180619-ddb21b71c0b4
 	github.com/vbauerster/mpb/v8 v8.10.1
 	github.com/vincent-petithory/dataurl v1.0.0
@@ -97,7 +102,6 @@ require (
 	github.com/chavacava/garif v0.1.0 // indirect
 	github.com/ckaznocha/intrange v0.3.0 // indirect
 	github.com/containerd/console v1.0.4 // indirect
-	github.com/containerd/containerd/api v1.8.0 // indirect
 	github.com/containerd/containerd/v2 v2.0.5 // indirect
 	github.com/containerd/continuity v0.4.5 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
@@ -112,7 +116,6 @@ require (
 	github.com/daixiang0/gci v0.13.5 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/denis-tingaikin/go-header v0.5.0 // indirect
-	github.com/distribution/reference v0.6.0 // indirect
 	github.com/dnephin/pflag v1.0.7 // indirect
 	github.com/docker/distribution v2.8.3+incompatible // indirect
 	github.com/docker/docker-credential-helpers v0.9.3 // indirect
@@ -156,7 +159,6 @@ require (
 	github.com/golangci/unconvert v0.0.0-20240309020433-c5143eacb3ed // indirect
 	github.com/google/go-cmp v0.7.0 // indirect
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/gordonklaus/ineffassign v0.1.0 // indirect
 	github.com/gostaticanalysis/analysisutil v0.7.1 // indirect
 	github.com/gostaticanalysis/comment v1.5.0 // indirect
@@ -216,7 +218,6 @@ require (
 	github.com/nishanths/predeclared v0.2.2 // indirect
 	github.com/nunnatsa/ginkgolinter v0.19.1 // indirect
 	github.com/olekukonko/tablewriter v0.0.5 // indirect
-	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.4 // indirect
 	github.com/perimeterx/marshmallow v1.1.5 // indirect
 	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
@@ -267,7 +268,6 @@ require (
 	github.com/tklauser/numcpus v0.6.1 // indirect
 	github.com/tomarrell/wrapcheck/v2 v2.10.0 // indirect
 	github.com/tommy-muehle/go-mnd/v2 v2.5.1 // indirect
-	github.com/tonistiigi/fsutil v0.0.0-20250417144416-3f76f8130144 // indirect
 	github.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea // indirect
 	github.com/tonistiigi/vt100 v0.0.0-20240514184818-90bafcd6abab // indirect
 	github.com/ultraware/funlen v0.2.0 // indirect

pkg/cli/build.go

@@ -12,8 +12,7 @@ import (
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/http"
-	"github.com/replicate/cog/pkg/image"
-	"github.com/replicate/cog/pkg/registry"
+	"github.com/replicate/cog/pkg/model/factory"
 	"github.com/replicate/cog/pkg/util/console"
 )
 
@@ -97,32 +96,22 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 		logClient.EndBuild(ctx, err, logCtx)
 		return err
 	}
-	registryClient := registry.NewRegistryClient()
-	if err := image.Build(
-		ctx,
-		cfg,
-		projectDir,
-		imageName,
-		buildSecrets,
-		buildNoCache,
-		buildSeparateWeights,
-		buildUseCudaBaseImage,
-		buildProgressOutput,
-		buildSchemaFile,
-		buildDockerfileFile,
-		DetermineUseCogBaseImage(cmd),
-		buildStrip,
-		buildPrecompile,
-		buildFast,
-		nil,
-		buildLocalImage,
-		dockerClient,
-		registryClient); err != nil {
+
+	modelFactory, err := factory.New(dockerClient)
+	if err != nil {
+		logClient.EndBuild(ctx, err, logCtx)
+		return err
+	}
+
+	settings := buildSettings(cmd, cfg, false, projectDir)
+
+	model, _, err := modelFactory.Build(ctx, settings)
+	if err != nil {
 		logClient.EndBuild(ctx, err, logCtx)
 		return err
 	}
 
-	console.Infof("\nImage built as %s", imageName)
+	console.Infof("\nImage built as %s", model.ImageRef())
 	logClient.EndBuild(ctx, nil, logCtx)
 
 	return nil

pkg/cli/image_helpers.go

@@ -0,0 +1,35 @@
+package cli
+
+import (
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/spf13/cobra"
+
+	"github.com/replicate/cog/pkg/config"
+	"github.com/replicate/cog/pkg/model/factory"
+)
+
+func buildSettings(cmd *cobra.Command, cfg *config.Config, isPredict bool, projectDir string) factory.BuildSettings {
+	return factory.BuildSettings{
+		Tag:        config.DockerImageName(projectDir),
+		WorkingDir: projectDir,
+		Config:     cfg,
+		Platform: ocispec.Platform{
+			Architecture: "amd64",
+			OS:           "linux",
+		},
+		Monobase:         buildFast || cfg.Build.Fast,
+		NoCache:          buildNoCache,
+		BuildSecrets:     buildSecrets,
+		SeparateWeights:  buildSeparateWeights,
+		UseCudaBaseImage: buildUseCudaBaseImage,
+		SchemaFile:       buildSchemaFile,
+		DockerfileFile:   buildDockerfileFile,
+		Precompile:       buildPrecompile,
+		ProgressOutput:   buildProgressOutput,
+		Strip:            buildStrip,
+		UseCogBaseImage:  DetermineUseCogBaseImage(cmd),
+		LocalImage:       buildLocalImage,
+		PredictBuild:     isPredict,
+		Annotations:      map[string]string{},
+	}
+}