flux-fast/experiments_kontext.sh at main · eudiamonia/flux-fast · GitHub

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
88
89
90
91
92
93
94
95
96
#!/bin/bash

CKPT="black-forest-labs/FLUX.1-Kontext-dev"
IMAGE="yarn-art-pikachu.png"
PROMPT="Make Pikachu hold a sign that says 'Black Forest Labs is awesome', yarn art style, detailed, vibrant colors"
CACHE_DIR="/fsx/sayak/.cache"

# bfloat16
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode disabled \
    --disable_fused_projections \
    --disable_channels_last \
    --disable_fa3 \
    --disable_quant \
    --disable_inductor_tuning_flags \
    --output-file bf16.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16.txt 2>&1

# bfloat16 + torch.compile
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode compile \
    --disable_fused_projections \
    --disable_channels_last \
    --disable_fa3 \
    --disable_quant \
    --disable_inductor_tuning_flags \
    --output-file bf16_compile.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16_compile.txt 2>&1

# bfloat16 + torch.compile + qkv projection
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode compile \
    --disable_channels_last \
    --disable_fa3 \
    --disable_quant \
    --disable_inductor_tuning_flags \
    --output-file bf16_compile_qkv.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16_compile_qkv.txt 2>&1

# bfloat16 + torch.compile + qkv projection + channels_last
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode compile \
    --disable_fa3 \
    --disable_quant \
    --disable_inductor_tuning_flags \
    --output-file bf16_compile_qkv_chan.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16_compile_qkv_chan.txt 2>&1

# bfloat16 + torch.compile + qkv projection + channels_last + FA3
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode compile \
    --disable_quant \
    --disable_inductor_tuning_flags \
    --output-file bf16_compile_qkv_chan_fa3.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16_compile_qkv_chan_fa3.txt 2>&1

# bfloat16 + torch.compile + qkv projection + channels_last + FA3 + float8 quant
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode compile \
    --disable_inductor_tuning_flags \
    --output-file bf16_compile_qkv_chan_fa3_quant.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16_compile_qkv_chan_fa3_quant.txt 2>&1

# bfloat16 + torch.compile + qkv projection + channels_last + FA3 + float8 quant + inductor flags
python run_benchmark.py \
    --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --compile_export_mode compile \
    --output-file bf16_compile_qkv_chan_fa3_quant_flags.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > bf16_compile_qkv_chan_fa3_quant_flags.txt 2>&1

# fully optimized (torch.export + AOTI to address cold start)
python run_benchmark.py --ckpt $CKPT --image $IMAGE --prompt "$PROMPT" \
    --output-file fully_optimized.png \
    --num_inference_steps 28 \
    --cache-dir $CACHE_DIR \
    > fully_optimized.txt 2>&1