File size: 10,778 Bytes
da1a4ff | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | import json
import tempfile
from pathlib import Path
import unittest
from inference_profile import (
InferenceProfile,
InferenceProfileError,
load_checkpoint_capabilities,
load_inference_profile,
resolve_model_directory,
)
GENERATION = {
"schema_version": 1,
"inference_profile": "generation_edit",
"alignment_padding_mode": "zero_masked",
"multi_frame_output": False,
"vae_input_channels": 4,
"vae_sample_mode": "argmax",
}
LAYER = {
"schema_version": 1,
"inference_profile": "layer_decompose",
"alignment_padding_mode": "learned",
"multi_frame_output": True,
"vae_input_channels": 4,
"vae_sample_mode": "argmax",
}
class InferenceProfileTest(unittest.TestCase):
def test_all_profile_fields_are_required(self):
for field in GENERATION:
raw = dict(GENERATION)
raw.pop(field)
with self.assertRaisesRegex(InferenceProfileError, "missing required fields"):
InferenceProfile.from_dict(raw)
def test_profile_rejects_unknown_fields(self):
raw = dict(GENERATION, inferred_from_directory_name=True)
with self.assertRaisesRegex(InferenceProfileError, "unsupported fields"):
InferenceProfile.from_dict(raw)
def test_generation_profile_task_matrix(self):
profile = InferenceProfile.from_dict(GENERATION)
self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
self.assertEqual(profile.resolve_sampling_parameters().cfg, 1.0)
profile.validate_task("text-to-image", has_reference_image=False)
profile.validate_task("image-edit", has_reference_image=True)
with self.assertRaisesRegex(InferenceProfileError, "layer_decompose checkpoint"):
profile.validate_task("layer-decompose", has_reference_image=True, num_layers=4)
def test_layer_profile_task_matrix(self):
profile = InferenceProfile.from_dict(LAYER)
self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
self.assertEqual(profile.resolve_sampling_parameters().cfg, 2.0)
profile.validate_task("layer-decompose", has_reference_image=True, num_layers=4)
with self.assertRaisesRegex(InferenceProfileError, "generation_edit checkpoint"):
profile.validate_task("image-edit", has_reference_image=True)
def test_generation_profile_requires_qwen_vae_contract(self):
with self.assertRaisesRegex(InferenceProfileError, "vae_input_channels=4"):
InferenceProfile.from_dict(dict(GENERATION, vae_input_channels=3))
with self.assertRaisesRegex(InferenceProfileError, "vae_sample_mode='argmax'"):
InferenceProfile.from_dict(dict(GENERATION, vae_sample_mode="sample"))
def test_layer_profile_requires_qwen_vae_contract(self):
with self.assertRaisesRegex(InferenceProfileError, "vae_sample_mode='argmax'"):
InferenceProfile.from_dict(dict(LAYER, vae_sample_mode="sample"))
def test_sampling_overrides_are_validated(self):
profile = InferenceProfile.from_dict(GENERATION)
resolved = profile.resolve_sampling_parameters(steps=18, cfg=1.5)
self.assertEqual((resolved.steps, resolved.cfg), (18, 1.5))
with self.assertRaisesRegex(InferenceProfileError, "steps"):
profile.resolve_sampling_parameters(steps=0)
with self.assertRaisesRegex(InferenceProfileError, "CFG"):
profile.resolve_sampling_parameters(cfg=float("inf"))
def test_profile_is_loaded_from_local_model_directory(self):
with self.subTest("local profile"):
import tempfile
from pathlib import Path
with tempfile.TemporaryDirectory() as directory:
model_directory = Path(directory)
(model_directory / "inference_profile.json").write_text(
json.dumps(LAYER), encoding="utf-8"
)
self.assertEqual(
load_inference_profile(model_directory).inference_profile,
"layer_decompose",
)
self.assertEqual(
resolve_model_directory(model_directory), model_directory.resolve()
)
def test_missing_profile_is_an_error(self):
import tempfile
with tempfile.TemporaryDirectory() as directory:
with self.assertRaisesRegex(InferenceProfileError, "must contain"):
load_inference_profile(directory)
def test_missing_absolute_local_path_is_not_treated_as_hub_id(self):
import tempfile
from pathlib import Path
with tempfile.TemporaryDirectory() as directory:
with self.assertRaises(FileNotFoundError):
resolve_model_directory(Path(directory) / "missing")
VAE_QWEN_4CH = {"_class_name": "AutoencoderKLQwenImage", "input_channels": 4}
def build_component_package(
root: Path,
transformer_extra=None,
vae=None,
profile=None,
) -> None:
(root / "transformer").mkdir(parents=True, exist_ok=True)
(root / "vae").mkdir(parents=True, exist_ok=True)
transformer = {"_class_name": "DiffusionTransformer", "dim": 4}
transformer.update(transformer_extra or {})
(root / "transformer" / "config.json").write_text(json.dumps(transformer))
(root / "vae" / "config.json").write_text(json.dumps(vae or VAE_QWEN_4CH))
if profile is not None:
(root / "inference_profile.json").write_text(json.dumps(profile))
class ComponentCapabilityTest(unittest.TestCase):
def capabilities(self, **kwargs):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
build_component_package(root, **kwargs)
return load_checkpoint_capabilities(root)
def test_generation_pair_without_profile(self):
profile = self.capabilities(
transformer_extra={
"alignment_padding_mode": "zero_masked",
"multi_frame_output": False,
}
)
self.assertEqual(profile.inference_profile, "generation_edit")
self.assertEqual(profile.vae_input_channels, 4)
self.assertEqual(profile.vae_sample_mode, "argmax")
self.assertEqual(profile.resolve_sampling_parameters().steps, 12)
self.assertEqual(profile.resolve_sampling_parameters().cfg, 1.0)
def test_layer_pair_without_profile(self):
profile = self.capabilities(
transformer_extra={
"alignment_padding_mode": "learned",
"multi_frame_output": True,
}
)
self.assertEqual(profile.inference_profile, "layer_decompose")
self.assertEqual(profile.resolve_sampling_parameters().cfg, 2.0)
def test_partial_pair_is_rejected(self):
for field in ("alignment_padding_mode", "multi_frame_output"):
with self.subTest(only=field):
pair = {
"alignment_padding_mode": "zero_masked",
"multi_frame_output": False,
}
with self.assertRaisesRegex(InferenceProfileError, "declared together"):
self.capabilities(transformer_extra={field: pair[field]})
def test_every_invalid_pair_is_rejected(self):
invalid = [
{"alignment_padding_mode": "zero_masked", "multi_frame_output": True},
{"alignment_padding_mode": "learned", "multi_frame_output": False},
{"alignment_padding_mode": "zero", "multi_frame_output": False},
{"alignment_padding_mode": "learned", "multi_frame_output": 1},
{"alignment_padding_mode": None, "multi_frame_output": False},
]
for pair in invalid:
with self.subTest(pair=pair):
with self.assertRaises(InferenceProfileError):
self.capabilities(transformer_extra=pair)
def test_vae_class_must_be_qwen(self):
with self.assertRaisesRegex(InferenceProfileError, "AutoencoderKLQwenImage"):
self.capabilities(
transformer_extra={
"alignment_padding_mode": "zero_masked",
"multi_frame_output": False,
},
vae={"_class_name": "AutoencoderKL", "in_channels": 4},
)
def test_vae_channel_fields_must_agree_and_be_four(self):
base = {"alignment_padding_mode": "zero_masked", "multi_frame_output": False}
with self.assertRaisesRegex(InferenceProfileError, "disagree"):
self.capabilities(
transformer_extra=base,
vae=dict(VAE_QWEN_4CH, in_channels=16),
)
with self.assertRaisesRegex(InferenceProfileError, "4-channel"):
self.capabilities(
transformer_extra=base,
vae={"_class_name": "AutoencoderKLQwenImage", "input_channels": 3},
)
with self.assertRaisesRegex(InferenceProfileError, "must declare"):
self.capabilities(
transformer_extra=base,
vae={"_class_name": "AutoencoderKLQwenImage"},
)
# in_channels alone is accepted when it agrees by itself.
profile = self.capabilities(
transformer_extra=base,
vae={"_class_name": "AutoencoderKLQwenImage", "in_channels": 4},
)
self.assertEqual(profile.vae_input_channels, 4)
def test_legacy_profile_fallback(self):
profile = self.capabilities(profile=LAYER)
self.assertEqual(profile.inference_profile, "layer_decompose")
self.assertEqual(profile.alignment_padding_mode, "learned")
def test_legacy_fallback_requires_profile_file(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
build_component_package(root)
with self.assertRaisesRegex(InferenceProfileError, "must contain"):
load_checkpoint_capabilities(root)
def test_component_fields_and_agreeing_legacy_profile(self):
profile = self.capabilities(
transformer_extra={
"alignment_padding_mode": "zero_masked",
"multi_frame_output": False,
},
profile=GENERATION,
)
self.assertEqual(profile.inference_profile, "generation_edit")
def test_component_fields_disagreeing_with_legacy_profile(self):
with self.assertRaisesRegex(InferenceProfileError, "disagree"):
self.capabilities(
transformer_extra={
"alignment_padding_mode": "learned",
"multi_frame_output": True,
},
profile=GENERATION,
)
if __name__ == "__main__":
unittest.main()
|