Update app.py

app.py

@@ -636,13 +636,23 @@ def convert_ebook_to_audio(ebook_file, target_voice_file, language, use_custom_m
     remove_folder_with_contents(full_folder_working_files)
     remove_folder_with_contents(output_audio_directory)
 
+    # If the language argument is set use it instead
+    language = args.language if args.language else language
+
     if use_custom_model and custom_model_file and custom_config_file and custom_vocab_file:
         custom_model = {
             'model': custom_model_file.name,
             'config': custom_config_file.name,
             'vocab': custom_vocab_file.name
         }
-    if use_custom_model and custom_model_url:
+    # If headless is used with the custom model arguments
+    if args.use_custom_model and args.custom_model and args.custom_config and args.custom_vocab:
+        custom_model = {
+            'model': args.custom_model,
+            'config': args.custom_config,
+            'vocab': args.custom_vocab
+        }
+    if (use_custom_model and custom_model_url) or (args.use_custom_model and custom_model_url):
         print(f"Received custom model URL: {custom_model_url}")
         download_dir = os.path.join(".", "Working_files", "custom_model")
         download_and_extract_zip(custom_model_url, download_dir)
@@ -791,6 +801,9 @@ def run_gradio_interface():
 
 # Check if running in headless mode
 if args.headless:
+    # If the arg.custom_model_url exists then use it as the custom_model_url lol
+    custom_model_url = args.custom_model_url if args.custom_model_url else None
+    
     if not args.ebook:
         print("Error: In headless mode, you must specify an ebook file using --ebook.")
         exit(1)
@@ -800,14 +813,28 @@ if args.headless:
     custom_model = None
 
     if args.use_custom_model:
-        custom_model = {
-            'model': args.custom_model,
-            'config': args.custom_config,
-            'vocab': args.custom_vocab
-        }
+        # Check if custom_model_url is provided
+        if args.custom_model_url:
+            # Download the custom model from the provided URL
+            custom_model_url = args.custom_model_url
+        else:
+            # If no URL is provided, ensure all custom model files are provided
+            if not args.custom_model or not args.custom_config or not args.custom_vocab:
+                print("Error: You must provide either a --custom_model_url or all of the following arguments:")
+                print("--custom_model, --custom_config, and --custom_vocab")
+                exit(1)
+            else:
+                # Assign the custom model files
+                custom_model = {
+                    'model': args.custom_model,
+                    'config': args.custom_config,
+                    'vocab': args.custom_vocab
+                }
+
+
 
     # Example headless execution
-    convert_ebook_to_audio(ebook_file_path, target_voice, args.language, args.use_custom_model, args.custom_model, args.custom_config, args.custom_vocab)
+    convert_ebook_to_audio(ebook_file_path, target_voice, args.language, args.use_custom_model, args.custom_model, args.custom_config, args.custom_vocab, custom_model_url)
 
 else:
     # Launch Gradio UI