102 lines
5.3 KiB
Python
102 lines
5.3 KiB
Python
"""Minimal direct prompt-only H3 video sampler for visual smoke previews."""
|
|
|
|
import time
|
|
|
|
import torch
|
|
|
|
from .packing import H3PromptPacker, unpatchify_video
|
|
|
|
|
|
def shifted_sigma(base: torch.Tensor, shift: float) -> torch.Tensor:
|
|
"""H3 flow-SNR shift used by Comfy's ModelSamplingDiscreteFlow."""
|
|
return shift * base / (1 + (shift - 1) * base)
|
|
|
|
|
|
def beta_sigmas(steps: int, *, device: torch.device | str, alpha: float = 0.6, beta: float = 0.6) -> torch.Tensor:
|
|
"""Comfy's discrete beta scheduler over H3's 1,000-entry shift-12 table."""
|
|
from scipy.stats import beta as beta_distribution
|
|
|
|
table = shifted_sigma(torch.arange(1, 1001, device=device, dtype=torch.float32) / 1000, 12.0)
|
|
fractions = 1.0 - torch.arange(steps, device=device, dtype=torch.float64).cpu().numpy() / steps
|
|
indices = torch.from_numpy((999 * beta_distribution.ppf(fractions, alpha, beta)).round().astype("int64")).to(device)
|
|
indices = torch.unique_consecutive(indices)
|
|
return torch.cat((table[indices], table.new_zeros(1)))
|
|
|
|
|
|
def res_multistep_update(x: torch.Tensor, denoised: torch.Tensor, sigma: torch.Tensor, sigma_down: torch.Tensor, old_denoised: torch.Tensor | None, old_sigma_down: torch.Tensor | None, previous_sigma: torch.Tensor | None) -> torch.Tensor:
|
|
"""Deterministic H3/Comfy RES multistep update."""
|
|
if sigma_down == 0 or old_denoised is None:
|
|
return x + (x - denoised) / sigma * (sigma_down - sigma)
|
|
t, t_old, t_next, t_prev = sigma.log().neg(), old_sigma_down.log().neg(), sigma_down.log().neg(), previous_sigma.log().neg()
|
|
h = t_next - t
|
|
c2 = (t_prev - t_old) / h
|
|
phi1 = torch.expm1(-h) / -h
|
|
phi2 = (phi1 - 1.0) / -h
|
|
b1 = torch.nan_to_num(phi1 - phi2 / c2, nan=0.0)
|
|
b2 = torch.nan_to_num(phi2 / c2, nan=0.0)
|
|
return torch.exp(-h) * x + h * (b1 * denoised + b2 * old_denoised)
|
|
|
|
|
|
def _unpack_audio(rows: torch.Tensor) -> torch.Tensor:
|
|
steps = rows.shape[0] // 2
|
|
return rows.reshape(2, steps, 32).permute(2, 0, 1).unsqueeze(0)
|
|
|
|
|
|
def _audio_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
|
|
base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
|
|
return 3.0 * base / (1.0 + (3.0 - 1.0) * base)
|
|
|
|
|
|
@torch.inference_mode()
|
|
def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, *, steps: int = 12) -> torch.Tensor:
|
|
"""Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
|
|
sigmas = beta_sigmas(steps, device=video.device)
|
|
audio_carried = audio
|
|
video_history = audio_history = None
|
|
video_history_sigma = audio_history_sigma = None
|
|
total_steps = len(sigmas) - 1
|
|
started = time.perf_counter()
|
|
for index, sigma in enumerate(sigmas[:-1], start=1):
|
|
step_started = time.perf_counter()
|
|
previous_index = index - 1
|
|
sigma_down = sigmas[index]
|
|
sigma_audio = _audio_sigma(sigma)
|
|
native_audio = audio_carried * (sigma_audio / sigma)
|
|
hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
|
|
raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
|
|
raw_video = raw_video.to(torch.bfloat16).float()
|
|
raw_audio = raw_audio.to(torch.bfloat16)
|
|
velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
|
|
carry = sigma_audio / sigma
|
|
velocity_audio = (
|
|
(1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
|
|
+ (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
|
|
).float()
|
|
video_denoised = video - sigma * velocity_video
|
|
audio_denoised = audio_carried - sigma * velocity_audio
|
|
previous_sigma = sigmas[previous_index - 1] if previous_index else None
|
|
video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
|
|
audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)
|
|
video_history, audio_history = video_denoised, audio_denoised
|
|
video_history_sigma = audio_history_sigma = sigma_down
|
|
elapsed = time.perf_counter() - started
|
|
eta = elapsed / index * (total_steps - index)
|
|
print(
|
|
f"sampling step {index}/{total_steps}: "
|
|
f"{time.perf_counter() - step_started:.1f}s, elapsed {elapsed:.1f}s, eta {eta:.1f}s",
|
|
flush=True,
|
|
)
|
|
return video
|
|
|
|
|
|
@torch.inference_mode()
|
|
def sample_video_euler(model, packer: H3PromptPacker, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, *, steps: int = 2) -> torch.Tensor:
|
|
"""Use Euler updates to obtain a visual-only H3 preview, not parity sampling."""
|
|
sigmas = beta_sigmas(steps, device=video.device)
|
|
for index in range(steps):
|
|
sigma = sigmas[index]
|
|
hidden, timesteps, segments, positions, video_segment, audio_segment = packer(text, video, audio, float(sigma))
|
|
velocity, _ = model(hidden, timesteps, positions, segments, video_segment, audio_segment)
|
|
velocity = unpatchify_video(velocity, video.shape[2], video.shape[-2], video.shape[-1])
|
|
video.add_(velocity.to(video.dtype), alpha=float(sigmas[index + 1] - sigma))
|
|
return video
|