⚡️ Speed up method CogView4DenoiseInvocation._convert_timesteps_to_sigmas by 24%

[codeflash-ai]

📄 24% (0.24x) speedup for CogView4DenoiseInvocation._convert_timesteps_to_sigmas in invokeai/app/invocations/cogview4_denoise.py

⏱️ Runtime : 1.53 milliseconds → 1.24 milliseconds (best of 224 runs)

📝 Explanation and details

The optimization achieves a 24% speedup by eliminating function call overhead and streamlining the computation flow. The key changes are:

What was optimized:

• Inlined nested functions: The original code defined two nested functions (calculate_timestep_shift and time_shift_linear) inside _convert_timesteps_to_sigmas that were called once per invocation. The optimized version moves this logic directly inline.
• Eliminated function call overhead: From the line profiler, 85.4% of the original execution time was spent in the time_shift_linear call, with additional overhead from calculate_timestep_shift.
• Hardcoded constants: Replaced parameterized values with their actual constants (base_seq_len=256, base_shift=0.25, max_shift=0.75, sigma=1.0) since these were never varied.

Why this leads to speedup:
In Python, function calls have significant overhead due to frame creation, argument passing, and namespace lookups. The line profiler shows that the nested function calls consumed ~87% of the total runtime. By inlining the computation (m = (image_seq_len / 256) ** 0.5; mu = m * 0.75 + 0.25; sigmas = mu / (mu + (1.0 / timesteps - 1.0))), we eliminate this overhead entirely.

Test case performance:
The optimization consistently improves performance across all test scenarios by 16-36%, with particularly strong gains on:

• Multiple timesteps (25-30% faster) - benefits from eliminating repeated function calls
• Edge cases with extreme values (17-28% faster) - simpler execution path handles edge cases more efficiently
• Large tensor operations (19-26% faster) - vectorized operations scale better without function call overhead

The optimization maintains identical numerical results and error handling while significantly reducing computational overhead, making it especially valuable for image processing pipelines where this function may be called frequently during the denoising process.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest # used for our unit tests
import torch
from invokeai.app.invocations.cogview4_denoise import CogView4DenoiseInvocation

unit tests

Helper: Create an instance of the class for test calls

@pytest.fixture(scope="module")
def cogview_invocation():
return CogView4DenoiseInvocation()

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

def test_single_timestep_middle(cogview_invocation):
# Test with a single mid-range timestep
image_seq_len = 256
timesteps = torch.tensor([0.5])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 65.2μs -> 53.6μs (21.7% faster)

def test_multiple_timesteps_typical(cogview_invocation):
# Test with several typical timesteps
image_seq_len = 256
timesteps = torch.tensor([0.1, 0.5, 0.9])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 49.2μs -> 39.1μs (25.7% faster)
# Check that all values are between 0 and 1
for s in sigmas:
pass

def test_typical_nondefault_image_seq_len(cogview_invocation):
# Test with a different image_seq_len
image_seq_len = 1024
timesteps = torch.tensor([0.5])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 45.1μs -> 36.2μs (24.6% faster)
# mu = sqrt(1024/256) * 0.75 + 0.25 = 2 * 0.75 + 0.25 = 1.5 + 0.25 = 1.75
# sigma = 1.0
# t = 0.5
# sig = 1.75 / (1.75 + (1/0.5 - 1)) = 1.75 / (1.75 + 1) = 1.75 / 2.75
expected = 1.75 / 2.75

def test_typical_float_timesteps(cogview_invocation):
# Test with float32 tensor
image_seq_len = 256
timesteps = torch.tensor([0.2, 0.4, 0.6, 0.8], dtype=torch.float32)
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 46.1μs -> 36.1μs (27.4% faster)

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

def test_timestep_at_one(cogview_invocation):
# t=1 should be handled (1/1-1=0), so sigma should be mu/(mu+0)=1.0
image_seq_len = 256
timesteps = torch.tensor([1.0])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 65.7μs -> 56.0μs (17.3% faster)

def test_timestep_very_close_to_zero(cogview_invocation):
# t very close to zero (but not zero)
image_seq_len = 256
t = 1e-8
timesteps = torch.tensor([t])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 46.4μs -> 37.5μs (23.6% faster)

def test_timestep_very_close_to_one(cogview_invocation):
# t very close to one (but not one)
image_seq_len = 256
t = 1.0 - 1e-8
timesteps = torch.tensor([t])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 44.9μs -> 36.0μs (24.6% faster)

def test_negative_timestep_raises(cogview_invocation):
# Negative timesteps should raise or produce nan
image_seq_len = 256
timesteps = torch.tensor([-0.1])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 44.7μs -> 34.8μs (28.3% faster)

def test_image_seq_len_not_multiple_of_32(cogview_invocation):
# image_seq_len not a multiple of 32
image_seq_len = 123
timesteps = torch.tensor([0.5])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 44.8μs -> 35.6μs (26.0% faster)

def test_empty_timesteps(cogview_invocation):
# Empty timesteps tensor should return an empty list
image_seq_len = 256
timesteps = torch.tensor([])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 64.4μs -> 54.7μs (17.7% faster)

def test_large_image_seq_len(cogview_invocation):
# Very large image_seq_len
image_seq_len = 10000
timesteps = torch.tensor([0.5])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 47.5μs -> 37.4μs (27.0% faster)

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

def test_many_timesteps_linear(cogview_invocation):
# Test with a large number of timesteps, evenly spaced
image_seq_len = 256
n = 1000
timesteps = torch.linspace(1e-6, 1.0, n)
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 57.6μs -> 48.3μs (19.1% faster)
# All values should be between 0 and 1
for s in sigmas:
pass

def test_many_timesteps_random(cogview_invocation):
# Test with a large number of random timesteps in (0,1]
image_seq_len = 256
n = 1000
torch.manual_seed(42)
timesteps = torch.rand(n) * 0.999999 + 1e-6 # avoid zero
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 47.5μs -> 37.3μs (27.5% faster)
# All values should be between 0 and 1
for s in sigmas:
pass

def test_large_image_seq_len_and_many_timesteps(cogview_invocation):
# Large image_seq_len and many timesteps
image_seq_len = 9999
n = 500
timesteps = torch.linspace(1e-6, 1.0, n)
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 49.3μs -> 42.3μs (16.4% faster)
# All values should be between 0 and 1
for s in sigmas:
pass

def test_extremely_large_image_seq_len(cogview_invocation):
# Test with image_seq_len at upper bound for 32-bit float precision
image_seq_len = 2**20 # 1048576
timesteps = torch.tensor([0.5])
codeflash_output = cogview_invocation._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 44.5μs -> 36.0μs (23.5% faster)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import pytest # used for our unit tests
import torch
from invokeai.app.invocations.cogview4_denoise import CogView4DenoiseInvocation

unit tests

@pytest.fixture
def cogview4():
# Fixture to provide a fresh instance for each test
return CogView4DenoiseInvocation()

--- BASIC TEST CASES ---

def test_single_timestep_middle(cogview4):
# Test with a single timestep in the middle of the range (0.5)
image_seq_len = 256
timesteps = torch.tensor([0.5])
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 49.1μs -> 38.7μs (27.0% faster)
# Calculate expected value manually
mu = ((256/256)**0.5) * 0.75 + 0.25 # = 1.0
expected = mu / (mu + (1/0.5 - 1)**1.0) # (1.0 / (1.0 + 1.0)) = 0.5

def test_multiple_timesteps_basic(cogview4):
# Test with a few timesteps from 0.1 to 0.9
image_seq_len = 256
timesteps = torch.tensor([0.1, 0.5, 0.9])
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 49.2μs -> 37.6μs (30.9% faster)
mu = 1.0
expected = [
mu / (mu + (1/0.1 - 1)**1.0), # 1/(1+9) = 0.1
mu / (mu + (1/0.5 - 1)**1.0), # 1/(1+1) = 0.5
mu / (mu + (1/0.9 - 1)**1.0), # 1/(1+0.111...) = ~0.9
]
for s, e in zip(sigmas, expected):
pass

def test_image_seq_len_different(cogview4):
# Test with a different image_seq_len
image_seq_len = 1024
timesteps = torch.tensor([0.5])
mu = ((1024/256)**0.5) * 0.75 + 0.25 # sqrt(4)=2, 2*0.75=1.5+0.25=1.75
expected = mu / (mu + (1/0.5 - 1)**1.0) # 1.75/(1.75+1)=1.75/2.75
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 45.9μs -> 37.8μs (21.4% faster)

def test_timesteps_tensor_float32_and_float64(cogview4):
# Test with float32 and float64 tensors
image_seq_len = 256
for dtype in [torch.float32, torch.float64]:
timesteps = torch.tensor([0.5, 0.9], dtype=dtype)
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 70.5μs -> 51.8μs (36.1% faster)
mu = 1.0
expected = [
mu / (mu + (1/0.5 - 1)**1.0),
mu / (mu + (1/0.9 - 1)**1.0),
]
for s, e in zip(sigmas, expected):
pass

--- EDGE TEST CASES ---

def test_timesteps_zero_and_one(cogview4):
# Test with timesteps at the edges: 0.0 and 1.0
image_seq_len = 256
timesteps = torch.tensor([0.0, 1.0])
mu = 1.0
# For t=0, (1/0-1) is inf, so denominator is inf, so sigma=0
# For t=1, (1/1-1)=0, so denominator is mu, so sigma=1
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 45.2μs -> 37.1μs (21.9% faster)

def test_timesteps_near_zero_and_one(cogview4):
# Test with timesteps very close to 0 and 1
image_seq_len = 256
eps = 1e-8
timesteps = torch.tensor([eps, 1-eps])
mu = 1.0
# For t near 0, denominator is huge, sigma ~0
# For t near 1, denominator is ~mu, sigma ~1
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 44.7μs -> 37.9μs (17.9% faster)

def test_large_and_small_image_seq_len(cogview4):
# Test with very small and very large image_seq_len
for image_seq_len in [1, 10000]:
timesteps = torch.tensor([0.5])
mu = ((image_seq_len/256)**0.5) * 0.75 + 0.25
expected = mu / (mu + (1/0.5 - 1)**1.0)
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 68.1μs -> 53.8μs (26.4% faster)

def test_negative_and_zero_image_seq_len(cogview4):
# image_seq_len=0 or negative should result in mu=0.25 or nan/complex
# image_seq_len=0: m=0, mu=0.25
image_seq_len = 0
timesteps = torch.tensor([0.5])
mu = 0.25
expected = mu / (mu + (1/0.5 - 1)**1.0) # 0.25/(0.25+1)=0.25/1.25
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 46.4μs -> 36.8μs (26.1% faster)

# image_seq_len negative: sqrt of negative is nan, so expect nan
image_seq_len = -1
timesteps = torch.tensor([0.5])
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 28.6μs -> 24.9μs (14.8% faster)

def test_empty_timesteps(cogview4):
# Test with an empty timesteps tensor
image_seq_len = 256
timesteps = torch.tensor([])
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 43.1μs -> 35.1μs (22.6% faster)

def test_non_1d_timesteps_tensor(cogview4):
# Test with a 2D tensor
image_seq_len = 256
timesteps = torch.tensor([[0.1, 0.2], [0.3, 0.4]])
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 48.6μs -> 38.2μs (27.2% faster)
# Should flatten to a list of correct length and values
mu = 1.0
expected = [
mu / (mu + (1/0.1 - 1)**1.0),
mu / (mu + (1/0.2 - 1)**1.0),
mu / (mu + (1/0.3 - 1)**1.0),
mu / (mu + (1/0.4 - 1)**1.0),
]
for s, e in zip(sigmas, expected):
pass

def test_timesteps_with_nan_inf(cogview4):
# Test with timesteps containing nan and inf values
image_seq_len = 256
timesteps = torch.tensor([float('nan'), float('inf'), -float('inf')])
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 47.9μs -> 38.5μs (24.4% faster)

--- LARGE SCALE TEST CASES ---

def test_large_timesteps_tensor(cogview4):
# Test with a large number of timesteps (1000 elements)
image_seq_len = 256
timesteps = torch.linspace(0, 1, steps=1000)
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 64.0μs -> 50.4μs (26.9% faster)
# Check monotonicity: sigmas should increase as timesteps increase
for i in range(1, len(sigmas)):
pass

def test_large_image_seq_len_and_timesteps(cogview4):
# Test with large image_seq_len and large timesteps tensor
image_seq_len = 8192
timesteps = torch.linspace(0.01, 0.99, steps=500)
mu = ((8192/256)**0.5) * 0.75 + 0.25
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 55.0μs -> 45.6μs (20.6% faster)
# Check some representative values
for idx in [0, 249, 499]:
t = timesteps[idx].item()
expected = mu / (mu + (1/t - 1)**1.0)

def test_performance_large_tensor(cogview4):
# Test that the function runs efficiently for a large, but not huge, tensor
image_seq_len = 256
timesteps = torch.rand(1000)
# This shouldn't take significant time or memory
codeflash_output = cogview4._convert_timesteps_to_sigmas(image_seq_len, timesteps); sigmas = codeflash_output # 65.4μs -> 52.0μs (25.6% faster)
# All values should be between 0 and 1
for s in sigmas:
pass

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-CogView4DenoiseInvocation._convert_timesteps_to_sigmas-mhvmf9gi and push.