tests/gpu_roundtrip.py

"""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()