Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
108 lines
4.3 KiB
Python
108 lines
4.3 KiB
Python
"""Second pass nodes: loading H3's upscaler, sharpen/upscale, and pins that still fit."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Imports comfy."""
|
|
|
|
|
|
def _small_resizer():
|
|
from modules.models.minimax_h3.latent_upscaler import H3LatentResizer3D
|
|
torch.manual_seed(0)
|
|
return H3LatentResizer3D(channels=32, in_blocks=1, out_blocks=1).eval().requires_grad_(False)
|
|
|
|
|
|
@pytest.fixture
|
|
def upscaler_dir(tmp_path, monkeypatch):
|
|
import folder_paths
|
|
from safetensors.torch import save_file
|
|
save_file(_small_resizer().state_dict(), str(tmp_path / "h3_up.safetensors"))
|
|
save_file({"conv_in.weight": torch.zeros(2, 2)}, str(tmp_path / "other.safetensors"))
|
|
save_file({"base.low_blocks.0.conv1.weight": torch.zeros(2, 2)}, str(tmp_path / "attn.safetensors"))
|
|
monkeypatch.setitem(folder_paths.folder_names_and_paths, "latent_upscale_models",
|
|
([str(tmp_path)], {".safetensors"}))
|
|
folder_paths.filename_list_cache.pop("latent_upscale_models", None)
|
|
return tmp_path
|
|
|
|
|
|
def test_the_loader_builds_h3s_upscaler_from_its_weights(upscaler_dir):
|
|
from modules.latent.second_pass.nodes import FunUILatentUpscalerLoader
|
|
model = FunUILatentUpscalerLoader.execute("h3_up.safetensors").result[0]
|
|
assert callable(model.funui_latent_upscale)
|
|
|
|
|
|
def test_an_unknown_file_is_refused_with_where_to_go(upscaler_dir):
|
|
from modules.latent.second_pass.nodes import FunUILatentUpscalerLoader
|
|
with pytest.raises(RuntimeError, match="Load Latent Upscale Model"):
|
|
FunUILatentUpscalerLoader.execute("other.safetensors")
|
|
|
|
|
|
def test_the_newer_attention_upscaler_is_named_as_such(upscaler_dir):
|
|
from modules.latent.second_pass.nodes import FunUILatentUpscalerLoader
|
|
with pytest.raises(RuntimeError, match="attention-based"):
|
|
FunUILatentUpscalerLoader.execute("attn.safetensors")
|
|
|
|
|
|
def _nested(h=4, w=4):
|
|
from comfy.nested_tensor import NestedTensor
|
|
torch.manual_seed(1)
|
|
return NestedTensor((torch.randn(1, 24, 2, h, w), torch.randn(1, 32, 2, 3)))
|
|
|
|
|
|
def _pinned(h=4, w=4):
|
|
z = torch.randn(1, 24, 1, h, w)
|
|
return [[torch.zeros(1, 5, 48), {"minimax_keyframes": [{"resolved_frame_index": 0, "latent": z}]}]]
|
|
|
|
|
|
def test_sharpen_keeps_the_size_changes_the_video_and_leaves_audio(comfyui):
|
|
from modules.latent.second_pass.nodes import FunUILatentResample
|
|
samples = _nested()
|
|
pos = _pinned()
|
|
out, p, _n, status = FunUILatentResample.execute(
|
|
{"samples": samples, "noise_mask": torch.ones(1)}, _small_resizer(), "sharpen", 2.0,
|
|
positive=pos).result
|
|
video, audio = out["samples"].unbind()
|
|
assert video.shape == samples.unbind()[0].shape
|
|
assert not torch.allclose(video, samples.unbind()[0])
|
|
assert torch.equal(audio, samples.unbind()[1])
|
|
assert p is pos and "noise_mask" in out and "sharpen" in status
|
|
|
|
|
|
def test_upscale_doubles_the_grid_and_resizes_the_pin(comfyui):
|
|
from modules.latent.second_pass.nodes import FunUILatentResample
|
|
out, pos, _n, status = FunUILatentResample.execute(
|
|
{"samples": _nested()}, _small_resizer(), "upscale", 2.0, positive=_pinned()).result
|
|
assert tuple(out["samples"].unbind()[0].shape[-2:]) == (8, 8)
|
|
assert tuple(pos[0][1]["minimax_keyframes"][0]["latent"].shape[-2:]) == (8, 8)
|
|
assert "1 anchor pin(s) resized" in status
|
|
|
|
|
|
def test_a_resized_pin_fits_a_real_h3_forward_and_a_stale_one_does_not(tiny_h3):
|
|
"""The reason rescaling exists: H3 refuses a pin sized to the old grid."""
|
|
from modules.models.minimax_h3.anchor import rescale_pins
|
|
dm = tiny_h3.patcher.model.diffusion_model
|
|
video = torch.randn(1, 24, 2, 8, 8)
|
|
stale = _pinned(4, 4)
|
|
fresh, n = rescale_pins(stale, 8, 8)
|
|
assert n == 1
|
|
|
|
def fwd(cond):
|
|
kf = cond[0][1]["minimax_keyframes"]
|
|
payload = {"keyframes": kf, "cond_video_latents": [k["latent"] for k in kf]}
|
|
with torch.inference_mode():
|
|
return dm([video, tiny_h3.audio], torch.tensor([500.0]), tiny_h3.context,
|
|
transformer_options={}, minimax_payload=payload)
|
|
|
|
with pytest.raises(RuntimeError, match="shape mismatch"):
|
|
fwd(stale)
|
|
assert fwd(fresh)[0].shape == video.shape
|
|
|
|
|
|
def test_conditioning_no_model_module_recognises_passes_through():
|
|
from modules.latent.second_pass.nodes import rescale
|
|
plain = [[torch.zeros(1, 3, 8), {}]]
|
|
assert rescale(plain, 8, 8) == (plain, 0)
|