From 24d94c13a8efb48d7f0e29b5cc2523a638ead097 Mon Sep 17 00:00:00 2001 From: Davids048 Date: Mon, 31 Aug 2026 23:55:07 +0000 Subject: [PATCH 1/2] Dreamverse: add FastH3 support --- apps/dreamverse/README.md | 25 +- apps/dreamverse/arch.md | 13 +- apps/dreamverse/dreamverse/config.py | 23 +- .../dreamverse/generation_contracts.py | 46 +++ .../dreamverse/generation_worker.py | 96 ++++++ apps/dreamverse/dreamverse/gpu_pool.py | 27 +- .../dreamverse/minimax_h3_generation.py | 297 ++++++++++++++++++ .../dreamverse/session/controller.py | 4 +- .../dreamverse/tests/test_config.py | 40 ++- .../dreamverse/tests/test_gpu_pool.py | 29 +- .../tests/test_minimax_h3_generation.py | 247 +++++++++++++++ .../dreamverse/dreamverse/video_generation.py | 26 +- fastvideo/models/loader/lora_patch.py | 5 +- fastvideo/tests/loader/test_lora_patch.py | 16 +- 14 files changed, 848 insertions(+), 46 deletions(-) create mode 100644 apps/dreamverse/dreamverse/generation_contracts.py create mode 100644 apps/dreamverse/dreamverse/generation_worker.py create mode 100644 apps/dreamverse/dreamverse/minimax_h3_generation.py create mode 100644 apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py diff --git a/apps/dreamverse/README.md b/apps/dreamverse/README.md index 783fa58e54..e464224aec 100644 --- a/apps/dreamverse/README.md +++ b/apps/dreamverse/README.md @@ -97,13 +97,33 @@ dreamverse-server --port 8009 dreamverse-mock-server --port 8009 ``` +### Run Dreamverse with FastH3 + +Select the VSA data-free FastH3 Preview profile when you start the backend: + +```bash +DREAMVERSE_MODEL_ID=fast-h3 dreamverse-server --port 8009 +``` + +The `fast-h3` profile uses four visible GPUs by default. It loads the `MiniMaxAI/MiniMax-H3` base checkpoint and the +`vsa-datafree/adapter_model.safetensors` adapter from +`FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA`. Each request generates a 124-frame, 768×1344 video with +synchronized audio and five sigma-grid points. Dreamverse uses the last frame of each segment as first-frame +conditioning for the following segment. + +Set `CUDA_VISIBLE_DEVICES` when you need to choose the four physical GPUs: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 DREAMVERSE_MODEL_ID=fast-h3 dreamverse-server --port 8009 +``` + > **Expect a slow first boot.** With `torch.compile` and startup warmup enabled > (the default), the backend compiles the segment 1 and segment 2 inference > paths before it reports ready — this can take **tens of minutes on a cold > cache**, regardless of how you deploy (local, server, Docker, or Modal). > `/healthz` responds as soon as the process is up; `/readyz` stays `503` until -> warmup finishes. For a faster, uncompiled startup while testing, set -> `FASTVIDEO_ENABLE_STARTUP_WARMUP=0` before starting the backend. +> warmup finishes. To defer compilation until the first generated request while +> testing, set `FASTVIDEO_ENABLE_STARTUP_WARMUP=0` before starting the backend. ## Frontend Setup @@ -219,6 +239,7 @@ selection, and mock-server behavior: pytest apps/dreamverse/dreamverse/tests/test_config.py \ apps/dreamverse/dreamverse/tests/test_entrypoints.py \ apps/dreamverse/dreamverse/tests/test_gpu_pool.py \ + apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py \ apps/dreamverse/dreamverse/tests/test_mock_server.py -q ``` diff --git a/apps/dreamverse/arch.md b/apps/dreamverse/arch.md index 9a527a2541..64b0aff7b0 100644 --- a/apps/dreamverse/arch.md +++ b/apps/dreamverse/arch.md @@ -139,7 +139,18 @@ session. - startup warmup - user join/leave commands - `USER_STEP` execution for each segment -- continuation state between segments +- generation-command routing and stream-result delivery + +Model generation has a separate ownership boundary inside each GPU process: + +- `apps/dreamverse/dreamverse/generation_worker.py` selects the backend that the active model profile declares and owns + the backend lifecycle. +- `apps/dreamverse/dreamverse/video_generation.py` owns LTX-2 generator configuration, video and audio continuation, and + runtime LoRA application. +- `apps/dreamverse/dreamverse/minimax_h3_generation.py` owns the VSA data-free FastH3 adapter, FastH3 generator and + request configuration, and last-frame continuation through MiniMax H3 first-frame conditioning. +- `apps/dreamverse/dreamverse/generation_contracts.py` defines the decoded media and stream-trimming result that both + model backends return to `apps/dreamverse/dreamverse/gpu_pool.py`. `apps/dreamverse/dreamverse/prompt_enhancer.py` manages: diff --git a/apps/dreamverse/dreamverse/config.py b/apps/dreamverse/dreamverse/config.py index 4802deee03..0c65f70c42 100644 --- a/apps/dreamverse/dreamverse/config.py +++ b/apps/dreamverse/dreamverse/config.py @@ -1,5 +1,6 @@ import os from pathlib import Path +from typing import cast _REPO_ROOT = Path(__file__).resolve().parents[1] _SERVER_ROOT = Path(__file__).resolve().parent @@ -55,16 +56,34 @@ def _resolve_frontend_static_dir_candidates() -> tuple[str, ...]: MODEL_REGISTRY = { "fast-ltx2": { "name": "FastLTX2", + "generation_backend": "ltx2", + "default_sp_size": 1, "model_path": "FastVideo/LTX2-Distilled-Diffusers", "config_model_path": "FastVideo/LTX2-Distilled-Diffusers", "lora_repo": "FastVideo/LTX2-OmniNFT-LoRA", }, "fast-ltx23": { "name": "FastLTX23", + "generation_backend": "ltx2", + "default_sp_size": 1, "model_path": "FastVideo/LTX-2.3-Distilled-Diffusers", "config_model_path": "FastVideo/LTX-2.3-Distilled-Diffusers", "lora_repo": "FastVideo/LTX-2.3-OmniNFT-LoRA", }, + "fast-h3": { + "name": "FastH3", + "generation_backend": "minimax_h3", + "default_sp_size": 4, + "model_path": "MiniMaxAI/MiniMax-H3", + "adapter_repo": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", + "adapter_filename": "vsa-datafree/adapter_model.safetensors", + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "height": 768, + "width": 1344, + "num_frames": 124, + "num_inference_steps": 5, + "seed": 1000, + }, } DEFAULT_MODEL_ID = "fast-ltx2" @@ -171,7 +190,7 @@ def _optional_env(*names: str) -> str | None: DEVTOOLS_ENABLED = _env_bool("FASTVIDEO_ENABLE_DEVTOOLS", False) PROMPT_SAFETY_ENABLED = _env_bool("FASTVIDEO_ENABLE_PROMPT_SAFETY", False) DREAMVERSE_MAX_AUTOTUNE = _env_bool("DREAMVERSE_MAX_AUTOTUNE", True) -DREAMVERSE_SP_SIZE = max(1, _env_int("DREAMVERSE_SP_SIZE", 1)) +DREAMVERSE_SP_SIZE = max(1, _env_int("DREAMVERSE_SP_SIZE", cast(int, MODEL_CONFIG["default_sp_size"]))) DREAMVERSE_MODEL_PATH = (os.getenv("DREAMVERSE_MODEL_PATH", "").strip() or None) if DREAMVERSE_MODEL_PATH: @@ -213,7 +232,7 @@ def _resolve_lora_spec(spec: str) -> str | None: if not spec: return None if spec.lower() == "omninft": - return MODEL_CONFIG.get("lora_repo") + return cast(str | None, MODEL_CONFIG.get("lora_repo")) if spec.lower() in AVAILABLE_LORAS: return AVAILABLE_LORAS[spec.lower()]["repo"] return spec diff --git a/apps/dreamverse/dreamverse/generation_contracts.py b/apps/dreamverse/dreamverse/generation_contracts.py new file mode 100644 index 0000000000..782b23304d --- /dev/null +++ b/apps/dreamverse/dreamverse/generation_contracts.py @@ -0,0 +1,46 @@ +"""Shared contract between DreamVerse generation backends and GPU workers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Protocol + + +@dataclass +class StepResult: + """Decoded media and stream-trimming metadata for one DreamVerse segment.""" + + frames: list + audio: Any + audio_sample_rate: int | None + timings: dict[str, float] + head_trim_frames: int + head_trim_audio_frames: int + + +class GenerationBackend(Protocol): + """Model-owned generation operations used by one GPU worker process.""" + + def initialize(self, model_config: dict | None = None) -> None: + ... + + def shutdown(self) -> None: + ... + + def clear_conditioning(self) -> None: + ... + + def generate_step( + self, + prompt: str, + segment_idx: int, + image_path: str | None, + reset_conditioning: bool, + ) -> StepResult: + ... + + def warmup(self, prompt: str) -> dict[str, float]: + ... + + def apply_lora_stack(self, stack: list[tuple[str, float]]) -> tuple[str | None, str | None]: + ... diff --git a/apps/dreamverse/dreamverse/generation_worker.py b/apps/dreamverse/dreamverse/generation_worker.py new file mode 100644 index 0000000000..e06a9aa179 --- /dev/null +++ b/apps/dreamverse/dreamverse/generation_worker.py @@ -0,0 +1,96 @@ +"""Select and own one model-specific generation backend per GPU process.""" + +from __future__ import annotations + +from dreamverse.config import MODEL_CONFIG +from dreamverse.generation_contracts import GenerationBackend, StepResult + + +def _create_generation_backend(backend_name: str, gpu_id: int) -> GenerationBackend: + """Construct the backend that owns the selected model family's behavior.""" + if backend_name == "ltx2": + from dreamverse.video_generation import LTX2GenerationBackend + + return LTX2GenerationBackend(gpu_id) + if backend_name == "minimax_h3": + from dreamverse.minimax_h3_generation import MiniMaxH3GenerationBackend + + return MiniMaxH3GenerationBackend(gpu_id) + raise ValueError(f"Unsupported DreamVerse generation backend: {backend_name!r}") + + +class VideoGenerationWorker: + """Delegate GPU lifecycle and generation calls to the active model backend.""" + + def __init__(self, gpu_id: int): + self.gpu_id = gpu_id + self.model_config: dict = dict(MODEL_CONFIG) + self.backend_name: str | None = None + self.backend: GenerationBackend | None = None + + def initialize(self, model_config: dict | None = None) -> None: + """Load the requested model through its generation backend. + + Model selection belongs here so the GPU process and streaming layers + use one stable media contract without importing model-specific code. + """ + requested_model_config = dict(model_config) if model_config is not None else dict(self.model_config) + backend_name = requested_model_config.get("generation_backend") + if not isinstance(backend_name, str) or not backend_name: + raise ValueError("DreamVerse model configuration requires `generation_backend`.") + + candidate_backend = self.backend + if candidate_backend is None or self.backend_name != backend_name: + if candidate_backend is not None: + candidate_backend.shutdown() + candidate_backend = _create_generation_backend(backend_name, self.gpu_id) + + try: + candidate_backend.initialize(requested_model_config) + except Exception: + try: + candidate_backend.shutdown() + except Exception as shutdown_error: + print(f"[GPU {self.gpu_id}] Backend cleanup after initialization failure: {shutdown_error}") + self.backend = None + self.backend_name = None + raise + + self.model_config = requested_model_config + self.backend = candidate_backend + self.backend_name = backend_name + + def _require_backend(self) -> GenerationBackend: + """Return the initialized backend or fail before processing a command.""" + if self.backend is None: + raise RuntimeError("Generation backend is not initialized.") + return self.backend + + def shutdown(self) -> None: + """Release model resources owned by the selected backend.""" + if self.backend is not None: + self.backend.shutdown() + + def clear_conditioning(self) -> None: + self._require_backend().clear_conditioning() + + def generate_step( + self, + prompt: str, + segment_idx: int, + image_path: str | None, + reset_conditioning: bool, + ) -> StepResult: + """Generate one segment through the selected model backend.""" + return self._require_backend().generate_step( + prompt, + segment_idx, + image_path, + reset_conditioning, + ) + + def warmup(self, prompt: str) -> dict[str, float]: + return self._require_backend().warmup(prompt) + + def apply_lora_stack(self, stack: list[tuple[str, float]]) -> tuple[str | None, str | None]: + return self._require_backend().apply_lora_stack(stack) diff --git a/apps/dreamverse/dreamverse/gpu_pool.py b/apps/dreamverse/dreamverse/gpu_pool.py index 3a3b98e173..a524fe16c7 100644 --- a/apps/dreamverse/dreamverse/gpu_pool.py +++ b/apps/dreamverse/dreamverse/gpu_pool.py @@ -12,7 +12,7 @@ from multiprocessing import Process, Queue from dreamverse.config import ( - DEFAULT_MODEL_ID, + ACTIVE_MODEL_ID, DREAMVERSE_SP_SIZE, MODEL_REGISTRY, STARTUP_WARMUP_ENABLED, @@ -54,7 +54,7 @@ def _parse_requested_gpu_limit() -> int | None: raw_value = os.getenv("FASTVIDEO_GPU_COUNT", "").strip().lower() if not raw_value: - return 1 + return DREAMVERSE_SP_SIZE if raw_value == "all": return None try: @@ -164,12 +164,12 @@ def gpu_worker_process( os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" - from dreamverse.video_generation import VideoGenerationWorker + from dreamverse.generation_worker import VideoGenerationWorker worker = VideoGenerationWorker(gpu_id) def event_loop(first_cmd: Command = None): - """Blocking event loop for LTX2; dispatches user commands.""" + """Block on generation commands after the model is initialized.""" print(f"[GPU {gpu_id}] Entering event loop") def handle_command(cmd: Command): @@ -435,7 +435,7 @@ def __init__(self, gpu_id: int, cuda_device: str): self._response_reader_task: asyncio.Task | None = None self._active: bool = False self._reader_lock: asyncio.Lock | None = None - self.current_model_id: str = DEFAULT_MODEL_ID + self.current_model_id: str | None = ACTIVE_MODEL_ID self.shared_stream_buffer = None self.shared_stream_buffer_size = SHARED_STREAM_BUFFER_BYTES @@ -690,7 +690,7 @@ def unregister_stream_queue(self, user_id: str) -> None: async def join_user(self, user_id: str, model_id: str = None) -> JoinAck: """Add a user to this GPU.""" if model_id is None: - model_id = DEFAULT_MODEL_ID + model_id = ACTIVE_MODEL_ID # Reload model if a different one is requested if model_id != self.current_model_id and model_id in MODEL_REGISTRY: @@ -705,16 +705,23 @@ async def join_user(self, user_id: str, model_id: str = None) -> JoinAck: self.connected_users.clear() model_config = MODEL_REGISTRY[model_id] - reload_response = await self._send_command(Command(CommandType.RELOAD_MODEL, - payload=ReloadModelPayload(model_config=model_config), - user_id="__reload__"), - timeout=600.0) + try: + reload_response = await self._send_command(Command( + CommandType.RELOAD_MODEL, + payload=ReloadModelPayload(model_config=model_config), + user_id="__reload__"), + timeout=600.0) + except Exception: + self.current_model_id = None + raise match reload_response: case ReloadAck(): pass case WorkerError(message=msg): + self.current_model_id = None raise RuntimeError(f"Model reload failed: {msg}") case _: + self.current_model_id = None raise RuntimeError(f"Unexpected reload response: " f"{type(reload_response).__name__}") diff --git a/apps/dreamverse/dreamverse/minimax_h3_generation.py b/apps/dreamverse/dreamverse/minimax_h3_generation.py new file mode 100644 index 0000000000..6c99b40b72 --- /dev/null +++ b/apps/dreamverse/dreamverse/minimax_h3_generation.py @@ -0,0 +1,297 @@ +"""FastH3 model lifecycle and first-frame continuation for DreamVerse.""" + +from __future__ import annotations + +import gc +import os +import time +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +from dreamverse.config import DREAMVERSE_SP_SIZE +from dreamverse.generation_contracts import StepResult + +if TYPE_CHECKING: + from PIL.Image import Image + + +def _required_config_str(model_config: dict, field_name: str) -> str: + """Read one required non-empty string from a DreamVerse model profile.""" + value = model_config.get(field_name) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"FastH3 model configuration requires `{field_name}`.") + return value.strip() + + +class MiniMaxH3GenerationBackend: + """Run the VSA data-free FastH3 adapter and retain one continuation frame.""" + + def __init__(self, gpu_id: int): + self.gpu_id = gpu_id + self.generator: Any | None = None + self.model_config: dict = {} + self.continuation_image: Image | None = None + + def _gpu_mem(self) -> str: + allocated_gib = torch.cuda.memory_allocated() / 1024**3 + reserved_gib = torch.cuda.memory_reserved() / 1024**3 + return f"alloc={allocated_gib:.2f}GiB, reserved={reserved_gib:.2f}GiB" + + @staticmethod + def _configure_environment(attention_backend: str) -> None: + """Apply the fixed boot-time switches from the FastH3 reference recipe.""" + os.environ.update({ + "FASTVIDEO_ATTENTION_BACKEND": attention_backend, + "FASTVIDEO_FA4": "1", + "FASTVIDEO_MINIMAX_H3_FUSIONS": "all", + "FASTVIDEO_VSA_SM100A": "0", + }) + os.environ.pop("FASTVIDEO_INFERENCE_TORCH_COMPILE", None) + + def initialize(self, model_config: dict | None = None) -> None: + """Download the fixed Preview adapter and load the FastH3 generator. + + The model profile owns the base checkpoint, adapter file, attention + backend, and generation geometry. The backend translates that profile + into FastVideo's typed generator configuration. + """ + if model_config is not None: + self.model_config = dict(model_config) + if not self.model_config: + raise ValueError("FastH3 initialization requires a model configuration.") + + if self.generator is not None: + self.generator.shutdown() + self.generator = None + gc.collect() + torch.cuda.empty_cache() + + self.clear_conditioning() + model_path = _required_config_str(self.model_config, "model_path") + adapter_repo = _required_config_str(self.model_config, "adapter_repo") + adapter_filename = _required_config_str(self.model_config, "adapter_filename") + attention_backend = _required_config_str(self.model_config, "attention_backend") + self._configure_environment(attention_backend) + + from huggingface_hub import hf_hub_download + + from fastvideo import VideoGenerator + from fastvideo.api import ( + CompileConfig, + ComponentConfig, + EngineConfig, + GeneratorConfig, + OffloadConfig, + ParallelismConfig, + PipelineSelection, + ) + + adapter_path = hf_hub_download(repo_id=adapter_repo, filename=adapter_filename) + experimental = { + "attention_backend": attention_backend, + "inference_torch_compile": attention_backend == "FLASH_ATTN", + "vae_parallel_decode": True, + "vae_parallel_decode_strategy": "gather", + } + if attention_backend == "VIDEO_SPARSE_ATTN_H3": + experimental.update({ + "VSA_sparsity": 0.9, + "VSA_tile_size": 64, + }) + generator_config = GeneratorConfig( + model_path=model_path, + pipeline=PipelineSelection( + components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0), + experimental=experimental, + ), + engine=EngineConfig( + num_gpus=DREAMVERSE_SP_SIZE, + parallelism=ParallelismConfig(tp_size=1, sp_size=DREAMVERSE_SP_SIZE), + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=True, + image_encoder=True, + vae=True, + pin_cpu_memory=True, + ), + compile=CompileConfig(enabled=False, vae_enabled=True), + use_fsdp_inference=False, + ), + ) + + print(f"[GPU {self.gpu_id}] Loading FastH3 model: {model_path}") + print(f"[GPU {self.gpu_id}] FastH3 adapter: {adapter_repo}/{adapter_filename}") + print(f"[GPU {self.gpu_id}] Before model load: {self._gpu_mem()}") + self.generator = VideoGenerator.from_config(generator_config) + print(f"[GPU {self.gpu_id}] FastH3 loaded: {self._gpu_mem()} (warmup pending)") + + def shutdown(self) -> None: + """Release the FastVideo generator and cached continuation image.""" + self.clear_conditioning() + if self.generator is not None: + self.generator.shutdown() + self.generator = None + + def clear_conditioning(self) -> None: + """Release the first-frame image retained for the next segment.""" + if self.continuation_image is not None: + self.continuation_image.close() + self.continuation_image = None + + @staticmethod + def _load_rgb_image(image_path: str) -> Image: + """Load an image into an independent RGB buffer with no open file handle.""" + from PIL import Image + + with Image.open(image_path) as image: + return image.convert("RGB").copy() + + def _select_conditioning_image( + self, + segment_idx: int, + image_path: str | None, + reset_conditioning: bool, + ) -> tuple[Image | None, bool]: + """Select the initial upload or retained last frame for one segment.""" + if reset_conditioning: + self.clear_conditioning() + if segment_idx > 1 and self.continuation_image is not None: + return self.continuation_image.copy(), True + if segment_idx > 1 and not reset_conditioning: + raise RuntimeError(f"FastH3 segment {segment_idx} requires a retained continuation frame.") + if segment_idx == 1 and image_path: + return self._load_rgb_image(image_path), False + return None, False + + def _build_request(self, prompt: str, conditioning_image: Image | None): + """Build the typed FastVideo request owned by the FastH3 profile.""" + from fastvideo.api import GenerationRequest, InputConfig, OutputConfig, SamplingConfig + + return GenerationRequest( + prompt=prompt, + negative_prompt="", + inputs=InputConfig(pil_image=conditioning_image), + sampling=SamplingConfig( + height=int(self.model_config["height"]), + width=int(self.model_config["width"]), + num_frames=int(self.model_config["num_frames"]), + fps=24, + num_inference_steps=int(self.model_config["num_inference_steps"]), + guidance_scale=1.0, + batch_cfg=False, + seed=int(self.model_config["seed"]), + ), + output=OutputConfig(save_video=False, return_frames=True), + ) + + def _save_continuation_frame(self, frames: list) -> None: + """Retain the last decoded frame as first-frame conditioning.""" + from PIL import Image + + self.clear_conditioning() + self.continuation_image = Image.fromarray(np.ascontiguousarray(frames[-1])).convert("RGB") + + def generate_step( + self, + prompt: str, + segment_idx: int, + image_path: str | None, + reset_conditioning: bool, + ) -> StepResult: + """Generate one synchronized FastH3 segment and retain its last frame. + + Later segments use MiniMax H3's first-frame-to-video path. The first + conditioned frame and its matching audio duration are trimmed before + streaming so adjacent segments do not duplicate media. + """ + if self.generator is None: + raise RuntimeError("FastH3 generator is not initialized.") + conditioning_image, uses_continuation = self._select_conditioning_image( + segment_idx, + image_path, + reset_conditioning, + ) + request = self._build_request(prompt, conditioning_image) + started = time.perf_counter() + try: + result = self.generator.generate(request) + finally: + if conditioning_image is not None: + conditioning_image.close() + torch.cuda.synchronize() + generation_ms = (time.perf_counter() - started) * 1000.0 + + if isinstance(result, list): + raise RuntimeError("FastH3 returned multiple results for one DreamVerse segment.") + frames = result.frames + if not isinstance(frames, list) or not frames: + raise RuntimeError("FastH3 generation did not return decoded frames.") + audio = result.audio + audio_sample_rate = result.audio_sample_rate + if audio is not None and audio_sample_rate is None: + raise RuntimeError("FastH3 returned audio without an audio sample rate.") + + save_started = time.perf_counter() + self._save_continuation_frame(frames) + save_conditioning_ms = (time.perf_counter() - save_started) * 1000.0 + timings = { + "generation_ms": generation_ms, + "generation_time_ms": float(result.generation_time or 0.0) * 1000.0, + "save_conditioning_ms": save_conditioning_ms, + "e2e_latency_ms": (time.perf_counter() - started) * 1000.0, + } + trim_frames = 1 if uses_continuation else 0 + print(f"[GPU {self.gpu_id}] FastH3 segment {segment_idx}: " + f"{len(frames)} frames, gen={generation_ms:.0f}ms, " + f"save_conditioning={save_conditioning_ms:.0f}ms, " + f"e2e={timings['e2e_latency_ms']:.0f}ms") + return StepResult( + frames=frames, + audio=audio, + audio_sample_rate=audio_sample_rate, + timings=timings, + head_trim_frames=trim_frames, + head_trim_audio_frames=trim_frames, + ) + + def warmup(self, prompt: str) -> dict[str, float]: + """Compile the FastH3 text and first-frame paths before readiness.""" + warmup_prompt = (prompt or "").strip() + if not warmup_prompt: + raise RuntimeError("Startup warmup prompt must be non-empty.") + print(f"[GPU {self.gpu_id}] FastH3 startup warmup starting " + "(synthetic segments: text-to-video, first-frame-to-video)") + started = time.perf_counter() + text_result = self.generate_step( + warmup_prompt, + segment_idx=1, + image_path=None, + reset_conditioning=True, + ) + first_frame_result = self.generate_step( + warmup_prompt, + segment_idx=2, + image_path=None, + reset_conditioning=False, + ) + total_ms = (time.perf_counter() - started) * 1000.0 + self.clear_conditioning() + text_ms = float(text_result.timings.get("e2e_latency_ms", 0.0)) + first_frame_ms = float(first_frame_result.timings.get("e2e_latency_ms", 0.0)) + print(f"[GPU {self.gpu_id}] FastH3 startup warmup complete: " + f"text_to_video={text_ms:.0f}ms, " + f"first_frame_to_video={first_frame_ms:.0f}ms, " + f"total={total_ms:.0f}ms") + return { + "warmup_text_to_video_ms": text_ms, + "warmup_first_frame_to_video_ms": first_frame_ms, + "warmup_total_ms": total_ms, + } + + def apply_lora_stack(self, stack: list[tuple[str, float]]) -> tuple[str | None, str | None]: + """Reject runtime LoRA mutation because FastH3 uses one startup adapter.""" + del stack + raise RuntimeError("FastH3 uses its fixed startup adapter and does not support runtime LoRA changes.") diff --git a/apps/dreamverse/dreamverse/session/controller.py b/apps/dreamverse/dreamverse/session/controller.py index 352ca3e0c9..e3985938d8 100644 --- a/apps/dreamverse/dreamverse/session/controller.py +++ b/apps/dreamverse/dreamverse/session/controller.py @@ -30,7 +30,7 @@ from dreamverse.worker_ipc import MediaChunk, MediaComplete, MediaInit from dreamverse.config import ( - DEFAULT_MODEL_ID, + ACTIVE_MODEL_ID, GENERATION_SEGMENT_CAP, PROMPT_AUTO_SLEEP_MS, PROMPT_AUTO_TIMEOUT_MS, @@ -264,7 +264,7 @@ async def cancel_task(task: asyncio.Task | None): timeout_task = asyncio.create_task(session_timeout()) # Join the engine on this GPU. - await slot.join_user(client_id, model_id=DEFAULT_MODEL_ID) + await slot.join_user(client_id, model_id=ACTIVE_MODEL_ID) # Notify client they're connected to a GPU. await ws_send_json({ diff --git a/apps/dreamverse/dreamverse/tests/test_config.py b/apps/dreamverse/dreamverse/tests/test_config.py index 755bb3b40f..382794a28d 100644 --- a/apps/dreamverse/dreamverse/tests/test_config.py +++ b/apps/dreamverse/dreamverse/tests/test_config.py @@ -2,13 +2,14 @@ import importlib.util from pathlib import Path +from types import ModuleType import pytest SERVER_DIR = Path(__file__).resolve().parents[1] -def _load_config_module(): +def _load_config_module() -> ModuleType: spec = importlib.util.spec_from_file_location( "server_config_test_module", SERVER_DIR / "config.py", @@ -20,7 +21,7 @@ def _load_config_module(): return module -def _set_required_prompt_keys(monkeypatch): +def _set_required_prompt_keys(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("CEREBRAS_API_KEY", "cerebras-key") monkeypatch.setenv("GROQ_API_KEY", "groq-key") @@ -150,3 +151,38 @@ def test_config_rejects_invalid_prompt_provider(monkeypatch): with pytest.raises(RuntimeError, match="Invalid FASTVIDEO_PROMPT_PROVIDER"): _load_config_module() + + +def test_config_registers_vsa_datafree_fasth3_profile(monkeypatch): + """The FastH3 registry entry owns the complete fixed Preview recipe.""" + _set_required_prompt_keys(monkeypatch) + + module = _load_config_module() + + assert module.MODEL_REGISTRY["fast-h3"] == { + "name": "FastH3", + "generation_backend": "minimax_h3", + "default_sp_size": 4, + "model_path": "MiniMaxAI/MiniMax-H3", + "adapter_repo": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", + "adapter_filename": "vsa-datafree/adapter_model.safetensors", + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "height": 768, + "width": 1344, + "num_frames": 124, + "num_inference_steps": 5, + "seed": 1000, + } + + +def test_config_uses_fasth3_sequence_parallel_default(monkeypatch): + """Selecting FastH3 defaults DreamVerse to its four-GPU topology.""" + _set_required_prompt_keys(monkeypatch) + monkeypatch.setenv("DREAMVERSE_MODEL_ID", "fast-h3") + monkeypatch.delenv("DREAMVERSE_SP_SIZE", raising=False) + + module = _load_config_module() + + assert module.ACTIVE_MODEL_ID == "fast-h3" + assert module.MODEL_CONFIG["generation_backend"] == "minimax_h3" + assert module.DREAMVERSE_SP_SIZE == 4 diff --git a/apps/dreamverse/dreamverse/tests/test_gpu_pool.py b/apps/dreamverse/dreamverse/tests/test_gpu_pool.py index e676bd14cc..27a7638650 100644 --- a/apps/dreamverse/dreamverse/tests/test_gpu_pool.py +++ b/apps/dreamverse/dreamverse/tests/test_gpu_pool.py @@ -63,6 +63,14 @@ def test_get_available_gpus_defaults_to_first_visible_device(monkeypatch): assert gpu_pool.get_available_gpus() == [3] +def test_get_available_gpus_defaults_to_active_model_sequence_parallel_size(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4") + monkeypatch.delenv("FASTVIDEO_GPU_COUNT", raising=False) + monkeypatch.setattr(gpu_pool, "DREAMVERSE_SP_SIZE", 4) + + assert gpu_pool.get_available_gpus() == [0, 1, 2, 3] + + def test_get_available_gpus_rejects_invalid_gpu_count(monkeypatch): monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "zero") @@ -71,6 +79,23 @@ def test_get_available_gpus_rejects_invalid_gpu_count(monkeypatch): gpu_pool.get_available_gpus() +def test_join_user_failed_reload_marks_model_uninitialized(monkeypatch): + """A failed model reload forces the next join to reload a model.""" + slot = gpu_pool.GPUSlot(gpu_id=0, cuda_device="0") + slot.current_model_id = "fast-ltx2" + + async def fake_send_command(command, timeout): + del command, timeout + return gpu_pool.WorkerError(user_id="__reload__", message="load failed") + + monkeypatch.setattr(slot, "_send_command", fake_send_command) + + with pytest.raises(RuntimeError, match="Model reload failed"): + asyncio.run(slot.join_user("client-id", model_id="fast-h3")) + + assert slot.current_model_id is None + + def test_send_command_raises_on_worker_death(): """A worker that consumes a command and exits without replying must surface as RuntimeError via sentinel detection, not after the long @@ -92,9 +117,9 @@ def test_send_command_raises_on_worker_death(): ready = resp_q.get(timeout=30.0) assert ready == "READY" - async def runner(): + async def runner() -> None: slot = gpu_pool.GPUSlot(gpu_id=0, cuda_device="0") - slot.process = proc + slot.process = proc # type: ignore[assignment] slot.command_queue = cmd_q slot.response_queue = resp_q diff --git a/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py b/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py new file mode 100644 index 0000000000..c8d77c86eb --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_minimax_h3_generation.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import os +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +import dreamverse.generation_worker as generation_worker +from dreamverse.minimax_h3_generation import MiniMaxH3GenerationBackend + + +FASTH3_MODEL_CONFIG = { + "name": "FastH3", + "generation_backend": "minimax_h3", + "default_sp_size": 4, + "model_path": "MiniMaxAI/MiniMax-H3", + "adapter_repo": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", + "adapter_filename": "vsa-datafree/adapter_model.safetensors", + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "height": 768, + "width": 1344, + "num_frames": 124, + "num_inference_steps": 5, + "seed": 1000, +} + + +class _RecordingGenerator: + """Record typed requests and return small synchronized media fixtures.""" + + def __init__(self) -> None: + self.requests: list[Any] = [] + self.conditioning_pixels: list[np.ndarray | None] = [] + + def generate(self, request): + """Capture the request and return two tiny video frames with audio.""" + self.requests.append(request) + conditioning_image = request.inputs.pil_image + self.conditioning_pixels.append( + None if conditioning_image is None else np.asarray(conditioning_image).copy()) + frames = [ + np.full((2, 3, 3), 10, dtype=np.uint8), + np.full((2, 3, 3), 20, dtype=np.uint8), + ] + return SimpleNamespace( + frames=frames, + audio=np.zeros((2, 16), dtype=np.float32), + audio_sample_rate=44100, + generation_time=0.25, + ) + + +def test_initialize_builds_vsa_datafree_fasth3_generator(monkeypatch): + """Initialization translates the DreamVerse profile into typed FastVideo config.""" + from fastvideo import VideoGenerator + + captured = {} + fake_generator = SimpleNamespace(shutdown=lambda: None) + + def fake_from_config(config): + captured["config"] = config + return fake_generator + + def fake_download(**kwargs): + captured["download"] = kwargs + return f"/models/{kwargs['filename']}" + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download) + monkeypatch.setattr(VideoGenerator, "from_config", fake_from_config) + monkeypatch.setattr("dreamverse.minimax_h3_generation.DREAMVERSE_SP_SIZE", 4) + monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "test-attention") + monkeypatch.setenv("FASTVIDEO_FA4", "0") + monkeypatch.setenv("FASTVIDEO_MINIMAX_H3_FUSIONS", "0") + monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1") + monkeypatch.setenv("FASTVIDEO_INFERENCE_TORCH_COMPILE", "1") + + backend = MiniMaxH3GenerationBackend(gpu_id=0) + monkeypatch.setattr(backend, "_gpu_mem", lambda: "alloc=0.00GiB, reserved=0.00GiB") + backend.initialize(FASTH3_MODEL_CONFIG) + + config = captured["config"] + assert captured["download"] == { + "repo_id": "FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", + "filename": "vsa-datafree/adapter_model.safetensors", + } + assert config.model_path == "MiniMaxAI/MiniMax-H3" + assert config.pipeline.components.lora_path.endswith("vsa-datafree/adapter_model.safetensors") + assert config.pipeline.components.lora_strength == 1.0 + assert config.pipeline.experimental == { + "attention_backend": "VIDEO_SPARSE_ATTN_H3", + "inference_torch_compile": False, + "vae_parallel_decode": True, + "vae_parallel_decode_strategy": "gather", + "VSA_sparsity": 0.9, + "VSA_tile_size": 64, + } + assert config.engine.num_gpus == 4 + assert config.engine.parallelism.tp_size == 1 + assert config.engine.parallelism.sp_size == 4 + assert config.engine.offload.dit is False + assert config.engine.offload.dit_layerwise is False + assert config.engine.offload.text_encoder is True + assert config.engine.offload.vae is True + assert config.engine.compile.vae_enabled is True + assert config.engine.use_fsdp_inference is False + assert os.environ["FASTVIDEO_ATTENTION_BACKEND"] == "VIDEO_SPARSE_ATTN_H3" + assert os.environ["FASTVIDEO_FA4"] == "1" + assert os.environ["FASTVIDEO_MINIMAX_H3_FUSIONS"] == "all" + assert os.environ["FASTVIDEO_VSA_SM100A"] == "0" + assert "FASTVIDEO_INFERENCE_TORCH_COMPILE" not in os.environ + + +def test_initialize_selects_declared_generation_backend(monkeypatch): + """The GPU worker constructs the backend that the active model profile declares.""" + from unittest.mock import Mock + + selected_backend = Mock() + monkeypatch.setattr( + generation_worker, + "_create_generation_backend", + lambda backend_name, gpu_id: selected_backend, + ) + worker = generation_worker.VideoGenerationWorker(gpu_id=3) + + worker.initialize(FASTH3_MODEL_CONFIG) + + assert worker.backend_name == "minimax_h3" + assert worker.backend is selected_backend + selected_backend.initialize.assert_called_once_with(FASTH3_MODEL_CONFIG) + + +def test_initialize_failure_clears_backend_ownership(monkeypatch): + """A failed family change leaves the GPU worker explicitly uninitialized.""" + ltx_backend = SimpleNamespace(initialize=lambda config: None, shutdown=lambda: None) + + def fail_initialize(config): + del config + raise RuntimeError("load failed") + + fasth3_backend = SimpleNamespace( + initialize=fail_initialize, + shutdown=lambda: None, + ) + backends = { + "ltx2": ltx_backend, + "minimax_h3": fasth3_backend, + } + monkeypatch.setattr( + generation_worker, + "_create_generation_backend", + lambda backend_name, gpu_id: backends[backend_name], + ) + worker = generation_worker.VideoGenerationWorker(gpu_id=3) + worker.initialize({"generation_backend": "ltx2"}) + + with pytest.raises(RuntimeError, match="load failed"): + worker.initialize(FASTH3_MODEL_CONFIG) + + assert worker.backend is None + assert worker.backend_name is None + assert worker.model_config == {"generation_backend": "ltx2"} + + +def test_generate_step_uses_last_frame_for_continuation(monkeypatch): + """A later segment receives the prior segment's last decoded frame.""" + backend = MiniMaxH3GenerationBackend(gpu_id=0) + backend.model_config = dict(FASTH3_MODEL_CONFIG) + backend.generator = _RecordingGenerator() + monkeypatch.setattr("dreamverse.minimax_h3_generation.torch.cuda.synchronize", lambda: None) + + first_result = backend.generate_step( + "first prompt", + segment_idx=1, + image_path=None, + reset_conditioning=True, + ) + second_result = backend.generate_step( + "second prompt", + segment_idx=2, + image_path=None, + reset_conditioning=False, + ) + + first_request = backend.generator.requests[0] + assert first_request.inputs.pil_image is None + assert first_request.negative_prompt == "" + assert first_request.sampling.height == 768 + assert first_request.sampling.width == 1344 + assert first_request.sampling.num_frames == 124 + assert first_request.sampling.num_inference_steps == 5 + assert first_request.sampling.fps == 24 + assert first_request.sampling.guidance_scale == 1.0 + assert first_request.sampling.batch_cfg is False + assert first_request.sampling.seed == 1000 + assert first_request.output.save_video is False + assert first_request.output.return_frames is True + assert backend.generator.conditioning_pixels[1].tolist() == np.full((2, 3, 3), 20).tolist() + assert first_result.head_trim_frames == 0 + assert first_result.head_trim_audio_frames == 0 + assert second_result.head_trim_frames == 1 + assert second_result.head_trim_audio_frames == 1 + assert second_result.audio_sample_rate == 44100 + + +def test_generate_step_reset_uses_text_to_video_path(monkeypatch): + """Resetting continuation produces an unconditioned text-to-video request.""" + backend = MiniMaxH3GenerationBackend(gpu_id=0) + backend.model_config = dict(FASTH3_MODEL_CONFIG) + backend.generator = _RecordingGenerator() + monkeypatch.setattr("dreamverse.minimax_h3_generation.torch.cuda.synchronize", lambda: None) + + backend.generate_step("first prompt", 1, None, True) + reset_result = backend.generate_step("reset prompt", 2, None, True) + + assert backend.generator.requests[-1].inputs.pil_image is None + assert reset_result.head_trim_frames == 0 + assert reset_result.head_trim_audio_frames == 0 + + +def test_generate_step_missing_continuation_frame(monkeypatch): + """A later segment fails when no reset or retained frame defines its input.""" + backend = MiniMaxH3GenerationBackend(gpu_id=0) + backend.model_config = dict(FASTH3_MODEL_CONFIG) + backend.generator = _RecordingGenerator() + + with pytest.raises(RuntimeError, match="requires a retained continuation frame"): + backend.generate_step("later prompt", 2, None, False) + + assert backend.generator.requests == [] + + +def test_warmup_exercises_text_and_first_frame_paths(monkeypatch): + """Warmup covers both request shapes used by a DreamVerse session.""" + backend = MiniMaxH3GenerationBackend(gpu_id=0) + backend.model_config = dict(FASTH3_MODEL_CONFIG) + backend.generator = _RecordingGenerator() + monkeypatch.setattr("dreamverse.minimax_h3_generation.torch.cuda.synchronize", lambda: None) + + timings = backend.warmup("warmup prompt") + + assert backend.generator.conditioning_pixels[0] is None + assert backend.generator.conditioning_pixels[1] is not None + assert backend.continuation_image is None + assert "warmup_text_to_video_ms" in timings + assert "warmup_first_frame_to_video_ms" in timings diff --git a/apps/dreamverse/dreamverse/video_generation.py b/apps/dreamverse/dreamverse/video_generation.py index 88dd493435..59f74680c1 100644 --- a/apps/dreamverse/dreamverse/video_generation.py +++ b/apps/dreamverse/dreamverse/video_generation.py @@ -1,9 +1,9 @@ -"""LTX2 model lifecycle and continuation conditioning. +"""LTX-2 model lifecycle and continuation conditioning. Runs inside a GPU worker subprocess. Owns the model, the audio encoder, and the per-session continuation state carried across segments. Callers must set ``os.environ["CUDA_VISIBLE_DEVICES"]`` -before constructing ``VideoGenerationWorker`` — all ``fastvideo.*`` +before constructing ``LTX2GenerationBackend`` — all ``fastvideo.*`` imports are deferred to method bodies so nothing touches CUDA at module import time. """ @@ -14,9 +14,6 @@ import os import re import time -from dataclasses import dataclass -from typing import Any - import numpy as np import torch @@ -35,6 +32,7 @@ DREAMVERSE_LORA_STACK, _resolve_lora_spec, ) +from dreamverse.generation_contracts import StepResult # Multi-frame decoded continuation defaults from # examples/inference/basic/basic_ltx2_distilled_video_continuation.py. @@ -80,22 +78,6 @@ def _reset_lora_registry(worker) -> dict: return {"status": "lora_registry_reset"} -@dataclass -class StepResult: - """Output of one generation step. - - ``head_trim_frames`` / ``head_trim_audio_frames`` are derived here - so downstream AV streaming never needs to import conditioning - constants. - """ - frames: list - audio: Any - audio_sample_rate: int | None - timings: dict - head_trim_frames: int - head_trim_audio_frames: int - - class ContinuationState: """Per-session video + audio conditioning carried across segments.""" @@ -202,7 +184,7 @@ def save_audio_latents(self, latents: torch.Tensor | None) -> None: self.audio_latents = latents.detach().clone().cpu() -class VideoGenerationWorker: +class LTX2GenerationBackend: """Single-GPU LTX2 generator with continuation state. Caller must set ``os.environ["CUDA_VISIBLE_DEVICES"]`` before diff --git a/fastvideo/models/loader/lora_patch.py b/fastvideo/models/loader/lora_patch.py index 40925020ef..a28f0ec14f 100644 --- a/fastvideo/models/loader/lora_patch.py +++ b/fastvideo/models/loader/lora_patch.py @@ -182,7 +182,10 @@ def apply_to(self, param_name: str, tensor: torch.Tensor) -> torch.Tensor: raise ValueError(f"LoRA diff for {param_name} has shape {tuple(delta.shape)}, " f"but the parameter is {tuple(tensor.shape)}") self._applied.add(param_name) - return tensor.to(torch.float32) + delta.to(torch.float32) * self._strength + # Safetensors returns the adapter delta on CPU, while inference can stream the + # base checkpoint directly onto its target GPU when CPU offload is disabled. + delta = delta.to(device=tensor.device, dtype=torch.float32) + return tensor.to(torch.float32) + delta * self._strength def provides(self, param_name: str) -> bool: """Whether the adapter carries this parameter whole, without reading it.""" diff --git a/fastvideo/tests/loader/test_lora_patch.py b/fastvideo/tests/loader/test_lora_patch.py index 55e33a38da..02b7b89835 100644 --- a/fastvideo/tests/loader/test_lora_patch.py +++ b/fastvideo/tests/loader/test_lora_patch.py @@ -1,7 +1,7 @@ """Whole-parameter LoRA payload: key classification, application, and reporting. -CPU-only. Nothing here loads a model -- the point is the key algebra and the arithmetic, -which is where a wrong answer is silent rather than loud. +Nothing here loads a model. Most tests exercise key algebra and arithmetic on the CPU; +one regression test verifies cross-device adapter loading when CUDA is available. """ import logging @@ -134,6 +134,18 @@ def test_apply_to_adds_the_delta(tmp_path): assert torch.allclose(out, torch.full((4, ), 1.25)) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for cross-device adapter loading") +def test_apply_to_moves_cpu_delta_to_base_tensor_device(tmp_path): + """A CPU safetensors delta follows a directly loaded CUDA checkpoint tensor.""" + path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.full((4, ), 0.25)}) + patch = DenseLoRAPatch.from_adapter(path) + + out = patch.apply_to("blocks.0.norm1.weight", torch.ones(4, device="cuda")) + + assert out.device.type == "cuda" + assert torch.allclose(out.cpu(), torch.full((4, ), 1.25)) + + def test_apply_to_leaves_unrelated_parameters_untouched(tmp_path): path = write_adapter(tmp_path, {"blocks.0.norm1.diff": torch.full((4, ), 0.25)}) patch = DenseLoRAPatch.from_adapter(path) From 879742bc63104d417287f2cf89daf4d5dc97910b Mon Sep 17 00:00:00 2001 From: Davids048 Date: Mon, 31 Aug 2026 23:58:53 +0000 Subject: [PATCH 2/2] rename LTX2 generation backend. --- apps/dreamverse/arch.md | 2 +- .../dreamverse/benchmarks/benchmark_pipeline.py | 2 +- apps/dreamverse/dreamverse/generation_worker.py | 2 +- .../{video_generation.py => ltx2_generation.py} | 0 .../dreamverse/tests/test_generation_worker.py | 9 +++++++++ .../dreamverse/tests/test_import_contract.py | 4 ++-- apps/dreamverse/gpu-pool.drawio | 14 +++++++------- apps/dreamverse/gpu-pool.svg | 2 +- 8 files changed, 22 insertions(+), 13 deletions(-) rename apps/dreamverse/dreamverse/{video_generation.py => ltx2_generation.py} (100%) create mode 100644 apps/dreamverse/dreamverse/tests/test_generation_worker.py diff --git a/apps/dreamverse/arch.md b/apps/dreamverse/arch.md index 64b0aff7b0..eb532f17d8 100644 --- a/apps/dreamverse/arch.md +++ b/apps/dreamverse/arch.md @@ -145,7 +145,7 @@ Model generation has a separate ownership boundary inside each GPU process: - `apps/dreamverse/dreamverse/generation_worker.py` selects the backend that the active model profile declares and owns the backend lifecycle. -- `apps/dreamverse/dreamverse/video_generation.py` owns LTX-2 generator configuration, video and audio continuation, and +- `apps/dreamverse/dreamverse/ltx2_generation.py` owns LTX-2 generator configuration, video and audio continuation, and runtime LoRA application. - `apps/dreamverse/dreamverse/minimax_h3_generation.py` owns the VSA data-free FastH3 adapter, FastH3 generator and request configuration, and last-frame continuation through MiniMax H3 first-frame conditioning. diff --git a/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py b/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py index 984b7e8d9d..fdfa462557 100644 --- a/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py +++ b/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py @@ -1,6 +1,6 @@ """Benchmark the LTX-2 generation pipeline driven by the dreamverse Python SDK path. -Mirrors how ``apps/dreamverse/dreamverse/video_generation.py`` constructs +Mirrors how ``apps/dreamverse/dreamverse/ltx2_generation.py`` constructs ``GeneratorConfig`` and calls ``VideoGenerator.generate()``, then captures per-stage timings via the ``FASTVIDEO_STAGE_LOGGING=1`` log hooks (same mechanism as ``FastVideo-internal/examples/inference/basic/ diff --git a/apps/dreamverse/dreamverse/generation_worker.py b/apps/dreamverse/dreamverse/generation_worker.py index e06a9aa179..1585940ae1 100644 --- a/apps/dreamverse/dreamverse/generation_worker.py +++ b/apps/dreamverse/dreamverse/generation_worker.py @@ -9,7 +9,7 @@ def _create_generation_backend(backend_name: str, gpu_id: int) -> GenerationBackend: """Construct the backend that owns the selected model family's behavior.""" if backend_name == "ltx2": - from dreamverse.video_generation import LTX2GenerationBackend + from dreamverse.ltx2_generation import LTX2GenerationBackend return LTX2GenerationBackend(gpu_id) if backend_name == "minimax_h3": diff --git a/apps/dreamverse/dreamverse/video_generation.py b/apps/dreamverse/dreamverse/ltx2_generation.py similarity index 100% rename from apps/dreamverse/dreamverse/video_generation.py rename to apps/dreamverse/dreamverse/ltx2_generation.py diff --git a/apps/dreamverse/dreamverse/tests/test_generation_worker.py b/apps/dreamverse/dreamverse/tests/test_generation_worker.py new file mode 100644 index 0000000000..c98bf65441 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_generation_worker.py @@ -0,0 +1,9 @@ +from dreamverse.generation_worker import _create_generation_backend +from dreamverse.ltx2_generation import LTX2GenerationBackend + + +def test_create_generation_backend_ltx2_module_import(): + backend = _create_generation_backend("ltx2", gpu_id=3) + + assert isinstance(backend, LTX2GenerationBackend) + assert backend.gpu_id == 3 diff --git a/apps/dreamverse/dreamverse/tests/test_import_contract.py b/apps/dreamverse/dreamverse/tests/test_import_contract.py index f6cee5493e..5840835480 100644 --- a/apps/dreamverse/dreamverse/tests/test_import_contract.py +++ b/apps/dreamverse/dreamverse/tests/test_import_contract.py @@ -17,11 +17,11 @@ ) ALLOWED_INTERNAL_IMPORTS = { ( - "video_generation.py", + "ltx2_generation.py", "fastvideo.models.audio.ltx2_audio_processing", ), ( - "video_generation.py", + "ltx2_generation.py", "fastvideo.models.loader.component_loader", ), } diff --git a/apps/dreamverse/gpu-pool.drawio b/apps/dreamverse/gpu-pool.drawio index ea9fe67b56..9a9d4d3852 100644 --- a/apps/dreamverse/gpu-pool.drawio +++ b/apps/dreamverse/gpu-pool.drawio @@ -70,7 +70,7 @@ - + @@ -79,13 +79,13 @@ - + - + @@ -250,7 +250,7 @@ - + @@ -389,10 +389,10 @@ - + - + @@ -534,7 +534,7 @@ - + diff --git a/apps/dreamverse/gpu-pool.svg b/apps/dreamverse/gpu-pool.svg index aa379cf605..1259670865 100644 --- a/apps/dreamverse/gpu-pool.svg +++ b/apps/dreamverse/gpu-pool.svg @@ -1 +1 @@ -
gpu_pool.py — runtime architecture
gpu_pool.py — runtime architecture
Client (WebSocket)
Client (WebSocket)
GPUPool
GPUPool
GPUSlot (one per GPU)
GPUSlot (one per GPU)
start()
start()
shutdown()
shutdown()
join_user()
join_user()
user_step()
user_step()
leave_user()
leave_user()
_send_command()

puts on command_queue,
blocks on response_queue.get

**for user agonistic commands
_send_command()...
_send_command_tagged()

creates Future; registers in
_pending_futures[user_id];
ensures response reader; awaits future

**for user specific commands
_send_command_tagged()...
_response_reader()

a while loop that reads response_queue;
dispatch WorkerEvent:
• media → _stream_queues
• acks/errors → _pending_futures
_response_reader()...
«dict» _pending_futures

{user_id → asyncio.Future}

created by _send_command_tagged;
resolved + popped by _response_reader;
cleared on model reload or leave
«dict» _pending_futures...
«dict» _stream_queues: 
each is a queue of metadata that tell the consumer which part of the shared 
stream buffer to read next.

{user_id → asyncio.Queue}

created by register_stream_queue()
written by _response_reader (MediaInit / Chunk / Complete);
consumed by  AV loop → ws.send_bytes / send_json;
popped by leave_user(); cleared on reload
«dict» _stream_queues:...
── process boundary (mp.Queue + mp.RawArray) ──
── process boundary (mp.Queue + mp.RawArray) ──
command_queue
command_queue
response_queue
response_queue
shared_stream_buffer
(mp.RawArray, 256 MiB)

zero-copy mp4 bytes
shared_stream_buffer...
worker_ipc.py — shared types (imported on both sides of the process boundary)

Command(type: CommandType, payload: CommandPayload | None, user_id: str | None)
CommandPayload = UserStepPayload | WarmupPayload | ReloadModelPayload
WorkerEvent = StepComplete | WorkerError | JoinAck | LeaveAck | ReloadAck | WarmupComplete | MediaInit | MediaChunk | MediaComplete | InitAck | ShutdownAck
worker_ipc.py — shared types (imported on both sides of the process boundary)...
Worker Subprocess - in each GPUSlot
Worker Subprocess - in each GPUSlot
command dispatcher

gpu_worker_process() branches on
CommandType; asserts payload type

INIT / WARMUP / RELOAD_MODEL
USER_JOIN / USER_STEP / USER_LEAVE
SHUTDOWN
command dispatcher...
VideoGenerationWorker.generate_step()
video_generation.py:380

reads + updates ContinuationState,
calls generator
VideoGenerationWorker.generate_step()...
stream_fmp4()
av_streaming.py:121

trims overlap, pipes to ffmpeg,
publishes StreamInit / StreamChunk /
StreamComplete via callback
stream_fmp4()...
VideoGenerator (fastvideo)

LTX2 DiT + refine upsampler
FP4 quant, torch.compile

owned by VideoGenerationWorker
video_generation.py:211
VideoGenerator (fastvideo)...
ffmpeg subprocess

libx264 / *_nvenc
fragmented mp4
ffmpeg subprocess...
ContinuationState
video_generation.py:89

• video_images: list[PIL.Image]
• audio_latents: torch.Tensor (CPU)

carried across segments
ContinuationState...
acquire
acquire
assigns slot
assigns slot
register future
register future
set future result
set future result
put(MediaEvent)
put(MediaEvent)
put Command
put Command
put Command
put Command
WorkerEvent
WorkerEvent
blocking .get
blocking .get
USER_STEP
USER_STEP
StepComplete / ack types
StepComplete / ack types
generator.generate_video()
generator.generate_video()
read / write
read / write
frames + audio
frames + audio
rawvideo + wav
rawvideo + wav
write mp4 bytes
at shared_write_offset
write mp4 bytes...
bytes from the buffer are sent through websocket 
based on info from the stream queue
bytes from the buffer are sent through websocket...
session AV loop (main.py:2124):
ws.send_bytes(chunk) — binary mp4 frames
ws.send_json({type: media_init | media_segment_complete | ...})
session AV loop (main.py:2124):...
MediaChunk (offset, len)
MediaChunk (offset, len)
Legend

■ blue client / external
■ green main-process pool/slot
(methods — italic label)
■ yellow containers (routing state)
■ red IPC primitives (mp.Queue, mp.RawArray)

Worker subprocess modules:
■ orange gpu_pool.py (dispatcher)
■ lavender video_generation.py
■ teal av_streaming.py
■ gray worker_ipc.py (shared types)

Flow:
client → pool → slot
→ _send_command(_tagged) → command_queue
→ dispatcher → generate_step()
→ stream_fmp4() → ffmpeg
→ shared_buf + response_queue
→ _response_reader → futures / stream_queues
→ client awaits (via main.py AV loop)
Legend...
FastVideo video_generator
FastVideo video_genera...
reload model
reload model
Queue: owned by GPUSlot, passed to worker process as argument.
Queue: owned by GPUSlot, passed to w...Text is not SVG - cannot display \ No newline at end of file +
gpu_pool.py — runtime architecture
gpu_pool.py — runtime architecture
Client (WebSocket)
Client (WebSocket)
GPUPool
GPUPool
GPUSlot (one per GPU)
GPUSlot (one per GPU)
start()
start()
shutdown()
shutdown()
join_user()
join_user()
user_step()
user_step()
leave_user()
leave_user()
_send_command()

puts on command_queue,
blocks on response_queue.get

**for user agonistic commands
_send_command()...
_send_command_tagged()

creates Future; registers in
_pending_futures[user_id];
ensures response reader; awaits future

**for user specific commands
_send_command_tagged()...
_response_reader()

a while loop that reads response_queue;
dispatch WorkerEvent:
• media → _stream_queues
• acks/errors → _pending_futures
_response_reader()...
«dict» _pending_futures

{user_id → asyncio.Future}

created by _send_command_tagged;
resolved + popped by _response_reader;
cleared on model reload or leave
«dict» _pending_futures...
«dict» _stream_queues: 
each is a queue of metadata that tell the consumer which part of the shared 
stream buffer to read next.

{user_id → asyncio.Queue}

created by register_stream_queue()
written by _response_reader (MediaInit / Chunk / Complete);
consumed by  AV loop → ws.send_bytes / send_json;
popped by leave_user(); cleared on reload
«dict» _stream_queues:...
── process boundary (mp.Queue + mp.RawArray) ──
── process boundary (mp.Queue + mp.RawArray) ──
command_queue
command_queue
response_queue
response_queue
shared_stream_buffer
(mp.RawArray, 256 MiB)

zero-copy mp4 bytes
shared_stream_buffer...
worker_ipc.py — shared types (imported on both sides of the process boundary)

Command(type: CommandType, payload: CommandPayload | None, user_id: str | None)
CommandPayload = UserStepPayload | WarmupPayload | ReloadModelPayload
WorkerEvent = StepComplete | WorkerError | JoinAck | LeaveAck | ReloadAck | WarmupComplete | MediaInit | MediaChunk | MediaComplete | InitAck | ShutdownAck
worker_ipc.py — shared types (imported on both sides of the process boundary)...
Worker Subprocess - in each GPUSlot
Worker Subprocess - in each GPUSlot
command dispatcher

gpu_worker_process() branches on
CommandType; asserts payload type

INIT / WARMUP / RELOAD_MODEL
USER_JOIN / USER_STEP / USER_LEAVE
SHUTDOWN
command dispatcher...
VideoGenerationWorker.generate_step()
ltx2_generation.py:380

reads + updates ContinuationState,
calls generator
VideoGenerationWorker.generate_step()...
stream_fmp4()
av_streaming.py:121

trims overlap, pipes to ffmpeg,
publishes StreamInit / StreamChunk /
StreamComplete via callback
stream_fmp4()...
VideoGenerator (fastvideo)

LTX2 DiT + refine upsampler
FP4 quant, torch.compile

owned by VideoGenerationWorker
ltx2_generation.py:211
VideoGenerator (fastvideo)...
ffmpeg subprocess

libx264 / *_nvenc
fragmented mp4
ffmpeg subprocess...
ContinuationState
ltx2_generation.py:89

• video_images: list[PIL.Image]
• audio_latents: torch.Tensor (CPU)

carried across segments
ContinuationState...
acquire
acquire
assigns slot
assigns slot
register future
register future
set future result
set future result
put(MediaEvent)
put(MediaEvent)
put Command
put Command
put Command
put Command
WorkerEvent
WorkerEvent
blocking .get
blocking .get
USER_STEP
USER_STEP
StepComplete / ack types
StepComplete / ack types
generator.generate_video()
generator.generate_video()
read / write
read / write
frames + audio
frames + audio
rawvideo + wav
rawvideo + wav
write mp4 bytes
at shared_write_offset
write mp4 bytes...
bytes from the buffer are sent through websocket 
based on info from the stream queue
bytes from the buffer are sent through websocket...
session AV loop (main.py:2124):
ws.send_bytes(chunk) — binary mp4 frames
ws.send_json({type: media_init | media_segment_complete | ...})
session AV loop (main.py:2124):...
MediaChunk (offset, len)
MediaChunk (offset, len)
Legend

■ blue client / external
■ green main-process pool/slot
(methods — italic label)
■ yellow containers (routing state)
■ red IPC primitives (mp.Queue, mp.RawArray)

Worker subprocess modules:
■ orange gpu_pool.py (dispatcher)
■ lavender ltx2_generation.py
■ teal av_streaming.py
■ gray worker_ipc.py (shared types)

Flow:
client → pool → slot
→ _send_command(_tagged) → command_queue
→ dispatcher → generate_step()
→ stream_fmp4() → ffmpeg
→ shared_buf + response_queue
→ _response_reader → futures / stream_queues
→ client awaits (via main.py AV loop)
Legend...
FastVideo video_generator
FastVideo video_genera...
reload model
reload model
Queue: owned by GPUSlot, passed to worker process as argument.
Queue: owned by GPUSlot, passed to w...Text is not SVG - cannot display \ No newline at end of file