feat: add SDXL Turbo image generator wrapper
This commit is contained in:
@@ -0,0 +1,52 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
from PIL import Image
|
||||
|
||||
from server.asset_generator import AssetGenerator
|
||||
|
||||
|
||||
class TestAssetGenerator:
|
||||
@patch("server.asset_generator.AutoPipelineForText2Image")
|
||||
def test_init_loads_model(self, mock_pipeline_cls):
|
||||
"""Generator loads the SDXL Turbo pipeline on init."""
|
||||
mock_pipe = MagicMock()
|
||||
mock_pipeline_cls.from_pretrained.return_value = mock_pipe
|
||||
mock_pipe.to.return_value = mock_pipe
|
||||
|
||||
gen = AssetGenerator(device="cpu")
|
||||
mock_pipeline_cls.from_pretrained.assert_called_once()
|
||||
|
||||
@patch("server.asset_generator.AutoPipelineForText2Image")
|
||||
def test_generate_returns_bytes(self, mock_pipeline_cls):
|
||||
"""Generate returns PNG bytes."""
|
||||
mock_pipe = MagicMock()
|
||||
mock_pipeline_cls.from_pretrained.return_value = mock_pipe
|
||||
mock_pipe.to.return_value = mock_pipe
|
||||
|
||||
# Mock pipeline output
|
||||
fake_image = Image.new("RGB", (512, 512), color="black")
|
||||
mock_result = MagicMock()
|
||||
mock_result.images = [fake_image]
|
||||
mock_pipe.return_value = mock_result
|
||||
|
||||
gen = AssetGenerator(device="cpu")
|
||||
data = gen.generate("dark void, horror")
|
||||
assert isinstance(data, bytes)
|
||||
assert len(data) > 0
|
||||
|
||||
@patch("server.asset_generator.AutoPipelineForText2Image")
|
||||
def test_generate_uses_negative_prompt(self, mock_pipeline_cls):
|
||||
"""Generate passes the negative prompt to the pipeline."""
|
||||
mock_pipe = MagicMock()
|
||||
mock_pipeline_cls.from_pretrained.return_value = mock_pipe
|
||||
mock_pipe.to.return_value = mock_pipe
|
||||
|
||||
fake_image = Image.new("RGB", (512, 512), color="black")
|
||||
mock_result = MagicMock()
|
||||
mock_result.images = [fake_image]
|
||||
mock_pipe.return_value = mock_result
|
||||
|
||||
gen = AssetGenerator(device="cpu")
|
||||
gen.generate("test prompt")
|
||||
|
||||
call_kwargs = mock_pipe.call_args
|
||||
assert "negative_prompt" in call_kwargs.kwargs
|
||||
Reference in New Issue
Block a user