diff --git a/backend/src/nodes/properties/inputs/file_inputs.py b/backend/src/nodes/properties/inputs/file_inputs.py index 3933058e7..b7a25a04b 100644 --- a/backend/src/nodes/properties/inputs/file_inputs.py +++ b/backend/src/nodes/properties/inputs/file_inputs.py @@ -111,6 +111,9 @@ def PthFileInput(primary_input: bool = False) -> FileInput: file_kind="pth", filetypes=[".pt", ".pth", ".ckpt", ".safetensors"], primary_input=primary_input, + ).with_docs( + "Supports PyTorch state dict (.pth), TorchScript (.pt), Checkpoint (.ckpt), or SafeTensors (.safetensors) files.", + "Use SafeTensors files for better security (the other file types might be able to execute arbitrary code when loaded).", ) diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/io/load_model.py b/backend/src/packages/chaiNNer_pytorch/pytorch/io/load_model.py index 8422da570..745fe6682 100644 --- a/backend/src/packages/chaiNNer_pytorch/pytorch/io/load_model.py +++ b/backend/src/packages/chaiNNer_pytorch/pytorch/io/load_model.py @@ -34,10 +34,7 @@ def parse_ckpt_state_dict(checkpoint: dict): schema_id="chainner:pytorch:load_model", name="Load Model", description=[ - ( - "Load PyTorch state dict (.pth), TorchScript (.pt), or Checkpoint (.ckpt) files into an" - " auto-detected supported model architecture." - ), + ("Load PyTorch model file into an auto-detected supported model architecture."), ( "- For Super-Resolution, we support most variations of the RRDB" " architecture (ESRGAN, Real-ESRGAN, RealSR, BSRGAN, SPSR), Real-ESRGAN's"