51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from onnx_web.chain.blend_denoise_localstd import BlendDenoiseLocalStdStage
|
||
|
|
from onnx_web.chain.result import ImageMetadata, StageResult
|
||
|
|
from onnx_web.params import ImageParams, Size
|
||
|
|
|
||
|
|
|
||
|
|
class TestBlendDenoiseLocalStdStage(unittest.TestCase):
|
||
|
|
def test_run(self):
|
||
|
|
# Create a dummy image
|
||
|
|
image = Image.new("RGB", (64, 64), color="white")
|
||
|
|
|
||
|
|
# Create a dummy StageResult object
|
||
|
|
sources = StageResult.from_images(
|
||
|
|
[image],
|
||
|
|
metadata=[
|
||
|
|
ImageMetadata(
|
||
|
|
ImageParams("test", "txt2img", "ddim", "test", 5.0, 25, 0),
|
||
|
|
Size(64, 64),
|
||
|
|
)
|
||
|
|
],
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create an instance of BlendDenoiseLocalStdStage
|
||
|
|
stage = BlendDenoiseLocalStdStage()
|
||
|
|
|
||
|
|
# Call the run method with dummy parameters
|
||
|
|
result = stage.run(
|
||
|
|
_worker=None,
|
||
|
|
_server=None,
|
||
|
|
_stage=None,
|
||
|
|
_params=None,
|
||
|
|
sources=sources,
|
||
|
|
strength=5,
|
||
|
|
range=4,
|
||
|
|
stage_source=None,
|
||
|
|
callback=None,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Assert that the result is an instance of StageResult
|
||
|
|
self.assertIsInstance(result, StageResult)
|
||
|
|
|
||
|
|
# Assert that the result contains the denoised image
|
||
|
|
self.assertEqual(len(result), 1)
|
||
|
|
self.assertEqual(result.size(), Size(64, 64))
|
||
|
|
|
||
|
|
# Assert that the metadata is preserved
|
||
|
|
self.assertEqual(result.metadata, sources.metadata)
|