soulxpodcast 0.1.0__tar.gz

soulxpodcast-0.1.0/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: soulxpodcast
+Version: 0.1.0
+Summary: soulx podcast
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: accelerate==1.10.1
+Requires-Dist: diffusers==0.37.1
+Requires-Dist: einops==0.8.2
+Requires-Dist: librosa==0.11.0
+Requires-Dist: numpy==2.4.6
+Requires-Dist: onnxruntime==1.26.0
+Requires-Dist: onnxruntime-gpu==1.26.0
+Requires-Dist: s3tokenizer==0.3.0
+Requires-Dist: scipy==1.17.1
+Requires-Dist: sympy==1.13.1
+Requires-Dist: torch==2.5.1+cu121
+Requires-Dist: torchaudio==2.5.1+cu121
+Requires-Dist: torchvision==0.20.1+cu121
+Requires-Dist: transformers==4.57.1

soulxpodcast-0.1.0/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "soulxpodcast"
+version = "0.1.0"
+description = "soulx podcast"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "accelerate==1.10.1",
+    "diffusers==0.37.1",
+    "einops==0.8.2",
+    "librosa==0.11.0",
+    "numpy==2.4.6",
+    "onnxruntime==1.26.0",
+    "onnxruntime-gpu==1.26.0",
+    "s3tokenizer==0.3.0",
+    "scipy==1.17.1",
+    "sympy==1.13.1",
+    "torch==2.5.1+cu121",
+    "torchaudio==2.5.1+cu121",
+    "torchvision==0.20.1+cu121",
+    "transformers==4.57.1",
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu121"
+url = "https://download.pytorch.org/whl/cu121"
+explicit = true
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu121" }
+torchvision = { index = "pytorch-cu121" }
+torchaudio = { index = "pytorch-cu121" }

soulxpodcast-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

soulxpodcast-0.1.0/src/soulxpodcast/config.py

@@ -0,0 +1,142 @@
+import os
+from dataclasses import dataclass, field, fields, is_dataclass, asdict
+from typing import Any, Dict, List, Optional
+from pathlib import Path
+import json
+
+import torch
+from transformers import AutoConfig
+from transformers import PretrainedConfig
+
+
+@dataclass
+class SoulXPodcastLLMConfig:
+    architectures: list[str] = field(default_factory=lambda: ["Qwen3ForCausalLM"])
+    attention_dropout: float = 0.0
+    bos_token_id: int = 151643
+    eos_token_id: int = 151675  # speech eos
+    hidden_act: str = "silu"
+    hidden_size: int = 2048
+    initializer_range: float = 0.02
+    intermediate_size: int = 6144
+    max_position_embeddings: int = 40960
+    max_window_layers: int = 28
+    model_type: str = "qwen3"
+    num_attention_heads: int = 16
+    num_hidden_layers: int = 28
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    rms_norm_eps: float = 1e-06
+    rope_scaling: dict | None = None
+    rope_theta: float = 1000000.0
+    sliding_window: int = 32768
+    tie_word_embeddings: bool = True
+    torch_dtype: str = "bfloat16"
+    transformers_version: str = "4.52.3"
+    use_cache: bool = True
+    use_sliding_window: bool = False
+    vocab_size: int = 159488  # text_vocab_size + speech_vocab_size + 2 (eos and task_id)
+    lm_head_bias: bool = False
+    qkv_bias: bool = False
+    fp16_flow: bool = False
+    speech_token_offset: int = 152927
+
+    @classmethod
+    def from_initial_and_json(
+        cls, 
+        initial_values: Dict[str, Any] = None, 
+        json_file: Optional[str] = None
+    ):
+        """
+        Create an instance from initial values and JSON data.
+        
+        Args:
+            initial_values: Dictionary of initial values (highest priority)
+            json_file: Path to JSON file
+
+        Returns:
+            SoulXPodcastLLMConfig instance
+        """
+        # Merge all data sources
+        merged_data = {}
+        
+        # 1. Load from JSON file first (lowest priority)
+        if json_file and os.path.exists(json_file):
+            file_data = cls._load_json_file(json_file)
+            merged_data.update(file_data)
+        
+         # 2. Overwrite with initial values (highest priority)
+        if initial_values:
+            merged_data.update(initial_values)
+        
+        # Filter dataclass fields
+        valid_fields = {f.name for f in fields(cls)}
+        init_data = {k: v for k, v in merged_data.items() if k in valid_fields}
+        
+        return cls(**init_data)
+
+    @staticmethod
+    def _load_json_file(file_path: str) -> Dict[str, Any]:
+        """Load data from a JSON file"""
+        path = Path(file_path)
+        if not path.exists():
+            return {}
+        with open(path, 'r', encoding='utf-8') as f:
+            return json.load(f)
+
+class AutoPretrainedConfig(PretrainedConfig):
+    model_type = "qwen3"
+    
+    def __init__(self, **kwargs):
+        # Filter out non-configuration parameters
+        config_kwargs = {k: v for k, v in kwargs.items() 
+                        if not k.startswith('_') and k != 'self'}
+        super().__init__(**config_kwargs)
+    
+    @classmethod
+    def from_dataclass(cls, dataclass_config):
+        """Automatically create configuration from any dataclass"""
+        if not is_dataclass(dataclass_config):
+            raise ValueError("Input must be a dataclass instance")
+        
+        dataclass_dict = asdict(dataclass_config)
+        return cls(**dataclass_dict)
+
+
+@dataclass
+class SamplingParams:
+    temperature: float = 0.6
+    repetition_penalty: float = 1.25
+    top_k: int = 100
+    top_p: float = 0.9
+    min_tokens: int = 8
+    max_tokens: int = 3000
+    stop_token_ids: list[int] = field(default_factory=lambda: [151675])
+    # RasSampler parameters
+    use_ras: bool = True
+    win_size: int = 25
+    tau_r: float = 0.2
+
+
+@dataclass
+class Config:
+    model: str
+    max_model_len: int = 8192  # 15s prompt + 30s generated audio for 25hz audio tokenizer
+    gpu_memory_utilization: float = 0.9
+    tensor_parallel_size: int = 1
+    enforce_eager: bool = False
+    hf_config: SoulXPodcastLLMConfig | AutoConfig = field(default_factory=SoulXPodcastLLMConfig)
+    eos: int = -1
+    llm_engine: str = "hf" # support hf, nano-vllm
+    max_turn_size: int = 10
+    turn_tokens_threshold: int = 6192
+    
+    prompt_context: int = 2 # default to 2 for two-speaker podcast;
+    history_context: int = 2
+    history_text_context: int = 2
+    
+    def __post_init__(self):
+        assert os.path.isdir(self.model)
+
+        max_pos = getattr(self.hf_config, "max_position_embeddings", 8192)
+        self.max_model_len = min(self.max_model_len, max_pos)

soulxpodcast-0.1.0/src/soulxpodcast/engine/llm_engine.py

@@ -0,0 +1,114 @@
+import os
+import types
+import atexit
+from time import perf_counter
+from functools import partial
+from dataclasses import fields, asdict
+
+import torch
+import torch.multiprocessing as mp
+from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteriaList
+from transformers import EosTokenCriteria, RepetitionPenaltyLogitsProcessor
+try:    
+    from vllm import LLM
+    from vllm import SamplingParams as VllmSamplingParams
+    from vllm.inputs import TokensPrompt as TokensPrompt
+    SUPPORT_VLLM = True
+except ImportError:
+    SUPPORT_VLLM = False
+
+from soulxpodcast.config import Config, SamplingParams
+from soulxpodcast.models.modules.sampler import _ras_sample_hf_engine
+
+class HFLLMEngine:
+
+    def __init__(self, model, **kwargs):
+        config_fields = {field.name for field in fields(Config)}
+        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
+        config = Config(model, **config_kwargs)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model, use_fast=True)
+        config.eos = config.hf_config.eos_token_id # speech eos token;
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.model = AutoModelForCausalLM.from_pretrained(model, torch_dtype=torch.bfloat16, device_map=self.device)
+        self.config = config
+        self.pad_token_id = self.tokenizer.pad_token_id
+
+    def generate(
+        self,
+        prompt: list[str],
+        sampling_param: SamplingParams,
+        past_key_values=None,
+    ) -> dict:
+        
+        stopping_criteria = StoppingCriteriaList([EosTokenCriteria(eos_token_id=self.config.hf_config.eos_token_id)])
+        if sampling_param.use_ras:
+            sample_hf_engine_handler = partial(_ras_sample_hf_engine, 
+                    use_ras=sampling_param.use_ras, 
+                    win_size=sampling_param.win_size, tau_r=sampling_param.tau_r)
+        else:
+            sample_hf_engine_handler = None
+        rep_pen_processor = RepetitionPenaltyLogitsProcessor(
+            penalty=sampling_param.repetition_penalty,
+            prompt_ignore_length=len(prompt)
+        ) # exclude the input prompt, consistent with vLLM implementation;
+        with torch.no_grad(): 
+            input_len = len(prompt)
+            generated_ids = self.model.generate(
+                input_ids = torch.tensor([prompt], dtype=torch.int64).to(self.device),
+                do_sample=True,
+                top_k=sampling_param.top_k,
+                top_p=sampling_param.top_p,
+                min_new_tokens=sampling_param.min_tokens,
+                max_new_tokens=sampling_param.max_tokens,
+                temperature=sampling_param.temperature,
+                stopping_criteria=stopping_criteria,
+                past_key_values=past_key_values,
+                custom_generate=sample_hf_engine_handler,
+                use_cache=True,
+                logits_processor=[rep_pen_processor]
+            )
+            generated_ids = generated_ids[:, input_len:].cpu().numpy().tolist()[0]
+        output = {
+            "text": self.tokenizer.decode(generated_ids),
+            "token_ids": generated_ids,
+        }
+        return output
+
+class VLLMEngine:
+
+    def __init__(self, model, **kwargs):
+        
+        config_fields = {field.name for field in fields(Config)}
+        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
+        config = Config(model, **config_kwargs)
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
+        config.eos = config.hf_config.eos_token_id # speech eos token;
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        os.environ["VLLM_USE_V1"] = "0"
+        if SUPPORT_VLLM:
+            self.model = LLM(model=model, enforce_eager=True, dtype="bfloat16", max_model_len=8192, enable_prefix_caching=True,)
+        else:
+            raise ImportError("Not Support VLLM now!!!")
+        self.config = config
+        self.pad_token_id = self.tokenizer.pad_token_id
+
+    def generate(
+        self,
+        prompt: list[str],
+        sampling_param: SamplingParams,
+        past_key_values=None,
+    ) -> dict:
+        sampling_param.stop_token_ids = [self.config.hf_config.eos_token_id]
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                TokensPrompt(prompt_token_ids=prompt), 
+                VllmSamplingParams(**asdict(sampling_param)),
+                use_tqdm=False,
+            )[0].outputs[0].token_ids
+        output = {
+            "text": self.tokenizer.decode(generated_ids),
+            "token_ids": list(generated_ids),
+        }
+        return output

soulxpodcast-0.1.0/src/soulxpodcast/models/modules/flow.py

@@ -0,0 +1,197 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from soulxpodcast.models.modules.flow_components.estimator import \
+    CausalConditionalDecoder
+from soulxpodcast.models.modules.flow_components.upsample_encoder import (
+    UpsampleConformerEncoder, make_pad_mask)
+
+
+@dataclass
+class CfmParams:
+    sigma_min: float = 1e-6
+    solver: str = "euler"
+    t_scheduler: str = "cosine"
+    training_cfg_rate: float = 0.2
+    inference_cfg_rate: float = 0.7
+
+
+class CausalConditionalCFM(torch.nn.Module):
+    def __init__(self, in_channels=320, cfm_params=CfmParams(), n_spks=1, spk_emb_dim=80, estimator: torch.nn.Module = None):
+        super().__init__()
+        self.n_feats = in_channels
+        self.n_spks = n_spks
+        self.spk_emb_dim = spk_emb_dim
+        self.solver = cfm_params.solver
+        if hasattr(cfm_params, "sigma_min"):
+            self.sigma_min = cfm_params.sigma_min
+        else:
+            self.sigma_min = 1e-4
+        self.t_scheduler = cfm_params.t_scheduler
+        self.training_cfg_rate = cfm_params.training_cfg_rate
+        self.inference_cfg_rate = cfm_params.inference_cfg_rate
+        in_channels = in_channels + (spk_emb_dim if n_spks > 0 else 0)
+        # Just change the architecture of the estimator here
+        self.estimator = CausalConditionalDecoder() if estimator is None else estimator
+
+    @torch.inference_mode()
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, streaming=False):
+        """Forward diffusion
+
+        Args:
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            n_timesteps (int): number of diffusion steps
+            temperature (float, optional): temperature for scaling noise. Defaults to 1.0.
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+
+        Returns:
+            sample: generated mel-spectrogram
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature
+        # fix prompt and overlap part mu and z
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
+        if self.t_scheduler == 'cosine':
+            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond, streaming=streaming), None
+
+    def solve_euler(self, x, t_span, mu, mask, spks, cond, streaming=False):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+        """
+        batch_size = x.size(0)
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+
+        # I am storing this because I can later plot it by putting a debugger here and saving it to a file
+        # Or in future might add like a return_all_steps flag
+        sol = []
+
+        # Do not use concat, it may cause memory format changed and trt infer with wrong results!
+        # Create tensors with double batch size for CFG (conditional + unconditional)
+        x_in = torch.zeros([batch_size * 2, x.size(1), x.size(2)], device=x.device, dtype=x.dtype)
+        mask_in = torch.zeros([batch_size * 2, mask.size(1), mask.size(2)], device=x.device, dtype=x.dtype)
+        mu_in = torch.zeros([batch_size * 2, mu.size(1), mu.size(2)], device=x.device, dtype=x.dtype)
+        t_in = torch.zeros([batch_size * 2], device=x.device, dtype=x.dtype)
+        spks_in = torch.zeros([batch_size * 2, spks.size(1)], device=x.device, dtype=x.dtype)
+        cond_in = torch.zeros([batch_size * 2, cond.size(1), cond.size(2)], device=x.device, dtype=x.dtype)
+
+        for step in range(1, len(t_span)):
+            # Classifier-Free Guidance inference introduced in VoiceBox
+            # Copy conditional and unconditional input
+            x_in[:batch_size] = x
+            x_in[batch_size:] = x
+            mask_in[:batch_size] = mask
+            mask_in[batch_size:] = mask
+            mu_in[:batch_size] = mu
+            # Unconditional part remains 0
+            t_in.fill_(t)
+            spks_in[:batch_size] = spks
+            cond_in[:batch_size] = cond
+
+            dphi_dt = self.estimator(
+                x_in, mask_in,
+                mu_in, t_in,
+                spks_in,
+                cond_in,
+                streaming
+            )
+            dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [batch_size, batch_size], dim=0)
+            dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt)
+            x = x + dt * dphi_dt
+            t = t + dt
+            sol.append(x)
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+        return sol[-1].float()
+
+
+class CausalMaskedDiffWithXvec(torch.nn.Module):
+    def __init__(
+        self,
+        input_size: int = 512,
+        output_size: int = 80,
+        spk_embed_dim: int = 192,
+        output_type: str = "mel",
+        vocab_size: int = 6561,
+        input_frame_rate: int = 25,
+        token_mel_ratio: int = 2,
+        pre_lookahead_len: int = 3,
+        encoder: torch.nn.Module = None,
+        decoder: torch.nn.Module = None,
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.vocab_size = vocab_size
+        self.output_type = output_type
+        self.input_frame_rate = input_frame_rate
+        self.input_embedding = nn.Embedding(vocab_size, input_size)
+        self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
+        self.encoder = UpsampleConformerEncoder() if encoder is None else encoder
+        self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size)
+        self.decoder = CausalConditionalCFM() if decoder is None else decoder
+        self.token_mel_ratio = token_mel_ratio
+        self.pre_lookahead_len = pre_lookahead_len
+
+    @torch.inference_mode()
+    def forward(self,
+                token,
+                token_len,
+                prompt_feat,
+                prompt_feat_len,
+                embedding,
+                streaming,
+                finalize):
+        # xvec projection
+        embedding = F.normalize(embedding, dim=1)
+        embedding = self.spk_embed_affine_layer(embedding)
+
+        # concat text and prompt_text
+        mask = (~make_pad_mask(token_len, max_len=token.shape[1])).unsqueeze(-1).to(embedding)
+        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        # text encode
+        if finalize is True:
+            h, h_lengths = self.encoder(token, token_len, streaming=streaming)
+        else:
+            token, context = token[:, :-self.pre_lookahead_len], token[:, -self.pre_lookahead_len:]
+            h, h_lengths = self.encoder(token, token_len, context=context, streaming=streaming)
+        h = self.encoder_proj(h)
+
+        # get conditions
+        conds = torch.zeros_like(h, device=token.device)
+        for i, j in enumerate(prompt_feat_len):
+            conds[i, :j] = prompt_feat[i, :j]
+        conds = conds.transpose(1, 2)
+
+        h_lengths = h_lengths.sum(dim=-1).squeeze(dim=1)
+        mask = (~make_pad_mask(h_lengths, max_len=h.shape[1])).to(h)
+        feat, _ = self.decoder(
+            mu=h.transpose(1, 2).contiguous(),
+            mask=mask.unsqueeze(1),
+            spks=embedding,
+            cond=conds,
+            n_timesteps=15,
+            streaming=streaming
+        )  # [B, num_mels, T]
+        return feat.float(), h_lengths