73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
"""Public installation resolves model IDs and exposes supported commands."""
|
|
import json
|
|
import pytest
|
|
|
|
from yue2 import cli, pipeline
|
|
|
|
|
|
@pytest.mark.parametrize("vae,repo", [
|
|
("standard", "m-a-p/YuE2-Vae"),
|
|
("legacy", "m-a-p/YuE2-Vae-legacy"),
|
|
])
|
|
def test_cli_resolves_public_model_and_decoder(vae, repo, monkeypatch, tmp_path):
|
|
monkeypatch.setenv("YUE2_KIT", str(tmp_path))
|
|
args = cli.parser().parse_args(["generate", "--vae", vae])
|
|
assert cli.model_paths(args) == ("m-a-p/YuE2-3B", repo)
|
|
|
|
|
|
def test_cli_accepts_explicit_model_and_decoder(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("YUE2_KIT", str(tmp_path))
|
|
args = cli.parser().parse_args([
|
|
"generate", "--model", "example/song-model", "--vae", "example/decoder",
|
|
])
|
|
assert cli.model_paths(args) == ("example/song-model", "example/decoder")
|
|
|
|
|
|
@pytest.mark.parametrize("command", ["verify", "bench", "eval"])
|
|
def test_cli_rejects_commands_not_in_the_public_package(command):
|
|
with pytest.raises(SystemExit) as exc:
|
|
cli.parser().parse_args([command])
|
|
assert exc.value.code == 2
|
|
|
|
|
|
def test_pipeline_defaults_and_explicit_revisions_reach_hub(monkeypatch, tmp_path):
|
|
resolved = []
|
|
|
|
def resolve(repo, **kwargs):
|
|
resolved.append((repo, kwargs))
|
|
return tmp_path / repo.rsplit("/", 1)[-1]
|
|
|
|
class Pipeline(pipeline.YuE2Pipeline):
|
|
def __init__(self, model, vae, **kwargs):
|
|
self.model, self.vae, self.options = model, vae, kwargs
|
|
self.load_timing = {}
|
|
|
|
monkeypatch.setattr(pipeline, "resolve_model", resolve)
|
|
result = Pipeline.from_pretrained(
|
|
revision="a" * 40, vae_revision="b" * 40,
|
|
token=False, cache_dir=str(tmp_path), progress=False,
|
|
)
|
|
assert [repo for repo, _ in resolved] == ["m-a-p/YuE2-3B", "m-a-p/YuE2-Vae"]
|
|
assert [kwargs["revision"] for _, kwargs in resolved] == ["a" * 40, "b" * 40]
|
|
assert all(kwargs["token"] is False for _, kwargs in resolved)
|
|
assert result.options["progress"] is False
|
|
|
|
|
|
def test_saved_pipeline_uses_its_own_decoder(monkeypatch, tmp_path):
|
|
(tmp_path / "pipeline.json").write_text(json.dumps({
|
|
"model": "model", "vae": "decoder", "generation_config": {},
|
|
}))
|
|
resolved = []
|
|
|
|
def resolve(path, **kwargs):
|
|
resolved.append(path)
|
|
return path
|
|
|
|
class Pipeline(pipeline.YuE2Pipeline):
|
|
def __init__(self, *args, **kwargs):
|
|
self.load_timing = {}
|
|
|
|
monkeypatch.setattr(pipeline, "resolve_model", resolve)
|
|
Pipeline.from_pretrained(tmp_path, progress=False, local_files_only=True)
|
|
assert resolved == [tmp_path / "model", tmp_path / "decoder"]
|