# ComfyUI-DiaTTS/dia_lib/model.py import time from enum import Enum import dac import numpy as np import torch import torchaudio from tqdm.auto import tqdm 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(): """Gets the default torch device.""" if torch.cuda.is_available(): return torch.device("cuda") if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def _sample_next_token(logits_BCxV, temperature, top_p, cfg_filter_top_k=None) -> torch.Tensor: """Samples the next token based on logits, temperature, and top_p.""" exec_device = logits_BCxV.device if temperature == 0.0: return torch.argmax(logits_BCxV, dim=-1) logits_BCxV = logits_BCxV / temperature if cfg_filter_top_k is not None: cfg_filter_top_k = min(cfg_filter_top_k, logits_BCxV.shape[-1]) if cfg_filter_top_k > 0: _, top_k_indices_BCxV = torch.topk(logits_BCxV, k=cfg_filter_top_k, dim=-1) mask = torch.ones_like(logits_BCxV, dtype=torch.bool, device=exec_device) 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, device=exec_device) 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) final_probs_BCxV = torch.nan_to_num(final_probs_BCxV) if torch.sum(final_probs_BCxV) == 0: final_probs_BCxV.fill_(1.0) final_probs_BCxV = final_probs_BCxV / final_probs_BCxV.shape[-1] sampled_indices_BC = torch.multinomial(final_probs_BCxV, num_samples=1) return sampled_indices_BC.squeeze(-1) class ComputeDtype(str, Enum): """Enum for compute dtypes.""" FLOAT32 = "float32"; FLOAT16 = "float16"; BFLOAT16 = "bfloat16" def to_dtype(self) -> torch.dtype: """Converts enum value to torch.dtype.""" return getattr(torch, self.value) class Dia: """Main class for Dia TTS model loading and generation.""" def __init__(self, config: DiaConfig, compute_dtype=ComputeDtype.FLOAT32, device=None): """Initializes the Dia model, ensuring placement on the specified device.""" 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.model.to(self.device) self.dac_model = None def _devices_equal(self, device1: torch.device, device2: torch.device) -> bool: """Robustly compares two torch.device objects.""" if device1 is None or device2 is None: return False if device1.type != device2.type: return False if device1.type == 'cuda': index1 = device1.index if device1.index is not None else torch.cuda.current_device() index2 = device2.index if device2.index is not None else torch.cuda.current_device() return index1 == index2 return True def _load_dac_model(self): """Loads the Descript Audio Codec model to the instance's device.""" if self.dac_model is not None: try: dac_device = next(self.dac_model.parameters()).device if not self._devices_equal(dac_device, self.device): print(f"Moving existing DAC model from {dac_device} to {self.device}...") self.dac_model.to(self.device) except StopIteration: print("DAC model has no parameters, forcing reload.") self.dac_model = None except Exception as e: print(f"Error checking/moving DAC model device: {e}. Reloading.") self.dac_model = None if self.dac_model is not None: return try: print(f"Loading DAC model to {self.device}...") dac_model_path = dac.utils.download() self.dac_model = dac.DAC.load(dac_model_path) self.dac_model.to(self.device) self.dac_model.eval() dac_device_check = next(self.dac_model.parameters()).device if not self._devices_equal(dac_device_check, self.device): print(f"Warning: DAC model loaded to {dac_device_check} instead of requested {self.device}. Retrying move.") self.dac_model.to(self.device) dac_device_check_retry = next(self.dac_model.parameters()).device if not self._devices_equal(dac_device_check_retry, self.device): raise RuntimeError(f"Failed to move DAC model to {self.device}") print(f"DAC model loaded successfully on {self.device}.") except Exception as e: self.dac_model = None raise RuntimeError("Failed to load DAC model") from e def _prepare_text_input(self, text: str) -> torch.Tensor: """Encodes text prompt, pads, and creates tensor on the model's device.""" text_pad_value = self.config.data.text_pad_value; max_len = self.config.data.text_length try: 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) except Exception as e: print(f"Error encoding text: {e}") raise ValueError("Failed to encode input text.") from e current_len = len(text_tokens) padding_needed = max_len - current_len if padding_needed < 0: print(f"Warning: Input text truncated from {current_len} to {max_len} bytes.") text_tokens = text_tokens[:max_len] padded_text_np = np.array(text_tokens, dtype=np.uint8) elif padding_needed == 0: padded_text_np = np.array(text_tokens, dtype=np.uint8) else: padded_text_np = np.pad(text_tokens, (0, padding_needed), 'constant', constant_values=text_pad_value).astype(np.uint8) return torch.from_numpy(padded_text_np).to(dtype=torch.long, device=self.device).unsqueeze(0) def encode_audio_prompt(self, waveform: torch.Tensor, sample_rate: int) -> torch.Tensor: """Encodes a raw waveform tensor into DAC codes for use as an audio prompt.""" if self.dac_model is None: self._load_dac_model() if waveform.device != self.device: waveform = waveform.to(self.device) if waveform.dtype != torch.float32: original_dtype = waveform.dtype waveform = waveform.to(torch.float32) if not torch.is_floating_point(original_dtype): max_val = torch.iinfo(original_dtype).max waveform = waveform / max_val if waveform.ndim == 2 and waveform.shape[0] == 1: waveform = waveform.unsqueeze(1) elif waveform.ndim == 1: waveform = waveform.unsqueeze(0).unsqueeze(0) elif waveform.ndim == 3 and waveform.shape[0] == 1 and waveform.shape[1] > 1: print(f"Dia.encode_audio_prompt: Input prompt has {waveform.shape[1]} channels. Averaging to mono.") waveform = torch.mean(waveform, dim=1, keepdim=True) elif not (waveform.ndim == 3 and waveform.shape[0] == 1 and waveform.shape[1] == 1): raise ValueError(f"Unsupported waveform shape for audio prompt: {waveform.shape}.") if sample_rate != DEFAULT_SAMPLE_RATE: print(f"Dia.encode_audio_prompt: Resampling prompt from {sample_rate} Hz to {DEFAULT_SAMPLE_RATE} Hz.") waveform = torchaudio.functional.resample(waveform, orig_freq=sample_rate, new_freq=DEFAULT_SAMPLE_RATE) waveform = torch.clamp(waveform, -1.0, 1.0) try: audio_data = self.dac_model.preprocess(waveform, DEFAULT_SAMPLE_RATE) _, codes, _, _, _ = self.dac_model.encode(audio_data) prompt_codes = codes.squeeze(0).transpose(0, 1).contiguous() return prompt_codes.to(torch.int) except Exception as e: print(f"Error during DAC encoding of prompt: {e}") import traceback traceback.print_exc() raise RuntimeError("Failed to encode audio prompt using DAC.") from e def _prepare_audio_prompt(self, audio_prompt_codes: torch.Tensor | None) -> tuple[torch.Tensor, int]: """Prepares the initial audio tokens (BOS and optional prompt codes) on the model's device.""" 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_codes is not None: if audio_prompt_codes.device != self.device: audio_prompt_codes = audio_prompt_codes.to(self.device) if not torch.is_tensor(audio_prompt_codes) or audio_prompt_codes.dtype not in [torch.int, torch.long]: audio_prompt_codes = audio_prompt_codes.to(torch.int) if audio_prompt_codes.ndim != 2 or audio_prompt_codes.shape[1] != num_channels: raise ValueError(f"Audio prompt codes have unexpected shape {audio_prompt_codes.shape}.") prompt_len = audio_prompt_codes.shape[0] prefill_step += prompt_len prefill = torch.cat([prefill, audio_prompt_codes], dim=0) print(f"Prepared audio prompt with {prompt_len} timesteps.") 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) # Pass device to build_delay_indices delay_precomp = build_delay_indices(B=1, T=prefill.shape[0], C=num_channels, delay_pattern=delay_pattern, device=self.device) prefill_delayed = apply_audio_delay(prefill.unsqueeze(0), audio_pad_value, audio_bos_value, delay_precomp).squeeze(0) return prefill_delayed, prefill_step def _prepare_generation(self, text: str, audio_prompt_codes: torch.Tensor | None, verbose: bool): """Prepares encoder/decoder states and initial inputs, ensuring device consistency.""" try: param_device = next(self.model.parameters()).device if not self._devices_equal(param_device, self.device): print(f"Warning: Model parameters on {param_device}, but expected {self.device}. Moving...") self.model.to(self.device) if not self._devices_equal(next(self.model.parameters()).device, self.device): raise RuntimeError(f"Failed to move model parameters to {self.device}") except StopIteration: raise RuntimeError("Model has no parameters!") 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) prefill, prefill_step = self._prepare_audio_prompt(audio_prompt_codes) if verbose: print(f"generate: data loaded. Prefill steps: {prefill_step}. Prefill tensor shape: {prefill.shape}") 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: if verbose: print(f"generate: Prefilling decoder state for {dec_step} steps...") 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) if verbose: print("generate: Decoder prefill complete.") return dec_state, dec_output def _decoder_step(self, tokens_Bx1xC, dec_state, cfg_scale, temperature, top_p, cfg_filter_top_k) -> torch.Tensor: """Performs a single autoregressive decoding step.""" 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 return _sample_next_token(logits_CxV.float(), temperature, top_p, cfg_filter_top_k) def _generate_output(self, generated_codes: torch.Tensor) -> np.ndarray: """Reverts delay pattern and decodes audio codes using DAC.""" if self.dac_model is None: self._load_dac_model() try: dac_device = next(self.dac_model.parameters()).device if not self._devices_equal(dac_device, self.device): self.dac_model.to(self.device) except StopIteration: raise RuntimeError("DAC model has no parameters after load attempt.") 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) # Pass the device of generated_codes to build_revert_indices codes_device = generated_codes.device revert_precomp = build_revert_indices(B=1, T=seq_length, C=num_channels, delay_pattern=delay_pattern, device=codes_device) codebook_delayed = generated_codes.unsqueeze(0) # Ensure generated codes are on the same device as precomputed indices # This shouldn't be necessary if build_revert_indices uses codes_device, but belts and braces if codebook_delayed.device != codes_device: codebook_delayed = codebook_delayed.to(codes_device) codebook_reverted = revert_audio_delay(codebook_delayed, audio_pad_value, revert_precomp, seq_length) codebook_trimmed = codebook_reverted[:, :-max_delay_pattern, :] min_valid_index = 0; max_valid_index = 1023 invalid_mask = (codebook_trimmed < min_valid_index) | (codebook_trimmed > max_valid_index) if torch.any(invalid_mask): print(f"Warning: Clamping {torch.sum(invalid_mask)} invalid codebook indices to 0.") codebook_trimmed[invalid_mask] = 0 # codebook_for_dac should now be on the correct device (codes_device) codebook_for_dac = codebook_trimmed.transpose(1, 2) # decode function will handle moving codes to dac_model's device if necessary audio_tensor = decode(self.dac_model, codebook_for_dac) return audio_tensor.squeeze().cpu().numpy() @torch.inference_mode() def generate(self, text: str, max_tokens=None, cfg_scale=3.0, temperature=1.3, top_p=0.95, use_torch_compile=False, cfg_filter_top_k=35, audio_prompt=None, verbose=False, pbar=None ) -> np.ndarray | None: """Generates audio waveform from text input, updating progress bars.""" 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 clamped_max_tokens = self.config.data.audio_length if max_tokens is None else min(max_tokens, self.config.data.audio_length) max_delay_pattern = max(delay_pattern) self.model.eval() 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 step_fn = torch.compile(self._decoder_step, mode="default") if use_torch_compile else self._decoder_step tqdm_pbar = None if verbose: tqdm_total = clamped_max_tokens - (dec_step + 1) if clamped_max_tokens > dec_step + 1 else 0 tqdm_pbar = tqdm(total=tqdm_total, desc="Dia Generating Tokens", unit="token", initial=0) while dec_step < clamped_max_tokens -1 : 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) is_last_generatable_step = (dec_step == clamped_max_tokens - max_delay_pattern - 1) if (not eos_detected and pred_C[0] == audio_eos_value) or is_last_generatable_step: if not eos_detected: if verbose: print(f"\nEOS detected at step {dec_step + 1}.") eos_detected = True; eos_countdown = max_delay_pattern if is_last_generatable_step and not eos_detected: if verbose: print(f"\nApproaching max_tokens ({clamped_max_tokens}), forcing EOS generation.") pred_C[0] = audio_eos_value 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, apply_mask=(bos_countdown > 0)) if eos_countdown == 0: if verbose: print("EOS padding complete. Stopping generation.") break dec_step += 1 if pbar: pbar.update(1) if tqdm_pbar: tqdm_pbar.update(1) if tqdm_pbar: tqdm_pbar.close() final_step_index = dec_step if dec_output.prefill_step > final_step_index: print("Dia: Warning - Nothing generated beyond prefill/prompt.") return None generated_codes = dec_output.generated_tokens[dec_output.prefill_step : final_step_index + 1, :] if generated_codes.shape[0] == 0: print("Dia: Warning - Generated codes array is empty.") return None if verbose: print(f"Generated {generated_codes.shape[0]} token steps.") return self._generate_output(generated_codes)