diff --git a/lightx2v/models/runners/cosmos3/cosmos3_runner.py b/lightx2v/models/runners/cosmos3/cosmos3_runner.py index c1feb3910..5eba294ec 100644 --- a/lightx2v/models/runners/cosmos3/cosmos3_runner.py +++ b/lightx2v/models/runners/cosmos3/cosmos3_runner.py @@ -196,7 +196,7 @@ def load_text_encoder(self): if not os.path.exists(tokenizer_path): tokenizer_path = self.config["model_path"] logger.info(f"Loading Cosmos3 tokenizer from {tokenizer_path}") - return AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + return AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=self.config.get("trust_remote_code", False)) def load_image_encoder(self): return None diff --git a/tests/test_cosmos3_trust_remote_code.py b/tests/test_cosmos3_trust_remote_code.py new file mode 100644 index 000000000..606649e6c --- /dev/null +++ b/tests/test_cosmos3_trust_remote_code.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from unittest.mock import patch + +from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Runner + + +def _runner_with_config(config: dict) -> Cosmos3Runner: + runner = Cosmos3Runner.__new__(Cosmos3Runner) + runner.config = config + return runner + + +def test_load_text_encoder_defaults_trust_remote_code_to_false(tmp_path): + config = {"model_path": str(tmp_path)} + runner = _runner_with_config(config) + with patch("lightx2v.models.runners.cosmos3.cosmos3_runner.AutoTokenizer.from_pretrained") as from_pretrained: + runner.load_text_encoder() + assert from_pretrained.call_args.kwargs["trust_remote_code"] is False + + +def test_load_text_encoder_honors_explicit_trust_remote_code_opt_in(tmp_path): + config = {"model_path": str(tmp_path), "trust_remote_code": True} + runner = _runner_with_config(config) + with patch("lightx2v.models.runners.cosmos3.cosmos3_runner.AutoTokenizer.from_pretrained") as from_pretrained: + runner.load_text_encoder() + assert from_pretrained.call_args.kwargs["trust_remote_code"] is True