Coverage for tests/test_wrapper_flux2.py: 100%
89 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3import sys
4import pytest
6from typing import Any
8from unittest.mock import patch
9from unittest.mock import MagicMock
10from tests.torch_mock import TorchMock
11from tests.diffusers_mock import DiffusersMock
13from PIL import Image
15mock_torch = TorchMock()
16mock_diffusers = DiffusersMock()
18sys.path.append("wrapper")
19sys.path.append("wrapper/flux2")
20sys.path.append("wrapper/flux")
22mock_modules = {
23 'nvidia_smi': MagicMock(),
24 'imageio': MagicMock(),
25 'cv2': MagicMock(),
26 'torch': mock_torch,
27 'xfuser': MagicMock(),
28 'xfuser.config': MagicMock(),
29 'xfuser.core': MagicMock(),
30 'xfuser.core.distributed': MagicMock(),
31 'xfuser.model_executor': MagicMock(),
32 'xfuser.model_executor.models': MagicMock(),
33 'xfuser.model_executor.models.transformers.transformer_flux': MagicMock(),
34 'xfuser.model_executor.models.transformers.transformer_flux2': MagicMock(),
35 'xfuser.model_executor.layers': MagicMock(),
36 'xfuser.model_executor.layers.attention_processor': MagicMock(),
37}
38mock_modules.update(mock_torch.get_sub_modules())
39mock_modules.update(mock_diffusers.get_sub_modules())
41with patch.dict(sys.modules, mock_modules):
42 from flux2.wrapper_flux2 import Flux2Generation
45@pytest.mark.asyncio
46async def test_wrapper_flux2() -> None:
47 model = Flux2Generation()
48 assert model is not None
49 assert model.model_name == "flux2"
50 assert model.status == "initializing"
52 with pytest.raises(ValueError, match="Model not initialized"):
53 await model.generate(
54 width=128,
55 height=80,
56 prompt="test prompt")
58 # Capture the mock Flux2Pipeline and transformer wrapper so we can assert
59 # the new sharding behaviour: device_map="balanced" on both the transformer
60 # and the pipeline (so VAE and text encoders are distributed too), and
61 # pipeline.to() is never called.
62 mock_pipeline_cls = mock_modules['diffusers'].Flux2Pipeline
63 mock_transformer_cls = mock_modules[
64 'xfuser.model_executor.models.transformers.transformer_flux2'
65 ].xFuserFlux2Transformer2DWrapper
67 # Pre-access mock sub-component attributes so we hold stable references for
68 # assertions after init() calls .to() on each of them.
69 mock_pipeline_instance = mock_pipeline_cls.from_pretrained.return_value
71 model.init()
72 assert model.status == "ok"
74 # Verify transformer was loaded with device_map="balanced"
75 _, transformer_kwargs = mock_transformer_cls.from_pretrained.call_args
76 assert transformer_kwargs.get("device_map") == "balanced", (
77 "Transformer must be loaded with device_map='balanced' to shard across GPUs"
78 )
80 # Verify the pipeline was also loaded with device_map="balanced" so that
81 # VAE and text encoders are distributed rather than crammed onto one GPU.
82 _, pipeline_kwargs = mock_pipeline_cls.from_pretrained.call_args
83 assert pipeline_kwargs.get("device_map") == "balanced", (
84 "Pipeline must be loaded with device_map='balanced' to distribute VAE and text encoders"
85 )
87 # Verify the full pipeline was NOT moved to a single device (would cause OOM)
88 mock_pipeline_instance.to.assert_not_called()
90 health = model.get_health()
91 assert health is not None
92 timestamps = model.get_timestamps()
93 assert timestamps is not None
95 with pytest.raises(ValueError, match="Missing JSON body"):
96 await model.get_rest_args(None)
97 with pytest.raises(ValueError, match="Missing 'prompt' parameter"):
98 await model.get_rest_args({})
99 args = await model.get_rest_args({
100 "job_id": "unittest",
101 "prompt": "Test prompt",
102 "width": 80,
103 "height": 60,
104 })
105 assert "args" in args
107 await model.warmup()
109 image = await model.generate(
110 width=256,
111 height=320,
112 prompt="Test prompt")
113 assert image is not None
114 assert isinstance(image, Image.Image)
115 assert image.size == (256, 320)
117 image = await model.generate(
118 width=480,
119 height=320,
120 prompt="Test prompt")
121 assert image is not None
122 assert isinstance(image, Image.Image)
123 assert image.size == (480, 320)
125 # 48x48 not supported for 2 GPUs (latent shape 9, odd).
126 model.world_size = 2
127 with pytest.raises(ValueError, match="48x48 not supported for 2 GPUs"):
128 await model.generate(
129 width=48,
130 height=48,
131 prompt="Test prompt")
133 del model
136@pytest.mark.asyncio
137async def test_additional_coverage() -> None:
138 """Cover seed path, step callbacks, parallelism init, and compile-disabled paths."""
139 model = Flux2Generation()
140 model.init()
141 assert model.status == "ok"
143 image = await model.generate(
144 width=256,
145 height=320,
146 prompt="Seed coverage test",
147 seed=42)
148 assert isinstance(image, Image.Image)
150 pipeline_instance = model.pipeline
152 def _pipeline_with_callback(*args: Any, **kwargs: Any) -> Any:
153 n_steps = kwargs.get("num_inference_steps", 2)
154 callback = kwargs.get("callback_on_step_end")
155 if callback:
156 for step in range(n_steps):
157 callback(pipeline_instance, step, 0, {})
158 out = MagicMock()
159 out.images = [Image.new("RGB", (kwargs.get("width", 64), kwargs.get("height", 64)))]
160 return out
162 pipeline_instance.side_effect = _pipeline_with_callback
163 image = await model.generate(
164 width=256,
165 height=320,
166 prompt="Callback coverage test",
167 sampling_steps=2)
168 assert isinstance(image, Image.Image)
170 model.world_size = 2
171 model.init_model_parallelism()
173 model.torch_compile = False
174 model.model_compile()
176 del model
179def test_model_compile_no_pipeline() -> None:
180 """model_compile() with pipeline=None returns early (pipeline not yet loaded)."""
181 model = Flux2Generation()
182 assert model.pipeline is None
183 model.model_compile()
184 del model