digital-garbageandClaude Opus 5.5
106163d039
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
115 lines
4.8 KiB
Python
115 lines
4.8 KiB
Python
"""STAS (stas.py + samplers._install_stas): the right rows and channels are set to
|
|
alpha x peak x sign, only on the first steps, only on the picture rows."""
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
import stas # noqa: E402
|
|
|
|
|
|
def test_target_rows_are_frame0_plus_each_frames_edges():
|
|
rows = stas.target_rows(frames=3, per_frame=25, edge=0.08) # k = 2
|
|
expect = set(range(25)) | {25, 26, 48, 49, 50, 51, 73, 74}
|
|
assert set(rows.tolist()) == expect
|
|
|
|
|
|
def test_only_massive_channels_are_found_and_steered():
|
|
v = torch.randn(200, 128) * 0.1
|
|
v[:, 5] = 20.0 * torch.sign(torch.randn(200)) # massive on every token
|
|
v[3, 5] = -40.0 # its peak
|
|
dims, peaks, _ = stas.ma_dims(v)
|
|
assert dims.tolist() == [5] and float(peaks[0]) == 40.0
|
|
before = v.clone()
|
|
stas.steer(v, torch.tensor([0, 1, 2]), dims, peaks, 2.0)
|
|
assert torch.equal(v[:3, 5], 80.0 * torch.sign(before[:3, 5]))
|
|
assert torch.equal(v[3:], before[3:]) and torch.equal(v[:, :5], before[:, :5])
|
|
|
|
|
|
def test_first_steps():
|
|
assert stas.first_steps(50) == 20 and stas.first_steps(7) == 3 and stas.first_steps(1) == 1
|
|
|
|
|
|
class _Net:
|
|
patch_size = (1, 2, 2)
|
|
|
|
def __init__(self):
|
|
self.blocks = [lambda h: h] * 4
|
|
|
|
|
|
class _Model:
|
|
def __init__(self):
|
|
self.model_options = {"transformer_options": {}}
|
|
self.model = type("M", (), {})()
|
|
self.model.diffusion_model = _Net()
|
|
|
|
def clone(self):
|
|
m = _Model()
|
|
m.model_options = {k: (dict(v) if isinstance(v, dict) else v)
|
|
for k, v in self.model_options.items()}
|
|
return m
|
|
|
|
|
|
def _call(patched, h, step, total, video_rows):
|
|
import samplers
|
|
hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][("double_block", 2)]
|
|
sched = torch.linspace(1.0, 0.0, total + 1)
|
|
args = {"img": h, "mod_segments": video_rows,
|
|
"transformer_options": {"sample_sigmas": sched, "sigmas": sched[step:step + 1]}}
|
|
return hook(args, {"original_block": lambda a: {"img": a["img"]}})["img"], samplers
|
|
|
|
|
|
def test_hook_steers_picture_rows_on_early_steps_only(monkeypatch):
|
|
import samplers
|
|
import h3_repr_steering as rs
|
|
# 10 text rows, then 2 frames x (4x6 latent -> 2x3 = 6 tokens) = 12 video rows, then audio
|
|
def mask(video_rows, seq_len, device):
|
|
m = torch.zeros(seq_len, dtype=torch.bool)
|
|
m[10:22] = True
|
|
return m
|
|
monkeypatch.setattr(rs, "video_mask_from_mod_segments", mask)
|
|
samples = torch.zeros(1, 24, 2, 4, 6)
|
|
patched, report = samplers.FunPackLTXAVSceneChainSampler()._install_stas(_Model(), 2, 2.0, samples)
|
|
h = torch.randn(30, 64) * 0.1
|
|
h[10:22, 3] = 30.0
|
|
h[12, 3] = 50.0
|
|
out, _ = _call(patched, h.clone(), step=0, total=7, video_rows=None)
|
|
# every video row is in S here (frame 0 = all of frame 0; edges k=1 of 6 per frame)
|
|
targets = [10, 11, 12, 13, 14, 15, 16, 21]
|
|
assert torch.allclose(out[targets, 3], torch.full((8,), 100.0))
|
|
assert torch.equal(out[:10], h[:10]) and torch.equal(out[22:], h[22:]) # text/audio
|
|
assert torch.equal(out[17:21], h[17:21]) # interior
|
|
late, _ = _call(patched, h.clone(), step=5, total=7, video_rows=None)
|
|
assert torch.equal(late, h) # past 40%
|
|
assert "channels 3" in report()[0] and report()[1] == 1
|
|
|
|
|
|
def test_unreadable_layout_is_declared(monkeypatch):
|
|
import samplers
|
|
import h3_repr_steering as rs
|
|
monkeypatch.setattr(rs, "video_mask_from_mod_segments", lambda *a: None)
|
|
patched, report = samplers.FunPackLTXAVSceneChainSampler()._install_stas(
|
|
_Model(), 2, 2.0, torch.zeros(1, 24, 2, 4, 6))
|
|
h = torch.randn(30, 8)
|
|
out, _ = _call(patched, h.clone(), step=0, total=7, video_rows=None)
|
|
assert torch.equal(out, h) and "nothing steered" in report()[0] and report()[1] == 0
|
|
|
|
|
|
def test_weak_branch_calls_are_steered_but_not_counted(monkeypatch):
|
|
import samplers
|
|
import h3_repr_steering as rs
|
|
monkeypatch.setattr(rs, "video_mask_from_mod_segments",
|
|
lambda v, n, d: torch.arange(n) < 12)
|
|
patched, report = samplers.FunPackLTXAVSceneChainSampler()._install_stas(
|
|
_Model(), 2, 2.0, torch.zeros(1, 24, 2, 4, 6))
|
|
hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][("double_block", 2)]
|
|
sched = torch.linspace(1.0, 0.0, 8)
|
|
h = torch.randn(20, 64) * 0.1
|
|
h[:12, 3] = 30.0
|
|
to = {"sample_sigmas": sched, "sigmas": sched[:1], samplers.WEAK_BRANCH_FLAG: True}
|
|
out = hook({"img": h.clone(), "mod_segments": None, "transformer_options": to},
|
|
{"original_block": lambda a: {"img": a["img"]}})["img"]
|
|
assert float(out[0, 3]) == 60.0 and report()[1] == 0
|