Nynxz
f1a98ee8d1
The forwarding of downscale_ratio_spacial/temporal into fix_empty_latent_channels already landed on main (v2.3.22). What was still missing: stock common_ksampler pops both tags from the returned latent after resampling. Without the pop, a sampled latent keeps a stale 8x tag after being rescaled to the model's native ratio, so a downstream sampler (or a chained KSampler) would rescale it a second time. Also assert the forwarding in tests: the fix_empty_latent_channels mock now captures the ratio args, and TestDownscaleRatioForwarding checks both KSamplers forward the tag and strip it from the output.
283 lines
10 KiB
Python
283 lines
10 KiB
Python
"""
|
|
Unit tests for ROCmKSamplerAdvanced
|
|
"""
|
|
import types
|
|
import pytest
|
|
import torch
|
|
|
|
import sys
|
|
|
|
|
|
def _install_comfy_sample_mocks(monkeypatch, captured):
|
|
import rocm_nodes.core.sampler as sampler_module
|
|
comfy = sampler_module.comfy
|
|
|
|
def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None,
|
|
downscale_ratio_temporal=None):
|
|
captured['downscale_ratio_spacial'] = downscale_ratio_spacial
|
|
captured['downscale_ratio_temporal'] = downscale_ratio_temporal
|
|
return latent_image
|
|
|
|
def prepare_noise(latent, seed, batch_inds=None):
|
|
torch.manual_seed(seed % (2**31))
|
|
return torch.randn_like(latent)
|
|
|
|
def sample(model, noise, steps, cfg, sampler_name, scheduler,
|
|
positive, negative, latent_image, denoise,
|
|
noise_mask=None, callback=None, disable_pbar=False, seed=0,
|
|
start_step=0, last_step=None, force_full_denoise=True,
|
|
disable_noise=False):
|
|
captured.update({
|
|
'noise': noise,
|
|
'steps': steps,
|
|
'cfg': cfg,
|
|
'sampler_name': sampler_name,
|
|
'scheduler': scheduler,
|
|
'latent_image': latent_image,
|
|
'denoise': denoise,
|
|
'start_step': start_step,
|
|
'last_step': last_step,
|
|
'seed': seed,
|
|
})
|
|
return torch.zeros_like(latent_image)
|
|
|
|
monkeypatch.setattr(comfy.sample, 'fix_empty_latent_channels', fix_empty_latent_channels, raising=False)
|
|
monkeypatch.setattr(comfy.sample, 'prepare_noise', prepare_noise, raising=False)
|
|
monkeypatch.setattr(comfy.sample, 'sample', sample, raising=False)
|
|
|
|
|
|
def _make_model(latent_channels=4):
|
|
"""Create a minimal model mock with architecture info for detect_model_sampling_type."""
|
|
lat_fmt = types.SimpleNamespace(
|
|
latent_channels=latent_channels,
|
|
spacial_downscale_ratio=8,
|
|
)
|
|
model_type = types.SimpleNamespace(
|
|
name="FLOW",
|
|
)
|
|
inner = types.SimpleNamespace(
|
|
model_type=model_type,
|
|
latent_format=lat_fmt,
|
|
memory_usage_factor=1.0,
|
|
)
|
|
return types.SimpleNamespace(
|
|
model=inner,
|
|
model_dtype=lambda: torch.float32,
|
|
load_device='cpu',
|
|
)
|
|
|
|
|
|
class TestROCmKSamplerAdvanced:
|
|
def test_add_noise_disabled_uses_latent_as_noise(self, monkeypatch, sample_conditioning):
|
|
from rocm_nodes.core.sampler import ROCmKSamplerAdvanced
|
|
|
|
import comfy
|
|
comfy.utils.PROGRESS_BAR_ENABLED = False
|
|
sys.modules.setdefault('latent_preview', types.ModuleType('latent_preview'))
|
|
import latent_preview
|
|
def prepare_callback(model, steps):
|
|
return None
|
|
setattr(latent_preview, 'prepare_callback', prepare_callback)
|
|
|
|
captured = {}
|
|
_install_comfy_sample_mocks(monkeypatch, captured)
|
|
|
|
node = ROCmKSamplerAdvanced()
|
|
latent = {"samples": torch.randn(1, 4, 32, 32)}
|
|
model = _make_model()
|
|
|
|
out, = node.sample(
|
|
model=model,
|
|
add_noise="disable",
|
|
noise_seed=123,
|
|
steps=4,
|
|
cfg=1.0,
|
|
sampler_name='euler',
|
|
scheduler='simple',
|
|
positive=sample_conditioning,
|
|
negative=sample_conditioning,
|
|
latent_image=latent,
|
|
start_at_step=0,
|
|
end_at_step=10000,
|
|
return_with_leftover_noise="disable",
|
|
denoise=1.0,
|
|
compatibility_mode=False,
|
|
)
|
|
|
|
assert isinstance(out, dict)
|
|
assert 'samples' in out
|
|
assert isinstance(captured['noise'], torch.Tensor)
|
|
assert tuple(captured['noise'].shape) == tuple(latent['samples'].shape)
|
|
# last_step normalization happens inside comfy.sample.sample (mocked away);
|
|
# our node passes end_at_step through as-is
|
|
assert captured['last_step'] == 10000
|
|
|
|
def test_step_bounds_normalization(self, monkeypatch, sample_conditioning):
|
|
from rocm_nodes.core.sampler import ROCmKSamplerAdvanced
|
|
|
|
import comfy
|
|
comfy.utils.PROGRESS_BAR_ENABLED = False
|
|
sys.modules.setdefault('latent_preview', types.ModuleType('latent_preview'))
|
|
import latent_preview
|
|
setattr(latent_preview, 'prepare_callback', lambda model, steps: None)
|
|
|
|
captured = {}
|
|
_install_comfy_sample_mocks(monkeypatch, captured)
|
|
|
|
node = ROCmKSamplerAdvanced()
|
|
latent = {"samples": torch.randn(1, 4, 16, 16)}
|
|
model = _make_model()
|
|
|
|
out, = node.sample(
|
|
model=model,
|
|
add_noise="enable",
|
|
noise_seed=0,
|
|
steps=10,
|
|
cfg=1.0,
|
|
sampler_name='euler',
|
|
scheduler='simple',
|
|
positive=sample_conditioning,
|
|
negative=sample_conditioning,
|
|
latent_image=latent,
|
|
start_at_step=-5,
|
|
end_at_step=10,
|
|
return_with_leftover_noise="disable",
|
|
denoise=1.0,
|
|
compatibility_mode=False,
|
|
)
|
|
|
|
assert isinstance(out, dict)
|
|
# start_step / end_at_step normalization happens inside comfy.sample.sample
|
|
# (mocked away); our node passes values through as-is
|
|
assert captured['start_step'] == -5
|
|
assert captured['last_step'] == 10
|
|
|
|
def test_model_detection_flow(self, monkeypatch, sample_conditioning):
|
|
"""Verify flow-matching model is detected and logged."""
|
|
from rocm_nodes.core.sampler import ROCmKSampler
|
|
|
|
import comfy
|
|
comfy.utils.PROGRESS_BAR_ENABLED = False
|
|
sys.modules.setdefault('latent_preview', types.ModuleType('latent_preview'))
|
|
import latent_preview
|
|
setattr(latent_preview, 'prepare_callback', lambda model, steps: None)
|
|
|
|
captured = {}
|
|
_install_comfy_sample_mocks(monkeypatch, captured)
|
|
|
|
node = ROCmKSampler()
|
|
latent = {"samples": torch.randn(1, 4, 32, 32)}
|
|
model = _make_model(latent_channels=16)
|
|
|
|
out, = node.sample(
|
|
model=model,
|
|
seed=0,
|
|
steps=4,
|
|
cfg=1.0,
|
|
sampler_name='euler',
|
|
scheduler='simple',
|
|
positive=sample_conditioning,
|
|
negative=sample_conditioning,
|
|
latent_image=latent,
|
|
denoise=1.0,
|
|
)
|
|
|
|
assert isinstance(out, dict)
|
|
assert 'samples' in out
|
|
assert isinstance(out['samples'], torch.Tensor)
|
|
|
|
def test_model_detection_pixel_space(self, monkeypatch, sample_conditioning):
|
|
"""Verify pixel-space model is detected correctly."""
|
|
from rocm_nodes.core.sampler import ROCmKSampler
|
|
|
|
import comfy
|
|
comfy.utils.PROGRESS_BAR_ENABLED = False
|
|
sys.modules.setdefault('latent_preview', types.ModuleType('latent_preview'))
|
|
import latent_preview
|
|
setattr(latent_preview, 'prepare_callback', lambda model, steps: None)
|
|
|
|
captured = {}
|
|
_install_comfy_sample_mocks(monkeypatch, captured)
|
|
|
|
node = ROCmKSampler()
|
|
latent = {"samples": torch.randn(1, 3, 64, 64)}
|
|
|
|
lat_fmt = types.SimpleNamespace(
|
|
latent_channels=3,
|
|
spacial_downscale_ratio=1,
|
|
)
|
|
model_type = types.SimpleNamespace(name="FLOW")
|
|
inner = types.SimpleNamespace(
|
|
model_type=model_type,
|
|
latent_format=lat_fmt,
|
|
memory_usage_factor=0.03,
|
|
)
|
|
model = types.SimpleNamespace(
|
|
model=inner,
|
|
model_dtype=lambda: torch.float32,
|
|
load_device='cpu',
|
|
)
|
|
|
|
out, = node.sample(
|
|
model=model,
|
|
seed=0,
|
|
steps=4,
|
|
cfg=1.0,
|
|
sampler_name='euler',
|
|
scheduler='simple',
|
|
positive=sample_conditioning,
|
|
negative=sample_conditioning,
|
|
latent_image=latent,
|
|
denoise=1.0,
|
|
)
|
|
|
|
assert isinstance(out, dict)
|
|
assert 'samples' in out
|
|
|
|
|
|
class TestDownscaleRatioForwarding:
|
|
"""EmptyLatentImage tags latents with downscale_ratio_spacial; the KSamplers
|
|
must forward it to fix_empty_latent_channels and strip it from the output,
|
|
like stock common_ksampler (non-8x models sample at the wrong size without it)."""
|
|
|
|
def _run(self, monkeypatch, sample_conditioning, node_cls, **kwargs):
|
|
import comfy
|
|
import sys
|
|
import types
|
|
import latent_preview
|
|
comfy.utils.PROGRESS_BAR_ENABLED = False
|
|
sys.modules.setdefault('latent_preview', types.ModuleType('latent_preview'))
|
|
setattr(latent_preview, 'prepare_callback', lambda model, steps: None)
|
|
|
|
captured = {}
|
|
_install_comfy_sample_mocks(monkeypatch, captured)
|
|
|
|
latent = {"samples": torch.zeros(1, 4, 32, 32),
|
|
"downscale_ratio_spacial": 8}
|
|
out, = node_cls().sample(
|
|
model=_make_model(),
|
|
latent_image=latent,
|
|
positive=sample_conditioning,
|
|
negative=sample_conditioning,
|
|
**kwargs,
|
|
)
|
|
return captured, out
|
|
|
|
def test_ksampler_forwards_and_strips(self, monkeypatch, sample_conditioning):
|
|
from rocm_nodes.core.sampler import ROCmKSampler
|
|
captured, out = self._run(monkeypatch, sample_conditioning, ROCmKSampler,
|
|
seed=0, steps=4, cfg=1.0, sampler_name='euler',
|
|
scheduler='simple', denoise=1.0)
|
|
assert captured['downscale_ratio_spacial'] == 8
|
|
assert 'downscale_ratio_spacial' not in out
|
|
|
|
def test_ksampler_advanced_forwards_and_strips(self, monkeypatch, sample_conditioning):
|
|
from rocm_nodes.core.sampler import ROCmKSamplerAdvanced
|
|
captured, out = self._run(monkeypatch, sample_conditioning, ROCmKSamplerAdvanced,
|
|
add_noise="enable", noise_seed=0, steps=4, cfg=1.0,
|
|
sampler_name='euler', scheduler='simple',
|
|
start_at_step=0, end_at_step=10000,
|
|
return_with_leftover_noise="disable", denoise=1.0)
|
|
assert captured['downscale_ratio_spacial'] == 8
|
|
assert 'downscale_ratio_spacial' not in out
|