⚡️ Speed up method EmbedMaxDct.encode by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for EmbedMaxDct.encode in invokeai/backend/image_util/imwatermark/vendor.py

⏱️ Runtime : 105 milliseconds → 99.2 milliseconds (best of 58 runs)

📝 Explanation and details

The optimized code achieves a 6% speedup through several key performance improvements in the image watermarking pipeline:

What optimizations were applied:

1. Precomputed slice boundaries - Moved rows4 = row // 4 * 4 and cols4 = col // 4 * 4 calculations outside the channel loop to avoid redundant computation
2. Local variable caching - Cached frequently accessed attributes (self._block, self._wmLen, self._watermarks) as local variables to reduce attribute lookup overhead
3. Optimized watermark indexing - Replaced the separate num counter with direct calculation (i * num_blocks_col + j) % wmLen, eliminating an increment operation per block
4. Eliminated redundant assignment - Removed the diffusedBlock = self.diffuse_dct_matrix(...) assignment since the method modifies blocks in-place, avoiding unnecessary variable creation and assignment overhead

Why these optimizations improve performance:

• Reduced Python overhead: Local variable access is faster than attribute lookup, especially in tight loops with ~21K iterations
• Better loop efficiency: Direct index calculation eliminates the need to maintain and increment a counter variable
• Memory allocation reduction: Removing the intermediate diffusedBlock variable reduces temporary object creation in the hot path

Key performance impact:
The line profiler shows the most significant improvement in encode_frame (from 177ms to 146ms), which processes the majority of blocks. The diffuse_dct_matrix call remains the bottleneck at ~80% of runtime, but the loop overhead optimizations provide measurable gains.

Test case performance:
The optimizations show consistent 2-9% improvements across various scenarios, with larger images (256x256, 512x512) benefiting most due to the multiplicative effect of loop optimizations across thousands of blocks. Small images see modest gains due to lower absolute loop counts.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 82 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import cv2
import numpy as np
# imports
import pytest
from invokeai.backend.image_util.imwatermark.vendor import EmbedMaxDct

# --------------------------
# Unit tests for EmbedMaxDct.encode
# --------------------------

# --------- Basic Test Cases ----------

def test_encode_identity_on_zero_watermark():
    """
    Test that encoding with all-zero watermark and zero scale does not change the image.
    """
    img = np.zeros((8, 8, 3), dtype=np.uint8)
    watermarks = [0] * 8
    encoder = EmbedMaxDct(watermarks=watermarks, wmLen=8, scales=[0, 0, 0], block=4)
    codeflash_output = encoder.encode(img.copy()); result = codeflash_output # 12.9μs -> 12.9μs (0.070% faster)

def test_encode_identity_on_no_channels_scaled():
    """
    Test that encoding with all-zero scales leaves the image unchanged.
    """
    img = np.random.randint(0, 256, (16, 16, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[0, 0, 0], block=4)
    codeflash_output = encoder.encode(img.copy()); result = codeflash_output # 9.84μs -> 10.4μs (5.24% slower)

def test_encode_changes_image_with_nonzero_scale():
    """
    Test that encoding with nonzero scale and watermark bits changes the image.
    """
    img = np.ones((8, 8, 3), dtype=np.uint8) * 128
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[36, 0, 0], block=4)
    codeflash_output = encoder.encode(img.copy()); result = codeflash_output # 150μs -> 148μs (1.60% faster)

def test_encode_preserves_shape_and_dtype():
    """
    Test that the output image has the same shape and dtype as the input.
    """
    img = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[36, 36, 0], block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 346μs -> 329μs (4.99% faster)

def test_encode_different_watermarks_give_different_results():
    """
    Test that different watermark bits produce different encoded images.
    """
    img = np.ones((16, 16, 3), dtype=np.uint8) * 128
    encoder1 = EmbedMaxDct(watermarks=[0]*8, wmLen=8, scales=[36, 36, 0], block=4)
    encoder2 = EmbedMaxDct(watermarks=[1]*8, wmLen=8, scales=[36, 36, 0], block=4)
    codeflash_output = encoder1.encode(img.copy()); result1 = codeflash_output # 229μs -> 219μs (4.58% faster)
    codeflash_output = encoder2.encode(img.copy()); result2 = codeflash_output # 176μs -> 171μs (3.06% faster)

# --------- Edge Test Cases ----------

def test_encode_minimum_size_image_multiple_of_block():
    """
    Test with the minimum image size that is a multiple of block size (4x4).
    """
    img = np.random.randint(0, 256, (4, 4, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[1], wmLen=1, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 163μs -> 160μs (1.67% faster)

def test_encode_minimum_size_image_not_multiple_of_block():
    """
    Test with the minimum image size that is NOT a multiple of block size (5x5).
    The function should only process the largest multiple-of-block region.
    """
    img = np.random.randint(0, 256, (5, 5, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[1], wmLen=1, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 160μs -> 157μs (2.09% faster)

def test_encode_empty_watermark_bits():
    """
    Test with empty watermark list (should raise IndexError).
    """
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[], wmLen=0, scales=[36, 36, 36], block=4)
    with pytest.raises(IndexError):
        encoder.encode(img)

def test_encode_one_channel_scaled_only():
    """
    Test that only the first channel is affected if only the first scale is nonzero.
    """
    img = np.ones((8, 8, 3), dtype=np.uint8) * 128
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[36, 0, 0], block=4)
    codeflash_output = encoder.encode(img.copy()); result = codeflash_output # 156μs -> 154μs (1.25% faster)

def test_encode_all_channels_scaled():
    """
    Test that all channels are affected if all scales are nonzero.
    """
    img = np.ones((8, 8, 3), dtype=np.uint8) * 128
    encoder = EmbedMaxDct(watermarks=[1,1,1,1,1,1,1,1], wmLen=8, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img.copy()); result = codeflash_output # 205μs -> 201μs (2.34% faster)
    # All channels should differ from the original
    for c in range(3):
        pass

def test_encode_with_block_size_larger_than_image():
    """
    If block size is larger than image, encode_frame should not process any blocks.
    """
    img = np.random.randint(0, 256, (3, 3, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[1], wmLen=1, scales=[36, 36, 36], block=4)
    # Should not raise, and should return an image of same shape
    codeflash_output = encoder.encode(img); result = codeflash_output

def test_encode_with_negative_scale():
    """
    Negative scales should be treated as zero (channel skipped).
    """
    img = np.ones((8, 8, 3), dtype=np.uint8) * 128
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[-1, 0, 0], block=4)
    codeflash_output = encoder.encode(img.copy()); result = codeflash_output # 13.6μs -> 13.9μs (2.41% slower)

def test_encode_with_nonstandard_dtype():
    """
    Test with float32 image. Output should be uint8 since cv2.cvtColor returns uint8.
    """
    img = np.random.rand(8, 8, 3).astype(np.float32) * 255
    img = img.astype(np.uint8)
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 231μs -> 228μs (1.26% faster)

# --------- Large Scale Test Cases ----------

def test_encode_large_image_performance():
    """
    Test with a large image (512x512x3). Should complete in reasonable time and preserve shape/dtype.
    """
    img = np.random.randint(0, 256, (512, 512, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[0,1]*32, wmLen=64, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 38.5ms -> 36.3ms (6.18% faster)

def test_encode_large_watermark_list():
    """
    Test with a large watermark list and large image. Ensures watermark cycling works.
    """
    img = np.random.randint(0, 256, (256, 256, 3), dtype=np.uint8)
    watermark_bits = [0,1]*500  # 1000 bits
    encoder = EmbedMaxDct(watermarks=watermark_bits, wmLen=1000, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 8.96ms -> 8.20ms (9.33% faster)

def test_encode_many_channels_random_scales():
    """
    Test with random scales for each channel and a medium image.
    """
    img = np.random.randint(0, 256, (128, 128, 3), dtype=np.uint8)
    scales = [36, 0, 18]
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=scales, block=4)
    codeflash_output = encoder.encode(img); result = codeflash_output # 1.23ms -> 1.12ms (9.11% faster)

def test_encode_consistency_multiple_calls():
    """
    Test that calling encode multiple times on the same image and encoder gives the same result.
    """
    img = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
    encoder = EmbedMaxDct(watermarks=[1,0,1,0,1,0,1,0], wmLen=8, scales=[36, 36, 36], block=4)
    codeflash_output = encoder.encode(img.copy()); result1 = codeflash_output # 717μs -> 671μs (6.93% faster)
    codeflash_output = encoder.encode(img.copy()); result2 = codeflash_output # 659μs -> 601μs (9.52% faster)

def test_encode_different_block_sizes():
    """
    Test with different block sizes (2, 4, 8) on the same image.
    """
    img = np.random.randint(0, 256, (16, 16, 3), dtype=np.uint8)
    for block_size in [2, 4, 8]:
        encoder = EmbedMaxDct(watermarks=[1,0,1,0], wmLen=4, scales=[36, 36, 36], block=block_size)
        codeflash_output = encoder.encode(img); result = codeflash_output # 635μs -> 609μs (4.26% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import cv2
import numpy as np
# imports
import pytest  # used for our unit tests
from invokeai.backend.image_util.imwatermark.vendor import EmbedMaxDct

# unit tests

# ----------- BASIC TEST CASES -----------

def test_encode_identity_on_zero_scales():
    # When scales are zero, encode should not modify the image
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[0, 1, 0, 1], wmLen=4, scales=[0, 0, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 8.80μs -> 9.33μs (5.65% slower)

def test_encode_modifies_image_with_nonzero_scale():
    # When scale is nonzero, encode should modify the image
    img = np.ones((8, 8, 3), dtype=np.uint8) * 128
    embedder = EmbedMaxDct(watermarks=[1, 0, 1, 0], wmLen=4, scales=[36, 0, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 126μs -> 123μs (2.79% faster)

def test_encode_returns_uint8_image():
    # Output should be uint8 type
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0, 1, 0], wmLen=4, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 197μs -> 192μs (2.73% faster)

def test_encode_preserves_shape():
    # Output shape should match input shape
    img = np.random.randint(0, 256, (16, 16, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0], wmLen=2, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 231μs -> 224μs (3.34% faster)

def test_encode_with_different_block_sizes():
    # Test with block size 2
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1], wmLen=1, scales=[36, 36, 0], block=2)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 217μs -> 212μs (2.72% faster)

# ----------- EDGE TEST CASES -----------

def test_encode_minimum_size_block():
    # Minimum size for block=4 is 4x4
    img = np.random.randint(0, 256, (4, 4, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[0], wmLen=1, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 163μs -> 158μs (2.66% faster)

def test_encode_non_divisible_shape():
    # Shape not divisible by block size, should process up to divisible part
    img = np.random.randint(0, 256, (10, 10, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0], wmLen=2, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 191μs -> 185μs (3.36% faster)

def test_encode_all_zeros_image():
    # All zeros input should not cause errors
    img = np.zeros((8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 1, 1, 1], wmLen=4, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 190μs -> 185μs (2.65% faster)

def test_encode_all_ones_image():
    # All ones input should not cause errors
    img = np.ones((8, 8, 3), dtype=np.uint8) * 255
    embedder = EmbedMaxDct(watermarks=[0, 0, 0, 0], wmLen=4, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 183μs -> 183μs (0.251% slower)

def test_encode_single_watermark_bit():
    # Watermark length 1 should repeat over all blocks
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1], wmLen=1, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 183μs -> 188μs (2.80% slower)

def test_encode_with_negative_scale():
    # Negative scale disables encoding for that channel
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0], wmLen=2, scales=[-1, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 110μs -> 112μs (1.87% slower)

def test_encode_with_empty_watermark():
    # Empty watermark should raise IndexError in encode_frame
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[], wmLen=0, scales=[36, 36, 0], block=4)
    with pytest.raises(ZeroDivisionError):
        # modulo by zero in encode_frame
        embedder.encode(img) # 58.3μs -> 58.9μs (0.963% slower)

def test_encode_with_large_watermark():
    # Watermark longer than number of blocks should not error
    img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8)
    watermark = [1] * 100
    embedder = EmbedMaxDct(watermarks=watermark, wmLen=100, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 202μs -> 197μs (2.95% faster)

# ----------- LARGE SCALE TEST CASES -----------

def test_encode_large_image():
    # Large image should be processed without error
    img = np.random.randint(0, 256, (256, 256, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0, 1, 0], wmLen=4, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 8.98ms -> 8.22ms (9.20% faster)

def test_encode_large_watermark_and_image():
    # Large watermark and image
    watermark = [1, 0] * 500  # 1000 bits
    img = np.random.randint(0, 256, (128, 128, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=watermark, wmLen=1000, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 2.34ms -> 2.14ms (9.21% faster)

def test_encode_performance_large_image():
    # Performance: should complete within reasonable time for 512x512 image
    import time
    img = np.random.randint(0, 256, (512, 512, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0, 1, 0], wmLen=4, scales=[36, 36, 0], block=4)
    start = time.time()
    codeflash_output = embedder.encode(img); encoded = codeflash_output # 38.2ms -> 36.4ms (4.96% faster)
    duration = time.time() - start

def test_encode_consistency():
    # Encoding same image with same watermark should yield same result
    img = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8)
    embedder = EmbedMaxDct(watermarks=[1, 0, 1, 0], wmLen=4, scales=[36, 36, 0], block=4)
    codeflash_output = embedder.encode(img); encoded1 = codeflash_output # 355μs -> 340μs (4.36% faster)
    codeflash_output = embedder.encode(img); encoded2 = codeflash_output # 280μs -> 261μs (7.06% faster)

def test_encode_different_watermarks_produce_different_results():
    # Different watermarks should yield different encoded images
    img = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8)
    embedder1 = EmbedMaxDct(watermarks=[1, 0, 1, 0], wmLen=4, scales=[36, 36, 0], block=4)
    embedder2 = EmbedMaxDct(watermarks=[0, 1, 0, 1], wmLen=4, scales=[36, 36, 0], block=4)
    codeflash_output = embedder1.encode(img); encoded1 = codeflash_output # 325μs -> 311μs (4.41% faster)
    codeflash_output = embedder2.encode(img); encoded2 = codeflash_output # 276μs -> 258μs (7.01% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-EmbedMaxDct.encode-mhwxdium and push.