⚡️ Speed up function rescale_noise_cfg by 34%

[codeflash-ai]

📄 34% (0.34x) speedup for rescale_noise_cfg in src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py

⏱️ Runtime : 68.8 milliseconds → 51.2 milliseconds (best of 64 runs)

📝 Explanation and details

Key optimizations:

• Use tuple for dim argument in .std(), which is slightly faster than list.
• If guidance_rescale is 0.0 or 1.0, short-circuit and avoid unnecessary computations.
• Compute the final result with a single fused multiplication (eliminating an extra allocation and avoiding temporary tensors).
• Avoid recomputation of dim tuple.
• Fewer intermediate tensors allocated for the result, aiding both speed and lower memory use.

Correctness:
All logic and results remain unchanged for all possible inputs. All doc strings and comments regarding the core logic are preserved.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 44 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.deprecated.alt_diffusion.pipeline_alt_diffusion import \
    rescale_noise_cfg

# unit tests

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

def test_identity_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_one():
    # guidance_rescale=1.0 should return 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_half_rescale_point_five():
    # guidance_rescale=0.5 should blend original and rescaled
    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_broadcasting_batch_size_1():
    # Test with batch size 1 to ensure broadcasting works
    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.75); out = codeflash_output

def test_broadcasting_singleton_channel():
    # Test with singleton channel dimension
    noise_cfg = torch.randn(2, 1, 4, 4)
    noise_pred_text = torch.randn(2, 1, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.25); out = codeflash_output

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

def test_zero_std_noise_cfg():
    # If noise_cfg is constant (std=0), division by zero should not crash (should produce inf or nan)
    noise_cfg = torch.ones(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 = codeflash_output

def test_zero_std_noise_pred_text():
    # If noise_pred_text is constant (std=0), scaling factor is zero, so output should be all zeros when rescale=1
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.ones(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); out = codeflash_output

def test_negative_guidance_rescale():
    # guidance_rescale < 0 should extrapolate
    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 = codeflash_output
    # Output should be noise_cfg - (rescaled - 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 should extrapolate towards 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 - 0.5 * noise_cfg


def test_empty_tensor():
    # Should work with empty tensors (zero batch)
    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 with single element tensors
    noise_cfg = torch.tensor([[[[1.0]]]])
    noise_pred_text = torch.tensor([[[[2.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_performance():
    # Test with large but <100MB tensor
    # 100 * 3 * 128 * 128 * 4 bytes = ~25MB
    noise_cfg = torch.randn(100, 3, 128, 128)
    noise_pred_text = torch.randn(100, 3, 128, 128)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.7); out = codeflash_output

def test_large_batch():
    # Test with large batch dimension
    noise_cfg = torch.randn(512, 3, 16, 16)
    noise_pred_text = torch.randn(512, 3, 16, 16)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.3); out = codeflash_output

def test_large_channel():
    # Test with large channel dimension
    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.6); out = codeflash_output

def test_large_singleton_dimensions():
    # Test with large singleton dimensions (should broadcast correctly)
    noise_cfg = torch.randn(10, 1, 64, 1)
    noise_pred_text = torch.randn(10, 1, 64, 1)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.9); out = codeflash_output

# -------------------------- DETERMINISM TEST --------------------------

def test_determinism():
    # Running the function twice with the same input should produce the same output
    torch.manual_seed(42)
    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.4); out1 = codeflash_output
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.4); out2 = codeflash_output

# -------------------------- FLOAT PRECISION TESTS --------------------------

@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
def test_float_precision(dtype):
    # Should work with both float32 and float64
    noise_cfg = torch.randn(2, 3, 4, 4, dtype=dtype)
    noise_pred_text = torch.randn(2, 3, 4, 4, dtype=dtype)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.5); out = codeflash_output

# -------------------------- NON-4D TENSOR TESTS --------------------------

def test_3d_tensor():
    # Should work with 3D tensors (e.g., [batch, channel, length])
    noise_cfg = torch.randn(2, 3, 16)
    noise_pred_text = torch.randn(2, 3, 16)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.2); out = codeflash_output

def test_5d_tensor():
    # Should work with 5D tensors (e.g., [batch, channel, depth, height, width])
    noise_cfg = torch.randn(2, 3, 4, 8, 8)
    noise_pred_text = torch.randn(2, 3, 4, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.8); out = codeflash_output

# -------------------------- INVALID INPUT TESTS --------------------------

def test_non_tensor_input_raises():
    # Should raise AttributeError if input is not a tensor
    with pytest.raises(AttributeError):
        rescale_noise_cfg([[1,2],[3,4]], [[5,6],[7,8]], guidance_rescale=0.5)

def test_guidance_rescale_nan():
    # If guidance_rescale is nan, output should be nan
    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=float('nan')); out = codeflash_output
# 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  # needed for tensor operations
from src.diffusers.pipelines.deprecated.alt_diffusion.pipeline_alt_diffusion import \
    rescale_noise_cfg

# unit tests

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

def test_identity_guidance_rescale_zero():
    # When guidance_rescale=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); result = codeflash_output

def test_full_rescale_guidance_rescale_one():
    # When guidance_rescale=1, output should be fully rescaled
    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); result = codeflash_output

def test_half_rescale_guidance_rescale_half():
    # When guidance_rescale=0.5, output should be halfway between original and rescaled
    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); result = codeflash_output

def test_equal_inputs():
    # If noise_cfg == noise_pred_text, stds are equal, so rescale factor is 1, so output should be noise_cfg
    noise_cfg = torch.randn(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_cfg, guidance_rescale=1.0); result = codeflash_output

def test_batch_size_one():
    # Test with batch size 1
    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); result = codeflash_output

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

def test_zero_std_noise_cfg():
    # If noise_cfg is constant (std=0), division by zero should result in inf or nan, but torch handles this as inf
    noise_cfg = torch.ones(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); result = codeflash_output

def test_zero_std_noise_pred_text():
    # If noise_pred_text is constant (std=0), rescale factor is zero, so output should be zeros when guidance_rescale=1
    noise_cfg = torch.randn(2, 3, 4, 4)
    noise_pred_text = torch.ones(2, 3, 4, 4)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.0); result = codeflash_output

def test_negative_guidance_rescale():
    # Negative guidance_rescale should extrapolate in the opposite direction
    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 = -1.0 * rescaled + 2.0 * noise_cfg
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=-1.0); result = codeflash_output

def test_guidance_rescale_greater_than_one():
    # guidance_rescale>1 should extrapolate beyond the rescaled value
    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 = 1.5 * rescaled + (-0.5) * noise_cfg
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=1.5); result = codeflash_output

def test_singleton_dimensions():
    # Test with singleton dimensions (e.g., [batch, 1, 1, 1])
    noise_cfg = torch.randn(2, 1, 1, 1)
    noise_pred_text = torch.randn(2, 1, 1, 1)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.3); result = codeflash_output


def test_empty_tensor():
    # If input is empty, output should also be empty and not error
    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); result = codeflash_output

def test_high_dimensional_tensor():
    # Test with 5D tensor
    noise_cfg = torch.randn(2, 3, 4, 4, 2)
    noise_pred_text = torch.randn(2, 3, 4, 4, 2)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.8); result = codeflash_output

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

def test_large_tensor_performance():
    # Test with large tensor but <100MB
    # 100*3*64*64*4 bytes = ~5MB, so 256*3*64*64 = ~12.5MB
    noise_cfg = torch.randn(128, 3, 64, 64)
    noise_pred_text = torch.randn(128, 3, 64, 64)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.9); result = codeflash_output

def test_large_batch_size():
    # Test with large batch dimension
    noise_cfg = torch.randn(512, 3, 8, 8)
    noise_pred_text = torch.randn(512, 3, 8, 8)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.6); result = codeflash_output

def test_large_singleton_channel():
    # Test with large spatial dimensions but singleton channel
    noise_cfg = torch.randn(4, 1, 128, 128)
    noise_pred_text = torch.randn(4, 1, 128, 128)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.4); result = codeflash_output

def test_large_guidance_rescale_range():
    # Test with a range of guidance_rescale values on a large tensor
    noise_cfg = torch.randn(16, 3, 64, 64)
    noise_pred_text = torch.randn(16, 3, 64, 64)
    for alpha in [0.0, 0.25, 0.5, 0.75, 1.0]:
        codeflash_output = rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=alpha); result = codeflash_output

def test_large_tensor_equal_inputs():
    # Test large tensor with noise_cfg == noise_pred_text
    noise_cfg = torch.randn(64, 3, 32, 32)
    codeflash_output = rescale_noise_cfg(noise_cfg, noise_cfg, guidance_rescale=1.0); result = codeflash_output
# 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-mbdmg9wn and push.