ai/live: Configure parameters from auth callback

server/ai_live_video.go

@@ -20,7 +20,7 @@
 	if err != nil {
 		slog.Info("error publishing trickle", "err", err)
 	}
-	params.segmentReader.SwitchReader(func(reader io.Reader) {
+	params.liveParams.segmentReader.SwitchReader(func(reader io.Reader) {
 		// check for end of stream
 		if _, eos := reader.(*media.EOSReader); eos {
 			if err := publisher.Close(); err != nil {
@@ -67,7 +67,7 @@
 		ffmpeg.Transcode3(&ffmpeg.TranscodeOptionsIn{
 			Fname: fmt.Sprintf("pipe:%d", r.Fd()),
 		}, []ffmpeg.TranscodeOptions{{
-			Oname:        params.outputRTMPURL,
+			Oname:        params.liveParams.outputRTMPURL,
 			AudioEncoder: ffmpeg.ComponentOptions{Name: "copy"},
 			VideoEncoder: ffmpeg.ComponentOptions{Name: "copy"},
 			Muxer:        ffmpeg.ComponentOptions{Name: "flv"},
@@ -88,15 +88,21 @@
 
 func startControlPublish(control *url.URL, params aiRequestParams) {
 	controlPub, err := trickle.NewTricklePublisher(control.String())
+	stream := params.liveParams.stream
 	if err != nil {
-		slog.Info("error starting control publisher", "stream", params.stream, "err", err)
+		slog.Info("error starting control publisher", "stream", stream, "err", err)
 		return
 	}
 	params.node.LiveMu.Lock()
 	defer params.node.LiveMu.Unlock()
-	params.node.LivePipelines[params.stream] = &core.LivePipeline{ControlPub: controlPub}
+	params.node.LivePipelines[stream] = &core.LivePipeline{ControlPub: controlPub}
 }
 
+const (
+	mediaMTXControlPort = "9997"
+	mediaMTXControlUser = "admin"
+)
+
 func (ls *LivepeerServer) kickInputConnection(mediaMTXHost, sourceID, sourceType string) error {
 	var apiPath string
 	switch sourceType {

server/ai_mediaserver.go

@@ -420,19 +420,45 @@
 			return
 		}
 
-		err = authenticateAIStream(AuthWebhookURL, AIAuthRequest{
-			Stream:      streamName,
-			Type:        sourceTypeStr,
-			QueryParams: queryParams,
-		})
-		if err != nil {
-			kickErr := ls.kickInputConnection(remoteHost, sourceID, sourceType)
-			if kickErr != nil {
-				clog.Errorf(ctx, "failed to kick input connection: %s", kickErr.Error())
+		// if auth webhook returns pipeline config these will be replaced
+		pipeline := qp.Get("pipeline")
+		rawParams := qp.Get("params")
+		var pipelineParams map[string]interface{}
+		if rawParams != "" {
+			if err := json.Unmarshal([]byte(rawParams), &pipelineParams); err != nil {
+				clog.Errorf(ctx, "Invalid pipeline params: %s", err)
+				http.Error(w, "Invalid model params", http.StatusBadRequest)
+				return
+			}
+		}
+
+		if AuthWebhookURL != nil {
+			authResp, err := authenticateAIStream(AuthWebhookURL, AIAuthRequest{
+				Stream:      streamName,
+				Type:        sourceTypeStr,
+				QueryParams: queryParams,
+			})
+			if err != nil {
+				kickErr := ls.kickInputConnection(remoteHost, sourceID, sourceType)
+				if kickErr != nil {
+					clog.Errorf(ctx, "failed to kick input connection: %s", kickErr.Error())
+				}
+				clog.Errorf(ctx, "Live AI auth failed: %s", err.Error())
+				http.Error(w, "Forbidden", http.StatusForbidden)
+				return
+			}
+
+			if authResp.RTMPOutputURL != "" {
+				outputURL = authResp.RTMPOutputURL
+			}
+
+			if authResp.Pipeline != "" {
+				pipeline = authResp.Pipeline
+			}
+
+			if len(authResp.paramsMap) > 0 {
+				pipelineParams = authResp.paramsMap
 			}
-			clog.Errorf(ctx, "Live AI auth failed: %s", err.Error())
-			http.Error(w, "Forbidden", http.StatusForbidden)
-			return
 		}
 
 		requestID := string(core.RandomManifestID())
@@ -449,16 +475,20 @@
 		}()
 
 		params := aiRequestParams{
-			node:          ls.LivepeerNode,
-			os:            drivers.NodeStorage.NewSession(requestID),
-			sessManager:   ls.AISessionManager,
-			segmentReader: ssr,
-			outputRTMPURL: outputURL,
-			stream:        streamName,
+			node:        ls.LivepeerNode,
+			os:          drivers.NodeStorage.NewSession(requestID),
+			sessManager: ls.AISessionManager,
+
+			liveParams: liveRequestParams{
+				segmentReader: ssr,
+				outputRTMPURL: outputURL,
+				stream:        streamName,
+			},
 		}
 
 		req := worker.GenLiveVideoToVideoJSONRequestBody{
-			// TODO set model and initial parameters here if necessary (eg, prompt)
+			ModelId: &pipeline,
+			Params:  &pipelineParams,
 		}
 		processAIRequest(ctx, params, req)
 	})
@@ -523,8 +553,3 @@
 		}
 	}
 }
-
-const (
-	mediaMTXControlPort = "9997"
-	mediaMTXControlUser = "admin"
-)

server/ai_process.go

@@ -86,7 +86,11 @@
 	os          drivers.OSSession
 	sessManager *AISessionManager
 
-	// For live video pipelines
+	liveParams liveRequestParams
+}
+
+// For live video pipelines
+type liveRequestParams struct {
 	segmentReader *media.SwitchableSegmentReader
 	outputRTMPURL string
 	stream        string
@@ -1402,7 +1406,7 @@
 	case worker.GenLiveVideoToVideoJSONRequestBody:
 		cap = core.Capability_LiveVideoToVideo
 		modelID = defaultLiveVideoToVideoModelID
-		if v.ModelId != nil {
+		if v.ModelId != nil && *v.ModelId != "" {
 			modelID = *v.ModelId
 		}
 		submitFn = func(ctx context.Context, params aiRequestParams, sess *AISession) (interface{}, error) {

server/auth.go

@@ -109,26 +109,39 @@
 	// TODO not sure what params we need yet
 }
 
-func authenticateAIStream(authURL *url.URL, req AIAuthRequest) error {
+// Contains the configuration parameters for this AI job
+type AIAuthResponse struct {
+	// Where to send the output video
+	RTMPOutputURL string `json:"rtmp_output_url""`
+
+	// Name of the pipeline to run
+	Pipeline string `json:"pipeline"`
+
+	// Parameters for the pipeline
+	PipelineParams json.RawMessage        `json:"pipeline_parameters"`
+	paramsMap      map[string]interface{} // unmarshaled params
+}
+
+func authenticateAIStream(authURL *url.URL, req AIAuthRequest) (*AIAuthResponse, error) {
 	if authURL == nil {
-		return nil
+		return nil, fmt.Errorf("No auth URL configured")
 	}
 	started := time.Now()
 
 	jsonValue, err := json.Marshal(req)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	resp, err := http.Post(authURL.String(), "application/json", bytes.NewBuffer(jsonValue))
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	rbody, err := io.ReadAll(resp.Body)
 	resp.Body.Close()
 	if resp.StatusCode != 200 {
-		return fmt.Errorf("status=%d error=%s", resp.StatusCode, string(rbody))
+		return nil, fmt.Errorf("status=%d error=%s", resp.StatusCode, string(rbody))
 	}
 
 	took := time.Since(started)
@@ -137,5 +150,15 @@
 		monitor.AuthWebhookFinished(took)
 	}
 
-	return nil
+	var authResp AIAuthResponse
+	if err := json.Unmarshal(rbody, &authResp); err != nil {
+		return nil, err
+	}
+	if len(authResp.PipelineParams) > 0 {
+		if err := json.Unmarshal([]byte(authResp.PipelineParams), &authResp.paramsMap); err != nil {
+			return nil, err
+		}
+	}
+
+	return &authResp, nil
 }

server/auth_test.go

@@ -72,10 +72,11 @@ func TestAILiveAuthSucceeds(t *testing.T) {
 	s, serverURL := stubAuthServer(t, http.StatusOK, `{}`)
 	defer s.Close()
 
-	err := authenticateAIStream(serverURL, AIAuthRequest{
+	resp, err := authenticateAIStream(serverURL, AIAuthRequest{
 		Stream: "stream",
 	})
 	require.NoError(t, err)
+	require.Equal(t, AIAuthResponse{}, *resp)
 }
 
 func TestNoErrorWhenTranscodeAuthHeaderNotPassed(t *testing.T) {