SYMBOL INDEX (88 symbols across 7 files)

FILE: app.py
  function set_seed (line 57) | def set_seed(seed: int):
  function run_inference (line 69) | def run_inference(

FILE: cli.py
  function set_seed (line 12) | def set_seed(seed: int):
  function main (line 25) | def main():

FILE: dia/audio.py
  function build_delay_indices (line 6) | def build_delay_indices(B: int, T: int, C: int, delay_pattern: tp.List[i...
  function apply_audio_delay (line 44) | def apply_audio_delay(
  function build_revert_indices (line 88) | def build_revert_indices(B: int, T: int, C: int, delay_pattern: tp.List[...
  function revert_audio_delay (line 125) | def revert_audio_delay(

FILE: dia/config.py
  class EncoderConfig (line 21) | class EncoderConfig(BaseModel, frozen=True):
  class DecoderConfig (line 57) | class DecoderConfig(BaseModel, frozen=True):
  class DiaConfig (line 103) | class DiaConfig(BaseModel, frozen=True):
    method save (line 140) | def save(self, path: str) -> None:
    method load (line 157) | def load(cls, path: str) -> "DiaConfig | None":

FILE: dia/layers.py
  function _normalize_axes (line 12) | def _normalize_axes(axes: tuple[int, ...], ndim: int) -> tuple[int, ...]:
  class DenseGeneral (line 16) | class DenseGeneral(nn.Module):
    method __init__ (line 32) | def __init__(
    method forward (line 49) | def forward(self, inputs: Tensor) -> Tensor:
  class MlpBlock (line 61) | class MlpBlock(nn.Module):
    method __init__ (line 64) | def __init__(self, embed_dim: int, intermediate_dim: int, compute_dtyp...
    method forward (line 82) | def forward(self, x: torch.Tensor) -> torch.Tensor:
  class RotaryEmbedding (line 95) | class RotaryEmbedding(nn.Module):
    method __init__ (line 98) | def __init__(
    method forward (line 118) | def forward(self, inputs: torch.Tensor, position: torch.Tensor):
    method apply_rope (line 132) | def apply_rope(self, inputs: torch.Tensor, sin: torch.Tensor, cos: tor...
  function custom_scaled_dot_product_attention (line 139) | def custom_scaled_dot_product_attention(
  class CrossAttention (line 192) | class CrossAttention(nn.Module):
    method __init__ (line 195) | def __init__(
    method forward (line 249) | def forward(
  class FusedQKV (line 313) | class FusedQKV(nn.Module):
    method __init__ (line 314) | def __init__(
    method forward (line 333) | def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.T...
  class SelfAttention (line 345) | class SelfAttention(nn.Module):
    method __init__ (line 348) | def __init__(
    method get_linear_weight (line 406) | def get_linear_weight(self, dense: DenseGeneral):
    method patch_fused_qkv (line 420) | def patch_fused_qkv(self):
    method forward (line 439) | def forward(
  class EncoderLayer (line 531) | class EncoderLayer(nn.Module):
    method __init__ (line 534) | def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
    method forward (line 567) | def forward(
  class Encoder (line 591) | class Encoder(nn.Module):
    method __init__ (line 594) | def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
    method forward (line 612) | def forward(
  class DecoderLayer (line 626) | class DecoderLayer(nn.Module):
    method __init__ (line 629) | def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
    method forward (line 684) | def forward(
  class Decoder (line 730) | class Decoder(nn.Module):
    method __init__ (line 733) | def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
    method precompute_cross_attn_cache (line 763) | def precompute_cross_attn_cache(
    method decode_step (line 784) | def decode_step(
    method forward (line 819) | def forward(self, tgt_ids_BxTxC: torch.Tensor, state: DecoderInference...
  class DiaModel (line 869) | class DiaModel(
    method __init__ (line 884) | def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):

FILE: dia/model.py
  function _get_default_device (line 20) | def _get_default_device():
  function _sample_next_token (line 28) | def _sample_next_token(
  class ComputeDtype (line 79) | class ComputeDtype(str, Enum):
    method to_dtype (line 84) | def to_dtype(self) -> torch.dtype:
  class Dia (line 95) | class Dia:
    method __init__ (line 96) | def __init__(
    method from_local (line 132) | def from_local(
    method from_pretrained (line 177) | def from_pretrained(
    method _load_dac_model (line 221) | def _load_dac_model(self):
    method _encode_text (line 240) | def _encode_text(self, text: str) -> torch.Tensor:
    method _pad_text_input (line 265) | def _pad_text_input(self, text_tokens: list[torch.Tensor]) -> torch.Te...
    method _prepare_audio_prompt (line 282) | def _prepare_audio_prompt(self, audio_prompts: list[torch.Tensor | Non...
    method _prepare_generation (line 343) | def _prepare_generation(
    method _decoder_step (line 399) | def _decoder_step(
    method _generate_output (line 469) | def _generate_output(self, generated_codes: torch.Tensor, lengths_Bx: ...
    method _encode (line 528) | def _encode(self, audio: torch.Tensor) -> torch.Tensor:
    method _decode (line 540) | def _decode(self, audio_codes: torch.Tensor) -> torch.Tensor:
    method load_audio (line 550) | def load_audio(self, audio_path: str) -> torch.Tensor:
    method save_audio (line 579) | def save_audio(self, path: str, audio: np.ndarray):
    method generate (line 594) | def generate(

FILE: dia/state.py
  function create_attn_mask (line 9) | def create_attn_mask(
  class EncoderInferenceState (line 43) | class EncoderInferenceState:
    method new (line 53) | def new(cls, config: DiaConfig, cond_src: torch.Tensor) -> "EncoderInf...
  class KVCache (line 72) | class KVCache(torch.nn.Module):
    method __init__ (line 76) | def __init__(
    method from_kv (line 95) | def from_kv(cls, k: torch.Tensor, v: torch.Tensor) -> "KVCache":
    method update (line 107) | def update(self, k: torch.Tensor, v: torch.Tensor, current_idx: torch....
    method prefill (line 113) | def prefill(self, k: torch.Tensor, v: torch.Tensor):
  class DecoderInferenceState (line 120) | class DecoderInferenceState:
    method new (line 134) | def new(
    method prepare_step (line 177) | def prepare_step(self, step_from: int, step_to: int | None = None) -> ...
  class DecoderOutput (line 184) | class DecoderOutput:
    method new (line 189) | def new(cls, batch_size: int, config: DiaConfig, device: torch.device)...
    method get_tokens_at (line 201) | def get_tokens_at(self, step_from: int, step_to: int | None = None) ->...
    method update_one (line 206) | def update_one(self, dec_out: torch.Tensor, step: int, apply_mask: boo...
    method prefill (line 214) | def prefill(self, dec_out: torch.Tensor, prefill_steps: list[int]):