⚡️ Speed up function rescale_noise_cfg by 44%

[codeflash-ai]

📄 44% (0.44x) speedup for rescale_noise_cfg in src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py

⏱️ Runtime : 6.65 milliseconds → 4.61 milliseconds (best of 320 runs)

📝 Explanation and details

Here is an optimized version of the provided program.
Key performance recommendations based on the line profiler.

1. Avoid repeated construction of axes:
   list(range(1, x.ndim)) is a minor but avoidable overhead, especially when applied twice. Store it once.

2. Minimize Python-side operations:
   Use tuple for axes directly and avoid redundant list constructions.

3. Move computation of axes outside to avoid recompute on every call:
   Since axes are always tuple(range(1, x.ndim)), define a tiny helper for this, but to keep the single-function signature, inline it in each call.

4. Use in-place math when possible:
   While PyTorch Tensors will not always benefit from in-place ops due to autograd, the operations here do not require gradients, so we can consider in-place modification, but for safety, stick to the out-of-place as it’s already vectorized.

5. Avoid duplicate computation when guidance_rescale==0.0:
   If guidance_rescale is 0, just return original input.
   Similarly, if it's 1.0, shortcut to fully rescaled output.

6. Early return to minimize computation on defaults.

Here’s the optimized code.

Summary of optimizations:

• Avoids list allocation for axis in every call.
• Fast-path for default or trivial guidance_rescale to minimize unnecessary computation.
• Preserves all function behavior and output.
• Remains compatible with PyTorch's tensor ops, and does not introduce new dependencies.

This implementation will reduce both Python-overhead and runtime, especially with frequent small tensor calls.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import pytest  # used for our unit tests
import torch  # used for tensor operations
from src.diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint import \
    rescale_noise_cfg

# unit tests

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

def test_identity_guidance_rescale_zero():
    # guidance_rescale=0.0 should return noise_cfg unchanged
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0); out = codeflash_output

def test_full_rescale_guidance_rescale_one():
    # guidance_rescale=1.0 should return the fully rescaled noise_cfg
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    expected = noise_cfg * (std_text / std_cfg)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); out = codeflash_output

def test_halfway_guidance_rescale_point_five():
    # guidance_rescale=0.5 should blend original and rescaled outputs equally
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    expected = 0.5 * rescaled + 0.5 * noise_cfg
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.5); out = codeflash_output

def test_different_shapes_same_batch():
    # noise_cfg and noise_pred_text can be different, but must have same shape
    noise_cfg = torch.randn(4, 3, 8, 8)
    noise_pred_text = torch.randn(4, 3, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.3); out = codeflash_output

def test_broadcasting_batch_size_one():
    # Single batch, should work
    noise_cfg = torch.randn(1, 3, 8, 8)
    noise_pred_text = torch.randn(1, 3, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.7); out = codeflash_output

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

def test_zero_variance_noise_cfg():
    # std=0 in noise_cfg: division by zero produces inf, but torch handles it as inf
    noise_cfg = torch.ones(2, 3, 4, 4) * 5.0
    noise_pred_text = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); out = codeflash_output

def test_zero_variance_noise_pred_text():
    # std=0 in noise_pred_text: scaling factor is zero, so output should be all zeros if guidance_rescale=1.0
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.ones(2, 3, 4, 4) * 2.0
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); out = codeflash_output

def test_zero_variance_both():
    # Both std=0: division by zero, result should be nan
    noise_cfg = torch.ones(2, 3, 4, 4) * 7.0
    noise_pred_text = torch.ones(2, 3, 4, 4) * -3.0
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); out = codeflash_output

def test_negative_guidance_rescale():
    # Negative guidance_rescale should extrapolate away from noise_cfg
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=-1.0); out_neg = codeflash_output
    # Output should be: -1*rescaled + 2*noise_cfg
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    expected = -1.0 * rescaled + 2.0 * noise_cfg

def test_guidance_rescale_greater_than_one():
    # guidance_rescale > 1.0 should extrapolate toward rescaled
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.5); out = codeflash_output
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    expected = 1.5 * rescaled + (1 - 1.5) * noise_cfg

def test_non_float_guidance_rescale():
    # Should accept integer guidance_rescale (implicitly cast to float)
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1); out = codeflash_output
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    expected = noise_cfg * (std_text / std_cfg)

def test_high_dimensional_tensors():
    # Should work for 5D tensor (e.g., batch, channels, depth, height, width)
    noise_cfg = torch.randn(2, 3, 4, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.5); out = codeflash_output

def test_empty_tensor():
    # Should work (no error) for empty tensors
    noise_cfg = torch.empty(0, 3, 4, 4)
    noise_pred_text = torch.empty(0, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.5); out = codeflash_output

def test_single_element_tensor():
    # Should work for a single element (degenerate case)
    noise_cfg = torch.tensor([[[[42.0]]]])
    noise_pred_text = torch.tensor([[[[3.0]]]])
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.5); out = codeflash_output

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

def test_large_tensor_batch():
    # Test with a large batch size, but <100MB
    # (e.g., 100 x 3 x 32 x 32 floats = 100*3*32*32*4 = 1,228,800 bytes ~1.2MB)
    noise_cfg = torch.randn(100, 3, 32, 32)
    noise_pred_text = torch.randn(100, 3, 32, 32)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.25); out = codeflash_output

def test_large_channel_tensor():
    # Test with a large number of channels
    noise_cfg = torch.randn(2, 512, 8, 8)
    noise_pred_text = torch.randn(2, 512, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.3); out = codeflash_output

def test_large_spatial_tensor():
    # Test with large spatial dimensions, but keep under 100MB
    # (e.g., 1 x 3 x 128 x 256 = 98,304 floats * 4 = 393,216 bytes ~0.4MB)
    noise_cfg = torch.randn(1, 3, 128, 256)
    noise_pred_text = torch.randn(1, 3, 128, 256)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.8); out = codeflash_output

def test_large_5d_tensor():
    # 5D tensor, e.g., batch x channel x depth x height x width
    # (e.g., 2 x 3 x 8 x 8 x 8 = 3,072 floats * 4 = 12,288 bytes)
    noise_cfg = torch.randn(2, 3, 8, 8, 8)
    noise_pred_text = torch.randn(2, 3, 8, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.6); out = codeflash_output

def test_performance_large_tensor():
    # Test that function runs in reasonable time for large but <100MB tensor
    noise_cfg = torch.randn(32, 3, 64, 64)  # 393,216 floats * 4 = 1.57MB
    noise_pred_text = torch.randn(32, 3, 64, 64)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.9); out = codeflash_output

# -------------- ERROR HANDLING TESTS --------------


def test_non_tensor_inputs_raise():
    # Should raise an error if inputs are not torch tensors
    with pytest.raises(AttributeError):
        rescale_noise_cfg([[1.0]], [[2.0]], guidance_rescale=0.5)
# 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  # required for tensor creation and manipulation
from src.diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint import \
    rescale_noise_cfg

# unit tests

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

def test_identity_when_guidance_rescale_zero():
    # If guidance_rescale is 0, output should be exactly noise_cfg
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0); output = codeflash_output

def test_full_rescale_when_guidance_rescale_one():
    # If guidance_rescale is 1, output should be the fully rescaled version
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.randn(2, 3, 4, 4)
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    expected = noise_cfg * (std_text / std_cfg)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); output = codeflash_output

def test_half_rescale_when_guidance_rescale_half():
    # When guidance_rescale=0.5, output should be halfway between original and rescaled
    noise_cfg = torch.randn(1, 2, 2, 2)
    noise_pred_text = torch.randn(1, 2, 2, 2)
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    expected = 0.5 * rescaled + 0.5 * noise_cfg
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.5); output = codeflash_output

def test_output_shape_matches_input():
    # Output should have the same shape as noise_cfg
    noise_cfg = torch.randn(4, 2, 8, 8)
    noise_pred_text = torch.randn(4, 2, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.3); output = codeflash_output

def test_batch_independence():
    # Each batch should be rescaled independently
    noise_cfg = torch.ones(2, 2, 2, 2)
    noise_pred_text = torch.zeros(2, 2, 2, 2)
    # The std of ones is 0, so division by zero, but torch returns nan. Let's use different values to avoid this.
    noise_cfg[0] = 1.0
    noise_cfg[1] = 2.0
    noise_pred_text[0] = 1.0
    noise_pred_text[1] = 10.0
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); output = codeflash_output
    # The ratio of stds for each batch should be applied independently
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    expected = noise_cfg * (std_text / std_cfg)

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

def test_zero_std_noise_cfg():
    # If noise_cfg is constant, std_cfg is zero, so division by zero yields inf or nan
    noise_cfg = torch.ones(1, 2, 2, 2)
    noise_pred_text = torch.randn(1, 2, 2, 2)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); output = codeflash_output

def test_zero_std_noise_pred_text():
    # If noise_pred_text is constant, std_text is zero, so output should be all zeros if guidance_rescale=1.0
    noise_cfg = torch.randn(1, 2, 2, 2)
    noise_pred_text = torch.ones(1, 2, 2, 2)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); output = codeflash_output

def test_negative_guidance_rescale():
    # Negative guidance_rescale should extrapolate between rescaled and original
    noise_cfg = torch.randn(1, 2, 2, 2)
    noise_pred_text = torch.randn(1, 2, 2, 2)
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    g = -0.5
    expected = g * rescaled + (1 - g) * noise_cfg
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=g); output = codeflash_output

def test_guidance_rescale_greater_than_one():
    # guidance_rescale > 1 should extrapolate beyond the rescaled value
    noise_cfg = torch.randn(1, 2, 2, 2)
    noise_pred_text = torch.randn(1, 2, 2, 2)
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    g = 2.0
    expected = g * rescaled + (1 - g) * noise_cfg
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=g); output = codeflash_output


def test_single_pixel_tensor():
    # Test with a tensor of shape (1,1,1,1)
    noise_cfg = torch.tensor([[[[1.0]]]])
    noise_pred_text = torch.tensor([[[[2.0]]]])
    # stds are zero, so output should be nan or inf when guidance_rescale=1
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); output = codeflash_output

def test_non_float_tensor():
    # Test with integer tensor types (should work, but output will be float)
    noise_cfg = torch.randint(0, 10, (1, 2, 2, 2), dtype=torch.int32)
    noise_pred_text = torch.randint(0, 10, (1, 2, 2, 2), dtype=torch.int32)
    codeflash_output = rescale_noise_cfg(noise_cfg.float(), noise_pred_text.float(), guidance_rescale=0.5); output = codeflash_output

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

def test_large_tensor_performance_and_correctness():
    # Test with a large tensor (but < 100MB)
    # (100, 3, 32, 32) of float32 is about 1.2MB, well within the limit
    shape = (100, 3, 32, 32)
    noise_cfg = torch.randn(*shape)
    noise_pred_text = torch.randn(*shape)
    # Just ensure it runs and output shape matches
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.7); output = codeflash_output
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    rescaled = noise_cfg * (std_text / std_cfg)
    expected = 0.7 * rescaled + 0.3 * noise_cfg

def test_large_batch_dimension():
    # Test with a large batch size
    shape = (512, 1, 8, 8)  # 512*1*8*8*4 = 131072 bytes = 0.125MB
    noise_cfg = torch.randn(*shape)
    noise_pred_text = torch.randn(*shape)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.8); output = codeflash_output

def test_large_channel_dimension():
    # Test with a large channel size
    shape = (2, 512, 4, 4)  # 2*512*4*4*4 = 65,536 bytes = 0.0625MB
    noise_cfg = torch.randn(*shape)
    noise_pred_text = torch.randn(*shape)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.9); output = codeflash_output

def test_large_spatial_dimensions():
    # Test with large spatial dimensions
    shape = (1, 3, 128, 128)  # 1*3*128*128*4 = 196,608 bytes = 0.1875MB
    noise_cfg = torch.randn(*shape)
    noise_pred_text = torch.randn(*shape)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.6); output = codeflash_output

def test_large_tensor_guidance_rescale_zero_and_one():
    # For large tensors, test guidance_rescale=0 and 1
    shape = (100, 3, 32, 32)
    noise_cfg = torch.randn(*shape)
    noise_pred_text = torch.randn(*shape)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0); out_zero = codeflash_output
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); out_one = codeflash_output
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    expected_one = noise_cfg * (std_text / std_cfg)
# 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-rescale_noise_cfg-mbdv2rmj and push.