Spaces:
Running on Zero
Running on Zero
| """Single-GPU Ref2VA inference with a fixed first frame and Retake34 continuation.""" | |
| from __future__ import annotations | |
| import gc | |
| import os | |
| from pathlib import Path | |
| import time | |
| from typing import Any | |
| import uuid | |
| from .cache import cache_paths, validate_cache | |
| from .config import RunConfig | |
| from .constants import ( | |
| AUDIO_CONDITION_CLEAN, | |
| AUDIO_LATENT_SHAPE, | |
| AUDIO_PREFIX_LATENTS, | |
| AUDIO_SHIFT, | |
| FPS, | |
| OVERLAP_AUDIO_SAMPLES, | |
| OVERLAP_FRAMES, | |
| RECOVERY_SCHEMA, | |
| REFERENCE_LATENT_SHAPE, | |
| REFERENCE_RUNTIME_SHAPE, | |
| RESULT_SCHEMA, | |
| STRIDE_FRAMES, | |
| TARGET_HEIGHT, | |
| TARGET_WIDTH, | |
| VIDEO_LATENT_SHAPE, | |
| VIDEO_PREFIX_LATENTS, | |
| VIDEO_SHIFT, | |
| VISUAL_CONDITION_CLEAN, | |
| WINDOW_AUDIO_SAMPLES, | |
| WINDOW_FRAMES, | |
| ) | |
| from .io import atomic_json | |
| from .runtime import build_runtime_batch, strict_merge_lora | |
| def _require_one_gpu(torch: Any) -> Any: | |
| if not torch.cuda.is_available() or torch.cuda.device_count() != 1: | |
| raise RuntimeError("generation requires exactly one visible CUDA GPU") | |
| torch.cuda.set_device(0) | |
| return torch.device("cuda:0") | |
| def _unique_key(handle: Any, *, exact: str | None = None, prefix: str | None = None) -> str: | |
| keys = list(handle.keys()) | |
| selected = [key for key in keys if key == exact] if exact else [key for key in keys if key.startswith(str(prefix))] | |
| if len(selected) != 1: | |
| raise ValueError(f"expected one tensor key, got {keys}") | |
| return selected[0] | |
| class Ref2VAEngine: | |
| def __init__(self, config: RunConfig) -> None: | |
| import torch | |
| from musubi_tuner.minimax_h3.audio_vae import load_audio_vae | |
| from musubi_tuner.minimax_h3.model import load_h3_transformer | |
| from musubi_tuner.minimax_h3.video_vae import load_video_vae | |
| from musubi_tuner.modules.custom_offloading_utils import BlockSwapConfig | |
| self.config = config | |
| self.torch = torch | |
| self.device = _require_one_gpu(torch) | |
| transformer = load_h3_transformer( | |
| config.models.dit, | |
| device="cpu", | |
| dtype=torch.bfloat16, | |
| attn_mode=config.attention_mode, | |
| split_attn=False, | |
| disable_mmap=False, | |
| ) | |
| self.lora_gate = strict_merge_lora( | |
| transformer, config.models.lora, config.expected_lora_modules | |
| ) | |
| swap = BlockSwapConfig(device=self.device, supports_backward=False, use_pinned_memory=True) | |
| transformer.enable_block_swap(config.blocks_to_swap, swap) | |
| transformer.move_to_device_except_swap_blocks(self.device) | |
| transformer.prepare_block_swap_before_forward() | |
| transformer.switch_block_swap_for_inference() | |
| transformer.eval().requires_grad_(False) | |
| self.transformer = transformer | |
| self.video_vae = load_video_vae( | |
| config.models.video_vae, device=self.device, dtype=torch.float16, disable_mmap=False | |
| ).eval().requires_grad_(False) | |
| self.audio_vae = load_audio_vae( | |
| config.models.audio_vae, device=self.device, dtype=torch.float32, disable_mmap=False | |
| ).eval().requires_grad_(False) | |
| def _load_window_cache(self, index: int): | |
| import torch | |
| from safetensors import safe_open | |
| paths = cache_paths(self.config, index) | |
| with safe_open(str(paths["anchor"]), framework="pt", device="cpu") as handle: | |
| anchor = handle.get_tensor(_unique_key(handle, prefix="latents_ref_")) | |
| with safe_open(str(paths["reference"]), framework="pt", device="cpu") as handle: | |
| reference = handle.get_tensor(_unique_key(handle, prefix="latents_ref_")) | |
| with safe_open(str(paths["qwen"]), framework="pt", device="cpu") as handle: | |
| hidden = handle.get_tensor(_unique_key(handle, exact="varlen_mmh3_hidden_states_bfloat16")) | |
| tags = handle.get_tensor(_unique_key(handle, exact="varlen_mmh3_token_tags_int64")) | |
| if tuple(reference.shape) != REFERENCE_LATENT_SHAPE or reference.dtype != torch.float32: | |
| raise ValueError(f"window {index} DUV reference cache shape/dtype changed") | |
| reference = torch.nn.functional.pad(reference.unsqueeze(0), (0, 1, 0, 0), value=0.0).contiguous() | |
| if tuple(reference.shape) != REFERENCE_RUNTIME_SHAPE: | |
| raise ValueError(f"window {index} padded DUV reference shape changed") | |
| if anchor.ndim != 4 or tuple(anchor.shape[:2]) != (24, 1) or anchor.dtype != torch.float32: | |
| raise ValueError("anchor cache shape/dtype changed") | |
| if hidden.ndim != 2 or hidden.shape[1] != 5120 or hidden.dtype != torch.bfloat16: | |
| raise ValueError(f"window {index} Qwen hidden states changed") | |
| if tags.shape != (hidden.shape[0],) or tags.dtype != torch.int64: | |
| raise ValueError(f"window {index} Qwen token tags changed") | |
| if not all(torch.isfinite(value).all().item() for value in (anchor, reference, hidden)): | |
| raise ValueError(f"window {index} cache contains non-finite values") | |
| return anchor.unsqueeze(0), reference, hidden.unsqueeze(0), tags.unsqueeze(0) | |
| def _load_input_video(self): | |
| import torch | |
| from safetensors import safe_open | |
| path = cache_paths(self.config, 0)["input_video"] | |
| with safe_open(str(path), framework="pt", device="cpu") as handle: | |
| latent = handle.get_tensor(_unique_key(handle, exact="input_video_latents")) | |
| if tuple(latent.shape) != VIDEO_LATENT_SHAPE[1:] or latent.dtype != torch.float32: | |
| raise ValueError("fixed-first-frame latent shape/dtype changed") | |
| if torch.count_nonzero(latent[:, 1:]).item() != 0: | |
| raise ValueError("fixed-first-frame latent must contain data only at temporal index 0") | |
| latent = latent.unsqueeze(0).to(self.device) | |
| mask = torch.ones_like(latent, dtype=torch.bool) | |
| mask[:, :, 0] = False | |
| return latent, mask | |
| def _recovery_path(self, index: int) -> Path: | |
| return self.config.cache_dir / "recovery" / f"window_{index:03d}.safetensors" | |
| def _load_recovery(self, index: int): | |
| import torch | |
| from safetensors import safe_open | |
| path = self._recovery_path(index) | |
| with safe_open(str(path), framework="pt", device="cpu") as handle: | |
| metadata = handle.metadata() or {} | |
| if metadata.get("schema") != RECOVERY_SCHEMA or metadata.get("window_index") != str(index): | |
| raise ValueError(f"recovery identity changed: {path}") | |
| video = handle.get_tensor("video") | |
| audio = handle.get_tensor("audio") | |
| if tuple(video.shape) != VIDEO_LATENT_SHAPE[1:] or tuple(audio.shape) != AUDIO_LATENT_SHAPE[1:]: | |
| raise ValueError(f"recovery shape changed: {path}") | |
| if video.dtype != torch.float32 or audio.dtype != torch.float32: | |
| raise ValueError(f"recovery dtype changed: {path}") | |
| if not torch.isfinite(video).all().item() or not torch.isfinite(audio).all().item(): | |
| raise ValueError(f"recovery contains non-finite values: {path}") | |
| return video, audio | |
| def _retake_inputs(self, previous: tuple[Any, Any]): | |
| torch = self.torch | |
| previous_video, previous_audio = previous | |
| torch.cuda.empty_cache() | |
| with torch.inference_mode(): | |
| decoded_video = self.video_vae.decode( | |
| previous_video.unsqueeze(0).to(self.device, torch.float16) | |
| )[:, :, :WINDOW_FRAMES] | |
| if tuple(decoded_video.shape) != (1, 3, WINDOW_FRAMES, TARGET_HEIGHT, TARGET_WIDTH): | |
| raise RuntimeError("Retake34 decoded video shape changed") | |
| tail = decoded_video[:, :, -OVERLAP_FRAMES:] | |
| padded = torch.cat((tail, tail[:, :, -1:].repeat(1, 1, STRIDE_FRAMES, 1, 1)), dim=2) | |
| clean_video = self.video_vae.encode(padded.to(torch.float16)).float() | |
| decoded_audio = self.audio_vae.decode( | |
| previous_audio.unsqueeze(0).to(self.device, torch.float32) | |
| )[:, :, :WINDOW_AUDIO_SAMPLES] | |
| if tuple(decoded_audio.shape) != (1, 2, WINDOW_AUDIO_SAMPLES): | |
| raise RuntimeError("Retake34 decoded audio shape changed") | |
| prefix = decoded_audio[:, :, -OVERLAP_AUDIO_SAMPLES:] | |
| prefix_audio = self.audio_vae.encode(prefix.to(torch.float32)).float() | |
| if prefix_audio.shape[-1] != AUDIO_PREFIX_LATENTS: | |
| raise RuntimeError("Retake34 audio prefix shape changed") | |
| clean_audio = torch.cat( | |
| (prefix_audio, prefix_audio[..., -1:].repeat(1, 1, 1, AUDIO_LATENT_SHAPE[-1] - AUDIO_PREFIX_LATENTS)), | |
| dim=-1, | |
| ) | |
| if tuple(clean_video.shape) != VIDEO_LATENT_SHAPE or tuple(clean_audio.shape) != AUDIO_LATENT_SHAPE: | |
| raise RuntimeError("Retake34 clean latent shape changed") | |
| video_mask = torch.ones_like(clean_video, dtype=torch.bool) | |
| video_mask[:, :, :VIDEO_PREFIX_LATENTS] = False | |
| audio_mask = torch.ones_like(clean_audio, dtype=torch.bool) | |
| audio_mask[..., :AUDIO_PREFIX_LATENTS] = False | |
| return clean_video, video_mask, clean_audio, audio_mask | |
| def sample_window(self, index: int) -> Path: | |
| import torch | |
| from musubi_tuner.minimax_h3.sampling import initialize_target_latents, sample_joint_av | |
| from safetensors.torch import save_file | |
| output = self._recovery_path(index) | |
| if output.exists(): | |
| self._load_recovery(index) | |
| return output | |
| anchor, reference, hidden, tags = self._load_window_cache(index) | |
| runtime = build_runtime_batch(anchor, reference, hidden, tags) | |
| window_seed = self.config.seed + index | |
| initial_video, initial_audio = initialize_target_latents( | |
| video_shape=VIDEO_LATENT_SHAPE, | |
| audio_shape=AUDIO_LATENT_SHAPE, | |
| seed=window_seed, | |
| device=self.device, | |
| video_dtype=torch.float32, | |
| audio_dtype=torch.float32, | |
| ) | |
| ref = reference.to(self.device, torch.float32) | |
| ref_generator = torch.Generator(device="cpu").manual_seed(self.config.seed + 1000 + index) | |
| ref_noise = torch.randn(tuple(ref.shape), generator=ref_generator, dtype=torch.float32).to(self.device) | |
| augmented_ref = VISUAL_CONDITION_CLEAN * ref + (1.0 - VISUAL_CONDITION_CLEAN) * ref_noise | |
| anchor_ref = anchor.to(self.device, torch.float32) | |
| anchor_generator = torch.Generator(device="cpu").manual_seed(self.config.seed + 3000 + index) | |
| anchor_noise = torch.randn(tuple(anchor_ref.shape), generator=anchor_generator, dtype=torch.float32).to(self.device) | |
| augmented_anchor = VISUAL_CONDITION_CLEAN * anchor_ref + (1.0 - VISUAL_CONDITION_CLEAN) * anchor_noise | |
| retake = self._load_input_video() + (None, None) if index == 0 else self._retake_inputs(self._load_recovery(index - 1)) | |
| with torch.inference_mode(): | |
| sampled = sample_joint_av( | |
| self.transformer, | |
| layout=runtime.layout, | |
| text_hidden_states=runtime.text_hidden_states.to(self.device, torch.bfloat16), | |
| text_token_tags=runtime.text_token_tags.to(self.device), | |
| initial_video=initial_video, | |
| initial_audio=initial_audio, | |
| steps=self.config.steps, | |
| video_shift=VIDEO_SHIFT, | |
| audio_shift=AUDIO_SHIFT, | |
| visual_condition_latents=(augmented_anchor, augmented_ref), | |
| audio_condition_latents=(), | |
| visual_condition_clean=VISUAL_CONDITION_CLEAN, | |
| audio_condition_clean=AUDIO_CONDITION_CLEAN, | |
| input_video_latents=retake[0], | |
| video_denoise_mask=retake[1], | |
| input_audio_latents=retake[2], | |
| audio_denoise_mask=retake[3], | |
| ) | |
| video = sampled.video[0].float().cpu().contiguous() | |
| audio = sampled.audio[0].float().cpu().contiguous() | |
| if not torch.isfinite(video).all().item() or not torch.isfinite(audio).all().item(): | |
| raise RuntimeError(f"window {index} produced non-finite latents") | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = output.with_name(f".{output.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp") | |
| try: | |
| save_file( | |
| {"video": video, "audio": audio}, | |
| str(temporary), | |
| metadata={ | |
| "schema": RECOVERY_SCHEMA, | |
| "window_index": str(index), | |
| "seed": str(window_seed), | |
| "steps": str(self.config.steps), | |
| }, | |
| ) | |
| os.replace(temporary, output) | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| del sampled, video, audio, anchor, reference, hidden, tags, runtime | |
| del initial_video, initial_audio, ref, ref_noise, augmented_ref | |
| del anchor_ref, anchor_noise, augmented_anchor, retake | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| return output | |
| def write_video(self) -> dict[str, Any]: | |
| import av | |
| import torch | |
| output = self.config.output | |
| if output.exists(): | |
| raise FileExistsError(f"refusing to overwrite output: {output}") | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = output.with_name(f".{output.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp.mp4") | |
| frame_count = 0 | |
| try: | |
| with av.open(str(temporary), mode="w") as container: | |
| stream = container.add_stream("libx264", rate=FPS) | |
| stream.width = TARGET_WIDTH | |
| stream.height = TARGET_HEIGHT | |
| stream.pix_fmt = "yuv420p" | |
| for index in range(len(self.config.windows)): | |
| video, _audio = self._load_recovery(index) | |
| torch.cuda.empty_cache() | |
| with torch.inference_mode(): | |
| decoded = self.video_vae.decode( | |
| video.unsqueeze(0).to(self.device, torch.float16) | |
| )[0, :, :WINDOW_FRAMES].cpu() | |
| pixels = ( | |
| ((decoded.float().clamp(-1, 1) + 1.0) * 127.5) | |
| .round() | |
| .to(torch.uint8) | |
| .permute(1, 2, 3, 0) | |
| .contiguous() | |
| ) | |
| if index: | |
| pixels = pixels[OVERLAP_FRAMES:] | |
| for value in pixels: | |
| frame = av.VideoFrame.from_ndarray(value.numpy(), format="rgb24") | |
| for packet in stream.encode(frame): | |
| container.mux(packet) | |
| frame_count += 1 | |
| del video, decoded, pixels | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| for packet in stream.encode(): | |
| container.mux(packet) | |
| if frame_count != self.config.output_frames: | |
| raise RuntimeError(f"stitched frame count mismatch: {frame_count}") | |
| os.replace(temporary, output) | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| receipt = { | |
| "schema": RESULT_SCHEMA, | |
| "solver": "explicit_euler", | |
| "steps": self.config.steps, | |
| "seed": self.config.seed, | |
| "window_frames": WINDOW_FRAMES, | |
| "overlap_frames": OVERLAP_FRAMES, | |
| "output_frames": frame_count, | |
| "fps": FPS, | |
| "resolution": [TARGET_WIDTH, TARGET_HEIGHT], | |
| "audio": False, | |
| "output": str(output), | |
| "bytes": output.stat().st_size, | |
| "lora_gate": self.lora_gate, | |
| } | |
| atomic_json(output.with_suffix(output.suffix + ".json"), receipt) | |
| return receipt | |
| def generate(config: RunConfig) -> dict[str, Any]: | |
| validate_cache(config) | |
| started = time.monotonic() | |
| engine = Ref2VAEngine(config) | |
| for index in range(len(config.windows)): | |
| engine.sample_window(index) | |
| result = engine.write_video() | |
| result["elapsed_seconds"] = time.monotonic() - started | |
| return result | |