23 lines
808 B
Python
23 lines
808 B
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from flux2_resolution_guard.inference.image import analytic_compand_correction_tensor
|
|
from flux2_resolution_guard.models import SMICCorrectionModel
|
|
|
|
|
|
def test_analytic_compand_shape():
|
|
image = torch.rand(1, 3, 64, 64)
|
|
anchor = torch.rand(1, 3, 64, 64)
|
|
mask = torch.ones(1, 1, 64, 64)
|
|
out = analytic_compand_correction_tensor(image=image, anchor_image=anchor, mask=mask, mp_ratio=1.6, strength=0.8)
|
|
assert out.shape == image.shape
|
|
assert torch.isfinite(out).all()
|
|
|
|
|
|
def test_checkpointless_model_correction():
|
|
model = SMICCorrectionModel()
|
|
image = torch.rand(1, 3, 64, 64)
|
|
out = model(edited=image, original=image, anchor=image, mask=torch.ones(1, 1, 64, 64), mp_ratio=1.5)
|
|
assert out["corrected"].shape == image.shape
|