| """Guard the single-repository download and mounted-checkpoint loading paths.""" |
| import sys |
| from types import SimpleNamespace |
| from unittest.mock import MagicMock, patch |
|
|
| import pytest |
|
|
| from flux3_engine import Engine, download_weights |
| from model_config import GGUF_FILENAME, Settings |
|
|
|
|
| @pytest.mark.parametrize('mode', ['q8_0', 'bf16']) |
| def test_download_only_uses_selected_repository(monkeypatch, mode): |
| hub = MagicMock() |
| monkeypatch.setitem(sys.modules, 'huggingface_hub', hub) |
| settings = Settings('example/complete-copy', 'pinned-commit', mode) |
| download_weights(settings) |
| hub.snapshot_download.assert_called_once() |
| call = hub.snapshot_download.call_args |
| assert call.args == (settings.model_id,) |
| assert call.kwargs['revision'] == settings.revision |
| assert (GGUF_FILENAME in call.kwargs['allow_patterns']) == (mode == 'q8_0') |
| assert ('transformer/*.safetensors' in call.kwargs['allow_patterns']) == (mode == 'bf16') |
| hub.hf_hub_download.assert_not_called() |
|
|
|
|
| @pytest.mark.parametrize('missing_gguf', [False, True]) |
| def test_mounted_q8_loads_locally_without_hub_cache(monkeypatch, tmp_path, missing_gguf): |
| checkpoint = tmp_path / GGUF_FILENAME |
| if not missing_gguf: |
| checkpoint.parent.mkdir() |
| checkpoint.write_bytes(b'test-checkpoint') |
| torch, diffusers, transformers, hub = (MagicMock() for _ in range(4)) |
| for name, module in [('torch', torch), ('diffusers', diffusers), |
| ('transformers', transformers), ('huggingface_hub', hub)]: |
| monkeypatch.setitem(sys.modules, name, module) |
| pipe = MagicMock() |
| blocks = diffusers.ModularPipeline.from_pretrained.return_value.blocks |
| |
| blocks.init_pipeline.side_effect = lambda path: pipe |
|
|
| def update_components(**components): |
| for name, component in components.items(): |
| setattr(pipe, name, component) |
| component.parameters.return_value = [SimpleNamespace(device='cuda:0')] |
|
|
| pipe.update_components.side_effect = update_components |
| with patch('flux3_engine.download_weights') as download: |
| engine = Engine(settings=Settings(quantization='q8_0'), model_path=tmp_path) |
| if missing_gguf: |
| with pytest.raises(FileNotFoundError, match='complete Socialtrait checkpoint'): |
| engine.load() |
| diffusers.ModularPipeline.from_pretrained.assert_not_called() |
| else: |
| engine.load() |
| args, kwargs = diffusers.Flux3Transformer2DModel.from_single_file.call_args |
| assert args == (str(checkpoint),) |
| assert kwargs['config'] == str(tmp_path) |
| assert kwargs['local_files_only'] is True |
| diffusers.ModularPipeline.from_pretrained.assert_called_once_with( |
| str(tmp_path), local_files_only=True) |
| assert engine.pipe is pipe |
| download.assert_not_called() |
| assert not hub.mock_calls |
|
|
