| #!/usr/bin/env python3 | |
| """Quick sanity check: build TAIC/C-TAIC and run one forward pass.""" | |
| import os | |
| import sys | |
| REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) | |
| sys.path.insert(0, REPO_ROOT) | |
| import torch | |
| from flexicm.models import TAIC, CTAIC | |
| def main(): | |
| device = "cpu" | |
| x = torch.rand(1, 3, 256, 256, device=device) | |
| taic = TAIC(out_channels=128).to(device) | |
| taic.freeze_base_codec() | |
| out = taic(x) | |
| assert out["h"].shape == (1, 128, 64, 64), out["h"].shape | |
| assert out["y_hat"].shape == (1, 192, 16, 16), out["y_hat"].shape | |
| ctaic = CTAIC(out_channels=128).to(device) | |
| ctaic.freeze_for_stage2() | |
| out2 = ctaic(x, y_b_hat=out["y_hat"], use_condition=True) | |
| assert out2["h"].shape == (1, 128, 64, 64) | |
| print("sanity check passed") | |
| if __name__ == "__main__": | |
| main() | |