Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
image_diffusion.sh 3.77 KB
Newer Older
ishandhanani's avatar
feat: sglang to 0.5.9 + updated docs (#6518)  
ishandhanani committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87 
#!/bin/bash
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Image diffusion worker (text-to-image). Default model: FLUX.1-dev (~38 GB VRAM).
# GPUs: 1

set -e

# Setup cleanup trap
cleanup() {
    echo "Cleaning up background processes..."
    kill $FRONTEND_PID 2>/dev/null || true
    wait $FRONTEND_PID 2>/dev/null || true
    echo "Cleanup complete."
}
trap cleanup EXIT INT TERM

# Defaults
MODEL_PATH="black-forest-labs/FLUX.1-dev"
FS_URL="file:///tmp/dynamo_media"
HTTP_URL=""
HTTP_PORT="${HTTP_PORT:-8000}"

# Parse command line arguments
EXTRA_ARGS=()
while [[ $# -gt 0 ]]; do
    case $1 in
        --model-path)
            MODEL_PATH="$2"
            shift 2
            ;;
        --fs-url)
            FS_URL="$2"
            shift 2
            ;;
        --http-url)
            HTTP_URL="$2"
            shift 2
            ;;
        --http-port)
            HTTP_PORT="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Launch a Dynamo image diffusion worker."
            echo ""
            echo "Options:"
            echo "  --model-path <path>          Model path (default: black-forest-labs/FLUX.1-dev)"
            echo "  --fs-url <url>               Filesystem URL for image storage (default: file:///tmp/dynamo_media)"
            echo "  --http-url <url>             Base URL for serving images over HTTP (optional)"
            echo "  --http-port <port>           Frontend HTTP port (default: 8000)"
            echo "  -h, --help                   Show this help message"
            echo ""
            echo "Additional flags are forwarded to dynamo.sglang."
            echo ""
            echo "Examples:"
            echo "  # Local file storage"
            echo "  $0 --model-path black-forest-labs/FLUX.1-dev --fs-url file:///tmp/images"
            echo ""
            echo "  # S3 storage (set FSSPEC_S3_KEY, FSSPEC_S3_SECRET, optionally FSSPEC_S3_ENDPOINT_URL)"
            echo "  $0 --fs-url s3://my-bucket/images"
            exit 0
            ;;
        *)
            EXTRA_ARGS+=("$1")
            shift
            ;;
    esac
done

echo "=========================================="
echo "Launching Image Diffusion Worker"
echo "=========================================="
echo "Model:       $MODEL_PATH"
echo "Frontend:    http://localhost:$HTTP_PORT"
echo "FS URL:      $FS_URL"
[ -n "$HTTP_URL" ] && echo "HTTP URL:    $HTTP_URL"
echo "=========================================="
echo ""
echo "Example test command:"
echo ""
echo "  curl http://localhost:${HTTP_PORT}/v1/images/generations \\"
echo "    -H 'Content-Type: application/json' \\"
echo "    -d '{"
Alec's avatar
docs: restructure vLLM docs and add startup banners to launch scripts (#6698)  
Alec committed
88 
echo "      \"prompt\": \"Explain why Roger Federer is considered one of the greatest tennis players of all time\","
ishandhanani's avatar
feat: sglang to 0.5.9 + updated docs (#6518)  
ishandhanani committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124 
echo "      \"model\": \"${MODEL_PATH}\","
echo "      \"size\": \"1024x1024\","
echo "      \"response_format\": \"url\","
echo "      \"nvext\": {"
echo "        \"num_inference_steps\": 15"
echo "      }"
echo "    }'"
echo ""
echo "=========================================="

# Build optional HTTP URL arg
HTTP_URL_ARGS=()
if [ -n "$HTTP_URL" ]; then
    HTTP_URL_ARGS=(--media-output-http-url "$HTTP_URL")
fi

# Launch frontend
echo "Starting Dynamo Frontend on port $HTTP_PORT..."
python3 -m dynamo.frontend \
    --http-port "$HTTP_PORT" &
FRONTEND_PID=$!

sleep 2

# Launch image diffusion worker
echo "Starting Image Diffusion Worker..."
python3 -m dynamo.sglang \
    --model-path "$MODEL_PATH" \
    --served-model-name "$MODEL_PATH" \
    --image-diffusion-worker \
    --media-output-fs-url "$FS_URL" \
    "${HTTP_URL_ARGS[@]}" \
    --trust-remote-code \
    --skip-tokenizer-init \
    --enable-metrics \
    "${EXTRA_ARGS[@]}"