mkv-episode-matcher 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

mkv_episode_matcher/__main__.py

@@ -108,75 +108,53 @@ def main():
     )
     args = parser.parse_args()
     if args.check_gpu:
-        from mkv_episode_matcher.speech_to_text import check_gpu_support
+        from mkv_episode_matcher.utils import check_gpu_support
         check_gpu_support()
         return
     logger.debug(f"Command-line arguments: {args}")
-    open_subtitles_api_key = ""
-    open_subtitles_user_agent = ""
-    open_subtitles_username = ""
-    open_subtitles_password = ""
-    # Check if API key is provided via command-line argument
-    tmdb_api_key = args.tmdb_api_key
-
-    # If API key is not provided, try to get it from the cache
-    if not tmdb_api_key:
-        cached_config = get_config(CONFIG_FILE)
-        if cached_config:
-            tmdb_api_key = cached_config.get("tmdb_api_key")
 
-    # If API key is still not available, prompt the user to input it
+    # Load configuration once
+    config = get_config(CONFIG_FILE)
+
+    # Get TMDb API key
+    tmdb_api_key = args.tmdb_api_key or config.get("tmdb_api_key")
     if not tmdb_api_key:
         tmdb_api_key = input("Enter your TMDb API key: ")
-        # Cache the API key
-
     logger.debug(f"TMDb API Key: {tmdb_api_key}")
+    
     logger.debug("Getting OpenSubtitles API key")
-    cached_config = get_config(CONFIG_FILE)
-    try:
-        open_subtitles_api_key = cached_config.get("open_subtitles_api_key")
-        open_subtitles_user_agent = cached_config.get("open_subtitles_user_agent")
-        open_subtitles_username = cached_config.get("open_subtitles_username")
-        open_subtitles_password = cached_config.get("open_subtitles_password")
-    except:
-        pass
+    open_subtitles_api_key = config.get("open_subtitles_api_key")
+    open_subtitles_user_agent = config.get("open_subtitles_user_agent")
+    open_subtitles_username = config.get("open_subtitles_username")
+    open_subtitles_password = config.get("open_subtitles_password")
+    
     if args.get_subs:
         if not open_subtitles_api_key:
             open_subtitles_api_key = input("Enter your OpenSubtitles API key: ")
-
         if not open_subtitles_user_agent:
             open_subtitles_user_agent = input("Enter your OpenSubtitles User Agent: ")
-
         if not open_subtitles_username:
             open_subtitles_username = input("Enter your OpenSubtitles Username: ")
-
         if not open_subtitles_password:
             open_subtitles_password = input("Enter your OpenSubtitles Password: ")
-
-    # If show directory is provided via command-line argument, use it
-    show_dir = args.show_dir
+    
+    # Use config for show directory and tesseract path
+    show_dir = args.show_dir or config.get("show_dir")
+    if not show_dir:
+        show_dir = input("Enter the main directory of the show:")
+    logger.info(f"Show Directory: {show_dir}")
     if not show_dir:
-        show_dir = cached_config.get("show_dir")
-        if not show_dir:
-            # If show directory is not provided, prompt the user to input it
-            show_dir = input("Enter the main directory of the show:")
-        logger.info(f"Show Directory: {show_dir}")
-        # if the user does not provide a show directory, make the default show directory the current working directory
-        if not show_dir:
-            show_dir = os.getcwd()
+        show_dir = os.getcwd()
+    
     if not args.tesseract_path:
-        tesseract_path = cached_config.get("tesseract_path")
-
+        tesseract_path = config.get("tesseract_path")
         if not tesseract_path:
-            tesseract_path = input(
-                r"Enter the path to the tesseract executable: ['C:\Program Files\Tesseract-OCR\tesseract.exe']"
-            )
-
+            tesseract_path = input(r"Enter the path to the tesseract executable: ['C:\Program Files\Tesseract-OCR\tesseract.exe']")
     else:
         tesseract_path = args.tesseract_path
     logger.debug(f"Teesseract Path: {tesseract_path}")
     logger.debug(f"Show Directory: {show_dir}")
-
+    
     # Set the configuration
     set_config(
         tmdb_api_key,

mkv_episode_matcher/episode_identification.py

@@ -18,7 +18,7 @@ class EpisodeMatcher:
         self.cache_dir = Path(cache_dir)
         self.min_confidence = min_confidence
         self.show_name = show_name
-        self.chunk_duration = 300  # 5 minutes
+        self.chunk_duration = 30
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.temp_dir = Path(tempfile.gettempdir()) / "whisper_chunks"
         self.temp_dir.mkdir(exist_ok=True)
@@ -44,7 +44,9 @@ class EpisodeMatcher:
                 '-ss', str(start_time),
                 '-t', str(self.chunk_duration),
                 '-i', mkv_file,
-                '-vn',
+                '-vn',  # Disable video
+                '-sn',  # Disable subtitles
+                '-dn',  # Disable data streams
                 '-acodec', 'pcm_s16le',
                 '-ar', '16000',
                 '-ac', '1',
@@ -80,31 +82,73 @@ class EpisodeMatcher:
         except Exception as e:
             logger.error(f"Error loading reference chunk from {srt_file}: {e}")
             return ''
-
-    def identify_episode(self, video_file, temp_dir, season_number):
-        try:
-            # Get video duration
-            duration = float(subprocess.check_output([
-                'ffprobe', '-v', 'error',
-                '-show_entries', 'format=duration',
-                '-of', 'default=noprint_wrappers=1:nokey=1',
-                video_file
-            ]).decode())
+    def _try_match_with_model(self, video_file, model_name, max_duration, reference_files):
+        """
+        Attempt to match using specified model, checking multiple 30-second chunks up to max_duration.
+        
+        Args:
+            video_file: Path to the video file
+            model_name: Name of the Whisper model to use
+            max_duration: Maximum duration in seconds to check
+            reference_files: List of reference subtitle files
+        """
+        # Use cached model
+        model = get_whisper_model(model_name, self.device)
+        
+        # Calculate number of chunks to check (30 seconds each)
+        num_chunks = max_duration // self.chunk_duration
+        
+        for chunk_idx in range(num_chunks):
+            start_time = chunk_idx * self.chunk_duration
+            logger.debug(f"Trying {model_name} model at {start_time} seconds")
             
-            total_chunks = int(np.ceil(duration / self.chunk_duration))
+            audio_path = self.extract_audio_chunk(video_file, start_time)
             
-            # Load Whisper model
-            model = whisper.load_model("base", device=self.device)
+            result = model.transcribe(
+                audio_path,
+                task="transcribe",
+                language="en"
+            )
             
-            # Get season-specific reference files using multiple patterns
-            reference_dir = self.cache_dir / "data" / self.show_name
+            chunk_text = result["text"]
+            best_confidence = 0
+            best_match = None
             
-            # Create season patterns for different formats
+            # Compare with reference chunks
+            for ref_file in reference_files:
+                ref_text = self.load_reference_chunk(ref_file, chunk_idx)
+                confidence = self.chunk_score(chunk_text, ref_text)
+                
+                if confidence > best_confidence:
+                    best_confidence = confidence
+                    best_match = ref_file
+                    
+                if confidence > self.min_confidence:
+                    season_ep = re.search(r'S(\d+)E(\d+)', best_match.stem)
+                    if season_ep:
+                        season, episode = map(int, season_ep.groups())
+                        return {
+                            'season': season,
+                            'episode': episode,
+                            'confidence': best_confidence,
+                            'reference_file': str(best_match),
+                            'matched_at': start_time
+                        }
+            
+            logger.debug(f"No match found at {start_time} seconds (best confidence: {best_confidence:.2f})")
+        
+        return None
+
+    def identify_episode(self, video_file, temp_dir, season_number):
+        """Progressive episode identification with faster initial attempt."""
+        try:
+            # Get reference files first
+            reference_dir = self.cache_dir / "data" / self.show_name
             patterns = [
-                f"S{season_number:02d}E",  # S01E01
-                f"S{season_number}E",      # S1E01
-                f"{season_number:02d}x",   # 01x01
-                f"{season_number}x",       # 1x01
+                f"S{season_number:02d}E",
+                f"S{season_number}E",
+                f"{season_number:02d}x",
+                f"{season_number}x",
             ]
             
             reference_files = []
@@ -114,55 +158,43 @@ class EpisodeMatcher:
                         for p in patterns)]
                 reference_files.extend(files)
             
-            # Remove duplicates while preserving order
             reference_files = list(dict.fromkeys(reference_files))
             
             if not reference_files:
                 logger.error(f"No reference files found for season {season_number}")
                 return None
-                
-            # Process chunks until match found
-            for chunk_idx in range(min(3, total_chunks)):  # Only try first 3 chunks
-                start_time = chunk_idx * self.chunk_duration
-                audio_path = self.extract_audio_chunk(video_file, start_time)
-                
-                # Transcribe chunk
-                result = model.transcribe(
-                    audio_path,
-                    task="transcribe",
-                    language="en"
-                )
-                
-                chunk_text = result["text"]
-                best_confidence = 0
-                best_match = None
-                
-                # Compare with reference chunks
-                for ref_file in reference_files:
-                    ref_text = self.load_reference_chunk(ref_file, chunk_idx)
-                    confidence = self.chunk_score(chunk_text, ref_text)
-                    
-                    if confidence > best_confidence:
-                        best_confidence = confidence
-                        best_match = ref_file
-                        
-                    if confidence > self.min_confidence:
-                        season_ep = re.search(r'S(\d+)E(\d+)', best_match.stem)
-                        if season_ep:
-                            season, episode = map(int, season_ep.groups())
-                            return {
-                                'season': season,
-                                'episode': episode,
-                                'confidence': best_confidence,
-                                'reference_file': str(best_match),
-                            }
+
+            # Try with tiny model first (fastest) - check first 2 minutes
+            logger.info("Attempting match with tiny model (first 2 minutes)...")
+            match = self._try_match_with_model(video_file, "tiny", 120, reference_files)
+            if match and match['confidence'] > 0.65:  # Slightly lower threshold for tiny
+                logger.info(f"Successfully matched with tiny model at {match['matched_at']}s (confidence: {match['confidence']:.2f})")
+                return match
+            
+            # If unsuccessful with tiny, try base model on first 3 minutes
+            logger.info("Tiny model match failed, trying base model (first 3 minutes)...")
+            match = self._try_match_with_model(video_file, "base", 180, reference_files)
+            if match and match['confidence'] > self.min_confidence:
+                logger.info(f"Successfully matched with base model at {match['matched_at']}s (confidence: {match['confidence']:.2f})")
+                return match
             
+            # If still no match, try base model on up to 10 minutes
+            logger.info("No match in first 3 minutes, extending base model search to 10 minutes...")
+            match = self._try_match_with_model(video_file, "base", 600, reference_files)
+            if match:
+                logger.info(f"Successfully matched with base model at {match['matched_at']}s (confidence: {match['confidence']:.2f})")
+                return match
+            
+            logger.info("Speech recognition match failed")
             return None
             
         finally:
             # Cleanup temp files
             for file in self.temp_dir.glob("chunk_*.wav"):
-                file.unlink()
+                try:
+                    file.unlink()
+                except Exception as e:
+                    logger.warning(f"Failed to delete temp file {file}: {e}")
 
 def detect_file_encoding(file_path):
     """
@@ -279,4 +311,19 @@ class SubtitleReader:
                 logger.warning(f"Error parsing subtitle block: {e}")
                 continue
                 
-        return text_lines
+        return text_lines
+    
+_whisper_models = {}
+
+def get_whisper_model(model_name="tiny", device=None):
+    """Cache whisper models to avoid reloading."""
+    global _whisper_models
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+    key = f"{model_name}_{device}"
+    if key not in _whisper_models:
+        _whisper_models[key] = whisper.load_model(model_name, device=device)
+        logger.info(f"Loaded {model_name} model on {device}")
+    
+    return _whisper_models[key]

mkv_episode_matcher/episode_matcher.py

@@ -19,7 +19,6 @@ from mkv_episode_matcher.utils import (
     process_srt_files,
     compare_and_rename_files,get_valid_seasons,rename_episode_file
 )
-from mkv_episode_matcher.speech_to_text import process_speech_to_text
 from mkv_episode_matcher.episode_identification import EpisodeMatcher
 
 def process_show(season=None, dry_run=False, get_subs=False):
@@ -32,7 +31,7 @@ def process_show(season=None, dry_run=False, get_subs=False):
     # Early check for reference files
     reference_dir = Path(CACHE_DIR) / "data" / show_name
     reference_files = list(reference_dir.glob("*.srt"))
-    if not reference_files:
+    if (not get_subs) and (not reference_files):
         logger.error(f"No reference subtitle files found in {reference_dir}")
         logger.info("Please download reference subtitles first")
         return
@@ -67,7 +66,7 @@ def process_show(season=None, dry_run=False, get_subs=False):
             if get_subs:
                 show_id = fetch_show_id(matcher.show_name)
                 if show_id:
-                    get_subtitles(show_id, seasons={season_num})
+                    get_subtitles(show_id, seasons={season_num}, config=config)
                     
             unmatched_files = []
             for mkv_file in mkv_files:
@@ -76,8 +75,6 @@ def process_show(season=None, dry_run=False, get_subs=False):
                 
                 if match:
                     new_name = f"{matcher.show_name} - S{match['season']:02d}E{match['episode']:02d}.mkv"
-                    new_path = os.path.join(season_path, new_name)
-                    
                     logger.info(f"Speech matched {os.path.basename(mkv_file)} to {new_name} "
                               f"(confidence: {match['confidence']:.2f})")
                     
@@ -105,4 +102,4 @@ def process_show(season=None, dry_run=False, get_subs=False):
         finally:
             if not dry_run:
                 shutil.rmtree(temp_dir)
-                cleanup_ocr_files(show_dir)
+                cleanup_ocr_files(show_dir)

mkv_episode_matcher/utils.py

@@ -2,7 +2,7 @@
 import os
 import re
 import shutil
-
+import torch
 import requests
 from loguru import logger
 from opensubtitlescom import OpenSubtitles
@@ -121,16 +121,17 @@ def rename_episode_file(original_file_path, new_filename):
         logger.error(f"Failed to rename file: {e}")
         return None
         
-def get_subtitles(show_id, seasons: set[int]):
+def get_subtitles(show_id, seasons: set[int], config=None):
     """
     Retrieves and saves subtitles for a given TV show and seasons.
 
     Args:
         show_id (int): The ID of the TV show.
         seasons (Set[int]): A set of season numbers for which subtitles should be retrieved.
+        config (Config object, optional): Preloaded configuration.
     """
-    logger.info(f"Getting subtitles for show ID {show_id}")
-    config = get_config(CONFIG_FILE)
+    if config is None:
+        config = get_config(CONFIG_FILE)
     show_dir = config.get("show_dir")
     series_name = sanitize_filename(os.path.basename(show_dir))
     tmdb_api_key = config.get("tmdb_api_key")
@@ -388,4 +389,11 @@ def compare_text(text1, text2):
 
     # Compare the two lists of text lines
     matching_lines = set(flat_text1).intersection(flat_text2)
-    return len(matching_lines)
+    return len(matching_lines)
+
+def check_gpu_support():
+    logger.info('Checking GPU support...')
+    if torch.cuda.is_available():
+        logger.info(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
+    else:
+        logger.warning("CUDA not available. Using CPU. Refer to https://pytorch.org/get-started/locally/ for GPU support.")

{mkv_episode_matcher-0.4.1.dist-info → mkv_episode_matcher-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: mkv-episode-matcher
-Version: 0.4.1
+Version: 0.5.0
 Summary: The MKV Episode Matcher is a tool for identifying TV series episodes from MKV files and renaming the files accordingly.
 Home-page: https://github.com/Jsakkos/mkv-episode-matcher
 Author: Jonathan Sakkos

{mkv_episode_matcher-0.4.1.dist-info → mkv_episode_matcher-0.5.0.dist-info}/RECORD

@@ -1,14 +1,13 @@
 mkv_episode_matcher/.gitattributes,sha256=Gh2-F2vCM7SZ01pX23UT8pQcmauXWfF3gwyRSb6ZAFs,66
 mkv_episode_matcher/__init__.py,sha256=aNlpgTo1kHVrBcR3SH6wRmCgKu8KjNTki1ZvFfAud6s,240
-mkv_episode_matcher/__main__.py,sha256=0ZwFXw30744lZ8HIMzUN4F2abpTRCT4glBCM9XQdD3k,7328
+mkv_episode_matcher/__main__.py,sha256=swYnLA2T8hvYMuNmK-EVRPBYzUdMxLP7pb8vxLuAnmc,6508
 mkv_episode_matcher/config.py,sha256=zDDKBcsDt5fME9BRqiTi7yWKeast1pZh36BNYMvIBYM,2419
-mkv_episode_matcher/episode_identification.py,sha256=_6M1UJkq1RGfmLI32u9dNOVvgp5Vf2MjqW2MTx0Gl8E,10329
-mkv_episode_matcher/episode_matcher.py,sha256=vunYpHQxyXo3l88BUScXa7_kMYMCV1pXpQxaLa-plZA,4325
+mkv_episode_matcher/episode_identification.py,sha256=jpDWvb16YAHNUzn9fuiHNJ_TB9EYmNg1ahdp361zSf4,12671
+mkv_episode_matcher/episode_matcher.py,sha256=BjPdPQwEHJWx_EOqj_AjKTsEFumdWHGNh7ERP-gfJ2g,4204
 mkv_episode_matcher/mkv_to_srt.py,sha256=4yxBHRVhgVby0UtQ2aTXGuoQpid8pkgjMIaHU6GCdzc,10857
-mkv_episode_matcher/speech_to_text.py,sha256=wVDrFFR7oASGMyq5cfOWmInEIeU9b3MPCLs9EyJrOMw,3128
 mkv_episode_matcher/subtitle_utils.py,sha256=rYSbd393pKYQW0w4sXgals02WFGqMYYYkQHDbEkWF8c,2666
 mkv_episode_matcher/tmdb_client.py,sha256=LbMCgjmp7sCbrQo_CDlpcnryKPz5S7inE24YY9Pyjk4,4172
-mkv_episode_matcher/utils.py,sha256=VASbougN3rb2iu40iZWkGjKIbahW713TOrFBo_TR9wo,14269
+mkv_episode_matcher/utils.py,sha256=bw2-cQsA4tdL9E1HNVTBuCkjXWDYR1And_1k2_BqdMg,14651
 mkv_episode_matcher/libraries/pgs2srt/.gitignore,sha256=mt3uxWYZaFurMw_yGE258gWhtGKPVR7e3Ll4ALJpyj4,23
 mkv_episode_matcher/libraries/pgs2srt/README.md,sha256=olb25G17tj0kxPgp_LcH5I2QWXjgP1m8JFyjYRGz4UU,1374
 mkv_episode_matcher/libraries/pgs2srt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -19,8 +18,8 @@ mkv_episode_matcher/libraries/pgs2srt/requirements.txt,sha256=sg87dqWw_qpbwciw-M
 mkv_episode_matcher/libraries/pgs2srt/Libraries/SubZero/SubZero.py,sha256=geT1LXdVd8yED9zoJ9K1XfP2JzGcM7u1SslHYrJI09o,10061
 mkv_episode_matcher/libraries/pgs2srt/Libraries/SubZero/post_processing.py,sha256=GKtVy_Lxv-z27mkRG8pJF2znKWXwZTot7jL6kN-zIxM,10503
 mkv_episode_matcher/libraries/pgs2srt/Libraries/SubZero/dictionaries/data.py,sha256=AlJHUYXl85J95OzGRik-AHVfzDd7Q8BJCvD4Nr8kRIk,938598
-mkv_episode_matcher-0.4.1.dist-info/METADATA,sha256=IWKnIxGcG_W1EEtPfYasyMCfyZL7-NL_ScYfNl0HVlg,5579
-mkv_episode_matcher-0.4.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-mkv_episode_matcher-0.4.1.dist-info/entry_points.txt,sha256=IglJ43SuCZq2eQ3shMFILCkmQASJHnDCI3ogohW2Hn4,64
-mkv_episode_matcher-0.4.1.dist-info/top_level.txt,sha256=XRLbd93HUaedeWLtkyTvQjFcE5QcBRYa3V-CfHrq-OI,20
-mkv_episode_matcher-0.4.1.dist-info/RECORD,,
+mkv_episode_matcher-0.5.0.dist-info/METADATA,sha256=3U2-ciHxqaP2BvtZ-awK5siPcCE9nFjopBID493NgBs,5579
+mkv_episode_matcher-0.5.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+mkv_episode_matcher-0.5.0.dist-info/entry_points.txt,sha256=IglJ43SuCZq2eQ3shMFILCkmQASJHnDCI3ogohW2Hn4,64
+mkv_episode_matcher-0.5.0.dist-info/top_level.txt,sha256=XRLbd93HUaedeWLtkyTvQjFcE5QcBRYa3V-CfHrq-OI,20
+mkv_episode_matcher-0.5.0.dist-info/RECORD,,

mkv_episode_matcher/speech_to_text.py

@@ -1,96 +0,0 @@
-# mkv_episode_matcher/speech_to_text.py
-
-import os
-import subprocess
-from pathlib import Path
-import whisper
-import torch
-from loguru import logger
-
-def process_speech_to_text(mkv_file, output_dir):
-    """
-    Convert MKV file to transcript using Whisper.
-    
-    Args:
-        mkv_file (str): Path to MKV file
-        output_dir (str): Directory to save transcript files
-    """
-    # Extract audio if not already done
-    wav_file = extract_audio(mkv_file, output_dir)
-    if not wav_file:
-        return None
-
-    # Load model
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    if device == "cuda":
-        logger.info(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
-    else:
-        logger.info("CUDA not available. Using CPU.")
-    
-    model = whisper.load_model("base", device=device)
-    
-    # Generate transcript
-    segments_file = os.path.join(output_dir, f"{Path(mkv_file).stem}.segments.json")
-    if not os.path.exists(segments_file):
-        try:
-            result = model.transcribe(
-                wav_file,
-                task="transcribe",
-                language="en",
-            )
-            
-            # Save segments
-            import json
-            with open(segments_file, 'w', encoding='utf-8') as f:
-                json.dump(result["segments"], f, indent=2)
-                
-            logger.info(f"Transcript saved to {segments_file}")
-            
-        except Exception as e:
-            logger.error(f"Error during transcription: {e}")
-            return None
-    else:
-        logger.info(f"Using existing transcript: {segments_file}")
-    
-    return segments_file
-
-def extract_audio(mkv_file, output_dir):
-    """
-    Extract audio from MKV file using FFmpeg.
-    
-    Args:
-        mkv_file (str): Path to MKV file
-        output_dir (str): Directory to save WAV file
-        
-    Returns:
-        str: Path to extracted WAV file
-    """
-    wav_file = os.path.join(output_dir, f"{Path(mkv_file).stem}.wav")
-    
-    if not os.path.exists(wav_file):
-        logger.info(f"Extracting audio from {mkv_file}")
-        try:
-            cmd = [
-                'ffmpeg',
-                '-i', mkv_file,
-                '-vn',  # Disable video
-                '-acodec', 'pcm_s16le',  # Convert to PCM format
-                '-ar', '16000',  # Set sample rate to 16kHz
-                '-ac', '1',  # Convert to mono
-                wav_file
-            ]
-            subprocess.run(cmd, check=True, capture_output=True)
-            logger.info(f"Audio extracted to {wav_file}")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"Error extracting audio: {e}")
-            return None
-    else:
-        logger.info(f"Audio file {wav_file} already exists, skipping extraction")
-    
-    return wav_file
-def check_gpu_support():
-    logger.info('Checking GPU support...')
-    if torch.cuda.is_available():
-        logger.info(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
-    else:
-        logger.warning("CUDA not available. Using CPU. Refer to https://pytorch.org/get-started/locally/ for GPU support.")