19 lines
615 B
Python
19 lines
615 B
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from flux2_resolution_guard.models import SMICCorrectionModel, SMICConfig
|
|
|
|
|
|
def test_model_forward_shapes():
|
|
model = SMICCorrectionModel(SMICConfig())
|
|
edited = torch.rand(2, 3, 64, 64)
|
|
original = torch.rand(2, 3, 64, 64)
|
|
anchor = torch.rand(2, 3, 64, 64)
|
|
mask = torch.rand(2, 1, 64, 64)
|
|
|
|
out = model(edited=edited, original=original, anchor=anchor, mask=mask, mp_ratio=torch.tensor([1.2, 1.8]))
|
|
assert out["corrected"].shape == (2, 3, 64, 64)
|
|
assert out["flow"].shape == (2, 2, 64, 64)
|
|
assert out["gate"].shape == (2, 1, 64, 64)
|