182 lines
8.9 KiB
Python
182 lines
8.9 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)
|
|
|
|
|
|
def _model_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
|
|
"""Comfy BaseModel's flow timestep round-trip seen by the H3 diffusion model."""
|
|
return (video_sigma * 1000.0).float() / 1000.0
|
|
|
|
|
|
def _decode_audio_latent(audio_carried: torch.Tensor, *, shift_video: float = 12.0, shift_audio: float = 3.0) -> torch.Tensor:
|
|
"""Convert Comfy's AV sampler-carried audio state back to native audio VAE latents."""
|
|
return audio_carried * (shift_audio / shift_video)
|
|
|
|
|
|
def _cache_sample(x: torch.Tensor, factor: int) -> torch.Tensor:
|
|
if factor <= 1:
|
|
return x
|
|
if x.ndim == 5:
|
|
return x[..., ::factor, ::factor]
|
|
return x[..., ::factor]
|
|
|
|
|
|
@torch.inference_mode()
|
|
def sample_video_res_multistep(
|
|
model,
|
|
packer: H3PromptPacker,
|
|
text: torch.Tensor,
|
|
video: torch.Tensor,
|
|
audio: torch.Tensor,
|
|
*,
|
|
steps: int = 12,
|
|
model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None,
|
|
return_audio: bool = False,
|
|
progress: bool = False,
|
|
cache_mode: str | None = None,
|
|
cache_threshold: float = 0.0,
|
|
cache_start_percent: float = 0.0,
|
|
cache_end_percent: float = 1.0,
|
|
cache_subsample_factor: int = 2,
|
|
cache_stats: dict | None = None,
|
|
) -> 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()
|
|
cache_mode = None if cache_mode in {None, "", "disabled", "none"} else cache_mode
|
|
if cache_mode not in {None, "easycache", "h3_cache"}:
|
|
raise ValueError(f"Unsupported cache mode: {cache_mode}")
|
|
if cache_stats is not None:
|
|
cache_stats.update({"mode": cache_mode, "threshold": cache_threshold, "skipped_steps": 0, "rates": []})
|
|
cache = {
|
|
"video_diff": None,
|
|
"audio_diff": None,
|
|
"video_prev": None,
|
|
"audio_prev": None,
|
|
"prev_norm": None,
|
|
"cumulative_rate": 0.0,
|
|
}
|
|
for index, sigma in enumerate(sigmas[:-1], start=1):
|
|
step_started = time.perf_counter()
|
|
previous_index = index - 1
|
|
sigma_down = sigmas[index]
|
|
current_percent = previous_index / total_steps
|
|
can_cache = cache_mode is not None and cache_threshold > 0 and cache_start_percent <= current_percent <= cache_end_percent and cache["video_diff"] is not None
|
|
skipped = False
|
|
if can_cache:
|
|
video_now = _cache_sample(video, cache_subsample_factor)
|
|
audio_now = _cache_sample(audio_carried, cache_subsample_factor)
|
|
input_change = (video_now - cache["video_prev"]).flatten().abs().mean() + (audio_now - cache["audio_prev"]).flatten().abs().mean()
|
|
input_norm = cache["prev_norm"].clamp_min(1e-8)
|
|
rate = (input_change / input_norm).item()
|
|
if cache_mode == "easycache":
|
|
cache["cumulative_rate"] += rate
|
|
skipped = cache["cumulative_rate"] < cache_threshold
|
|
if not skipped:
|
|
cache["cumulative_rate"] = 0.0
|
|
else:
|
|
skipped = rate < cache_threshold
|
|
if cache_stats is not None:
|
|
cache_stats["rates"].append({"step": previous_index, "rate": rate, "skipped": skipped})
|
|
if skipped:
|
|
video_denoised = video + cache["video_diff"]
|
|
audio_denoised = audio_carried + cache["audio_diff"]
|
|
if cache_stats is not None:
|
|
cache_stats["skipped_steps"] += 1
|
|
else:
|
|
sigma_audio = _audio_sigma(sigma)
|
|
carry = sigma_audio / sigma
|
|
native_audio = audio_carried.to(torch.bfloat16) * carry
|
|
step_timesteps = None if model_timesteps is None else model_timesteps[previous_index]
|
|
hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps)
|
|
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])
|
|
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
|
|
if cache_mode is not None:
|
|
cache["video_diff"] = (video_denoised - video).detach()
|
|
cache["audio_diff"] = (audio_denoised - audio_carried).detach()
|
|
cache["video_prev"] = _cache_sample(video, cache_subsample_factor).detach().clone()
|
|
cache["audio_prev"] = _cache_sample(audio_carried, cache_subsample_factor).detach().clone()
|
|
cache["prev_norm"] = video.flatten().abs().mean() + audio_carried.flatten().abs().mean()
|
|
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
|
|
if progress:
|
|
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, _decode_audio_latent(audio_carried)) if return_audio else 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, _model_sigma(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
|