815 lines
26 KiB
Bash
Executable file
815 lines
26 KiB
Bash
Executable file
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
cd "$ROOT_DIR"
|
|
|
|
WORKFLOW="${FLUX_WORKFLOW:-examples/comfyui/flux_dev_full_text_to_image.json}"
|
|
MODEL_ROOT="${FLUX_MODEL_ROOT:-local/diffusion_models}"
|
|
WIDTH="${FLUX_WIDTH:-256}"
|
|
HEIGHT="${FLUX_HEIGHT:-}"
|
|
STEPS="${FLUX_STEPS:-20}"
|
|
|
|
DUMP_DIR="${FLUX_DUMP_DIR:-/tmp/flux_cond_ref}"
|
|
STEP_DIR="${FLUX_STEP_DIR:-/tmp/flux_step_ref}"
|
|
REF_OUTPUT="${FLUX_REF_OUTPUT:-/tmp/sdcpp_flux_cond_dump.png}"
|
|
NATIVE_OUTPUT="${FLUX_NATIVE_OUTPUT:-/tmp/flux_native_oracle.png}"
|
|
HANDOFF_OUTPUT="${FLUX_HANDOFF_OUTPUT:-/tmp/flux_handoff_oracle.png}"
|
|
WARM_OUTPUT="${FLUX_WARM_OUTPUT:-}"
|
|
|
|
DIFFUSION_MODEL="${FLUX_REF_DIFFUSION_MODEL:-}"
|
|
VAE_MODEL="${FLUX_REF_VAE_MODEL:-}"
|
|
CLIP_L_MODEL="${FLUX_REF_CLIP_L_MODEL:-}"
|
|
T5XXL_MODEL="${FLUX_REF_T5XXL_MODEL:-}"
|
|
|
|
REF_PROMPT="${FLUX_REF_PROMPT:-}"
|
|
REF_SEED="${FLUX_REF_SEED:-}"
|
|
REF_CFG_SCALE="${FLUX_REF_CFG_SCALE:-}"
|
|
DUMP_STEP_INDEX="${FLUX_DUMP_STEP_INDEX:-1}"
|
|
COND_DIR="${FLUX_COND_DIR:-}"
|
|
COMPARE_COND_DIR="${FLUX_COMPARE_COND_DIR:-}"
|
|
T5_DEBUG_DIR="${FLUX_T5_DEBUG_DIR:-}"
|
|
T5_STAGE_LAYER="${FLUX_T5_DEBUG_LAYER:-}"
|
|
T5_MODE="${FLUX_T5_MODE:-lazy}"
|
|
WARMUP_RUNS="${FLUX_WARMUP_RUNS:-1}"
|
|
MEASURED_RUNS="${FLUX_MEASURED_RUNS:-1}"
|
|
LHS_DIR="${FLUX_LHS_DIR:-}"
|
|
RHS_DIR="${FLUX_RHS_DIR:-}"
|
|
T5_CPU_MATH="${FLUX_T5_FORCE_CPU_MATH:-0}"
|
|
T5_CPU_ATTN="${FLUX_T5_FORCE_CPU_ATTN:-0}"
|
|
T5_F32_LINEAR="${FLUX_T5_FORCE_F32_LINEAR:-0}"
|
|
|
|
CARGO_MANIFEST="libs/diffusion/Cargo.toml"
|
|
SDCPP_DIR="local/stable-diffusion.cpp"
|
|
SDCPP_BUILD_DIR="$SDCPP_DIR/build"
|
|
SDCLI_BIN="$SDCPP_BUILD_DIR/bin/sd-cli"
|
|
REF_WARM_BENCH_SRC="tools/flux_ref_warm_bench.cpp"
|
|
REF_WARM_BENCH_BIN="$SDCPP_BUILD_DIR/bin/flux-ref-warm-bench"
|
|
CMAKE_BIN="${CMAKE_BIN:-/Applications/CMake.app/Contents/bin/cmake}"
|
|
|
|
finalize_args() {
|
|
if [[ -z "$HEIGHT" ]]; then
|
|
HEIGHT="$WIDTH"
|
|
fi
|
|
if [[ -z "$DIFFUSION_MODEL" ]]; then
|
|
DIFFUSION_MODEL="$MODEL_ROOT/unet/flux1-dev.safetensors"
|
|
fi
|
|
if [[ -z "$VAE_MODEL" ]]; then
|
|
VAE_MODEL="$MODEL_ROOT/vae/ae.safetensors"
|
|
fi
|
|
if [[ -z "$CLIP_L_MODEL" ]]; then
|
|
CLIP_L_MODEL="$MODEL_ROOT/text_encoders/clip_l.safetensors"
|
|
fi
|
|
if [[ -z "$T5XXL_MODEL" ]]; then
|
|
T5XXL_MODEL="$MODEL_ROOT/text_encoders/t5xxl_fp16.safetensors"
|
|
fi
|
|
}
|
|
|
|
usage() {
|
|
cat <<'EOF'
|
|
usage: tools/flux_debug.sh <command> [options]
|
|
|
|
commands:
|
|
check
|
|
ref-build
|
|
ref-generate
|
|
ref-warm-bench
|
|
text-smoke
|
|
t5-smoke
|
|
t5-stage-compare
|
|
t5-debug-compare
|
|
transformer-smoke
|
|
native-generate
|
|
warm-bench
|
|
ref-dump
|
|
ref-step-dump
|
|
handoff-generate
|
|
transformer-ref-compare
|
|
transformer-ref-stage-compare
|
|
transformer-ref-compare-f32
|
|
oracle
|
|
|
|
options:
|
|
--workflow PATH comfy workflow json
|
|
--model-root PATH model root with unet/, vae/, text_encoders/
|
|
--width N output width
|
|
--height N output height
|
|
--steps N denoise step count
|
|
--dump-dir PATH reference conditioning dump directory
|
|
--step-dir PATH reference step dump directory
|
|
--dump-step-index N reference step index to dump
|
|
--ref-output PATH output image path for reference sd-cli run
|
|
--native-output PATH output image path for Rust native generation
|
|
--warm-output PATH output image path for the final measured warm bench run
|
|
--handoff-output PATH output image path for Rust handoff generation
|
|
--warmup-runs N in-process warmup run count for warm-bench
|
|
--measured-runs N in-process measured run count for warm-bench
|
|
--cond-dir PATH conditioning override directory for Rust runs
|
|
--compare-cond-dir PATH reference conditioning directory to diff against
|
|
--t5-debug-dir PATH dump per-block T5 hidden states to this directory
|
|
--t5-stage-layer N dump stage tensors for this T5 layer when stage debugging is enabled
|
|
--t5-mode MODE T5 backend: lazy or compiled
|
|
--t5-cpu-math force the lazy Rust T5 path to run with CPU math for debugging
|
|
--t5-cpu-attn force the lazy Rust T5 attention kernels onto the CPU
|
|
--t5-f32-linear decode T5 linear weights to F32 before lazy Metal matmuls
|
|
--lhs-dir PATH left-hand debug directory for compare commands
|
|
--rhs-dir PATH right-hand debug directory for compare commands
|
|
--ref-prompt TEXT override prompt used for reference sd-cli runs
|
|
--ref-seed N override seed used for reference sd-cli runs
|
|
--ref-cfg-scale N override cfg scale used for reference sd-cli runs
|
|
|
|
examples:
|
|
tools/flux_debug.sh oracle
|
|
tools/flux_debug.sh transformer-smoke --width 384 --height 384
|
|
tools/flux_debug.sh ref-dump --ref-prompt test
|
|
tools/flux_debug.sh ref-step-dump --dump-step-index 2
|
|
tools/flux_debug.sh native-generate --compare-cond-dir /tmp/flux_cond_ref
|
|
EOF
|
|
}
|
|
|
|
log() {
|
|
printf '[flux-debug] %s\n' "$*"
|
|
}
|
|
|
|
run() {
|
|
printf '+'
|
|
for arg in "$@"; do
|
|
printf ' %q' "$arg"
|
|
done
|
|
printf '\n'
|
|
"$@"
|
|
}
|
|
|
|
require_file() {
|
|
local path="$1"
|
|
if [[ ! -f "$path" ]]; then
|
|
printf 'missing file: %s\n' "$path" >&2
|
|
exit 1
|
|
fi
|
|
}
|
|
|
|
require_dir() {
|
|
local path="$1"
|
|
if [[ ! -d "$path" ]]; then
|
|
printf 'missing directory: %s\n' "$path" >&2
|
|
exit 1
|
|
fi
|
|
}
|
|
|
|
parse_args() {
|
|
while [[ $# -gt 0 ]]; do
|
|
case "$1" in
|
|
--workflow)
|
|
WORKFLOW="$2"
|
|
shift 2
|
|
;;
|
|
--model-root)
|
|
MODEL_ROOT="$2"
|
|
shift 2
|
|
;;
|
|
--width)
|
|
WIDTH="$2"
|
|
shift 2
|
|
;;
|
|
--height)
|
|
HEIGHT="$2"
|
|
shift 2
|
|
;;
|
|
--steps)
|
|
STEPS="$2"
|
|
shift 2
|
|
;;
|
|
--dump-dir)
|
|
DUMP_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--step-dir|--ref-step-dir)
|
|
STEP_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--dump-step-index)
|
|
DUMP_STEP_INDEX="$2"
|
|
shift 2
|
|
;;
|
|
--ref-output)
|
|
REF_OUTPUT="$2"
|
|
shift 2
|
|
;;
|
|
--native-output)
|
|
NATIVE_OUTPUT="$2"
|
|
shift 2
|
|
;;
|
|
--warm-output)
|
|
WARM_OUTPUT="$2"
|
|
shift 2
|
|
;;
|
|
--handoff-output)
|
|
HANDOFF_OUTPUT="$2"
|
|
shift 2
|
|
;;
|
|
--warmup-runs)
|
|
WARMUP_RUNS="$2"
|
|
shift 2
|
|
;;
|
|
--measured-runs)
|
|
MEASURED_RUNS="$2"
|
|
shift 2
|
|
;;
|
|
--cond-dir)
|
|
COND_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--compare-cond-dir)
|
|
COMPARE_COND_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--t5-debug-dir)
|
|
T5_DEBUG_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--t5-stage-layer)
|
|
T5_STAGE_LAYER="$2"
|
|
shift 2
|
|
;;
|
|
--t5-mode)
|
|
T5_MODE="$2"
|
|
shift 2
|
|
;;
|
|
--t5-cpu-math)
|
|
T5_CPU_MATH=1
|
|
shift
|
|
;;
|
|
--t5-cpu-attn)
|
|
T5_CPU_ATTN=1
|
|
shift
|
|
;;
|
|
--t5-f32-linear)
|
|
T5_F32_LINEAR=1
|
|
shift
|
|
;;
|
|
--lhs-dir)
|
|
LHS_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--rhs-dir)
|
|
RHS_DIR="$2"
|
|
shift 2
|
|
;;
|
|
--ref-prompt)
|
|
REF_PROMPT="$2"
|
|
shift 2
|
|
;;
|
|
--ref-seed)
|
|
REF_SEED="$2"
|
|
shift 2
|
|
;;
|
|
--ref-cfg-scale)
|
|
REF_CFG_SCALE="$2"
|
|
shift 2
|
|
;;
|
|
-h|--help)
|
|
usage
|
|
exit 0
|
|
;;
|
|
--)
|
|
shift
|
|
break
|
|
;;
|
|
*)
|
|
printf 'unknown option: %s\n\n' "$1" >&2
|
|
usage >&2
|
|
exit 1
|
|
;;
|
|
esac
|
|
done
|
|
if [[ $# -gt 0 ]]; then
|
|
printf 'unexpected extra arguments: %s\n\n' "$*" >&2
|
|
usage >&2
|
|
exit 1
|
|
fi
|
|
finalize_args
|
|
}
|
|
|
|
resolve_ref_prompt() {
|
|
if [[ -n "$REF_PROMPT" ]]; then
|
|
printf '%s\n' "$REF_PROMPT"
|
|
return
|
|
fi
|
|
cargo run --release --quiet --manifest-path "$CARGO_MANIFEST" --bin flux-tokenize -- "$WORKFLOW" \
|
|
| sed -n 's/^prompt\.t5xxl: //p' \
|
|
| head -n 1
|
|
}
|
|
|
|
resolve_ref_seed() {
|
|
if [[ -n "$REF_SEED" ]]; then
|
|
printf '%s\n' "$REF_SEED"
|
|
return
|
|
fi
|
|
sed -n 's/.*"seed"[[:space:]]*:[[:space:]]*\([-0-9][0-9]*\).*/\1/p' "$WORKFLOW" | head -n 1
|
|
}
|
|
|
|
resolve_ref_cfg_scale() {
|
|
if [[ -n "$REF_CFG_SCALE" ]]; then
|
|
printf '%s\n' "$REF_CFG_SCALE"
|
|
return
|
|
fi
|
|
sed -n 's/.*"cfg"[[:space:]]*:[[:space:]]*\([-0-9.][0-9.]*\).*/\1/p' "$WORKFLOW" | head -n 1
|
|
}
|
|
|
|
ensure_common_inputs() {
|
|
require_file "$WORKFLOW"
|
|
require_dir "$MODEL_ROOT"
|
|
}
|
|
|
|
ensure_reference_inputs() {
|
|
ensure_common_inputs
|
|
require_file "$DIFFUSION_MODEL"
|
|
require_file "$VAE_MODEL"
|
|
require_file "$CLIP_L_MODEL"
|
|
require_file "$T5XXL_MODEL"
|
|
}
|
|
|
|
ensure_ref_build() {
|
|
require_file "$CMAKE_BIN"
|
|
if [[ ! -f "$SDCPP_BUILD_DIR/CMakeCache.txt" ]]; then
|
|
run "$CMAKE_BIN" -S "$SDCPP_DIR" -B "$SDCPP_BUILD_DIR" -DCMAKE_BUILD_TYPE=Release
|
|
fi
|
|
run "$CMAKE_BIN" --build "$SDCPP_BUILD_DIR" --config Release --target sd-cli -j4
|
|
require_file "$SDCLI_BIN"
|
|
}
|
|
|
|
ensure_ref_warm_bench_bin() {
|
|
ensure_ref_build
|
|
require_file "$REF_WARM_BENCH_SRC"
|
|
if [[ ! -x "$REF_WARM_BENCH_BIN" || "$REF_WARM_BENCH_SRC" -nt "$REF_WARM_BENCH_BIN" || "$SDCPP_BUILD_DIR/libstable-diffusion.a" -nt "$REF_WARM_BENCH_BIN" ]]; then
|
|
local -a cmd=(
|
|
/usr/bin/c++
|
|
-O3
|
|
-DNDEBUG
|
|
-std=c++17
|
|
-arch arm64
|
|
-I "$SDCPP_DIR/include"
|
|
"$REF_WARM_BENCH_SRC"
|
|
-o "$REF_WARM_BENCH_BIN"
|
|
"$SDCPP_BUILD_DIR/libstable-diffusion.a"
|
|
"$SDCPP_BUILD_DIR/thirdparty/libwebp/libwebp.a"
|
|
"$SDCPP_BUILD_DIR/thirdparty/libwebp/libwebpmux.a"
|
|
"$SDCPP_BUILD_DIR/thirdparty/libwebm/libwebm.a"
|
|
"$SDCPP_BUILD_DIR/ggml/src/libggml.a"
|
|
"$SDCPP_BUILD_DIR/ggml/src/libggml-cpu.a"
|
|
"$SDCPP_BUILD_DIR/ggml/src/ggml-blas/libggml-blas.a"
|
|
-framework Accelerate
|
|
"$SDCPP_BUILD_DIR/ggml/src/ggml-metal/libggml-metal.a"
|
|
"$SDCPP_BUILD_DIR/ggml/src/libggml-base.a"
|
|
-lm
|
|
-framework Foundation
|
|
-framework Metal
|
|
-framework MetalKit
|
|
"$SDCPP_BUILD_DIR/thirdparty/libwebp/libwebp.a"
|
|
"$SDCPP_BUILD_DIR/thirdparty/libwebp/libsharpyuv.a"
|
|
)
|
|
run "${cmd[@]}"
|
|
fi
|
|
require_file "$REF_WARM_BENCH_BIN"
|
|
}
|
|
|
|
cmd_check() {
|
|
ensure_common_inputs
|
|
run cargo check --manifest-path "$CARGO_MANIFEST"
|
|
}
|
|
|
|
cmd_text_smoke() {
|
|
ensure_common_inputs
|
|
if [[ -n "$T5_DEBUG_DIR" ]]; then
|
|
mkdir -p "$T5_DEBUG_DIR"
|
|
fi
|
|
local -a env_args=()
|
|
if [[ -n "$COND_DIR" ]]; then
|
|
env_args+=("FLUX_COND_DIR=$COND_DIR")
|
|
fi
|
|
if [[ -n "$T5_DEBUG_DIR" ]]; then
|
|
env_args+=("FLUX_T5_DEBUG_DIR=$T5_DEBUG_DIR")
|
|
fi
|
|
if [[ -n "$T5_STAGE_LAYER" ]]; then
|
|
env_args+=("FLUX_T5_DEBUG_LAYER=$T5_STAGE_LAYER")
|
|
fi
|
|
if [[ -n "$T5_MODE" ]]; then
|
|
env_args+=("FLUX_T5_MODE=$T5_MODE")
|
|
fi
|
|
if [[ "$T5_CPU_MATH" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_MATH=1")
|
|
fi
|
|
if [[ "$T5_CPU_ATTN" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_ATTN=1")
|
|
fi
|
|
if [[ "$T5_F32_LINEAR" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_F32_LINEAR=1")
|
|
fi
|
|
if [[ ${#env_args[@]} -gt 0 ]]; then
|
|
run env "${env_args[@]}" cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-text-smoke -- "$WORKFLOW" "$MODEL_ROOT"
|
|
else
|
|
run cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-text-smoke -- "$WORKFLOW" "$MODEL_ROOT"
|
|
fi
|
|
}
|
|
|
|
cmd_t5_smoke() {
|
|
ensure_common_inputs
|
|
local -a env_args=("FLUX_T5_MODE=$T5_MODE")
|
|
if [[ "$T5_CPU_MATH" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_MATH=1")
|
|
fi
|
|
if [[ "$T5_CPU_ATTN" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_ATTN=1")
|
|
fi
|
|
if [[ "$T5_F32_LINEAR" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_F32_LINEAR=1")
|
|
fi
|
|
if [[ ${#env_args[@]} -gt 0 ]]; then
|
|
run env "${env_args[@]}" cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-t5-smoke -- "$WORKFLOW" "$MODEL_ROOT"
|
|
else
|
|
run cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-t5-smoke -- "$WORKFLOW" "$MODEL_ROOT"
|
|
fi
|
|
}
|
|
|
|
cmd_t5_stage_compare() {
|
|
ensure_common_inputs
|
|
local debug_dir="${T5_DEBUG_DIR:-/tmp/flux_t5_ggml_dbg}"
|
|
local stage_layer="${T5_STAGE_LAYER:-0}"
|
|
mkdir -p "$debug_dir"
|
|
rm -f "$debug_dir"/*.bin "$debug_dir"/*.txt
|
|
local -a env_args=(
|
|
"FLUX_T5_MODE=$T5_MODE"
|
|
"FLUX_T5_DEBUG_DIR=$debug_dir"
|
|
"FLUX_T5_DEBUG_STAGES=1"
|
|
"FLUX_T5_DEBUG_LAYER=$stage_layer"
|
|
)
|
|
if [[ "$T5_CPU_MATH" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_MATH=1")
|
|
fi
|
|
if [[ "$T5_CPU_ATTN" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_ATTN=1")
|
|
fi
|
|
if [[ "$T5_F32_LINEAR" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_F32_LINEAR=1")
|
|
fi
|
|
run env "${env_args[@]}" cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-text-smoke -- "$WORKFLOW" "$MODEL_ROOT"
|
|
run cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-t5-stage-compare -- "$WORKFLOW" "$MODEL_ROOT" "$debug_dir" "$stage_layer"
|
|
}
|
|
|
|
cmd_t5_debug_compare() {
|
|
local lhs_dir="${LHS_DIR:-}"
|
|
local rhs_dir="${RHS_DIR:-$T5_DEBUG_DIR}"
|
|
if [[ -z "$lhs_dir" || -z "$rhs_dir" ]]; then
|
|
printf 't5-debug-compare needs --lhs-dir and --rhs-dir (or --t5-debug-dir for rhs)\n' >&2
|
|
exit 1
|
|
fi
|
|
require_dir "$lhs_dir"
|
|
require_dir "$rhs_dir"
|
|
run cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux_t5_debug_compare -- "$lhs_dir" "$rhs_dir"
|
|
}
|
|
|
|
cmd_transformer_smoke() {
|
|
ensure_common_inputs
|
|
run cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-transformer-smoke -- "$WORKFLOW" "$MODEL_ROOT" "$WIDTH" "$HEIGHT"
|
|
}
|
|
|
|
cmd_native_generate() {
|
|
ensure_common_inputs
|
|
local -a env_args=()
|
|
if [[ -n "$COND_DIR" ]]; then
|
|
env_args+=("FLUX_COND_DIR=$COND_DIR")
|
|
fi
|
|
if [[ -n "$COMPARE_COND_DIR" ]]; then
|
|
env_args+=("FLUX_COMPARE_COND_DIR=$COMPARE_COND_DIR")
|
|
fi
|
|
if [[ -n "$T5_MODE" ]]; then
|
|
env_args+=("FLUX_T5_MODE=$T5_MODE")
|
|
fi
|
|
if [[ "$T5_CPU_MATH" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_MATH=1")
|
|
fi
|
|
if [[ "$T5_CPU_ATTN" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_CPU_ATTN=1")
|
|
fi
|
|
if [[ "$T5_F32_LINEAR" != "0" ]]; then
|
|
env_args+=("FLUX_T5_FORCE_F32_LINEAR=1")
|
|
fi
|
|
if [[ ${#env_args[@]} -gt 0 ]]; then
|
|
run env "${env_args[@]}" cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-generate -- "$WORKFLOW" "$MODEL_ROOT" "$NATIVE_OUTPUT" "$WIDTH" "$HEIGHT" "$STEPS"
|
|
else
|
|
run cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-generate -- "$WORKFLOW" "$MODEL_ROOT" "$NATIVE_OUTPUT" "$WIDTH" "$HEIGHT" "$STEPS"
|
|
fi
|
|
run file "$NATIVE_OUTPUT"
|
|
}
|
|
|
|
cmd_warm_bench() {
|
|
ensure_common_inputs
|
|
local -a cmd=(
|
|
cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-warm-bench --
|
|
"$WORKFLOW" "$MODEL_ROOT" "$WIDTH" "$HEIGHT" "$STEPS" "$WARMUP_RUNS" "$MEASURED_RUNS"
|
|
)
|
|
local -a env_args=()
|
|
if [[ -n "$T5_MODE" ]]; then
|
|
env_args+=("FLUX_T5_MODE=$T5_MODE")
|
|
fi
|
|
if [[ -n "$WARM_OUTPUT" ]]; then
|
|
cmd+=("$WARM_OUTPUT")
|
|
fi
|
|
if [[ ${#env_args[@]} -gt 0 ]]; then
|
|
run env "${env_args[@]}" "${cmd[@]}"
|
|
else
|
|
run "${cmd[@]}"
|
|
fi
|
|
if [[ -n "$WARM_OUTPUT" ]]; then
|
|
run file "$WARM_OUTPUT"
|
|
fi
|
|
}
|
|
|
|
cmd_ref_dump() {
|
|
ensure_reference_inputs
|
|
ensure_ref_build
|
|
mkdir -p "$DUMP_DIR"
|
|
if [[ -n "$T5_DEBUG_DIR" ]]; then
|
|
mkdir -p "$T5_DEBUG_DIR"
|
|
fi
|
|
rm -f "$DUMP_DIR"/*.bin "$DUMP_DIR"/*.txt
|
|
local prompt
|
|
local seed
|
|
local cfg_scale
|
|
prompt="$(resolve_ref_prompt)"
|
|
seed="$(resolve_ref_seed)"
|
|
cfg_scale="$(resolve_ref_cfg_scale)"
|
|
if [[ -z "$prompt" ]]; then
|
|
printf 'could not resolve reference prompt from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$seed" ]]; then
|
|
printf 'could not resolve reference seed from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$cfg_scale" ]]; then
|
|
printf 'could not resolve reference cfg scale from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
log "reference prompt: $prompt"
|
|
log "reference seed: $seed cfg_scale: $cfg_scale"
|
|
if [[ -n "$T5_DEBUG_DIR" ]]; then
|
|
run env FLUX_DUMP_TEXT_COND_DIR="$DUMP_DIR" FLUX_DUMP_T5_DEBUG_DIR="$T5_DEBUG_DIR" "$SDCLI_BIN" \
|
|
--diffusion-model "$DIFFUSION_MODEL" \
|
|
--vae "$VAE_MODEL" \
|
|
--clip_l "$CLIP_L_MODEL" \
|
|
--t5xxl "$T5XXL_MODEL" \
|
|
-p "$prompt" \
|
|
-s "$seed" \
|
|
--cfg-scale "$cfg_scale" \
|
|
-W "$WIDTH" \
|
|
-H "$HEIGHT" \
|
|
--steps 1 \
|
|
-o "$REF_OUTPUT"
|
|
else
|
|
run env FLUX_DUMP_TEXT_COND_DIR="$DUMP_DIR" "$SDCLI_BIN" \
|
|
--diffusion-model "$DIFFUSION_MODEL" \
|
|
--vae "$VAE_MODEL" \
|
|
--clip_l "$CLIP_L_MODEL" \
|
|
--t5xxl "$T5XXL_MODEL" \
|
|
-p "$prompt" \
|
|
-s "$seed" \
|
|
--cfg-scale "$cfg_scale" \
|
|
-W "$WIDTH" \
|
|
-H "$HEIGHT" \
|
|
--steps 1 \
|
|
-o "$REF_OUTPUT"
|
|
fi
|
|
require_file "$DUMP_DIR/flux_clip_pooled.bin"
|
|
require_file "$DUMP_DIR/flux_t5_hidden.bin"
|
|
require_file "$DUMP_DIR/flux_t5_meta.txt"
|
|
run file "$REF_OUTPUT"
|
|
log "conditioning dump: $DUMP_DIR"
|
|
}
|
|
|
|
cmd_ref_generate() {
|
|
ensure_reference_inputs
|
|
ensure_ref_build
|
|
local prompt
|
|
local seed
|
|
local cfg_scale
|
|
prompt="$(resolve_ref_prompt)"
|
|
seed="$(resolve_ref_seed)"
|
|
cfg_scale="$(resolve_ref_cfg_scale)"
|
|
if [[ -z "$prompt" ]]; then
|
|
printf 'could not resolve reference prompt from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$seed" ]]; then
|
|
printf 'could not resolve reference seed from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$cfg_scale" ]]; then
|
|
printf 'could not resolve reference cfg scale from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
log "reference prompt: $prompt"
|
|
log "reference seed: $seed cfg_scale: $cfg_scale"
|
|
run "$SDCLI_BIN" \
|
|
--diffusion-model "$DIFFUSION_MODEL" \
|
|
--vae "$VAE_MODEL" \
|
|
--clip_l "$CLIP_L_MODEL" \
|
|
--t5xxl "$T5XXL_MODEL" \
|
|
-p "$prompt" \
|
|
-s "$seed" \
|
|
--cfg-scale "$cfg_scale" \
|
|
-W "$WIDTH" \
|
|
-H "$HEIGHT" \
|
|
--steps "$STEPS" \
|
|
-o "$REF_OUTPUT"
|
|
run file "$REF_OUTPUT"
|
|
}
|
|
|
|
cmd_ref_warm_bench() {
|
|
ensure_reference_inputs
|
|
ensure_ref_warm_bench_bin
|
|
local prompt
|
|
local seed
|
|
local cfg_scale
|
|
prompt="$(resolve_ref_prompt)"
|
|
seed="$(resolve_ref_seed)"
|
|
cfg_scale="$(resolve_ref_cfg_scale)"
|
|
if [[ -z "$prompt" ]]; then
|
|
printf 'could not resolve reference prompt from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$seed" ]]; then
|
|
printf 'could not resolve reference seed from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$cfg_scale" ]]; then
|
|
printf 'could not resolve reference cfg scale from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
log "reference prompt: $prompt"
|
|
log "reference seed: $seed cfg_scale: $cfg_scale"
|
|
run "$REF_WARM_BENCH_BIN" \
|
|
--diffusion-model "$DIFFUSION_MODEL" \
|
|
--vae "$VAE_MODEL" \
|
|
--clip_l "$CLIP_L_MODEL" \
|
|
--t5xxl "$T5XXL_MODEL" \
|
|
--prompt "$prompt" \
|
|
--seed "$seed" \
|
|
--cfg-scale "$cfg_scale" \
|
|
--width "$WIDTH" \
|
|
--height "$HEIGHT" \
|
|
--steps "$STEPS" \
|
|
--warmup-runs "$WARMUP_RUNS" \
|
|
--measured-runs "$MEASURED_RUNS"
|
|
}
|
|
|
|
cmd_ref_step_dump() {
|
|
ensure_reference_inputs
|
|
ensure_ref_build
|
|
mkdir -p "$DUMP_DIR" "$STEP_DIR"
|
|
rm -f "$DUMP_DIR"/*.bin "$DUMP_DIR"/*.txt "$STEP_DIR"/*.bin "$STEP_DIR"/*.txt
|
|
local prompt
|
|
local seed
|
|
local cfg_scale
|
|
local dump_step_index
|
|
prompt="$(resolve_ref_prompt)"
|
|
seed="$(resolve_ref_seed)"
|
|
cfg_scale="$(resolve_ref_cfg_scale)"
|
|
dump_step_index="$DUMP_STEP_INDEX"
|
|
if [[ -z "$prompt" ]]; then
|
|
printf 'could not resolve reference prompt from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$seed" ]]; then
|
|
printf 'could not resolve reference seed from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
if [[ -z "$cfg_scale" ]]; then
|
|
printf 'could not resolve reference cfg scale from workflow: %s\n' "$WORKFLOW" >&2
|
|
exit 1
|
|
fi
|
|
log "reference prompt: $prompt"
|
|
log "reference seed: $seed cfg_scale: $cfg_scale"
|
|
run env FLUX_DUMP_TEXT_COND_DIR="$DUMP_DIR" FLUX_DUMP_STEP_DIR="$STEP_DIR" "$SDCLI_BIN" \
|
|
--diffusion-model "$DIFFUSION_MODEL" \
|
|
--vae "$VAE_MODEL" \
|
|
--clip_l "$CLIP_L_MODEL" \
|
|
--t5xxl "$T5XXL_MODEL" \
|
|
-p "$prompt" \
|
|
-s "$seed" \
|
|
--cfg-scale "$cfg_scale" \
|
|
-W "$WIDTH" \
|
|
-H "$HEIGHT" \
|
|
--steps "$dump_step_index" \
|
|
-o "$REF_OUTPUT"
|
|
require_file "$DUMP_DIR/flux_clip_pooled.bin"
|
|
require_file "$DUMP_DIR/flux_t5_hidden.bin"
|
|
require_file "$DUMP_DIR/flux_t5_meta.txt"
|
|
require_file "$STEP_DIR/flux_noised_input.bin"
|
|
require_file "$STEP_DIR/flux_cond_out.bin"
|
|
require_file "$STEP_DIR/flux_step_meta.txt"
|
|
run file "$REF_OUTPUT"
|
|
log "conditioning dump: $DUMP_DIR"
|
|
log "step dump: $STEP_DIR"
|
|
}
|
|
|
|
cmd_handoff_generate() {
|
|
ensure_common_inputs
|
|
local cond_dir="${COND_DIR:-$DUMP_DIR}"
|
|
require_dir "$cond_dir"
|
|
require_file "$cond_dir/flux_clip_pooled.bin"
|
|
require_file "$cond_dir/flux_t5_hidden.bin"
|
|
require_file "$cond_dir/flux_t5_meta.txt"
|
|
run env FLUX_COND_DIR="$cond_dir" cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-generate -- "$WORKFLOW" "$MODEL_ROOT" "$HANDOFF_OUTPUT" "$WIDTH" "$HEIGHT" "$STEPS"
|
|
run file "$HANDOFF_OUTPUT"
|
|
}
|
|
|
|
cmd_transformer_ref_compare() {
|
|
ensure_common_inputs
|
|
local cond_dir="${COND_DIR:-$DUMP_DIR}"
|
|
require_dir "$cond_dir"
|
|
require_dir "$STEP_DIR"
|
|
require_file "$cond_dir/flux_clip_pooled.bin"
|
|
require_file "$cond_dir/flux_t5_hidden.bin"
|
|
require_file "$cond_dir/flux_t5_meta.txt"
|
|
require_file "$STEP_DIR/flux_noised_input.bin"
|
|
require_file "$STEP_DIR/flux_cond_out.bin"
|
|
require_file "$STEP_DIR/flux_step_meta.txt"
|
|
run env FLUX_COND_DIR="$cond_dir" FLUX_REF_STEP_DIR="$STEP_DIR" cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-transformer-smoke -- "$WORKFLOW" "$MODEL_ROOT" "$WIDTH" "$HEIGHT"
|
|
}
|
|
|
|
cmd_transformer_ref_stage_compare() {
|
|
ensure_common_inputs
|
|
local cond_dir="${COND_DIR:-$DUMP_DIR}"
|
|
require_dir "$cond_dir"
|
|
require_dir "$STEP_DIR"
|
|
require_file "$cond_dir/flux_clip_pooled.bin"
|
|
require_file "$cond_dir/flux_t5_hidden.bin"
|
|
require_file "$cond_dir/flux_t5_meta.txt"
|
|
require_file "$STEP_DIR/flux_noised_input.bin"
|
|
require_file "$STEP_DIR/flux_cond_out.bin"
|
|
require_file "$STEP_DIR/flux_step_meta.txt"
|
|
run env FLUX_COND_DIR="$cond_dir" FLUX_REF_STEP_DIR="$STEP_DIR" FLUX_DEBUG_TRANSFORMER_STAGES=1 cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-transformer-smoke -- "$WORKFLOW" "$MODEL_ROOT" "$WIDTH" "$HEIGHT"
|
|
}
|
|
|
|
cmd_transformer_ref_compare_f32() {
|
|
ensure_common_inputs
|
|
local cond_dir="${COND_DIR:-$DUMP_DIR}"
|
|
require_dir "$cond_dir"
|
|
require_dir "$STEP_DIR"
|
|
require_file "$cond_dir/flux_clip_pooled.bin"
|
|
require_file "$cond_dir/flux_t5_hidden.bin"
|
|
require_file "$cond_dir/flux_t5_meta.txt"
|
|
require_file "$STEP_DIR/flux_noised_input.bin"
|
|
require_file "$STEP_DIR/flux_cond_out.bin"
|
|
require_file "$STEP_DIR/flux_step_meta.txt"
|
|
run env FLUX_COND_DIR="$cond_dir" FLUX_REF_STEP_DIR="$STEP_DIR" FLUX_FORCE_F32_WEIGHTS=1 cargo run --release --manifest-path "$CARGO_MANIFEST" --bin flux-transformer-smoke -- "$WORKFLOW" "$MODEL_ROOT" "$WIDTH" "$HEIGHT"
|
|
}
|
|
|
|
cmd_oracle() {
|
|
cmd_text_smoke
|
|
cmd_t5_smoke
|
|
cmd_ref_step_dump
|
|
cmd_transformer_ref_compare
|
|
cmd_native_generate
|
|
cmd_handoff_generate
|
|
log "reference image: $REF_OUTPUT"
|
|
log "native image: $NATIVE_OUTPUT"
|
|
log "handoff image: $HANDOFF_OUTPUT"
|
|
log "conditioning dump: $DUMP_DIR"
|
|
log "step dump: $STEP_DIR"
|
|
}
|
|
|
|
main() {
|
|
local command="${1:-}"
|
|
shift || true
|
|
parse_args "$@"
|
|
case "$command" in
|
|
check) cmd_check ;;
|
|
ref-build) ensure_ref_build ;;
|
|
ref-generate) cmd_ref_generate ;;
|
|
ref-warm-bench) cmd_ref_warm_bench ;;
|
|
text-smoke) cmd_text_smoke ;;
|
|
t5-smoke) cmd_t5_smoke ;;
|
|
t5-stage-compare) cmd_t5_stage_compare ;;
|
|
t5-debug-compare) cmd_t5_debug_compare ;;
|
|
transformer-smoke) cmd_transformer_smoke ;;
|
|
native-generate) cmd_native_generate ;;
|
|
warm-bench) cmd_warm_bench ;;
|
|
ref-dump) cmd_ref_dump ;;
|
|
ref-step-dump) cmd_ref_step_dump ;;
|
|
handoff-generate) cmd_handoff_generate ;;
|
|
transformer-ref-compare) cmd_transformer_ref_compare ;;
|
|
transformer-ref-stage-compare) cmd_transformer_ref_stage_compare ;;
|
|
transformer-ref-compare-f32) cmd_transformer_ref_compare_f32 ;;
|
|
oracle) cmd_oracle ;;
|
|
-h|--help|help|"") usage ;;
|
|
*)
|
|
printf 'unknown command: %s\n\n' "$command" >&2
|
|
usage >&2
|
|
exit 1
|
|
;;
|
|
esac
|
|
}
|
|
|
|
main "$@"
|