Update modeling_spark_tts.py

modeling_spark_tts.py

@@ -3092,31 +3092,32 @@ class SparkTTSModel(PreTrainedModel, GenerationMixin):
         # Check for trust_remote_code - needed for config loading if custom code involved there too
         trust_remote_code = model_kwargs.pop("trust_remote_code", False) # Important
 
-        # Determine actual model directory (could be cache path)
-        if pretrained_model_name_or_path is not None:
-             resolved_model_path = Path(pretrained_model_name_or_path)
-             if not resolved_model_path.is_dir():
-                  # Attempt to download and resolve cache path if it's an ID
-                  # This requires internet connection if not cached
-                  try:
-                       resolved_model_path = Path(cached_file(
-                            pretrained_model_name_or_path,
-                            filename=cls.config_class.config_files[0], # e.g., "config.json"
-                            cache_dir=cache_dir,
-                            force_download=force_download,
-                            local_files_only=local_files_only,
-                            token=token,
-                            revision=revision,
-                            _raise_exceptions_for_missing_entries=False,
-                            _raise_exceptions_for_connection_errors=False,
-                       )).parent
-                  except Exception as e:
-                       logger.warning(f"Could not resolve cache path for {pretrained_model_name_or_path}: {e}. Assuming it's a local path.")
-                       resolved_model_path = Path(pretrained_model_name_or_path) # Fallback
-             if not resolved_model_path.is_dir():
-                  raise EnvironmentError(f"Cannot find model directory at {resolved_model_path}")
+
+        # NEW IMPROVED PATH RESOLUTION
+        from huggingface_hub import snapshot_download
+        import os
+        # Check if it's a local path first
+        if os.path.isdir(pretrained_model_name_or_path):
+            resolved_model_path = Path(pretrained_model_name_or_path)
         else:
-             raise ValueError("pretrained_model_name_or_path must be provided.")
+            # Try to get from Hugging Face Hub
+            try:
+                logger.info(f"Downloading/locating model from Hugging Face Hub: {pretrained_model_name_or_path}")
+                # This will download the model if needed and return the cached path
+                resolved_model_path = Path(snapshot_download(
+                    pretrained_model_name_or_path,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    local_files_only=local_files_only,
+                    token=token,
+                ))
+            except Exception as e:
+                logger.error(f"Error downloading model: {e}")
+                raise EnvironmentError(f"Failed to find or download model '{pretrained_model_name_or_path}': {e}")
+
+        if not resolved_model_path.is_dir():
+            raise EnvironmentError(f"Cannot find model directory at {resolved_model_path}")
 
 
         # Helper function to resolve paths relative to the main model directory