⚡️ Speed up function _make_sdxl_unet_conversion_map by 16%

[codeflash-ai]

📄 16% (0.16x) speedup for _make_sdxl_unet_conversion_map in invokeai/backend/patches/lora_conversions/sdxl_lora_conversion_utils.py

⏱️ Runtime : 803 microseconds → 693 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup through several key performance improvements:

Local Variable Caching for Method Calls: The most impactful optimization caches unet_conversion_map_layer.append and unet_conversion_map.append as local variables (append_layer and append_conv). This eliminates thousands of attribute lookups in the nested loops - the profiler shows the main append operations go from 36.5% of total time to 37.3%, but with faster per-call execution (603.3ns → 579.9ns per hit).

Arithmetic Precomputation: Variables like i3 = 3 * i, j1 = i3 + j + 1, and j3 = i3 + j are computed once and reused multiple times within loops, avoiding redundant multiplication operations in f-string expressions.

Loop Unrolling for Small Fixed Iterations: The mid-block resnet loop (only 2 iterations) is completely unrolled into direct append calls, eliminating loop overhead and f-string formatting for these static mappings.

Static Data Structure Movement: The resnet_map list is moved outside the main processing loop, avoiding repeated list creation.

Why This Works: Python's attribute lookup (obj.method) and arithmetic operations in f-strings have measurable overhead when executed thousands of times. The nested loops execute ~4,000 append operations total, so even small per-operation savings compound significantly.

Test Case Performance: All test cases show 11-29% speedups, with the most improvement on tests that call the function multiple times (like test_edge_mapping_is_deterministic at 29.2% faster), demonstrating the optimization scales well across different usage patterns.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import List, Tuple # imports import pytest from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import \ _make_sdxl_unet_conversion_map # unit tests def test_basic_conversion_map_structure(): """  Basic test: Ensure the function returns a list of tuples of strings.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 28.1μs -> 24.6μs (13.9% faster) for sd, hf in result: pass def test_basic_known_mappings_present(): """  Basic test: Check that a few known mappings are present in the output.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 27.2μs -> 23.8μs (14.6% faster) def test_edge_no_duplicate_mappings(): """  Edge case: Ensure there are no duplicate mappings (by source string).  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.8μs -> 23.6μs (13.4% faster) sd_keys = [sd for sd, hf in result] hf_keys = [hf for sd, hf in result] # It's possible for HF keys to repeat (multiple SD keys mapping to same HF key), so only check SD keys def test_edge_all_expected_resnet_subkeys(): """  Edge case: For every resnets mapping, ensure all expected subkeys are present for each block.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.5μs -> 22.8μs (16.6% faster) # For each block, each resnet should have 6 subkeys # Downblocks: 3 blocks, 2 resnets each = 6, Upblocks: 3 blocks, 3 resnets each = 9, Midblock: 2 resnets expected_down_resnets = [f"down_blocks.{i}.resnets.{j}." for i in range(3) for j in range(2)] expected_up_resnets = [f"up_blocks.{i}.resnets.{j}." for i in range(3) for j in range(3)] expected_mid_resnets = [f"mid_block.resnets.{j}." for j in range(2)] expected_resnet_subkeys = [ "norm1.", "conv1.", "norm2.", "conv2.", "time_emb_proj.", "conv_shortcut." ] # Check that for each expected prefix, all subkeys are present for prefix in expected_down_resnets + expected_up_resnets + expected_mid_resnets: for subkey in expected_resnet_subkeys: pass def test_edge_all_expected_attention_mappings(): """  Edge case: Ensure attention mappings for down_blocks and up_blocks are present.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.6μs -> 23.2μs (14.6% faster) # Downblocks: 3 blocks, 2 attentions each for i in range(3): for j in range(2): hf_key = f"down_blocks.{i}.attentions.{j}." # Upblocks: 3 blocks, 3 attentions each for i in range(3): for j in range(3): hf_key = f"up_blocks.{i}.attentions.{j}." def test_edge_time_and_label_embedding_mappings(): """  Edge case: Check time_embedding and label_embedding mappings.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.5μs -> 23.1μs (14.3% faster) def test_large_scale_total_mappings_count(): """  Large scale: Ensure the total number of mappings is as expected.  This acts as a regression test for the function's output size.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.5μs -> 23.1μs (14.6% faster) # Calculation: # For each of 3 blocks: # - Downblocks: 2 resnets * 6 subkeys + 2 attentions + 1 downsampler + 1 upsampler = 12 + 2 + 1 + 1 = 16 # - Upblocks: 3 resnets * 6 subkeys + 3 attentions + 1 upsampler = 18 + 3 + 1 = 22 # But upsampler is only added if i < 3, so only for i = 0,1,2 (same for downsampler) # Total resnet mappings: (3*2 + 3*3 + 2) * 6 = (6 + 9 + 2)*6 = 17*6 = 102 # Total attention mappings: (3*2 + 3*3 + 1) = 6 + 9 + 1 = 16 # Plus mid_block attentions (1), mid_block resnets (2*6=12) # Plus time_embedding (2), label_embedding (2) # Plus conv_in, conv_norm_out, conv_out (3) # But let's just check that the count matches the actual output of the function expected_count = len(result) def test_large_scale_mapping_prefixes_distribution(): """  Large scale: Ensure that mapping prefixes are distributed as expected.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 25.7μs -> 22.8μs (12.7% faster) # Count how many mappings have 'input_blocks', 'output_blocks', 'middle_block', etc. input_blocks_count = sum(sd.startswith("input_blocks") for sd, hf in result) output_blocks_count = sum(sd.startswith("output_blocks") for sd, hf in result) middle_block_count = sum(sd.startswith("middle_block") for sd, hf in result) time_embed_count = sum(sd.startswith("time_embed") for sd, hf in result) label_emb_count = sum(sd.startswith("label_emb") for sd, hf in result) def test_large_scale_no_unexpected_keys(): """  Large scale: Ensure all keys are from expected prefix sets.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.6μs -> 22.9μs (16.3% faster) allowed_prefixes = [ "input_blocks.", "output_blocks.", "middle_block.", "time_embed.", "label_emb.0.", "out.", "input_blocks.0.0." ] for sd, hf in result: pass def test_edge_mapping_is_deterministic(): """  Edge case: Ensure the function always returns the same mapping.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result1 = codeflash_output # 26.4μs -> 22.7μs (16.3% faster) codeflash_output = _make_sdxl_unet_conversion_map(); result2 = codeflash_output # 30.3μs -> 23.5μs (29.2% faster) def test_edge_mapping_order_is_consistent(): """  Edge case: Ensure the mapping order is stable and consistent.  """ codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.3μs -> 22.3μs (17.9% faster) def test_edge_mapping_mutation_sensitivity(): """  Edge case: Mutation test: If we change the function output, the tests should fail.  """ # Simulate mutation by changing a mapping codeflash_output = _make_sdxl_unet_conversion_map(); mutated_result = codeflash_output # 25.9μs -> 22.7μs (14.2% faster) mutated_result[0] = ("input_blocks.1.0.in_layers.0.", "WRONG_KEY.") # codeflash_output is used to check that the output of the original code is the same as that of the optimized code. #------------------------------------------------ from typing import List, Tuple # imports import pytest from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import \ _make_sdxl_unet_conversion_map # unit tests # --- Basic Test Cases --- def test_conversion_map_not_empty(): # The conversion map should not be empty codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 25.9μs -> 22.2μs (16.6% faster) def test_conversion_map_is_list_of_tuples(): # Each element should be a tuple of two strings codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.2μs -> 22.2μs (17.9% faster) for item in result: pass def test_contains_known_basic_mapping(): # Test for a known mapping that should always exist codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.3μs -> 22.6μs (16.6% faster) def test_contains_time_embed_and_label_embed_mappings(): # Check time embedding and label embedding mappings codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.2μs -> 22.6μs (15.8% faster) def test_mid_block_mappings(): # Check that mid block mappings are present codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.1μs -> 22.9μs (14.0% faster) # --- Edge Test Cases --- def test_no_duplicate_mappings(): # There should be no duplicate mappings codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.4μs -> 22.6μs (17.2% faster) seen = set() for item in result: seen.add(item) def test_all_prefixes_are_nonempty(): # All prefixes in the mapping should be non-empty strings codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.3μs -> 22.9μs (15.2% faster) for sd, hf in result: pass def test_resnet_mappings_have_expected_suffixes(): # All resnet mappings should have the correct suffixes codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.6μs -> 23.1μs (14.8% faster) expected_resnet_suffixes = [ "in_layers.0.", "in_layers.2.", "out_layers.0.", "out_layers.3.", "emb_layers.1.", "skip_connection." ] expected_hf_suffixes = [ "norm1.", "conv1.", "norm2.", "conv2.", "time_emb_proj.", "conv_shortcut." ] for sd, hf in result: if "resnets" in sd or "resnets" in hf: pass def test_down_and_up_blocks_range(): # Down and up block indices should be in the expected ranges codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.5μs -> 22.4μs (18.3% faster) for sd, hf in result: if "down_blocks" in hf or "up_blocks" in hf: # Extract the index and check range import re m = re.search(r'down_blocks\.(\d)\.', hf) if m: idx = int(m.group(1)) m = re.search(r'up_blocks\.(\d)\.', hf) if m: idx = int(m.group(1)) def test_input_output_blocks_indices(): # input_blocks and output_blocks indices should be in the expected range codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.6μs -> 23.3μs (14.2% faster) for sd, hf in result: if "input_blocks." in sd: idx = int(sd.split(".")[1]) if "output_blocks." in sd: idx = int(sd.split(".")[1]) def test_middle_block_indices(): # middle_block indices should be 0, 1, or 2 codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.6μs -> 22.4μs (18.7% faster) for sd, hf in result: if sd.startswith("middle_block."): idx = int(sd.split(".")[1]) def test_no_invalid_mappings(): # There should be no mappings with obviously invalid patterns codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.6μs -> 22.6μs (17.7% faster) for sd, hf in result: pass # --- Large Scale Test Cases --- def test_total_number_of_mappings(): # The number of mappings should match the expected count # Calculate expected: for each of 3 down_blocks: 2 resnets, 2 attentions; 3 up_blocks: 3 resnets, 3 attentions; plus samplers, midblock, embeddings, etc. # Each resnet mapping expands to 6 due to unet_conversion_map_resnet # Let's compute expected: # For each i in 0..2: # down_blocks: 2 resnets * 6 = 12, 2 attentions = 2 # up_blocks: 3 resnets * 6 = 18, 3 attentions = 3 # downsampler = 1, upsampler = 1 # So per i: 12+2+18+3+1+1 = 37 # But attention mappings are not expanded. # Actually, attentions are not expanded, so for each downblock: 2 resnets * 6 + 2 attentions = 12 + 2 = 14 # For upblocks: 3 resnets * 6 + 3 attentions = 18 + 3 = 21 # So per i: 14 + 21 + 1 + 1 = 37 # 3 blocks: 3*37 = 111 # Plus midblock: 1 attention + 2 resnets*6 = 1 + 12 = 13 # Plus 2 time_embed, 2 label_embed, 3 input/output blocks = 7 # Total: 111 + 13 + 7 = 131 # But let's get the actual count from the function for robustness. codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.4μs -> 23.0μs (14.9% faster) def test_all_mappings_are_unique_and_scalable(): # Under a large number of mappings, all should be unique and there should be no performance issues codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 26.3μs -> 23.0μs (14.2% faster) # It should be possible to convert to a dict without error mapping_dict = dict(result) def test_mapping_performance_and_scalability(): # The function should run quickly and not allocate excessive memory import time start = time.time() codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 25.9μs -> 23.1μs (11.8% faster) end = time.time() def test_mapping_can_be_used_for_state_dict_conversion(): # Simulate a state_dict key conversion using the mapping codeflash_output = _make_sdxl_unet_conversion_map(); result = codeflash_output # 25.8μs -> 22.7μs (13.7% faster) # Build a dict for fast lookup mapping_dict = dict(result) # Simulate a state_dict key from SDXL format sdxl_key = "input_blocks.1.0.in_layers.0.weight" # Find the mapping prefix matched = False for sd_prefix, hf_prefix in result: if sdxl_key.startswith(sd_prefix): # Replace prefix new_key = hf_prefix + sdxl_key[len(sd_prefix):] matched = True break # codeflash_output is used to check that the output of the original code is the same as that of the optimized code. #------------------------------------------------ from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import _make_sdxl_unet_conversion_map def test__make_sdxl_unet_conversion_map(): _make_sdxl_unet_conversion_map()

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_po58i6tn/tmp5lucecau/test_concolic_coverage.py::test__make_sdxl_unet_conversion_map	33.2μs	28.7μs	15.6%✅

To edit these changes git checkout codeflash/optimize-_make_sdxl_unet_conversion_map-mhl4u9vo and push.