nodes/BobRandomNumber-ComfyUI-DiaTTS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
BobRandomNumber
2025-04-26 02:38:21 -04:00
committed by GitHub
parent 013d2b9d4b
commit 9ee960e77e
10 changed files with 1951 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+60
View File
@@ -0,0 +1,60 @@
# ComfyUI DiaTest TTS Node
This node pack integrates the [Nari Labs Dia](https://github.com/nari-labs/dia) text-to-speech model into ComfyUI using a single node for loading (onto GPU) and generation.
Dia allows generating dialogue with speaker tags (`[S1]`, `[S2]`) and non-verbal sounds (`(laughs)`, etc.). This node loads the model from Hugging Face Hub and generates audio directly using float32 precision. It requires a CUDA-enabled GPU.
## Installation
1. Ensure you have a CUDA-enabled GPU and the necessary NVIDIA drivers installed.
2. Navigate to your `ComfyUI/custom_nodes/` directory.
3. Clone this repository:
```bash
git clone <repository_url> ComfyUI-DiaTest
```
(Replace `<repository_url>` with the actual URL).
Alternatively, download the ZIP and extract it into `custom_nodes` as `ComfyUI-DiaTest`.
4. Install the required dependencies:
* Activate ComfyUI's Python environment (e.g., `source ./venv/bin/activate`).
* Navigate to the node directory: `cd ComfyUI/custom_nodes/ComfyUI-DiaTest`
* Install requirements: `pip install -r requirements.txt`
5. Restart ComfyUI.
## Node
### Dia TTS Generate (`DiaGenerate`)
Loads the specified Dia model from Hugging Face Hub onto the GPU (if not already cached) and generates audio based on the provided text and parameters. Uses float32 precision internally. Requires CUDA.
**Inputs:**
* `repo_id`: The Hugging Face repository ID (default: `nari-labs/Dia-1.6B`).
* `text`: The main text transcript to generate audio for. Use `[S1]`, `[S2]` for speaker turns and parentheses for non-verbals like `(laughs)`.
* `max_tokens`: Maximum number of audio tokens to generate (controls length).
* `cfg_scale`: Classifier-Free Guidance scale (higher means stronger adherence to text).
* `temperature`: Sampling temperature (higher means more randomness).
* `top_p`: Nucleus sampling probability.
* `cfg_filter_top_k`: Top-K filtering applied during CFG.
* `speed_factor`: Adjusts the speed of the generated audio (1.0 = original speed, <1.0 = slower, >1.0 = faster). Default: 0.94.
* `seed`: Random seed for reproducibility.
**Outputs:**
* `audio`: The generated audio (`AUDIO` format: `{'waveform': tensor[B,C,T], 'sample_rate': sr}`), ready to be saved or previewed. Sample rate is 44100 Hz.
## Usage Example
1. Add the `Dia TTS Generate` node from the `audio/DiaTest` category.
2. Configure the `repo_id` if needed.
3. Enter your dialogue script into the `text` input.
4. Adjust generation parameters (`cfg_scale`, `temperature`, `speed_factor`, etc.) as needed.
5. Connect the `audio` output to a `SaveAudio` or `PreviewAudio` node.
6. Queue the prompt.
## Notes
* This node **requires a CUDA-enabled GPU**. It will fail to load if CUDA is not detected.
* The first time you run the node for a specific `repo_id`, it will download the model files from Hugging Face Hub, which may take some time. Subsequent runs will use the cached model.
* Changing the `repo_id` will trigger a model reload.
* The model uses float32 precision internally.
* The Descript Audio Codec (DAC) dependency (`descript-audio-codec`) must be installed via `requirements.txt`.
__init__.py
+9
View File
@@ -0,0 +1,9 @@
# ComfyUI-DiaTest/__init__.py
# Import the final mappings from nodes.py
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
# Expose them for ComfyUI
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
print("### Loading: ComfyUI-DiaTest Nodes ###")
dia_lib/__init__.py
View File
dia_lib/audio.py
+185
View File
@@ -0,0 +1,185 @@
import typing as tp
import torch
def build_delay_indices(B: int, T: int, C: int, delay_pattern: tp.List[int]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
"""
Precompute (t_idx_BxTxC, indices_BTCx3) so that out[t, c] = in[t - delay[c], c].
Negative t_idx => BOS; t_idx >= T => PAD.
"""
delay_arr = torch.tensor(delay_pattern, dtype=torch.int32)
t_idx_BxT = torch.broadcast_to(
torch.arange(T, dtype=torch.int32)[None, :],
[B, T],
)
t_idx_BxTx1 = t_idx_BxT[..., None]
t_idx_BxTxC = t_idx_BxTx1 - delay_arr.view(1, 1, C)
b_idx_BxTxC = torch.broadcast_to(
torch.arange(B, dtype=torch.int32).view(B, 1, 1),
[B, T, C],
)
c_idx_BxTxC = torch.broadcast_to(
torch.arange(C, dtype=torch.int32).view(1, 1, C),
[B, T, C],
)
# We must clamp time indices to [0..T-1] so gather_nd equivalent won't fail
t_clamped_BxTxC = torch.clamp(t_idx_BxTxC, 0, T - 1)
indices_BTCx3 = torch.stack(
[
b_idx_BxTxC.reshape(-1),
t_clamped_BxTxC.reshape(-1),
c_idx_BxTxC.reshape(-1),
],
dim=1,
).long() # Ensure indices are long type for indexing
return t_idx_BxTxC, indices_BTCx3
def apply_audio_delay(
audio_BxTxC: torch.Tensor,
pad_value: int,
bos_value: int,
precomp: tp.Tuple[torch.Tensor, torch.Tensor],
) -> torch.Tensor:
"""
Applies the delay pattern to batched audio tokens using precomputed indices,
inserting BOS where t_idx < 0 and PAD where t_idx >= T.
Args:
audio_BxTxC: [B, T, C] int16 audio tokens (or int32/float)
pad_value: the padding token
bos_value: the BOS token
precomp: (t_idx_BxTxC, indices_BTCx3) from build_delay_indices
Returns:
result_BxTxC: [B, T, C] delayed audio tokens
"""
device = audio_BxTxC.device # Get device from input tensor
t_idx_BxTxC, indices_BTCx3 = precomp
t_idx_BxTxC = t_idx_BxTxC.to(device) # Move precomputed indices to device
indices_BTCx3 = indices_BTCx3.to(device)
# Equivalent of tf.gather_nd using advanced indexing
# Ensure indices are long type if not already (build_delay_indices should handle this)
gathered_flat = audio_BxTxC[indices_BTCx3[:, 0], indices_BTCx3[:, 1], indices_BTCx3[:, 2]]
gathered_BxTxC = gathered_flat.view(audio_BxTxC.shape)
# Create masks on the correct device
mask_bos = t_idx_BxTxC < 0 # => place bos_value
mask_pad = t_idx_BxTxC >= audio_BxTxC.shape[1] # => place pad_value
# Create scalar tensors on the correct device
bos_tensor = torch.tensor(bos_value, dtype=audio_BxTxC.dtype, device=device)
pad_tensor = torch.tensor(pad_value, dtype=audio_BxTxC.dtype, device=device)
# If mask_bos, BOS; else if mask_pad, PAD; else original gather
# All tensors should now be on the same device
result_BxTxC = torch.where(mask_bos, bos_tensor, torch.where(mask_pad, pad_tensor, gathered_BxTxC))
return result_BxTxC
def build_revert_indices(B: int, T: int, C: int, delay_pattern: tp.List[int]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
"""
Precompute indices for the revert operation using PyTorch.
Returns:
A tuple (t_idx_BxTxC, indices_BTCx3) where:
- t_idx_BxTxC is a tensor of shape [B, T, C] computed as time indices plus the delay.
- indices_BTCx3 is a tensor of shape [B*T*C, 3] used for gathering, computed from:
batch indices, clamped time indices, and channel indices.
"""
# Use default device unless specified otherwise; assumes inputs might define device later
device = None # Or determine dynamically if needed, e.g., from a model parameter
delay_arr = torch.tensor(delay_pattern, dtype=torch.int32, device=device)
t_idx_BT1 = torch.broadcast_to(torch.arange(T, device=device).unsqueeze(0), [B, T])
t_idx_BT1 = t_idx_BT1.unsqueeze(-1)
t_idx_BxTxC = torch.minimum(
t_idx_BT1 + delay_arr.view(1, 1, C),
torch.tensor(T - 1, device=device),
)
b_idx_BxTxC = torch.broadcast_to(torch.arange(B, device=device).view(B, 1, 1), [B, T, C])
c_idx_BxTxC = torch.broadcast_to(torch.arange(C, device=device).view(1, 1, C), [B, T, C])
indices_BTCx3 = torch.stack(
[
b_idx_BxTxC.reshape(-1),
t_idx_BxTxC.reshape(-1),
c_idx_BxTxC.reshape(-1),
],
axis=1,
).long() # Ensure indices are long type
return t_idx_BxTxC, indices_BTCx3
def revert_audio_delay(
audio_BxTxC: torch.Tensor,
pad_value: int,
precomp: tp.Tuple[torch.Tensor, torch.Tensor],
T: int,
) -> torch.Tensor:
"""
Reverts a delay pattern from batched audio tokens using precomputed indices (PyTorch version).
Args:
audio_BxTxC: Input delayed audio tensor
pad_value: Padding value for out-of-bounds indices
precomp: Precomputed revert indices tuple containing:
- t_idx_BxTxC: Time offset indices tensor
- indices_BTCx3: Gather indices tensor for original audio
T: Original sequence length before padding
Returns:
Reverted audio tensor with same shape as input
"""
t_idx_BxTxC, indices_BTCx3 = precomp
device = audio_BxTxC.device # Get device from input tensor
# Move precomputed indices to the same device as audio_BxTxC if they aren't already
t_idx_BxTxC = t_idx_BxTxC.to(device)
indices_BTCx3 = indices_BTCx3.to(device)
# Using PyTorch advanced indexing (equivalent to tf.gather_nd or np equivalent)
gathered_flat = audio_BxTxC[indices_BTCx3[:, 0], indices_BTCx3[:, 1], indices_BTCx3[:, 2]]
gathered_BxTxC = gathered_flat.view(audio_BxTxC.size()) # Use .size() for robust reshaping
# Create pad_tensor on the correct device
pad_tensor = torch.tensor(pad_value, dtype=audio_BxTxC.dtype, device=device)
# Create T tensor on the correct device for comparison
T_tensor = torch.tensor(T, device=device)
result_BxTxC = torch.where(t_idx_BxTxC >= T_tensor, pad_tensor, gathered_BxTxC) # Changed np.where to torch.where
return result_BxTxC
@torch.no_grad()
@torch.inference_mode()
def decode(
model,
audio_codes,
):
"""
Decodes the given frames into an output audio waveform
"""
if len(audio_codes) != 1:
raise ValueError(f"Expected one frame, got {len(audio_codes)}")
try:
audio_values = model.quantizer.from_codes(audio_codes)
audio_values = model.decode(audio_values[0])
return audio_values
except Exception as e:
print(f"Error in decode method: {str(e)}")
raise
dia_lib/config.py
+187
View File
@@ -0,0 +1,187 @@
"""Configuration management module for the Dia model.
This module provides comprehensive configuration management for the Dia model,
utilizing Pydantic for validation. It defines configurations for data processing,
model architecture (encoder and decoder), and training settings.
Key components:
- DataConfig: Parameters for data loading and preprocessing.
- EncoderConfig: Architecture details for the encoder module.
- DecoderConfig: Architecture details for the decoder module.
- ModelConfig: Combined model architecture settings.
- TrainingConfig: Training hyperparameters and settings.
- DiaConfig: Master configuration combining all components.
"""
import os
from typing import Annotated
from pydantic import BaseModel, BeforeValidator, Field
class DataConfig(BaseModel, frozen=True):
"""Configuration for data loading and preprocessing.
Attributes:
text_length: Maximum length of text sequences (must be multiple of 128).
audio_length: Maximum length of audio sequences (must be multiple of 128).
channels: Number of audio channels.
text_pad_value: Value used for padding text sequences.
audio_eos_value: Value representing the end of audio sequences.
audio_bos_value: Value representing the beginning of audio sequences.
audio_pad_value: Value used for padding audio sequences.
delay_pattern: List of delay values for each audio channel.
"""
text_length: Annotated[int, BeforeValidator(lambda x: (x + 127) // 128 * 128)] = Field(gt=0, multiple_of=128)
audio_length: Annotated[int, BeforeValidator(lambda x: (x + 127) // 128 * 128)] = Field(gt=0, multiple_of=128)
channels: int = Field(default=9, gt=0, multiple_of=1)
text_pad_value: int = Field(default=0)
audio_eos_value: int = Field(default=1024)
audio_pad_value: int = Field(default=1025)
audio_bos_value: int = Field(default=1026)
delay_pattern: list[Annotated[int, Field(ge=0)]] = Field(default_factory=lambda: [0, 8, 9, 10, 11, 12, 13, 14, 15])
def __hash__(self) -> int:
"""Generate a hash based on all fields of the config."""
return hash(
(
self.text_length,
self.audio_length,
self.channels,
self.text_pad_value,
self.audio_pad_value,
self.audio_bos_value,
self.audio_eos_value,
tuple(self.delay_pattern),
)
)
class EncoderConfig(BaseModel, frozen=True):
"""Configuration for the encoder component of the Dia model.
Attributes:
n_layer: Number of transformer layers.
n_embd: Embedding dimension.
n_hidden: Hidden dimension size in the MLP layers.
n_head: Number of attention heads.
head_dim: Dimension per attention head.
"""
n_layer: int = Field(gt=0)
n_embd: int = Field(gt=0)
n_hidden: int = Field(gt=0)
n_head: int = Field(gt=0)
head_dim: int = Field(gt=0)
class DecoderConfig(BaseModel, frozen=True):
"""Configuration for the decoder component of the Dia model.
Attributes:
n_layer: Number of transformer layers.
n_embd: Embedding dimension.
n_hidden: Hidden dimension size in the MLP layers.
gqa_query_heads: Number of query heads for grouped-query self-attention.
kv_heads: Number of key/value heads for grouped-query self-attention.
gqa_head_dim: Dimension per query head for grouped-query self-attention.
cross_query_heads: Number of query heads for cross-attention.
cross_head_dim: Dimension per cross-attention head.
"""
n_layer: int = Field(gt=0)
n_embd: int = Field(gt=0)
n_hidden: int = Field(gt=0)
gqa_query_heads: int = Field(gt=0)
kv_heads: int = Field(gt=0)
gqa_head_dim: int = Field(gt=0)
cross_query_heads: int = Field(gt=0)
cross_head_dim: int = Field(gt=0)
class ModelConfig(BaseModel, frozen=True):
"""Main configuration container for the Dia model architecture.
Attributes:
encoder: Configuration for the encoder component.
decoder: Configuration for the decoder component.
src_vocab_size: Size of the source (text) vocabulary.
tgt_vocab_size: Size of the target (audio code) vocabulary.
dropout: Dropout probability applied within the model.
normalization_layer_epsilon: Epsilon value for normalization layers (e.g., LayerNorm).
weight_dtype: Data type for model weights (e.g., "float32", "bfloat16").
rope_min_timescale: Minimum timescale for Rotary Positional Embeddings (RoPE).
rope_max_timescale: Maximum timescale for Rotary Positional Embeddings (RoPE).
"""
encoder: EncoderConfig
decoder: DecoderConfig
src_vocab_size: int = Field(default=128, gt=0)
tgt_vocab_size: int = Field(default=1028, gt=0)
dropout: float = Field(default=0.0, ge=0.0, lt=1.0)
normalization_layer_epsilon: float = Field(default=1.0e-5, ge=0.0)
weight_dtype: str = Field(default="float32", description="Weight precision")
rope_min_timescale: int = Field(default=1, description="Timescale For global Attention")
rope_max_timescale: int = Field(default=10_000, description="Timescale For global Attention")
class TrainingConfig(BaseModel, frozen=True):
pass
class DiaConfig(BaseModel, frozen=True):
"""Master configuration for the Dia model.
Combines all sub-configurations into a single validated object.
Attributes:
version: Configuration version string.
model: Model architecture configuration.
training: Training process configuration (precision settings).
data: Data loading and processing configuration.
"""
version: str = Field(default="1.0")
model: ModelConfig
# TODO: remove training. this is just for backwards-compatability
training: TrainingConfig
data: DataConfig
def save(self, path: str) -> None:
"""Save the current configuration instance to a JSON file.
Ensures the parent directory exists and the file has a .json extension.
Args:
path: The target file path to save the configuration.
Raises:
ValueError: If the path is not a file with a .json extension.
"""
os.makedirs(os.path.dirname(path), exist_ok=True)
config_json = self.model_dump_json(indent=2)
with open(path, "w") as f:
f.write(config_json)
@classmethod
def load(cls, path: str) -> "DiaConfig | None":
"""Load and validate a Dia configuration from a JSON file.
Args:
path: The path to the configuration file.
Returns:
A validated DiaConfig instance if the file exists and is valid,
otherwise None if the file is not found.
Raises:
ValueError: If the path does not point to an existing .json file.
pydantic.ValidationError: If the JSON content fails validation against the DiaConfig schema.
"""
try:
with open(path, "r") as f:
content = f.read()
return cls.model_validate_json(content)
except FileNotFoundError:
return None
dia_lib/layers.py
+616
View File
@@ -0,0 +1,616 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.nn import RMSNorm
from .config import DiaConfig
from .state import DecoderInferenceState, EncoderInferenceState, KVCache
def _normalize_axes(axes: tuple[int, ...], ndim: int) -> tuple[int, ...]:
return tuple(ax if ax >= 0 else ndim + ax for ax in axes)
class DenseGeneral(nn.Module):
"""
PyTorch equivalent of flax.linen.DenseGeneral with shapes defined at init.
Stores weights (`kernel`) in the same layout as Jax and uses torch.tensordot
for the generalized matrix multiplication. Weight/bias shapes are calculated
and parameters created during initialization based on config.
`load_weights` validates shapes and copies data.
Attributes:
axis (Tuple[int, ...]): Input axis or axes to contract.
in_shapes (Tuple[int, ...]): Sizes of the input dimensions specified by `axis`.
out_features (Tuple[int, ...]): Shape of the output features (non-contracted dims).
use_bias (bool): Whether to add a bias term.
weight (nn.Parameter): The kernel parameter.
bias (Optional[nn.Parameter]): The bias parameter (if use_bias=True).
"""
def __init__(
self,
in_shapes: tuple[int, ...],
out_features: tuple[int, ...],
axis: tuple[int, ...] = (-1,),
weight_dtype: torch.dtype | None = None,
device: torch.device | None = None,
):
super().__init__()
self.in_shapes = in_shapes
self.out_features = out_features
self.axis = axis
self.kernel_shape = self.in_shapes + self.out_features
factory_kwargs = {"device": device, "dtype": weight_dtype}
self.weight = nn.Parameter(torch.empty(self.kernel_shape, **factory_kwargs))
self.register_parameter("bias", None)
def forward(self, inputs: Tensor) -> Tensor:
norm_axis = _normalize_axes(self.axis, inputs.ndim)
kernel_contract_axes = tuple(range(len(norm_axis)))
output = torch.tensordot(
inputs.to(self.weight.dtype),
self.weight,
dims=(norm_axis, kernel_contract_axes),
).to(inputs.dtype)
return output
class MlpBlock(nn.Module):
"""MLP block using DenseGeneral."""
def __init__(self, embed_dim: int, intermediate_dim: int, compute_dtype: torch.dtype):
super().__init__()
self.dtype = compute_dtype
self.wi_fused = DenseGeneral(
in_shapes=(embed_dim,),
out_features=(2, intermediate_dim),
axis=(-1,),
weight_dtype=compute_dtype,
)
self.wo = DenseGeneral(
in_shapes=(intermediate_dim,),
out_features=(embed_dim,),
axis=(-1,),
weight_dtype=compute_dtype,
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass."""
fused_x = self.wi_fused(x)
gate = fused_x[..., 0, :]
up = fused_x[..., 1, :]
hidden = torch.mul(F.silu(gate), up).to(self.dtype)
output = self.wo(hidden)
return output
class RotaryEmbedding(nn.Module):
"""Rotary Position Embedding (RoPE) implementation in PyTorch."""
def __init__(
self,
embedding_dims: int,
min_timescale: int = 1,
max_timescale: int = 10000,
dtype: torch.dtype = torch.float32,
):
super().__init__()
if embedding_dims % 2 != 0:
raise ValueError("Embedding dim must be even for RoPE.")
self.embedding_dims = embedding_dims
self.min_timescale = min_timescale
self.max_timescale = max_timescale
self.dtype = dtype
half_embedding_dim = embedding_dims // 2
fraction = (2.0 * torch.arange(0, half_embedding_dim)) / embedding_dims
self.register_buffer(
"timescale",
self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction,
persistent=False,
)
def extra_repr(self) -> str:
s = f"{self.timescale.shape}"
return s
def forward(self, inputs: torch.Tensor, position: torch.Tensor):
"""Applies RoPE."""
position = position.unsqueeze(-1).unsqueeze(-1)
timescale = self.timescale.to(inputs.device)
sinusoid_inp = position / timescale
sin = torch.sin(sinusoid_inp).to(inputs.dtype)
cos = torch.cos(sinusoid_inp).to(inputs.dtype)
first_half, second_half = torch.chunk(inputs, 2, dim=-1)
first_part = first_half * cos - second_half * sin
second_part = second_half * cos + first_half * sin
return torch.cat((first_part, second_part), dim=-1)
class Attention(nn.Module):
"""Attention using DenseGeneral."""
def __init__(
self,
config: DiaConfig,
q_embed_dim: int,
kv_embed_dim: int,
num_query_heads: int,
num_kv_heads: int,
head_dim: int,
compute_dtype: torch.dtype,
is_cross_attn: bool = False,
out_embed_dim: int | None = None,
):
super().__init__()
self.num_query_heads = num_query_heads
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
self.is_cross_attn = is_cross_attn
self.output_dim = out_embed_dim if out_embed_dim is not None else q_embed_dim
self.projected_query_dim = num_query_heads * head_dim
if num_query_heads % num_kv_heads != 0:
raise ValueError(f"num_query_heads ({num_query_heads}) must be divisible by num_kv_heads ({num_kv_heads})")
self.num_gqa_groups = num_query_heads // num_kv_heads
# --- Projection Layers using DenseGeneral ---
self.q_proj = DenseGeneral(
in_shapes=(q_embed_dim,),
out_features=(num_query_heads, head_dim),
axis=(-1,),
weight_dtype=compute_dtype,
)
self.k_proj = DenseGeneral(
in_shapes=(kv_embed_dim,),
out_features=(num_kv_heads, head_dim),
axis=(-1,),
weight_dtype=compute_dtype,
)
self.v_proj = DenseGeneral(
in_shapes=(kv_embed_dim,),
out_features=(num_kv_heads, head_dim),
axis=(-1,),
weight_dtype=compute_dtype,
)
self.o_proj = DenseGeneral(
in_shapes=(num_query_heads, head_dim),
out_features=(self.output_dim,),
axis=(-2, -1),
weight_dtype=compute_dtype,
)
# --- Rotary Embedding ---
self.rotary_emb = RotaryEmbedding(
embedding_dims=self.head_dim,
min_timescale=config.model.rope_min_timescale,
max_timescale=config.model.rope_max_timescale,
dtype=compute_dtype,
)
def forward(
self,
Xq: torch.Tensor, # (B, T, D) T = 1 in AR generation
Xkv: torch.Tensor, # (B, S, E) S = 1 in AR generation
q_positions: torch.Tensor, # (B, T)
kv_positions: torch.Tensor | None = None, # (B, S)
attn_mask: torch.Tensor | None = None, # None in Decoder Self Attention, Valid mask in Others
cache: KVCache | None = None, # None in Encoder, KVCache in Decoder
prefill: bool = False,
is_causal: bool = False,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
"""
Performs attention calculation with optional KV caching.
Args:
Xq: Query tensor (B, T, D). T=1 during single-step decoding.
Xkv: Key/Value source tensor (B, S, E). S=1 during single-step decoding for self-attn.
q_positions: Positions for queries (B, T).
kv_positions: Positions for keys/values (B, S). If None, uses q_positions.
attn_mask: Attention mask.
cache: KVCache.
prefill: If True, use prefill mode.
Returns:
A tuple containing:
- output: The attention output tensor (B, T, output_dim).
- present_kv: The K/V state to be cached for the next step ((B, N, S_new, H), (B, N, S_new, H)). For self-attn, S_new = S_past + S. For cross-attn, S_new = S_kv.
"""
if kv_positions is None:
kv_positions = q_positions
original_dtype = Xq.dtype
Xq_BxTxNxH = self.q_proj(Xq)
Xq_BxTxNxH = self.rotary_emb(Xq_BxTxNxH, position=q_positions)
Xq_BxNxTxH = Xq_BxTxNxH.transpose(1, 2)
attn_k: torch.Tensor | None = None
attn_v: torch.Tensor | None = None
if self.is_cross_attn:
attn_k, attn_v = cache.k, cache.v
else:
Xk_BxSxKxH = self.k_proj(Xkv) # (B, S, K, H)
Xv_BxSxKxH = self.v_proj(Xkv) # (B, S, K, H)
Xk_BxSxKxH = self.rotary_emb(Xk_BxSxKxH, position=kv_positions) # (B, S, K, H)
Xk_BxKxSxH = Xk_BxSxKxH.transpose(1, 2) # (B, K, S, H)
Xv_BxKxSxH = Xv_BxSxKxH.transpose(1, 2) # (B, K, S, H)
if cache is None:
attn_k = Xk_BxKxSxH
attn_v = Xv_BxKxSxH
else:
if prefill:
attn_k, attn_v = Xk_BxKxSxH, Xv_BxKxSxH
cache.prefill(attn_k, attn_v)
else:
attn_k, attn_v = cache.update(Xk_BxKxSxH, Xv_BxKxSxH)
attn_output = F.scaled_dot_product_attention(
Xq_BxNxTxH,
attn_k,
attn_v,
attn_mask=attn_mask,
scale=1.0,
enable_gqa=self.num_gqa_groups > 1,
is_causal=is_causal,
)
attn_output = attn_output.transpose(1, 2).contiguous() # (B, T, N, H)
output = self.o_proj(attn_output)
return output.to(original_dtype)
class EncoderLayer(nn.Module):
"""Transformer Encoder Layer using DenseGeneral."""
def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
super().__init__()
self.config = config
model_config = config.model
enc_config = config.model.encoder
embed_dim = enc_config.n_embd
self.pre_sa_norm = RMSNorm(
embed_dim,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
self.self_attention = Attention(
config,
q_embed_dim=embed_dim,
kv_embed_dim=embed_dim,
num_query_heads=enc_config.n_head,
num_kv_heads=enc_config.n_head,
head_dim=enc_config.head_dim,
compute_dtype=compute_dtype,
is_cross_attn=False,
out_embed_dim=embed_dim,
)
self.post_sa_norm = RMSNorm(
embed_dim,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
self.mlp = MlpBlock(embed_dim=embed_dim, intermediate_dim=enc_config.n_hidden, compute_dtype=compute_dtype)
def forward(
self,
x: torch.Tensor,
state: EncoderInferenceState,
) -> torch.Tensor:
residual = x
x_norm = self.pre_sa_norm(x)
sa_out = self.self_attention(
Xq=x_norm,
Xkv=x_norm,
q_positions=state.positions,
kv_positions=state.positions,
attn_mask=state.attn_mask,
)
x = residual + sa_out
residual = x
x_norm = self.post_sa_norm(x)
mlp_out = self.mlp(x_norm)
x = residual + mlp_out
return x
class Encoder(nn.Module):
"""Transformer Encoder Stack using DenseGeneral."""
def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
super().__init__()
self.config = config
model_config = config.model
enc_config = config.model.encoder
self.embedding = nn.Embedding(
model_config.src_vocab_size,
enc_config.n_embd,
dtype=compute_dtype,
)
self.layers = nn.ModuleList([EncoderLayer(config, compute_dtype) for _ in range(enc_config.n_layer)])
self.norm = RMSNorm(
enc_config.n_embd,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
def forward(
self,
x_ids: torch.Tensor,
state: EncoderInferenceState,
) -> torch.Tensor:
x = self.embedding(x_ids)
for layer in self.layers:
x = layer(x, state)
x = self.norm(x)
return x
class DecoderLayer(nn.Module):
"""Transformer Decoder Layer using DenseGeneral."""
def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
super().__init__()
self.config = config
model_config = config.model
dec_config = config.model.decoder
enc_config = config.model.encoder
dec_embed_dim = dec_config.n_embd
enc_embed_dim = enc_config.n_embd
# Norms
self.pre_sa_norm = RMSNorm(
dec_embed_dim,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
self.pre_ca_norm = RMSNorm(
dec_embed_dim,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
self.pre_mlp_norm = RMSNorm(
dec_embed_dim,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
# Self-Attention (GQA) with Causal Masking
self.self_attention = Attention(
config,
q_embed_dim=dec_embed_dim,
kv_embed_dim=dec_embed_dim,
num_query_heads=dec_config.gqa_query_heads,
num_kv_heads=dec_config.kv_heads,
head_dim=dec_config.gqa_head_dim,
compute_dtype=compute_dtype,
is_cross_attn=False,
out_embed_dim=dec_embed_dim,
)
# Cross-Attention (MHA)
self.cross_attention = Attention(
config=config,
q_embed_dim=dec_embed_dim,
kv_embed_dim=enc_embed_dim, # Note kv_embed_dim
num_query_heads=dec_config.cross_query_heads,
num_kv_heads=dec_config.cross_query_heads,
head_dim=dec_config.cross_head_dim,
compute_dtype=compute_dtype,
is_cross_attn=True,
out_embed_dim=dec_embed_dim,
)
# MLP
self.mlp = MlpBlock(
embed_dim=dec_embed_dim,
intermediate_dim=dec_config.n_hidden,
compute_dtype=compute_dtype,
)
def forward(
self,
x: torch.Tensor,
state: DecoderInferenceState,
self_attn_cache: KVCache | None = None,
cross_attn_cache: KVCache | None = None,
prefill: bool = False,
) -> torch.Tensor:
residual = x
x_norm = self.pre_sa_norm(x)
sa_out = self.self_attention(
Xq=x_norm, # (2, 1, D)
Xkv=x_norm, # (2, 1, D)
q_positions=state.dec_positions, # (2, 1)
kv_positions=state.dec_positions, # (2, 1)
attn_mask=None,
cache=self_attn_cache,
prefill=prefill,
is_causal=prefill,
)
x = residual + sa_out
residual = x
x_norm = self.pre_ca_norm(x)
ca_out = self.cross_attention(
Xq=x_norm,
Xkv=state.enc_out,
q_positions=state.dec_positions,
kv_positions=state.enc_positions,
attn_mask=state.dec_cross_attn_mask,
cache=cross_attn_cache,
)
x = residual + ca_out
residual = x
x_norm = self.pre_mlp_norm(x)
mlp_out = self.mlp(x_norm)
x = residual + mlp_out
return x
class Decoder(nn.Module):
"""Transformer Decoder Stack using DenseGeneral."""
def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
super().__init__()
self.config = config
model_config = config.model
dec_config = config.model.decoder
data_config = config.data
self.num_channels = data_config.channels
self.num_layers = dec_config.n_layer
self.embeddings = nn.ModuleList(
[
nn.Embedding(model_config.tgt_vocab_size, dec_config.n_embd, dtype=compute_dtype)
for _ in range(self.num_channels)
]
)
self.layers = nn.ModuleList(
[DecoderLayer(config=config, compute_dtype=compute_dtype) for _ in range(self.num_layers)]
)
self.norm = RMSNorm(
dec_config.n_embd,
eps=model_config.normalization_layer_epsilon,
dtype=torch.float32,
)
self.logits_dense = DenseGeneral(
in_shapes=(dec_config.n_embd,),
out_features=(self.num_channels, model_config.tgt_vocab_size),
axis=(-1,),
weight_dtype=compute_dtype,
)
def precompute_cross_attn_cache(
self,
enc_out: torch.Tensor, # (B, S, E)
enc_positions: torch.Tensor, # (B, S)
) -> list[KVCache]:
"""
Computes the Key and Value tensors for cross-attention for each layer from the encoder output.
"""
per_layer_kv_cache: list[KVCache] = []
for layer in self.layers:
cross_attn_module = layer.cross_attention
k_proj = cross_attn_module.k_proj(enc_out)
v_proj = cross_attn_module.v_proj(enc_out)
k_proj = cross_attn_module.rotary_emb(k_proj, position=enc_positions)
k = k_proj.transpose(1, 2)
v = v_proj.transpose(1, 2)
per_layer_kv_cache.append(KVCache.from_kv(k, v))
return per_layer_kv_cache
def decode_step(
self,
tgt_ids_Bx1xC: torch.Tensor, # [B, 1, C]
state: DecoderInferenceState,
) -> torch.Tensor:
"""
Performs a single decoding step, managing KV caches layer by layer.
Returns:
A tuple containing:
- logits_Bx1xCV: The final output logits for the current step (B, 1, C*V), cast to float32.
"""
x = None
for i in range(self.num_channels):
channel_tokens = tgt_ids_Bx1xC[..., i]
channel_embed = self.embeddings[i](channel_tokens)
x = channel_embed if x is None else x + channel_embed
for i, layer in enumerate(self.layers):
self_cache = state.self_attn_cache[i]
cross_cache = state.cross_attn_cache[i]
x = layer(
x, # (2, 1, D)
state,
self_attn_cache=self_cache,
cross_attn_cache=cross_cache,
)
x = self.norm(x)
logits_Bx1xCxV = self.logits_dense(x)
return logits_Bx1xCxV.to(torch.float32)
def forward(self, tgt_ids_BxTxC: torch.Tensor, state: DecoderInferenceState) -> torch.Tensor:
"""
Forward pass for the Decoder stack, managing KV caches.
Args:
tgt_ids_BxTxC: Target token IDs (B, T, C).
encoder_out: Output from the encoder (B, S, E).
tgt_positions: Positions for target sequence (B, T).
src_positions: Positions for source sequence (B, S).
self_attn_mask: Mask for self-attention.
cross_attn_mask: Mask for cross-attention.
past_key_values: List containing the self-attention KV cache for each layer
from the previous decoding step. `len(past_key_values)` should
equal `num_layers`.
precomputed_cross_attn_kv: A single tuple containing the pre-computed K/V cache
derived from `encoder_out`. This is passed identically
to all layers.
Returns:
A tuple containing:
- logits: The final output logits (B, T, C * V), cast to float32.
- present_key_values: A list containing the updated self-attention KV cache
for each layer for the *current* decoding step.
"""
_, _, num_channels_in = tgt_ids_BxTxC.shape
assert num_channels_in == self.num_channels, "Input channels mismatch"
# Embeddings
x = None
for i in range(self.num_channels):
channel_tokens = tgt_ids_BxTxC[..., i]
channel_embed = self.embeddings[i](channel_tokens)
x = channel_embed if x is None else x + channel_embed
for i, layer in enumerate(self.layers):
self_cache = state.self_attn_cache[i]
cross_cache = state.cross_attn_cache[i]
x = layer(x, state, self_attn_cache=self_cache, cross_attn_cache=cross_cache, prefill=True)
# Final Norm
x = self.norm(x)
logits_BxTxCxV = self.logits_dense(x)
return logits_BxTxCxV.to(torch.float32)
class DiaModel(nn.Module):
"""PyTorch Dia Model using DenseGeneral."""
def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
super().__init__()
self.config = config
self.encoder = Encoder(config, compute_dtype)
self.decoder = Decoder(config, compute_dtype)
dia_lib/model.py
+447
View File
@@ -0,0 +1,447 @@
import time
from enum import Enum
import dac
import numpy as np
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from .audio import apply_audio_delay, build_delay_indices, build_revert_indices, decode, revert_audio_delay
from .config import DiaConfig
from .layers import DiaModel
from .state import DecoderInferenceState, DecoderOutput, EncoderInferenceState
DEFAULT_SAMPLE_RATE = 44100
def _get_default_device():
if torch.cuda.is_available():
return torch.device("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def _sample_next_token(
logits_BCxV: torch.Tensor,
temperature: float,
top_p: float,
cfg_filter_top_k: int | None = None,
) -> torch.Tensor:
if temperature == 0.0:
return torch.argmax(logits_BCxV, dim=-1)
logits_BCxV = logits_BCxV / temperature
if cfg_filter_top_k is not None:
_, top_k_indices_BCxV = torch.topk(logits_BCxV, k=cfg_filter_top_k, dim=-1)
mask = torch.ones_like(logits_BCxV, dtype=torch.bool)
mask.scatter_(dim=-1, index=top_k_indices_BCxV, value=False)
logits_BCxV = logits_BCxV.masked_fill(mask, -torch.inf)
if top_p < 1.0:
probs_BCxV = torch.softmax(logits_BCxV, dim=-1)
sorted_probs_BCxV, sorted_indices_BCxV = torch.sort(probs_BCxV, dim=-1, descending=True)
cumulative_probs_BCxV = torch.cumsum(sorted_probs_BCxV, dim=-1)
sorted_indices_to_remove_BCxV = cumulative_probs_BCxV > top_p
sorted_indices_to_remove_BCxV[..., 1:] = sorted_indices_to_remove_BCxV[..., :-1].clone()
sorted_indices_to_remove_BCxV[..., 0] = 0
indices_to_remove_BCxV = torch.zeros_like(sorted_indices_to_remove_BCxV)
indices_to_remove_BCxV.scatter_(dim=-1, index=sorted_indices_BCxV, src=sorted_indices_to_remove_BCxV)
logits_BCxV = logits_BCxV.masked_fill(indices_to_remove_BCxV, -torch.inf)
final_probs_BCxV = torch.softmax(logits_BCxV, dim=-1)
sampled_indices_BC = torch.multinomial(final_probs_BCxV, num_samples=1)
sampled_indices_C = sampled_indices_BC.squeeze(-1)
return sampled_indices_C
class ComputeDtype(str, Enum):
FLOAT32 = "float32"
FLOAT16 = "float16"
BFLOAT16 = "bfloat16"
def to_dtype(self) -> torch.dtype:
if self == ComputeDtype.FLOAT32:
return torch.float32
elif self == ComputeDtype.FLOAT16:
return torch.float16
elif self == ComputeDtype.BFLOAT16:
return torch.bfloat16
else:
raise ValueError(f"Unsupported compute dtype: {self}")
class Dia:
def __init__(
self,
config: DiaConfig,
compute_dtype: str | ComputeDtype = ComputeDtype.FLOAT32,
device: torch.device | None = None,
):
"""Initializes the Dia model.
Args:
config: The configuration object for the model.
device: The device to load the model onto. If None, will automatically select the best available device.
Raises:
RuntimeError: If there is an error loading the DAC model.
"""
super().__init__()
self.config = config
self.device = device if device is not None else _get_default_device()
if isinstance(compute_dtype, str):
compute_dtype = ComputeDtype(compute_dtype)
self.compute_dtype = compute_dtype.to_dtype()
self.model = DiaModel(config, self.compute_dtype)
self.dac_model = None
@classmethod
def from_local(
cls,
config_path: str,
checkpoint_path: str,
compute_dtype: str | ComputeDtype = ComputeDtype.FLOAT32,
device: torch.device | None = None,
) -> "Dia":
"""Loads the Dia model from local configuration and checkpoint files.
Args:
config_path: Path to the configuration JSON file.
checkpoint_path: Path to the model checkpoint (.pth) file.
device: The device to load the model onto. If None, will automatically select the best available device.
Returns:
An instance of the Dia model loaded with weights and set to eval mode.
Raises:
FileNotFoundError: If the config or checkpoint file is not found.
RuntimeError: If there is an error loading the checkpoint.
"""
config = DiaConfig.load(config_path)
if config is None:
raise FileNotFoundError(f"Config file not found at {config_path}")
dia = cls(config, compute_dtype, device)
try:
state_dict = torch.load(checkpoint_path, map_location=dia.device)
dia.model.load_state_dict(state_dict)
except FileNotFoundError:
raise FileNotFoundError(f"Checkpoint file not found at {checkpoint_path}")
except Exception as e:
raise RuntimeError(f"Error loading checkpoint from {checkpoint_path}") from e
dia.model.to(dia.device)
dia.model.eval()
dia._load_dac_model()
return dia
@classmethod
def from_pretrained(
cls,
model_name: str = "nari-labs/Dia-1.6B",
compute_dtype: str | ComputeDtype = ComputeDtype.FLOAT32,
device: torch.device | None = None,
) -> "Dia":
"""Loads the Dia model from a Hugging Face Hub repository.
Downloads the configuration and checkpoint files from the specified
repository ID and then loads the model.
Args:
model_name: The Hugging Face Hub repository ID (e.g., "NariLabs/Dia-1.6B").
device: The device to load the model onto. If None, will automatically select the best available device.
Returns:
An instance of the Dia model loaded with weights and set to eval mode.
Raises:
FileNotFoundError: If config or checkpoint download/loading fails.
RuntimeError: If there is an error loading the checkpoint.
"""
config_path = hf_hub_download(repo_id=model_name, filename="config.json")
checkpoint_path = hf_hub_download(repo_id=model_name, filename="dia-v0_1.pth")
return cls.from_local(config_path, checkpoint_path, compute_dtype, device)
def _load_dac_model(self):
try:
dac_model_path = dac.utils.download()
dac_model = dac.DAC.load(dac_model_path).to(self.device)
except Exception as e:
raise RuntimeError("Failed to load DAC model") from e
self.dac_model = dac_model
def _prepare_text_input(self, text: str) -> torch.Tensor:
"""Encodes text prompt, pads, and creates attention mask and positions."""
text_pad_value = self.config.data.text_pad_value
max_len = self.config.data.text_length
byte_text = text.encode("utf-8")
replaced_bytes = byte_text.replace(b"[S1]", b"\x01").replace(b"[S2]", b"\x02")
text_tokens = list(replaced_bytes)
current_len = len(text_tokens)
padding_needed = max_len - current_len
if padding_needed <= 0:
text_tokens = text_tokens[:max_len]
padded_text_np = np.array(text_tokens, dtype=np.uint8)
else:
padded_text_np = np.pad(
text_tokens,
(0, padding_needed),
mode="constant",
constant_values=text_pad_value,
).astype(np.uint8)
src_tokens = torch.from_numpy(padded_text_np).to(torch.long).to(self.device).unsqueeze(0) # [1, S]
return src_tokens
def _prepare_audio_prompt(self, audio_prompt: torch.Tensor | None) -> tuple[torch.Tensor, int]:
num_channels = self.config.data.channels
audio_bos_value = self.config.data.audio_bos_value
audio_pad_value = self.config.data.audio_pad_value
delay_pattern = self.config.data.delay_pattern
max_delay_pattern = max(delay_pattern)
prefill = torch.full(
(1, num_channels),
fill_value=audio_bos_value,
dtype=torch.int,
device=self.device,
)
prefill_step = 1
if audio_prompt is not None:
prefill_step += audio_prompt.shape[0]
prefill = torch.cat([prefill, audio_prompt], dim=0)
delay_pad_tensor = torch.full(
(max_delay_pattern, num_channels), fill_value=-1, dtype=torch.int, device=self.device
)
prefill = torch.cat([prefill, delay_pad_tensor], dim=0)
delay_precomp = build_delay_indices(
B=1,
T=prefill.shape[0],
C=num_channels,
delay_pattern=delay_pattern,
)
prefill = apply_audio_delay(
audio_BxTxC=prefill.unsqueeze(0),
pad_value=audio_pad_value,
bos_value=audio_bos_value,
precomp=delay_precomp,
).squeeze(0)
return prefill, prefill_step
def _prepare_generation(self, text: str, audio_prompt: str | torch.Tensor | None, verbose: bool):
enc_input_cond = self._prepare_text_input(text)
enc_input_uncond = torch.zeros_like(enc_input_cond)
enc_input = torch.cat([enc_input_uncond, enc_input_cond], dim=0)
if isinstance(audio_prompt, str):
audio_prompt = self.load_audio(audio_prompt)
prefill, prefill_step = self._prepare_audio_prompt(audio_prompt)
if verbose:
print("generate: data loaded")
enc_state = EncoderInferenceState.new(self.config, enc_input_cond)
encoder_out = self.model.encoder(enc_input, enc_state)
dec_cross_attn_cache = self.model.decoder.precompute_cross_attn_cache(encoder_out, enc_state.positions)
dec_state = DecoderInferenceState.new(
self.config, enc_state, encoder_out, dec_cross_attn_cache, self.compute_dtype
)
dec_output = DecoderOutput.new(self.config, self.device)
dec_output.prefill(prefill, prefill_step)
dec_step = prefill_step - 1
if dec_step > 0:
dec_state.prepare_step(0, dec_step)
tokens_BxTxC = dec_output.get_tokens_at(0, dec_step).unsqueeze(0).expand(2, -1, -1)
self.model.decoder.forward(tokens_BxTxC, dec_state)
return dec_state, dec_output
def _decoder_step(
self,
tokens_Bx1xC: torch.Tensor,
dec_state: DecoderInferenceState,
cfg_scale: float,
temperature: float,
top_p: float,
cfg_filter_top_k: int,
) -> torch.Tensor:
audio_eos_value = self.config.data.audio_eos_value
logits_Bx1xCxV = self.model.decoder.decode_step(tokens_Bx1xC, dec_state)
logits_last_BxCxV = logits_Bx1xCxV[:, -1, :, :]
uncond_logits_CxV = logits_last_BxCxV[0, :, :]
cond_logits_CxV = logits_last_BxCxV[1, :, :]
logits_CxV = cond_logits_CxV + cfg_scale * (cond_logits_CxV - uncond_logits_CxV)
logits_CxV[:, audio_eos_value + 1 :] = -torch.inf
logits_CxV[1:, audio_eos_value:] = -torch.inf
pred_C = _sample_next_token(
logits_CxV.float(),
temperature=temperature,
top_p=top_p,
cfg_filter_top_k=cfg_filter_top_k,
)
return pred_C
def _generate_output(self, generated_codes: torch.Tensor) -> np.ndarray:
num_channels = self.config.data.channels
seq_length = generated_codes.shape[0]
delay_pattern = self.config.data.delay_pattern
audio_pad_value = self.config.data.audio_pad_value
max_delay_pattern = max(delay_pattern)
revert_precomp = build_revert_indices(
B=1,
T=seq_length,
C=num_channels,
delay_pattern=delay_pattern,
)
codebook = revert_audio_delay(
audio_BxTxC=generated_codes.unsqueeze(0),
pad_value=audio_pad_value,
precomp=revert_precomp,
T=seq_length,
)[:, :-max_delay_pattern, :]
min_valid_index = 0
max_valid_index = 1023
invalid_mask = (codebook < min_valid_index) | (codebook > max_valid_index)
codebook[invalid_mask] = 0
audio = decode(self.dac_model, codebook.transpose(1, 2))
return audio.squeeze().cpu().numpy()
def load_audio(self, audio_path: str) -> torch.Tensor:
audio, sr = torchaudio.load(audio_path, channels_first=True) # C, T
if sr != DEFAULT_SAMPLE_RATE:
audio = torchaudio.functional.resample(audio, sr, DEFAULT_SAMPLE_RATE)
audio = audio.to(self.device).unsqueeze(0) # 1, C, T
audio_data = self.dac_model.preprocess(audio, DEFAULT_SAMPLE_RATE)
_, encoded_frame, _, _, _ = self.dac_model.encode(audio_data) # 1, C, T
return encoded_frame.squeeze(0).transpose(0, 1)
def save_audio(self, path: str, audio: np.ndarray):
import soundfile as sf
sf.write(path, audio, DEFAULT_SAMPLE_RATE)
@torch.inference_mode()
def generate(
self,
text: str,
max_tokens: int | None = None,
cfg_scale: float = 3.0,
temperature: float = 1.3,
top_p: float = 0.95,
use_torch_compile: bool = False,
cfg_filter_top_k: int = 35,
audio_prompt: str | torch.Tensor | None = None,
audio_prompt_path: str | None = None,
use_cfg_filter: bool | None = None,
verbose: bool = False,
) -> np.ndarray:
audio_eos_value = self.config.data.audio_eos_value
audio_pad_value = self.config.data.audio_pad_value
delay_pattern = self.config.data.delay_pattern
max_tokens = self.config.data.audio_length if max_tokens is None else max_tokens
max_delay_pattern = max(delay_pattern)
self.model.eval()
if audio_prompt_path:
print("Warning: audio_prompt_path is deprecated. Use audio_prompt instead.")
audio_prompt = audio_prompt_path
if use_cfg_filter is not None:
print("Warning: use_cfg_filter is deprecated.")
if verbose:
total_start_time = time.time()
dec_state, dec_output = self._prepare_generation(text, audio_prompt, verbose)
dec_step = dec_output.prefill_step - 1
bos_countdown = max_delay_pattern
eos_detected = False
eos_countdown = -1
if use_torch_compile:
step_fn = torch.compile(self._decoder_step, mode="default")
else:
step_fn = self._decoder_step
if verbose:
print("generate: starting generation loop")
if use_torch_compile:
print("generate: by using use_torch_compile=True, the first step would take long")
start_time = time.time()
while dec_step < max_tokens:
dec_state.prepare_step(dec_step)
tokens_Bx1xC = dec_output.get_tokens_at(dec_step).unsqueeze(0).expand(2, -1, -1)
pred_C = step_fn(
tokens_Bx1xC,
dec_state,
cfg_scale,
temperature,
top_p,
cfg_filter_top_k,
)
if (not eos_detected and pred_C[0] == audio_eos_value) or dec_step == max_tokens - max_delay_pattern - 1:
eos_detected = True
eos_countdown = max_delay_pattern
if eos_countdown > 0:
step_after_eos = max_delay_pattern - eos_countdown
for i, d in enumerate(delay_pattern):
if step_after_eos == d:
pred_C[i] = audio_eos_value
elif step_after_eos > d:
pred_C[i] = audio_pad_value
eos_countdown -= 1
bos_countdown = max(0, bos_countdown - 1)
dec_output.update_one(pred_C, dec_step + 1, bos_countdown > 0)
if eos_countdown == 0:
break
dec_step += 1
if verbose and dec_step % 86 == 0:
duration = time.time() - start_time
print(
f"generate step {dec_step}: speed={86 / duration:.3f} tokens/s, realtime factor={1 / duration:.3f}x"
)
start_time = time.time()
if dec_output.prefill_step >= dec_step + 1:
print("Warning: Nothing generated")
return None
generated_codes = dec_output.generated_tokens[dec_output.prefill_step : dec_step + 1, :]
if verbose:
total_step = dec_step + 1 - dec_output.prefill_step
total_duration = time.time() - total_start_time
print(f"generate: total step={total_step}, total duration={total_duration:.3f}s")
return self._generate_output(generated_codes)
dia_lib/state.py
+203
View File
@@ -0,0 +1,203 @@
from dataclasses import dataclass
import torch
from .config import DiaConfig
def create_attn_mask(
q_padding_mask_1d: torch.Tensor,
k_padding_mask_1d: torch.Tensor,
device: torch.device,
is_causal: bool = False,
) -> torch.Tensor:
"""
Creates the attention mask (self or cross) mimicking JAX segment ID logic.
"""
B1, Tq = q_padding_mask_1d.shape
B2, Tk = k_padding_mask_1d.shape
assert B1 == B2, "Query and key batch dimensions must match"
p_mask_q = q_padding_mask_1d.unsqueeze(2) # Shape [B, Tq, 1]
p_mask_k = k_padding_mask_1d.unsqueeze(1) # Shape [B, 1, Tk]
# Condition A: Non-padding query attends to non-padding key
non_pad_attends_non_pad = p_mask_q & p_mask_k # Shape [B, Tq, Tk]
# Condition B: Padding query attends to padding key
pad_attends_pad = (~p_mask_q) & (~p_mask_k) # Shape [B, Tq, Tk]
# Combine: True if padding status is compatible (both non-pad OR both pad)
mask = non_pad_attends_non_pad | pad_attends_pad # Shape [B, Tq, Tk]
if is_causal:
assert Tq == Tk, "Causal mask requires query and key sequence lengths to be equal"
causal_mask_2d = torch.tril(torch.ones((Tq, Tk), dtype=torch.bool, device=device)) # Shape [Tq, Tk]
causal_mask = mask & causal_mask_2d # Shape [B, Tq, Tk]
return causal_mask.unsqueeze(1) # Shape [B, 1, Tq, Tk]
else:
return mask.unsqueeze(1) # Shape [B, 1, Tq, Tk]
@dataclass
class EncoderInferenceState:
"""Parameters specifically for encoder inference."""
max_seq_len: int
device: torch.device
positions: torch.Tensor
padding_mask: torch.Tensor
attn_mask: torch.Tensor
@classmethod
def new(cls, config: DiaConfig, cond_src: torch.Tensor) -> "EncoderInferenceState":
"""Creates EtorchrInferenceParams from DiaConfig and a device."""
device = cond_src.device
positions = torch.arange(config.data.text_length, device=device).to(torch.long).unsqueeze(0).expand(2, -1)
padding_mask = (cond_src != config.data.text_pad_value).to(device).expand(2, -1)
attn_mask = create_attn_mask(padding_mask, padding_mask, device, is_causal=False)
return cls(
max_seq_len=config.data.text_length,
device=device,
positions=positions,
padding_mask=padding_mask,
attn_mask=attn_mask,
)
class KVCache:
def __init__(
self,
num_heads: int,
max_len: int,
head_dim: int,
dtype: torch.dtype,
device: torch.device,
k: torch.Tensor | None = None,
v: torch.Tensor | None = None,
):
self.k = torch.zeros((2, num_heads, max_len, head_dim), dtype=dtype, device=device) if k is None else k
self.v = torch.zeros((2, num_heads, max_len, head_dim), dtype=dtype, device=device) if v is None else v
self.current_idx = torch.tensor(0)
@classmethod
def from_kv(cls, k: torch.Tensor, v: torch.Tensor) -> "KVCache":
return cls(
num_heads=k.shape[1],
max_len=k.shape[2],
head_dim=k.shape[3],
dtype=k.dtype,
device=k.device,
k=k,
v=v,
)
def update(self, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
self.k[:, :, self.current_idx : self.current_idx + 1, :] = k
self.v[:, :, self.current_idx : self.current_idx + 1, :] = v
self.current_idx += 1
return self.k[:, :, : self.current_idx, :], self.v[:, :, : self.current_idx, :]
def prefill(self, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
prefill_len = k.shape[2]
self.k[:, :, :prefill_len, :] = k
self.v[:, :, :prefill_len, :] = v
self.current_idx = prefill_len - 1
@dataclass
class DecoderInferenceState:
"""Parameters specifically for decoder inference."""
device: torch.device
dtype: torch.dtype
enc_out: torch.Tensor
enc_positions: torch.Tensor
dec_positions: torch.Tensor
dec_cross_attn_mask: torch.Tensor
self_attn_cache: list[KVCache]
cross_attn_cache: list[KVCache]
@classmethod
def new(
cls,
config: DiaConfig,
enc_state: EncoderInferenceState,
enc_out: torch.Tensor,
dec_cross_attn_cache: list[KVCache],
compute_dtype: torch.dtype,
) -> "DecoderInferenceState":
"""Creates DecoderInferenceParams from DiaConfig and a device."""
device = enc_out.device
max_audio_len = config.data.audio_length
dec_positions = torch.full((2, 1), fill_value=0, dtype=torch.long, device=device)
tgt_padding_mask = torch.ones((2, 1), dtype=torch.bool, device=device)
dec_cross_attn_mask = create_attn_mask(tgt_padding_mask, enc_state.padding_mask, device, is_causal=False)
self_attn_cache = [
KVCache(
config.model.decoder.kv_heads,
max_audio_len,
config.model.decoder.gqa_head_dim,
compute_dtype,
device,
)
for _ in range(config.model.decoder.n_layer)
]
return cls(
device=device,
dtype=compute_dtype,
enc_out=enc_out,
enc_positions=enc_state.positions,
dec_positions=dec_positions,
dec_cross_attn_mask=dec_cross_attn_mask,
self_attn_cache=self_attn_cache,
cross_attn_cache=dec_cross_attn_cache,
)
def prepare_step(self, step_from: int, step_to: int | None = None) -> None:
if step_to is None:
step_to = step_from + 1
self.dec_positions = torch.arange(step_from, step_to, device=self.device).unsqueeze(0).expand(2, -1)
@dataclass
class DecoderOutput:
generated_tokens: torch.Tensor
prefill_step: int
@classmethod
def new(cls, config: DiaConfig, device: torch.device) -> "DecoderOutput":
max_audio_len = config.data.audio_length
return cls(
generated_tokens=torch.full(
(max_audio_len, config.data.channels),
fill_value=-1,
dtype=torch.int,
device=device,
),
prefill_step=0,
)
def get_tokens_at(self, step_from: int, step_to: int | None = None) -> torch.Tensor:
if step_to is None:
step_to = step_from + 1
return self.generated_tokens[step_from:step_to, :]
def update_one(self, dec_out: torch.Tensor, step: int, apply_mask: bool = False):
if apply_mask:
mask = self.generated_tokens[step : step + 1, :] == -1
self.generated_tokens[step : step + 1, :] = torch.where(
mask, dec_out, self.generated_tokens[step : step + 1, :]
)
else:
self.generated_tokens[step : step + 1, :] = dec_out
def prefill(self, dec_out: torch.Tensor, prefill_step: int):
length = dec_out.shape[0]
self.generated_tokens[0:length, :] = dec_out
self.prefill_step = prefill_step
nodes.py
+240
View File
@@ -0,0 +1,240 @@
# ComfyUI-DiaTest/nodes.py
import os
import torch
import numpy as np
import folder_paths
import torchaudio # Kept for potential internal use by dia_lib
from huggingface_hub import hf_hub_download
import traceback
import gc
# --- Import Dia library components ---
try:
from .dia_lib.model import Dia, ComputeDtype, DEFAULT_SAMPLE_RATE
from .dia_lib.config import DiaConfig
except ImportError as e:
print("ComfyUI-DiaTest: Error importing Dia library components.")
print(f"Ensure the 'dia_lib' folder exists in '{os.path.dirname(__file__)}'.")
print(f"Import error: {e}")
raise e
# --- End Dia library imports ---
# --- Helper Functions ---
def get_torch_device():
"""Checks for CUDA availability and returns the CUDA device."""
if torch.cuda.is_available():
return torch.device("cuda")
else:
# If CUDA is not available, raise an error as GPU is required.
raise RuntimeError("CUDA device not available. This node requires a CUDA-enabled GPU.")
# --- End Helper Functions ---
# --- Global model cache ---
loaded_dia_model = None
loaded_model_key = None
class DiaGenerate:
"""Loads the Dia TTS model from Hub onto GPU and generates audio."""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
# Model Loading Params
"repo_id": ("STRING", {"default": "nari-labs/Dia-1.6B"}),
# No device override - GPU is forced
# Generation Params
"text": ("STRING", {"multiline": True, "dynamicPrompts": False, "default": "[S1] Hello world. [S2] This is a test."}),
"max_tokens": ("INT", {"default": 1720, "min": 860, "max": 3072, "step": 10}),
"cfg_scale": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 7.0, "step": 0.1}),
"temperature": ("FLOAT", {"default": 1.3, "min": 0.1, "max": 1.5, "step": 0.05}),
"top_p": ("FLOAT", {"default": 0.95, "min": 0.1, "max": 1.0, "step": 0.01}),
"cfg_filter_top_k": ("INT", {"default": 35, "min": 1, "max": 100, "step": 1}),
"speed_factor": ("FLOAT", {"default": 0.94, "min": 0.5, "max": 1.5, "step": 0.01}), # Added speed factor
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
}
RETURN_TYPES = ("AUDIO",)
RETURN_NAMES = ("audio",)
FUNCTION = "load_and_generate"
CATEGORY = "audio/DiaTest"
def load_and_generate(self, repo_id, text: str, max_tokens: int, cfg_scale: float, temperature: float, top_p: float, cfg_filter_top_k: int, speed_factor: float, seed: int):
global loaded_dia_model, loaded_model_key
# --- Model Loading Logic ---
# Force GPU device
device = get_torch_device() # This will raise error if no CUDA
# Use float32 internally
compute_dtype_str = "float32"
print(f"DiaTestGenerate: Target device: {device}, Compute dtype: {compute_dtype_str}")
# Cache key includes repo and dtype (device is fixed to CUDA)
current_key = (repo_id, compute_dtype_str, str(device))
dia_model = None
# Check cache
if loaded_dia_model is not None and current_key == loaded_model_key:
print("DiaTestGenerate: Using cached model.")
dia_model = loaded_dia_model
# Defensive check: ensure cached model is indeed on CUDA
if dia_model.device != device:
print(f"DiaTestGenerate: Warning: Cached model not on expected device ({dia_model.device}). Moving to {device}.")
try:
dia_model.model.to(device)
if dia_model.dac_model: dia_model.dac_model.to(device)
dia_model.device = device
except Exception as move_e:
print(f"DiaTestGenerate: Error moving cached model: {move_e}")
loaded_dia_model = None; loaded_model_key = None; raise move_e
else:
# Clear previous model if config changed
if loaded_dia_model is not None:
print(f"DiaTestGenerate: Configuration changed. Clearing previous model...")
try:
if hasattr(loaded_dia_model, 'model'): del loaded_dia_model.model
if hasattr(loaded_dia_model, 'dac_model'): del loaded_dia_model.dac_model
del loaded_dia_model
except Exception as del_e: print(f"DiaTestGenerate: Error deleting previous model: {del_e}")
loaded_dia_model = None; loaded_model_key = None; gc.collect()
torch.cuda.empty_cache(); print("DiaTestGenerate: Cleared CUDA cache.")
# Load model from Hub
print(f"DiaTestGenerate: Loading model from Hugging Face Hub: repo_id='{repo_id}'")
try:
# Pre-download files (optional)
try:
cache_dir = os.path.join(folder_paths.models_dir, "huggingface")
os.makedirs(cache_dir, exist_ok=True)
hf_hub_download(repo_id=repo_id, filename="config.json", cache_dir=cache_dir, resume_download=True, etag_timeout=10)
hf_hub_download(repo_id=repo_id, filename="dia-v0_1.pth", cache_dir=cache_dir, resume_download=True, etag_timeout=10)
print(f"DiaTestGenerate: Ensured model files are cached.")
except Exception as download_e:
print(f"DiaTestGenerate: Warning during file pre-check/download: {download_e}")
# Load the model
dia_model = Dia.from_pretrained(
model_name=repo_id,
compute_dtype=compute_dtype_str, # Use string 'float32'
device=device # Load directly onto GPU
)
print("DiaTestGenerate: Model loaded successfully.")
loaded_dia_model = dia_model
loaded_model_key = current_key
except Exception as e:
print(f"DiaTestGenerate: Error loading model: {e}")
traceback.print_exc()
loaded_dia_model = None; loaded_model_key = None
raise e
# --- End Model Loading Logic ---
# --- Generation Logic ---
if not text or text.isspace(): raise ValueError("Input text cannot be empty.")
# Seed setting
MAX_SEED_NUMPY = 2**32 - 1
seed_torch = seed; seed_numpy = seed % MAX_SEED_NUMPY
torch.manual_seed(seed_torch); np.random.seed(seed_numpy)
torch.cuda.manual_seed_all(seed_torch) # Seed CUDA
print(f"DiaTestGenerate: Using ComfyUI seed {seed} (Torch: {seed_torch}, NumPy: {seed_numpy})")
text_for_generate = text
try:
print(f"DiaTestGenerate: Starting generation...")
# Generation Call
with torch.inference_mode():
output_np = dia_model.generate(
text=text_for_generate, max_tokens=max_tokens, cfg_scale=cfg_scale, temperature=temperature,
top_p=top_p, cfg_filter_top_k=cfg_filter_top_k,
# No audio_prompt
use_torch_compile=False, verbose=True
)
# Handle failed generation
if output_np is None or output_np.size == 0:
print("DiaTestGenerate: Warning - Generation returned None or empty array. Outputting silence.")
silent_tensor = torch.zeros((1, 1, DEFAULT_SAMPLE_RATE), dtype=torch.float32)
result = {'waveform': silent_tensor, 'sample_rate': DEFAULT_SAMPLE_RATE}
return (result,)
print(f"DiaTestGenerate: Raw generation complete. Shape: {output_np.shape}")
# --- Apply Speed Factor ---
if speed_factor != 1.0:
# Ensure speed_factor is valid
speed_factor = max(0.1, min(speed_factor, 5.0)) # Clamp to reasonable range
original_len = len(output_np)
target_len = int(original_len / speed_factor)
if target_len > 0 and target_len != original_len:
print(f"DiaTestGenerate: Applying speed factor {speed_factor:.2f}x (length {original_len} -> {target_len})")
x_original = np.arange(original_len)
x_resampled = np.linspace(0, original_len - 1, target_len)
# Ensure float input for interp if not already
if not np.issubdtype(output_np.dtype, np.floating):
output_np = output_np.astype(np.float32)
resampled_audio_np = np.interp(x_resampled, x_original, output_np)
output_np = resampled_audio_np # Use the resampled audio
else:
print(f"DiaTestGenerate: Skipping speed adjustment (factor: {speed_factor:.2f}).")
# --- End Speed Factor ---
# --- Output Formatting ---
try:
# Convert numpy to tensor, ensure float32
output_tensor = torch.from_numpy(output_np.astype(np.float32))
# Ensure shape [channels, samples]
if output_tensor.ndim == 1: # Mono [samples] -> [1, samples]
output_tensor = output_tensor.unsqueeze(0)
elif output_tensor.ndim != 2: # Should only be 1D or 2D at this point
raise ValueError(f"Unexpected audio array dimension after speed factor: {output_tensor.ndim}.")
# Assuming 2D is already [channels, samples] - interpolation keeps channel dim first if input was 2D
# Add batch dimension -> [1, channels, samples]
output_tensor = output_tensor.unsqueeze(0)
output_tensor = output_tensor.contiguous()
# Final log & sanity checks
final_shape = output_tensor.shape; final_dtype = output_tensor.dtype
print(f"DiaTestGenerate: Final audio tensor shape: {final_shape}, dtype: {final_dtype}")
if len(final_shape) != 3: raise ValueError(f"Internal Error: Final tensor dim not 3! Shape: {final_shape}")
if final_shape[0] != 1: print(f"DiaTestGenerate: Warning - final batch size not 1: {final_shape[0]}")
if final_shape[1] == 0 or final_shape[2] == 0: raise ValueError(f"Internal Error: Final tensor zero dim! Shape: {final_shape}")
# Create dictionary for ComfyUI AUDIO output type
result = {'waveform': output_tensor, 'sample_rate': DEFAULT_SAMPLE_RATE}
return (result,) # Return dict inside tuple
except Exception as format_e:
print(f"DiaTestGenerate: Error formatting output: {format_e}")
traceback.print_exc()
raise format_e
# --- End Output Formatting ---
except Exception as e:
print(f"DiaTestGenerate: Error during generation: {e}")
traceback.print_exc()
raise e
# --- Node Mappings ---
NODE_CLASS_MAPPINGS = {
"DiaGenerate": DiaGenerate,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DiaGenerate": "Dia TTS Generate",
}
# --- Print message on load ---
print("### Loading: ComfyUI-DiaTest Nodes ###")
requirements.txt
+4
View File
@@ -0,0 +1,4 @@
# ComfyUI-DiaTest/requirements.txt
descript-audio-codec
huggingface_hub