Skip to content

Commit

Permalink
fix(ai): fix ai-worker client bindings
Browse files Browse the repository at this point in the history 
This commit ensures that the right golang client bindings response and
request types are used. It also cleans up the codebase a bit.
  • Loading branch information
@rickstaa
rickstaa committed Jun 10, 2024
1 parent 1b9dc5a commit 44da6b4
Show file tree
Hide file tree
Showing 10 changed files with 189 additions and 186 deletions.
2 changes: 1 addition & 1 deletion ai/file_worker.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ func (w *FileWorker) ImageToVideo(ctx context.Context, req worker.ImageToVideoMu
return &resp, nil
}

func (w *FileWorker) Upscale(ctx context.Context, req worker.UpscaleImageMultipartRequestBody) (*worker.ImageResponse, error) {
func (w *FileWorker) Upscale(ctx context.Context, req worker.UpscaleMultipartRequestBody) (*worker.ImageResponse, error) {
fname, ok := w.files["upscale"]
if !ok {
return nil, errors.New("upscale response file not found")
Expand Down
24 changes: 12 additions & 12 deletions cmd/livepeer/starter/starter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -589,18 +589,6 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
constraints[core.Capability_ImageToImage].Models[config.ModelID] = modelConstraint

n.SetBasePriceForCap("default", core.Capability_ImageToImage, config.ModelID, big.NewRat(config.PricePerUnit, config.PixelsPerUnit))
case "upscale":
_, ok := constraints[core.Capability_Upscale]
if !ok {
aiCaps = append(aiCaps, core.Capability_Upscale)
constraints[core.Capability_Upscale] = &core.Constraints{
Models: make(map[string]*core.ModelConstraint),
}
}

constraints[core.Capability_Upscale].Models[config.ModelID] = modelConstraint

n.SetBasePriceForCap("default", core.Capability_Upscale, config.ModelID, big.NewRat(config.PricePerUnit, config.PixelsPerUnit))
case "image-to-video":
_, ok := constraints[core.Capability_ImageToVideo]
if !ok {
Expand All @@ -613,6 +601,18 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
constraints[core.Capability_ImageToVideo].Models[config.ModelID] = modelConstraint

n.SetBasePriceForCap("default", core.Capability_ImageToVideo, config.ModelID, big.NewRat(config.PricePerUnit, config.PixelsPerUnit))
case "upscale":
_, ok := constraints[core.Capability_Upscale]
if !ok {
aiCaps = append(aiCaps, core.Capability_Upscale)
constraints[core.Capability_Upscale] = &core.Constraints{
Models: make(map[string]*core.ModelConstraint),
}
}

constraints[core.Capability_Upscale].Models[config.ModelID] = modelConstraint

n.SetBasePriceForCap("default", core.Capability_Upscale, config.ModelID, big.NewRat(config.PricePerUnit, config.PixelsPerUnit))
}

if len(aiCaps) > 0 {
Expand Down
2 changes: 1 addition & 1 deletion core/ai.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@ import (
type AI interface {
TextToImage(context.Context, worker.TextToImageJSONRequestBody) (*worker.ImageResponse, error)
ImageToImage(context.Context, worker.ImageToImageMultipartRequestBody) (*worker.ImageResponse, error)
Upscale(context.Context, worker.UpscaleImageMultipartRequestBody) (*worker.ImageResponse, error)
ImageToVideo(context.Context, worker.ImageToVideoMultipartRequestBody) (*worker.VideoResponse, error)
Upscale(context.Context, worker.UpscaleMultipartRequestBody) (*worker.ImageResponse, error)
Warm(context.Context, string, string, worker.RunnerEndpoint, worker.OptimizationFlags) error
Stop(context.Context) error
HasCapacity(pipeline, modelID string) bool
Expand Down
11 changes: 7 additions & 4 deletions core/orchestrator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,13 +117,15 @@ func (orch *orchestrator) TextToImage(ctx context.Context, req worker.TextToImag
func (orch *orchestrator) ImageToImage(ctx context.Context, req worker.ImageToImageMultipartRequestBody) (*worker.ImageResponse, error) {
return orch.node.imageToImage(ctx, req)
}
func (orch *orchestrator) Upscale(ctx context.Context, req worker.UpscaleImageMultipartRequestBody) (*worker.ImageResponse, error) {
return orch.node.upscale(ctx, req)
}

func (orch *orchestrator) ImageToVideo(ctx context.Context, req worker.ImageToVideoMultipartRequestBody) (*worker.ImageResponse, error) {
return orch.node.imageToVideo(ctx, req)
}

func (orch *orchestrator) Upscale(ctx context.Context, req worker.UpscaleMultipartRequestBody) (*worker.ImageResponse, error) {
return orch.node.upscale(ctx, req)
}

func (orch *orchestrator) ProcessPayment(ctx context.Context, payment net.Payment, manifestID ManifestID) error {
if orch.node == nil || orch.node.Recipient == nil {
return nil
Expand Down Expand Up @@ -938,7 +940,8 @@ func (n *LivepeerNode) textToImage(ctx context.Context, req worker.TextToImageJS
func (n *LivepeerNode) imageToImage(ctx context.Context, req worker.ImageToImageMultipartRequestBody) (*worker.ImageResponse, error) {
return n.AIWorker.ImageToImage(ctx, req)
}
func (n *LivepeerNode) upscale(ctx context.Context, req worker.UpscaleImageMultipartRequestBody) (*worker.ImageResponse, error) {

func (n *LivepeerNode) upscale(ctx context.Context, req worker.UpscaleMultipartRequestBody) (*worker.ImageResponse, error) {
return n.AIWorker.Upscale(ctx, req)
}

Expand Down
34 changes: 17 additions & 17 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,13 @@ require (
contrib.go.opencensus.io/exporter/prometheus v0.4.2
github.com/cenkalti/backoff v2.2.1+incompatible
github.com/ethereum/go-ethereum v1.13.5
github.com/getkin/kin-openapi v0.122.0
github.com/getkin/kin-openapi v0.124.0
github.com/golang/glog v1.1.1
github.com/golang/mock v1.6.0
github.com/golang/protobuf v1.5.3
github.com/jaypipes/ghw v0.10.0
github.com/jaypipes/pcidb v1.0.0
github.com/livepeer/ai-worker v0.0.3
github.com/livepeer/ai-worker v0.0.8
github.com/livepeer/go-tools v0.3.6-0.20240130205227-92479de8531b
github.com/livepeer/livepeer-data v0.7.5-0.20231004073737-06f1f383fb18
github.com/livepeer/lpms v0.0.0-20240120150405-de94555cdc69
Expand All @@ -25,12 +25,12 @@ require (
github.com/peterbourgon/ff/v3 v3.4.0
github.com/pkg/errors v0.9.1
github.com/prometheus/client_golang v1.14.0
github.com/stretchr/testify v1.8.4
github.com/stretchr/testify v1.9.0
github.com/testcontainers/testcontainers-go v0.9.0
github.com/urfave/cli v1.22.12
go.opencensus.io v0.24.0
go.uber.org/goleak v1.3.0
golang.org/x/net v0.19.0
golang.org/x/net v0.25.0
google.golang.org/grpc v1.57.1
google.golang.org/protobuf v1.31.0
pgregory.net/rapid v1.1.0
Expand Down Expand Up @@ -60,7 +60,7 @@ require (
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/deckarep/golang-set/v2 v2.1.0 // indirect
github.com/decred/dcrd/dcrec/secp256k1/v4 v4.1.0 // indirect
github.com/deepmap/oapi-codegen/v2 v2.1.0 // indirect
github.com/deepmap/oapi-codegen/v2 v2.2.0 // indirect
github.com/distribution/reference v0.5.0 // indirect
github.com/dlclark/regexp2 v1.7.0 // indirect
github.com/docker/cli v24.0.5+incompatible // indirect
Expand All @@ -75,26 +75,26 @@ require (
github.com/fsnotify/fsnotify v1.6.0 // indirect
github.com/gballet/go-libpcsclite v0.0.0-20190607065134-2772fd86a8ff // indirect
github.com/ghodss/yaml v1.0.0 // indirect
github.com/go-chi/chi/v5 v5.0.11 // indirect
github.com/go-chi/chi/v5 v5.0.12 // indirect
github.com/go-kit/log v0.2.1 // indirect
github.com/go-logfmt/logfmt v0.5.1 // indirect
github.com/go-logr/logr v1.2.4 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/go-ole/go-ole v1.2.6 // indirect
github.com/go-openapi/jsonpointer v0.19.6 // indirect
github.com/go-openapi/swag v0.22.4 // indirect
github.com/go-openapi/jsonpointer v0.20.2 // indirect
github.com/go-openapi/swag v0.22.8 // indirect
github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
github.com/go-stack/stack v1.8.1 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/golang/snappy v0.0.5-0.20220116011046-fa5810519dcb // indirect
github.com/google/go-cmp v0.5.9 // indirect
github.com/google/go-cmp v0.6.0 // indirect
github.com/google/pprof v0.0.0-20230207041349-798e818bf904 // indirect
github.com/google/s2a-go v0.1.4 // indirect
github.com/google/uuid v1.5.0 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.2.3 // indirect
github.com/googleapis/gax-go/v2 v2.10.0 // indirect
github.com/gorilla/mux v1.8.0 // indirect
github.com/gorilla/mux v1.8.1 // indirect
github.com/gorilla/websocket v1.4.2 // indirect
github.com/hashicorp/golang-lru v0.5.4 // indirect
github.com/holiman/uint256 v1.2.3 // indirect
Expand Down Expand Up @@ -172,7 +172,7 @@ require (
github.com/sirupsen/logrus v1.9.3 // indirect
github.com/spaolacci/murmur3 v1.1.0 // indirect
github.com/status-im/keycard-go v0.2.0 // indirect
github.com/stretchr/objx v0.5.0 // indirect
github.com/stretchr/objx v0.5.2 // indirect
github.com/supranational/blst v0.3.11 // indirect
github.com/tklauser/go-sysconf v0.3.12 // indirect
github.com/tklauser/numcpus v0.6.1 // indirect
Expand All @@ -184,14 +184,14 @@ require (
go.uber.org/atomic v1.11.0 // indirect
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.24.0 // indirect
golang.org/x/crypto v0.17.0 // indirect
golang.org/x/crypto v0.23.0 // indirect
golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
golang.org/x/mod v0.12.0 // indirect
golang.org/x/mod v0.17.0 // indirect
golang.org/x/oauth2 v0.8.0 // indirect
golang.org/x/sync v0.3.0 // indirect
golang.org/x/sys v0.15.0 // indirect
golang.org/x/text v0.14.0 // indirect
golang.org/x/tools v0.13.0 // indirect
golang.org/x/sync v0.7.0 // indirect
golang.org/x/sys v0.20.0 // indirect
golang.org/x/text v0.15.0 // indirect
golang.org/x/tools v0.21.0 // indirect
golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
google.golang.org/api v0.125.0 // indirect
google.golang.org/appengine v1.6.7 // indirect
Expand Down
Loading

0 comments on commit 44da6b4

Please sign in to comment.