Add resident FL2VA Turbo variants
This commit is contained in:
Daniel Maddern
parent
6d8c9ca4cf
commit
837c16871d
10 changed files with 440 additions and 35 deletions
|
|
@ -60,6 +60,10 @@ docker compose -f compose.spark.yml up -d h3-hot-runtime
|
|||
|
||||
Use `GET /ready` to confirm resident model readiness. Use `POST /generate` with JSON fields like `prompt`, `output`, `width`, `height`, `frames`, `steps`, `seed`, and optional `attention`. Supported request-level attention values are reported by `/ready`; switching attention does not reload model weights. The hot image includes Sage2, forced cuDNN SDPA, and Comfy Kitchen INT8 attention. Sage2 is the default based on the 960x544x124 GB10 benchmark and the existing parity baseline.
|
||||
|
||||
The Spark hot service also keeps the official FL2VA Turbo adapters resident. Set `turbo` to `"4step"` for v1.1 768p (shift 6/3) or `"8step"` for v1.0 (shift 12/3). The matching step count is selected by default and enforced when `steps` is supplied. Set `turbo` to `null` or `"none"` for the base beta/RES path. Turbo uses its separate uniform training-Euler schedule and cannot be combined with denoiser caching.
|
||||
|
||||
See [`TURBO.md`](TURBO.md) for artifact hashes, implementation details, API examples, validation evidence, and matched GB10 performance results.
|
||||
|
||||
Exact memory/lifetime options:
|
||||
|
||||
- `attention: "kj_head_sliced"` slices attention heads and runs the slice backend from `H3_HEAD_SLICE_BACKEND` (`sage2` by default) with `H3_HEAD_SLICE_SIZE` heads per slice (`8` by default).
|
||||
|
|
|
|||
94
TURBO.md
Normal file
94
TURBO.md
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
# FL2VA Turbo
|
||||
|
||||
The hot runtime can keep both official LightX2V FL2VA Turbo adapters resident and switch them per request without rebuilding the packed NVFP4 base model.
|
||||
|
||||
## Variants
|
||||
|
||||
| API value | Official artifact | NFE | Video/audio shift | Schedule |
|
||||
| --- | --- | ---: | --- | --- |
|
||||
| `4step` | `minimax_h3_fl2v_turbo_4step_v1.1_768p_comfyui_bf16.safetensors` | 4 | 6/3 | Uniform shifted training-Euler |
|
||||
| `8step` | `minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors` | 8 | 12/3 | Uniform shifted training-Euler |
|
||||
|
||||
Downloaded from Hugging Face revision `ec01fa4c86263832faa0bd1d6d8f36a281eaabb2`:
|
||||
|
||||
| Artifact | SHA-256 |
|
||||
| --- | --- |
|
||||
| 4-step v1.1 Comfy BF16 | `449d80f301ac571622c72e28b8fd72a4b3681b7a8df8a92f17c8f6ec43f56558` |
|
||||
| 8-step v1.0 Comfy BF16 | `2339acdf19bfe123f46b971ea35d367a84adb85de43627e1eceafa5a5b2b111e` |
|
||||
|
||||
The older unversioned 744 MB adapters remain on Spark but are not loaded by the hot service.
|
||||
|
||||
## Implementation
|
||||
|
||||
Each official file contains 624 tensors covering 208 linear targets: four projections in each of 50 denoiser blocks and two token-refiner blocks. The runtime validates the complete key set, dimensions, rank, and alpha while loading.
|
||||
|
||||
The adapters are applied dynamically:
|
||||
|
||||
```text
|
||||
output = NVFP4_base(x) + strength * (alpha / rank) * B(A(x))
|
||||
```
|
||||
|
||||
The BF16 branch receives the original activation. It does not receive the base checkpoint's `pre_quant_scale`-modified activation, and the LoRA delta is never merged into packed NVFP4 weights. Fused QKV uses rank/alpha 384; the other released projections use rank/alpha 128.
|
||||
|
||||
Turbo sampling maintains independent native video and audio states. The 4-step and 8-step variants use their own shifted sigma grids and positive data-ward training-Euler updates. The original beta/RES sampler remains unchanged when Turbo is disabled.
|
||||
|
||||
## Hot API
|
||||
|
||||
The compose service registers both adapters with repeated startup arguments:
|
||||
|
||||
```text
|
||||
--turbo-lora 4step=/turbo/minimax_h3_fl2v_turbo_4step_v1.1_768p_comfyui_bf16.safetensors
|
||||
--turbo-lora 8step=/turbo/minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors
|
||||
```
|
||||
|
||||
Select one per request:
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": "A calm restaurant at dusk.",
|
||||
"output": "/output/example.mp4",
|
||||
"width": 960,
|
||||
"height": 544,
|
||||
"frames": 124,
|
||||
"seed": 440408,
|
||||
"turbo": "4step"
|
||||
}
|
||||
```
|
||||
|
||||
The server chooses and enforces four or eight steps. Use `null` or `"none"` for the base path. Turbo and denoiser caching cannot be combined. `/ready` reports `available_turbos` and `current_turbo`.
|
||||
|
||||
## Matched GB10 Benchmark
|
||||
|
||||
All runs used the same Posly first frame, prompt, seed `440408`, 960x544 canvas, 124 frames, Sage2, FP16 tiled VAE, and no audio output.
|
||||
|
||||
| Mode | Steps | Sampling | Total wall | Speedup vs base 8-step | Sampling reduction |
|
||||
| --- | ---: | ---: | ---: | ---: | ---: |
|
||||
| Base beta/RES | 8 | 148.1 s | 179.3 s | 1.00x | baseline |
|
||||
| Turbo v1.0 | 8 | 121.1 s | 152.3 s | 1.18x | 18.2% |
|
||||
| Turbo v1.1 768p | 4 | 60.4 s | 91.2 s | 1.97x | 59.2% |
|
||||
|
||||
The 4-step run was 1.67x faster end-to-end than the 8-step Turbo run. These are observed hot-service request times, not isolated steady-state kernel timings; shape-specific compilation can affect the first request.
|
||||
|
||||
## Validation
|
||||
|
||||
- 22 unit/contract tests cover disabled exactness, dynamic branch arithmetic, adapter selection, both official sigma grids, and independent AV training-Euler updates.
|
||||
- Both official files loaded with all 624 expected keys and all 208 target dimensions validated.
|
||||
- Both variants generated matched 960x544x124 FL2VA proofs.
|
||||
- Both variants completed joint AV smoke runs and produced stereo 32 kHz AAC muxed outputs.
|
||||
- Base warmup remains bit-preserving while no adapter is selected.
|
||||
|
||||
Comparison outputs:
|
||||
|
||||
```text
|
||||
\\192.168.1.162\StoryStudioAssets\H3-output\posly\proofs-960x540\01-posly-hero-turbo4-v1.1-seed440408-5s.mp4
|
||||
\\192.168.1.162\StoryStudioAssets\H3-output\posly\proofs-960x540\01-posly-hero-turbo8-v1.0-seed440408-5s.mp4
|
||||
\\192.168.1.162\StoryStudioAssets\H3-output\posly\proofs-960x540\01-posly-hero-turbo4-left-turbo8-right-seed440408.mp4
|
||||
```
|
||||
|
||||
## Quality Notes
|
||||
|
||||
For the Posly restaurant source, both Turbo variants avoided the severe central lens flare produced by the base-model trials. The 4-step result appeared slightly more compositionally stable; the 8-step result showed more subject motion. Both still brightened the exterior over time, so Turbo does not by itself satisfy the fixed-exposure requirement.
|
||||
|
||||
The adapters were released against the BF16 H3 base. Dynamic use over the pruned NVFP4 checkpoint is operationally validated here but is not an upstream-certified pairing. A future strict parity gate should compare per-step denoiser outputs and final latents against the latest LightX2V BF16 reference.
|
||||
|
||||
SLA Turbo is intentionally excluded. It requires LightX2V's separate 85% dynamic sparse-attention path and should not be represented as equivalent to these dense variants.
|
||||
|
|
@ -10,6 +10,7 @@ services:
|
|||
- /home/daniel/aeon-spark-test/h3/comfy-models/text_encoders:/text-encoders:ro
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
|
||||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/aeon-spark-test/h3/models/MiniMax-H3-Turbo-FL2VA/latest:/turbo:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
command: ["sleep", "infinity"]
|
||||
h3-hot-runtime:
|
||||
|
|
@ -23,6 +24,7 @@ services:
|
|||
- /home/daniel/aeon-spark-test/h3/comfy-models/text_encoders:/text-encoders:ro
|
||||
- /home/daniel/aeon-spark-test/h3/comfy-models/vae:/vae:ro
|
||||
- /home/daniel/aeon-spark-test/h3/h3-blackwell-runtime/artifacts:/artifacts:ro
|
||||
- /home/daniel/aeon-spark-test/h3/models/MiniMax-H3-Turbo-FL2VA/latest:/turbo:ro
|
||||
- /home/daniel/StoryStudioAssets/H3-output:/output
|
||||
ports:
|
||||
- "8001:8000"
|
||||
|
|
@ -31,4 +33,4 @@ services:
|
|||
H3_NVFP4_SCALE_BACKEND: "vortex"
|
||||
H3_NVFP4_SCALE_VERSION: "1"
|
||||
H3_SOL_QKV_LAYOUT: "native"
|
||||
command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--warmup"]
|
||||
command: ["python", "/opt/h3-blackwell-runtime/tools/serve_hot_runtime.py", "--host", "0.0.0.0", "--port", "8000", "--attention", "sage2", "--turbo-lora", "4step=/turbo/minimax_h3_fl2v_turbo_4step_v1.1_768p_comfyui_bf16.safetensors", "--turbo-lora", "8step=/turbo/minimax_h3_fl2v_turbo_8step_v1.0_comfyui_bf16.safetensors", "--warmup"]
|
||||
|
|
|
|||
95
src/h3_blackwell_runtime/lora.py
Normal file
95
src/h3_blackwell_runtime/lora.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
"""Resident dynamic LoRA branches for the quantized H3 denoiser and refiner."""
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as functional
|
||||
from safetensors import safe_open
|
||||
from torch import nn
|
||||
|
||||
|
||||
class LoraBranch(nn.Module):
|
||||
def __init__(self, down: torch.Tensor, up: torch.Tensor, alpha: float):
|
||||
super().__init__()
|
||||
if down.ndim != 2 or up.ndim != 2 or down.shape[0] != up.shape[1]:
|
||||
raise ValueError(f"Invalid LoRA shapes: down={tuple(down.shape)}, up={tuple(up.shape)}")
|
||||
self.scale = float(alpha) / down.shape[0]
|
||||
self.register_buffer("down", down.contiguous(), persistent=False)
|
||||
self.register_buffer("up", up.contiguous(), persistent=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = x.to(self.down.dtype)
|
||||
return functional.linear(functional.linear(x, self.down), self.up) * self.scale
|
||||
|
||||
|
||||
class DynamicLoraMixin:
|
||||
"""Small mixin for linears that can host resident request-selectable LoRAs."""
|
||||
|
||||
def _init_dynamic_lora(self) -> None:
|
||||
self.lora_branches = nn.ModuleDict()
|
||||
self.active_lora: str | None = None
|
||||
self.lora_strength = 0.0
|
||||
|
||||
def add_lora(self, name: str, down: torch.Tensor, up: torch.Tensor, alpha: float) -> None:
|
||||
if name in self.lora_branches:
|
||||
raise ValueError(f"LoRA {name!r} is already attached")
|
||||
if down.shape[1] != self.in_features or up.shape[0] != self.out_features:
|
||||
raise ValueError(
|
||||
f"LoRA {name!r} dimensions {tuple(down.shape)}, {tuple(up.shape)} do not match "
|
||||
f"linear [{self.out_features}, {self.in_features}]"
|
||||
)
|
||||
self.lora_branches[name] = LoraBranch(down, up, alpha)
|
||||
|
||||
def set_lora(self, name: str | None, strength: float = 1.0) -> None:
|
||||
if name is not None and name not in self.lora_branches:
|
||||
raise ValueError(f"LoRA {name!r} is not attached")
|
||||
self.active_lora = name
|
||||
self.lora_strength = float(strength) if name is not None else 0.0
|
||||
|
||||
def _apply_lora(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
|
||||
if self.active_lora is None or self.lora_strength == 0.0:
|
||||
return base
|
||||
delta = self.lora_branches[self.active_lora](x)
|
||||
return base + delta.to(base.dtype) * self.lora_strength
|
||||
|
||||
|
||||
def iter_lora_targets(model: nn.Module, refiner: nn.Module) -> Iterator[tuple[str, DynamicLoraMixin]]:
|
||||
for index, block in enumerate(model.backbone.blocks):
|
||||
yield f"blocks.{index}.attn.qkv_proj", block.attention.qkv_proj
|
||||
yield f"blocks.{index}.attn.out_proj", block.attention.out_proj
|
||||
yield f"blocks.{index}.mlp.fc1", block.mlp.fc1
|
||||
yield f"blocks.{index}.mlp.fc2", block.mlp.fc2
|
||||
for index, block in enumerate(refiner.blocks):
|
||||
yield f"token_refiner.blocks.{index}.attn.qkv_proj", block.qkv
|
||||
yield f"token_refiner.blocks.{index}.attn.out_proj", block.out
|
||||
yield f"token_refiner.blocks.{index}.mlp.fc1", block.fc1
|
||||
yield f"token_refiner.blocks.{index}.mlp.fc2", block.fc2
|
||||
|
||||
|
||||
def load_lora_adapter(model: nn.Module, refiner: nn.Module, name: str, path: str, device: str) -> int:
|
||||
targets = list(iter_lora_targets(model, refiner))
|
||||
expected = {
|
||||
f"diffusion_model.{target}.{suffix}"
|
||||
for target, _module in targets
|
||||
for suffix in ("alpha", "lora_A.weight", "lora_B.weight")
|
||||
}
|
||||
with safe_open(path, framework="pt", device=device) as checkpoint:
|
||||
actual = set(checkpoint.keys())
|
||||
if actual != expected:
|
||||
missing = sorted(expected - actual)[:8]
|
||||
unexpected = sorted(actual - expected)[:8]
|
||||
raise ValueError(f"LoRA key mismatch: missing={missing}, unexpected={unexpected}")
|
||||
for target, module in targets:
|
||||
prefix = f"diffusion_model.{target}"
|
||||
module.add_lora(
|
||||
name,
|
||||
checkpoint.get_tensor(f"{prefix}.lora_A.weight"),
|
||||
checkpoint.get_tensor(f"{prefix}.lora_B.weight"),
|
||||
checkpoint.get_tensor(f"{prefix}.alpha").item(),
|
||||
)
|
||||
return len(targets)
|
||||
|
||||
|
||||
def set_active_lora(model: nn.Module, refiner: nn.Module, name: str | None, strength: float = 1.0) -> None:
|
||||
for _target, module in iter_lora_targets(model, refiner):
|
||||
module.set_lora(name, strength)
|
||||
|
|
@ -7,6 +7,8 @@ import torch
|
|||
import torch.nn.functional as functional
|
||||
from torch import nn
|
||||
|
||||
from .lora import DynamicLoraMixin
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Nvfp4LinearTensors:
|
||||
|
|
@ -28,7 +30,7 @@ def parse_quant_sidecar(sidecar: torch.Tensor) -> dict:
|
|||
return metadata
|
||||
|
||||
|
||||
class Nvfp4Linear(nn.Module):
|
||||
class Nvfp4Linear(DynamicLoraMixin, nn.Module):
|
||||
"""Execute a packed Comfy NVFP4 linear with Comfy Kitchen's CUDA 13 kernel."""
|
||||
|
||||
def __init__(self, tensors: Nvfp4LinearTensors, output_dtype=torch.bfloat16):
|
||||
|
|
@ -49,6 +51,7 @@ class Nvfp4Linear(nn.Module):
|
|||
self.register_buffer("weight_scale_2", tensors.weight_scale_2.to(torch.float32).contiguous(), persistent=False)
|
||||
self.register_buffer("bias", tensors.bias.contiguous() if tensors.bias is not None else None, persistent=False)
|
||||
self.register_buffer("pre_quant_scale", tensors.pre_quant_scale.contiguous() if tensors.pre_quant_scale is not None else None, persistent=False)
|
||||
self._init_dynamic_lora()
|
||||
|
||||
def _packed_weight(self):
|
||||
from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout
|
||||
|
|
@ -73,6 +76,7 @@ class Nvfp4Linear(nn.Module):
|
|||
from comfy_kitchen.tensor import QuantizedTensor
|
||||
|
||||
original_shape = x.shape[:-1]
|
||||
original_x = x
|
||||
flat_x = x.reshape(-1, self.in_features).contiguous()
|
||||
if self.pre_quant_scale is not None:
|
||||
flat_x = flat_x * self.pre_quant_scale.to(flat_x)
|
||||
|
|
@ -81,12 +85,13 @@ class Nvfp4Linear(nn.Module):
|
|||
if self.full_precision_matrix_mult:
|
||||
weight = packed_weight.dequantize().to(flat_x)
|
||||
output = functional.linear(flat_x, weight, bias)
|
||||
return output.reshape(*original_shape, self.out_features)
|
||||
if x.dtype == torch.float32:
|
||||
raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.")
|
||||
packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
|
||||
output = functional.linear(packed_x, packed_weight, bias)
|
||||
return output[:flat_x.shape[0], :self.out_features].reshape(*original_shape, self.out_features)
|
||||
else:
|
||||
if x.dtype == torch.float32:
|
||||
raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.")
|
||||
packed_x = QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")
|
||||
output = functional.linear(packed_x, packed_weight, bias)[:flat_x.shape[0], :self.out_features]
|
||||
base = output.reshape(*original_shape, self.out_features)
|
||||
return self._apply_lora(original_x, base)
|
||||
|
||||
|
||||
def load_nvfp4_linear(tensors: dict[str, torch.Tensor], prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
|
||||
|
|
|
|||
|
|
@ -15,16 +15,23 @@ from .attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND
|
|||
from .block import configure_mlp_chunking
|
||||
from .checkpoint import H3Checkpoint
|
||||
from .denoiser import H3PackedDenoiser
|
||||
from .lora import load_lora_adapter, set_active_lora
|
||||
from .packing import H3PromptPacker
|
||||
from .qwen3vl_text import Qwen3VLPromptConditioner
|
||||
from .qwen3vl_vision import build_fl2va_presentation, Qwen3VL32BVision
|
||||
from .sampler import sample_video_res_multistep
|
||||
from .sampler import sample_video_res_multistep, sample_video_turbo
|
||||
from .t2v import random_av_latents
|
||||
from .token_refiner import H3TokenRefiner
|
||||
from .vae_decoder import MiniMaxH3VideoVAE, dtype_from_name
|
||||
from .vae_encoder import MiniMaxH3VideoVAEEncoder
|
||||
|
||||
|
||||
TURBO_VARIANTS = {
|
||||
"4step": {"steps": 4, "video_shift": 6.0},
|
||||
"8step": {"steps": 8, "video_shift": 12.0},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RuntimeConfig:
|
||||
model_path: str = "/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"
|
||||
|
|
@ -37,6 +44,7 @@ class RuntimeConfig:
|
|||
vae_tile_size: int = 256
|
||||
mlp_chunks: int = 1
|
||||
mlp_chunk_threshold: int = 4096
|
||||
turbo_loras: tuple[tuple[str, str], ...] = ()
|
||||
device: str = "cuda"
|
||||
|
||||
|
||||
|
|
@ -62,6 +70,7 @@ class H3HotRuntime:
|
|||
def __init__(self, config: RuntimeConfig):
|
||||
self.config = config
|
||||
self.attention = config.attention
|
||||
self.turbo: str | None = None
|
||||
self.loaded_at = time.time()
|
||||
self.load_stages: list[dict] = []
|
||||
|
||||
|
|
@ -78,6 +87,15 @@ class H3HotRuntime:
|
|||
"token_refiner_loaded",
|
||||
lambda: H3TokenRefiner(self.checkpoint, attention_backend=_refiner_attention_backend(config.attention)).eval(),
|
||||
)
|
||||
self.turbo_loras = dict(config.turbo_loras)
|
||||
unknown_turbos = sorted(set(self.turbo_loras) - set(TURBO_VARIANTS))
|
||||
if unknown_turbos:
|
||||
raise ValueError(f"Unknown Turbo variants: {unknown_turbos}")
|
||||
for name, path in self.turbo_loras.items():
|
||||
self._timed_load(
|
||||
f"turbo_{name}_loaded",
|
||||
lambda name=name, path=path: load_lora_adapter(self.model, self.refiner, name, path, config.device),
|
||||
)
|
||||
self.packer = H3PromptPacker(self.checkpoint)
|
||||
self.video_vae = self._timed_load(
|
||||
"video_vae_loaded",
|
||||
|
|
@ -123,6 +141,8 @@ class H3HotRuntime:
|
|||
"ready": True,
|
||||
"initial_attention": self.config.attention,
|
||||
"current_attention": self.attention,
|
||||
"available_turbos": sorted(self.turbo_loras),
|
||||
"current_turbo": self.turbo,
|
||||
"vae_dtype": self.config.vae_dtype,
|
||||
"vae_tile_size": self.config.vae_tile_size,
|
||||
"mlp_chunks": self.config.mlp_chunks,
|
||||
|
|
@ -144,6 +164,16 @@ class H3HotRuntime:
|
|||
block.attention_backend = _refiner_attention_backend(attention)
|
||||
self.attention = attention
|
||||
|
||||
@torch.inference_mode()
|
||||
def set_turbo(self, turbo: str | None) -> None:
|
||||
turbo = None if turbo in {None, "", "none"} else turbo
|
||||
if turbo is not None and turbo not in self.turbo_loras:
|
||||
raise ValueError(f"Turbo variant {turbo!r} is not loaded")
|
||||
if turbo == self.turbo:
|
||||
return
|
||||
set_active_lora(self.model, self.refiner, turbo)
|
||||
self.turbo = turbo
|
||||
|
||||
def _build_fl2va(self, prompt: str, first_frame: torch.Tensor | None, last_frame: torch.Tensor | None, width: int, height: int, aligned_frames: int):
|
||||
from .qwen3vl_vision import build_fl2va_presentation, resize_keyframe
|
||||
|
||||
|
|
@ -193,6 +223,7 @@ class H3HotRuntime:
|
|||
steps: int,
|
||||
seed: int,
|
||||
attention: str | None = None,
|
||||
turbo: str | None = None,
|
||||
first_frame: torch.Tensor | None = None,
|
||||
last_frame: torch.Tensor | None = None,
|
||||
mux_audio: bool = True,
|
||||
|
|
@ -217,8 +248,18 @@ class H3HotRuntime:
|
|||
|
||||
output = Path(output)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
requested_turbo = None if turbo in {None, "", "none"} else turbo
|
||||
if requested_turbo is not None:
|
||||
if requested_turbo not in self.turbo_loras:
|
||||
raise ValueError(f"Turbo variant {requested_turbo!r} is not loaded")
|
||||
expected_steps = TURBO_VARIANTS[requested_turbo]["steps"]
|
||||
if steps != expected_steps:
|
||||
raise ValueError(f"Turbo {requested_turbo} requires exactly {expected_steps} steps")
|
||||
if cache_mode not in {None, "", "disabled", "none"}:
|
||||
raise ValueError("Turbo sampling does not support denoiser caching")
|
||||
if attention is not None:
|
||||
self.set_attention(attention)
|
||||
self.set_turbo(requested_turbo)
|
||||
|
||||
video, audio, aligned_frames = timed(
|
||||
"latents_initialized",
|
||||
|
|
@ -238,31 +279,24 @@ class H3HotRuntime:
|
|||
"cond_latents": cond_latents,
|
||||
"cond_frame_indices": [kf["resolved_frame_index"] for kf in presentation.keyframes],
|
||||
"frame_count": frame_count,
|
||||
"seed": seed,
|
||||
}
|
||||
else:
|
||||
text = timed("text_conditioned", lambda: self.refiner(self.conditioner(prompt)))
|
||||
pack_kwargs = {}
|
||||
sampled = timed(
|
||||
"sampled",
|
||||
lambda: sample_video_res_multistep(
|
||||
self.model,
|
||||
self.packer,
|
||||
text,
|
||||
video,
|
||||
audio,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
return_audio=mux_audio,
|
||||
cache_mode=cache_mode,
|
||||
cache_threshold=cache_threshold,
|
||||
cache_start_percent=cache_start_percent,
|
||||
cache_end_percent=cache_end_percent,
|
||||
cache_subsample_factor=cache_subsample_factor,
|
||||
cache_stats=cache_stats,
|
||||
**pack_kwargs,
|
||||
),
|
||||
)
|
||||
if self.turbo is None:
|
||||
sample = lambda: sample_video_res_multistep(
|
||||
self.model, self.packer, text, video, audio, steps=steps, seed=seed,
|
||||
return_audio=mux_audio, cache_mode=cache_mode, cache_threshold=cache_threshold,
|
||||
cache_start_percent=cache_start_percent, cache_end_percent=cache_end_percent,
|
||||
cache_subsample_factor=cache_subsample_factor, cache_stats=cache_stats, **pack_kwargs,
|
||||
)
|
||||
else:
|
||||
sample = lambda: sample_video_turbo(
|
||||
self.model, self.packer, text, video, audio, steps=steps,
|
||||
video_shift=TURBO_VARIANTS[self.turbo]["video_shift"], seed=seed,
|
||||
return_audio=mux_audio, **pack_kwargs,
|
||||
)
|
||||
sampled = timed("sampled", sample)
|
||||
if mux_audio:
|
||||
latent, audio_latent = sampled
|
||||
else:
|
||||
|
|
@ -327,6 +361,7 @@ class H3HotRuntime:
|
|||
"height": height,
|
||||
"seed": seed,
|
||||
"attention": self.attention,
|
||||
"turbo": self.turbo,
|
||||
"vae_dtype": self.config.vae_dtype,
|
||||
"vae_tile_size": self.config.vae_tile_size,
|
||||
"stages": stages,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,14 @@ def beta_sigmas(steps: int, *, device: torch.device | str, alpha: float = 0.6, b
|
|||
return torch.FloatTensor([*sigmas, 0.0]).to(device)
|
||||
|
||||
|
||||
def turbo_sigmas(steps: int, shift: float, *, device: torch.device | str) -> torch.Tensor:
|
||||
"""Uniform shifted sigma grid used to train the distilled Turbo adapters."""
|
||||
if steps < 1 or shift <= 0:
|
||||
raise ValueError("Turbo steps and shift must be positive")
|
||||
base = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float32, device="cpu")
|
||||
return shifted_sigma(base, shift).to(device)
|
||||
|
||||
|
||||
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:
|
||||
|
|
@ -192,6 +200,51 @@ def sample_video_res_multistep(
|
|||
return (video, _decode_audio_latent(audio_carried)) if return_audio else video
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def sample_video_turbo(
|
||||
model,
|
||||
packer: H3PromptPacker,
|
||||
text: torch.Tensor,
|
||||
video: torch.Tensor,
|
||||
audio: torch.Tensor,
|
||||
*,
|
||||
steps: int,
|
||||
video_shift: float,
|
||||
audio_shift: float = 3.0,
|
||||
return_audio: bool = False,
|
||||
seed: int = 0,
|
||||
text_token_tags: torch.Tensor | None = None,
|
||||
cond_latents: list[torch.Tensor] | None = None,
|
||||
cond_frame_indices: list[int] | None = None,
|
||||
frame_count: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run the distilled H3 training-Euler schedule on independent AV clocks."""
|
||||
video_schedule = turbo_sigmas(steps, video_shift, device=video.device)
|
||||
audio_schedule = turbo_sigmas(steps, audio_shift, device=audio.device)
|
||||
for index in range(steps):
|
||||
video_sigma = video_schedule[index]
|
||||
audio_sigma = audio_schedule[index]
|
||||
model_timesteps = torch.stack((1.0 - video_sigma, 1.0 - audio_sigma))
|
||||
hidden, times, segments, positions, video_segment, audio_segment = packer(
|
||||
text,
|
||||
video,
|
||||
audio,
|
||||
video_sigma,
|
||||
model_timesteps,
|
||||
text_token_tags=text_token_tags,
|
||||
cond_latents=cond_latents,
|
||||
cond_frame_indices=cond_frame_indices,
|
||||
frame_count=frame_count,
|
||||
seed=seed,
|
||||
)
|
||||
raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
|
||||
video_output = unpatchify_video(raw_video.to(torch.bfloat16).float(), video.shape[2], video.shape[-2], video.shape[-1])
|
||||
audio_output = _unpack_audio(raw_audio.to(torch.bfloat16)).float()
|
||||
video = video.float() + (video_sigma - video_schedule[index + 1]) * video_output
|
||||
audio = audio.float() + (audio_sigma - audio_schedule[index + 1]) * audio_output
|
||||
return (video, audio) 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."""
|
||||
|
|
|
|||
|
|
@ -6,15 +6,19 @@ from torch import nn
|
|||
|
||||
from .attention import rms_norm, run_attention
|
||||
from .checkpoint import H3Checkpoint
|
||||
from .lora import DynamicLoraMixin
|
||||
|
||||
|
||||
class _Linear(nn.Module):
|
||||
class _Linear(DynamicLoraMixin, nn.Module):
|
||||
def __init__(self, checkpoint: H3Checkpoint, prefix: str, dtype: torch.dtype):
|
||||
super().__init__()
|
||||
self.register_buffer("weight", checkpoint.tensor(f"{prefix}.weight", dtype=dtype), persistent=False)
|
||||
self.in_features = self.weight.shape[1]
|
||||
self.out_features = self.weight.shape[0]
|
||||
self._init_dynamic_lora()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return functional.linear(x, self.weight)
|
||||
return self._apply_lora(x, functional.linear(x, self.weight))
|
||||
|
||||
|
||||
class _RefinerBlock(nn.Module):
|
||||
|
|
|
|||
92
tests/test_turbo.py
Normal file
92
tests/test_turbo.py
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from h3_blackwell_runtime.lora import DynamicLoraMixin
|
||||
from h3_blackwell_runtime.sampler import sample_video_turbo, turbo_sigmas
|
||||
|
||||
|
||||
class _Linear(DynamicLoraMixin, nn.Module):
|
||||
def __init__(self, weight):
|
||||
super().__init__()
|
||||
self.register_buffer("weight", weight)
|
||||
self.in_features = weight.shape[1]
|
||||
self.out_features = weight.shape[0]
|
||||
self._init_dynamic_lora()
|
||||
|
||||
def forward(self, x):
|
||||
return self._apply_lora(x, torch.nn.functional.linear(x, self.weight))
|
||||
|
||||
|
||||
class TurboLoraContracts(unittest.TestCase):
|
||||
def test_dynamic_lora_matches_unmerged_formula(self):
|
||||
base_weight = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
||||
down = torch.tensor([[2.0, -1.0]])
|
||||
up = torch.tensor([[3.0], [-2.0]])
|
||||
x = torch.tensor([[0.5, -1.0]])
|
||||
linear = _Linear(base_weight)
|
||||
linear.add_lora("turbo", down, up, alpha=0.5)
|
||||
linear.set_lora("turbo", strength=0.75)
|
||||
expected = torch.nn.functional.linear(x, base_weight)
|
||||
expected += 0.75 * 0.5 * torch.nn.functional.linear(torch.nn.functional.linear(x, down), up)
|
||||
torch.testing.assert_close(linear(x), expected)
|
||||
|
||||
def test_disabled_lora_returns_exact_base_result(self):
|
||||
weight = torch.randn(3, 2)
|
||||
x = torch.randn(4, 2)
|
||||
linear = _Linear(weight)
|
||||
linear.add_lora("turbo", torch.randn(1, 2), torch.randn(3, 1), alpha=1.0)
|
||||
expected = torch.nn.functional.linear(x, weight)
|
||||
self.assertTrue(torch.equal(linear(x), expected))
|
||||
linear.set_lora("turbo", strength=0.0)
|
||||
self.assertTrue(torch.equal(linear(x), expected))
|
||||
|
||||
def test_unknown_lora_is_rejected(self):
|
||||
linear = _Linear(torch.randn(3, 2))
|
||||
with self.assertRaisesRegex(ValueError, "not attached"):
|
||||
linear.set_lora("missing")
|
||||
|
||||
def test_four_step_shift_six_schedule(self):
|
||||
actual = turbo_sigmas(4, 6.0, device="cpu")
|
||||
expected = torch.tensor([1.0, 0.9473684211, 0.8571428571, 0.6666666667, 0.0])
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
def test_eight_step_shift_twelve_schedule_has_nine_grid_points(self):
|
||||
actual = turbo_sigmas(8, 12.0, device="cpu")
|
||||
self.assertEqual(actual.shape, (9,))
|
||||
self.assertEqual(actual[0].item(), 1.0)
|
||||
self.assertEqual(actual[-1].item(), 0.0)
|
||||
self.assertTrue(bool(torch.all(actual[:-1] > actual[1:])))
|
||||
|
||||
def test_turbo_sampler_uses_positive_dataward_updates_and_independent_clocks(self):
|
||||
video = torch.zeros(1, 1, 1, 1, 1)
|
||||
audio = torch.zeros(1, 1, 1, 1)
|
||||
calls = []
|
||||
|
||||
def packer(text, current_video, current_audio, sigma, model_timesteps, **kwargs):
|
||||
calls.append((current_video.clone(), current_audio.clone(), model_timesteps.clone()))
|
||||
return (None, None, None, None, None, None)
|
||||
|
||||
def model(*args):
|
||||
return torch.ones(1), torch.ones(1)
|
||||
|
||||
with (
|
||||
patch("h3_blackwell_runtime.sampler.unpatchify_video", return_value=torch.full_like(video, 2.0)),
|
||||
patch("h3_blackwell_runtime.sampler._unpack_audio", return_value=torch.full_like(audio, 3.0)),
|
||||
):
|
||||
result_video, result_audio = sample_video_turbo(
|
||||
model, packer, torch.empty(0), video, audio,
|
||||
steps=2, video_shift=6.0, audio_shift=3.0, return_audio=True,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(result_video, torch.full_like(video, 2.0))
|
||||
torch.testing.assert_close(result_audio, torch.full_like(audio, 3.0))
|
||||
torch.testing.assert_close(calls[1][0], torch.full_like(video, 2.0 / 7.0))
|
||||
torch.testing.assert_close(calls[1][1], torch.full_like(audio, 0.75))
|
||||
torch.testing.assert_close(calls[1][2], torch.tensor([1.0 / 7.0, 0.25]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -14,7 +14,7 @@ import torch
|
|||
from PIL import Image
|
||||
|
||||
from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND, attention_backend_status
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
|
||||
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig, TURBO_VARIANTS
|
||||
|
||||
|
||||
def _load_image(value) -> torch.Tensor | None:
|
||||
|
|
@ -56,11 +56,19 @@ parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), d
|
|||
parser.add_argument("--vae-tile-size", type=int, default=256)
|
||||
parser.add_argument("--mlp-chunks", type=int, default=1)
|
||||
parser.add_argument("--mlp-chunk-threshold", type=int, default=4096)
|
||||
parser.add_argument("--turbo-lora", action="append", default=[], metavar="NAME=PATH", help="Load a resident 4step or 8step Turbo adapter.")
|
||||
parser.add_argument("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.")
|
||||
parser.add_argument("--warmup-output", type=Path, default=Path("/output/h3-blackwell-runtime/hot-runtime-warmup.mp4"))
|
||||
args = parser.parse_args()
|
||||
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold))
|
||||
turbo_loras = []
|
||||
for value in args.turbo_lora:
|
||||
if "=" not in value:
|
||||
parser.error("--turbo-lora must use NAME=PATH")
|
||||
name, path = value.split("=", 1)
|
||||
turbo_loras.append((name, path))
|
||||
|
||||
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold, turbo_loras=tuple(turbo_loras)))
|
||||
runtime_lock = threading.Lock()
|
||||
warmup_result = None
|
||||
if args.warmup:
|
||||
|
|
@ -118,7 +126,16 @@ class Handler(BaseHTTPRequestHandler):
|
|||
width = int(payload.get("width", 960))
|
||||
height = int(payload.get("height", 544))
|
||||
frames = int(payload.get("frames", 124))
|
||||
steps = int(payload.get("steps", 12))
|
||||
turbo = payload.get("turbo")
|
||||
if turbo in {"", "none"}:
|
||||
turbo = None
|
||||
if turbo is not None and turbo not in runtime.turbo_loras:
|
||||
write_json(self, 400, {"error": "unsupported turbo", "turbo": turbo, "available": sorted(runtime.turbo_loras)})
|
||||
return
|
||||
steps = int(payload.get("steps", TURBO_VARIANTS[turbo]["steps"] if turbo else 12))
|
||||
if turbo is not None and steps != TURBO_VARIANTS[turbo]["steps"]:
|
||||
write_json(self, 400, {"error": "invalid turbo steps", "turbo": turbo, "required_steps": TURBO_VARIANTS[turbo]["steps"]})
|
||||
return
|
||||
seed = int(payload.get("seed", 440407))
|
||||
attention = payload.get("attention")
|
||||
if attention is not None and attention not in AVAILABLE_BACKENDS:
|
||||
|
|
@ -130,6 +147,9 @@ class Handler(BaseHTTPRequestHandler):
|
|||
last_frame = _load_image(payload.get("last_frame"))
|
||||
save_latent = payload.get("save_latent")
|
||||
cache_mode = payload.get("cache_mode")
|
||||
if turbo is not None and cache_mode not in {None, "", "disabled", "none"}:
|
||||
write_json(self, 400, {"error": "turbo does not support denoiser caching", "turbo": turbo})
|
||||
return
|
||||
cache_threshold = float(payload.get("cache_threshold", 0.0))
|
||||
cache_start_percent = float(payload.get("cache_start_percent", 0.0))
|
||||
cache_end_percent = float(payload.get("cache_end_percent", 1.0))
|
||||
|
|
@ -145,6 +165,7 @@ class Handler(BaseHTTPRequestHandler):
|
|||
steps=steps,
|
||||
seed=seed,
|
||||
attention=attention,
|
||||
turbo=turbo,
|
||||
first_frame=first_frame,
|
||||
last_frame=last_frame,
|
||||
mux_audio=mux_audio,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue