Spaces:
Running
Running
| import os | |
| import time | |
| import logging | |
| import torch | |
| import numpy as np | |
| import soundfile as sf | |
| from pathlib import Path | |
| from typing import Optional | |
| from dia.model import Dia | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Constants | |
| DEFAULT_SAMPLE_RATE = 44100 | |
| DEFAULT_MODEL_NAME = "nari-labs/Dia-1.6B" | |
| # Global model instance (lazy loaded) | |
| _model = None | |
| def _get_model() -> Dia: | |
| """Lazy-load the Dia model to avoid loading it until needed""" | |
| global _model | |
| if _model is None: | |
| logger.info("Loading Dia model...") | |
| try: | |
| # Check if torch is available with correct version | |
| logger.info(f"PyTorch version: {torch.__version__}") | |
| logger.info(f"CUDA available: {torch.cuda.is_available()}") | |
| if torch.cuda.is_available(): | |
| logger.info(f"CUDA version: {torch.version.cuda}") | |
| logger.info(f"GPU device: {torch.cuda.get_device_name(0)}") | |
| # Check if model path exists | |
| logger.info(f"Attempting to load model from: {DEFAULT_MODEL_NAME}") | |
| # Load the model with detailed logging | |
| logger.info("Initializing Dia model...") | |
| _model = Dia.from_pretrained(DEFAULT_MODEL_NAME, compute_dtype="float16") | |
| # Log model details | |
| logger.info(f"Dia model loaded successfully") | |
| logger.info(f"Model type: {type(_model).__name__}") | |
| logger.info(f"Model device: {next(_model.parameters()).device}") | |
| except ImportError as import_err: | |
| logger.error(f"Import error loading Dia model: {import_err}") | |
| logger.error(f"This may indicate missing dependencies") | |
| raise | |
| except FileNotFoundError as file_err: | |
| logger.error(f"File not found error loading Dia model: {file_err}") | |
| logger.error(f"Model path may be incorrect or inaccessible") | |
| raise | |
| except Exception as e: | |
| logger.error(f"Error loading Dia model: {e}", exc_info=True) | |
| logger.error(f"Error type: {type(e).__name__}") | |
| logger.error(f"This may indicate incompatible versions or missing CUDA support") | |
| raise | |
| return _model | |
| def generate_speech(text: str, language: str = "zh") -> str: | |
| """Public interface for TTS generation using Dia model | |
| Args: | |
| text (str): Input text to synthesize | |
| language (str): Language code (not used in Dia model, kept for API compatibility) | |
| Returns: | |
| str: Path to the generated audio file | |
| """ | |
| logger.info(f"Generating speech for text length: {len(text)}") | |
| logger.info(f"Text content (first 50 chars): {text[:50]}...") | |
| try: | |
| # Create output directory if it doesn't exist | |
| output_dir = "temp/outputs" | |
| logger.info(f"Ensuring output directory exists: {output_dir}") | |
| try: | |
| os.makedirs(output_dir, exist_ok=True) | |
| logger.info(f"Output directory ready: {output_dir}") | |
| except PermissionError as perm_err: | |
| logger.error(f"Permission error creating output directory: {perm_err}") | |
| raise | |
| except Exception as dir_err: | |
| logger.error(f"Error creating output directory: {dir_err}") | |
| raise | |
| # Generate unique output path | |
| timestamp = int(time.time()) | |
| output_path = f"{output_dir}/output_{timestamp}.wav" | |
| logger.info(f"Output will be saved to: {output_path}") | |
| # Get the model | |
| logger.info("Retrieving Dia model instance") | |
| try: | |
| model = _get_model() | |
| logger.info("Successfully retrieved Dia model instance") | |
| except Exception as model_err: | |
| logger.error(f"Failed to get Dia model: {model_err}") | |
| logger.error(f"Error type: {type(model_err).__name__}") | |
| raise | |
| # Generate audio | |
| logger.info("Starting audio generation with Dia model") | |
| start_time = time.time() | |
| try: | |
| with torch.inference_mode(): | |
| logger.info("Calling model.generate() with inference_mode") | |
| output_audio_np = model.generate( | |
| text, | |
| max_tokens=None, # Use default from model config | |
| cfg_scale=3.0, | |
| temperature=1.3, | |
| top_p=0.95, | |
| cfg_filter_top_k=35, | |
| use_torch_compile=False, # Keep False for stability | |
| verbose=False | |
| ) | |
| logger.info("Model.generate() completed") | |
| except RuntimeError as rt_err: | |
| logger.error(f"Runtime error during generation: {rt_err}") | |
| if "CUDA out of memory" in str(rt_err): | |
| logger.error("CUDA out of memory error - consider reducing batch size or model size") | |
| raise | |
| except Exception as gen_err: | |
| logger.error(f"Error during audio generation: {gen_err}") | |
| logger.error(f"Error type: {type(gen_err).__name__}") | |
| raise | |
| end_time = time.time() | |
| generation_time = end_time - start_time | |
| logger.info(f"Generation finished in {generation_time:.2f} seconds") | |
| # Process the output | |
| if output_audio_np is not None: | |
| logger.info(f"Generated audio array shape: {output_audio_np.shape}, dtype: {output_audio_np.dtype}") | |
| logger.info(f"Audio stats - min: {output_audio_np.min():.4f}, max: {output_audio_np.max():.4f}, mean: {output_audio_np.mean():.4f}") | |
| # Apply a slight slowdown for better quality (0.94x speed) | |
| speed_factor = 0.94 | |
| original_len = len(output_audio_np) | |
| target_len = int(original_len / speed_factor) | |
| logger.info(f"Applying speed adjustment factor: {speed_factor}") | |
| if target_len != original_len and target_len > 0: | |
| try: | |
| x_original = np.arange(original_len) | |
| x_resampled = np.linspace(0, original_len - 1, target_len) | |
| output_audio_np = np.interp(x_resampled, x_original, output_audio_np) | |
| logger.info(f"Resampled audio from {original_len} to {target_len} samples for {speed_factor:.2f}x speed") | |
| except Exception as resample_err: | |
| logger.error(f"Error during audio resampling: {resample_err}") | |
| logger.warning("Using original audio without resampling") | |
| # Save the audio file | |
| logger.info(f"Saving audio to file: {output_path}") | |
| try: | |
| sf.write(output_path, output_audio_np, DEFAULT_SAMPLE_RATE) | |
| logger.info(f"Audio successfully saved to {output_path}") | |
| except Exception as save_err: | |
| logger.error(f"Error saving audio file: {save_err}") | |
| logger.error(f"Error type: {type(save_err).__name__}") | |
| raise | |
| return output_path | |
| else: | |
| logger.warning("Generation produced no output (None returned from model)") | |
| logger.warning("This may indicate a model configuration issue or empty input text") | |
| dummy_path = f"{output_dir}/dummy_{timestamp}.wav" | |
| logger.warning(f"Returning dummy audio path: {dummy_path}") | |
| return dummy_path | |
| except Exception as e: | |
| logger.error(f"TTS generation failed: {str(e)}", exc_info=True) | |
| logger.error(f"Error type: {type(e).__name__}") | |
| # Log additional diagnostic information based on error type | |
| if isinstance(e, ImportError): | |
| logger.error(f"Import error - missing dependency: {e.__class__.__module__}.{e.__class__.__name__}") | |
| logger.error("Check if all required packages are installed correctly") | |
| elif isinstance(e, RuntimeError) and "CUDA" in str(e): | |
| logger.error("CUDA-related runtime error - check GPU compatibility and memory") | |
| elif isinstance(e, AttributeError): | |
| logger.error(f"Attribute error - likely API incompatibility or incorrect module version") | |
| if hasattr(e, '__traceback__'): | |
| tb = e.__traceback__ | |
| while tb.tb_next: | |
| tb = tb.tb_next | |
| logger.error(f"Error occurred in file: {tb.tb_frame.f_code.co_filename}, line {tb.tb_lineno}") | |
| elif isinstance(e, FileNotFoundError): | |
| logger.error(f"File not found - check if model files exist and are accessible") | |
| # Return dummy path in case of error | |
| return "temp/outputs/dummy.wav" |