mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-04-10 15:00:23 +00:00
265 lines
9.0 KiB
Python
265 lines
9.0 KiB
Python
import pytest
|
|
import torch
|
|
import numpy as np
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Import the functions we're testing
|
|
from library.custom_train_functions import (
|
|
apply_snr_weight,
|
|
scale_v_prediction_loss_like_noise_prediction,
|
|
get_snr_scale,
|
|
add_v_prediction_like_loss,
|
|
apply_debiased_estimation,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def loss():
|
|
return torch.ones(2, 1)
|
|
|
|
|
|
@pytest.fixture
|
|
def timesteps():
|
|
return torch.tensor([[200, 600]], dtype=torch.int32)
|
|
|
|
|
|
@pytest.fixture
|
|
def noise_scheduler():
|
|
scheduler = MagicMock()
|
|
scheduler.get_snr_for_timestep = MagicMock(return_value=torch.tensor([0.294, 0.39]))
|
|
scheduler.all_snr = torch.tensor([0.05, 0.1, 0.15, 0.2, 0.25, 0.5, 1.0])
|
|
return scheduler
|
|
|
|
|
|
# Tests for apply_snr_weight
|
|
def test_apply_snr_weight_with_get_snr_method(loss, timesteps, noise_scheduler):
|
|
image_size = 64
|
|
gamma = 5.0
|
|
|
|
result = apply_snr_weight(
|
|
loss,
|
|
timesteps,
|
|
noise_scheduler,
|
|
gamma,
|
|
v_prediction=False,
|
|
image_size=image_size,
|
|
)
|
|
|
|
expected_result = torch.tensor([[1.0, 1.0]])
|
|
|
|
assert torch.allclose(result, expected_result, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_apply_snr_weight_with_all_snr(loss, timesteps):
|
|
gamma = 5.0
|
|
|
|
# Modify the mock to not use get_snr_for_timestep
|
|
mock_noise_scheduler_no_method = MagicMock()
|
|
mock_noise_scheduler_no_method.get_snr_for_timestep = None
|
|
mock_noise_scheduler_no_method.all_snr = torch.tensor([0.05, 0.1, 0.15, 0.2, 0.25, 0.5, 1.0])
|
|
|
|
result = apply_snr_weight(loss, timesteps, mock_noise_scheduler_no_method, gamma, v_prediction=False)
|
|
|
|
expected_result = torch.tensor([1.0, 1.0])
|
|
assert torch.allclose(result, expected_result, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_apply_snr_weight_with_v_prediction(loss, timesteps, noise_scheduler):
|
|
gamma = 5.0
|
|
|
|
result = apply_snr_weight(loss, timesteps, noise_scheduler, gamma, v_prediction=True)
|
|
|
|
expected_result = torch.tensor([[0.2272, 0.2806], [0.2272, 0.2806]])
|
|
|
|
assert torch.allclose(result, expected_result, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
# Tests for scale_v_prediction_loss_like_noise_prediction
|
|
def test_scale_v_prediction_loss(loss, timesteps, noise_scheduler):
|
|
with patch("library.custom_train_functions.get_snr_scale") as mock_get_snr_scale:
|
|
mock_get_snr_scale.return_value = torch.tensor([0.9, 0.8])
|
|
|
|
result = scale_v_prediction_loss_like_noise_prediction(loss, timesteps, noise_scheduler)
|
|
|
|
mock_get_snr_scale.assert_called_once_with(timesteps, noise_scheduler, None)
|
|
|
|
# Apply broadcasting for multiplication
|
|
scale = torch.tensor([[0.9, 0.8], [0.9, 0.8]])
|
|
expected_result = loss * scale
|
|
assert torch.allclose(result, expected_result)
|
|
|
|
|
|
# Tests for get_snr_scale
|
|
def test_get_snr_scale_with_get_snr_method(timesteps, noise_scheduler):
|
|
image_size = 64
|
|
|
|
result = get_snr_scale(timesteps, noise_scheduler, image_size)
|
|
|
|
# Verify the method was called correctly
|
|
noise_scheduler.get_snr_for_timestep.assert_called_once_with(timesteps, image_size)
|
|
|
|
expected_scale = torch.tensor([0.2272, 0.2806])
|
|
|
|
assert torch.allclose(result, expected_scale, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_get_snr_scale_with_all_snr(timesteps):
|
|
# Create a scheduler that only has all_snr
|
|
mock_scheduler_all_snr = MagicMock()
|
|
mock_scheduler_all_snr.get_snr_for_timestep = None
|
|
mock_scheduler_all_snr.all_snr = torch.tensor([0.05, 0.1, 0.15, 0.2, 0.25, 0.5, 0.75, 1.0])
|
|
|
|
result = get_snr_scale(timesteps, mock_scheduler_all_snr)
|
|
|
|
expected_scale = torch.tensor([[0.5000, 0.5000]])
|
|
|
|
assert torch.allclose(result, expected_scale, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_get_snr_scale_with_large_snr(timesteps, noise_scheduler):
|
|
# Set up the mock with a very large SNR value
|
|
noise_scheduler.get_snr_for_timestep.return_value = torch.tensor([2000.0, 5.0])
|
|
|
|
result = get_snr_scale(timesteps, noise_scheduler)
|
|
|
|
expected_scale = torch.tensor([0.9990, 0.8333])
|
|
|
|
assert torch.allclose(result, expected_scale, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
# Tests for add_v_prediction_like_loss
|
|
def test_add_v_prediction_like_loss(loss, timesteps, noise_scheduler):
|
|
v_pred_like_loss = torch.tensor([0.3, 0.2]).view(2, 1)
|
|
|
|
with patch("library.custom_train_functions.get_snr_scale") as mock_get_snr_scale:
|
|
mock_get_snr_scale.return_value = torch.tensor([0.9, 0.8])
|
|
|
|
result = add_v_prediction_like_loss(loss, timesteps, noise_scheduler, v_pred_like_loss)
|
|
|
|
mock_get_snr_scale.assert_called_once_with(timesteps, noise_scheduler, None)
|
|
|
|
expected_result = torch.tensor([[1.3333, 1.3750], [1.2222, 1.2500]])
|
|
assert torch.allclose(result, expected_result, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
# Tests for apply_debiased_estimation
|
|
def test_apply_debiased_estimation_no_snr(loss, timesteps):
|
|
# Create a scheduler without SNR methods
|
|
scheduler_without_snr = MagicMock()
|
|
# Need to explicitly set attribute to None instead of deleting
|
|
scheduler_without_snr.get_snr_for_timestep = None
|
|
|
|
result = apply_debiased_estimation(loss, timesteps, scheduler_without_snr)
|
|
|
|
# When no SNR methods are available, the function should return the loss unchanged
|
|
assert torch.equal(result, loss)
|
|
|
|
|
|
def test_apply_debiased_estimation_with_get_snr_method(loss, timesteps, noise_scheduler):
|
|
# Test with v_prediction=False
|
|
result_no_v = apply_debiased_estimation(loss, timesteps, noise_scheduler, v_prediction=False)
|
|
|
|
expected_result_no_v = torch.tensor([[1.8443, 1.6013], [1.8443, 1.6013]])
|
|
|
|
assert torch.allclose(result_no_v, expected_result_no_v, rtol=1e-4, atol=1e-4)
|
|
|
|
# Test with v_prediction=True
|
|
result_v = apply_debiased_estimation(loss, timesteps, noise_scheduler, v_prediction=True)
|
|
|
|
expected_result_v = torch.tensor([[0.7728, 0.7194], [0.7728, 0.7194]])
|
|
|
|
assert torch.allclose(result_v, expected_result_v, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_apply_debiased_estimation_with_all_snr(loss, timesteps):
|
|
# Create a scheduler that only has all_snr
|
|
mock_scheduler_all_snr = MagicMock()
|
|
mock_scheduler_all_snr.get_snr_for_timestep = None
|
|
mock_scheduler_all_snr.all_snr = torch.tensor([0.05, 0.1, 0.15, 0.2, 0.25, 0.5, 1.0])
|
|
|
|
result = apply_debiased_estimation(loss, timesteps, mock_scheduler_all_snr, v_prediction=False)
|
|
|
|
expected_result = torch.tensor([[1.0, 1.0]])
|
|
|
|
assert torch.allclose(result, expected_result, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
def test_apply_debiased_estimation_with_large_snr(loss, timesteps, noise_scheduler):
|
|
# Set up the mock with a very large SNR value
|
|
noise_scheduler.get_snr_for_timestep.return_value = torch.tensor([2000.0, 5.0])
|
|
|
|
result = apply_debiased_estimation(loss, timesteps, noise_scheduler, v_prediction=False)
|
|
|
|
expected_result = torch.tensor([[0.0316, 0.4472], [0.0316, 0.4472]])
|
|
|
|
assert torch.allclose(result, expected_result, rtol=1e-4, atol=1e-4)
|
|
|
|
|
|
# Additional edge cases
|
|
def test_empty_tensors(noise_scheduler):
|
|
# Test with empty tensors
|
|
loss = torch.tensor([], dtype=torch.float32)
|
|
timesteps = torch.tensor([], dtype=torch.int32)
|
|
|
|
assert isinstance(timesteps, torch.IntTensor)
|
|
|
|
noise_scheduler.get_snr_for_timestep.return_value = torch.tensor([], dtype=torch.float32)
|
|
|
|
result = apply_snr_weight(loss, timesteps, noise_scheduler, gamma=5.0)
|
|
|
|
assert result.shape == loss.shape
|
|
assert result.dtype == loss.dtype
|
|
|
|
|
|
def test_different_device_compatibility(loss, timesteps, noise_scheduler):
|
|
gamma = 5.0
|
|
device = torch.device("cpu")
|
|
|
|
# For a real device test, we need to create actual tensors on devices
|
|
loss_on_device = loss.to(device)
|
|
|
|
# Mock the SNR tensor that would be returned with proper device handling
|
|
snr_tensor = torch.tensor([0.204, 0.294], device=device)
|
|
noise_scheduler.get_snr_for_timestep.return_value = snr_tensor
|
|
|
|
result = apply_snr_weight(loss_on_device, timesteps, noise_scheduler, gamma)
|
|
|
|
# Additional tests for new functionality
|
|
def test_apply_snr_weight_with_image_size(loss, timesteps, noise_scheduler):
|
|
"""Test SNR weight application with image size consideration"""
|
|
gamma = 5.0
|
|
image_sizes = [None, 64, (256, 256)]
|
|
|
|
for image_size in image_sizes:
|
|
result = apply_snr_weight(
|
|
loss,
|
|
timesteps,
|
|
noise_scheduler,
|
|
gamma,
|
|
v_prediction=False,
|
|
image_size=image_size
|
|
)
|
|
|
|
# Allow for broadcasting
|
|
assert result.shape[0] == loss.shape[0]
|
|
assert result.dtype == loss.dtype
|
|
|
|
def test_apply_debiased_estimation_variations(loss, timesteps, noise_scheduler):
|
|
"""Test debiased estimation with different image sizes and prediction types"""
|
|
image_sizes = [None, 64, (256, 256)]
|
|
prediction_types = [True, False]
|
|
|
|
for image_size in image_sizes:
|
|
for v_prediction in prediction_types:
|
|
result = apply_debiased_estimation(
|
|
loss,
|
|
timesteps,
|
|
noise_scheduler,
|
|
v_prediction=v_prediction,
|
|
image_size=image_size
|
|
)
|
|
|
|
# Allow for broadcasting
|
|
assert result.shape[0] == loss.shape[0]
|
|
assert result.dtype == loss.dtype
|