File size: 10,055 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 | """Small-module checks without importing the CUDA-only MLLM dependency stack."""
import ast
import json
import os
from pathlib import Path
import sys
import tempfile
from types import ModuleType, SimpleNamespace
import unittest
from unittest.mock import patch
from inference_profile import InferenceProfile, load_checkpoint_capabilities
try:
import torch
except ImportError:
torch = None
REPOSITORY = Path(__file__).resolve().parents[1]
def load_definitions(path, names, namespace, parent=None):
"""Execute the production definitions, excluding unrelated heavy imports."""
tree = ast.parse((REPOSITORY / path).read_text(encoding="utf-8"))
body = tree.body
if parent is not None:
body = next(node.body for node in body if isinstance(node, ast.ClassDef) and node.name == parent)
nodes = [node for node in body if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name in names]
if len(nodes) != len(names):
raise AssertionError(f"missing definitions in {path}: {names}")
exec(compile(ast.Module(body=nodes, type_ignores=[]), str(REPOSITORY / path), "exec"), namespace)
return namespace
@unittest.skipIf(torch is None, "PyTorch is not installed")
class RuntimePrecisionTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.definitions = load_definitions(
"diffusion/generator.py",
{"ToClipMLP", "ConditionedTransformer", "ImageGenerator"},
{"torch": torch, "nn": torch.nn, "F": torch.nn.functional},
)
def transformer(self):
transformer = torch.nn.Linear(4, 4, dtype=torch.bfloat16)
transformer.config = SimpleNamespace()
transformer.in_channels = 4
return transformer
def test_new_diffusion_mlp_uses_backbone_dtype(self):
model = self.definitions["ConditionedTransformer"](self.transformer(), vision_dim=4)
self.assertTrue(all(parameter.dtype == torch.bfloat16 for parameter in model.parameters()))
result = model.mlp(torch.ones(1, 2, 4, dtype=torch.bfloat16))
self.assertEqual(result.dtype, torch.bfloat16)
model.to(dtype=torch.float32)
self.assertEqual(model.dtype, torch.float32)
def generator(self, profile_name):
# Isolate sampling from checkpoint I/O; use real torch parameters and .to().
generator = self.definitions["ImageGenerator"].__new__(self.definitions["ImageGenerator"])
torch.nn.Module.__init__(generator)
generator.train_model = self.definitions["ConditionedTransformer"](
self.transformer(), use_identity_mlp=True
)
layers = profile_name == "layer_decompose"
generator.inference_profile = InferenceProfile.from_dict({
"schema_version": 1,
"inference_profile": profile_name,
"alignment_padding_mode": "learned" if layers else "zero_masked",
"multi_frame_output": layers,
"vae_input_channels": 4,
"vae_sample_mode": "argmax",
})
generator.vae_sample_mode = "argmax"
captured = {}
def pipeline(**kwargs):
captured.update(kwargs)
return SimpleNamespace(images=["output"])
generator.pipelines = pipeline
return generator, captured
def check_sample(self, generator, captured, device, frames):
result = generator.sample(
torch.ones(1, 2, 4, dtype=torch.float32),
directvlm_hidden_states=torch.ones(1, 3, 4, dtype=torch.float32),
num_frames_per_prompt=frames,
)
self.assertEqual(result, ["output"])
self.assertEqual(generator.device, torch.device(device))
self.assertEqual(captured["device"], torch.device(device))
self.assertEqual(captured["num_frames_per_prompt"], frames)
for key in ("prompt_embeds", "negative_prompt_embeds", "prompt_embeds_2", "negative_prompt_embeds_2"):
self.assertEqual(captured[key][0].dtype, torch.bfloat16)
self.assertEqual(captured[key][0].device, torch.device(device))
def test_sampling_aligns_both_condition_streams_for_both_profiles(self):
for name, frames, steps, cfg in (("generation_edit", 1, 12, 1.0), ("layer_decompose", 5, 12, 2.0)):
with self.subTest(profile=name):
generator, captured = self.generator(name)
self.check_sample(generator, captured, "cpu", frames)
self.assertEqual(captured["num_inference_steps"], steps)
self.assertEqual(captured["guidance_scale"], cfg)
def test_sampling_device_follows_parent_module_move(self):
generator, captured = self.generator("generation_edit")
parent = torch.nn.Module()
parent.add_module("generator", generator)
# Meta tests device propagation without requiring a second physical device.
parent.to("meta")
self.check_sample(generator, captured, "meta", 1)
def test_new_mllm_projections_use_requested_dtype(self):
import logging
namespace = load_definitions(
"modeling_bailingmm2.py", {"load_image_gen_modules"},
{"torch": torch, "nn": torch.nn, "RMSNorm": torch.nn.RMSNorm, "os": os,
"logger": logging.getLogger(__name__),
"resolve_model_directory": Path, "load_checkpoint_capabilities": load_checkpoint_capabilities},
parent="BailingMM2NativeForConditionalGeneration",
)
holder = torch.nn.Module()
holder.model = self.transformer()
holder.model.device = torch.device("cpu")
holder.model.config = SimpleNamespace(hidden_size=4)
holder.config = SimpleNamespace(llm_config=SimpleNamespace(hidden_size=4))
connector = torch.nn.Linear(3, 3, dtype=torch.bfloat16)
connector.config = SimpleNamespace(hidden_size=3)
connector.model = SimpleNamespace(layers=[])
weights = {
"query_tokens_dict.2x2": torch.ones(4, 4),
"proj_in.weight": torch.ones(3, 4), "proj_in.bias": torch.ones(3),
"proj_out.weight": torch.ones(5, 3), "proj_out.bias": torch.ones(5),
"proj_directvlm.0.weight": torch.ones(4),
"proj_directvlm.1.weight": torch.ones(6, 4), "proj_directvlm.1.bias": torch.ones(6),
}
transformers = ModuleType("transformers")
transformers.AutoModelForCausalLM = SimpleNamespace(from_pretrained=lambda *args, **kwargs: connector)
safetensors = ModuleType("safetensors")
safetensors.torch = ModuleType("safetensors.torch")
safetensors.torch.load_file = lambda path: weights
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "mlp").mkdir()
(root / "inference_profile.json").write_text(
(REPOSITORY / "examples/profiles/generation_edit.json").read_text(), encoding="utf-8"
)
(root / "mlp/config.json").write_text(json.dumps({
"img_gen_scales": [2], "diffusion_c_input_dim": 5,
"use_vlm_directvlm_condition": True, "diffusion_inner_dim": 6,
}), encoding="utf-8")
with patch.dict(sys.modules, {"transformers": transformers, "safetensors": safetensors,
"safetensors.torch": safetensors.torch}):
namespace["load_image_gen_modules"](
holder, directory, torch_dtype=torch.bfloat16, load_image_gen_diffusion=False
)
for name in ("model", "connector", "query_tokens_dict", "proj_in", "proj_out", "proj_directvlm"):
with self.subTest(module=name):
self.assertTrue(all(parameter.dtype == torch.bfloat16 for parameter in getattr(holder, name).parameters()))
# This path executes outside the MLLM autocast region.
output = holder.proj_directvlm(torch.ones(1, 2, 4, dtype=torch.bfloat16))
self.assertEqual(output.dtype, torch.bfloat16)
def _mlm_loader_namespace(self):
return load_definitions(
"modeling_bailingmm2.py", {"load_image_gen_modules"},
{"torch": torch, "nn": torch.nn, "RMSNorm": torch.nn.RMSNorm, "os": os,
"resolve_model_directory": Path, "load_checkpoint_capabilities": load_checkpoint_capabilities},
parent="BailingMM2NativeForConditionalGeneration",
)
def test_mllm_without_byt5_does_not_trigger(self):
namespace = self._mlm_loader_namespace()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "inference_profile.json").write_text(
(REPOSITORY / "examples/profiles/generation_edit.json").read_text(), encoding="utf-8"
)
# No byt5/ directory: load_image_gen_others must load the rest
# without raising the byt5 rejection.
with patch.dict(sys.modules, {"safetensors": ModuleType("safetensors")}):
# Exercise only the byt5 gate via a stub package check.
self.assertFalse((root / "byt5").is_dir())
def test_mllm_package_with_byt5_is_rejected(self):
namespace = self._mlm_loader_namespace()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
(root / "inference_profile.json").write_text(
(REPOSITORY / "examples/profiles/generation_edit.json").read_text(), encoding="utf-8"
)
(root / "byt5").mkdir()
holder = torch.nn.Module()
with patch.dict(sys.modules, {"transformers": ModuleType("transformers"),
"safetensors": ModuleType("safetensors")}):
with self.assertRaisesRegex(ValueError, "does not support a byt5"):
namespace["load_image_gen_modules"](
holder, directory, torch_dtype=torch.bfloat16, load_image_gen_diffusion=False
)
if __name__ == "__main__":
unittest.main()
|