Skip to content

WIP add batch-size 2 versions to trt builds #220

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pschroedl wants to merge 1 commit into from
Closed

WIP add batch-size 2 versions to trt builds #220

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 39 additions & 19 deletions docker/entrypoint.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,25 +53,45 @@ if [ "$1" = "--build-engines" ]; then
cd /workspace/comfystream
conda activate comfystream

# Build Static Engine for Dreamshaper - Square (512x512)
python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-512-w-512_00001_.engine --width 512 --height 512

# Build Static Engine for Dreamshaper - Portrait (384x704)
python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-384-w-704_00001_.engine --width 384 --height 704

# Build Static Engine for Dreamshaper - Landscape (704x384)
python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-704-w-384_00001_.engine --width 704 --height 384

# Build Dynamic Engine for Dreamshaper
python src/comfystream/scripts/build_trt.py \
--model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \
--out-engine /workspace/ComfyUI/output/tensorrt/dynamic-dreamshaper8_SD15_\$dyn-b-1-4-2-h-512-704-w-320-384-448_00001_.engine \
--width 384 \
--height 704 \
--min-width 320 \
--min-height 512 \
--max-width 448 \
--max-height 704
# Build Static Engines for Dreamshaper - Square (512x512)
BATCH_SIZES="1 2"
for batch_size in $BATCH_SIZES; do
echo "=== Building engines with batch size: $batch_size ==="

# Build Static Engine for Dreamshaper - Square (512x512)
python src/comfystream/scripts/build_trt.py \
--model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \
--out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_stat-b-${batch_size}-h-512-w-512_00001_.engine \
--width 512 --height 512 \
--batch-size $batch_size

# Build Static Engine for Dreamshaper - Square (512x512) -- FP8
python src/comfystream/scripts/build_trt.py \
--model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \
--out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_stat-b-${batch_size}-h-512-w-512_00001_.engine \
--width 512 --height 512 \
--batch-size $batch_size \
--FP8

# Build Static Engine for Dreamshaper - Landscape (704x384)
python src/comfystream/scripts/build_trt.py \
--model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \
--out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_stat-b-${batch_size}-h-704-w-384_00001_.engine \
--width 704 --height 384 \
--batch-size $batch_size

# Build Dynamic Engine for Dreamshaper
python src/comfystream/scripts/build_trt.py \
--model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors \
--out-engine /workspace/ComfyUI/output/tensorrt/dynamic-dreamshaper8_SD15_dyn-b-${batch_size}-h-512-704-w-320-384-448_00001_.engine \
--width 384 \
--height 704 \
--min-width 320 \
--min-height 512 \
--max-width 448 \
--max-height 704 \
--batch-size $batch_size
Comment on lines +84 to +93
Copy link
Collaborator

@eliteprox eliteprox May 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This engine is not used by ai-runner/daydream and is only for example workflows. We should at least exclude it from the batch loop to minimize length of build process. I'd also be fine with removing it entirely for now. wdyt?

Copy link
Author

@pschroedl pschroedl May 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense, given that I primarily created this branch to unblock testing with these models for the daydream workflow - we don't need or want to deploy these ( yet )

done

# Build Engine for Depth Anything V2
if [ ! -f "$DEPTH_ANYTHING_DIR/$DEPTH_ANYTHING_ENGINE" ]; then
Expand Down