Fixed typing_extensions requirement #184

Fixed requirements check removes torch CUDA packages #184
Improved compatibility with Python 3.11 #191

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 All notable changes to this project will be documented in this file, starting with version 0.17.7.
 
+## [0.24.3] - 2024-08-28
+This version improves compatibility with Python 3.11 and fixes a few package-related issues.
+We're slowly moving towards Python 3.11 as the default version for the tool.
+
+### Fixes
+- Improved compatibility with Python 3.11
+- Fixed a bug that was removing CUDA support when the tool was checking if the required packages were installed
+- Other package-related bug fixes
+
 ## [0.24.2] - 2024-06-01
 
 ### Added

INSTALLATION.md

@@ -31,11 +31,11 @@ scenario, some stuff might not work at all on your computer, and you'll need pro
 
 #### Requirements
 
-Our installations are on MacOS 12.6+ running on M1 and Windows 10 machines in our editing room, 
+Our installations are on MacOS 12.6+, MacOS 14.4.1+ running on M1, and Windows 10 machines in our editing room, 
 but the scripts should run fine on other CPUs and GPUs. 
-For both production and development we're currently using Python 3.10.11. 
+For both production and development we're mostly using Python 3.10.11, but slowly migrating to 3.11.
 
-_Note: The tool worked fine on Python 3.9, but some packages are now optimized for Python 3.10. 
+_Note: The tool worked fine on Python 3.9, but some packages are now optimized for Python 3.10 and Python 3.11 
 Python 3.9 support will no longer be possible in the very near future._
 
 **The Resolve API integration only works on Resolve Studio 18 (not on the free version, and certainly not earlier 
@@ -74,6 +74,9 @@ So, install that using `xcode-select --install`.
     brew install ffmpeg
     brew install rust
 
+_Note: we're slowly migrating to Python 3.11, so please keep an eye on the requirements in the future. 
+If you're feeling brave, you can try installing Python 3.11 instead of 3.10 by modifying the above and below commands._
+
 #### 3. Download StoryToolkitAI:
 
 Download from GitHub using this command:
@@ -129,6 +132,9 @@ Download the latest Python 3.10 version from [the official Python website](https
 _Note: Do not use the Python installers available from the Windows Store. Only use other Python installers / versions
 if you know what you're doing._
 
+_Note 2: we're slowly migrating to Python 3.11, so please keep an eye on the requirements in the future. 
+If you're feeling brave, you can try use 3.11 instead of 3.10 below._
+
 Then simply install it on your machine using the default settings.
 
 To check if you installed the right version, open the Command Prompt and run:

requirements.txt

@@ -20,10 +20,11 @@ timecode==1.3.1
 tokenizers
 torch
 torchaudio
+torchvision
 onnxruntime
 tqdm
 transformers>4.33.3
-typing_extensions==4.7
+typing_extensions==4.12.2
 urllib3
 openai-whisper @ git+https://github.com/openai/whisper.git@ba3f3cd54b0e5b8ce1ab3de13e32122d0d5f98ab
 dtw-python

storytoolkitai/__main__.py

@@ -1,10 +1,27 @@
-
 import sys
 import argparse
 import platform
 import time
 import subprocess
 import os
+import ctypes
+import ctypes.util
+
+
+def is_windows():
+    return platform.system() == "Windows"
+
+
+def is_cuda_available():
+    if not is_windows():
+        return False
+
+    try:
+        cuda = ctypes.CDLL(ctypes.util.find_library("nvcuda") or "nvcuda.dll")
+        cuda.cuInit(0)
+        return True
+    except:
+        return False
 
 # add content root to sys.path
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
@@ -14,21 +31,19 @@
 # signal the start of the session in the log by adding some info about the machine
 logger.debug('\n--------------\n'
              'Platform: {} {}\n Platform version: {}\n OS: {} \n running Python {}'
-             '\n--------------'.format(
-    platform.system(), platform.release(),
-    platform.version(),
-    ' '.join(map(str, platform.win32_ver() + platform.mac_ver())),
-    '.'.join(map(str, sys.version_info))))
+             '\n--------------'
+             .format(platform.system(), platform.release(), platform.version(),
+                     ' '.join(map(str, platform.win32_ver() + platform.mac_ver())),
+                     '.'.join(map(str, sys.version_info))))
 
 # check if the user is running any version of Python 3.10
-if sys.version_info.major != 3 or sys.version_info.minor != 10:
+if sys.version_info.major != 3 or (sys.version_info.minor != 10 and sys.version_info.minor != 11):
 
     logger.warning('You are running Python {}.{}.{}.\n'.format(*sys.version_info) +
-                   'StoryToolkitAI is now optimized to run on Python 3.10.\n' +
-                   'Please download and install the latest version of Python 3.10 '
-                   'from: https://www.python.org/downloads/\nand then re-install the '
-                   'tool with a new environment.\n '
-                   'More info: https://github.com/octimot/StoryToolkitAI/blob/main/INSTALLATION.mdn\n')
+                   'StoryToolkitAI is now optimized to run on Python 3.10 and runs well on Python 3.11.\n' +
+                   'Please download and install the latest version of Python 3.10 or 3.11 '
+                   'from: https://www.python.org/downloads/\nand then re-install the tool with a new environment.\n '
+                   'More info: https://github.com/octimot/StoryToolkitAI/blob/main/INSTALLATION.md\n')
 
     # keep this message in the console for a bit
     time.sleep(5)
@@ -75,19 +90,75 @@
         logger.warning("Packages missing from the installation: {}".format(e))
         requirements_failed = True
 
-    except:
+    except Exception as e:
         # log the error and show the warning
-        logger.warning("There's something wrong with the packages installed in your Python environment:", exc_info=True)
+        logger.warning("There's something wrong with the packages installed in your Python environment: {}".format(e),
+                       exc_info=True)
         requirements_failed = True
 
-    # no matter what, we need to check if the user has the correct version of Python installed
+    # if the requirements are not met, we need to install them
     if requirements_failed:
 
         logger.warning('Some of the packages required to run StoryToolkitAI are missing from your Python environment.')
 
         # try to install the requirements automatically
         logger.warning('Attempting to automatically install the missing packages...')
 
+        # if we are on Windows and CUDA is available, we need to make sure we're using the correct PyTorch version
+        if is_windows() and is_cuda_available():
+
+            import shutil
+
+            # create the path for the new requirements file
+            windows_requirements_file_path = \
+                os.path.abspath(os.path.join(os.path.dirname(file_path), '..', 'requirements_windows_cuda.txt'))
+
+            # copy the original file
+            shutil.copy2(requirements_file_path, windows_requirements_file_path)
+
+            # read the contents of the new file
+            with open(windows_requirements_file_path, 'r') as f:
+                lines = f.readlines()
+
+            # modify the PyTorch related lines
+            pytorch_packages = ['torch', 'torchvision', 'torchaudio']
+            modified_lines = []
+            for line in lines:
+
+                # if it's either torch, torchvision or torchaudio
+                if line.startswith('torchaudio'):
+                    # modify the line to include the correct version
+                    line = 'torchaudio==2.0.1+cu118\n'
+                    modified_lines.append(line)
+
+                elif line.startswith('torchvision'):
+                    # modify the line to include the correct version
+                    line = 'torchvision==0.15.1+cu118\n'
+                    modified_lines.append(line)
+
+                elif line.startswith('torch'):
+                    # modify the line to include the correct version
+                    line = 'torch==2.0.0+cu118\n'
+                    modified_lines.append(line)
+
+                # otherwise, just add the line as is
+                else:
+                    modified_lines.append(line)
+
+            # add this to the beginning of the file
+            modified_lines.insert(0, '--find-links https://download.pytorch.org/whl/torch_stable.html\n')
+
+            # write the modified contents back to the file
+            with open(windows_requirements_file_path, 'w') as f:
+                f.writelines(modified_lines)
+
+            logger.debug(f"Created Windows CUDA requirements file: {windows_requirements_file_path}")
+
+            requirements_file_path = windows_requirements_file_path
+
+        else:
+            windows_requirements_file_path = None
+
         # get the relative path to the requirements file
         requirements_file_path_abs = os.path.abspath(requirements_file_path)
 
@@ -103,9 +174,9 @@
             try:
                 # restart the app while passing all the arguments
                 os.execl(sys.executable, sys.executable, *sys.argv)
-            except:
-                logger.error('Could not restart StoryToolkitAI. Please restart the app manually.')
-
+            except Exception as e:
+                logger.error('Could not restart StoryToolkitAI: {}'.format(e))
+                logger.error('Please restart the tool manually.')
 
         else:
             # let the user know that the automatic installation failed
@@ -121,6 +192,22 @@
         # keep this message in the console for a bit
         time.sleep(2)
 
+        # if we created a windows requirements file, we can delete it until next time
+        if windows_requirements_file_path:
+            try:
+                # Ensure all file handles are closed
+                import gc
+
+                gc.collect()
+
+                # Try to remove the file
+                os.remove(windows_requirements_file_path)
+                logger.debug(f"Successfully removed temporary file: {windows_requirements_file_path}")
+            except PermissionError:
+                logger.warning(f"Could not remove temporary file: {windows_requirements_file_path}. It may be in use.")
+            except Exception as e:
+                logger.warning(f"An error occurred while trying to remove {windows_requirements_file_path}: {str(e)}")
+
 from storytoolkitai.core.storytoolkitai import StoryToolkitAI
 
 StoryToolkitAI.check_ffmpeg()
@@ -130,6 +217,7 @@
 from storytoolkitai.ui.toolkit_ui import run_gui
 from storytoolkitai.ui.toolkit_cli import run_cli
 
+
 def main():
 
     parser = argparse.ArgumentParser(description="Story Toolkit AI")

storytoolkitai/core/toolkit_ops/toolkit_ops.py

@@ -4,6 +4,7 @@
 import json
 import yaml
 import subprocess
+import platform
 
 from threading import Thread
 
@@ -45,6 +46,10 @@
 from .media import MediaItem, VideoFileClip, AudioFileClip
 
 
+def is_arm64_mac():
+    return platform.system() == 'Darwin' and platform.machine() == 'arm64'
+
+
 class NLE:
     """
     Use this class to store any data that is supposed to be shared with the NLE (for eg. Resolve)
@@ -2260,7 +2265,7 @@ def process_transcription_metadata(self, other_options, transcription: Transcrip
 
     def classify_segments(self, segments: list, labels: list,
                           min_confidence: int or list = 0.55, multi_label_pass: list = None, **kwargs):
-        '''
+        """
         Classifies segments into different types using the transformers zero-shot-classification pipeline
         :param segments: the segments to classify
         :param labels: the labels to use for classification, if a list of lists is provided,
@@ -2269,7 +2274,7 @@ def classify_segments(self, segments: list, labels: list,
                                 if a list is provided, then the confidence is calculated for multi_label_pass label group
         :param multi_label_pass: a list of groups of labels that need to be passed together,
                                 so that the segment stays in the result
-        '''
+        """
 
         if segments is None or len(segments) == 0:
             logger.debug('No segments to classify.')
@@ -2301,10 +2306,16 @@ def classify_segments(self, segments: list, labels: list,
         model_name = self.stAI.get_app_setting('text_classifier_model', default_if_none='facebook/bart-large-mnli')
 
         logger.debug('Loading text classifier model: {}'.format(model_name))
+
         # get the zero-shot-classification pipeline
-        classifier = pipeline('zero-shot-classification',
-                              model=model_name,
-                              )
+        # if this is an arm64 mac use mps as device
+        if is_arm64_mac():
+            classifier = pipeline('zero-shot-classification', model=model_name, device='mps',)
+        elif torch.cuda.is_available():
+            logger.debug('Using GPU for classification.')
+            classifier = pipeline('zero-shot-classification', model=model_name, device='cuda')
+        else:
+            classifier = pipeline('zero-shot-classification', model=model_name)
 
         logger.debug('Classifying segments using the following labels: {}'.format(labels))
 

version.py

@@ -1 +1 @@
-__version__ = "0.24.2"
+__version__ = "0.24.3"