File size: 1,055 Bytes
6f8c80a | 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 | """Explicit GPU integration test; run with python -m tests.gpu_roundtrip."""
import tempfile
from pathlib import Path
import torch
from diffusers import QwenImage21Transformer2DModel
from scripts.quantize import quantize_component
def main():
with tempfile.TemporaryDirectory() as d:
root = Path(d)
model = QwenImage21Transformer2DModel(
num_layers=1, num_attention_heads=2, attention_head_dim=32,
axes_dims_rope=(8, 12, 12), context_in_dim=64)
model.save_pretrained(root / 'source' / 'transformer')
del model
report = quantize_component(root / 'source', root / 'target', 'transformer')
assert report['int8_parameters'] > 0, report
loaded = QwenImage21Transformer2DModel.from_pretrained(
root / 'target' / 'transformer', device_map={'': 0}, dtype=torch.bfloat16)
assert any(p.dtype == torch.int8 for p in loaded.parameters())
print('PASS: real QwenImage21Transformer2DModel INT8 save/reload', report)
if __name__ == '__main__':
main()
|