feat(diffusion): select Wan VAE decode mode automatically

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-11 09:41:05 +08:00
1 parent c27623c29c
commit 810112189b
16 files changed
+613 -186

No files matched your search

+20 -1
View File
@@ -728,6 +728,19 @@ def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
@pytest.mark.core_model
@pytest.mark.cpu
def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
memory_pool_active = False
preparation_pool_states = []
@contextmanager
def _memory_pool_context(tag):
nonlocal memory_pool_active
assert tag == "weights"
memory_pool_active = True
try:
yield
finally:
memory_pool_active = False
class _DummyLoader:
def __init__(self, load_config, od_config=None):
del load_config, od_config
@@ -775,16 +788,22 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
monkeypatch.setattr(
model_runner_module,
"prepare_pipeline_wan_spatial_shard_decode",
lambda pipeline: preparation_pool_states.append(memory_pool_active),
)
monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
monkeypatch.setattr(
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
)
DiffusionModelRunner.load_model(runner)
DiffusionModelRunner.load_model(runner, memory_pool_context_fn=_memory_pool_context)
assert runner.cache_backend is None
assert runner.od_config.cache_backend is None
assert dummy_cache_backend.enabled is False
assert preparation_pool_states == [True]
@pytest.mark.core_model