Upload folder using huggingface_hub

__init__.py

@@ -0,0 +1 @@
+"""Zen Translator test suite."""

conftest.py

@@ -0,0 +1,60 @@
+"""Pytest configuration and fixtures."""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+
+@pytest.fixture
+def sample_audio():
+    """Generate sample audio data for testing."""
+    # 3 seconds of audio at 16kHz
+    duration_seconds = 3.0
+    sample_rate = 16000
+    samples = int(duration_seconds * sample_rate)
+
+    # Generate a simple sine wave
+    t = np.linspace(0, duration_seconds, samples)
+    audio = np.sin(2 * np.pi * 440 * t).astype(np.float32)
+
+    return audio, sample_rate
+
+
+@pytest.fixture
+def sample_video_frame():
+    """Generate sample video frame for testing."""
+    # RGB frame 256x256
+    frame = np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8)
+    return frame
+
+
+@pytest.fixture
+def temp_audio_file(tmp_path, sample_audio):
+    """Create a temporary audio file."""
+    import soundfile as sf
+
+    audio, sr = sample_audio
+    audio_path = tmp_path / "test_audio.wav"
+    sf.write(str(audio_path), audio, sr)
+
+    return audio_path
+
+
+@pytest.fixture
+def translator_config():
+    """Create test translator configuration."""
+    from zen_translator.config import TranslatorConfig
+
+    return TranslatorConfig(
+        device="cpu",
+        dtype="float32",
+        enable_lip_sync=False,  # Disable for faster tests
+        use_flash_attention=False,
+    )
+
+
+@pytest.fixture
+def test_data_dir():
+    """Get test data directory."""
+    return Path(__file__).parent / "data"

test_config.py

@@ -0,0 +1,78 @@
+"""Tests for configuration module."""
+
+
+class TestTranslatorConfig:
+    """Tests for TranslatorConfig."""
+
+    def test_default_config(self):
+        """Test default configuration values."""
+        from zen_translator.config import TranslatorConfig
+
+        config = TranslatorConfig()
+
+        assert config.target_language == "en"
+        assert config.device == "cuda"
+        assert config.dtype == "bfloat16"
+        assert config.enable_lip_sync is True
+        assert config.voice_reference_seconds == 3.0
+
+    def test_config_from_env(self, monkeypatch):
+        """Test configuration from environment variables."""
+        from zen_translator.config import TranslatorConfig
+
+        monkeypatch.setenv("ZEN_TRANSLATOR_TARGET_LANGUAGE", "es")
+        monkeypatch.setenv("ZEN_TRANSLATOR_DEVICE", "cpu")
+
+        config = TranslatorConfig()
+
+        assert config.target_language == "es"
+        assert config.device == "cpu"
+
+    def test_supported_languages(self):
+        """Test supported language lists."""
+        from zen_translator.config import TranslatorConfig
+
+        config = TranslatorConfig()
+
+        # Check input languages
+        assert "en" in config.supported_input_languages
+        assert "zh" in config.supported_input_languages
+        assert "ja" in config.supported_input_languages
+        assert "yue" in config.supported_input_languages  # Cantonese
+
+        # Check output languages
+        assert "en" in config.supported_output_languages
+        assert "zh" in config.supported_output_languages
+        assert len(config.supported_output_languages) == 10
+
+    def test_lip_sync_quality_options(self):
+        """Test lip sync quality options."""
+        from zen_translator.config import TranslatorConfig
+
+        for quality in ["fast", "balanced", "quality"]:
+            config = TranslatorConfig(lip_sync_quality=quality)
+            assert config.lip_sync_quality == quality
+
+
+class TestNewsAnchorConfig:
+    """Tests for NewsAnchorConfig."""
+
+    def test_default_config(self):
+        """Test default news anchor config."""
+        from zen_translator.config import NewsAnchorConfig
+
+        config = NewsAnchorConfig()
+
+        assert config.min_clip_duration == 5.0
+        assert config.max_clip_duration == 30.0
+        assert len(config.target_anchors) > 0
+
+    def test_training_settings(self):
+        """Test training hyperparameters."""
+        from zen_translator.config import NewsAnchorConfig
+
+        config = NewsAnchorConfig()
+
+        assert config.batch_size == 4
+        assert config.learning_rate == 2e-5
+        assert config.num_epochs == 3

test_pipeline.py

@@ -0,0 +1,66 @@
+"""Tests for translation pipeline."""
+
+
+class TestTranslationPipeline:
+    """Tests for TranslationPipeline."""
+
+    def test_pipeline_initialization(self, translator_config):
+        """Test pipeline can be initialized."""
+        from zen_translator.pipeline import TranslationPipeline
+
+        pipeline = TranslationPipeline(translator_config)
+
+        assert pipeline.config == translator_config
+        assert pipeline.translator is not None
+        assert pipeline.voice_cloner is not None
+        assert pipeline._loaded is False
+
+    def test_get_supported_languages(self, translator_config):
+        """Test getting supported languages."""
+        from zen_translator.pipeline import TranslationPipeline
+
+        pipeline = TranslationPipeline(translator_config)
+        languages = pipeline.get_supported_languages()
+
+        assert "input" in languages
+        assert "output" in languages
+        assert len(languages["input"]) >= 18
+        assert len(languages["output"]) == 10
+
+
+class TestBatchTranslationPipeline:
+    """Tests for BatchTranslationPipeline."""
+
+    def test_batch_pipeline_initialization(self, translator_config):
+        """Test batch pipeline can be initialized."""
+        from zen_translator.pipeline import BatchTranslationPipeline
+
+        pipeline = BatchTranslationPipeline(translator_config)
+
+        assert pipeline.config == translator_config
+
+
+class TestPipelineConfig:
+    """Tests for pipeline configuration options."""
+
+    def test_default_config(self):
+        """Test default pipeline configuration."""
+        from zen_translator import TranslatorConfig
+
+        config = TranslatorConfig()
+
+        assert config.qwen3_omni_model == "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+        assert config.cosyvoice_model == "FunAudioLLM/CosyVoice2-0.5B"
+        assert config.wav2lip_model == "numz/wav2lip_studio"
+
+    def test_custom_model_paths(self):
+        """Test custom model path configuration."""
+        from zen_translator import TranslatorConfig
+
+        config = TranslatorConfig(
+            qwen3_omni_model="./local/qwen3-omni",
+            cosyvoice_model="./local/cosyvoice",
+        )
+
+        assert config.qwen3_omni_model == "./local/qwen3-omni"
+        assert config.cosyvoice_model == "./local/cosyvoice"

test_training.py

@@ -0,0 +1,136 @@
+"""Tests for training infrastructure."""
+
+
+class TestSwiftConfig:
+    """Tests for ms-swift training configuration."""
+
+    def test_default_config(self):
+        """Test default training config."""
+        from zen_translator.training import SwiftTrainingConfig
+
+        config = SwiftTrainingConfig()
+
+        assert config.model_type == "qwen3-omni"
+        assert config.train_type == "lora"
+        assert config.lora_rank == 64
+        assert config.lora_alpha == 128
+
+    def test_to_swift_args(self):
+        """Test conversion to swift CLI arguments."""
+        from zen_translator.training import SwiftTrainingConfig
+
+        config = SwiftTrainingConfig()
+        args = config.to_swift_args()
+
+        assert "--model_type=qwen3-omni" in args
+        assert "--train_type=lora" in args
+        assert "--lora_rank=64" in args
+
+    def test_to_yaml(self, tmp_path):
+        """Test YAML export."""
+        from zen_translator.training import SwiftTrainingConfig
+
+        config = SwiftTrainingConfig()
+        yaml_path = tmp_path / "config.yaml"
+
+        config.to_yaml(yaml_path)
+
+        assert yaml_path.exists()
+
+        # Verify content
+        import yaml
+
+        with open(yaml_path) as f:
+            saved = yaml.safe_load(f)
+
+        assert saved["model"]["type"] == "qwen3-omni"
+        assert saved["lora"]["rank"] == 64
+
+
+class TestZenIdentityConfig:
+    """Tests for Zen identity finetuning config."""
+
+    def test_identity_system_prompt(self):
+        """Test identity system prompt is set."""
+        from zen_translator.training import ZenIdentityConfig
+
+        config = ZenIdentityConfig()
+
+        assert "Zen Translator" in config.system_prompt
+        assert "Hanzo AI" in config.system_prompt
+
+
+class TestNewsAnchorConfig:
+    """Tests for news anchor training config."""
+
+    def test_anchor_names(self):
+        """Test anchor names are configured."""
+        from zen_translator.training import NewsAnchorConfig
+
+        config = NewsAnchorConfig()
+
+        assert len(config.anchor_names) > 0
+        assert "cnn" in config.anchor_names
+        assert "bbc" in config.anchor_names
+
+    def test_news_domains(self):
+        """Test news domains are configured."""
+        from zen_translator.training import NewsAnchorConfig
+
+        config = NewsAnchorConfig()
+
+        assert "politics" in config.news_domains
+        assert "technology" in config.news_domains
+
+
+class TestNewsChannels:
+    """Tests for predefined news channels."""
+
+    def test_channels_defined(self):
+        """Test news channels are defined."""
+        from zen_translator.training import NEWS_CHANNELS
+
+        assert len(NEWS_CHANNELS) > 0
+        assert "cnn" in NEWS_CHANNELS
+        assert "bbc" in NEWS_CHANNELS
+        assert "nhk" in NEWS_CHANNELS
+
+    def test_channel_urls(self):
+        """Test channel URLs are valid."""
+        from zen_translator.training import NEWS_CHANNELS
+
+        for name, url in NEWS_CHANNELS.items():
+            assert url.startswith("https://")
+            assert "youtube.com" in url
+
+
+class TestCreateTrainingDataset:
+    """Tests for dataset creation."""
+
+    def test_create_jsonl_dataset(self, tmp_path):
+        """Test JSONL dataset creation."""
+        from zen_translator.training import create_training_dataset
+
+        conversations = [
+            {
+                "conversations": [
+                    {"role": "user", "content": "Hello"},
+                    {"role": "assistant", "content": "Hi there!"},
+                ]
+            }
+        ]
+
+        output_path = tmp_path / "train.jsonl"
+        create_training_dataset(conversations, output_path, format="jsonl")
+
+        assert output_path.exists()
+
+        # Verify content
+        import json
+
+        with open(output_path) as f:
+            lines = f.readlines()
+
+        assert len(lines) == 1
+        data = json.loads(lines[0])
+        assert "conversations" in data

test_wav2lip_model.py

@@ -0,0 +1,134 @@
+"""Tests for Wav2Lip model architecture."""
+
+import torch
+
+
+class TestWav2LipModel:
+    """Tests for Wav2Lip neural network."""
+
+    def test_model_initialization(self):
+        """Test model can be initialized."""
+        from zen_translator.lip_sync.wav2lip_model import Wav2Lip
+
+        model = Wav2Lip()
+
+        assert model.audio_encoder is not None
+        assert model.face_encoder is not None
+        assert model.face_decoder is not None
+
+    def test_model_forward_shape(self):
+        """Test model forward pass produces correct output shape."""
+        from zen_translator.lip_sync.wav2lip_model import Wav2Lip
+
+        model = Wav2Lip()
+        model.eval()
+
+        # Create dummy inputs
+        batch_size = 2
+        mel_length = 16
+        mel_channels = 80
+
+        # Audio: (B, T, 1, 80, 16) -> mel spectrogram windows
+        audio = torch.randn(batch_size, 1, 1, mel_channels, mel_length)
+
+        # Face: (B, 6, 96, 96) -> half face + reference
+        face = torch.randn(batch_size, 6, 96, 96)
+
+        with torch.no_grad():
+            output = model(audio, face)
+
+        # Output should be (B, 3, 96, 96)
+        assert output.shape == (batch_size, 3, 96, 96)
+
+    def test_audio_encoder(self):
+        """Test audio encoder produces correct embedding."""
+        from zen_translator.lip_sync.wav2lip_model import AudioEncoder
+
+        encoder = AudioEncoder()
+        encoder.eval()
+
+        batch_size = 2
+        audio = torch.randn(batch_size, 1, 1, 80, 16)
+
+        with torch.no_grad():
+            embedding = encoder(audio)
+
+        # Should produce 512-dim embedding
+        assert embedding.shape[-3] == 512
+
+    def test_face_encoder(self):
+        """Test face encoder produces feature hierarchy."""
+        from zen_translator.lip_sync.wav2lip_model import FaceEncoder
+
+        encoder = FaceEncoder()
+        encoder.eval()
+
+        batch_size = 2
+        face = torch.randn(batch_size, 6, 96, 96)
+
+        with torch.no_grad():
+            features = encoder(face)
+
+        # Should produce 7 feature maps (one per block)
+        assert len(features) == 7
+
+
+class TestConvBlocks:
+    """Tests for convolution building blocks."""
+
+    def test_conv2d_block(self):
+        """Test Conv2d block."""
+        from zen_translator.lip_sync.wav2lip_model import Conv2d
+
+        block = Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
+        x = torch.randn(1, 3, 64, 64)
+
+        out = block(x)
+
+        assert out.shape == (1, 32, 64, 64)
+
+    def test_conv2d_residual(self):
+        """Test Conv2d with residual connection."""
+        from zen_translator.lip_sync.wav2lip_model import Conv2d
+
+        block = Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True)
+        x = torch.randn(1, 32, 64, 64)
+
+        out = block(x)
+
+        # With residual, output should be different from non-residual
+        assert out.shape == (1, 32, 64, 64)
+
+    def test_transpose_conv2d(self):
+        """Test ConvTranspose2d block."""
+        from zen_translator.lip_sync.wav2lip_model import ConvTranspose2d
+
+        block = ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1)
+        x = torch.randn(1, 32, 32, 32)
+
+        out = block(x)
+
+        # Should upsample by factor of 2
+        assert out.shape == (1, 16, 64, 64)
+
+
+class TestSyncDiscriminator:
+    """Tests for sync discriminator."""
+
+    def test_discriminator_output(self):
+        """Test sync discriminator produces probability."""
+        from zen_translator.lip_sync.wav2lip_model import SyncDiscriminator
+
+        discriminator = SyncDiscriminator()
+        discriminator.eval()
+
+        batch_size = 2
+        mel = torch.randn(batch_size, 80, 16)
+        face = torch.randn(batch_size, 3, 96, 96)
+
+        with torch.no_grad():
+            output = discriminator(mel, face)
+
+        # Should produce sync probability
+        assert output.shape == (batch_size, 1)
+        assert torch.all(output >= 0) and torch.all(output <= 1)