daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Match Comfy flow model sigma round trip

Browse source
This commit is contained in:
Daniel Maddern 2026-08-13 22:36:58 +07:00
parent d289b3c497
commit 4018fbea15
2 changed files with 9 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/h3_blackwell_runtime/packing.py
View file

@ -70,7 +70,7 @@ class H3PromptPacker:
text: torch.Tensor, text: torch.Tensor,
video: torch.Tensor, video: torch.Tensor,
audio: torch.Tensor, audio: torch.Tensor,
sigma: float, sigma: float | torch.Tensor,
model_timesteps: torch.Tensor | None = None, model_timesteps: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]: ) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]:
if text.shape[-1] == 5120: if text.shape[-1] == 5120:
@ -84,7 +84,7 @@ class H3PromptPacker:
text_length, audio_length = text_rows.shape[0], audio_rows.shape[0] text_length, audio_length = text_rows.shape[0], audio_rows.shape[0]
hidden = torch.cat((text_rows, audio_rows, video_rows)) hidden = torch.cat((text_rows, audio_rows, video_rows))
if model_timesteps is None: if model_timesteps is None:
video_sigma = torch.tensor(float(sigma), device=hidden.device).clamp(min=1e-6) video_sigma = torch.as_tensor(sigma, device=hidden.device, dtype=torch.float32).clamp(min=1e-6)
base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0)) base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base) audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base)
video_time, audio_time = (1.0 - video_sigma).item(), (1.0 - audio_sigma).item() video_time, audio_time = (1.0 - video_sigma).item(), (1.0 - audio_sigma).item()

9
src/h3_blackwell_runtime/sampler.py
View file

@ -47,6 +47,11 @@ def _audio_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
return 3.0 * base / (1.0 + (3.0 - 1.0) * base) 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
@torch.inference_mode() @torch.inference_mode()
def sample_video_res_multistep( def sample_video_res_multistep(
model, model,
@ -73,7 +78,7 @@ def sample_video_res_multistep(
carry = sigma_audio / sigma carry = sigma_audio / sigma
native_audio = audio_carried.to(torch.bfloat16) * carry native_audio = audio_carried.to(torch.bfloat16) * carry
step_timesteps = None if model_timesteps is None else model_timesteps[previous_index] 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, float(sigma), step_timesteps) 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_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
raw_video = raw_video.to(torch.bfloat16).float() raw_video = raw_video.to(torch.bfloat16).float()
raw_audio = raw_audio.to(torch.bfloat16) raw_audio = raw_audio.to(torch.bfloat16)
@ -105,7 +110,7 @@ def sample_video_euler(model, packer: H3PromptPacker, text: torch.Tensor, video:
sigmas = beta_sigmas(steps, device=video.device) sigmas = beta_sigmas(steps, device=video.device)
for index in range(steps): for index in range(steps):
sigma = sigmas[index] sigma = sigmas[index]
hidden, timesteps, segments, positions, video_segment, audio_segment = packer(text, video, audio, float(sigma)) 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, _ = model(hidden, timesteps, positions, segments, video_segment, audio_segment)
velocity = unpatchify_video(velocity, video.shape[2], video.shape[-2], video.shape[-1]) 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)) video.add_(velocity.to(video.dtype), alpha=float(sigmas[index + 1] - sigma))