mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 14:22:11 -04:00
Ettore Di Giacinto
7f9cf9c63d
Expose diffusers audio pipelines through the existing sound-generation RPC. AudioLDM2 can now return PCM WAV output from the model gallery without a separate backend. Assisted-by: Codex:gpt-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
512 lines
19 KiB
Python
512 lines
19 KiB
Python
"""
|
|
A test script to test the gRPC service and dynamic loader
|
|
"""
|
|
import unittest
|
|
import subprocess
|
|
import time
|
|
import os
|
|
import tempfile
|
|
import wave
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
# Import dynamic loader for testing (these don't need gRPC)
|
|
import backend
|
|
import diffusers_dynamic_loader as loader
|
|
from diffusers import DiffusionPipeline, StableDiffusionPipeline
|
|
|
|
# Try to import gRPC modules - may not be available during unit testing
|
|
try:
|
|
import grpc
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
GRPC_AVAILABLE = True
|
|
except ImportError:
|
|
GRPC_AVAILABLE = False
|
|
|
|
|
|
@unittest.skipUnless(GRPC_AVAILABLE, "gRPC modules not available")
|
|
class TestBackendServicer(unittest.TestCase):
|
|
"""
|
|
TestBackendServicer is the class that tests the gRPC service
|
|
"""
|
|
def setUp(self):
|
|
"""
|
|
This method sets up the gRPC service by starting the server
|
|
"""
|
|
self.service = subprocess.Popen(["python3", "backend.py", "--addr", "localhost:50051"])
|
|
|
|
def tearDown(self) -> None:
|
|
"""
|
|
This method tears down the gRPC service by terminating the server
|
|
"""
|
|
self.service.kill()
|
|
self.service.wait()
|
|
|
|
def test_server_startup(self):
|
|
"""
|
|
This method tests if the server starts up successfully
|
|
"""
|
|
time.sleep(20)
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.Health(backend_pb2.HealthMessage())
|
|
self.assertEqual(response.message, b'OK')
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("Server failed to start")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test_load_model(self):
|
|
"""
|
|
This method tests if the model is loaded successfully
|
|
"""
|
|
time.sleep(20)
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="Lykon/dreamshaper-8"))
|
|
self.assertTrue(response.success)
|
|
self.assertEqual(response.message, "Model loaded successfully")
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("LoadModel service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test(self):
|
|
"""
|
|
This method tests if the backend can generate images
|
|
"""
|
|
time.sleep(20)
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="Lykon/dreamshaper-8"))
|
|
print(response.message)
|
|
self.assertTrue(response.success)
|
|
image_req = backend_pb2.GenerateImageRequest(positive_prompt="cat", width=16,height=16, dst="test.jpg")
|
|
re = stub.GenerateImage(image_req)
|
|
self.assertTrue(re.success)
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("Image gen service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
|
|
class TestDiffusersDynamicLoader(unittest.TestCase):
|
|
"""Test cases for the diffusers dynamic loader functionality."""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""Set up test fixtures - clear caches to ensure fresh discovery."""
|
|
# Reset the caches to ensure fresh discovery
|
|
loader._pipeline_registry = None
|
|
loader._task_aliases = None
|
|
|
|
def test_camel_to_kebab_conversion(self):
|
|
"""Test CamelCase to kebab-case conversion."""
|
|
test_cases = [
|
|
("StableDiffusionPipeline", "stable-diffusion-pipeline"),
|
|
("StableDiffusionXLPipeline", "stable-diffusion-xl-pipeline"),
|
|
("FluxPipeline", "flux-pipeline"),
|
|
("DiffusionPipeline", "diffusion-pipeline"),
|
|
]
|
|
for input_val, expected in test_cases:
|
|
with self.subTest(input=input_val):
|
|
result = loader._camel_to_kebab(input_val)
|
|
self.assertEqual(result, expected)
|
|
|
|
def test_extract_task_keywords(self):
|
|
"""Test task keyword extraction from class names."""
|
|
# Test text-to-image detection
|
|
aliases = loader._extract_task_keywords("StableDiffusionPipeline")
|
|
self.assertIn("stable-diffusion", aliases)
|
|
|
|
# Test img2img detection
|
|
aliases = loader._extract_task_keywords("StableDiffusionImg2ImgPipeline")
|
|
self.assertIn("image-to-image", aliases)
|
|
self.assertIn("img2img", aliases)
|
|
|
|
# Test inpainting detection
|
|
aliases = loader._extract_task_keywords("StableDiffusionInpaintPipeline")
|
|
self.assertIn("inpainting", aliases)
|
|
self.assertIn("inpaint", aliases)
|
|
|
|
# Test depth2img detection
|
|
aliases = loader._extract_task_keywords("StableDiffusionDepth2ImgPipeline")
|
|
self.assertIn("depth-to-image", aliases)
|
|
|
|
def test_discover_pipelines_finds_known_classes(self):
|
|
"""Test that pipeline discovery finds at least one known pipeline class."""
|
|
registry = loader.get_pipeline_registry()
|
|
|
|
# Check that the registry is not empty
|
|
self.assertGreater(len(registry), 0, "Pipeline registry should not be empty")
|
|
|
|
# Check for known pipeline classes
|
|
known_pipelines = [
|
|
"StableDiffusionPipeline",
|
|
"DiffusionPipeline",
|
|
]
|
|
|
|
for pipeline_name in known_pipelines:
|
|
with self.subTest(pipeline=pipeline_name):
|
|
self.assertIn(
|
|
pipeline_name,
|
|
registry,
|
|
f"Expected to find {pipeline_name} in registry"
|
|
)
|
|
|
|
def test_discover_pipelines_caches_results(self):
|
|
"""Test that pipeline discovery results are cached."""
|
|
# Get registry twice
|
|
registry1 = loader.get_pipeline_registry()
|
|
registry2 = loader.get_pipeline_registry()
|
|
|
|
# Should be the same object (cached)
|
|
self.assertIs(registry1, registry2, "Registry should be cached")
|
|
|
|
def test_get_available_pipelines(self):
|
|
"""Test getting list of available pipelines."""
|
|
available = loader.get_available_pipelines()
|
|
|
|
# Should return a list
|
|
self.assertIsInstance(available, list)
|
|
|
|
# Should contain known pipelines
|
|
self.assertIn("StableDiffusionPipeline", available)
|
|
self.assertIn("DiffusionPipeline", available)
|
|
|
|
# Should be sorted
|
|
self.assertEqual(available, sorted(available))
|
|
|
|
def test_get_available_tasks(self):
|
|
"""Test getting list of available task aliases."""
|
|
tasks = loader.get_available_tasks()
|
|
|
|
# Should return a list
|
|
self.assertIsInstance(tasks, list)
|
|
|
|
# Should be sorted
|
|
self.assertEqual(tasks, sorted(tasks))
|
|
|
|
def test_resolve_pipeline_class_by_name(self):
|
|
"""Test resolving pipeline class by exact name."""
|
|
cls = loader.resolve_pipeline_class(class_name="StableDiffusionPipeline")
|
|
self.assertEqual(cls, StableDiffusionPipeline)
|
|
|
|
def test_resolve_pipeline_class_by_name_case_insensitive(self):
|
|
"""Test that class name resolution is case-insensitive."""
|
|
cls1 = loader.resolve_pipeline_class(class_name="StableDiffusionPipeline")
|
|
cls2 = loader.resolve_pipeline_class(class_name="stablediffusionpipeline")
|
|
self.assertEqual(cls1, cls2)
|
|
|
|
def test_resolve_pipeline_class_by_task(self):
|
|
"""Test resolving pipeline class by task alias."""
|
|
# Get the registry to find available tasks
|
|
aliases = loader.get_task_aliases()
|
|
|
|
# Test with a common task that should be available
|
|
if "stable-diffusion" in aliases:
|
|
cls = loader.resolve_pipeline_class(task="stable-diffusion")
|
|
self.assertIsNotNone(cls)
|
|
|
|
def test_resolve_pipeline_class_unknown_name_raises(self):
|
|
"""Test that resolving unknown class name raises ValueError with helpful message."""
|
|
with self.assertRaises(ValueError) as ctx:
|
|
loader.resolve_pipeline_class(class_name="NonExistentPipeline")
|
|
|
|
# Check that error message includes available pipelines
|
|
error_msg = str(ctx.exception)
|
|
self.assertIn("Unknown pipeline class", error_msg)
|
|
self.assertIn("Available pipelines", error_msg)
|
|
|
|
def test_resolve_pipeline_class_unknown_task_raises(self):
|
|
"""Test that resolving unknown task raises ValueError with helpful message."""
|
|
with self.assertRaises(ValueError) as ctx:
|
|
loader.resolve_pipeline_class(task="nonexistent-task-xyz")
|
|
|
|
# Check that error message includes available tasks
|
|
error_msg = str(ctx.exception)
|
|
self.assertIn("Unknown task", error_msg)
|
|
self.assertIn("Available tasks", error_msg)
|
|
|
|
def test_resolve_pipeline_class_no_params_raises(self):
|
|
"""Test that calling with no parameters raises helpful ValueError."""
|
|
with self.assertRaises(ValueError) as ctx:
|
|
loader.resolve_pipeline_class()
|
|
|
|
error_msg = str(ctx.exception)
|
|
self.assertIn("Must provide at least one of", error_msg)
|
|
|
|
def test_get_pipeline_info(self):
|
|
"""Test getting pipeline information."""
|
|
info = loader.get_pipeline_info("StableDiffusionPipeline")
|
|
|
|
self.assertEqual(info['name'], "StableDiffusionPipeline")
|
|
self.assertIsInstance(info['aliases'], list)
|
|
self.assertIsInstance(info['supports_single_file'], bool)
|
|
|
|
def test_get_pipeline_info_unknown_raises(self):
|
|
"""Test that getting info for unknown pipeline raises ValueError."""
|
|
with self.assertRaises(ValueError) as ctx:
|
|
loader.get_pipeline_info("NonExistentPipeline")
|
|
|
|
self.assertIn("Unknown pipeline", str(ctx.exception))
|
|
|
|
def test_discover_diffusers_classes_pipelines(self):
|
|
"""Test generic class discovery for DiffusionPipeline."""
|
|
classes = loader.discover_diffusers_classes("DiffusionPipeline")
|
|
|
|
# Should return a dict
|
|
self.assertIsInstance(classes, dict)
|
|
|
|
# Should contain known pipeline classes
|
|
self.assertIn("DiffusionPipeline", classes)
|
|
self.assertIn("StableDiffusionPipeline", classes)
|
|
|
|
def test_discover_diffusers_classes_caches_results(self):
|
|
"""Test that class discovery results are cached."""
|
|
classes1 = loader.discover_diffusers_classes("DiffusionPipeline")
|
|
classes2 = loader.discover_diffusers_classes("DiffusionPipeline")
|
|
|
|
# Should be the same object (cached)
|
|
self.assertIs(classes1, classes2)
|
|
|
|
def test_discover_diffusers_classes_exclude_base(self):
|
|
"""Test discovering classes without base class."""
|
|
classes = loader.discover_diffusers_classes("DiffusionPipeline", include_base=False)
|
|
|
|
# Should still contain subclasses
|
|
self.assertIn("StableDiffusionPipeline", classes)
|
|
|
|
def test_get_available_classes(self):
|
|
"""Test getting list of available classes for a base class."""
|
|
classes = loader.get_available_classes("DiffusionPipeline")
|
|
|
|
# Should return a sorted list
|
|
self.assertIsInstance(classes, list)
|
|
self.assertEqual(classes, sorted(classes))
|
|
|
|
# Should contain known classes
|
|
self.assertIn("StableDiffusionPipeline", classes)
|
|
|
|
|
|
class TestDiffusersDynamicLoaderWithMocks(unittest.TestCase):
|
|
"""Test cases using mocks to test edge cases."""
|
|
|
|
def test_load_pipeline_requires_model_id(self):
|
|
"""Test that load_diffusers_pipeline requires model_id."""
|
|
with self.assertRaises(ValueError) as ctx:
|
|
loader.load_diffusers_pipeline(class_name="StableDiffusionPipeline")
|
|
|
|
self.assertIn("model_id is required", str(ctx.exception))
|
|
|
|
def test_resolve_with_model_id_uses_diffusion_pipeline_fallback(self):
|
|
"""Test that resolving with only model_id falls back to DiffusionPipeline."""
|
|
# When model_id is provided, if hub lookup is not successful,
|
|
# should fall back to DiffusionPipeline.
|
|
# This tests the fallback behavior - the actual hub lookup may succeed
|
|
# or fail depending on network, but the fallback path should work.
|
|
cls = loader.resolve_pipeline_class(model_id="some/nonexistent/model")
|
|
self.assertEqual(cls, DiffusionPipeline)
|
|
|
|
|
|
@unittest.skipUnless(GRPC_AVAILABLE, "gRPC modules not available")
|
|
class TestGenerateImageOptionsKwargsMerge(unittest.TestCase):
|
|
"""Test that GenerateImage merges the options dict into pipeline kwargs.
|
|
|
|
The options dict holds image (PIL), negative_prompt, and
|
|
num_inference_steps. Without the merge, img2img pipelines never
|
|
receive the source image and fail with 'Input is in incorrect format'.
|
|
"""
|
|
|
|
def test_options_merged_into_pipeline_kwargs(self):
|
|
from backend import BackendServicer
|
|
from PIL import Image
|
|
import tempfile, os
|
|
|
|
svc = BackendServicer.__new__(BackendServicer)
|
|
# Minimal attributes the method reads
|
|
svc.pipe = MagicMock()
|
|
svc.pipe.return_value.images = [Image.new("RGB", (4, 4))]
|
|
svc.cfg_scale = 7.5
|
|
svc.controlnet = None
|
|
svc.img2vid = False
|
|
svc.txt2vid = False
|
|
svc.clip_skip = 0
|
|
svc.PipelineType = "StableDiffusionImg2ImgPipeline"
|
|
svc.options = {}
|
|
|
|
# Create a tiny source image for the request's src field
|
|
src_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
|
Image.new("RGB", (4, 4), color="red").save(src_file, format="PNG")
|
|
src_file.close()
|
|
|
|
dst_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
|
|
dst_file.close()
|
|
|
|
try:
|
|
request = MagicMock()
|
|
request.positive_prompt = "a test prompt"
|
|
request.negative_prompt = "bad quality"
|
|
request.step = 10
|
|
request.seed = 0
|
|
request.width = 0
|
|
request.height = 0
|
|
request.src = src_file.name
|
|
request.ref_images = []
|
|
request.dst = dst_file.name
|
|
|
|
svc.GenerateImage(request, context=None)
|
|
|
|
# The pipeline must have been called with the image kwarg
|
|
svc.pipe.assert_called_once()
|
|
_, call_kwargs = svc.pipe.call_args
|
|
self.assertIn("image", call_kwargs,
|
|
"source image must be passed to pipeline via kwargs")
|
|
self.assertIn("negative_prompt", call_kwargs,
|
|
"negative_prompt must be passed to pipeline via kwargs")
|
|
self.assertEqual(call_kwargs["num_inference_steps"], 10)
|
|
finally:
|
|
os.unlink(src_file.name)
|
|
os.unlink(dst_file.name)
|
|
|
|
def test_text_to_image_prompt_is_passed_by_keyword(self):
|
|
"""Test compatibility with pipelines that take image before prompt."""
|
|
import os
|
|
import tempfile
|
|
|
|
from PIL import Image
|
|
|
|
from backend import BackendServicer
|
|
|
|
class Flux2CompatiblePipeline:
|
|
"""Model the FLUX.2 call signature: image is before prompt."""
|
|
|
|
def __call__(self, image=None, prompt=None, **kwargs):
|
|
if prompt is None:
|
|
raise ValueError("prompt was not passed by keyword")
|
|
self.prompt = prompt
|
|
self.kwargs = kwargs
|
|
return MagicMock(images=[Image.new("RGB", (4, 4))])
|
|
|
|
pipeline = Flux2CompatiblePipeline()
|
|
svc = BackendServicer.__new__(BackendServicer)
|
|
svc.pipe = pipeline
|
|
svc.cfg_scale = 7.5
|
|
svc.controlnet = None
|
|
svc.img2vid = False
|
|
svc.txt2vid = False
|
|
svc.clip_skip = 0
|
|
svc.PipelineType = "Flux2KleinPipeline"
|
|
svc.options = {}
|
|
|
|
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as dst_file:
|
|
dst_path = dst_file.name
|
|
|
|
try:
|
|
request = MagicMock()
|
|
request.positive_prompt = "a red apple on a wooden table"
|
|
request.negative_prompt = ""
|
|
request.step = 4
|
|
request.seed = 0
|
|
request.width = 0
|
|
request.height = 0
|
|
request.src = ""
|
|
request.ref_images = []
|
|
request.dst = dst_path
|
|
|
|
svc.GenerateImage(request, context=None)
|
|
|
|
self.assertEqual(pipeline.prompt, request.positive_prompt)
|
|
self.assertEqual(pipeline.kwargs["num_inference_steps"], 4)
|
|
finally:
|
|
os.unlink(dst_path)
|
|
|
|
|
|
class TestDeviceSelection(unittest.TestCase):
|
|
"""Unit tests for backend.select_device (no GPU required)."""
|
|
|
|
def test_autodetect_cuda(self):
|
|
self.assertEqual(backend.select_device(False, None, True, False, False), "cuda")
|
|
|
|
def test_cpu_fallback(self):
|
|
self.assertEqual(backend.select_device(False, None, False, False, False), "cpu")
|
|
|
|
def test_forced_cuda(self):
|
|
self.assertEqual(backend.select_device(True, None, False, False, False), "cuda")
|
|
|
|
def test_device_option_wins(self):
|
|
self.assertEqual(backend.select_device(True, "cpu", True, True, True), "cpu")
|
|
|
|
def test_mps_overrides(self):
|
|
self.assertEqual(backend.select_device(False, None, True, False, True), "mps")
|
|
|
|
|
|
class TestWritePcmWav(unittest.TestCase):
|
|
def test_writes_clipped_float_samples_as_mono_pcm(self):
|
|
from audio_utils import write_pcm_wav
|
|
|
|
destination = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
|
destination.close()
|
|
|
|
try:
|
|
write_pcm_wav(destination.name, [0.0, 0.5, -0.5, 2.0], 16000)
|
|
|
|
with wave.open(destination.name, "rb") as generated:
|
|
self.assertEqual(generated.getframerate(), 16000)
|
|
self.assertEqual(generated.getnchannels(), 1)
|
|
self.assertEqual(generated.getsampwidth(), 2)
|
|
self.assertEqual(generated.getnframes(), 4)
|
|
self.assertEqual(
|
|
generated.readframes(4),
|
|
b"\x00\x00\x00@\x00\xc0\xff\x7f",
|
|
)
|
|
finally:
|
|
os.unlink(destination.name)
|
|
|
|
|
|
@unittest.skipUnless(GRPC_AVAILABLE, "gRPC modules not available")
|
|
class TestSoundGeneration(unittest.TestCase):
|
|
def test_maps_request_options_and_writes_pipeline_audio(self):
|
|
from backend import BackendServicer
|
|
|
|
service = BackendServicer.__new__(BackendServicer)
|
|
service.options = {"num_inference_steps": 200.0}
|
|
service.pipe = MagicMock()
|
|
service.pipe.return_value.audios = [[0.0, 0.5, -0.5]]
|
|
service.pipe.vae.config.sampling_rate = 16000
|
|
|
|
destination = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
|
destination.close()
|
|
|
|
try:
|
|
request = backend_pb2.SoundGenerationRequest(
|
|
text="ocean waves",
|
|
dst=destination.name,
|
|
duration=2.5,
|
|
temperature=0,
|
|
)
|
|
|
|
result = service.SoundGeneration(request, context=None)
|
|
|
|
self.assertTrue(result.success, result.message)
|
|
service.pipe.assert_called_once_with(
|
|
num_inference_steps=200,
|
|
prompt="ocean waves",
|
|
audio_length_in_s=2.5,
|
|
guidance_scale=0,
|
|
)
|
|
with wave.open(destination.name, "rb") as generated:
|
|
self.assertEqual(generated.getframerate(), 16000)
|
|
self.assertEqual(generated.getnframes(), 3)
|
|
finally:
|
|
os.unlink(destination.name)
|