Spaces:
Runtime error
Runtime error
File size: 2,908 Bytes
2857cf3 | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | from __future__ import annotations
from pathlib import Path
from PIL import Image
import pytest
import torch
import kneiff.utils.image.upscale as image_upscale
from kneiff.progress import ProgressUpdate
class _FakeUpscaleModel:
"""Small runtime descriptor used without loading model weights."""
device = torch.device("cpu")
dtype = torch.float32
scale = 4
def _install_fake_upscale_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(
image_upscale,
"pick_device",
lambda gpu, cpu: torch.device("cpu"),
)
monkeypatch.setattr(
image_upscale,
"resolve_weights_filename",
lambda **kwargs: "4x-test.safetensors",
)
monkeypatch.setattr(
image_upscale,
"download_weights",
lambda **kwargs: Path("cached-weights.safetensors"),
)
monkeypatch.setattr(
image_upscale,
"load_sr_model",
lambda **kwargs: _FakeUpscaleModel(),
)
monkeypatch.setattr(
image_upscale,
"upscale_pil_image",
lambda **kwargs: Image.new("RGB", (8, 8), color="green"),
)
def test_local_upscale_reports_loading_phases_and_late_image_total(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
_install_fake_upscale_runtime(monkeypatch)
first = tmp_path / "first.png"
second = tmp_path / "second.jpg"
Image.new("RGB", (2, 2), color="red").save(first)
Image.new("RGB", (2, 2), color="blue").save(second)
updates: list[ProgressUpdate] = []
result = image_upscale.run_local_upscale(
image_upscale.LocalUpscaleRequest(
inputs=[first, second],
out_dir=tmp_path / "output",
cpu=True,
),
progress_callback=updates.append,
)
assert [update.description for update in updates[:4]] == [
"Selecting upscale device",
"Resolving model weights",
"Downloading model weights",
"Loading upscale model",
]
assert sum(update.additional_total for update in updates) == 2
assert sum(update.advance for update in updates) == 2
assert len(result.results) == 2
assert all(item.output_path.is_file() for item in result.results)
def test_local_upscale_main_preserves_legacy_redirected_output(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
capsys: pytest.CaptureFixture[str],
) -> None:
_install_fake_upscale_runtime(monkeypatch)
source = tmp_path / "source.png"
output_root = tmp_path / "output"
Image.new("RGB", (2, 2), color="red").save(source)
image_upscale.main(inputs=[source], out_dir=output_root, cpu=True)
output = capsys.readouterr().out
assert "Model: Kim2091/UltraSharpV2/4x-test.safetensors" in output
assert "Device: cpu | dtype: torch.float32 | scale: 4x" in output
assert f"[1/1] Wrote: {output_root.resolve() / 'source.png'}" in output
|