Diya Penganddiyapeng
6cd1ff7e64
Signed-off-by: diyapeng <50388704+diyapeng@users.noreply.github.com> Co-authored-by: diyapeng <50388704+diyapeng@users.noreply.github.com>
51 lines
1.8 KiB
Python
51 lines
1.8 KiB
Python
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
|
|
Flux2Pipeline,
|
|
_resolve_component_quant_config,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_resolve_component_quant_config_routes_components():
|
|
text_encoder_config = object()
|
|
component_config = Mock()
|
|
component_config.resolve.side_effect = {
|
|
"text_encoder": text_encoder_config,
|
|
"transformer": None,
|
|
}.get
|
|
|
|
assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
|
|
assert _resolve_component_quant_config(component_config, "transformer") is None
|
|
|
|
|
|
def test_resolve_component_quant_config_preserves_global_config():
|
|
global_config = object()
|
|
assert _resolve_component_quant_config(global_config, "transformer") is global_config
|
|
assert _resolve_component_quant_config(None, "text_encoder") is None
|
|
|
|
|
|
def test_offload_transformer_before_text_encoder_weights():
|
|
pipeline = object.__new__(Flux2Pipeline)
|
|
pipeline.transformer = Mock()
|
|
weights = [
|
|
("transformer.layer.weight", torch.empty(1)),
|
|
("text_encoder.layer.weight", torch.empty(1)),
|
|
("text_encoder.layer.bias", torch.empty(1)),
|
|
]
|
|
|
|
iterator = iter(pipeline._offload_transformer_before_text_encoder(weights))
|
|
assert next(iterator)[0] == "transformer.layer.weight"
|
|
pipeline.transformer.to.assert_not_called()
|
|
|
|
with patch("vllm_omni.diffusion.models.flux2.pipeline_flux2.current_omni_platform.empty_cache") as empty_cache:
|
|
assert next(iterator)[0] == "text_encoder.layer.weight"
|
|
assert next(iterator)[0] == "text_encoder.layer.bias"
|
|
|
|
pipeline.transformer.to.assert_called_once_with("cpu")
|
|
empty_cache.assert_called_once_with()
|