path: root/app/tests/test_instruction_capabilities.py

"""Instruction-support and guidance regressions across the TTS backends.

Covers the capabilities the models actually implement (verified against
each backend's serving code) and what the clients send for them:
Breeze-TTS 2's recommended guidance strength with instructions, the
audio.cpp Qwen3-TTS variant split (CustomVoice reads instructions, the
Base cloner does not), the SGLang models that consume a separate style
instruction alongside their voice conditioning, and the Qwen demo's
CustomVoice instruction parameter.
"""

import io
import json
import tempfile
import unittest
import wave
from pathlib import Path
from unittest.mock import MagicMock, patch

from converter.clients import (
    BACKEND_QWEN, AudioCppTTSClient, QwenTTSClient,
    VOICE_MODE_CLONE, VOICE_MODE_CUSTOM, VOICE_MODE_DESIGN,
)
from converter.clients.audiocpp import (
    AUDIOCPP_FAMILY_BREEZE_TTS,
    AUDIOCPP_FAMILY_PROFILES,
    AUDIOCPP_VOICE_OPTIONAL,
    audiocpp_entry_supports_instructions,
    audiocpp_family_voice_policy,
)


_WAV_BYTES = b"RIFF\x18\x00\x00\x00WAVEfmt \x10\x00\x00\x00"


# ---------------------------------------------------------------------------
# Pure helpers
# ---------------------------------------------------------------------------

class VoicePolicyKnownFamiliesTests(unittest.TestCase):
    """Families whose verified policy must survive a stale local spec."""

    def test_breeze_is_tts_plus_clone_even_without_a_local_spec(self):
        # Remote Breeze entries against older local checkouts carry no
        # breeze_tts spec at all: the fallback keeps instructions-only
        # voice direction connectable instead of demanding a reference.
        self.assertEqual(
            audiocpp_family_voice_policy(AUDIOCPP_FAMILY_BREEZE_TTS),
            AUDIOCPP_VOICE_OPTIONAL)

    def test_vibevoice_accepts_reference_audio_despite_its_spec(self):
        # vibevoice.json declares only "tts", but the implementation
        # accepts reference audio: a mixed tts+clone family, so a picked
        # voice must not be silently dropped.
        self.assertEqual(audiocpp_family_voice_policy("vibevoice"),
                         AUDIOCPP_VOICE_OPTIONAL)


class EntryInstructionSupportTests(unittest.TestCase):
    """audiocpp_entry_supports_instructions: True/False/None per entry."""

    def test_breeze_supports_instructions(self):
        self.assertIs(
            audiocpp_entry_supports_instructions(
                AUDIOCPP_FAMILY_BREEZE_TTS, "tts", "Breeze-TTS-2-GGUF"),
            True)

    def test_qwen_customvoice_and_design_support_instructions(self):
        self.assertIs(
            audiocpp_entry_supports_instructions(
                "qwen3_tts", "tts", "Qwen3-TTS-12Hz-1.7B-CustomVoice-GGUF"),
            True)
        self.assertIs(
            audiocpp_entry_supports_instructions(
                "qwen3_tts", "vdes", "Qwen3-TTS-12Hz-1.7B-VoiceDesign"),
            True)

    def test_qwen_base_cloner_provably_does_not(self):
        self.assertIs(
            audiocpp_entry_supports_instructions(
                "qwen3_tts", "tts", "Qwen3-TTS-12Hz-1.7B-Base-GGUF"),
            False)

    def test_unknown_families_are_unknown_not_unsupported(self):
        self.assertIsNone(
            audiocpp_entry_supports_instructions("chatterbox",
                                                 "tts", "x"))


# ---------------------------------------------------------------------------
# audio.cpp request payloads (instructed Breeze runs)
# ---------------------------------------------------------------------------

def _breeze_client(instructions=None, request_options=None,
                   voice="narrator", seed=-1):
    """A fully-initialized Breeze client (no HTTP machinery touched)."""
    with patch.object(AudioCppTTSClient, "_connect"):
        client = AudioCppTTSClient(
            Path("."), voice=voice, instructions=instructions,
            request_options=request_options)
    client.api_url = "http://127.0.0.1:8080"
    client.model_id = "Breeze-TTS-2-GGUF"
    client.family = AUDIOCPP_FAMILY_BREEZE_TTS
    client.task = "tts"
    client.profile = AUDIOCPP_FAMILY_PROFILES[AUDIOCPP_FAMILY_BREEZE_TTS]
    client.design_mode = False
    client.instruction_voice = False
    client.plain_mode = False
    client.preset_mode = True
    client.speaker_mode = False
    client._seed = seed
    client._resolve_auto_guidance()
    return client


def _captured_payload(client):
    """The JSON body _request_wav sends, via a stubbed urlopen."""
    response = MagicMock()
    response.read.return_value = _WAV_BYTES
    response.__enter__ = lambda self: response
    response.__exit__ = lambda self, *exc: None
    with patch("converter.clients.audiocpp.urllib.request.urlopen") \
            as urlopen:
        urlopen.return_value = response
        client._request_wav("Hello there.")
    request = urlopen.call_args[0][0]
    return json.loads(request.data.decode("utf-8"))


class BreezeGuidanceDefaultTests(unittest.TestCase):
    """Breeze guidance: recommended 4 with instructions, otherwise none."""

    def test_instructed_clone_carries_guidance_4_and_the_instruction(self):
        client = _breeze_client(instructions="Screaming, crazed, yelling")
        self.assertEqual(client._auto_guidance_scale, 4.0)
        payload = _captured_payload(client)
        self.assertEqual(payload["guidance_scale"], 4.0)
        self.assertEqual(payload["voice"], "narrator")
        self.assertEqual(
            payload["options"],
            {"instruction": "Screaming, crazed, yelling"})
        self.assertNotIn("instructions", payload)

    def test_option_instruction_also_gets_the_guidance_default(self):
        client = _breeze_client(request_options={
            "instruction": "Read slowly and warmly."})
        self.assertEqual(client._auto_guidance_scale, 4.0)
        payload = _captured_payload(client)
        self.assertEqual(payload["guidance_scale"], 4.0)
        self.assertEqual(payload["options"]["instruction"],
                         "Read slowly and warmly.")

    def test_explicit_guidance_option_is_preserved(self):
        client = _breeze_client(
            instructions="Screaming, crazed, yelling",
            request_options={"guidance_scale": "2.5"})
        self.assertIsNone(client._auto_guidance_scale)
        payload = _captured_payload(client)
        self.assertNotIn("guidance_scale", payload)
        self.assertEqual(payload["options"]["guidance_scale"], "2.5")

    def test_guidance_0_override_still_counts_as_explicit(self):
        # 0 selects the instruction-free branch: a deliberate setting.
        client = _breeze_client(
            instructions="Screaming",
            request_options={"guidance_scale": "0"})
        payload = _captured_payload(client)
        self.assertNotIn("guidance_scale", payload)

    def test_plain_clone_without_instructions_uses_the_backend_default(self):
        client = _breeze_client()
        self.assertIsNone(client._auto_guidance_scale)
        payload = _captured_payload(client)
        self.assertNotIn("guidance_scale", payload)
        self.assertNotIn("options", payload)
        self.assertEqual(payload["voice"], "narrator")


class InstructionConflictTests(unittest.TestCase):
    """Two different instruction sources are refused before connecting."""

    def test_conflicting_instructions_and_option_raise_without_a_server(self):
        with self.assertRaises(RuntimeError) as ctx:
            _breeze_client(instructions="calm narration",
                           request_options={"instruction": "screaming"})
        self.assertIn("Two conflicting instructions",
                      str(ctx.exception))

    def test_identical_instructions_from_both_sources_are_accepted(self):
        client = _breeze_client(
            instructions="calm narration",
            request_options={"instruction": "calm narration"})
        self.assertEqual(client.instructions, "calm narration")

    def test_option_only_instruction_is_folded_into_the_reports(self):
        client = _breeze_client(
            request_options={"instruction": "calm narration"})
        self.assertEqual(client.instructions, "calm narration")


class SeedPrecisionTests(unittest.TestCase):
    """Full-range uint64 seeds travel as decimal strings (audio.cpp docs)."""

    def test_seed_above_2_pow_53_is_sent_as_a_string(self):
        seed = 2 ** 53 + 3  # beyond the exact JSON-number integer range
        client = _breeze_client(seed=seed)
        payload = _captured_payload(client)
        self.assertEqual(payload["seed"], str(seed))

    def test_ordinary_seeds_stay_numbers(self):
        client = _breeze_client(seed=42)
        payload = _captured_payload(client)
        self.assertEqual(payload["seed"], 42)


# ---------------------------------------------------------------------------
# SGLang-Omni: instructions on supported pipelines
# ---------------------------------------------------------------------------

class SgOmniInstructionTests(unittest.TestCase):
    """instructions reach the payload only where the serving code reads it."""

    _tmp_dir = None
    _REF = None

    @classmethod
    def setUpClass(cls):
        buffer = io.BytesIO()
        with wave.open(buffer, "wb") as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(24000)
            wav_file.writeframes(b"\x01\x00" * 16)
        cls._tmp_dir = tempfile.TemporaryDirectory()
        cls._REF = Path(cls._tmp_dir.name) / "narrator.wav"
        cls._REF.write_bytes(buffer.getvalue())
        cls.addClassCleanup(cls._tmp_dir.cleanup)

    def _client(self, model, **kwargs):
        from converter.clients import SgOmniTTSClient
        with patch.object(SgOmniTTSClient, "_connect"):
            client = SgOmniTTSClient(
                Path("."), model=model, ref_audio=str(self._REF),
                ref_text="Hello transcript.", instructions="screaming",
                **kwargs)
        entry = client.entry
        payload = client._request_payload("Hello there.")
        return entry, payload

    def test_qwen_base_clone_carries_ref_and_instruction(self):
        entry, payload = self._client("qwen3_tts_1_7b_base")
        self.assertIn("ref_audio", payload)
        self.assertNotEqual(payload.get("task_type"), "VoiceDesign")
        self.assertEqual(payload["instructions"], "screaming")

    def test_moss_clone_carries_ref_and_instruction(self):
        entry, payload = self._client("moss_tts")
        self.assertIn("ref_audio", payload)
        self.assertEqual(payload["instructions"], "screaming")

    def test_customvoice_speaker_with_instruction(self):
        entry, payload = self._client("qwen3_tts_0_6b_customvoice",
                                      voice="Vivian")
        self.assertEqual(payload["voice"], "Vivian")
        self.assertEqual(payload["instructions"], "screaming")
        self.assertNotIn("task_type", payload)

    def test_design_remains_voice_design_with_instruction(self):
        from converter.clients import SgOmniTTSClient
        with patch.object(SgOmniTTSClient, "_connect"):
            client = SgOmniTTSClient(
                Path("."), model="qwen3_tts_1_7b_voicedesign",
                instructions="a warm narrator")
        payload = client._request_payload("Hello there.")
        self.assertEqual(payload["task_type"], "VoiceDesign")
        self.assertEqual(payload["instructions"], "a warm narrator")
        self.assertNotIn("ref_audio", payload)

    def test_unsupported_model_refuses_instructions_at_connect(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client("higgs_audio_v3_tts")
        self.assertIn("does not consume style instructions",
                      str(ctx.exception))


# ---------------------------------------------------------------------------
# Qwen demo: CustomVoice instruction parameter
# ---------------------------------------------------------------------------

class QwenCustomVoiceInstructionTests(unittest.TestCase):
    """The run_instruct endpoint takes an ``instruct`` delivery control."""

    def test_run_instruct_sends_instruct_alongside_the_speaker(self):
        client = QwenTTSClient.__new__(QwenTTSClient)
        client.voice_mode = VOICE_MODE_CUSTOM
        client.speaker = "Vivian"
        client.language = "Auto"
        client.instructions = "screaming, crazed"
        client._seed = -1
        client.client = MagicMock()
        client._resolve_api_name = lambda *names: names[0]
        client._endpoint_accepts_param = MagicMock(return_value=True)
        client._generate_custom_voice("Hello there.")
        predict = client.client.predict
        predict.assert_called_once_with(
            text="Hello there.", lang_disp="Auto",
            spk_disp="Vivian", instruct="screaming, crazed",
            api_name="/run_instruct")

    def test_custom_voice_without_instructions_is_unchanged(self):
        client = QwenTTSClient.__new__(QwenTTSClient)
        client.voice_mode = VOICE_MODE_CUSTOM
        client.speaker = "Vivian"
        client.language = "Auto"
        client.instructions = ""
        client._seed = -1
        client.client = MagicMock()
        client._resolve_api_name = lambda *names: names[0]
        client._endpoint_accepts_param = MagicMock(return_value=True)
        client._generate_custom_voice("Hello there.")
        _, kwargs = client.client.predict.call_args
        self.assertNotIn("instruct", kwargs)


# ---------------------------------------------------------------------------
# audiobook voice-mode routing (qwen)
# ---------------------------------------------------------------------------

class CatalogInstructionFlagsTests(unittest.TestCase):
    """Only the verified pipelines carry supports_instructions."""

    def test_catalog_marks_only_the_verified_pipelines(self):
        from backends.sglomni.catalog import ENTRIES
        supported = {"qwen3_tts_0_6b_customvoice", "qwen3_tts_0_6b_base",
                     "qwen3_tts_1_7b_base", "qwen3_tts_1_7b_voicedesign",
                     "moss_tts", "moss_tts_local"}
        for entry in ENTRIES:
            with self.subTest(entry=entry.key):
                self.assertEqual(entry.supports_instructions,
                                 entry.key in supported)


class GradioPrefixProbeTests(unittest.TestCase):
    """Qwen demos under modern Gradio sit behind /gradio_api."""

    def _identify(self, modern_payload, legacy_payload=None):
        import backends.probe as probe
        seen = []

        def fake_get_json(url, timeout):
            seen.append(url)
            if url == "http://x/gradio_api/info":
                return modern_payload
            if url == "http://x/info":
                return legacy_payload
            return None

        with patch.object(probe, "_get_json", side_effect=fake_get_json):
            with patch.object(probe.common, "server_running",
                              return_value=True):
                identity = probe._identify_gradio("http://x", 1.0)
        return identity, seen

    def test_modern_prefix_is_probed_first_and_identifies(self):
        payload = {"named_endpoints": {"/run_instruct": {}}}
        identity, seen = self._identify(payload)
        self.assertEqual(identity, probe_identity("qwen-custom"))
        self.assertEqual(seen, ["http://x/gradio_api/info"])

    def test_legacy_info_still_identifies_older_gradio(self):
        payload = {"named_endpoints": {"/run_voice_clone": {}}}
        identity, seen = self._identify(None, payload)
        self.assertEqual(identity, probe_identity("qwen-clone"))
        self.assertEqual(seen, ["http://x/gradio_api/info",
                                "http://x/info"])

    def test_neither_prefix_answers_none(self):
        identity, _ = self._identify(None)
        self.assertIsNone(identity)


def probe_identity(name):
    """The probe's IDENTITY_* constant for a backend NAME (local import)."""
    import backends.probe as probe
    return {"qwen-custom": probe.IDENTITY_QWEN_CUSTOM,
            "qwen-clone": probe.IDENTITY_QWEN_CLONE,
            }[name]


# ---------------------------------------------------------------------------
# audiobook voice-mode routing (qwen)
# ---------------------------------------------------------------------------

class QwenVoiceModeRoutingTests(unittest.TestCase):
    """speaker + instructions is a directed CustomVoice run, not Design."""

    def test_voice_mode_for_qwen_combinations(self):
        from converter.converter import voice_mode_for
        cases = [
            (dict(voice=None, clone=None, instructions=None),
             VOICE_MODE_CUSTOM),
            (dict(voice=None, clone=None, instructions="screaming"),
             VOICE_MODE_DESIGN),
            (dict(voice="Vivian", clone=None, instructions="screaming"),
             VOICE_MODE_CUSTOM),
            (dict(voice=None, clone="ref.wav", instructions=None),
             VOICE_MODE_CLONE),
        ]
        for kwargs, expected in cases:
            with self.subTest(**kwargs):
                self.assertEqual(
                    voice_mode_for(BACKEND_QWEN, voice=kwargs["voice"],
                                   clone=kwargs["clone"],
                                   instructions=kwargs["instructions"]),
                    expected)


if __name__ == "__main__":
    unittest.main()