feat(runner): add diffusers pull script and upgrade to py3.11

Author: philwinder

.dockerignore

@@ -3,3 +3,4 @@ node_modules
 *.sw*
 venv
 vendor
+**/.venv/*

Dockerfile.runner

@@ -1,8 +1,20 @@
 #syntax=docker/dockerfile:1.4
 
-ARG TAG=2024-11-21a-empty
+ARG TAG=main-small
+ARG UV_VERSION="0.5.4"
 
-FROM ghcr.io/astral-sh/uv:0.5.4 as uv
+FROM ghcr.io/astral-sh/uv:${UV_VERSION} AS uv
+FROM ghcr.io/astral-sh/uv:${UV_VERSION}-bookworm-slim AS diffusers-build-env
+ENV UV_COMPILE_BYTECODE=1 UV_LINK_MODE=copy UV_PYTHON_INSTALL_DIR=/workspace/helix/runner/helix-diffusers/.python
+WORKDIR /workspace/helix/runner/helix-diffusers
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=runner/helix-diffusers/.python-version,target=.python-version \
+    --mount=type=bind,source=runner/helix-diffusers/uv.lock,target=uv.lock \
+    --mount=type=bind,source=runner/helix-diffusers/pyproject.toml,target=pyproject.toml \
+    uv sync --frozen --no-install-project --no-dev
+ADD runner/helix-diffusers /workspace/helix/runner/helix-diffusers
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv sync --frozen --no-dev
 
 ### BUILD
 
@@ -47,14 +59,10 @@ WORKDIR /workspace/helix
 # Copy runner directory from the repo
 COPY runner ./runner
 
-# We need to set this environment variable so that uv knows where
-# the virtual environment is to install packages
-ENV UV_PROJECT_ENVIRONMENT=/workspace/helix/runner/helix-diffusers/venv
-
-# Install the packages with uv using --mount=type=cache to cache the downloaded packages
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=from=uv,source=/uv,target=/usr/bin/uv \
-    cd /workspace/helix/runner/helix-diffusers && uv sync --no-dev
+# Copy the diffusers build environment including Python
+COPY --from=ghcr.io/astral-sh/uv:0.5.4 /uv /bin/uv
+COPY --from=diffusers-build-env /workspace/helix/runner/helix-diffusers /workspace/helix/runner/helix-diffusers
+ENV PATH="/workspace/helix/runner/helix-diffusers/.venv/bin:$PATH"
 
 # 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

api/pkg/model/models.go

@@ -122,7 +122,7 @@ func ProcessModelName(
 			}
 		}
 	case types.SessionTypeImage:
-		return Model_Diffusers_SD35, nil
+		return Model_Diffusers_SDTurbo, nil
 	}
 
 	// shouldn't get here
@@ -157,6 +157,7 @@ const (
 	Model_Axolotl_Mistral7b string = "mistralai/Mistral-7B-Instruct-v0.1"
 	Model_Cog_SDXL          string = "stabilityai/stable-diffusion-xl-base-1.0"
 	Model_Diffusers_SD35    string = "stabilityai/stable-diffusion-3.5-medium"
+	Model_Diffusers_SDTurbo string = "stabilityai/sd-turbo"
 
 	// We only need constants for _some_ ollama models that are hardcoded in
 	// various places (backward compat). Other ones can be added dynamically now.
@@ -170,10 +171,10 @@ const (
 func GetDefaultDiffusersModels() ([]*DiffusersGenericImage, error) {
 	return []*DiffusersGenericImage{
 		{
-			Id:          Model_Diffusers_SD35,
-			Name:        "Stable Diffusion 3.5 Medium",
-			Memory:      GB * 21,
-			Description: "Medium model, from Stability AI",
+			Id:          Model_Diffusers_SDTurbo,
+			Name:        "Stable Diffusion Turbo",
+			Memory:      GB * 5,
+			Description: "Turbo model, from Stability AI",
 			Hide:        false,
 		},
 	}, nil

api/pkg/runner/diffusers_model_instance.go

@@ -307,7 +307,7 @@ func (i *DiffusersModelInstance) Start(ctx context.Context) error {
 	if i.filter.Mode == types.SessionModeInference {
 		cmd = exec.CommandContext(
 			ctx,
-			"bash", "/workspace/helix/runner/venv_command.sh",
+			"uv", "run",
 			"uvicorn", "main:app",
 			"--host", "0.0.0.0",
 			"--port", strconv.Itoa(i.port),
@@ -321,14 +321,9 @@ func (i *DiffusersModelInstance) Start(ctx context.Context) error {
 	// Set the working directory to the runner dir (which makes relative path stuff easier)
 	cmd.Dir = "/workspace/helix/runner/helix-diffusers"
 
-	// 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)),
+		// Add the HF_TOKEN environment variable which is required by the diffusers library
+		fmt.Sprintf("HF_TOKEN=hf_ISxQhTIkdWkfZgUFPNUwVtHrCpMiwOYPIEKEN=%s", os.Getenv("HF_TOKEN")),
 		// Set python to be unbuffered so we get logs in real time
 		"PYTHONUNBUFFERED=1",
 	)

api/pkg/server/handlers.go

@@ -223,7 +223,7 @@ func (apiServer *HelixAPIServer) createSession(res http.ResponseWriter, req *htt
 			modelName = model.Model_Axolotl_Mistral7b
 		}
 	case types.SessionTypeImage:
-		modelName = model.Model_Diffusers_SD35
+		modelName = model.Model_Diffusers_SDTurbo
 	}
 
 	sessionID := system.GenerateUUID()

runner/helix-diffusers/.python-version

@@ -1 +1 @@
-3.10
+3.11

runner/helix-diffusers/main.py

@@ -9,9 +9,7 @@
 
 import PIL
 import torch
-from diffusers.pipelines.stable_diffusion_3 import (
-    StableDiffusion3Pipeline,
-)
+from diffusers import AutoPipelineForText2Image
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
@@ -25,7 +23,7 @@
 server_host = os.getenv("SERVER_HOST", "0.0.0.0")
 server_port = int(os.getenv("SERVER_PORT", 8000))
 server_url = f"http://{server_host}:{server_port}"
-model_id = os.getenv("MODEL_ID", "stabilityai/stable-diffusion-3.5-medium")
+model_id = os.getenv("MODEL_ID", "stabilityai/sd-turbo")
 
 
 class TextToImageInput(BaseModel):
@@ -46,16 +44,18 @@ def start(self, model_id: str):
             if torch.cuda.is_available():
                 logger.info("Loading CUDA")
                 self.device = "cuda"
-                self.pipeline = StableDiffusion3Pipeline.from_pretrained(
+                self.pipeline = AutoPipelineForText2Image.from_pretrained(
                     model_id,
                     torch_dtype=torch.bfloat16,
+                    local_files_only=True,
                 ).to(device=self.device)
             elif torch.backends.mps.is_available():
                 logger.info("Loading MPS for Mac M Series")
                 self.device = "mps"
-                self.pipeline = StableDiffusion3Pipeline.from_pretrained(
+                self.pipeline = AutoPipelineForText2Image.from_pretrained(
                     model_id,
                     torch_dtype=torch.bfloat16,
+                    local_files_only=True,
                 ).to(device=self.device)
             else:
                 raise Exception("No CUDA or MPS device available")

runner/helix-diffusers/pyproject.toml

@@ -3,15 +3,15 @@ name = "helix-diffusers"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
     "accelerate>=1.1.1",
     "diffusers>=0.31.0",
     "fastapi>=0.115.5",
     "httpx>=0.27.2",
     "protobuf>=5.28.3",
     "sentencepiece>=0.2.0",
-    "torch>=2.5.1",
+    "torch==2.5.1+cu124",
     "transformers>=4.46.3",
     "uvicorn>=0.32.1",
 ]
@@ -24,6 +24,14 @@ dev = [
     "ruff>=0.8.0",
 ]
 
+[tool.uv.sources]
+torch = { index = "pytorch-cu124" }
+
+[[tool.uv.index]]
+name = "pytorch-cu124"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true
+
 [tool.ruff]
 line-length = 100
 indent-width = 4
@@ -33,4 +41,4 @@ target-version = "py311"
 filterwarnings = [
     "ignore::UserWarning",
     "ignore::DeprecationWarning",
-]
+]