LocalAI/backend/python/diffusers/test.py

"""
A test script to test the gRPC service and dynamic loader
"""
import unittest
import subprocess
import time
from unittest.mock import patch, MagicMock

# Import dynamic loader for testing (these don't need gRPC)
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)