diff --git a/docs-site/content/cookbook/minimax-h3/open-h3-regenerate-2k.mdx b/docs-site/content/cookbook/minimax-h3/open-h3-regenerate-2k.mdx new file mode 100644 index 000000000..ad9ab233b --- /dev/null +++ b/docs-site/content/cookbook/minimax-h3/open-h3-regenerate-2k.mdx @@ -0,0 +1,7 @@ +--- +title: Open-H3-Regenerate-2K +description: Configure Open-H3-Regenerate-2K serving, which regenerates a 768P MiniMax-H3 draft at 2560x1440. +mode: frame +--- + + diff --git a/docs-site/islands/Cookbook.tsx b/docs-site/islands/Cookbook.tsx index 1bfba6ea3..a34d479a4 100644 --- a/docs-site/islands/Cookbook.tsx +++ b/docs-site/islands/Cookbook.tsx @@ -368,12 +368,16 @@ function Editor({ window.addEventListener("popstate", restore); return () => window.removeEventListener("popstate", restore); }, []); - const modes = mlx || catalog.runtime.interface === "websocket" ? ["serving"] : ["sdk", "inference", "serving"]; + const videoInput = catalog.workload === "v2v"; + const modes = + mlx || catalog.runtime.interface === "websocket" || videoInput ? ["serving"] : ["sdk", "inference", "serving"]; const modeNote = mlx ? "Apple Silicon runs through the native MLX server in this recipe. There is no CUDA CLI or Python SDK launch for this deployment." : catalog.runtime.interface === "websocket" ? "This deployment uses the native WebSocket server. Only serving is recorded for it." - : ""; + : videoInput + ? "This deployment needs a draft video. Send it to the server as video_reference; generated CLI and Python SDK configs cannot carry it yet." + : ""; const advancedGroups = groups.filter( (group) => group === "server" || group.includes("parallelism") || group.startsWith("streaming"), ); @@ -872,8 +876,9 @@ function modelModes(index: Index | undefined, model: DirectoryModel): string[] | function configurableModes(index: Index | undefined, modelId: string): string[] | undefined { const deployments = index?.models.find((model) => model.id === modelId)?.deployments; if (!deployments?.length) return undefined; - const servingOnly = (runtime: string) => runtime.includes("mlx") || runtime.includes("streaming"); - return deployments.some((deployment) => !servingOnly(deployment.runtime)) + const servingOnly = (deployment: { runtime: string; workload: string }) => + deployment.runtime.includes("mlx") || deployment.runtime.includes("streaming") || deployment.workload === "v2v"; + return deployments.some((deployment) => !servingOnly(deployment)) ? ["Python SDK", "CLI", "Serving"] : ["Serving"]; } diff --git a/docs-site/lib/cookbook-pages.ts b/docs-site/lib/cookbook-pages.ts index de5a2f916..af4e8e077 100644 --- a/docs-site/lib/cookbook-pages.ts +++ b/docs-site/lib/cookbook-pages.ts @@ -48,6 +48,15 @@ export const recipePages = [ tags: ["Reference to video", "Audio", "NVFP4"], featured: true, }, + { + id: "open-h3-regenerate-2k", + family: "MiniMax H3", + title: "Open-H3-Regenerate-2K", + mark: "H3", + path: "/cookbook/minimax-h3/open-h3-regenerate-2k", + description: "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with a three-step refiner.", + tags: ["Video upscale", "2K", "3-step"], + }, { id: "fastwan21", family: "Wan", @@ -125,6 +134,7 @@ const backendModels: Record = { "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2": "fasth3-preview", "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree": "fasth3-4step", "FastVideo/FastH3-OmniRef-NVFP4-QAD": "fasth3-omniref", + "Efficient-Large-Model/Open-H3-Regenerate-2K": "open-h3-regenerate-2k", "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": "fastwan21", "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": "wan21-i2v", "FastVideo/LTX2-Distilled-Diffusers": "ltx2-distilled", @@ -149,7 +159,7 @@ const TASK_GROUPS: Array<[RegExp, string]> = [ [/first\/last frame to video/, "Image to video"], [/text or image to video/, "Text/image to video"], [/image to video/, "Image to video"], - [/super-resolution/, "Video upscale"], + [/super-resolution|upscale/, "Video upscale"], [/video/, "Text to video"], ]; /** A small discovery vocabulary derived from the recipe task, not a capability claim. */ diff --git a/docs-site/navigation.json b/docs-site/navigation.json index ff98e5ba0..ab5a57185 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -170,6 +170,10 @@ "label": "Pika H3 Flex", "href": "/cookbook/minimax-h3/models/pika-h3-flex-t2va" }, + { + "label": "Open-H3-Regenerate-2K", + "href": "/cookbook/minimax-h3/open-h3-regenerate-2k" + }, { "label": "Recipes and setup", "href": "/cookbook/minimax-h3" diff --git a/docs-site/tests/cookbook.spec.ts b/docs-site/tests/cookbook.spec.ts index 523c20d31..154649a41 100644 --- a/docs-site/tests/cookbook.spec.ts +++ b/docs-site/tests/cookbook.spec.ts @@ -223,7 +223,7 @@ test("inherited values, explicit false, zero and null retain their meanings", as test("catalog discovery opens dedicated model recipes and filters by family and task", async ({ page }) => { await page.goto("/FastVideo/cookbook"); const cards = page.locator(".recipe-card"); - await expect(cards).toHaveCount(37); + await expect(cards).toHaveCount(38); await expect(cards.first()).toContainText("FastH3 Omni v1"); await page.getByLabel("Model family", { exact: true }).selectOption("Wan"); await page.getByLabel("Search recipes").fill("Wan2.1 I2V"); @@ -276,7 +276,7 @@ test("output tabs support keyboard navigation and precision chips use the backen await page.getByLabel("Search recipes").fill("does not exist"); await expect(page.getByRole("heading", { name: "No matching recipes" })).toBeVisible(); await page.getByRole("button", { name: "Clear filters" }).click(); - await expect(page.locator(".recipe-card")).toHaveCount(37); + await expect(page.locator(".recipe-card")).toHaveCount(38); }); test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => { @@ -751,7 +751,7 @@ test("MLX Omni discovery and its original route expose the MLX launch command", test("category counts match visible cards and exclude retired Preview", async ({ page }) => { await page.goto("/FastVideo/cookbook"); - await expect(page.locator(".recipe-card")).toHaveCount(37); + await expect(page.locator(".recipe-card")).toHaveCount(38); for (const category of ["All models", "Video", "Image", "Audio", "World"]) { const button = page.getByRole("button", { name: category, exact: true }); const count = Number(await button.locator(".recipe-filter-count").innerText()); @@ -819,3 +819,21 @@ test("WebSocket deployments expose only the serving interface", async ({ page }) await expect(modes.getByRole("button", { name: "CLI", exact: true })).toHaveCount(0); await expect(page.getByText(/Only serving is recorded for it/)).toBeVisible(); }); + +test("video-to-video deployments expose only the serving interface", async ({ page }) => { + await page.goto("/FastVideo/cookbook/minimax-h3/open-h3-regenerate-2k"); + const modes = page.getByRole("group", { name: "Run mode" }); + await expect(modes.getByRole("button")).toHaveCount(1); + await expect(modes.getByRole("button", { name: "Serving", exact: true })).toHaveAttribute("aria-pressed", "true"); + await expect(page.getByText(/Send it to the server as video_reference/)).toBeVisible(); +}); + +test("Open-H3-Regenerate-2K is discoverable as a video upscale with its serving page", async ({ page }) => { + await page.goto("/FastVideo/cookbook"); + await page.getByLabel("Task", { exact: true }).selectOption("Video upscale"); + const card = page.locator(".recipe-card").filter({ hasText: "Open-H3-Regenerate-2K" }); + await expect(card).toHaveCount(1); + await card.click(); + await expect(page).toHaveURL(/cookbook\/minimax-h3\/open-h3-regenerate-2k\/?$/); + await expect(page.getByRole("heading", { level: 1 })).toHaveText("Open-H3-Regenerate-2K"); +}); diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index ae6e0b726..26d614001 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -1083,6 +1083,22 @@ ], "related": ["pika-h3-flex-t2va", "pika-h3-flex-keyframes"] }, + { + "id": "open-h3-regenerate-2k", + "family": "minimax_h3", + "stage": "inference", + "task": "Video super-resolution (2K regeneration)", + "label": "Open-H3-Regenerate-2K", + "summary": "Regenerate a 768P MiniMax-H3 draft at 2560x1440 with NVlabs' three-step refiner. The output is silent.", + "model": "Efficient-Large-Model/Open-H3-Regenerate-2K", + "source": "examples/inference/basic/basic_open_h3_regenerate_2k.py", + "command": ": \"${INPUT_VIDEO:?Set INPUT_VIDEO to a 768P MiniMax-H3 draft}\"\npython examples/inference/basic/basic_open_h3_regenerate_2k.py --video \"$INPUT_VIDEO\" --prompt \"A red fox trots through fresh snow in a pine forest at golden hour.\"", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "2560x1440 MP4 at outputs/open_h3_regenerate_2k/regen2k_0.mp4", + "related": ["minimax-h3-t2v", "fasth3-8step-v2-cuda"] + }, { "id": "fasth3-lora-preview", "family": "minimax_h3", diff --git a/docs/cookbook/options.yaml b/docs/cookbook/options.yaml index 39007b94b..ce1d09b9f 100644 --- a/docs/cookbook/options.yaml +++ b/docs/cookbook/options.yaml @@ -44,7 +44,7 @@ runtimes: host: 0.0.0.0 port: 8000 served_model_name: null - workloads: [t2v, i2v] + workloads: [t2v, i2v, v2v] options: &cuda_options server.served_model_name: anyOf: diff --git a/docs/cookbook/recipes/open-h3-regenerate-2k.yaml b/docs/cookbook/recipes/open-h3-regenerate-2k.yaml new file mode 100644 index 000000000..b3f3aaf69 --- /dev/null +++ b/docs/cookbook/recipes/open-h3-regenerate-2k.yaml @@ -0,0 +1,18 @@ +title: Open-H3-Regenerate-2K +summary: Regenerate a 768P MiniMax-H3 draft at 2560x1440 with three refine steps. + +deployments: + cuda-rest: + runtime: fastvideo-cuda-rest + guide: docs/inference/open_h3_regenerate_2k.md + workload: v2v + defaults: examples/serving/openai_open_h3_regenerate_2k.yaml + requirements: + - 'UV_TORCH_BACKEND=cu130 uv pip install -e ".[fasth3]"' + - "One GPU with about 150 GB free (B200 or GB200). Send the draft as video_reference; local paths must sit under server.media_roots." + - "Optional NVFP4 decoder (Blackwell only): hf download FastVideo/FastH3-NVFP4-VAE-Decoder --include 'quality-36L/decoder-nvfp4.safetensors' --local-dir fasth3-nvfp4-vae" + options: + - server.* + - generator.engine.offload.text_encoder + - generator.pipeline.experimental.h3_vae_nvfp4_decoder + - default_request.sampling.seed diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 848fad33b..b46221869 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -153,6 +153,12 @@ surfaces: pipeline_config_extensions: preset_owned: + draft_size: + sources: + - fastvideo.configs.pipelines.open_h3_regenerate_2k.OpenH3Regenerate2KPipelineConfig + refine_start_sigma: + sources: + - fastvideo.configs.pipelines.open_h3_regenerate_2k.OpenH3Regenerate2KPipelineConfig audio_encoder_precision: sources: - fastvideo.configs.pipelines.wan.WanS2V14BConfig diff --git a/docs/inference/open_h3_regenerate_2k.md b/docs/inference/open_h3_regenerate_2k.md new file mode 100644 index 000000000..53da256e1 --- /dev/null +++ b/docs/inference/open_h3_regenerate_2k.md @@ -0,0 +1,135 @@ +# Open-H3-Regenerate-2K + +`OpenH3Regenerate2KPipeline` regenerates a 768P [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) draft video +(1344x768, 121 frames, 24 fps) at **2560x1440** with three refine steps. It is video to video: the draft is the input, +and the output is a silent 2K MP4. FastVideo ports NVlabs' +[Open-H3-Regenerate-2K](https://huggingface.co/Efficient-Large-Model/Open-H3-Regenerate-2K) +([project page](https://nvlabs.github.io/Sana/Open-H3-Regenerate-2K/), +[reference code](https://github.com/NVlabs/Sana/tree/sol-engine/models/Open-H3-Regenerate-2K)); the dense BF16 +`standard_fwd` path of that repo is the parity reference. The accelerated `sol_engine_fwd` path (Sol-Attn, MXFP8, CUDA +graphs) and the fast 24-layer VAE decoder are not ported. + +```text +draft video -> H3 VAE encode -> latent upscaler (2x, 3D conv) -> 3 Euler refine steps -> H3 VAE decode -> 2560x1440 + H3 DiT + merged LoRA + Sharp curve AdaLN branches +``` + +## What runs + +| Step | Details | +|---|---| +| Draft encode | `ffmpeg` resamples to 24 fps, cover-scales (bilinear) and centre-crops to 1344x768 (`draft_size: 768p`, or 1024x576 for `576p`), holds the last frame to 121 frames plus 3 copies. The FP32 H3 VAE posterior is sampled under seed 42. | +| Latent upscale | The author's 3D-conv upscaler (`fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py`), BF16, to the 160x90 latent grid. | +| Refine | Euler on sigmas `[0.909375, 0.857143, 0.666667, 0]`, CFG 1, model time `1 - sigma`, initial state `(1 - 0.909375) * upscaled + 0.909375 * noise`. `num_inference_steps: 2` uses `[0.7, 0.494118, 0]` (softer, about 30% faster). | +| Transformer | The base MiniMax-H3 DiT (the same module as the other H3 pipelines) with the released merged LoRA (FP32, rank 576 on 312 linears, scale 1) and 50 LMS Sharp curve AdaLN branches. | +| Decode | The released H3 VAE; the output is trimmed to 121 frames. | + +The merged LoRA loads through the standard H3 community-LoRA path as a constructor-time adapter: the pipeline sets +`lora_path` to `/lora/open_h3_regenerate_2k_lora.safetensors` when you do not pass one. The Sharp branches take an +8-D timestep curve (`lora/adaln_t_table.bin`, 1025x8) instead of the layer input, so they cannot be folded into weights; +`sharp_curve.py` adds them to the AdaLN projection outputs with forward hooks. + +**Audio.** The reference feeds an encoded-silence audio stream to the DiT as conditioning and writes the video with no +audio track (`-an`); it neither regenerates nor passes through the draft's audio. FastVideo does the same: the draft's +soundtrack is ignored and the output MP4 has no audio. + +## Weights + +The release repo keeps the unmodified MiniMax-H3 components (a copy of `MiniMaxAI/MiniMax-H3` at `6818f6c`) under +`base_model/` next to `lora/` and `latent_upsampler/`. Point `model_path` at the Hub id or at a local copy of that layout. +`adaln_t_table.bin` is base-derived and is not redistributed in the release repo: put the 32,800-byte file (SHA-256 +`c02a6c11888297688c1e6278185ea1f947023acfc69f9003bbcdcec9a229a8e7`) at `lora/adaln_t_table.bin`. The reference +`download.py` fetches it by byte range from `Comfy-Org/MiniMax-H3` at revision +`e5eb578a89295337b8ff433a035929ce0279e0b6` (tensor `adaln_t_table` of +`diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors`) and verifies the hash. + +A GB200/B200-class GPU is needed (about 150 GB with the DiT, the 2K activations and the VAE). + +## Usage + +```bash +python examples/inference/basic/basic_open_h3_regenerate_2k.py \ + --model-path Efficient-Large-Model/Open-H3-Regenerate-2K \ + --video draft_768p.mp4 --prompt "A red fox trots through fresh snow" --output outputs/regen2k +``` + +The prompt conditions the refine steps (use the draft's own prompt). Drafts from other video models are better encoded at +576P (`--draft-size 576p`). `--steps 2` runs the two-step variant. The output is always 2560x1440; the first 121 frames of +the draft are used. + +## Serving + +```bash +fastvideo serve --config examples/serving/openai_open_h3_regenerate_2k.yaml +``` + +`examples/serving/openai_open_h3_regenerate_2k.yaml` loads the release on one GPU, sets `media_roots`, and pins the 2K +defaults. The draft arrives as the typed `video_reference` of `POST /v1/videos/sync` (the response body is the MP4) or of +`POST /v1/videos` (an async job; poll `GET /v1/videos/{id}` and download `GET /v1/videos/{id}/content`): + +```bash +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H 'Content-Type: application/json' \ + -d '{ + "prompt": "A red fox trots through fresh snow", + "video_reference": {"video_url": "https://example.com/draft_768p.mp4"}, + "seed": 1766580932 + }' \ + -o regen_2k.mp4 -D - +``` + +The response headers carry `X-Request-Id`, `X-Model`, `X-Inference-Time-S` and `X-Stage-Durations`. Rules: + +- Exactly one `video_reference` (or the legacy `video_path` / `video_url`) and no image or audio references. +- `http(s)` URLs must resolve to public addresses unless `server.allow_private_media_urls` is set. A local path is accepted + only inside a directory listed in `server.media_roots`. `data:` URLs are refused for video references; this policy is + unchanged. +- `size`, `width` and `height` may be omitted or `2560x1440`; `num_frames` may be omitted or 124 (the frame count of the + VAE grid; the video has 121 frames); `fps` must be 24. `seconds` is not accepted. +- `num_inference_steps` is 3 (default) or 2; `seed` sets the refine noise. +- Serve with `attention_backend: TORCH_SDPA`: the refine stage selects cuDNN attention, which at about 133k tokens is + about 3x faster than PyTorch's default flash dispatch. The `FLASH_ATTN` backend bypasses that selection. + +## Parity with the reference + +Measured on one GB200 (BF16, `TORCH_SDPA` with cuDNN attention) against the reference `standard_fwd` with the released +VAE, same draft, seed (1766580932), prompt and VAE posterior seed (42). The reference's noise is a CUDA generator draw +that FastVideo reproduces (the same `torch.Generator(device="cuda")` seed), so no injected noise was needed. + +| Stage / quantity | Fox draft (FastVideo H3 t2v, 768P) | Dance draft (576P, project page) | +|---|---|---| +| Draft latent, upscaled latent, DiT input | bit-identical | bit-identical | +| One DiT forward (LoRA + Sharp), relative L2 | 1.1e-2 | 2.6e-2 | +| Refined latent, relative L2 | 1.3e-2 | 3.7e-2 | +| Decoded video PSNR / SSIM (21 frames) | 45.2 dB / 0.993 | 42.1 dB / 0.992 | + +The differences come from text-encoder rounding (relative L2 1e-3 on the prompt embeddings) and BF16 kernel choices; +the reference reports the same order (1.4% relative L2 on the refined latent between its own runs). + +Warm timings, one GB200 (seconds): + +| | Upscale | Refine (3 steps) | VAE decode | Total | +|---|---:|---:|---:|---:| +| Reference `standard_fwd` | 0.87 | 71.6 | 15.0 | 87.5 | +| FastVideo (end to end, incl. draft encode 5-10 s and decode to host) | 5.9-10.7 (encode + upscale + silence) | 69.5 | 25.0 | 107.5 | + +The first request builds nothing extra but loads weights (about 70 s once). The decode stage is slower than the +reference because it also copies 121 frames to host memory and the MP4 encode is outside both columns. + +## Licences + +The release is not covered by one licence; its `README.md` says so, and the runtime code's licence does not replace the +model licences. FastVideo only loads these weights and does not redistribute them. + +- The refiner LoRA, the fast H3 VAE decoder and `base_model/` (a copy of `MiniMaxAI/MiniMax-H3` at `6818f6c`: DiT, VAEs, + text encoder) follow the + [MiniMax H3 Community License](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/6818f6c32d12b210915e44ad56a4228c2608f160/LICENSE), + including its territory, use and redistribution terms (`base_model/LICENSE`, `NOTICE`). The Qwen3-VL text encoder's + Apache-2.0 text ships as `LICENSE-Qwen.txt`. +- The latent upscaler follows its own licence: + [LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler](https://github.com/LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler) + (MIT, per the reference code) and its [weights](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler) + (`latent_upsampler/LICENSE`, Apache-2.0). `minimax_h3_latent_upscaler.py` is a port of that node's network. +- The merged LoRA includes [LightX2V FL2VA Turbo](https://huggingface.co/lightx2v/Minimax-h3-Turbo) and + [LMS Sharp](https://huggingface.co/Alissonerdx/Minimax-H3-ComfyUI), which keep their own licences; the Sharp + `adaln_t_table` is read from [Comfy-Org/MiniMax-H3](https://huggingface.co/Comfy-Org/MiniMax-H3). diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 3aa8e43d2..e6b40458f 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -69,6 +69,7 @@ column links a runnable script in `examples/inference/basic/` where one exists. | matrixgame | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | I2V | [basic_matrixgame3.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_matrixgame3.py) | | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)
[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)
[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) | | pika_h3_flex | `Pika-Labs/Pika-H3-Flex` | T2V, I2V | [T2VA/I2VA/L2VA/FL2VA/Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pika_h3_flex.py) | +| open_h3_regenerate_2k | `Efficient-Large-Model/Open-H3-Regenerate-2K` | — | [basic_open_h3_regenerate_2k.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_open_h3_regenerate_2k.py) | | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) | | qwen_image21 | `Qwen/Qwen-Image-2.1` | T2I, I2I | [basic_qwen_image21.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_qwen_image21.py) | | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) | @@ -122,6 +123,9 @@ and distilled) generate video with audio from text, optionally plus an image (se `kandinsky6_sr` IDs upscale an existing video (see the [Video SR guide](kandinsky6_sr.md)). +**Note (Open-H3-Regenerate-2K)**: `open_h3_regenerate_2k` regenerates an existing 768P MiniMax-H3 draft video at +2560x1440 (video to video, silent output); see the [Open-H3-Regenerate-2K guide](open_h3_regenerate_2k.md). + **Note (LTX-2.5)**: T2V and I2V produce synchronized video and audio. The official `Lightricks/LTX-2.5` repository uses gated, split component files, so it must first be converted to FastVideo's component layout. See the diff --git a/docs/js/cookbook-workloads.mjs b/docs/js/cookbook-workloads.mjs index d2820a60a..243ee687c 100644 --- a/docs/js/cookbook-workloads.mjs +++ b/docs/js/cookbook-workloads.mjs @@ -32,6 +32,14 @@ export const WORKLOADS = Object.freeze({ extension: "wav", media_type: "audio/wav", }), + v2v: Object.freeze({ + input: "video", + output: "video", + endpoint: "/v1/videos/sync", + extension: "mp4", + media_type: "video/mp4", + reference: "video_reference", + }), v2a: Object.freeze({ input: "video", output: "audio", diff --git a/docs/tests/cookbook-config.test.mjs b/docs/tests/cookbook-config.test.mjs index 36c26869f..d95640fe5 100644 --- a/docs/tests/cookbook-config.test.mjs +++ b/docs/tests/cookbook-config.test.mjs @@ -488,6 +488,7 @@ test("sample requests follow shared workload media", () => { t2i: { endpoint: /\/v1\/images\/generations/, output: /--output output\.json/, reference: null }, t2a: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.wav/, reference: null }, v2a: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.wav/, reference: "video_reference" }, + v2v: { endpoint: /\/v1\/videos\/sync/, output: /--output output\.mp4/, reference: "video_reference" }, }; for (const [workload, expected] of Object.entries(expectations)) { const data = catalog(); diff --git a/examples/inference/basic/basic_open_h3_regenerate_2k.py b/examples/inference/basic/basic_open_h3_regenerate_2k.py new file mode 100644 index 000000000..cb2ff8c20 --- /dev/null +++ b/examples/inference/basic/basic_open_h3_regenerate_2k.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regenerate a 768P MiniMax-H3 draft video at 2560x1440 with Open-H3-Regenerate-2K. + +The draft is VAE-encoded (the first 121 frames at 24 fps, 1344x768), its latent is upscaled, and three refine steps +of the H3 transformer with the released LoRA run at 2560x1440. The output video is silent. See +docs/inference/open_h3_regenerate_2k.md. + + python examples/inference/basic/basic_open_h3_regenerate_2k.py --video draft.mp4 --prompt "..." --output out/ +""" + +from __future__ import annotations + +import argparse +import time +from pathlib import Path + +from fastvideo import VideoGenerator +from fastvideo.api import ( + ComponentConfig, + EngineConfig, + GenerationRequest, + GeneratorConfig, + InputConfig, + OffloadConfig, + OutputConfig, + PipelineSelection, + SamplingConfig, +) + +DEFAULT_MODEL = "Efficient-Large-Model/Open-H3-Regenerate-2K" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model-path", default=DEFAULT_MODEL, help="release repo id or a local copy of it") + parser.add_argument("--video", required=True, help="768P draft video (URL-free local path)") + parser.add_argument("--prompt", required=True, help="text the refine steps are conditioned on") + parser.add_argument("--output", type=Path, default=Path("outputs/open_h3_regenerate_2k")) + parser.add_argument("--steps", type=int, choices=(2, 3), default=3, help="refine steps (3 is the released setting)") + parser.add_argument("--seed", type=int, default=1766580932) + parser.add_argument("--attention-backend", default="TORCH_SDPA") + parser.add_argument("--draft-size", choices=("768p", "576p"), default="768p", + help="size the draft is encoded at; 576p for drafts from other video models") + parser.add_argument("--runs", type=int, default=1, help="repeat the request to measure warm timings") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + args.output.mkdir(parents=True, exist_ok=True) + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=args.model_path, + engine=EngineConfig(num_gpus=1, offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=False)), + pipeline=PipelineSelection( + components=ComponentConfig(), + experimental={ + "attention_backend": args.attention_backend, + "draft_size": args.draft_size, + }, + ), + )) + try: + for run in range(args.runs): + started = time.perf_counter() + result = generator.generate( + GenerationRequest( + prompt=args.prompt, + inputs=InputConfig(video_path=args.video), + sampling=SamplingConfig(height=1440, + width=2560, + num_frames=124, + fps=24, + num_inference_steps=args.steps, + guidance_scale=1.0, + seed=args.seed), + output=OutputConfig(output_path=str(args.output / f"regen2k_{run}.mp4"), + save_video=True, + return_frames=False), + )) + print(f"run {run}: {time.perf_counter() - started:.1f}s -> {result.video_path}") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/examples/serving/openai_open_h3_regenerate_2k.yaml b/examples/serving/openai_open_h3_regenerate_2k.yaml new file mode 100644 index 000000000..ffeafe536 --- /dev/null +++ b/examples/serving/openai_open_h3_regenerate_2k.yaml @@ -0,0 +1,43 @@ +# OpenAI-compatible Open-H3-Regenerate-2K server (one GPU with about 150 GB free, e.g. B200 / GB200). +# POST /v1/videos/sync (or /v1/videos) with a 768P H3 draft as video_reference; the response is the 2560x1440 MP4. +# Local draft paths are accepted only inside server.media_roots; http(s) URLs must resolve to public addresses unless +# server.allow_private_media_urls is set. data: URLs are always refused. See docs/inference/open_h3_regenerate_2k.md. +generator: + model_path: Efficient-Large-Model/Open-H3-Regenerate-2K + engine: + num_gpus: 1 + offload: + dit: false + dit_layerwise: false + text_encoder: true + vae: false + pin_cpu_memory: true + compile: + enabled: false + vae_enabled: false + pipeline: + experimental: + attention_backend: TORCH_SDPA + # 768p for MiniMax-H3 drafts; 576p encodes drafts from other video models at 1024x576. + draft_size: 768p + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_open_h3_regenerate_2k + served_model_name: open-h3-regenerate-2k + media_roots: + - /data/drafts + media_reference_max_bytes: 536870912 + +default_request: + sampling: + height: 1440 + width: 2560 + num_frames: 124 + fps: 24 + num_inference_steps: 3 + guidance_scale: 1.0 + seed: 1766580932 + output: + return_frames: false diff --git a/fastvideo/configs/pipelines/open_h3_regenerate_2k.py b/fastvideo/configs/pipelines/open_h3_regenerate_2k.py new file mode 100644 index 000000000..83dd12452 --- /dev/null +++ b/fastvideo/configs/pipelines/open_h3_regenerate_2k.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline configuration for Open-H3-Regenerate-2K (768P H3 draft -> 2560x1440).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + +DRAFT_SIZE_CHOICES = ("768p", "576p") + + +@dataclass +class OpenH3Regenerate2KPipelineConfig(MiniMaxH3PipelineConfig): + """MiniMax-H3 policy plus the refine-stage knobs. + + ``draft_size`` is the size the draft video is VAE-encoded at: 768p (1344x768) for MiniMax-H3 drafts, 576p + (1024x576) for drafts from other video models. ``refine_start_sigma`` re-noises the upscaled draft to another + level; the later steps then follow the shift-6 rule (None keeps the released grid). + """ + + draft_size: str = "768p" + refine_start_sigma: float | None = None + + def check_pipeline_config(self) -> None: + super().check_pipeline_config() + if self.draft_size not in DRAFT_SIZE_CHOICES: + raise ValueError(f"draft_size must be one of {DRAFT_SIZE_CHOICES}, got {self.draft_size!r}.") + if self.refine_start_sigma is not None and not 0.0 < self.refine_start_sigma < 1.0: + raise ValueError(f"refine_start_sigma must be in (0, 1), got {self.refine_start_sigma}.") diff --git a/fastvideo/entrypoints/openai/request_adapter.py b/fastvideo/entrypoints/openai/request_adapter.py index ed767e054..cd731e27e 100644 --- a/fastvideo/entrypoints/openai/request_adapter.py +++ b/fastvideo/entrypoints/openai/request_adapter.py @@ -41,6 +41,10 @@ class RequestAdaptationError(ValueError): """The transport request cannot be represented by the loaded pipeline.""" +OPEN_H3_REGENERATE_2K_FAMILY = "open_h3_regenerate_2k" +OPEN_H3_REGENERATE_2K_SIZE = (2560, 1440) +OPEN_H3_REGENERATE_2K_NUM_FRAMES = 124 + _VIDEO_SUFFIXES = frozenset({".avi", ".m4v", ".mkv", ".mov", ".mp4", ".mpeg", ".mpg", ".webm"}) @@ -381,6 +385,13 @@ def _apply_reference_inputs( if request.task is not None and model_family != "minimax_h3": raise RequestAdaptationError("The task selector is only defined for MiniMax-H3 servers.") + if model_family == OPEN_H3_REGENERATE_2K_FAMILY: + if len(videos) != 1 or images or audios: + raise RequestAdaptationError("Open-H3-Regenerate-2K takes exactly one video_reference (the 768P draft) " + "and no image or audio references.") + kwargs["video_path"] = videos[0] + return + if model_family == "minimax_h3" and (videos or audios): raise RequestAdaptationError("MiniMax-H3 video/audio references require a server configured with " "override_pipeline_cls_name=MiniMaxH3Ref2VAModularPipeline.") @@ -522,6 +533,16 @@ def build_generation_request( if height is not None and (not isinstance(height, int) or height <= 0): raise RequestAdaptationError(f"height must be a positive integer, got {height!r}") + if model_family == OPEN_H3_REGENERATE_2K_FAMILY: + if (width, height) not in ((None, None), OPEN_H3_REGENERATE_2K_SIZE): + raise RequestAdaptationError("Open-H3-Regenerate-2K always renders 2560x1440; omit size/width/height " + f"or pass 2560x1440, got {width}x{height}.") + if kwargs.get("num_frames", OPEN_H3_REGENERATE_2K_NUM_FRAMES) != OPEN_H3_REGENERATE_2K_NUM_FRAMES: + raise RequestAdaptationError("Open-H3-Regenerate-2K regenerates the draft's first 121 frames (24 fps); " + "omit seconds and num_frames.") + if kwargs["fps"] != 24: + raise RequestAdaptationError(f"Open-H3-Regenerate-2K requires fps=24, got {kwargs['fps']}.") + if model_family == "minimax_h3": from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_CANVAS_MULTIPLE, diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index a549c77f0..31b21d531 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -254,13 +254,7 @@ class BaseLayerWithLoRA(nn.Module): current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()).full_tensor() - # Apply LoRA with alpha scaling - lora_delta = ( - self.slice_lora_b_weights(self.lora_B).to(data) @ self.slice_lora_a_weights(self.lora_A).to(data)) - scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) - lora_delta *= scale * self.lora_strength - data += lora_delta - unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) + unsharded_base_layer.weight = nn.Parameter(self._merged_weight(data).to(current_device)) if isinstance(getattr(self.base_layer, "bias", None), DTensor): unsharded_base_layer.bias = nn.Parameter( self.base_layer.bias.to(get_local_torch_device(), @@ -277,16 +271,22 @@ class BaseLayerWithLoRA(nn.Module): current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()) - # Apply LoRA with alpha scaling - lora_delta = ( - self.slice_lora_b_weights(self.lora_B.to(data)) @ self.slice_lora_a_weights(self.lora_A.to(data))) - scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) - lora_delta *= scale * self.lora_strength - data += lora_delta - self.base_layer.weight.data = data.to(current_device, non_blocking=True) + self.base_layer.weight.data = self._merged_weight(data).to(current_device, non_blocking=True) self.merged = True + def _merged_weight(self, weight: torch.Tensor) -> torch.Tensor: + """``weight`` plus the scaled adapter delta, in ``weight``'s dtype. + + An adapter stored wider than the base (an FP32 adapter on a BF16 DiT) is multiplied and added at its own + precision and rounded once, as a merge-then-cast reference does; with FSDP or without. + """ + accumulate_dtype = torch.promote_types(weight.dtype, self.lora_A.dtype) + lora_B = self.slice_lora_b_weights(self.lora_B.to(device=weight.device, dtype=accumulate_dtype)) + lora_A = self.slice_lora_a_weights(self.lora_A.to(device=weight.device, dtype=accumulate_dtype)) + scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0) + return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * (scale * self.lora_strength)).to(weight.dtype) + @torch.no_grad() # @torch.compile(dynamic=True) def unmerge_lora_weights(self, *, if_merged_only: bool = True) -> None: diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 8cea9dde7..9c9ffb90f 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1655,6 +1655,7 @@ class PipelineComponentLoader: component_model_path, ) + module_name = module_name.rsplit("/", 1)[-1] # Get the appropriate loader for this module type loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) diff --git a/fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py b/fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py new file mode 100644 index 000000000..9ed762431 --- /dev/null +++ b/fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-H3 3D-convolution latent upscaler (LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler, MIT). + +Port of the upstream node's ``LatentResizer3D`` and its ``target dimensions`` size logic, without the ComfyUI shims, +attention blocks (the node forces them off) and temporal chunking (Open-H3-Regenerate-2K runs without it). Module names +match the upstream ``latent_upsampler/*.safetensors`` checkpoint, so it loads without remapping. + +The network works on a latent that was normalized with the statistics below, not the VAE-normalized latent the DiT +sees: ``upscale_latent`` therefore re-normalizes its input and de-normalizes the result with the same statistics. +""" +from __future__ import annotations + +import re + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch import nn + +VAE_DOWNSAMPLE = 16 +PIXEL_ALIGNMENT = 32 + +LATENTS_MEAN = ( + 0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949, + -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736, + 0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, + 0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279, + -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264) +LATENTS_STD = ( + 1.2223774194717407, 1.2767263650894165, 1.6831774711608887, 1.7549455165863037, 1.5636216402053833, + 2.194143533706665, 0.9653137922286987, 1.0569885969161987, 0.841948926448822, 0.7729952931404114, + 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049927, 0.7197399735450745, + 0.6936293244361877, 2.961095094680786, 2.7694199085235596, 3.0496184825897217, 2.1088054180145264, + 3.276226282119751, 3.1627357006073, 2.2816812992095947, 2.6127843856811523) + + +def _norm(channels: int) -> nn.GroupNorm: + return nn.GroupNorm(32, channels) + + +class ResBlockEmb3D(nn.Module): + + def __init__(self, channels: int, emb_channels: int) -> None: + super().__init__() + self.in_layers = nn.Sequential(_norm(channels), nn.SiLU(), nn.Conv3d(channels, channels, 3, padding=1)) + self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(emb_channels, 2 * channels)) + self.out_norm = _norm(channels) + self.out_layers = nn.Sequential(nn.SiLU(), nn.Dropout(p=0.0), nn.Conv3d(channels, channels, 3, padding=1)) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + h = self.in_layers(x) + scale, shift = self.emb_layers(emb).type(h.dtype)[..., None, None, None].chunk(2, dim=1) + h = self.out_norm(h) * (1 + scale) + shift + return x + self.out_layers(h) + + +class TemporalConv(nn.Module): + + def __init__(self, channels: int, kernel_size: int = 5) -> None: + super().__init__() + self.norm = _norm(channels) + self.dwconv = nn.Conv3d(channels, channels, (kernel_size, 1, 1), padding=(kernel_size // 2, 0, 0), + groups=channels) + self.pwconv = nn.Conv3d(channels, channels, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.pwconv(self.dwconv(F.silu(self.norm(x)))) + + +def _blocks(count: int, channels: int, emb_dim: int, temporal_every: int, kernel: int) -> nn.ModuleList: + blocks = nn.ModuleList() + for index in range(count): + blocks.append(ResBlockEmb3D(channels, emb_dim)) + if temporal_every > 0 and index % temporal_every == 0: + blocks.append(TemporalConv(channels, kernel)) + return blocks + + +class MiniMaxH3LatentUpscaler(nn.Module): + + def __init__(self, + in_channels: int = 24, + in_blocks: int = 12, + out_blocks: int = 12, + channels: int = 512, + temporal_every: int = 2, + temporal_kernel: int = 5) -> None: + super().__init__() + emb_dim = 64 + self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1) + self.embed = nn.Sequential(nn.Linear(1, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim)) + self.in_blocks = _blocks(in_blocks, channels, emb_dim, temporal_every, temporal_kernel) + self.out_blocks = _blocks(out_blocks, channels, emb_dim, temporal_every, temporal_kernel) + self.norm_out = _norm(channels) + self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1) + + def _run(self, blocks: nn.ModuleList, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + for block in blocks: + x = block(x, emb.expand(x.shape[0], -1)) if isinstance(block, ResBlockEmb3D) else block(x) + return x + + def forward(self, x: torch.Tensor, scale: float, target_size: tuple[int, int, int]) -> torch.Tensor: + """``x``: [B, C, T, h, w] normalized latent; returns [B, C, T, *target_size[1:]].""" + emb = self.embed(torch.tensor([[scale - 1.0]], dtype=x.dtype, device=x.device)) + x = self._run(self.in_blocks, self.conv_in(x), emb) + x = F.interpolate(x, size=target_size, mode="trilinear", align_corners=False) + x = self._run(self.out_blocks, x, emb) + return self.conv_out(F.silu(self.norm_out(x))) + + @classmethod + def from_checkpoint(cls, path: str) -> MiniMaxH3LatentUpscaler: + state = load_file(path, device="cpu") + model = cls(in_channels=state["conv_in.weight"].shape[1], + channels=state["conv_in.weight"].shape[0], + in_blocks=_count_blocks(state, "in_blocks"), + out_blocks=_count_blocks(state, "out_blocks"), + temporal_kernel=_temporal_kernel(state)) + model.load_state_dict(state, strict=True) + return model.eval().requires_grad_(False) + + +def _count_blocks(state: dict[str, torch.Tensor], prefix: str) -> int: + pattern = re.compile(rf"{prefix}\.(\d+)\.in_layers\.0\.weight") + return len({int(m.group(1)) for key in state if (m := pattern.match(key))}) + + +def _temporal_kernel(state: dict[str, torch.Tensor]) -> int: + return next(tensor.shape[2] for key, tensor in state.items() if key.endswith("dwconv.weight")) + + +def resolve_upscale_geometry(input_size: tuple[int, int], target_pixels: tuple[int, int]) -> tuple[tuple[int, int], float]: + """Latent (h, w) after aligning the target pixel size to 32, and the scale the network is conditioned on. + + ``input_size`` is the latent (h, w); ``target_pixels`` the requested (height, width). + """ + h_in, w_in = input_size + h_target, w_target = target_pixels + scale = (w_target / (w_in * VAE_DOWNSAMPLE) + h_target / (h_in * VAE_DOWNSAMPLE)) / 2.0 + aligned = (round(round(h_target / PIXEL_ALIGNMENT) * PIXEL_ALIGNMENT / VAE_DOWNSAMPLE), + round(round(w_target / PIXEL_ALIGNMENT) * PIXEL_ALIGNMENT / VAE_DOWNSAMPLE)) + return (max(1, aligned[0]), max(1, aligned[1])), scale + + +@torch.inference_mode() +def upscale_latent(model: MiniMaxH3LatentUpscaler, latent: torch.Tensor, target_pixels: tuple[int, int], + dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Upscale a [C, T, h, w] latent to ``target_pixels`` (height, width); FP32 [C, T, H, W] out. + + The input is returned unchanged when it already has the target latent size. + """ + channels, frames, h_in, w_in = latent.shape + (h_out, w_out), scale = resolve_upscale_geometry((h_in, w_in), target_pixels) + if (h_out, w_out) == (h_in, w_in): + return latent.float() + if scale < 1.0: + raise ValueError("The MiniMax-H3 latent upscaler only upscales (effective scale >= 1.0).") + mean = torch.tensor(LATENTS_MEAN, dtype=dtype, device=latent.device).view(1, -1, 1, 1, 1) + std = torch.tensor(LATENTS_STD, dtype=dtype, device=latent.device).view(1, -1, 1, 1, 1) + source = (latent[None].to(dtype) - mean) / std + result = model(source, scale, (frames, h_out, w_out)) + return (result * std + mean).float()[0] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/__init__.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/__init__.py new file mode 100644 index 000000000..2d9dc85f5 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.open_h3_regenerate_2k.open_h3_regenerate_2k_pipeline import ( + EntryClass, + OpenH3Regenerate2KPipeline, +) + +__all__ = ["EntryClass", "OpenH3Regenerate2KPipeline"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/draft_video.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/draft_video.py new file mode 100644 index 000000000..ce5efb6a4 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/draft_video.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Decode a draft video into the 124-frame pixel block the H3 VAE turns into 37 latent frames.""" + +from __future__ import annotations + +import subprocess + +import numpy as np +import torch + +import fastvideo.envs as envs + +DRAFT_SIZES: dict[str, tuple[int, int]] = {"768p": (1344, 768), "576p": (1024, 576)} +DRAFT_FRAMES = 121 +# The causal VAE maps 124 frames to 37 latent frames; the last frame is held for the three extra ones. +VAE_PADDING_FRAMES = 3 +# Seconds of held last frame ffmpeg appends so a short draft still yields DRAFT_FRAMES. +_TAIL_PAD_SECONDS = 6 +_DECODE_TIMEOUT_S = 600 + + +def ffmpeg_decode_command(path: str, width: int, height: int) -> list[str]: + """fps=24, cover-scale (bilinear) and centre-crop, last frame held, exactly DRAFT_FRAMES frames, no audio.""" + video_filter = (f"fps=24,scale={width}:{height}:force_original_aspect_ratio=increase:flags=bilinear," + f"crop={width}:{height},tpad=stop_mode=clone:stop_duration={_TAIL_PAD_SECONDS}") + return [ + envs.FASTVIDEO_FFMPEG_BIN.get(), "-v", "error", "-i", path, "-vf", video_filter, "-frames:v", + str(DRAFT_FRAMES), "-an", "-f", "rawvideo", "-pix_fmt", "rgb24", "pipe:1" + ] + + +def load_draft_pixels(path: str, size: str = "768p") -> torch.Tensor: + """uint8 pixels [1, 3, 124, H, W] on the CPU for the draft video at ``path``.""" + if size not in DRAFT_SIZES: + raise ValueError(f"draft size must be one of {sorted(DRAFT_SIZES)}, got {size!r}.") + width, height = DRAFT_SIZES[size] + try: + raw = subprocess.run(ffmpeg_decode_command(path, width, height), + capture_output=True, + check=True, + timeout=_DECODE_TIMEOUT_S).stdout + except FileNotFoundError as error: + raise RuntimeError(f"ffmpeg ({envs.FASTVIDEO_FFMPEG_BIN.get()!r}) is required to read the draft video; " + "set FASTVIDEO_FFMPEG_BIN.") from error + except subprocess.CalledProcessError as error: + raise ValueError( + f"Could not decode the draft video {path!r}: {error.stderr.decode(errors='replace')}") from error + expected = DRAFT_FRAMES * height * width * 3 + if len(raw) != expected: + raise ValueError(f"The draft video {path!r} decoded to {len(raw)} bytes, expected {expected} " + f"({DRAFT_FRAMES} frames at {width}x{height}).") + frames = torch.from_numpy(np.frombuffer(raw, np.uint8).reshape(DRAFT_FRAMES, height, width, 3).copy()) + pixels = frames.permute(3, 0, 1, 2)[None] + return torch.cat([pixels, pixels[:, :, -1:].expand(-1, -1, VAE_PADDING_FRAMES, -1, -1)], dim=2).contiguous() diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py new file mode 100644 index 000000000..d7f95a1a7 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Open-H3-Regenerate-2K: regenerate a 768P MiniMax-H3 draft video at 2560x1440. + +``draft video -> H3 VAE encode -> latent upscaler -> 3 fixed-sigma Euler steps (H3 DiT + merged LoRA + Sharp curve +AdaLN branches) -> H3 VAE decode``. The release repo holds the unmodified MiniMax-H3 components under ``base_model/`` +next to ``lora/`` and ``latent_upsampler/``; see ``docs/inference/open_h3_regenerate_2k.md``. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import torch + +from fastvideo import envs +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline +from fastvideo.pipelines.basic.minimax_h3.stages import MiniMaxH3VideoDecodingStage +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import SharpCurveBranches +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import ( + Regenerate2KDenoisingStage, + Regenerate2KFinalizeStage, + Regenerate2KLatentPreparationStage, +) +from fastvideo.utils import maybe_download_model, verify_model_config_and_directory + +logger = init_logger(__name__) + +BASE_MODEL_DIR = "base_model" +LORA_FILE = "lora/open_h3_regenerate_2k_lora.safetensors" +UPSCALER_FILE = "latent_upsampler/minimax_h3_latent_upscaler_3d_conv_v1_bf16.safetensors" +_BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae", "transformer", "scheduler", + "audio_scheduler") + + +def refuse_adaln_cache() -> None: + """Raise when an H3 AdaLN cache is enabled: it bypasses the projection modules that carry the Sharp branches.""" + if envs.FASTVIDEO_H3_ADALN_CACHE.get() or envs.FASTVIDEO_H3_ADALN_TABLE.get(): + raise ValueError("Open-H3-Regenerate-2K needs the AdaLN projections: unset FASTVIDEO_H3_ADALN_CACHE and " + "FASTVIDEO_H3_ADALN_TABLE, whose cached modulation skips the 50 Sharp curve branches.") + + +class OpenH3Regenerate2KPipeline(MiniMaxH3Pipeline): + """MiniMax-H3 T2VA layout with a draft-latent start, fixed sigmas and a silent output.""" + + _extra_config_module_map = {name: f"{BASE_MODEL_DIR}/{name}" for name in _BASE_COMPONENTS} + # The merged LoRA reaches only these linears; the AdaLN projections carry the Sharp curve branches instead. + lora_target_modules = ["attn.to_q", "attn.to_k", "attn.to_v", "attn.to_out", "ff.fc_in", "ff.fc_out"] + + @classmethod + def get_hf_download_allow_patterns(cls) -> list[str]: + components = [f"{BASE_MODEL_DIR}/{name}/**" for name in _BASE_COMPONENTS] + return [f"{BASE_MODEL_DIR}/model_index.json", *components, "lora/**", "latent_upsampler/*.safetensors"] + + def _load_config(self, model_path: str) -> dict[str, Any]: + refuse_adaln_cache() + revision = getattr(self.fastvideo_args, "revision", None) + self.model_path = maybe_download_model(model_path, + revision=revision, + allow_patterns=self.get_hf_download_allow_patterns()) + if self.fastvideo_args.lora_path is None: + self.fastvideo_args.lora_path = os.path.join(self.model_path, LORA_FILE) + return verify_model_config_and_directory(os.path.join(self.model_path, BASE_MODEL_DIR), + required_component_dirs=()) + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + super().initialize_pipeline(fastvideo_args) + self._upscaler: MiniMaxH3LatentUpscaler | None = None + self._curve: SharpCurveBranches | None = None + + def _load_upscaler(self) -> MiniMaxH3LatentUpscaler: + if self._upscaler is None: + model = MiniMaxH3LatentUpscaler.from_checkpoint(os.path.join(self.model_path, UPSCALER_FILE)) + self._upscaler = model.to(device=self.device, dtype=torch.bfloat16) + return self._upscaler + + def _load_curve(self) -> SharpCurveBranches: + if self._curve is None: + self._curve = SharpCurveBranches(Path(self.model_path)) + return self._curve + + def _add_denoise_stages(self, *, ref2va: bool) -> None: + config = self.fastvideo_args.pipeline_config + vae = self.get_module("vae") + audio_vae = self.get_module("audio_vae") + transformer = self.get_module("transformer") + if vae is None or audio_vae is None or transformer is None: + raise RuntimeError("Open-H3-Regenerate-2K needs the video VAE, audio VAE and transformer loaded.") + self.add_stage( + "latent_preparation_stage", + Regenerate2KLatentPreparationStage(vae, + audio_vae, + self.get_module("scheduler"), + load_upscaler=self._load_upscaler, + draft_size=config.draft_size, + start_sigma=config.refine_start_sigma)) + self.add_stage("denoising_stage", + Regenerate2KDenoisingStage(transformer, self._load_curve(), config.refine_start_sigma)) + self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae)) + self.add_stage("audio_decoding_stage", Regenerate2KFinalizeStage()) + self._denoise_stages_ready = True + + +EntryClass = OpenH3Regenerate2KPipeline + +__all__ = ["EntryClass", "OpenH3Regenerate2KPipeline"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/presets.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/presets.py new file mode 100644 index 000000000..1eb710e6a --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/presets.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference preset for Open-H3-Regenerate-2K.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_REFINE_STAGE = PresetStageSpec( + name="refine", + kind="denoising", + description="Latent upscale, then three fixed-sigma Euler refine steps", + allowed_overrides=frozenset({"num_inference_steps"}), +) + +OPEN_H3_REGENERATE_2K = InferencePreset( + name="open_h3_regenerate_2k", + version=1, + model_family="open_h3_regenerate_2k", + description="Open-H3-Regenerate-2K: regenerate a 768P MiniMax-H3 draft video at 2560x1440 (silent output)", + workload_type=None, + stage_schemas=(_REFINE_STAGE, ), + defaults={ + "height": 1440, + "width": 2560, + "num_frames": 124, + "fps": 24, + "guidance_scale": 1.0, + "batch_cfg": False, + "negative_prompt": "", + "num_inference_steps": 3, + "seed": 1766580932, + }, +) + +ALL_PRESETS = (OPEN_H3_REGENERATE_2K, ) + +__all__ = ["ALL_PRESETS", "OPEN_H3_REGENERATE_2K"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py new file mode 100644 index 000000000..017415fb5 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fixed Open-H3-Regenerate-2K refine schedule: Euler, CFG 1, model time ``1 - sigma``.""" + +from __future__ import annotations + +import torch + +REFINE_SHIFT = 6.0 +DEFAULT_REFINE_STEPS = 3 +# Start noise of the two-step variant; the three-step grid starts at 0.909375. +TWO_STEP_START_SIGMA = 0.7 + + +def shifted_grid(start: float, steps: int, shift: float = REFINE_SHIFT) -> list[float]: + """Euler grid from ``start`` to 0: even spacing in unshifted sigma, then ``s * x / (1 + (s - 1) * x)``.""" + if not 0.0 < start < 1.0: + raise ValueError(f"start sigma must be in (0, 1), got {start}.") + if steps < 1: + raise ValueError(f"steps must be positive, got {steps}.") + raw = start / (shift - (shift - 1) * start) + mapped = [shift * x / (1 + (shift - 1) * x) for x in (raw * (1 - i / steps) for i in range(1, steps))] + return [start, *mapped, 0.0] + + +REFINE_SCHEDULES: dict[int, tuple[float, ...]] = { + 3: (0.909375, 6 / 7, 2 / 3, 0.0), + 2: tuple(shifted_grid(TWO_STEP_START_SIGMA, 2)), +} + + +def refine_sigmas(steps: int = DEFAULT_REFINE_STEPS, start_sigma: float | None = None) -> list[float]: + """Sigma grid of ``steps`` Euler steps, ending at 0.""" + if start_sigma is not None: + return shifted_grid(start_sigma, steps) + if steps not in REFINE_SCHEDULES: + raise ValueError(f"Open-H3-Regenerate-2K supports {sorted(REFINE_SCHEDULES)} refine steps, got {steps}.") + return list(REFINE_SCHEDULES[steps]) + + +def mix_initial_state(upscaled: torch.Tensor, noise: torch.Tensor, sigma: float) -> torch.Tensor: + """Re-noise the upscaled draft latent to the first sigma: ``(1 - sigma) * upscaled + sigma * noise``.""" + return (1.0 - sigma) * upscaled + sigma * noise diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py new file mode 100644 index 000000000..a08efab22 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LMS Sharp curve AdaLN branches. + +Fifty AdaLN projections get an extra low-rank term whose input is an 8-D curve interpolated from the model timestep +(``adaln_t_table``, 1025 x 8) instead of the layer input, so it cannot be folded into the LoRA. The term is added to the +projection output through a forward hook; ``set_timesteps`` must be called with the per-forward unique timesteps first. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file + +CURVE_BRANCHES_FILE = "lora/sharp_curve_branches.safetensors" +CURVE_TABLE_FILE = "lora/adaln_t_table.bin" +CURVE_TABLE_SHAPE = (1025, 8) +NUM_CURVE_BRANCHES = 50 +_KEY_PREFIX = "transformer." +_A_SUFFIX = ".curve_A.weight" +_B_SUFFIX = ".curve_B.weight" + + +def interpolate_curve(table: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor: + """Linear interpolation of ``table`` rows at model times in [0, 1]; ``timesteps`` [U] -> [U, 8].""" + t = timesteps.float() + if t.numel() and (t.min() < 0 or t.max() > 1): + raise ValueError("Sharp curve timesteps must lie in [0, 1].") + table = table.to(t.device) + position = t * (len(table) - 1) + low = position.floor().long().clamp(max=len(table) - 2) + return torch.lerp(table[low], table[low + 1], (position - low).unsqueeze(-1)) + + +_ATTACHED_ATTR = "_open_h3_sharp_curve_branches" + + +class SharpCurveBranches: + + def __init__(self, weights_dir: str | Path) -> None: + root = Path(weights_dir) + raw = torch.frombuffer(bytearray((root / CURVE_TABLE_FILE).read_bytes()), dtype=torch.float32) + if raw.numel() != CURVE_TABLE_SHAPE[0] * CURVE_TABLE_SHAPE[1]: + raise ValueError(f"{root / CURVE_TABLE_FILE} must hold a {CURVE_TABLE_SHAPE} float32 table.") + self.table = raw.reshape(CURVE_TABLE_SHAPE).clone() + tensors = load_file(str(root / CURVE_BRANCHES_FILE), device="cpu") + self.factors: dict[str, tuple[torch.Tensor, torch.Tensor]] = { + key[len(_KEY_PREFIX):-len(_A_SUFFIX)]: (value.float(), tensors[key[:-len(_A_SUFFIX)] + _B_SUFFIX].float()) + for key, value in tensors.items() if key.endswith(_A_SUFFIX) + } + if len(self.factors) != NUM_CURVE_BRANCHES: + raise ValueError(f"Expected {NUM_CURVE_BRANCHES} Sharp curve branches, found {len(self.factors)}.") + self.coords: torch.Tensor | None = None + self._device_factors: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + + def set_timesteps(self, unique_timesteps: torch.Tensor) -> None: + self.coords = interpolate_curve(self.table, unique_timesteps) + + def attach(self, transformer: torch.nn.Module) -> None: + """Register the branch hooks on ``transformer`` unless this instance already carries them. + + The mark lives on the module, so a reloaded lazy transformer, a new instance, gets its own hooks. + """ + if getattr(transformer, _ATTACHED_ATTR, None) is self: + return + modules = dict(transformer.named_modules()) + missing = sorted(set(self.factors) - set(modules)) + if missing: + raise ValueError(f"Sharp curve branches match no transformer module: {missing[:5]}") + for name in self.factors: + modules[name].register_forward_hook(self._hook(name)) + setattr(transformer, _ATTACHED_ATTR, self) + + def _hook(self, name: str): + + def add_branch(module, inputs, output): + projected = output[0] if isinstance(output, tuple) else output + coords = self.coords + if coords is None or coords.shape[0] != projected.shape[0]: + raise RuntimeError("Sharp curve timesteps were not set for this forward.") + factors = self._device_factors.get(name) + if factors is None or factors[0].device != projected.device: + factors = tuple(f.to(projected.device) for f in self.factors[name]) + self._device_factors[name] = factors + delta = F.linear(F.linear(coords, factors[0]), factors[1]).to(projected.dtype) + projected = projected + delta + return (projected, *output[1:]) if isinstance(output, tuple) else projected + + return add_branch diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/__init__.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/__init__.py new file mode 100644 index 000000000..3c7a022bf --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages.denoising import Regenerate2KDenoisingStage +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages.finalize import Regenerate2KFinalizeStage +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages.latent_preparation import ( + Regenerate2KLatentPreparationStage, ) + +__all__ = [ + "Regenerate2KDenoisingStage", + "Regenerate2KFinalizeStage", + "Regenerate2KLatentPreparationStage", +] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py new file mode 100644 index 000000000..8e20a67ac --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fixed-sigma Euler refine loop of Open-H3-Regenerate-2K.""" + +from __future__ import annotations + +import contextlib +from typing import Any + +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout, build_row_timesteps +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import refine_sigmas +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import SharpCurveBranches +from fastvideo.pipelines.lazy_module import is_lazy_module +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class Regenerate2KDenoisingStage(PipelineStage): + """Euler steps over the refine sigmas with the video clock at ``1 - sigma`` and the silent audio held clean.""" + + performance_component_metric = "dit_time_s" + + def __init__(self, transformer: Any, curve: SharpCurveBranches, start_sigma: float | None) -> None: + super().__init__() + self.transformer = transformer + self.curve = curve + self.start_sigma = start_sigma + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors_dims(3)) + result.add_check("latents", batch.latents, V.with_dims(2)) + result.add_check("audio_latents", batch.audio_latents, V.with_dims(2)) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, V.with_dims(2)) + return result + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + transformer = self.transformer + self.curve.attach(transformer.materialize() if is_lazy_module(transformer) else transformer) + layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY) + if not isinstance(layout, MiniMaxH3PackedLayout): + raise ValueError("The MiniMax-H3 packed layout is missing before denoising.") + device = get_local_torch_device() + sigmas = refine_sigmas(batch.num_inference_steps, self.start_sigma) + state = batch.latents.to(device) + layerwise = bool(getattr(fastvideo_args, "dit_layerwise_offload", False)) + full_cpu_offload = (bool(getattr(fastvideo_args, "dit_cpu_offload", False)) and not layerwise + and not bool(getattr(fastvideo_args, "use_fsdp_inference", False))) + try: + if full_cpu_offload: + transformer.to(device) + with self._attention_kernel(device): + state = self._euler_loop(batch, layout, state, sigmas, device, transformer) + finally: + if layerwise: + manager = getattr(transformer, "_layerwise_offload_manager", None) + if manager is not None and getattr(manager, "enabled", False): + manager.release_all() + if full_cpu_offload: + transformer.to("cpu") + batch.latents = state + return batch + + @staticmethod + def _attention_kernel(device: torch.device) -> contextlib.AbstractContextManager[None]: + """cuDNN attention: at 133k tokens PyTorch's default SDPA dispatch picks its flash kernel, 3x slower on B200.""" + return sdpa_kernel(SDPBackend.CUDNN_ATTENTION) if device.type == "cuda" else contextlib.nullcontext() + + def _euler_loop(self, batch: ForwardBatch, layout: MiniMaxH3PackedLayout, state: torch.Tensor, sigmas: list[float], + device: torch.device, transformer: Any) -> torch.Tensor: + position_ids = layout.position_ids.to(device) + token_tags = layout.token_tags.to(device) + video_indices = layout.video_indices.to(device) + audio_indices = layout.audio_indices.to(device) + text_indices = layout.text_indices.to(device) + prompt_embeds = batch.prompt_embeds[0].to(device) + audio = batch.audio_latents.to(device=device, dtype=torch.bfloat16)[None] + for index, (sigma, following) in enumerate(zip(sigmas[:-1], sigmas[1:], strict=True)): + unique, inverse = build_row_timesteps(layout, + video_timestep=1.0 - sigma, + audio_timestep=1.0, + condition_video_timestep=1.0, + condition_audio_timestep=1.0) + unique, inverse = unique.to(device), inverse.to(device) + self.curve.set_timesteps(unique) + torch.compiler.cudagraph_mark_step_begin() + with set_forward_context(current_timestep=index, attn_metadata=None, forward_batch=batch): + velocity, _ = transformer(hidden_states=state.to(torch.bfloat16)[None], + audio_hidden_states=audio, + encoder_hidden_states=prompt_embeds, + timestep=unique, + timestep_indices=inverse, + token_tags=token_tags, + position_ids=position_ids, + video_indices=video_indices, + audio_indices=audio_indices, + text_indices=text_indices) + state = state + (sigma - following) * velocity[0].float() + batch.step_index = index + return state + + +__all__ = ["Regenerate2KDenoisingStage"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/finalize.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/finalize.py new file mode 100644 index 000000000..300caa4f8 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/finalize.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Trim the decoded video to the draft length and drop the per-request H3 state.""" + +from __future__ import annotations + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.open_h3_regenerate_2k.draft_video import DRAFT_FRAMES +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + + +class Regenerate2KFinalizeStage(PipelineStage): + """The released pipeline writes a silent video: the audio stream only conditions the DiT and is never decoded.""" + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("output", batch.output, V.with_dims(5)) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("output", batch.output, V.with_dims(5)) + return result + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + # The VAE decodes 124 frames (37 latent frames); the draft has 121. + batch.output = batch.output[:, :, :DRAFT_FRAMES] + batch.prompt_embeds = [] + batch.latents = None + batch.audio_latents = None + for key in tuple(batch.extra): + if key.startswith("minimax_h3_"): + batch.extra.pop(key) + return batch + + +__all__ = ["Regenerate2KFinalizeStage"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/latent_preparation.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/latent_preparation.py new file mode 100644 index 000000000..93f28431f --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/latent_preparation.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Draft video -> upscaled, re-noised H3 latent plus the fixed silent-audio rows.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models import pinned_offload +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler, upscale_latent +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_FPS, + MINIMAX_H3_KEYFRAME_ENCODE_SEED, + h3_dit_patch_size, + patchify_video_latents, +) +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import ( + MINIMAX_H3_LAYOUT_KEY, + MiniMaxH3LatentPreparationStage, +) +from fastvideo.pipelines.basic.open_h3_regenerate_2k.draft_video import load_draft_pixels +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import mix_initial_state, refine_sigmas +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + +DRAFT_FRAME_COUNT = 124 + + +class Regenerate2KLatentPreparationStage(MiniMaxH3LatentPreparationStage): + """Encode the draft with the H3 VAE, upscale its latent to the target grid and re-noise it to the first sigma.""" + + def __init__(self, vae: Any, audio_vae: Any, scheduler: Any, *, load_upscaler: Callable[[], + MiniMaxH3LatentUpscaler], + draft_size: str, start_sigma: float | None) -> None: + super().__init__(vae, audio_vae, scheduler, ref2va=False) + self.load_upscaler = load_upscaler + self.draft_size = draft_size + self.start_sigma = start_sigma + self._silence: tuple[tuple[str, int], torch.Tensor] | None = None + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = super().verify_input(batch, fastvideo_args) + result.add_check("video_path", batch.video_path, lambda value: isinstance(value, str) and bool(value)) + return result + + def _encode_draft(self, video_path: str, device: torch.device, fastvideo_args: FastVideoArgs) -> torch.Tensor: + """Normalized native latent [C, T, h, w] of the draft, with the VAE posterior sampled under a fixed seed.""" + pixels = load_draft_pixels(video_path, self.draft_size) + pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory) + try: + posterior = self.vae.encode_pixels(pixels).latent_dist + sample = posterior.sample(generator=torch.Generator("cpu").manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) + return self.vae.normalize_latents(sample.float())[0].to(device) + finally: + if fastvideo_args.vae_cpu_offload: + pinned_offload.unload(self.vae) + + def _silent_audio_rows(self, num_frames: int, device: torch.device, fastvideo_args: FastVideoArgs) -> torch.Tensor: + """The encoded-silence audio stream, cached: it is the same for every request of this length.""" + key = (str(device), num_frames) + if self._silence is not None and self._silence[0] == key: + return self._silence[1] + pinned_offload.load(self.audio_vae, device, pin=fastvideo_args.pin_cpu_memory) + try: + samples = round(num_frames / MINIMAX_H3_FPS * int(self.audio_vae.sampling_rate)) + silence = torch.zeros(MINIMAX_H3_AUDIO_CHANNELS, 1, samples, device=device) + posterior = self.audio_vae.encode(silence).latent_dist + latents = self.audio_vae.normalize_latents(posterior.mode().float()).transpose(1, 2) + rows = latents.reshape(-1, self.audio_vae.latent_channels).to(torch.bfloat16) + finally: + if fastvideo_args.vae_cpu_offload: + pinned_offload.unload(self.audio_vae) + self._silence = (key, rows) + return rows + + @staticmethod + def _noise_rows(shape: tuple[int, int], device: torch.device, seed: int) -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(seed) + return torch.randn((1, *shape), device=device, dtype=torch.float32, generator=generator)[0] + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.num_frames != DRAFT_FRAME_COUNT: + raise ValueError(f"Open-H3-Regenerate-2K refines {DRAFT_FRAME_COUNT} frames, got {batch.num_frames}.") + device = get_local_torch_device() + patch_size = h3_dit_patch_size(fastvideo_args) + _, channels, num_latent_frames, latent_height, latent_width = batch.raw_latent_shape + + draft = self._encode_draft(batch.video_path, device, fastvideo_args) + upscaler = self.load_upscaler() + upscaled = upscale_latent(upscaler, draft, (batch.height, batch.width)) + if tuple(upscaled.shape) != (channels, num_latent_frames, latent_height, latent_width): + raise ValueError(f"The upscaled latent has shape {tuple(upscaled.shape)}, expected " + f"{(channels, num_latent_frames, latent_height, latent_width)}.") + upscaled_rows = patchify_video_latents(upscaled[None], patch_size) + + if batch.latents is None: + noise_rows = self._noise_rows(tuple(upscaled_rows.shape), device, 0 if batch.seed is None else batch.seed) + else: + expected = (1, channels, num_latent_frames, latent_height, latent_width) + if tuple(batch.latents.shape) != expected: + raise ValueError(f"Injected noise must have shape {expected}, got {tuple(batch.latents.shape)}.") + noise_rows = patchify_video_latents(batch.latents.to(device=device, dtype=torch.float32), patch_size) + + layout = self._build_layout(batch, fastvideo_args) + audio_rows = self._silent_audio_rows(batch.num_frames, device, fastvideo_args) + if audio_rows.shape[0] != layout.audio_indices.numel(): + raise ValueError(f"The silent audio has {audio_rows.shape[0]} rows, the packed layout expects " + f"{layout.audio_indices.numel()} ({layout.num_audio_latents} latents).") + + batch.latents = mix_initial_state(upscaled_rows, noise_rows, + refine_sigmas(batch.num_inference_steps, self.start_sigma)[0]) + batch.audio_latents = audio_rows + batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout + for key in tuple(batch.extra): + if key.startswith("minimax_h3_") and key != MINIMAX_H3_LAYOUT_KEY: + batch.extra.pop(key) + return batch + + +__all__ = ["Regenerate2KLatentPreparationStage"] diff --git a/fastvideo/registry.py b/fastvideo/registry.py index d792c677e..0a0d18361 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -43,6 +43,7 @@ from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.configs.pipelines.pika_h3_flex import PikaH3FlexPipelineConfig +from fastvideo.configs.pipelines.open_h3_regenerate_2k import OpenH3Regenerate2KPipelineConfig from fastvideo.configs.pipelines.mmaudio import ( MMAudioLarge44kV2PipelineConfig, MMAudioLarge44kV2AConfig, @@ -218,6 +219,13 @@ def _resolve_class_name(config: dict[str, Any]) -> str | None: return class_name +OPEN_H3_REGENERATE_2K_HF_PATH = "Efficient-Large-Model/Open-H3-Regenerate-2K" + + +def _is_open_h3_regenerate_2k_release(model_path: str) -> bool: + return os.path.isfile(os.path.join(model_path, "lora", "open_h3_regenerate_2k_lora.safetensors")) + + def _get_config_info( model_path: str, *, @@ -249,6 +257,9 @@ def _get_config_info( return _CONFIG_REGISTRY.get(model_id) # 3. Use detectors (path or model_index pipeline name). + if os.path.exists(model_path) and _is_open_h3_regenerate_2k_release(model_path): + # The release keeps its Diffusers manifest under base_model/, so a renamed local copy has no root manifest. + return _CONFIG_REGISTRY.get(_MODEL_HF_PATH_TO_NAME[OPEN_H3_REGENERATE_2K_HF_PATH]) if os.path.exists(model_path): config = verify_model_config_and_directory(model_path, required_component_dirs=[]) elif allow_remote: @@ -1174,6 +1185,22 @@ def _register_configs() -> None: pipeline_cls_name="PikaH3FlexModularPipeline", ) + # Open-H3-Regenerate-2K (video -> video). Registered before MiniMax H3 so the first matching detector wins. There is + # no V2V WorkloadType, so, like the Kandinsky6 SR entries, it declares no workload. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=OpenH3Regenerate2KPipelineConfig, + workload_types=(), + hf_model_paths=[OPEN_H3_REGENERATE_2K_HF_PATH], + model_detectors=[ + lambda path: any(token in path.lower() + for token in ("open-h3-regenerate", "open_h3_regenerate", "openh3regenerate2kpipeline")), + ], + model_family="open_h3_regenerate_2k", + default_preset="open_h3_regenerate_2k", + pipeline_cls_name="OpenH3Regenerate2KPipeline", + ) + # MiniMax H3 register_configs( sampling_param_cls=None, @@ -1285,19 +1312,18 @@ def get_model_info( # is triggered just to log the replaced name). pipeline_name: str | None = override_pipeline_cls_name logger.info("Using override pipeline class name %s", pipeline_name) + elif config_info.pipeline_cls_name is not None: + # The resolved (path/detector-based) config pins the pipeline class, e.g. an I2V checkpoint whose + # `_class_name` would otherwise resolve to the T2V pipeline. The manifest is not read: a release such as + # Open-H3-Regenerate-2K keeps it in a subfolder. + pipeline_name = config_info.pipeline_cls_name + logger.info("Pinned pipeline class name %s", pipeline_name) else: if os.path.exists(model_path): config = verify_model_config_and_directory(model_path, required_component_dirs=[]) else: config = maybe_download_model_index(model_path, revision=revision) - pipeline_name = _resolve_class_name(config) - if config_info.pipeline_cls_name is not None: - # The resolved (path/detector-based) config pins the pipeline class, - # e.g. an I2V checkpoint whose `_class_name` would otherwise resolve - # to the T2V pipeline. - logger.info("Pinning pipeline class name from %s to %s", pipeline_name, config_info.pipeline_cls_name) - pipeline_name = config_info.pipeline_cls_name if pipeline_name is None: raise ValueError("Model config does not contain a _class_name attribute. " @@ -1374,6 +1400,8 @@ def _register_presets() -> None: ALL_PRESETS as MINIMAX_H3_PRESETS, ) from fastvideo.pipelines.basic.pika_h3_flex.presets import ( ALL_PRESETS as PIKA_H3_FLEX_PRESETS, ) + from fastvideo.pipelines.basic.open_h3_regenerate_2k.presets import ( + ALL_PRESETS as OPEN_H3_REGENERATE_2K_PRESETS, ) from fastvideo.pipelines.basic.mmaudio.presets import ( ALL_PRESETS as MMAUDIO_PRESETS, ) from fastvideo.pipelines.basic.sd35.presets import ( @@ -1408,6 +1436,7 @@ def _register_presets() -> None: MATRIXGAME3_PRESETS, MINIMAX_H3_PRESETS, PIKA_H3_FLEX_PRESETS, + OPEN_H3_REGENERATE_2K_PRESETS, MMAUDIO_PRESETS, SD35_PRESETS, STABLE_AUDIO_PRESETS, diff --git a/fastvideo/tests/entrypoints/test_openai_open_h3_regenerate_2k.py b/fastvideo/tests/entrypoints/test_openai_open_h3_regenerate_2k.py new file mode 100644 index 000000000..16a2530c5 --- /dev/null +++ b/fastvideo/tests/entrypoints/test_openai_open_h3_regenerate_2k.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Open-H3-Regenerate-2K through the OpenAI-compatible video routes, with a fake generator (no GPU).""" + +import asyncio +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi.testclient import TestClient +from PIL import Image + +from fastvideo.entrypoints.cli.inference_config import build_serve_config +from fastvideo.entrypoints.openai.api_server import create_app +from fastvideo.entrypoints.openai.media_policy import MediaPolicy +from fastvideo.entrypoints.openai.stores import VIDEO_STORE + +ROOT = Path(__file__).resolve().parents[3] +MODEL = "Efficient-Large-Model/Open-H3-Regenerate-2K" +ARTIFACT = b"fake-2k-video-for-transport-tests" +PROMPT = "a quiet street at dusk" + + +class FileGenerator: + + def __init__(self): + self.requests = [] + + def generate(self, request): + self.requests.append(request) + output = Path(request.output.output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(ARTIFACT) + return SimpleNamespace(video_path=str(output), generation_time=0.01) + + def shutdown(self): + pass + + +@pytest.fixture +def served(monkeypatch, tmp_path): + generator = FileGenerator() + monkeypatch.setattr("fastvideo.entrypoints.openai.api_server.VideoGenerator.from_fastvideo_args", + lambda args: generator) + drafts = tmp_path / "drafts" + drafts.mkdir() + (drafts / "draft.mp4").write_bytes(b"draft") + Image.new("RGB", (8, 8)).save(drafts / "still.png") + args = SimpleNamespace(model_path=MODEL, lora_path=None, lora_nickname="default", lora_strength=1.0, + override_pipeline_cls_name=None) + config = build_serve_config(SimpleNamespace(config=str(ROOT / "examples/serving/openai_open_h3_regenerate_2k.yaml")), + overrides=[]) + app = create_app(args, + str(tmp_path / "server"), + config.default_request, + config.server.served_model_name, + media_policy=MediaPolicy(roots=(str(drafts), ))) + with TestClient(app) as client: + yield client, generator, drafts + asyncio.run(VIDEO_STORE.clear()) + + +def _body(drafts, **extra): + return {"prompt": PROMPT, "video_reference": {"video_url": str(drafts / "draft.mp4")}, **extra} + + +def test_serving_yaml_pins_the_2k_defaults(): + config = build_serve_config(SimpleNamespace(config=str(ROOT / "examples/serving/openai_open_h3_regenerate_2k.yaml")), + overrides=[]) + sampling = config.default_request.sampling + assert (sampling.height, sampling.width, sampling.num_frames, sampling.num_inference_steps) == (1440, 2560, 124, 3) + assert config.generator.model_path == MODEL + + +def test_sync_route_returns_the_mp4_and_maps_the_draft_to_video_path(served): + client, generator, drafts = served + response = client.post("/v1/videos/sync", json=_body(drafts, seed=7)) + assert response.status_code == 200, response.text + assert response.content == ARTIFACT + assert response.headers["content-type"] == "video/mp4" + assert response.headers["x-model"] == "open-h3-regenerate-2k" + request = generator.requests[0] + assert request.prompt == PROMPT + assert request.inputs.video_path == str((drafts / "draft.mp4").resolve()) + assert (request.sampling.height, request.sampling.width, request.sampling.num_frames) == (1440, 2560, 124) + assert request.sampling.num_inference_steps == 3 and request.sampling.seed == 7 + + +def test_async_route_completes_and_serves_the_content(served): + client, generator, drafts = served + created = client.post("/v1/videos", json=_body(drafts)) + assert created.status_code == 200, created.text + video_id = created.json()["id"] + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + status = client.get(f"/v1/videos/{video_id}").json()["status"] + if status in ("completed", "failed"): + break + time.sleep(0.02) + assert status == "completed" + assert client.get(f"/v1/videos/{video_id}/content").content == ARTIFACT + assert generator.requests[0].inputs.video_path.endswith("draft.mp4") + + +def test_two_refine_steps_can_be_requested(served): + client, generator, drafts = served + assert client.post("/v1/videos/sync", json=_body(drafts, num_inference_steps=2)).status_code == 200 + assert generator.requests[0].sampling.num_inference_steps == 2 + + +@pytest.mark.parametrize("reference, message", [ + ({"video_url": "data:video/mp4;base64,AAAA"}, "Data video references are not supported"), + ({"video_url": "/etc/passwd"}, "Local media paths are not accepted"), + ({"file_id": "file_1"}, "file_id video references are not supported"), +]) +def test_unusable_video_references_are_rejected(served, reference, message): + client, generator, _ = served + response = client.post("/v1/videos/sync", json={"prompt": PROMPT, "video_reference": reference}) + assert response.status_code == 400 and message in response.text + assert not generator.requests + + +@pytest.mark.parametrize("extra, message", [ + ({"video_reference": None}, "exactly one video_reference"), + ({"image_reference": {"image_url": "STILL"}}, "exactly one video_reference"), + ({"size": "1920x1080"}, "always renders 2560x1440"), + ({"seconds": 5}, "first 121 frames"), + ({"fps": 30}, "requires fps=24"), +]) +def test_requests_the_pipeline_cannot_honor_are_rejected(served, extra, message): + client, generator, drafts = served + body = {**_body(drafts), **extra} + if "image_reference" in body: + body["image_reference"] = {"image_url": str(drafts / "still.png")} + body = {key: value for key, value in body.items() if value is not None} + response = client.post("/v1/videos/sync", json=body) + assert response.status_code == 400 and message in response.text + assert not generator.requests + + +def test_two_drafts_are_rejected(served): + client, generator, drafts = served + reference = {"video_url": str(drafts / "draft.mp4")} + response = client.post("/v1/videos/sync", json={"prompt": PROMPT, "video_reference": [reference, reference]}) + assert response.status_code == 400 and "exactly one video_reference" in response.text diff --git a/fastvideo/tests/layers/test_lora_merge_precision.py b/fastvideo/tests/layers/test_lora_merge_precision.py new file mode 100644 index 000000000..6bbf39ca0 --- /dev/null +++ b/fastvideo/tests/layers/test_lora_merge_precision.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contract: merging an adapter stored wider than the base rounds once, after a wide add.""" + +from __future__ import annotations + +import pytest +import torch + +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.lora import linear as lora_linear +from fastvideo.layers.lora.linear import get_lora_layer + +IN, OUT, RANK = 32, 24, 16 + + +@pytest.fixture(autouse=True) +def _cpu_merge(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(torch._dynamo.config, "disable", True) + monkeypatch.setattr(lora_linear, "get_local_torch_device", lambda: torch.device("cpu")) + + +def _merged(adapter_dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + torch.manual_seed(0) + layer = ReplicatedLinear(IN, OUT, bias=False, params_dtype=torch.bfloat16) + with torch.no_grad(): + layer.weight.copy_(torch.randn(OUT, IN)) + base = layer.weight.detach().clone() + wrapped = get_lora_layer(layer) + a, b = torch.randn(RANK, IN, dtype=adapter_dtype), torch.randn(OUT, RANK, dtype=adapter_dtype) * 0.1 + wrapped.set_lora_weights(a, b, lora_alpha=RANK, strength=1.0) + wrapped.merge_lora_weights() + return wrapped.base_layer.weight.detach(), base, a, b + + +def test_fp32_adapter_is_added_in_fp32_and_rounded_once() -> None: + merged, base, a, b = _merged(torch.float32) + expected = (base.float() + b @ a).to(torch.bfloat16) + assert merged.dtype == torch.bfloat16 + assert torch.equal(merged, expected) + + +def test_bf16_adapter_keeps_the_bf16_merge() -> None: + merged, base, a, b = _merged(torch.bfloat16) + assert torch.equal(merged, base + b @ a) + + +def test_merge_helper_used_with_and_without_fsdp_rounds_once() -> None: + torch.manual_seed(1) + layer = ReplicatedLinear(IN, OUT, bias=False, params_dtype=torch.bfloat16) + with torch.no_grad(): + layer.weight.copy_(torch.randn(OUT, IN)) + wrapped = get_lora_layer(layer) + a, b = torch.randn(RANK, IN), torch.randn(OUT, RANK) * 0.1 + wrapped.lora_A, wrapped.lora_B = torch.nn.Parameter(a), torch.nn.Parameter(b) + wrapped.lora_rank, wrapped.lora_alpha, wrapped.lora_strength = RANK, RANK, 0.5 + weight = layer.weight.detach().clone() + + merged = wrapped._merged_weight(weight) + + assert merged.dtype == torch.bfloat16 + assert torch.equal(merged, (weight.float() + 0.5 * (b @ a)).to(torch.bfloat16)) diff --git a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py new file mode 100644 index 000000000..e5becaa9d --- /dev/null +++ b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the Open-H3-Regenerate-2K building blocks.""" + +import numpy as np +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import ( + MiniMaxH3LatentUpscaler, + resolve_upscale_geometry, + upscale_latent, +) +from fastvideo.pipelines.basic.open_h3_regenerate_2k import draft_video +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import ( + mix_initial_state, + refine_sigmas, + shifted_grid, +) +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import ( + NUM_CURVE_BRANCHES, + SharpCurveBranches, + interpolate_curve, +) + + +def test_default_sigmas_match_the_released_grid(): + assert refine_sigmas(3) == pytest.approx([0.909375, 0.857143, 0.666667, 0.0], abs=1e-6) + + +def test_two_step_sigmas_start_at_0_7(): + assert refine_sigmas(2) == pytest.approx([0.7, 0.494118, 0.0], abs=1e-6) + + +def test_shifted_grid_from_the_released_start_is_the_original_four_step_grid(): + assert shifted_grid(0.909375, 4) == pytest.approx([0.909375, 0.841446, 0.732075, 0.526697, 0.0], abs=1e-6) + + +def test_custom_start_sigma_is_decreasing_and_ends_at_zero(): + grid = refine_sigmas(3, start_sigma=0.8) + assert grid[0] == 0.8 and grid[-1] == 0.0 + assert all(a > b for a, b in zip(grid, grid[1:])) + + +@pytest.mark.parametrize("steps", [0, 4]) +def test_unsupported_step_count_is_rejected(steps): + with pytest.raises(ValueError): + refine_sigmas(steps) + + +def test_initial_state_mix_weights_noise_by_sigma(): + upscaled, noise = torch.ones(2, 3), torch.zeros(2, 3) + assert torch.allclose(mix_initial_state(upscaled, noise, 0.909375), torch.full((2, 3), 0.090625)) + assert torch.equal(mix_initial_state(upscaled, noise, 0.0), upscaled) + + +def test_upscale_geometry_targets_2k_from_768p_and_576p(): + assert resolve_upscale_geometry((48, 84), (1440, 2560))[0] == (90, 160) + shape, scale = resolve_upscale_geometry((36, 64), (1440, 2560)) + assert shape == (90, 160) and scale == pytest.approx(2.5) + + +def test_tiny_upscaler_forward_shape_and_identity_passthrough(): + model = MiniMaxH3LatentUpscaler(channels=32, in_blocks=2, out_blocks=2).eval() + latent = torch.randn(24, 5, 4, 6) + out = upscale_latent(model, latent, (96, 160), dtype=torch.float32) + assert out.shape == (24, 5, 6, 10) and out.dtype == torch.float32 + assert torch.equal(upscale_latent(model, latent, (64, 96)), latent) + + +def test_upscaler_checkpoint_roundtrip(tmp_path): + model = MiniMaxH3LatentUpscaler(channels=32, in_blocks=3, out_blocks=2) + path = str(tmp_path / "upscaler.safetensors") + save_file(model.state_dict(), path) + loaded = MiniMaxH3LatentUpscaler.from_checkpoint(path) + assert len(loaded.in_blocks) == len(model.in_blocks) and len(loaded.out_blocks) == len(model.out_blocks) + for key, value in model.state_dict().items(): + assert torch.equal(loaded.state_dict()[key], value) + + +def test_curve_interpolation_hits_table_rows_and_midpoints(): + table = torch.arange(1025 * 8, dtype=torch.float32).reshape(1025, 8) + coords = interpolate_curve(table, torch.tensor([0.0, 1.0, 0.5, 1 / 2048])) + assert torch.equal(coords[0], table[0]) and torch.equal(coords[1], table[1024]) + assert torch.equal(coords[2], table[512]) + assert torch.allclose(coords[3], (table[0] + table[1]) / 2) + with pytest.raises(ValueError): + interpolate_curve(table, torch.tensor([1.5])) + + +class _Projection(torch.nn.Module): + + def __init__(self, out_features): + super().__init__() + self.linear = torch.nn.Linear(4, out_features) + + def forward(self, x): + return self.linear(x), None + + +def _write_curve_release(root, out_features, count): + (root / "lora").mkdir() + (root / "lora" / "adaln_t_table.bin").write_bytes(np.random.rand(1025 * 8).astype(np.float32).tobytes()) + tensors = {} + for index in range(count): + tensors[f"transformer.blocks.{index}.proj.curve_A.weight"] = torch.randn(3, 8) + tensors[f"transformer.blocks.{index}.proj.curve_B.weight"] = torch.randn(out_features, 3) + save_file(tensors, str(root / "lora" / "sharp_curve_branches.safetensors")) + + +def test_sharp_branches_require_all_fifty(tmp_path): + _write_curve_release(tmp_path, 6, 3) + with pytest.raises(ValueError, match="50"): + SharpCurveBranches(tmp_path) + + +def test_sharp_branch_adds_curve_term_to_projection_output(tmp_path): + _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES) + curve = SharpCurveBranches(tmp_path) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for _ in range(NUM_CURVE_BRANCHES): + block = torch.nn.Module() + block.proj = _Projection(6) + model.blocks.append(block) + curve.attach(model) + curve.attach(model) + timesteps = torch.tensor([0.1, 0.9]) + x = torch.randn(2, 4) + base, _ = model.blocks[0].proj.linear(x), None + curve.set_timesteps(timesteps) + out, extra = model.blocks[0].proj(x) + a, b = curve.factors["blocks.0.proj"] + expected = base + torch.nn.functional.linear(torch.nn.functional.linear(curve.coords, a), b) + assert extra is None and torch.allclose(out, expected, atol=1e-6) + + +def test_sharp_branch_without_timesteps_fails_loudly(tmp_path): + _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES) + curve = SharpCurveBranches(tmp_path) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for _ in range(NUM_CURVE_BRANCHES): + block = torch.nn.Module() + block.proj = _Projection(6) + model.blocks.append(block) + curve.attach(model) + with pytest.raises(RuntimeError, match="timesteps"): + model.blocks[0].proj(torch.randn(2, 4)) + + +def test_ffmpeg_command_cover_scales_and_caps_at_121_frames(): + command = draft_video.ffmpeg_decode_command("clip.mp4", 1344, 768) + video_filter = command[command.index("-vf") + 1] + assert "fps=24" in video_filter and "scale=1344:768:force_original_aspect_ratio=increase" in video_filter + assert "crop=1344:768" in video_filter + assert command[command.index("-frames:v") + 1] == "121" + + +def test_draft_pixels_pad_the_last_frame_to_124(monkeypatch): + width, height = draft_video.DRAFT_SIZES["576p"] + frames = np.arange(121, dtype=np.uint8).reshape(121, 1, 1, 1) * np.ones((121, height, width, 3), np.uint8) + + class _Done: + stdout = frames.tobytes() + + monkeypatch.setattr(draft_video.subprocess, "run", lambda *a, **k: _Done()) + pixels = draft_video.load_draft_pixels("clip.mp4", "576p") + assert pixels.shape == (1, 3, 124, height, width) and pixels.dtype == torch.uint8 + assert torch.equal(pixels[:, :, 120], pixels[:, :, 123]) + + +def test_draft_pixels_reject_truncated_decode(monkeypatch): + + class _Short: + stdout = b"\x00" * 10 + + monkeypatch.setattr(draft_video.subprocess, "run", lambda *a, **k: _Short()) + with pytest.raises(ValueError, match="decoded to 10 bytes"): + draft_video.load_draft_pixels("clip.mp4") + + +class _ConstantVelocityTransformer: + """Predicts a constant video velocity and records what the stage feeds it.""" + + def __init__(self, velocity): + self.velocity = velocity + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + rows = kwargs["hidden_states"].shape[1] + audio_rows = kwargs["audio_hidden_states"].shape[1] + return torch.full((1, rows, 8), self.velocity), torch.zeros(1, audio_rows, 4) + + +class _RecordingCurve: + + def __init__(self): + self.timesteps = [] + self.attached = 0 + self.targets = [] + + def attach(self, transformer): + self.attached += 1 + self.targets.append(transformer) + + def set_timesteps(self, unique): + self.timesteps.append(unique.clone()) + + +def _curve_model(out_features): + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for _ in range(NUM_CURVE_BRANCHES): + block = torch.nn.Module() + block.proj = _Projection(out_features) + model.blocks.append(block) + return model + + +def test_sharp_branches_reach_a_reloaded_transformer_instance(tmp_path): + _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES) + curve = SharpCurveBranches(tmp_path) + curve.set_timesteps(torch.tensor([0.1, 0.9])) + x = torch.randn(2, 4) + for model in (_curve_model(6), _curve_model(6)): + curve.attach(model) + out, _ = model.blocks[0].proj(x) + assert not torch.allclose(out, model.blocks[0].proj.linear(x)) + + +def test_stage_attaches_the_branches_to_each_materialized_lazy_transformer(monkeypatch): + from fastvideo.pipelines.lazy_module import LazyModule + + lazy = LazyModule("transformer", lambda: _ConstantVelocityTransformer(0.5)) + first = _denoise(monkeypatch, steps=2, transformer=lazy)[2] + loaded_first = first.targets[-1] + assert lazy.release() + second = _denoise(monkeypatch, steps=2, transformer=lazy, curve=first)[2] + + assert not isinstance(loaded_first, LazyModule) and not isinstance(second.targets[-1], LazyModule) + assert second.targets[-1] is not loaded_first + + +def _denoise(monkeypatch, steps, velocity=0.5, transformer=None, curve=None, fastvideo_args=None): + from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY + from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) + layout = build_packed_sequence(torch.ones(3, dtype=torch.long), 2, 2, 2, 2, (1, 2, 2)) + batch = ForwardBatch(data_type="video") + batch.num_inference_steps = steps + batch.prompt_embeds = [torch.zeros(1, 3, 4)] + batch.latents = torch.zeros(layout.video_indices.numel(), 8) + batch.audio_latents = torch.ones(layout.audio_indices.numel(), 4) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout + transformer = transformer if transformer is not None else _ConstantVelocityTransformer(velocity) + curve = curve if curve is not None else _RecordingCurve() + stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None) + return stage.forward(batch, fastvideo_args), transformer, curve + + +def test_euler_steps_integrate_the_constant_velocity_over_sigma_zero_to_start(monkeypatch): + batch, transformer, curve = _denoise(monkeypatch, steps=3, velocity=0.5) + assert len(transformer.calls) == 3 and curve.attached == 1 + assert torch.allclose(batch.latents, torch.full_like(batch.latents, 0.909375 * 0.5)) + + +def test_video_clock_is_one_minus_sigma_and_audio_is_held_clean(monkeypatch): + _, transformer, curve = _denoise(monkeypatch, steps=3) + for call, unique, sigma in zip(transformer.calls, curve.timesteps, [0.909375, 6 / 7, 2 / 3]): + assert torch.allclose(unique.sort().values, torch.tensor([1.0 - sigma, 1.0]).sort().values.float(), atol=1e-6) + assert call["hidden_states"].dtype == torch.bfloat16 and call["audio_hidden_states"].dtype == torch.bfloat16 + assert torch.equal(call["audio_hidden_states"].float(), torch.ones_like(call["audio_hidden_states"]).float()) + + +def test_two_step_request_runs_two_forwards(monkeypatch): + batch, transformer, _ = _denoise(monkeypatch, steps=2, velocity=1.0) + assert len(transformer.calls) == 2 + assert torch.allclose(batch.latents, torch.full_like(batch.latents, 0.7), atol=1e-6) + + +def test_registry_resolves_the_release_to_the_regenerate_pipeline(tmp_path): + from fastvideo.configs.pipelines.open_h3_regenerate_2k import OpenH3Regenerate2KPipelineConfig + from fastvideo.registry import get_model_info, get_pipeline_config_cls_from_name, get_preset_selection + + hub_id = "Efficient-Large-Model/Open-H3-Regenerate-2K" + (tmp_path / "lora").mkdir() + (tmp_path / "lora" / "open_h3_regenerate_2k_lora.safetensors").write_bytes(b"") + for path in (hub_id, str(tmp_path)): + assert get_pipeline_config_cls_from_name(path) is OpenH3Regenerate2KPipelineConfig + assert get_model_info(path).pipeline_cls.__name__ == "OpenH3Regenerate2KPipeline" + preset, family = get_preset_selection(hub_id) + assert (preset, family) == ("open_h3_regenerate_2k", "open_h3_regenerate_2k") + + +class _OffloadedTransformer(_ConstantVelocityTransformer): + + def __init__(self, velocity, fail=False): + super().__init__(velocity) + self.fail = fail + self.moves = [] + + def to(self, device): + self.moves.append(str(device)) + return self + + def __call__(self, **kwargs): + if self.fail: + raise RuntimeError("forward failed") + return super().__call__(**kwargs) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_full_cpu_offload_moves_the_transformer_back_to_cpu_after_the_refine(monkeypatch, fail): + from types import SimpleNamespace + + transformer = _OffloadedTransformer(0.5, fail=fail) + args = SimpleNamespace(dit_cpu_offload=True, dit_layerwise_offload=False, use_fsdp_inference=False) + if fail: + with pytest.raises(RuntimeError, match="forward failed"): + _denoise(monkeypatch, steps=2, transformer=transformer, fastvideo_args=args) + else: + _denoise(monkeypatch, steps=2, transformer=transformer, fastvideo_args=args) + + assert transformer.moves == ["cpu", "cpu"] # to the stage device (CPU in this test), then back to the host + + +@pytest.mark.parametrize("variable, value", [("FASTVIDEO_H3_ADALN_CACHE", True), + ("FASTVIDEO_H3_ADALN_TABLE", "/tmp/adaln_tables.pt")]) +def test_pipeline_refuses_an_adaln_cache_that_would_skip_the_sharp_branches(variable, value): + from fastvideo import envs + from fastvideo.pipelines.basic.open_h3_regenerate_2k.open_h3_regenerate_2k_pipeline import refuse_adaln_cache + + refuse_adaln_cache() + with getattr(envs, variable).override(value), pytest.raises(ValueError, match="Sharp curve branches"): + refuse_adaln_cache() diff --git a/tests/local_tests/open_h3_regenerate_2k/test_upscaler_parity.py b/tests/local_tests/open_h3_regenerate_2k/test_upscaler_parity.py new file mode 100644 index 000000000..b4bf1b155 --- /dev/null +++ b/tests/local_tests/open_h3_regenerate_2k/test_upscaler_parity.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in parity of the ported latent upscaler against the author's ComfyUI node. + +Needs the release weights (OPEN_H3_REGEN_WEIGHTS) and a checkout of NVlabs/Sana on the ``sol-engine`` branch +(OPEN_H3_REGEN_REF, the models/Open-H3-Regenerate-2K directory). GPU only: the node runs the network in bf16. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +import types +from pathlib import Path + +import pytest +import torch + +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler, upscale_latent + +WEIGHTS = os.environ.get("OPEN_H3_REGEN_WEIGHTS") +REFERENCE = os.environ.get("OPEN_H3_REGEN_REF") +CHECKPOINT = "latent_upsampler/minimax_h3_latent_upscaler_3d_conv_v1_bf16.safetensors" + +pytestmark = pytest.mark.skipif(not (WEIGHTS and REFERENCE and torch.cuda.is_available()), + reason="needs OPEN_H3_REGEN_WEIGHTS, OPEN_H3_REGEN_REF and a GPU") + + +def _load_author_node(checkpoint: str): + folder = types.ModuleType("folder_paths") + folder.folder_names_and_paths = {"latent_upscale_models": ([], set())} + folder.get_full_path_or_raise = lambda category, name: checkpoint + sys.modules["folder_paths"] = folder + sys.modules["comfy.model_management"] = None + sys.modules["comfy_api.latest"] = None + spec = importlib.util.spec_from_file_location("author_upscaler", + Path(REFERENCE) / "open_h3" / "upscaler_node.py") + node = importlib.util.module_from_spec(spec) + spec.loader.exec_module(node) + node.io.NodeOutput = lambda value: (value, ) + return node + + +@pytest.mark.parametrize("shape", [(24, 37, 48, 84), (24, 37, 36, 64)]) +def test_upscaler_matches_the_author_node(shape): + checkpoint = str(Path(WEIGHTS) / CHECKPOINT) + node = _load_author_node(checkpoint) + torch.manual_seed(0) + latent = torch.randn(shape) + mode = {"mode": node.UpscaleMode.TARGET_DIMENSIONS, "width": 2560, "height": 1440} + with torch.inference_mode(): + (expected, ) = node.MinimaxH3LatentUpscaler3D.execute({"samples": latent[None].float()}, + "upscaler", + mode, + align=32, + enable_temporal_chunking=False, + force_unload=False, + device="cuda", + precision="bf16") + model = MiniMaxH3LatentUpscaler.from_checkpoint(checkpoint).to("cuda", torch.bfloat16) + actual = upscale_latent(model, latent.cuda(), (1440, 2560)).cpu() + assert torch.equal(actual, expected["samples"].float()[0])