fix(api): quick fix to sdxl image inference

.drone.yml

@@ -162,7 +162,7 @@ steps:
       # Runner with no baked models = empty
       # See https://github.com/helixml/base-images
       # and https://github.com/helixml/base-images/releases
-      - TAG=2024-11-12a-empty
+      - TAG=2024-11-21a-empty
       - APP_VERSION=${DRONE_TAG:-${DRONE_COMMIT_SHA:-latest}}
     username: admin
     password:
@@ -204,7 +204,7 @@ steps:
       # Runner with small models = small
       # See https://github.com/helixml/base-images
       # and https://github.com/helixml/base-images/releases
-      - TAG=2024-11-12a-small
+      - TAG=2024-11-21a-small
       - APP_VERSION=${DRONE_TAG:-${DRONE_COMMIT_SHA:-latest}}
     username: admin
     password:
@@ -232,7 +232,7 @@ steps:
       # Runner with small models = small
       # See https://github.com/helixml/base-images
       # and https://github.com/helixml/base-images/releases
-      - TAG=2024-11-12a-small
+      - TAG=2024-11-21a-small
       - APP_VERSION=${DRONE_TAG:-${DRONE_COMMIT_SHA:-latest}}
     username: admin
     password:
@@ -277,7 +277,7 @@ steps:
       # Runner with large models = large
       # See https://github.com/helixml/base-images
       # and https://github.com/helixml/base-images/releases
-      - TAG=2024-11-12a-large
+      - TAG=2024-11-21a-large
       - APP_VERSION=${DRONE_TAG:-${DRONE_COMMIT_SHA:-latest}}
     username: admin
     password:
@@ -305,7 +305,7 @@ steps:
       # Runner with large models = large
       # See https://github.com/helixml/base-images
       # and https://github.com/helixml/base-images/releases
-      - TAG=2024-11-12a-large
+      - TAG=2024-11-21a-large
       - APP_VERSION=${DRONE_TAG:-${DRONE_COMMIT_SHA:-latest}}
     username: admin
     password:

Dockerfile.runner

@@ -1,6 +1,6 @@
 #syntax=docker/dockerfile:1.4
 
-ARG TAG=2024-11-12a-empty
+ARG TAG=2024-11-21a-empty
 
 ### BUILD
 
@@ -45,6 +45,9 @@ WORKDIR /workspace/helix
 # Copy runner directory from the repo
 COPY runner ./runner
 
+# Copy the cog wrapper, cog and cog-sdxl is installed in the base image, this is just the cog server
+COPY cog/helix_cog_wrapper.py /workspace/cog-sdxl/helix_cog_wrapper.py
+
 # So that the runner can function when run as non-root, symlink some stuff into
 # locations in /tmp (needed for locked down OpenShift support)
 RUN mkdir -p /tmp/helix/ollama /tmp/helix/src /tmp/helix/cache /tmp/helix/root-cache /tmp/helix/config /workspace/axolotl/dataset_cache && \

api/pkg/model/cog_sdxl.go

@@ -165,43 +165,44 @@ func (l *CogSDXL) GetCommand(ctx context.Context, sessionFilter types.SessionFil
 	}
 	var cmd *exec.Cmd
 	if sessionFilter.Mode == types.SessionModeInference {
-		args := []string{
-			"runner/venv_command.sh",
-			"python3", "-u",
-			"helix_cog_wrapper.py", "inference",
-		}
-
 		cmd = exec.CommandContext(
 			ctx,
-			"bash",
-			args...,
+			"bash", "/workspace/helix/runner/venv_command.sh",
+			"uvicorn", "helix_cog_wrapper:app",
+			"--host", "0.0.0.0",
+			"--port", strconv.Itoa(config.Port),
 		)
 	} else if sessionFilter.Mode == types.SessionModeFinetune {
 		cmd = exec.CommandContext(
 			ctx,
-			"bash", "runner/venv_command.sh",
+			"bash", "/workspace/helix/runner/venv_command.sh",
 			"python3", "-u",
 			"helix_cog_wrapper.py", "finetune",
 		)
 	} else {
 		return nil, fmt.Errorf("invalid session mode: %s", sessionFilter.Mode)
 	}
 
-	wd, err := os.Getwd()
-	if err != nil {
-		return nil, err
-	}
+	// Set the working directory to the runner dir (which makes relative path stuff easier)
+	cmd.Dir = "/workspace/cog-sdxl"
 
-	cmd.Env = []string{
-		fmt.Sprintf("APP_FOLDER=%s", path.Clean(path.Join(wd, "..", "cog-sdxl"))),
-		fmt.Sprintf("HELIX_NEXT_TASK_URL=%s", config.NextTaskURL),
-		fmt.Sprintf("HELIX_INITIAL_SESSION_URL=%s", config.InitialSessionURL),
+	// Inherit all the parent environment variables
+	cmd.Env = append(cmd.Env,
+		os.Environ()...,
+	)
+
+	cmd.Env = append(cmd.Env,
+		// Add the APP_FOLDER environment variable which is required by the old code
+		fmt.Sprintf("APP_FOLDER=%s", path.Clean(cmd.Dir)),
+		// Set python to be unbuffered so we get logs in real time
+		"PYTHONUNBUFFERED=1",
+		// Set the log level, which is a name, but must be uppercased
+		fmt.Sprintf("LOG_LEVEL=%s", strings.ToUpper(os.Getenv("LOG_LEVEL"))),
 		// cog likes to download LoRA from a URL, so we construct one for it
 		fmt.Sprintf("API_HOST=%s", API_HOST),
 		// one day it will need to auth to the API server to download LoRAs
 		fmt.Sprintf("API_TOKEN=%s", API_TOKEN),
-		"PYTHONUNBUFFERED=1",
-	}
+	)
 
 	return cmd, nil
 }

api/pkg/model/models.go

@@ -50,6 +50,9 @@ func (m ModelName) InferenceRuntime() types.InferenceRuntime {
 	if strings.Contains(m.String(), ":") {
 		return types.InferenceRuntimeOllama
 	}
+	if m.String() == Model_Cog_SDXL {
+		return types.InferenceRuntimeCog
+	}
 	// misnamed: axolotl runtime handles axolotl and cog/sd-scripts
 	return types.InferenceRuntimeAxolotl
 }