add pricing and latency calculations

livepeer · Aug 28, 2024 · b9c7771 · b9c7771
1 parent c2b046b
commit b9c7771
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/go.mod b/go.mod
@@ -13,7 +13,7 @@ require (
 	github.com/golang/protobuf v1.5.4
 	github.com/jaypipes/ghw v0.10.0
 	github.com/jaypipes/pcidb v1.0.0
-	github.com/livepeer/ai-worker v0.1.2
+	github.com/livepeer/ai-worker v0.1.3-0.20240828042752-36ebcce29442
 	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-20240819180416-f87352959b85

diff --git a/go.sum b/go.sum
@@ -623,8 +623,8 @@ github.com/libp2p/go-netroute v0.2.0 h1:0FpsbsvuSnAhXFnCY0VLFbJOzaK0VnP0r1QT/o4n
 github.com/libp2p/go-netroute v0.2.0/go.mod h1:Vio7LTzZ+6hoT4CMZi5/6CpY3Snzh2vgZhWgxMNwlQI=
 github.com/libp2p/go-openssl v0.1.0 h1:LBkKEcUv6vtZIQLVTegAil8jbNpJErQ9AnT+bWV+Ooo=
 github.com/libp2p/go-openssl v0.1.0/go.mod h1:OiOxwPpL3n4xlenjx2h7AwSGaFSC/KZvf6gNdOBQMtc=
-github.com/livepeer/ai-worker v0.1.2 h1:I73J4zJYad95QE1JFSrqrjKKCTqLHypDcoPq/zZM5aw=
-github.com/livepeer/ai-worker v0.1.2/go.mod h1:Xlnb0nFG2VsGeMG9hZmReVQXeFt0Dv28ODiUT2ooyLE=
+github.com/livepeer/ai-worker v0.1.3-0.20240828042752-36ebcce29442 h1:2JkKVkxlurXI6r/HtTQfC7asTRmFZdwF/Q1Zm0EtTFI=
+github.com/livepeer/ai-worker v0.1.3-0.20240828042752-36ebcce29442/go.mod h1:Xlnb0nFG2VsGeMG9hZmReVQXeFt0Dv28ODiUT2ooyLE=
 github.com/livepeer/go-tools v0.3.6-0.20240130205227-92479de8531b h1:VQcnrqtCA2UROp7q8ljkh2XA/u0KRgVv0S1xoUvOweE=
 github.com/livepeer/go-tools v0.3.6-0.20240130205227-92479de8531b/go.mod h1:hwJ5DKhl+pTanFWl+EUpw1H7ukPO/H+MFpgA7jjshzw=
 github.com/livepeer/joy4 v0.1.2-0.20191121080656-b2fea45cbded h1:ZQlvR5RB4nfT+cOQee+WqmaDOgGtP2oDMhcVvR4L0yA=

diff --git a/server/ai_http.go b/server/ai_http.go
@@ -313,7 +313,17 @@ func handleAIRequest(ctx context.Context, w http.ResponseWriter, r *http.Request
 			return orch.SegmentAnything2(ctx, v)
 		}
 
-		outPixels = 1000000 // Temp hardcoded 1MM - Convert to ...?
+		imageRdr, err := v.Image.Reader()
+		if err != nil {
+			respondWithError(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+		config, _, err := image.DecodeConfig(imageRdr)
+		if err != nil {
+			respondWithError(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+		outPixels = int64(config.Height) * int64(config.Width)
 	default:
 		respondWithError(w, "Unknown request type", http.StatusBadRequest)
 		return
@@ -385,6 +395,8 @@ func handleAIRequest(ctx context.Context, w http.ResponseWriter, r *http.Request
 			if err == nil {
 				latencyScore = CalculateAudioToTextLatencyScore(took, durationSeconds)
 			}
+		case worker.SegmentAnything2MultipartRequestBody:
+			latencyScore = CalculateSegmentAnything2LatencyScore(took, outPixels)
 		}
 
 		var pricePerAIUnit float64

diff --git a/server/ai_process.go b/server/ai_process.go
@@ -587,13 +587,21 @@ func submitUpscale(ctx context.Context, params aiRequestParams, sess *AISession,
 	return resp.JSON200, nil
 }
 
-func processSegmentAnything2(ctx context.Context, params aiRequestParams, req worker.SegmentAnything2MultipartRequestBody) (*worker.SegmentAnything2Response, error) {
+func CalculateSegmentAnything2LatencyScore(took time.Duration, outPixels int64) float64 {
+	if outPixels <= 0 {
+		return 0
+	}
+
+	return took.Seconds() / float64(outPixels)
+}
+
+func processSegmentAnything2(ctx context.Context, params aiRequestParams, req worker.SegmentAnything2MultipartRequestBody) (*worker.MasksResponse, error) {
 	resp, err := processAIRequest(ctx, params, req)
 	if err != nil {
 		return nil, err
 	}
 
-	txtResp := resp.(*worker.SegmentAnything2Response)
+	txtResp := resp.(*worker.MasksResponse)
 
 	return txtResp, nil
 }