39 lines
1.8 KiB
Python
39 lines
1.8 KiB
Python
"""Direct prompt-only MiniMax H3 T2V shape and latent helpers."""
|
|
|
|
import torch
|
|
|
|
|
|
FPS = 24
|
|
AUDIO_LATENT_FPS = 40
|
|
|
|
|
|
def align_frame_count(frames: int) -> int:
|
|
"""Snap to H3's valid 17k+5 temporal grid."""
|
|
frames = max(5, frames)
|
|
return frames + (5 - frames) % 17
|
|
|
|
|
|
def temporal_shape(frames: int) -> tuple[int, int, int]:
|
|
"""Return output frames, video-latent frames, and joint audio-latent steps."""
|
|
frames = align_frame_count(frames)
|
|
video_steps = 2 if frames <= 5 else ((frames - 5) // 17) * 5 + 2
|
|
return frames, video_steps, round(frames / FPS * AUDIO_LATENT_FPS)
|
|
|
|
|
|
def empty_av_latents(width: int, height: int, frames: int, *, device: torch.device | str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
|
|
"""Allocate H3's joint video/audio sampling state without ComfyUI objects."""
|
|
if width % 32 or height % 32:
|
|
raise ValueError("H3 T2V dimensions must be multiples of 32.")
|
|
frames, video_steps, audio_steps = temporal_shape(frames)
|
|
video = torch.zeros((1, 24, video_steps, height // 16, width // 16), device=device)
|
|
audio = torch.zeros((1, 32, 2, audio_steps), device=device)
|
|
return video, audio, frames
|
|
|
|
|
|
def random_av_latents(width: int, height: int, frames: int, seed: int, *, device: torch.device | str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
|
|
"""Generate Comfy-equivalent CPU-seeded joint AV noise latents."""
|
|
video, audio, frames = empty_av_latents(width, height, frames, device="cpu")
|
|
generator = torch.manual_seed(seed)
|
|
video = torch.randn(video.shape, dtype=torch.float32, generator=generator, device="cpu").to(video.dtype)
|
|
audio = torch.randn(audio.shape, dtype=torch.float32, generator=generator, device="cpu").to(audio.dtype)
|
|
return video.to(device), audio.to(device), frames
|