Spaces:
Running on Zero
Running on Zero
File size: 2,509 Bytes
41ff959 a6825eb 41ff959 a6825eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import pytest
import torch
import src.model.encoder.depth.depthpro.depthpro_wrapper as depthpro_wrapper
from src.demo import infer_single_image
from src.demo.infer_single_image import (
_extract_state_dict,
_resolve_checkpoint_path,
load_demo_encoder,
)
def test_extract_state_dict_uses_safe_loading(tmp_path, monkeypatch) -> None:
checkpoint = tmp_path / "model.ckpt"
expected = {"encoder.weight": torch.tensor([1.0])}
torch.save(expected, checkpoint)
original_load = torch.load
def tracked_load(*args, **kwargs):
assert kwargs["weights_only"] is True
return original_load(*args, **kwargs)
monkeypatch.setattr(torch, "load", tracked_load)
actual = _extract_state_dict(checkpoint)
assert actual.keys() == expected.keys()
assert torch.equal(actual["encoder.weight"], expected["encoder.weight"])
def test_resolve_checkpoint_path_accepts_local_file(tmp_path) -> None:
checkpoint = tmp_path / "model.ckpt"
checkpoint.touch()
assert _resolve_checkpoint_path(str(checkpoint)) == checkpoint
def test_resolve_checkpoint_path_rejects_missing_file(tmp_path) -> None:
checkpoint = tmp_path / "missing.ckpt"
with pytest.raises(FileNotFoundError, match="Checkpoint not found"):
_resolve_checkpoint_path(str(checkpoint))
def test_load_demo_encoder_ignores_decoder_weights(tmp_path, monkeypatch) -> None:
checkpoint = tmp_path / "model.ckpt"
torch.save(
{
"encoder.weight": torch.tensor([[3.0]]),
"decoder.unused": torch.tensor([7.0]),
},
checkpoint,
)
encoder = torch.nn.Linear(1, 1, bias=False)
monkeypatch.setattr(infer_single_image, "get_encoder", lambda _: encoder)
loaded = load_demo_encoder(
cfg=type("Cfg", (), {"model": type("Model", (), {"encoder": object()})()})(),
checkpoint_path=checkpoint,
device=torch.device("cpu"),
)
assert loaded is encoder
assert not loaded.training
assert torch.equal(loaded.weight, torch.tensor([[3.0]]))
def test_depthpro_does_not_require_external_checkpoint(monkeypatch) -> None:
captured = {}
def create_without_weights(config, **_kwargs):
captured["checkpoint_uri"] = config.checkpoint_uri
return torch.nn.Identity(), None
monkeypatch.setattr(
depthpro_wrapper,
"create_model_and_transforms",
create_without_weights,
)
depthpro_wrapper.DepthPro()
assert captured["checkpoint_uri"] is None
|