Merge pull request #475 from SilyNoMeta/feat/f5-pickletensor

Add initial support for pickletensor models to F5-TTS

system/tts_engines/f5tts/model_engine.py

@@ -403,6 +403,13 @@ class tts_class:
                 if not model_files:
                     # If no model_*.safetensors found, try any .safetensors file
                     model_files = list(model_dir.glob("*.safetensors"))
+                if not model_files:
+                # Try finding a pt model file as fallback
+                    # If no model_*.safetensors found, try finding a .pt model file
+                    model_files = list(model_dir.glob("model_*.pt"))
+                if not model_files:
+                    # If no model_*.safetensors found, try any .pt file
+                    model_files = list(model_dir.glob("*.pt"))
 
                 vocab_file = model_dir / "vocab.txt"
                 vocos_dir = model_dir / "vocos"
@@ -506,11 +513,28 @@ class tts_class:
             vocab_path = model_dir / "vocab.txt"
             vocos_path = model_dir / "vocos"
             
-            # Dynamically find the safetensors model file
+            # Dynamically find the safetensors or pickletensor model file
+            model_is_pickle = False
             model_files = list(model_dir.glob("model_*.safetensors"))
             if not model_files:
                 # Try finding any safetensors file as fallback
                 model_files = list(model_dir.glob("*.safetensors"))
+            if not model_files:
+                # Try finding the pt model file as fallback
+                model_files = list(model_dir.glob("model_*.pt"))
+                model_is_pickle = True
+            if not model_files:
+                # Try finding any pt file as fallback
+                model_files = list(model_dir.glob("*.pt"))
+                model_is_pickle = True
+
+            if model_is_pickle:
+                print(
+                    f"[{self.branding}ENG] \033[91mWarning\033[0m: The models found in '{model_dir}' are in Pickle format (.pt). "
+                    f"This format poses security risks due to potential arbitrary code execution. "
+                    f"Please ensure the source of the models is trusted. We recommend using 'safetensors' format for enhanced security. "
+                    f"For more information, visit: https://huggingface.co/docs/hub/en/security-pickle"
+                )
                 
             if not model_files:
                 print(f"[{self.branding}ENG] \033[91mError\033[0m: No model's safetensors file was found in the F5-TTS models directory.")