@aws/ml-container-creator 0.2.0

package/templates/diffusors/Dockerfile

@@ -0,0 +1,72 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+<% if (comments && comments.acceleratorInfo) { %>
+<%= comments.acceleratorInfo %>
+<% } %>
+
+<% if (comments && comments.validationInfo) { %>
+<%= comments.validationInfo %>
+<% } %>
+
+# vLLM-Omni: Diffusion model serving with OpenAI-compatible image generation API
+# Image source: https://hub.docker.com/r/vllm/vllm-omni (public, no NGC auth required)
+FROM vllm/vllm-omni:v0.16.0
+
+# Set a docker label to name this project, postpended with the build time
+LABEL project.name="<%= projectName %>-<%= buildTimestamp %>" \
+      project.base-name="<%= projectName %>" \
+      project.build-time="<%= buildTimestamp %>"
+
+# Set a docker label to enable container to use SAGEMAKER_BIND_TO_PORT environment variable if present
+LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
+
+# Set the model name for the diffusion model
+ENV VLLM_MODEL="<%= modelName %>"
+
+<% if (hfToken) { %>
+# Set HuggingFace authentication token
+ENV HF_TOKEN="<%= hfToken %>"
+<% } %>
+
+<% if (comments && comments.envVarExplanations && Object.keys(comments.envVarExplanations).length > 0) { %>
+# Environment Variables Configuration
+<% for (const [category, comment] of Object.entries(comments.envVarExplanations)) { %>
+<%= comment %>
+<% } %>
+<% } %>
+
+<% if (orderedEnvVars && orderedEnvVars.length > 0) { %>
+# Additional environment variables from configuration
+<% orderedEnvVars.forEach(({ key, value }) => { %>
+ENV <%= key %>=<%= value %>
+<% }); %>
+<% } %>
+
+# Patch vLLM-Omni encode_image_base64 to handle numpy.ndarray outputs
+# Some diffusion models return numpy arrays instead of PIL Images from their
+# pipeline, which causes an AttributeError on .save() during response encoding.
+# This patch adds a safe isinstance guard to convert ndarray → PIL.Image.
+COPY code/patch_image_api.py /tmp/patch_image_api.py
+RUN python3 /tmp/patch_image_api.py && rm /tmp/patch_image_api.py
+
+# Install nginx for SageMaker endpoint routing
+# Maps /invocations -> /v1/images/generations and /ping -> /health
+RUN apt-get update && apt-get install -y --no-install-recommends nginx \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy nginx configuration
+COPY nginx-diffusors.conf /etc/nginx/nginx.conf
+
+# Copy serve entrypoint and startup scripts
+COPY code/serve /usr/bin/serve
+RUN chmod 777 /usr/bin/serve
+
+COPY code/start_server.sh /usr/bin/start_server.sh
+RUN chmod +x /usr/bin/start_server.sh
+
+<% if (comments && comments.troubleshooting) { %>
+<%= comments.troubleshooting %>
+<% } %>
+
+ENTRYPOINT [ "/usr/bin/serve" ]

package/templates/diffusors/patch_image_api.py

@@ -0,0 +1,35 @@
+"""
+Patch vLLM-Omni image_api_utils.encode_image_base64 to handle numpy arrays.
+
+Some diffusion models (especially video models like Wan2.1) return numpy
+arrays with extra batch/frame dimensions and float32 dtype. The upstream
+encode_image_base64() expects a PIL Image. This patch normalises the array
+to (H, W, 3) uint8 before converting to PIL.
+"""
+
+import vllm_omni.entrypoints.openai.image_api_utils as mod
+
+path = mod.__file__
+source = open(path).read()
+
+# Match the save() call site — resilient to function signature changes
+old = "    image.save(buffer, format=\"PNG\")"
+new = """\
+    import numpy as np
+    from PIL import Image as _PILImage
+    if isinstance(image, np.ndarray):
+        # Squeeze batch / frame dims: (1,1,H,W,3) or (1,H,W,3) -> (H,W,3)
+        while image.ndim > 3:
+            image = image[0]
+        # float [0,1] -> uint8 [0,255]
+        if image.dtype != np.uint8:
+            image = np.clip(image * 255.0, 0, 255).astype(np.uint8)
+        image = _PILImage.fromarray(image)
+    image.save(buffer, format="PNG")"""
+
+if old not in source:
+    print("WARN: patch target not found — already patched or API changed")
+else:
+    source = source.replace(old, new, 1)
+    open(path, "w").write(source)
+    print(f"Patched {path}")

package/templates/diffusors/serve

@@ -0,0 +1,115 @@
+#!/bin/bash
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+echo "Starting vLLM-Omni server (diffusion model serving)"
+
+# Resolve model URI prefixes that engines cannot handle natively.
+# The generator's model-picker may store provider-specific URIs
+# (e.g. jumpstart://model-txt2img-stabilityai-stable-diffusion-v2-1-base)
+# as the model identifier. vLLM expects a HuggingFace repo ID or local path.
+_RAW_MODEL="${VLLM_MODEL:-}"
+if [[ "$_RAW_MODEL" == jumpstart://* ]] || [[ "$_RAW_MODEL" == jumpstart-hub://* ]] || [[ "$_RAW_MODEL" == registry://* ]]; then
+    if [ -d /opt/ml/model ] && [ "$(ls -A /opt/ml/model 2>/dev/null)" ]; then
+        echo "Resolved VLLM_MODEL='${_RAW_MODEL}' → /opt/ml/model (local artifacts found)"
+        export VLLM_MODEL="/opt/ml/model"
+    else
+        _BARE_ID="${_RAW_MODEL#*://}"
+        echo "Warning: VLLM_MODEL='${_RAW_MODEL}' has a provider prefix but /opt/ml/model is empty."
+        echo "Stripping prefix → '${_BARE_ID}' (engine will attempt to fetch from model hub)"
+        export VLLM_MODEL="${_BARE_ID}"
+    fi
+fi
+unset _RAW_MODEL _BARE_ID
+
+# Validate that the model name is set
+if [ -z "$VLLM_MODEL" ]; then
+    echo "Error: VLLM_MODEL environment variable is not set"
+    exit 1
+fi
+
+# Initialize server arguments with --omni flag
+# --omni activates vLLM-Omni diffusion/multi-stage support
+# port 8081 is the internal port; nginx on 8080 handles SageMaker routing
+#   /invocations -> /v1/images/generations
+#   /ping -> /health
+SERVER_ARGS=(--omni --host 0.0.0.0 --port 8081)
+
+# Define the prefix for environment variables to look for
+# Uses VLLM_OMNI_ prefix to avoid conflicts with base vLLM env vars
+PREFIX="VLLM_OMNI_"
+ARG_PREFIX="--"
+
+# Define environment variables to exclude from CLI flag conversion
+# VLLM_MODEL is used as the positional model argument, not a --flag
+EXCLUDE_VARS=("VLLM_MODEL")
+
+# Declare and populate array of matching environment variables
+mapfile -t env_vars < <(env | grep "^${PREFIX}")
+
+# Loop through the array and convert to command-line arguments
+for var in "${env_vars[@]}"; do
+    IFS='=' read -r key value <<< "$var"
+
+    # Skip excluded variables
+    skip=false
+    for exclude in "${EXCLUDE_VARS[@]}"; do
+        if [ "$key" = "$exclude" ]; then
+            skip=true
+            break
+        fi
+    done
+
+    if [ "$skip" = true ]; then
+        continue
+    fi
+
+    # Remove prefix, convert to lowercase, and replace underscores with dashes
+    arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-')
+    SERVER_ARGS+=("${ARG_PREFIX}${arg_name}")
+    if [ -n "$value" ]; then
+        SERVER_ARGS+=("$value")
+    fi
+done
+
+echo "-------------------------------------------------------------------"
+echo "vLLM-Omni engine args: vllm serve $VLLM_MODEL [${SERVER_ARGS[@]}]"
+echo "-------------------------------------------------------------------"
+
+# Launch vLLM-Omni on internal port (8081), then nginx on SageMaker port (8080)
+vllm serve "$VLLM_MODEL" "${SERVER_ARGS[@]}" &
+VLLM_PID=$!
+
+# Wait for vLLM-Omni to be ready before starting nginx
+echo "Waiting for vLLM-Omni server to start..."
+for i in {1..300}; do
+    if curl -s http://localhost:8081/health > /dev/null 2>&1; then
+        echo "vLLM-Omni server is ready!"
+        break
+    fi
+    if ! kill -0 $VLLM_PID 2>/dev/null; then
+        echo "Error: vLLM-Omni process exited unexpectedly"
+        exit 1
+    fi
+    if [ $i -eq 300 ]; then
+        echo "Error: vLLM-Omni server failed to start within 300 seconds"
+        exit 1
+    fi
+    sleep 1
+done
+
+echo "Starting nginx reverse proxy on port 8080..."
+nginx -c /etc/nginx/nginx.conf &
+NGINX_PID=$!
+
+# Wait for either process to exit (this keeps the container running)
+wait -n $VLLM_PID $NGINX_PID
+
+# If we get here, one process exited - this is an error condition
+EXIT_CODE=$?
+echo "Error: Process exited with code $EXIT_CODE"
+
+# Kill any remaining processes
+kill $VLLM_PID $NGINX_PID 2>/dev/null || true
+
+exit $EXIT_CODE

package/templates/diffusors/start_server.sh

@@ -0,0 +1,114 @@
+#!/bin/bash
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# start_server.sh — Alternative startup script for vLLM-Omni diffusion model serving.
+# This script builds the vllm serve command with explicit diffusion-specific CLI flags.
+# The primary entrypoint is the `serve` script; use this for manual customization.
+
+set -e
+
+echo "Starting vLLM-Omni server (diffusion model serving)"
+
+# ---------------------------------------------------------------
+# Validate required environment variables
+# ---------------------------------------------------------------
+if [ -z "$VLLM_MODEL" ]; then
+    echo "Error: VLLM_MODEL environment variable is not set."
+    echo "Set it to a HuggingFace diffusion model ID, e.g.:"
+    echo "  export VLLM_MODEL=stabilityai/stable-diffusion-3.5-medium"
+    exit 1
+fi
+
+# ---------------------------------------------------------------
+# Build the base vllm serve command
+# --omni: activates vLLM-Omni diffusion/multi-stage support
+# --port 8080: SageMaker requires containers to listen on port 8080
+#   https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html
+# ---------------------------------------------------------------
+CMD="vllm serve $VLLM_MODEL --omni --host 0.0.0.0 --port 8081"
+
+# ---------------------------------------------------------------
+# Diffusion-specific CLI flags (configurable via environment variables)
+# ---------------------------------------------------------------
+
+# --num-gpus: Number of GPUs for diffusion inference (tensor parallelism)
+# Example: export VLLM_NUM_GPUS=4
+if [ -n "$VLLM_NUM_GPUS" ]; then
+    CMD="$CMD --num-gpus $VLLM_NUM_GPUS"
+fi
+
+# --cache-backend: Diffusion acceleration backend
+# Options: tea_cache (hook-based adaptive caching), cache_dit (library-based), none
+# Example: export VLLM_CACHE_BACKEND=tea_cache
+if [ -n "$VLLM_CACHE_BACKEND" ]; then
+    CMD="$CMD --cache-backend $VLLM_CACHE_BACKEND"
+fi
+
+# --vae-use-tiling: Process VAE in tiles to reduce VRAM usage
+# Set to any non-empty value to enable
+# Example: export VLLM_VAE_USE_TILING=1
+if [ -n "$VLLM_VAE_USE_TILING" ]; then
+    CMD="$CMD --vae-use-tiling"
+fi
+
+# --ulysses-degree: Ulysses sequence parallelism degree for large models
+# Example: export VLLM_ULYSSES_DEGREE=2
+if [ -n "$VLLM_ULYSSES_DEGREE" ]; then
+    CMD="$CMD --ulysses-degree $VLLM_ULYSSES_DEGREE"
+fi
+
+# --ring-degree: Ring sequence parallelism degree
+# Example: export VLLM_RING_DEGREE=2
+if [ -n "$VLLM_RING_DEGREE" ]; then
+    CMD="$CMD --ring-degree $VLLM_RING_DEGREE"
+fi
+
+# --enable-cpu-offload: Offload model weights to CPU to save GPU memory
+# Set to any non-empty value to enable
+# Example: export VLLM_ENABLE_CPU_OFFLOAD=1
+if [ -n "$VLLM_ENABLE_CPU_OFFLOAD" ]; then
+    CMD="$CMD --enable-cpu-offload"
+fi
+
+echo "-------------------------------------------------------------------"
+echo "vLLM-Omni command: $CMD"
+echo "-------------------------------------------------------------------"
+
+# Launch vLLM-Omni on internal port (8081), then nginx on SageMaker port (8080)
+$CMD &
+VLLM_PID=$!
+
+# Wait for vLLM-Omni to be ready before starting nginx
+echo "Waiting for vLLM-Omni server to start..."
+for i in {1..300}; do
+    if curl -s http://localhost:8081/health > /dev/null 2>&1; then
+        echo "vLLM-Omni server is ready!"
+        break
+    fi
+    if ! kill -0 $VLLM_PID 2>/dev/null; then
+        echo "Error: vLLM-Omni process exited unexpectedly"
+        exit 1
+    fi
+    if [ $i -eq 300 ]; then
+        echo "Error: vLLM-Omni server failed to start within 300 seconds"
+        exit 1
+    fi
+    sleep 1
+done
+
+echo "Starting nginx reverse proxy on port 8080..."
+nginx -c /etc/nginx/nginx.conf &
+NGINX_PID=$!
+
+# Wait for either process to exit (this keeps the container running)
+wait -n $VLLM_PID $NGINX_PID
+
+# If we get here, one process exited - this is an error condition
+EXIT_CODE=$?
+echo "Error: Process exited with code $EXIT_CODE"
+
+# Kill any remaining processes
+kill $VLLM_PID $NGINX_PID 2>/dev/null || true
+
+exit $EXIT_CODE

package/templates/do/.gitkeep

@@ -0,0 +1 @@
+# This file ensures the do/ directory is tracked by git