diff --git a/apps/dreamverse/dreamverse/av_streaming.py b/apps/dreamverse/dreamverse/av_streaming.py new file mode 100644 index 000000000..ee9c53988 --- /dev/null +++ b/apps/dreamverse/dreamverse/av_streaming.py @@ -0,0 +1,445 @@ +# pyright: reportMissingTypeArgument=false, reportArgumentType=false, reportOptionalSubscript=false, reportOptionalMemberAccess=false, reportConstantRedefinition=false, reportCallIssue=false +# ruff: noqa: UP007, SIM108, SIM105 +# mypy: ignore-errors +"""ffmpeg fMP4 muxing with chunk-level event emission. + +Self-contained: spawns ffmpeg as a subprocess, pipes raw frames into +its stdin, reads fragmented-MP4 chunks from stdout, and publishes each +chunk as a ``StreamEvent`` via the caller-supplied ``publish`` +callback. Knows nothing about multiprocessing queues, the GPU pool, +or individual users — the caller decides what "publish" means. +""" +import fcntl +import os +import shutil +import subprocess +import tempfile +import threading +import time +import uuid +import wave +from dataclasses import dataclass +from typing import Union +from collections.abc import Callable + +import numpy as np +import torch + +FFMPEG_BIN = shutil.which(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg")) +AV_MEDIA_MIME = os.getenv( + "STREAM_MIME_TYPE", + 'video/mp4; codecs="avc1.42E01E,mp4a.40.2"', +) +AV_CHUNK_SIZE_BYTES = 1048576 +TARGET_FPS = 24 +AV_FRAGMENT_DURATION_US = int(os.getenv("FASTVIDEO_FRAG_US", "250000")) +X264_GOP_FRAMES = int(os.getenv("FASTVIDEO_X264_GOP", "12")) +X264_PROFILE = os.getenv("FASTVIDEO_X264_PROFILE", "baseline").strip().lower() +if X264_PROFILE not in {"baseline", "main", "high", "main10", "high10"}: + print(f"[WARN] Unsupported FASTVIDEO_X264_PROFILE={X264_PROFILE}; using baseline") + X264_PROFILE = "baseline" +USE_SHARED_STREAM_BUFFER = (os.getenv("FASTVIDEO_USE_SHARED_STREAM_BUFFER", "1").strip().lower() + not in {"0", "false", "no"}) +SHARED_STREAM_BUFFER_BYTES = int(os.getenv("FASTVIDEO_SHARED_STREAM_BUFFER_BYTES", str(256 * 1024 * 1024))) + + +@dataclass +class StreamInit: + """First event emitted — tells the consumer the stream is starting.""" + stream_id: str + mime: str + uses_shared_buffer: bool + + +@dataclass +class StreamChunk: + """One fMP4 chunk. Either ``chunk`` (raw bytes) or + ``chunk_offset``+``chunk_length`` (read from the shared buffer) + will be populated, never both.""" + stream_id: str + chunk: bytes | None = None + chunk_offset: int | None = None + chunk_length: int | None = None + uses_shared_buffer: bool = False + + +@dataclass +class StreamComplete: + """Final event emitted — muxing finished successfully.""" + stream_id: str + chunks: int + + +StreamEvent = Union[StreamInit, StreamChunk, StreamComplete] + + +def generate_stream_id(segment_idx: int) -> str: + """Convenience: build a stream id of the form ``seg007-abcd1234``.""" + return f"seg{segment_idx:03d}-{uuid.uuid4().hex[:8]}" + + +def _normalize_audio_tensor(audio: object) -> tuple[np.ndarray, int] | None: + """Convert audio tensor/array into int16 ndarray [samples, channels].""" + if audio is None: + return None + + if torch.is_tensor(audio): + audio_np = audio.detach().cpu().float().numpy() + else: + audio_np = np.asarray(audio, dtype=np.float32) + + if audio_np.ndim == 1: + audio_np = audio_np[:, None] + elif audio_np.ndim == 2: + if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]: + audio_np = audio_np.T + else: + return None + + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_int16 = (audio_np * 32767.0).astype(np.int16) + num_channels = audio_int16.shape[1] + return audio_int16, num_channels + + +def _write_audio_wav( + audio_int16: np.ndarray, + num_channels: int, + sample_rate: int, +) -> str: + """Write normalized int16 audio to a temporary WAV file.""" + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + wav_path = f.name + try: + with wave.open(wav_path, "wb") as wav_file: + wav_file.setnchannels(num_channels) + wav_file.setsampwidth(2) + wav_file.setframerate(sample_rate) + wav_file.writeframes(audio_int16.tobytes()) + except Exception: + try: + os.unlink(wav_path) + except FileNotFoundError: + pass + raise + return wav_path + + +def stream_fmp4( + *, + frames: list[np.ndarray], + audio: object, + audio_sample_rate: int | None, + stream_id: str, + timings: dict, + head_trim_frames: int = 0, + head_trim_audio_frames: int | None = None, + shared_buffer=None, + shared_buffer_bytes: int = 0, + publish: Callable[[StreamEvent], None], + log_prefix: str = "", +) -> tuple[bool, str | None]: + """Encode frames+audio with ffmpeg, publish each fMP4 chunk as an event. + + Args: + frames: RGB24 video frames as HxWx3 uint8 arrays. + audio: 1D/2D tensor or ndarray, float values in [-1, 1]. + audio_sample_rate: sample rate of ``audio``. + stream_id: caller-supplied identifier carried on every event. + timings: dict mutated in place with ffmpeg/stream timing metrics. + head_trim_frames: video frames to drop from the start + (conditioning overlap). + head_trim_audio_frames: video-frame-equivalent audio to drop. + Defaults to ``head_trim_frames``. + shared_buffer: optional ``mp.RawArray``-compatible object; when + provided, chunks are written into it and emitted by offset + rather than by bytes (avoids IPC copies). + shared_buffer_bytes: size of ``shared_buffer`` in bytes. + publish: callback invoked once per stream event. + log_prefix: prepended to warning prints (e.g. ``"[GPU 0]"``). + + Returns: + ``(True, None)`` on success, ``(False, error_message)`` on + failure. On mid-stream failure, a ``StreamInit`` may have + already been published — the caller is responsible for + handling that. + """ + if not frames: + return False, "no frames returned" + if audio is None: + return False, "audio is None" + if audio_sample_rate is None: + return False, "audio_sample_rate is None" + if FFMPEG_BIN is None: + return False, "ffmpeg not found" + + if head_trim_audio_frames is None: + head_trim_audio_frames = head_trim_frames + + normalized_audio = _normalize_audio_tensor(audio) + if normalized_audio is None: + shape_hint = getattr(audio, "shape", None) + return False, f"unsupported audio shape={shape_hint}" + audio_int16, num_channels = normalized_audio + + if head_trim_frames < 0: + return False, (f"head_trim_frames must be >= 0, " + f"got {head_trim_frames}") + if head_trim_frames >= len(frames): + return False, (f"head_trim_frames={head_trim_frames} removes " + f"all {len(frames)} frames in segment") + + out_frames = (frames[head_trim_frames:] if head_trim_frames > 0 else frames) + sample_rate = int(audio_sample_rate) + if head_trim_audio_frames > 0: + trim_start_samples = int(round((head_trim_audio_frames / float(TARGET_FPS)) * sample_rate)) + if trim_start_samples >= audio_int16.shape[0]: + return False, ("audio too short after overlap trim: " + f"trim_start_samples={trim_start_samples}" + f", audio_samples={audio_int16.shape[0]}") + keep_samples = int(round((len(out_frames) / float(TARGET_FPS)) * sample_rate)) + trim_end_samples = min( + audio_int16.shape[0], + trim_start_samples + keep_samples, + ) + if trim_end_samples <= trim_start_samples: + return False, ("invalid audio trim range: " + f"start={trim_start_samples}, " + f"end={trim_end_samples}") + audio_int16 = audio_int16[trim_start_samples:trim_end_samples] + + height = int(out_frames[0].shape[0]) + width = int(out_frames[0].shape[1]) + codec = os.getenv("FASTVIDEO_VIDEO_CODEC", "libx264") + t_wav_start = time.perf_counter() + wav_path = _write_audio_wav(audio_int16, num_channels, sample_rate) + wav_write_ms = (time.perf_counter() - t_wav_start) * 1000 + + cmd = [ + FFMPEG_BIN, + "-hide_banner", + "-loglevel", + "error", + "-y", + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "-s:v", + f"{width}x{height}", + "-r", + str(TARGET_FPS), + "-i", + "pipe:0", + "-i", + wav_path, + "-c:v", + codec, + ] + if codec.endswith("_nvenc"): + cmd += [ + "-preset", + os.getenv("FASTVIDEO_NVENC_PRESET", "p1"), + "-tune", + os.getenv("FASTVIDEO_NVENC_TUNE", "ull"), + "-rc", + os.getenv("FASTVIDEO_NVENC_RC", "constqp"), + "-qp", + os.getenv("FASTVIDEO_NVENC_QP", "28"), + "-bf", + os.getenv("FASTVIDEO_NVENC_BF", "0"), + ] + else: + cmd += [ + "-preset", + os.getenv("FASTVIDEO_X264_PRESET", "ultrafast"), + "-tune", + "zerolatency", + "-profile:v", + X264_PROFILE, + # Emit frequent keyframes so fragments are independently playable. + "-g", + str(X264_GOP_FRAMES), + "-keyint_min", + str(X264_GOP_FRAMES), + "-x264-params", + "scenecut=0", + ] + cmd += [ + "-c:a", + "aac", + "-pix_fmt", + os.getenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p"), + "-shortest", + "-movflags", + "+frag_keyframe+empty_moov+default_base_moof", + "-frag_duration", + str(AV_FRAGMENT_DURATION_US), + "-flush_packets", + "1", + "-muxdelay", + "0", + "-muxpreload", + "0", + "-f", + "mp4", + "pipe:1", + ] + + proc: subprocess.Popen | None = None + stderr_chunks: list[bytes] = [] + writer_error: list[Exception | None] = [None] + t_stream_start = time.perf_counter() + use_shared_buffer = (USE_SHARED_STREAM_BUFFER and shared_buffer is not None and shared_buffer_bytes > 0) + shared_write_offset = 0 + shared_buffer_fallback = False + shared_np = (np.frombuffer( + shared_buffer, + dtype=np.uint8, + count=shared_buffer_bytes, + ) if use_shared_buffer else None) + chunk_intervals_ms: list[float] = [] + chunk_publish_ms: list[float] = [] + chunk_read_ms: list[float] = [] + + try: + t_proc_spawn_start = time.perf_counter() + proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + bufsize=0, + ) + assert proc.stdin is not None + assert proc.stdout is not None + assert proc.stderr is not None + if hasattr(fcntl, "F_SETPIPE_SZ"): + fcntl.fcntl(proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, 1048576) + fcntl.fcntl(proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, 1048576) + ffmpeg_spawn_ms = (time.perf_counter() - t_proc_spawn_start) * 1000 + + def _write_frames(): + try: + for frame in out_frames: + proc.stdin.write(np.ascontiguousarray(frame).tobytes()) + proc.stdin.close() + except Exception as exc: + writer_error[0] = exc + try: + proc.stdin.close() + except Exception: + pass + + def _read_stderr(): + try: + while True: + data = proc.stderr.read(4096) + if not data: + break + stderr_chunks.append(data) + except Exception: + pass + + writer_thread = threading.Thread(target=_write_frames, daemon=True) + stderr_thread = threading.Thread(target=_read_stderr, daemon=True) + writer_thread.start() + stderr_thread.start() + + publish(StreamInit( + stream_id=stream_id, + mime=AV_MEDIA_MIME, + uses_shared_buffer=use_shared_buffer, + )) + + chunk_count = 0 + total_bytes = 0 + first_chunk_ms: float | None = None + last_chunk_emit_t = time.perf_counter() + while True: + t_read_start = time.perf_counter() + chunk = proc.stdout.read(AV_CHUNK_SIZE_BYTES) + t_read_end = time.perf_counter() + if not chunk: + break + chunk_read_ms.append((t_read_end - t_read_start) * 1000) + chunk_count += 1 + total_bytes += len(chunk) + if first_chunk_ms is None: + first_chunk_ms = (t_read_end - t_stream_start) * 1000 + chunk_intervals_ms.append((t_read_end - last_chunk_emit_t) * 1000) + t_publish_start = time.perf_counter() + if use_shared_buffer and not shared_buffer_fallback: + chunk_len = len(chunk) + write_end = shared_write_offset + chunk_len + if write_end <= shared_buffer_bytes: + shared_np[shared_write_offset:write_end] = np.frombuffer(chunk, dtype=np.uint8) + publish( + StreamChunk( + stream_id=stream_id, + chunk_offset=shared_write_offset, + chunk_length=chunk_len, + uses_shared_buffer=True, + )) + shared_write_offset = write_end + chunk_publish_ms.append((time.perf_counter() - t_publish_start) * 1000) + last_chunk_emit_t = time.perf_counter() + continue + shared_buffer_fallback = True + print(f"{log_prefix} Shared stream buffer exhausted at " + f"{shared_write_offset / (1024 * 1024):.1f}MB; " + "falling back to queue chunk bytes") + publish(StreamChunk( + stream_id=stream_id, + chunk=chunk, + )) + chunk_publish_ms.append((time.perf_counter() - t_publish_start) * 1000) + last_chunk_emit_t = time.perf_counter() + + writer_thread.join(timeout=5.0) + rc = proc.wait() + stderr_thread.join(timeout=1.0) + + if rc != 0: + stderr_tail = b"".join(stderr_chunks).decode(errors="ignore")[-1200:] + return False, f"ffmpeg av_fmp4 stream failed (rc={rc}): {stderr_tail}" + if writer_error[0] is not None: + return False, f"ffmpeg frame writer failed: {writer_error[0]}" + + timings["av_encode_stream_ms"] = (time.perf_counter() - t_stream_start) * 1000 + timings["av_stream_bytes"] = total_bytes + timings["av_trim_head_frames"] = head_trim_frames + timings["av_trim_head_audio_frames"] = head_trim_audio_frames + timings["av_frames_encoded"] = len(out_frames) + timings["av_shared_buffer_used"] = (bool(use_shared_buffer and not shared_buffer_fallback)) + timings["av_wav_write_ms"] = wav_write_ms + timings["av_ffmpeg_spawn_ms"] = ffmpeg_spawn_ms + timings["av_first_chunk_ms"] = first_chunk_ms or 0.0 + if chunk_intervals_ms: + timings["av_chunk_interval_ms_min"] = min(chunk_intervals_ms) + timings["av_chunk_interval_ms_median"] = float(np.median(chunk_intervals_ms)) + timings["av_chunk_interval_ms_p95"] = (float(np.percentile(chunk_intervals_ms, 95))) + timings["av_chunk_interval_ms_max"] = max(chunk_intervals_ms) + if chunk_publish_ms: + timings["av_chunk_publish_ms_median"] = float(np.median(chunk_publish_ms)) + timings["av_chunk_publish_ms_p95"] = float(np.percentile(chunk_publish_ms, 95)) + if chunk_read_ms: + timings["av_chunk_read_ms_median"] = float(np.median(chunk_read_ms)) + timings["av_chunk_read_ms_p95"] = float(np.percentile(chunk_read_ms, 95)) + publish(StreamComplete( + stream_id=stream_id, + chunks=chunk_count, + )) + return True, None + except Exception as exc: + return False, str(exc) + finally: + if proc is not None and proc.poll() is None: + try: + proc.kill() + except Exception: + pass + try: + os.remove(wav_path) + except OSError: + pass diff --git a/apps/dreamverse/dreamverse/benchmarks/__init__.py b/apps/dreamverse/dreamverse/benchmarks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/dreamverse/dreamverse/benchmarks/benchmark_av_streaming.py b/apps/dreamverse/dreamverse/benchmarks/benchmark_av_streaming.py new file mode 100644 index 000000000..caa4be00b --- /dev/null +++ b/apps/dreamverse/dreamverse/benchmarks/benchmark_av_streaming.py @@ -0,0 +1,288 @@ +"""Benchmark the AV streaming hot-path used by dreamverse-server. + +Measures wall-time, encoded byte volume, and chunk count for +``av_streaming.stream_fmp4`` over synthetic frames + audio at +production resolution. Sweeps codecs (default: ``libx264`` and +``h264_nvenc`` if the active ffmpeg supports it) and ffmpeg presets so +the deploy can pick a configuration that achieves a >=1.0 realtime +ratio (5.04s of generated video produced in <=5.04s wall-time). + +Usage:: + + python -m apps.dreamverse.server.benchmarks.benchmark_av_streaming + python -m apps.dreamverse.server.benchmarks.benchmark_av_streaming \\ + --frames 121 --width 1920 --height 1088 --runs 3 \\ + --codecs libx264 h264_nvenc --x264-preset ultrafast \\ + --nvenc-preset p1 + FASTVIDEO_FFMPEG_BIN=$HOME/opt/ffmpeg-native/bin/ffmpeg \\ + python -m apps.dreamverse.server.benchmarks.benchmark_av_streaming + +Skips ``h264_nvenc`` automatically if the binary lacks the encoder. +This is the regression guard documented in +`.agents/memory/dreamverse-integration/decisions-log.md` D-21. +""" +from __future__ import annotations + +import argparse +import os +import shutil +import statistics +import subprocess +import sys +import time +from collections.abc import Iterable +from dataclasses import dataclass + +import numpy as np +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_SERVER_ROOT = os.path.dirname(_HERE) +if _SERVER_ROOT not in sys.path: + sys.path.insert(0, _SERVER_ROOT) + +from dreamverse.av_streaming import stream_fmp4 # noqa: E402 + + +@dataclass +class BenchResult: + codec: str + preset: str + runs: int + wall_ms_min: float + wall_ms_median: float + wall_ms_p95: float + wall_ms_max: float + bytes_median: int + chunks_median: float + realtime_ratio_median: float + error: str | None = None + + +def _make_synthetic_frames(num: int, width: int, height: int, seed: int) -> list[np.ndarray]: + rng = np.random.default_rng(seed) + base = rng.integers(0, 255, size=(height, width, 3), dtype=np.uint8) + out: list[np.ndarray] = [] + for i in range(num): + f = base.copy() + f[:, :, 0] = (f[:, :, 0].astype(np.int32) + i * 2) % 256 + out.append(f) + return out + + +def _make_synthetic_audio(num_frames: int, fps: int, sample_rate: int, seed: int) -> torch.Tensor: + duration_s = num_frames / fps + samples = int(round(duration_s * sample_rate)) + rng = np.random.default_rng(seed + 1) + audio = (rng.uniform(-0.1, 0.1, (2, samples))).astype(np.float32) + return torch.from_numpy(audio) + + +def _ffmpeg_supports(codec: str, ffmpeg_bin: str) -> bool: + try: + out = subprocess.run( + [ffmpeg_bin, "-hide_banner", "-encoders"], + capture_output=True, + text=True, + check=False, + timeout=10, + ).stdout + except Exception: + return False + needle = f" {codec} " + return any(needle in line for line in out.splitlines()) + + +def _run_one(frames: list[np.ndarray], audio: torch.Tensor, sample_rate: int, codec: str, + preset: str) -> tuple[float, int, int, str | None]: + timings: dict = {} + chunks: list = [] + + def _publish(event): + chunks.append(event) + + prev_codec = os.environ.get("FASTVIDEO_VIDEO_CODEC") + prev_preset = os.environ.get("FASTVIDEO_X264_PRESET") + prev_nvenc_preset = os.environ.get("FASTVIDEO_NVENC_PRESET") + os.environ["FASTVIDEO_VIDEO_CODEC"] = codec + if codec.endswith("_nvenc"): + os.environ["FASTVIDEO_NVENC_PRESET"] = preset + else: + os.environ["FASTVIDEO_X264_PRESET"] = preset + + t0 = time.perf_counter() + try: + ok, err = stream_fmp4( + frames=frames, + audio=audio, + audio_sample_rate=sample_rate, + stream_id="bench", + timings=timings, + head_trim_frames=0, + head_trim_audio_frames=0, + shared_buffer=None, + shared_buffer_bytes=0, + publish=_publish, + log_prefix="[bench]", + ) + finally: + if prev_codec is None: + os.environ.pop("FASTVIDEO_VIDEO_CODEC", None) + else: + os.environ["FASTVIDEO_VIDEO_CODEC"] = prev_codec + if prev_preset is None: + os.environ.pop("FASTVIDEO_X264_PRESET", None) + else: + os.environ["FASTVIDEO_X264_PRESET"] = prev_preset + if prev_nvenc_preset is None: + os.environ.pop("FASTVIDEO_NVENC_PRESET", None) + else: + os.environ["FASTVIDEO_NVENC_PRESET"] = prev_nvenc_preset + wall_ms = (time.perf_counter() - t0) * 1000.0 + if not ok: + return wall_ms, 0, 0, err or "stream_fmp4 returned False" + total_bytes = int(timings.get("av_stream_bytes", 0)) + return wall_ms, total_bytes, len(chunks), None + + +def benchmark(codecs: Iterable[str], runs: int, frames_n: int, width: int, height: int, fps: int, sample_rate: int, + x264_preset: str, nvenc_preset: str, seed: int, ffmpeg_bin: str) -> list[BenchResult]: + print(f"[bench] ffmpeg_bin={ffmpeg_bin}") + print(f"[bench] frames={frames_n} {width}x{height} fps={fps} " + f"audio_sr={sample_rate} runs/codec={runs}") + frames = _make_synthetic_frames(frames_n, width, height, seed) + audio = _make_synthetic_audio(frames_n, fps, sample_rate, seed) + playable_s = frames_n / fps + print(f"[bench] playable={playable_s:.3f}s " + f"(realtime_ratio = playable / wall_time; >= 1.0 means no " + f"buffer drain)") + + results: list[BenchResult] = [] + for codec in codecs: + preset = nvenc_preset if codec.endswith("_nvenc") else x264_preset + if not _ffmpeg_supports(codec, ffmpeg_bin): + results.append( + BenchResult(codec=codec, + preset=preset, + runs=0, + wall_ms_min=0, + wall_ms_median=0, + wall_ms_p95=0, + wall_ms_max=0, + bytes_median=0, + chunks_median=0, + realtime_ratio_median=0, + error=f"{codec} not in ffmpeg")) + continue + walls: list[float] = [] + sizes: list[int] = [] + chunkcounts: list[int] = [] + last_err: str | None = None + for run in range(runs): + wall_ms, total_bytes, chunk_count, err = _run_one(frames, audio, sample_rate, codec, preset) + print(f"[bench] codec={codec:12s} preset={preset:9s} " + f"run={run + 1}/{runs} wall={wall_ms:7.1f}ms " + f"bytes={total_bytes:>9d} chunks={chunk_count:>3d} " + f"realtime={playable_s / (wall_ms / 1000.0):5.2f}x" + f"{' ERR=' + err if err else ''}") + if err is not None: + last_err = err + continue + walls.append(wall_ms) + sizes.append(total_bytes) + chunkcounts.append(chunk_count) + if not walls: + results.append( + BenchResult(codec=codec, + preset=preset, + runs=0, + wall_ms_min=0, + wall_ms_median=0, + wall_ms_p95=0, + wall_ms_max=0, + bytes_median=0, + chunks_median=0, + realtime_ratio_median=0, + error=last_err or "all runs failed")) + continue + walls_sorted = sorted(walls) + p95_idx = max(0, int(round(0.95 * (len(walls_sorted) - 1)))) + wall_med = statistics.median(walls) + results.append( + BenchResult( + codec=codec, + preset=preset, + runs=len(walls), + wall_ms_min=min(walls), + wall_ms_median=wall_med, + wall_ms_p95=walls_sorted[p95_idx], + wall_ms_max=max(walls), + bytes_median=int(statistics.median(sizes)), + chunks_median=statistics.median(chunkcounts), + realtime_ratio_median=playable_s / (wall_med / 1000.0), + )) + return results + + +def _print_summary(results: list[BenchResult]) -> None: + print() + print("=== summary ===") + header = (f"{'codec':14s} {'preset':10s} {'runs':>4s} " + f"{'wall_med_ms':>11s} {'wall_p95_ms':>11s} " + f"{'bytes_med':>10s} {'realtime':>8s} notes") + print(header) + print("-" * len(header)) + for r in results: + if r.error is not None: + print(f"{r.codec:14s} {r.preset:10s} {r.runs:>4d} " + f"{'-':>11s} {'-':>11s} {'-':>10s} {'-':>8s} " + f"ERR: {r.error}") + continue + print(f"{r.codec:14s} {r.preset:10s} {r.runs:>4d} " + f"{r.wall_ms_median:>11.1f} {r.wall_ms_p95:>11.1f} " + f"{r.bytes_median:>10d} {r.realtime_ratio_median:>7.2f}x " + f"{'OK' if r.realtime_ratio_median >= 1.0 else 'BUFFER DRAINS'}") + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--frames", type=int, default=121, help="num frames (default: 121, matches NUM_FRAMES)") + p.add_argument("--width", type=int, default=1920) + p.add_argument("--height", type=int, default=1088) + p.add_argument("--fps", type=int, default=24) + p.add_argument("--sample-rate", type=int, default=24000) + p.add_argument("--runs", type=int, default=3, help="runs per codec (default: 3 — 1 warmup + 2 timed in median)") + p.add_argument("--codecs", + nargs="+", + default=["libx264", "h264_nvenc"], + help="codecs to benchmark; missing ones are skipped") + p.add_argument("--x264-preset", default="ultrafast") + p.add_argument("--nvenc-preset", default="p1") + p.add_argument("--seed", type=int, default=0) + args = p.parse_args() + + ffmpeg_bin = shutil.which(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg")) + if ffmpeg_bin is None: + print("ffmpeg not found", file=sys.stderr) + return 2 + + results = benchmark( + codecs=args.codecs, + runs=args.runs, + frames_n=args.frames, + width=args.width, + height=args.height, + fps=args.fps, + sample_rate=args.sample_rate, + x264_preset=args.x264_preset, + nvenc_preset=args.nvenc_preset, + seed=args.seed, + ffmpeg_bin=ffmpeg_bin, + ) + _print_summary(results) + any_below = any(r.error is None and r.realtime_ratio_median < 1.0 for r in results) + return 1 if any_below else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py b/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py new file mode 100644 index 000000000..1be9e1608 --- /dev/null +++ b/apps/dreamverse/dreamverse/benchmarks/benchmark_pipeline.py @@ -0,0 +1,362 @@ +"""Benchmark the LTX-2 generation pipeline driven by the dreamverse Python SDK path. + +Mirrors how ``apps/dreamverse/server/video_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/ +basic_ltx2_distilled_i2v_two_stage_time.py``). + +Reports for each scenario: +* Total wall-time (median, p95) +* Per-stage execution_time (input_validation_stage, + prompt_encoding_stage, ltx2_refine_init_stage, latent_preparation_stage, + denoising_stage, ltx2_upsample_stage, ltx2_refine_lora_stage, + ltx2_refine_denoising_stage, audio_decoding_stage, decoding_stage) +* Realtime ratio (frames / fps / wall_time; >= 1.0 = no buffer drain) +* Peak GPU memory + +Sweep scenarios (default): +* compile=False, warmup=False → cold inference baseline +* compile=True, warmup=False → JIT compile mid-run +* compile=True, warmup=True → fully warmed (production mode) + +Skips compile / NVENC if the host lacks support. Does NOT exercise the +AV streaming path (use ``benchmark_av_streaming.py`` for that). + +Usage:: + + python -m apps.dreamverse.server.benchmarks.benchmark_pipeline + python -m apps.dreamverse.server.benchmarks.benchmark_pipeline \\ + --runs 3 --scenarios compile_warm cold --gpu 4 + +Cross-references D-21 / D-22 in +``.agents/memory/dreamverse-integration/decisions-log.md``. +""" +from __future__ import annotations + +import argparse +import contextlib +import json +import os +import statistics +import time +from collections import OrderedDict +from dataclasses import asdict, dataclass + +if "FASTVIDEO_STAGE_LOGGING" not in os.environ: + os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" +if "FASTVIDEO_ATTENTION_BACKEND" not in os.environ: + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" + +import torch # noqa: E402 + +from fastvideo import VideoGenerator # noqa: E402 +from fastvideo.api import ( # noqa: E402 + ComponentConfig, CompileConfig, EngineConfig, GeneratorConfig, OffloadConfig, PipelineSelection, QuantizationConfig, +) + +DEFAULT_PROMPT = ("A cinematic drone shot over coastal cliffs at sunrise, golden " + "light, gentle ocean waves, ultra detailed") +DEFAULT_MODEL = "FastVideo/LTX2-Distilled-Diffusers" + + +@dataclass +class ScenarioConfig: + name: str + enable_compile: bool + do_warmup: bool + nvenc: bool = False + + +DEFAULT_SCENARIOS: list[ScenarioConfig] = [ + ScenarioConfig(name="cold", enable_compile=False, do_warmup=False), + ScenarioConfig(name="compile_cold", enable_compile=True, do_warmup=False), + ScenarioConfig(name="compile_warm", enable_compile=True, do_warmup=True), +] + + +@dataclass +class RunResult: + wall_ms: float + stage_times_ms: OrderedDict[str, float] + peak_gpu_mb: float + error: str | None = None + + +@dataclass +class ScenarioResult: + name: str + enable_compile: bool + do_warmup: bool + runs: int + wall_ms_median: float + wall_ms_p95: float + stage_means_ms: OrderedDict[str, float] + realtime_ratio_median: float + peak_gpu_mb_max: float + error: str | None = None + + +def _build_generator_config(model_path: str, enable_compile: bool, num_gpus: int) -> GeneratorConfig: + components = ComponentConfig(config_root=model_path) + return GeneratorConfig( + model_path=model_path, + engine=EngineConfig( + num_gpus=num_gpus, + offload=OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False, vae=False, pin_cpu_memory=True), + compile=CompileConfig(enabled=enable_compile, + text_encoder_enabled=enable_compile, + backend="inductor", + fullgraph=True, + mode="max-autotune-no-cudagraphs", + dynamic=False), + use_fsdp_inference=False, + quantization=QuantizationConfig(transformer_quant="NVFP4"), + ), + pipeline=PipelineSelection( + components=components, + vae_tiling=False, + preset_overrides={ + "refine": { + "enabled": True, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "add_noise": True, + }, + }, + ), + ) + + +def _extract_stage_times(result: dict) -> OrderedDict[str, float]: + out: OrderedDict[str, float] = OrderedDict() + info = result.get("logging_info") if isinstance(result, dict) else None + if info is None: + return out + stages = getattr(info, "stages", None) + if not stages: + return out + for name, metrics in stages.items(): + exec_time = float(metrics.get("execution_time", 0.0)) + out[name] = exec_time * 1000.0 + return out + + +def _peak_gpu_mb() -> float: + if not torch.cuda.is_available(): + return 0.0 + try: + return torch.cuda.max_memory_allocated() / (1024 * 1024) + except Exception: + return 0.0 + + +def _reset_peak_gpu() -> None: + if torch.cuda.is_available(): + with contextlib.suppress(Exception): + torch.cuda.reset_peak_memory_stats() + + +def _do_one_run(generator: VideoGenerator, prompt: str, *, height: int, width: int, num_frames: int, seed: int, + num_inference_steps: int) -> RunResult: + _reset_peak_gpu() + t0 = time.perf_counter() + try: + result = generator.generate_video( + prompt=prompt, + negative_prompt="", + save_video=False, + height=height, + width=width, + num_frames=num_frames, + fps=24, + num_inference_steps=num_inference_steps, + guidance_scale=1.0, + seed=seed, + ltx2_image_crf=0.0, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize() + except Exception as exc: + return RunResult(wall_ms=(time.perf_counter() - t0) * 1000.0, + stage_times_ms=OrderedDict(), + peak_gpu_mb=0.0, + error=f"{type(exc).__name__}: {exc}") + wall_ms = (time.perf_counter() - t0) * 1000.0 + return RunResult( + wall_ms=wall_ms, + stage_times_ms=_extract_stage_times(result), + peak_gpu_mb=_peak_gpu_mb(), + ) + + +def benchmark_scenario(scenario: ScenarioConfig, model_path: str, num_gpus: int, prompt: str, num_runs: int, + num_frames: int, height: int, width: int, num_inference_steps: int, seed: int) -> ScenarioResult: + print() + print(f"=== scenario: {scenario.name} " + f"(compile={scenario.enable_compile} warmup={scenario.do_warmup}) ===") + + config = _build_generator_config(model_path, scenario.enable_compile, num_gpus) + generator = VideoGenerator.from_config(config) + + if scenario.do_warmup: + print(f"[{scenario.name}] warmup: 2 generate calls " + "(triggers compile + first-shape graphs)") + for warmup_idx in range(2): + t0 = time.perf_counter() + _do_one_run(generator, + prompt, + height=height, + width=width, + num_frames=num_frames, + seed=seed + 100 + warmup_idx, + num_inference_steps=num_inference_steps) + print(f"[{scenario.name}] warmup {warmup_idx + 1}: " + f"{(time.perf_counter() - t0) * 1000:.0f}ms") + + runs: list[RunResult] = [] + last_error: str | None = None + for run_idx in range(num_runs): + result = _do_one_run(generator, + prompt, + height=height, + width=width, + num_frames=num_frames, + seed=seed + run_idx, + num_inference_steps=num_inference_steps) + if result.error is not None: + print(f"[{scenario.name}] run {run_idx + 1}: " + f"ERROR {result.error}") + last_error = result.error + continue + playable_s = num_frames / 24.0 + rt = playable_s / (result.wall_ms / 1000.0) + print(f"[{scenario.name}] run {run_idx + 1}/{num_runs}: " + f"wall={result.wall_ms:.0f}ms peak={result.peak_gpu_mb:.0f}MB " + f"realtime={rt:.2f}x stages={len(result.stage_times_ms)}") + runs.append(result) + + if not runs: + return ScenarioResult(name=scenario.name, + enable_compile=scenario.enable_compile, + do_warmup=scenario.do_warmup, + runs=0, + wall_ms_median=0.0, + wall_ms_p95=0.0, + stage_means_ms=OrderedDict(), + realtime_ratio_median=0.0, + peak_gpu_mb_max=0.0, + error=last_error or "all runs failed") + + walls = [r.wall_ms for r in runs] + walls_sorted = sorted(walls) + p95_idx = max(0, int(round(0.95 * (len(walls_sorted) - 1)))) + stage_means: OrderedDict[str, float] = OrderedDict() + if runs: + for stage_name in runs[0].stage_times_ms: + vals = [r.stage_times_ms.get(stage_name, 0.0) for r in runs] + stage_means[stage_name] = sum(vals) / len(vals) + return ScenarioResult( + name=scenario.name, + enable_compile=scenario.enable_compile, + do_warmup=scenario.do_warmup, + runs=len(runs), + wall_ms_median=statistics.median(walls), + wall_ms_p95=walls_sorted[p95_idx], + stage_means_ms=stage_means, + realtime_ratio_median=(num_frames / 24.0) / (statistics.median(walls) / 1000.0), + peak_gpu_mb_max=max(r.peak_gpu_mb for r in runs), + ) + + +def _print_summary(results: list[ScenarioResult]) -> None: + print() + print("=== summary ===") + header = (f"{'scenario':14s} {'runs':>4s} " + f"{'wall_med_ms':>11s} {'wall_p95_ms':>11s} " + f"{'peak_mb':>8s} {'realtime':>8s}") + print(header) + print("-" * len(header)) + for r in results: + if r.error is not None: + print(f"{r.name:14s} {r.runs:>4d} ERR: {r.error}") + continue + print(f"{r.name:14s} {r.runs:>4d} " + f"{r.wall_ms_median:>11.0f} {r.wall_ms_p95:>11.0f} " + f"{r.peak_gpu_mb_max:>8.0f} {r.realtime_ratio_median:>7.2f}x") + for r in results: + if r.error is not None or not r.stage_means_ms: + continue + print() + print(f"=== {r.name} per-stage means (ms) ===") + for stage, mean_ms in r.stage_means_ms.items(): + pct = 100.0 * mean_ms / r.wall_ms_median if r.wall_ms_median > 0 \ + else 0.0 + print(f" {stage:35s} {mean_ms:>9.1f}ms ({pct:5.1f}%)") + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", default=DEFAULT_MODEL) + p.add_argument("--prompt", default=DEFAULT_PROMPT) + p.add_argument("--scenarios", + nargs="+", + default=[s.name for s in DEFAULT_SCENARIOS], + choices=[s.name for s in DEFAULT_SCENARIOS], + help="which scenarios to run") + p.add_argument("--runs", type=int, default=3) + p.add_argument("--num-frames", type=int, default=121) + p.add_argument("--height", type=int, default=1088) + p.add_argument("--width", type=int, default=1920) + p.add_argument("--num-inference-steps", type=int, default=5) + p.add_argument("--num-gpus", type=int, default=1) + p.add_argument("--gpu", type=int, default=None, help="set CUDA_VISIBLE_DEVICES to this single GPU index") + p.add_argument("--seed", type=int, default=10) + p.add_argument("--output-json", default=None, help="write structured results to this path") + args = p.parse_args() + + if args.gpu is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) + + selected = [s for s in DEFAULT_SCENARIOS if s.name in args.scenarios] + print(f"[bench] model={args.model} num_gpus={args.num_gpus}") + print(f"[bench] frames={args.num_frames} {args.width}x{args.height} " + f"steps={args.num_inference_steps} runs/scenario={args.runs}") + print(f"[bench] scenarios={[s.name for s in selected]}") + + results: list[ScenarioResult] = [] + for scenario in selected: + try: + result = benchmark_scenario(scenario, args.model, args.num_gpus, args.prompt, args.runs, args.num_frames, + args.height, args.width, args.num_inference_steps, args.seed) + except Exception as exc: + result = ScenarioResult(name=scenario.name, + enable_compile=scenario.enable_compile, + do_warmup=scenario.do_warmup, + runs=0, + wall_ms_median=0.0, + wall_ms_p95=0.0, + stage_means_ms=OrderedDict(), + realtime_ratio_median=0.0, + peak_gpu_mb_max=0.0, + error=f"{type(exc).__name__}: {exc}") + results.append(result) + + _print_summary(results) + + if args.output_json: + out = [] + for r in results: + entry = asdict(r) + entry["stage_means_ms"] = dict(r.stage_means_ms) + out.append(entry) + with open(args.output_json, "w") as f: + json.dump({"args": vars(args), "results": out}, f, indent=2) + print(f"[bench] wrote {args.output_json}") + + any_below = any(r.error is None and r.realtime_ratio_median < 1.0 for r in results) + return 1 if any_below else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/dreamverse/dreamverse/gpu_pool.py b/apps/dreamverse/dreamverse/gpu_pool.py new file mode 100644 index 000000000..0d1419103 --- /dev/null +++ b/apps/dreamverse/dreamverse/gpu_pool.py @@ -0,0 +1,961 @@ +# pyright: reportMissingTypeArgument=false, reportArgumentType=false, reportOptionalSubscript=false, reportAttributeAccessIssue=false, reportOptionalMemberAccess=false, reportUndefinedVariable=false +# ruff: noqa: UP038, SIM105, F821 +# mypy: ignore-errors +import asyncio +import multiprocessing as mp +import os +import subprocess +import time +import traceback +from dataclasses import dataclass +from enum import Enum +from multiprocessing import Process, Queue + +from dreamverse.config import ( + DEFAULT_MODEL_ID, + MODEL_REGISTRY, + STARTUP_WARMUP_ENABLED, + STARTUP_WARMUP_PROMPT, + STARTUP_WARMUP_TIMEOUT_SECONDS, +) +from dreamverse.av_streaming import ( + SHARED_STREAM_BUFFER_BYTES, + USE_SHARED_STREAM_BUFFER, + StreamChunk, + StreamComplete, + StreamEvent, + StreamInit, + generate_stream_id, + stream_fmp4, +) +from dreamverse.worker_ipc import ( + CommandPayload, + InitAck, + JoinAck, + LeaveAck, + MediaChunk, + MediaComplete, + MediaInit, + ReloadAck, + ReloadModelPayload, + ShutdownAck, + StepComplete, + UserStepPayload, + WarmupComplete, + WarmupPayload, + WorkerError, + WorkerEvent, +) + + +def _parse_requested_gpu_limit() -> int | None: + raw_value = os.getenv("FASTVIDEO_GPU_COUNT", "").strip().lower() + if not raw_value: + return 1 + if raw_value == "all": + return None + try: + requested = int(raw_value) + except ValueError as exc: + raise RuntimeError("Invalid FASTVIDEO_GPU_COUNT. Use a positive integer or 'all'.") from exc + if requested <= 0: + raise RuntimeError("Invalid FASTVIDEO_GPU_COUNT. Use a positive integer or 'all'.") + return requested + + +def _limit_gpu_ids(gpu_ids: list[int]) -> list[int]: + requested_limit = _parse_requested_gpu_limit() + print(f"[INFO] Using GPU limit={requested_limit}") + if requested_limit is None: + return gpu_ids + return gpu_ids[:requested_limit] or gpu_ids + + +class CommandType(Enum): + """Commands sent from main process to GPU worker.""" + INIT = "init" + WARMUP = "warmup" + SHUTDOWN = "shutdown" + USER_JOIN = "user_join" + USER_STEP = "user_step" + USER_LEAVE = "user_leave" + RELOAD_MODEL = "reload_model" + + +@dataclass +class Command: + """Command sent to GPU worker subprocess. + + Commands that carry data (USER_STEP, WARMUP, RELOAD_MODEL) + populate ``payload`` with a typed payload from ``worker_ipc``. + Commands that don't (INIT, SHUTDOWN, USER_JOIN, USER_LEAVE) + leave ``payload`` as ``None``. + """ + type: CommandType + payload: CommandPayload | None = None + user_id: str | None = None + + +def _stream_event_to_worker_event( + event: StreamEvent, + user_id: str, + segment_idx: int, +) -> WorkerEvent: + """Translate an av_streaming event into a typed worker event. + + ``StreamEvent`` (ffmpeg output layer) carries no routing info; + this adds ``user_id`` + ``segment_idx`` so the pool's + ``_response_reader`` can dispatch to the right per-user queue. + """ + match event: + case StreamInit(stream_id=sid, mime=m, uses_shared_buffer=u): + return MediaInit( + user_id=user_id, + segment_idx=segment_idx, + stream_id=sid, + mime=m, + uses_shared_buffer=u, + ) + case StreamChunk( + stream_id=sid, + chunk=c, + chunk_offset=co, + chunk_length=cl, + uses_shared_buffer=u, + ): + return MediaChunk( + user_id=user_id, + segment_idx=segment_idx, + stream_id=sid, + chunk=c, + chunk_offset=co, + chunk_length=cl, + uses_shared_buffer=u, + ) + case StreamComplete(stream_id=sid, chunks=n): + return MediaComplete( + user_id=user_id, + segment_idx=segment_idx, + stream_id=sid, + chunks=n, + ) + case _: + raise ValueError(f"unknown stream event: {type(event).__name__}") + + +def gpu_worker_process( + gpu_id: int, + cuda_device: str, + command_queue: Queue, + response_queue: Queue, + shared_stream_buffer=None, + shared_stream_buffer_bytes: int = 0, +): + """Worker process that runs on a single GPU. + + CUDA_VISIBLE_DEVICES must be set BEFORE importing VideoGenerationWorker + (which transitively touches CUDA). Delegates model lifecycle and + generation to VideoGenerationWorker; AV muxing to av_streaming. + """ + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" + + from dreamverse.video_generation import VideoGenerationWorker + + worker = VideoGenerationWorker(gpu_id) + + def event_loop(first_cmd: Command = None): + """Blocking event loop for LTX2; dispatches user commands.""" + print(f"[GPU {gpu_id}] Entering event loop") + + def handle_command(cmd: Command): + if cmd.type == CommandType.USER_JOIN: + print(f"[GPU {gpu_id}] User {cmd.user_id[:8]} joined") + worker.clear_conditioning() + response_queue.put(JoinAck(user_id=cmd.user_id)) + + elif cmd.type == CommandType.USER_STEP: + try: + assert isinstance(cmd.payload, UserStepPayload), (f"USER_STEP requires UserStepPayload, " + f"got {type(cmd.payload).__name__}") + payload = cmd.payload + segment_idx = payload.segment_idx + step_result = worker.generate_step( + payload.prompt, + segment_idx, + image_path=payload.image_path, + reset_conditioning=payload.reset_conditioning, + ) + head_trim_frames = step_result.head_trim_frames + head_trim_audio_frames = step_result.head_trim_audio_frames + if head_trim_frames > 0 or head_trim_audio_frames > 0: + print(f"[GPU {gpu_id}] Segment {segment_idx}: " + f"trimming video={head_trim_frames} " + f"audio={head_trim_audio_frames} " + f"overlap frames from AV output") + audio_shape = getattr(step_result.audio, "shape", None) + print(f"[GPU {gpu_id}] AV attempt segment " + f"{segment_idx}: " + f"audio_present={step_result.audio is not None}, " + f"audio_shape={audio_shape}, " + f"audio_sample_rate={step_result.audio_sample_rate}") + stream_id = generate_stream_id(segment_idx) + + def _publish(event: StreamEvent) -> None: + response_queue.put(_stream_event_to_worker_event(event, cmd.user_id, segment_idx)) + + av_ok, av_error = stream_fmp4( + frames=step_result.frames, + audio=step_result.audio, + audio_sample_rate=step_result.audio_sample_rate, + stream_id=stream_id, + timings=step_result.timings, + head_trim_frames=head_trim_frames, + head_trim_audio_frames=head_trim_audio_frames, + shared_buffer=shared_stream_buffer, + shared_buffer_bytes=shared_stream_buffer_bytes, + publish=_publish, + log_prefix=f"[GPU {gpu_id}]", + ) + if not av_ok: + raise RuntimeError(av_error or "worker av_fmp4 stream failed") + print(f"[GPU {gpu_id}] AV streamed segment {segment_idx}: " + f"encode_total={step_result.timings.get('av_encode_stream_ms', 0):.0f}ms " + f"wav_write={step_result.timings.get('av_wav_write_ms', 0):.1f}ms " + f"spawn={step_result.timings.get('av_ffmpeg_spawn_ms', 0):.1f}ms " + f"first_chunk={step_result.timings.get('av_first_chunk_ms', 0):.0f}ms " + f"chunk_interval_med={step_result.timings.get('av_chunk_interval_ms_median', 0):.1f}ms " + f"chunk_interval_p95={step_result.timings.get('av_chunk_interval_ms_p95', 0):.1f}ms " + f"publish_med={step_result.timings.get('av_chunk_publish_ms_median', 0):.2f}ms " + f"read_med={step_result.timings.get('av_chunk_read_ms_median', 0):.1f}ms") + step_result.timings["ipc_put_start_ns"] = time.time_ns() + response_queue.put( + StepComplete( + user_id=cmd.user_id, + segment_idx=segment_idx, + timings=step_result.timings, + )) + except Exception as e: + print(f"[GPU {gpu_id}] Step error: {e}") + traceback.print_exc() + response_queue.put(WorkerError( + user_id=cmd.user_id, + message=str(e), + )) + + elif cmd.type == CommandType.USER_LEAVE: + print(f"[GPU {gpu_id}] User {cmd.user_id[:8]} left") + worker.clear_conditioning() + response_queue.put(LeaveAck(user_id=cmd.user_id)) + + elif cmd.type == CommandType.SHUTDOWN: + return False # Signal to exit + + elif cmd.type == CommandType.WARMUP: + try: + assert isinstance(cmd.payload, WarmupPayload), (f"WARMUP requires WarmupPayload, " + f"got {type(cmd.payload).__name__}") + timings = worker.warmup(cmd.payload.prompt) + response_queue.put(WarmupComplete( + user_id=cmd.user_id, + timings=timings, + )) + except Exception as e: + print(f"[GPU {gpu_id}] Warmup error: {e}") + traceback.print_exc() + response_queue.put(WorkerError( + user_id=cmd.user_id, + message=str(e), + )) + + elif cmd.type == CommandType.RELOAD_MODEL: + try: + assert isinstance(cmd.payload, ReloadModelPayload), (f"RELOAD_MODEL requires ReloadModelPayload, " + f"got {type(cmd.payload).__name__}") + worker.initialize(cmd.payload.model_config) + response_queue.put(ReloadAck(user_id=cmd.user_id)) + except Exception as e: + print(f"[GPU {gpu_id}] Reload error: {e}") + traceback.print_exc() + response_queue.put(WorkerError( + user_id=cmd.user_id, + message=str(e), + )) + + return True # Continue loop + + if first_cmd is not None: + if first_cmd.type == CommandType.SHUTDOWN: + worker.shutdown() + response_queue.put(ShutdownAck()) + return + handle_command(first_cmd) + + import queue as queue_module + while True: + try: + cmd = command_queue.get(timeout=1.0) + except queue_module.Empty: + continue + + if cmd.type == CommandType.SHUTDOWN: + print(f"[GPU {gpu_id}] Event loop shutting down") + worker.shutdown() + response_queue.put(ShutdownAck()) + return + + if not handle_command(cmd): + return + + print(f"[GPU {gpu_id}] Worker process starting...") + + try: + while True: + cmd: Command = command_queue.get() + + if cmd.type == CommandType.SHUTDOWN: + print(f"[GPU {gpu_id}] Shutting down...") + worker.shutdown() + response_queue.put(ShutdownAck()) + break + + elif cmd.type == CommandType.INIT: + try: + worker.initialize() + response_queue.put(InitAck(success=True)) + except Exception as e: + print(f"[GPU {gpu_id}] Init error: {e}") + traceback.print_exc() + response_queue.put(InitAck(success=False, error=str(e))) + + elif cmd.type == CommandType.WARMUP: + try: + assert isinstance(cmd.payload, WarmupPayload), (f"WARMUP requires WarmupPayload, " + f"got {type(cmd.payload).__name__}") + timings = worker.warmup(cmd.payload.prompt) + response_queue.put(WarmupComplete( + user_id=cmd.user_id, + timings=timings, + )) + except Exception as e: + print(f"[GPU {gpu_id}] Warmup error: {e}") + traceback.print_exc() + response_queue.put(WorkerError( + user_id=cmd.user_id, + message=str(e), + )) + + elif cmd.type == CommandType.RELOAD_MODEL: + try: + assert isinstance(cmd.payload, ReloadModelPayload), (f"RELOAD_MODEL requires ReloadModelPayload, " + f"got {type(cmd.payload).__name__}") + worker.initialize(cmd.payload.model_config) + response_queue.put(ReloadAck(user_id=cmd.user_id)) + except Exception as e: + print(f"[GPU {gpu_id}] Reload error: {e}") + traceback.print_exc() + response_queue.put(WorkerError( + user_id=cmd.user_id, + message=str(e), + )) + + elif cmd.type in (CommandType.USER_JOIN, CommandType.USER_STEP, CommandType.USER_LEAVE): + event_loop(first_cmd=cmd) + break + + except Exception as e: + print(f"[GPU {gpu_id}] Worker crashed: {e}") + traceback.print_exc() + + print(f"[GPU {gpu_id}] Worker process exiting") + + +class GPUSlot: + """Manages a single GPU worker subprocess.""" + + def __init__(self, gpu_id: int, cuda_device: str): + self.gpu_id = gpu_id + self.cuda_device = cuda_device + self.process: Process | None = None + self.command_queue: Queue | None = None + self.response_queue: Queue | None = None + self.ready: bool = False + self.warmup_enabled: bool = STARTUP_WARMUP_ENABLED + self.warmup_success: bool = False + self.warmup_error: str | None = None + self.warmup_timings: dict[str, float] = {} + self._lock = asyncio.Lock() + + # Client state + self.connected_users: set[str] = set() + self._pending_futures: dict[str, asyncio.Future] = {} + self._stream_queues: dict[str, asyncio.Queue] = {} + 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.shared_stream_buffer = None + self.shared_stream_buffer_size = SHARED_STREAM_BUFFER_BYTES + + @property + def client_count(self) -> int: + return len(self.connected_users) + + @property + def is_available(self) -> bool: + """A GPU is available if it has no active users.""" + alive = self.ready and self.process is not None and self.process.is_alive() + if not alive: + return False + return len(self.connected_users) == 0 + + @property + def is_empty(self) -> bool: + return len(self.connected_users) == 0 + + async def start(self): + """Start the GPU worker subprocess.""" + self.ready = False + self.warmup_success = False + self.warmup_error = None + self.warmup_timings = {} + + ctx = mp.get_context("spawn") + self.command_queue = ctx.Queue() + self.response_queue = ctx.Queue() + if USE_SHARED_STREAM_BUFFER and self.shared_stream_buffer_size > 0: + # Fixed shared byte buffer to avoid per-chunk IPC payload copies. + self.shared_stream_buffer = mp.RawArray("B", self.shared_stream_buffer_size) + + self.process = ctx.Process( + target=gpu_worker_process, + args=( + self.gpu_id, + self.cuda_device, + self.command_queue, + self.response_queue, + self.shared_stream_buffer, + self.shared_stream_buffer_size, + ), + daemon=False, + ) + loop = asyncio.get_event_loop() + await loop.run_in_executor(None, self.process.start) + + # Send init command and wait for response + init_response = await self._send_command(Command(CommandType.INIT), timeout=600.0) + if not isinstance(init_response, InitAck) or not init_response.success: + error_msg = (init_response.error if isinstance(init_response, InitAck) else + f"unexpected init response: {type(init_response).__name__}") + raise RuntimeError(f"GPU {self.gpu_id} failed to initialize: {error_msg}") + + if self.warmup_enabled: + try: + warmup_response = await self._send_command( + Command( + CommandType.WARMUP, + payload=WarmupPayload(prompt=STARTUP_WARMUP_PROMPT), + user_id=f"__warmup_gpu_{self.gpu_id}__", + ), + timeout=float(STARTUP_WARMUP_TIMEOUT_SECONDS), + ) + except Exception as exc: + self.warmup_error = str(exc) + raise RuntimeError(f"GPU {self.gpu_id} warmup failed: {self.warmup_error}") from exc + match warmup_response: + case WarmupComplete(timings=timings): + self.warmup_timings = { + key: float(value) + for key, value in timings.items() if isinstance(value, (int, float)) + } + self.warmup_success = True + case WorkerError(message=msg): + self.warmup_error = msg or "Warmup failed." + raise RuntimeError(f"GPU {self.gpu_id} warmup failed: {self.warmup_error}") + case _: + self.warmup_error = (f"unexpected warmup response: " + f"{type(warmup_response).__name__}") + raise RuntimeError(f"GPU {self.gpu_id} warmup failed: {self.warmup_error}") + else: + print(f"[GPU {self.gpu_id}] Startup warmup disabled by " + "FASTVIDEO_ENABLE_STARTUP_WARMUP") + + self.ready = True + + async def _send_command(self, cmd: Command, timeout: float = 300.0) -> WorkerEvent: + """Send a command and wait for the response or worker death. + + Whichever arrives first wins. A worker dying (exit, signal, OOM, + segfault) flips the kernel-level sentinel fd readable, which + asyncio notices via add_reader — typically within ~10 ms. The + queue timeout still bounds hangs where the worker stays alive but + never replies. + """ + loop = asyncio.get_running_loop() + process = self.process + await loop.run_in_executor(None, self.command_queue.put, cmd) + + response_fut = loop.run_in_executor(None, lambda: self.response_queue.get(timeout=timeout)) + + death_fut: asyncio.Future | None = None + sentinel_fd = process.sentinel if process is not None else None + + if sentinel_fd is not None: + death_fut = loop.create_future() + + def _on_death() -> None: + try: + loop.remove_reader(sentinel_fd) + except (ValueError, OSError): + pass + if death_fut is not None and not death_fut.done(): + death_fut.set_result(None) + + loop.add_reader(sentinel_fd, _on_death) + + waiters = [response_fut] + ([death_fut] if death_fut is not None else []) + try: + done, _ = await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED) + finally: + if (sentinel_fd is not None and death_fut is not None and not death_fut.done()): + try: + loop.remove_reader(sentinel_fd) + except (ValueError, OSError): + pass + death_fut.cancel() + + if (death_fut is not None and death_fut in done and response_fut not in done): + try: + return self.response_queue.get_nowait() + except Exception: + pass + self.ready = False + pid = process.pid if process is not None else "?" + # Reap the process so .exitcode is populated. Sentinel + # readability means the kernel has already exited the + # process; join is non-blocking in practice. + if process is not None: + try: + process.join(timeout=1) + except Exception: + pass + exitcode = process.exitcode if process is not None else None + raise RuntimeError(f"GPU {self.gpu_id} worker died during command " + f"(pid={pid}, exitcode={exitcode})") + + return response_fut.result() + + async def _send_command_tagged(self, cmd: Command, timeout: float = 300.0) -> WorkerEvent: + """Send a tagged command and wait for the matching response. + + Uses the response reader background task to route responses. + """ + loop = asyncio.get_event_loop() + future = loop.create_future() + + # Register this request's pending future + self._pending_futures[cmd.user_id] = future + + # Send command + await loop.run_in_executor(None, self.command_queue.put, cmd) + + # Ensure response reader is running + await self._ensure_response_reader() + + # Wait for the response + try: + response = await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + if (cmd.user_id in self._pending_futures and self._pending_futures[cmd.user_id] is future): + self._pending_futures.pop(cmd.user_id, None) + raise + + return response + + async def _ensure_response_reader(self): + """Start the response reader background task if not running.""" + if self._reader_lock is None: + self._reader_lock = asyncio.Lock() + + async with self._reader_lock: + if (self._response_reader_task is None or self._response_reader_task.done()): + self._response_reader_task = asyncio.create_task(self._response_reader()) + + async def _response_reader(self): + """Background task that reads response queue and routes to per-user futures.""" + loop = asyncio.get_event_loop() + + while self._active: + try: + + def get_response_nonblocking(): + try: + return self.response_queue.get(timeout=0.5) + except Exception: + return None + + event = await loop.run_in_executor(None, get_response_nonblocking) + + if event is None: + continue + + # AV streaming events → route to the user's stream queue. + if isinstance(event, (MediaInit, MediaChunk, MediaComplete)): + stream_queue = self._stream_queues.get(event.user_id) + if stream_queue is not None: + await stream_queue.put(event) + else: + print(f"[GPU {self.gpu_id}] Unmatched stream event for user " + f"{event.user_id[:8]}") + continue + + # System-level acks shouldn't reach the tagged router; they + # belong to the untagged `_send_command` path. + if isinstance(event, (InitAck, ShutdownAck)): + print(f"[GPU {self.gpu_id}] System event leaked into " + f"tagged reader: {type(event).__name__}") + continue + + # Late-mutation of timings for observability. Only events + # that actually carry timings get this annotation. + if isinstance(event, (StepComplete, WarmupComplete)): + event.timings["ipc_get_done_ns"] = time.time_ns() + + user_id = event.user_id + if user_id and user_id in self._pending_futures: + future = self._pending_futures.pop(user_id) + if not future.done(): + future.set_result(event) + else: + print(f"[GPU {self.gpu_id}] Unmatched response for user " + f"{user_id[:8] if user_id else 'None'}") + except asyncio.CancelledError: + break + except Exception as e: + print(f"[GPU {self.gpu_id}] Response reader error: {e}") + await asyncio.sleep(0.01) + + def register_stream_queue(self, user_id: str) -> asyncio.Queue: + """Register stream-event queue for a specific user.""" + queue = asyncio.Queue() + self._stream_queues[user_id] = queue + return queue + + def unregister_stream_queue(self, user_id: str) -> None: + """Remove stream-event queue for a specific user.""" + self._stream_queues.pop(user_id, 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 + + # Reload model if a different one is requested + if model_id != self.current_model_id and model_id in MODEL_REGISTRY: + print(f"[GPU {self.gpu_id}] Model switch: " + f"{self.current_model_id} -> {model_id}") + + for uid, future in list(self._pending_futures.items()): + if not future.done(): + future.set_exception(RuntimeError("Model changed, session reset")) + self._pending_futures.clear() + self._stream_queues.clear() + 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) + match reload_response: + case ReloadAck(): + pass + case WorkerError(message=msg): + raise RuntimeError(f"Model reload failed: {msg}") + case _: + raise RuntimeError(f"Unexpected reload response: " + f"{type(reload_response).__name__}") + + self.current_model_id = model_id + print(f"[GPU {self.gpu_id}] Model reloaded: {model_id}") + + self._active = True + self.connected_users.add(user_id) + try: + response = await self._send_command_tagged(Command(CommandType.USER_JOIN, user_id=user_id), timeout=600.0) + match response: + case JoinAck() as ack: + return ack + case WorkerError(message=msg): + self.connected_users.discard(user_id) + raise RuntimeError(f"User join failed for {user_id[:8]}: {msg}") + case _: + self.connected_users.discard(user_id) + raise RuntimeError(f"Unexpected join response for {user_id[:8]}: " + f"{type(response).__name__}") + except Exception: + self.connected_users.discard(user_id) + raise + + async def user_step( + self, + user_id: str, + prompt: str, + segment_idx: int = 1, + image_path: str | None = None, + reset_conditioning: bool = False, + ) -> dict[str, float]: + """Execute a generation step for a specific user. + + Returns the timings dict. Frames/audio are no longer part of + this return — they stream asynchronously via the AV media + events (MediaInit/MediaChunk/MediaComplete) which the caller + consumes through ``register_stream_queue``. + """ + payload = UserStepPayload( + prompt=prompt, + segment_idx=segment_idx, + image_path=image_path, + reset_conditioning=bool(reset_conditioning), + ) + response = await self._send_command_tagged(Command(CommandType.USER_STEP, payload=payload, user_id=user_id), + timeout=1800.0) + match response: + case StepComplete(timings=timings): + return timings + case WorkerError(message=msg): + raise RuntimeError(f"User step failed for {user_id[:8]}: {msg}") + case _: + raise RuntimeError(f"Unexpected step response for {user_id[:8]}: " + f"{type(response).__name__}") + + async def leave_user(self, user_id: str) -> None: + """Remove a user from this GPU.""" + try: + await self._send_command_tagged(Command(CommandType.USER_LEAVE, user_id=user_id), timeout=30.0) + except Exception as e: + print(f"[GPU {self.gpu_id}] Leave user error: {e}") + finally: + self.connected_users.discard(user_id) + self._pending_futures.pop(user_id, None) + self._stream_queues.pop(user_id, None) + + async def shutdown(self): + """Shutdown the worker subprocess.""" + self._active = False + if self._response_reader_task and not self._response_reader_task.done(): + self._response_reader_task.cancel() + try: + await self._response_reader_task + except asyncio.CancelledError: + pass + + if self.process is not None: + if self.process.is_alive(): + try: + await self._send_command(Command(CommandType.SHUTDOWN), timeout=30.0) + except Exception: + pass + self.process.terminate() + self.process.join(timeout=5) + if self.process.is_alive(): + self.process.kill() + self.process.join(timeout=5) + else: + # Process already died (sentinel path). Reap it so the + # OS releases the PID. + try: + self.process.join(timeout=1) + except Exception: + pass + + for q in (self.command_queue, self.response_queue): + if q is not None: + try: + q.close() + except Exception: + pass + try: + q.join_thread() + except Exception: + pass + + +class GPUPool: + """Manages multiple GPU worker subprocesses.""" + + def __init__(self, gpu_ids: list[int]): + self.gpu_ids = gpu_ids + self.slots: dict[int, GPUSlot] = {gpu_id: GPUSlot(gpu_id, str(gpu_id)) for gpu_id in gpu_ids} + self.waiting_list: list[tuple[str, asyncio.Event, WebSocket]] = [] + self.client_gpu_map: dict[str, int] = {} + self._pool_lock = asyncio.Lock() + + async def initialize(self): + """Initialize all GPU workers and wait for them to be ready. + + Any per-GPU failure aborts startup so uvicorn refuses to serve + traffic with no functional GPUs. + """ + print(f"Initializing GPU pool with {len(self.gpu_ids)} GPUs: {self.gpu_ids}") + await asyncio.gather(*(self._init_gpu(gpu_id) for gpu_id in self.gpu_ids)) + + async def _init_gpu(self, gpu_id: int): + """Initialize a single GPU and assign any waiting clients.""" + try: + await self.slots[gpu_id].start() + except Exception as e: + print(f"GPU {gpu_id} failed to initialize: {e}") + await self.slots[gpu_id].shutdown() + raise + + print(f"GPU pool: {gpu_id} ready " + f"({sum(1 for s in self.slots.values() if s.ready)}/{len(self.gpu_ids)})") + + # Check if anyone is waiting for a GPU + async with self._pool_lock: + slot = self.slots[gpu_id] + if slot.is_available and self.waiting_list: + waiting_client_id, ready_event, _ = self.waiting_list.pop(0) + self.client_gpu_map[waiting_client_id] = gpu_id + print(f"Client {waiting_client_id[:8]} assigned GPU {gpu_id} from queue") + ready_event.set() + + await self._send_queue_updates() + + async def acquire(self, client_id: str, websocket=None) -> tuple[int, GPUSlot]: + """Acquire a GPU slot for a client.""" + async with self._pool_lock: + for gpu_id, slot in self.slots.items(): + if slot.is_available: + self.client_gpu_map[client_id] = gpu_id + print(f"Client {client_id[:8]} acquired GPU {gpu_id}") + return gpu_id, slot + + # No slot available, wait in queue + print(f"Client {client_id[:8]} waiting in queue " + f"(all {len(self.gpu_ids)} GPUs at capacity)") + ready_event = asyncio.Event() + async with self._pool_lock: + self.waiting_list.append((client_id, ready_event, websocket)) + await self._send_queue_updates() + + try: + await ready_event.wait() + except asyncio.CancelledError: + # Client disconnected while queued; remove stale queue entry. + async with self._pool_lock: + self.waiting_list = [ + item for item in self.waiting_list if not (item[0] == client_id and item[1] is ready_event) + ] + self.client_gpu_map.pop(client_id, None) + await self._send_queue_updates() + raise + + gpu_id = self.client_gpu_map.get(client_id) + if gpu_id is None: + raise RuntimeError(f"Client {client_id} was signaled but has no GPU assigned") + + return gpu_id, self.slots[gpu_id] + + async def release(self, client_id: str): + """Release a client from its GPU slot.""" + async with self._pool_lock: + # Remove stale queue entries if the client disconnected while waiting. + prev_wait_len = len(self.waiting_list) + self.waiting_list = [item for item in self.waiting_list if item[0] != client_id] + removed_from_queue = len(self.waiting_list) != prev_wait_len + + gpu_id = self.client_gpu_map.pop(client_id, None) + if gpu_id is None: + if removed_from_queue: + await self._send_queue_updates() + return + + slot = self.slots[gpu_id] + print(f"Client {client_id[:8]} released GPU {gpu_id}") + + try: + await slot.leave_user(client_id) + except Exception as e: + print(f"[GPU {gpu_id}] Leave user failed: {e}") + + # Assign to next waiting client if GPU has capacity + if slot.is_available and self.waiting_list: + waiting_client_id, ready_event, _ = self.waiting_list.pop(0) + self.client_gpu_map[waiting_client_id] = gpu_id + print(f"Client {waiting_client_id[:8]} assigned GPU {gpu_id} from queue") + ready_event.set() + + await self._send_queue_updates() + + async def _send_queue_updates(self): + """Send updated queue positions to all waiting clients.""" + for i, (cid, _, ws) in enumerate(self.waiting_list): + if ws is not None: + try: + await ws.send_json({ + "type": "queue_status", + "position": i + 1, + "total_gpus": len(self.gpu_ids), + "available_gpus": 0, + }) + except Exception: + pass + + async def shutdown(self): + """Shutdown all GPU workers.""" + print("Shutting down GPU pool...") + tasks = [slot.shutdown() for slot in self.slots.values()] + await asyncio.gather(*tasks, return_exceptions=True) + print("GPU pool shutdown complete") + + def get_status(self) -> dict: + """Get the current status of the GPU pool.""" + warmed_count = sum(1 for slot in self.slots.values() if slot.warmup_success) + warmup_failures = sum(1 for slot in self.slots.values() + if slot.warmup_enabled and slot.warmup_error is not None) + return { + "total_gpus": len(self.gpu_ids), + "available_gpus": sum(1 for slot in self.slots.values() if slot.is_available), + "queue_size": len(self.waiting_list), + "warmup_enabled": STARTUP_WARMUP_ENABLED, + "warmup_successful_gpus": warmed_count, + "warmup_failed_gpus": warmup_failures, + "gpu_status": { + gpu_id: { + "ready": slot.ready, + "available": slot.is_available, + "client_count": slot.client_count, + "current_model_id": slot.current_model_id, + "process_alive": (slot.process.is_alive() if slot.process else False), + "warmup_enabled": slot.warmup_enabled, + "warmup_success": slot.warmup_success, + "warmup_error": slot.warmup_error, + "warmup_timings": slot.warmup_timings, + } + for gpu_id, slot in self.slots.items() + } + } + + +def get_available_gpus() -> list[int]: + """Get list of available GPU IDs from environment or auto-detect.""" + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if cuda_visible: + visible_gpu_ids = [int(x.strip()) for x in cuda_visible.split(",") if x.strip()] + return _limit_gpu_ids(visible_gpu_ids) + + # Auto-detect available GPUs + try: + result = subprocess.run(["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], + capture_output=True, + text=True) + if result.returncode == 0: + detected_gpu_ids = [int(x.strip()) for x in result.stdout.strip().split("\n") if x.strip()] + print(f"Auto-detected GPU IDs: {detected_gpu_ids}") + return _limit_gpu_ids(detected_gpu_ids) + except Exception: + pass + + return _limit_gpu_ids([0]) diff --git a/apps/dreamverse/dreamverse/mock_server.py b/apps/dreamverse/dreamverse/mock_server.py new file mode 100644 index 000000000..6c58277c3 --- /dev/null +++ b/apps/dreamverse/dreamverse/mock_server.py @@ -0,0 +1,1164 @@ +# pyright: reportMissingTypeArgument=false, reportConstantRedefinition=false, reportGeneralTypeIssues=false +# ruff: noqa: UP022, F841, SIM105 +# mypy: ignore-errors +""" +Protocol-compatible mock server for the current LTX2 streaming UI. + +This mock speaks the same websocket API shape as `server/main.py` for: +- `session_init_v2` +- `append_prompt` +- `simple_generate` +- prompt window reset / rewrite control messages + +It streams a small fragmented MP4 clip generated by ffmpeg so the frontend can +exercise the real AV pipeline without needing a GPU backend. +""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +import json +import logging +import os +import shutil +import subprocess +import uuid + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.middleware.cors import CORSMiddleware +from fastapi.staticfiles import StaticFiles + +from dreamverse._deps import require_dreamverse_runtime_deps +from dreamverse.config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP +from dreamverse.session_init_image import cleanup_session_init_image, persist_session_init_image + +LATENCY_MS = 200 +SESSION_TIMEOUT_SECONDS = 300 +MOCK_FRAME_WIDTH = 640 +MOCK_FRAME_HEIGHT = 352 +MOCK_FPS = 24 +MOCK_DURATION_SECONDS = 5.0 +MOCK_STREAM_CHUNK_SIZE_BYTES = 256 * 1024 +MOCK_STREAM_CHUNK_DELAY_MS = 15 +MOCK_AV_MIME = 'video/mp4; codecs="avc1.42E01E"' +FFMPEG_BIN = shutil.which(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg")) +MOCK_SEGMENT_BYTES: bytes | None = None + + +class _HeartbeatAccessLogFilter(logging.Filter): + """Drop noisy access logs for frequent health/readiness probes.""" + + def filter(self, record: logging.LogRecord) -> bool: + message = record.getMessage() + return ('"GET /healthz ' not in message and '"GET /readyz ' not in message) + + +def _install_heartbeat_log_filter() -> None: + access_logger = logging.getLogger("uvicorn.access") + for existing in access_logger.filters: + if isinstance(existing, _HeartbeatAccessLogFilter): + return + access_logger.addFilter(_HeartbeatAccessLogFilter()) + + +def _build_mock_segment_bytes() -> bytes: + if not FFMPEG_BIN: + raise RuntimeError("ffmpeg is required to run the mock streaming server.") + + command = [ + FFMPEG_BIN, + "-hide_banner", + "-loglevel", + "error", + "-f", + "lavfi", + "-i", + f"testsrc2=size={MOCK_FRAME_WIDTH}x{MOCK_FRAME_HEIGHT}:rate={MOCK_FPS}", + "-t", + f"{MOCK_DURATION_SECONDS}", + "-an", + "-c:v", + "libx264", + "-preset", + "ultrafast", + "-tune", + "zerolatency", + "-pix_fmt", + "yuv420p", + "-profile:v", + "baseline", + "-level", + "3.0", + "-movflags", + "+empty_moov+default_base_moof+frag_keyframe", + "-frag_duration", + "250000", + "-f", + "mp4", + "pipe:1", + ] + result = subprocess.run( + command, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + if not result.stdout: + raise RuntimeError("Mock ffmpeg command did not return MP4 bytes.") + return result.stdout + + +class PromptSubmission: + + def __init__(self, prompt_id: str, raw_prompt: str): + self.prompt_id = prompt_id + self.raw_prompt = raw_prompt + + +class ReadyPrompt: + + def __init__( + self, + *, + prompt: str, + source: str, + prompt_id: str | None = None, + seed_prompt_index: int | None = None, + loop_iteration: int | None = None, + ): + self.prompt = prompt + self.source = source + self.prompt_id = prompt_id + self.seed_prompt_index = seed_prompt_index + self.loop_iteration = loop_iteration + + +@asynccontextmanager +async def lifespan(app: FastAPI): + global MOCK_SEGMENT_BYTES + del app + print(f"Mock server starting with {LATENCY_MS}ms simulated latency...", flush=True) + MOCK_SEGMENT_BYTES = _build_mock_segment_bytes() + print( + "Mock server ready " + f"({len(MOCK_SEGMENT_BYTES) / (1024 * 1024):.2f} MiB fMP4 clip)", + flush=True, + ) + yield + MOCK_SEGMENT_BYTES = None + print("Mock server shutting down...", flush=True) + + +app = FastAPI(lifespan=lifespan) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +@app.get("/healthz") +async def healthz(): + return { + "status": "ok", + "service": "ltx2-streaming-mock-server", + } + + +@app.get("/readyz") +async def readyz(): + return { + "status": "ready", + "service": "ltx2-streaming-mock-server", + } + + +@app.get("/status") +async def status(): + return { + "queue_size": 0, + "available_gpus": 1, + "total_gpus": 1, + } + + +async def _stream_mock_segment( + ws_send_json, + ws_send_bytes, + *, + segment_idx: int, +) -> None: + if not MOCK_SEGMENT_BYTES: + raise RuntimeError("Mock segment bytes were not initialized.") + + stream_id = f"mock-stream-{segment_idx}-{uuid.uuid4().hex[:8]}" + await ws_send_json({ + "type": "media_init", + "segment_idx": segment_idx, + "mime": MOCK_AV_MIME, + "stream_id": stream_id, + "mode": "av_fmp4", + }) + + chunk_count = 0 + for start in range(0, len(MOCK_SEGMENT_BYTES), MOCK_STREAM_CHUNK_SIZE_BYTES): + end = start + MOCK_STREAM_CHUNK_SIZE_BYTES + await ws_send_bytes(MOCK_SEGMENT_BYTES[start:end]) + chunk_count += 1 + if MOCK_STREAM_CHUNK_DELAY_MS > 0: + await asyncio.sleep(MOCK_STREAM_CHUNK_DELAY_MS / 1000.0) + + await ws_send_json({ + "type": "media_segment_complete", + "segment_idx": segment_idx, + "stream_id": stream_id, + "chunks": chunk_count, + }) + + +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + + client_id = str(uuid.uuid4()) + print(f"Client {client_id[:8]} connected") + + send_lock = asyncio.Lock() + stop_event = asyncio.Event() + session_init_image = None + + async def ws_send_json(payload: dict) -> None: + async with send_lock: + await websocket.send_json(payload) + + async def ws_send_bytes(payload: bytes) -> None: + async with send_lock: + await websocket.send_bytes(payload) + + async def session_timeout() -> None: + await asyncio.sleep(SESSION_TIMEOUT_SECONDS) + if stop_event.is_set(): + return + stop_event.set() + try: + await ws_send_json({ + "type": "session_timeout", + "message": f"Session expired after {SESSION_TIMEOUT_SECONDS} seconds", + }) + await websocket.close(code=1000, reason="Session timeout") + except Exception: + pass + + try: + await ws_send_json({ + "type": "queue_status", + "position": 0, + "total_gpus": 1, + "available_gpus": 1, + }) + + try: + init_data = await asyncio.wait_for(websocket.receive_json(), timeout=10.0) + except asyncio.TimeoutError: + init_data = {} + + init_type = init_data.get("type") + preset_id = init_data.get("preset_id") + preset_label = str(init_data.get("preset_label") or "").strip() + initial_rollout_prompt = str(init_data.get("initial_rollout_prompt") or "").strip() + enhancement_enabled = bool(init_data.get("enhancement_enabled", True)) + auto_extension_enabled = bool(init_data.get("auto_extension_enabled", False)) + loop_generation_enabled = bool(init_data.get("loop_generation_enabled", False)) + single_clip_mode = bool(init_data.get("single_clip_mode", False)) + generation_paused = False + + if init_type == "session_init_v2": + incoming_prompts = init_data.get("curated_prompts", []) + elif init_type == "select_model": + incoming_prompts = init_data.get("segment_prompts", []) + else: + incoming_prompts = [] + + curated_prompts = [prompt.strip() for prompt in incoming_prompts if isinstance(prompt, str) and prompt.strip()] + generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0) + + try: + session_init_image = persist_session_init_image(init_data.get("initial_image")) + except ValueError as exc: + await ws_send_json({ + "type": "error", + "message": str(exc), + }) + await websocket.close(code=1003, reason="Invalid initial image") + return + + timeout_task = asyncio.create_task(session_timeout()) + + await ws_send_json({ + "type": "gpu_assigned", + "gpu_id": 0, + "session_timeout": SESSION_TIMEOUT_SECONDS, + }) + + raw_prompt_queue: asyncio.Queue[PromptSubmission] = asyncio.Queue() + ready_prompt_queue: asyncio.Queue[ReadyPrompt] = asyncio.Queue() + seed_prompt_memory = list(curated_prompts) + curated_idx = 0 + segment_idx = 0 + loop_iteration = 0 if generation_paused else 1 + prompt_sources_blocked = False + pending_seed_reset = False + pending_seed_reset_reason = "" + pending_simple_submission: PromptSubmission | None = None + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + initial_rollout_waiting_for_rewrite = generation_paused + segment_generation_active = False + rewrite_restart_pending = False + project_active = True + project_stream_started = False + pending_project_end = False + + def replace_session_image(initial_image_payload: object) -> None: + nonlocal session_init_image + next_session_image = persist_session_init_image(initial_image_payload) + previous_session_image = session_init_image + session_init_image = next_session_image + if previous_session_image is not None: + cleanup_session_init_image(previous_session_image) + + async def send_stream_start(seed_reason: str) -> None: + await ws_send_json({ + "type": "ltx2_stream_start", + "total_segments": len(curated_prompts), + "preset_id": preset_id, + "stream_mode": "av_fmp4", + "live_mode": True, + "loop_generation_enabled": loop_generation_enabled, + "loop_iteration": loop_iteration, + "generation_segment_cap": 0, + }) + if seed_reason == "init": + await ws_send_json({ + "type": "seed_prompts_updated", + "prompts": seed_prompt_memory, + "preset_id": preset_id, + "preset_label": preset_label, + "reason": "init", + "seed_prompt_count": len(seed_prompt_memory), + }) + await ws_send_json({ + "type": "loop_generation_updated", + "enabled": loop_generation_enabled, + }) + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + else: + await ws_send_json({ + "type": "seed_prompts_reset_applied", + "reason": seed_reason or "manual_reset", + "loop_iteration": loop_iteration, + "seed_prompt_count": len(seed_prompt_memory), + "dropped_user_raw_queue": 0, + "dropped_user_ready_queue": 0, + }) + + async def promote_submission_to_ready(submission: PromptSubmission, ) -> None: + nonlocal prompt_sources_blocked + + await ws_send_json({ + "type": "prompt_received", + "prompt_id": submission.prompt_id, + "queue_depth": raw_prompt_queue.qsize() + 1, + }) + + if enhancement_enabled: + await ws_send_json({ + "type": "prompt_enhancing", + "prompt_id": submission.prompt_id, + }) + await asyncio.sleep(min(LATENCY_MS, 120) / 1000.0) + final_prompt = submission.raw_prompt + await ready_prompt_queue.put( + ReadyPrompt( + prompt=final_prompt, + source="user_enhanced", + prompt_id=submission.prompt_id, + loop_iteration=loop_iteration, + )) + await ws_send_json({ + "type": "prompt_ready", + "prompt_id": submission.prompt_id, + "prompt": final_prompt, + "source": "user_enhanced", + "latency_ms": float(min(LATENCY_MS, 120)), + }) + else: + final_prompt = submission.raw_prompt + await ready_prompt_queue.put( + ReadyPrompt( + prompt=final_prompt, + source="user_raw", + prompt_id=submission.prompt_id, + loop_iteration=loop_iteration, + )) + await ws_send_json({ + "type": "prompt_ready", + "prompt_id": submission.prompt_id, + "prompt": final_prompt, + "source": "user_raw", + "latency_ms": 0.0, + }) + + if prompt_sources_blocked: + prompt_sources_blocked = False + await ws_send_json({ + "type": "prompt_sources_resumed", + "segment_idx": segment_idx + 1, + "source": "user_prompt", + "prompt_id": submission.prompt_id, + }) + + async def apply_project_init_payload(payload: dict[str, object], ) -> bool: + nonlocal preset_id + nonlocal preset_label + nonlocal initial_rollout_prompt + nonlocal enhancement_enabled + nonlocal auto_extension_enabled + nonlocal loop_generation_enabled + nonlocal single_clip_mode + nonlocal generation_paused + nonlocal seed_prompt_memory + nonlocal curated_prompts + nonlocal curated_idx + nonlocal segment_idx + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal pending_simple_submission + nonlocal single_clip_waiting_for_request + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal loop_iteration + nonlocal project_active + nonlocal project_stream_started + nonlocal pending_project_end + + next_initial_rollout_prompt = str(payload.get("initial_rollout_prompt") or "").strip() + next_preset_id = str(payload.get("preset_id") or "").strip() + next_preset_label = str(payload.get("preset_label") or "").strip() + next_curated_prompts = [ + prompt.strip() for prompt in payload.get("curated_prompts", []) + if isinstance(prompt, str) and prompt.strip() + ] + + try: + replace_session_image(payload.get("initial_image")) + except ValueError as exc: + await ws_send_json({ + "type": "error", + "message": str(exc), + }) + return False + + if next_preset_id: + preset_id = next_preset_id + if next_preset_label: + preset_label = next_preset_label + + initial_rollout_prompt = next_initial_rollout_prompt + enhancement_enabled = bool(payload.get("enhancement_enabled", True)) + auto_extension_enabled = bool(payload.get("auto_extension_enabled", False)) + loop_generation_enabled = bool(payload.get("loop_generation_enabled", False)) + single_clip_mode = bool(payload.get("single_clip_mode", False)) + + seed_prompt_memory = list(next_curated_prompts) + curated_prompts = list(seed_prompt_memory) + curated_idx = 0 + segment_idx = 0 + pending_seed_reset = False + pending_seed_reset_reason = "" + pending_simple_submission = None + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + generation_paused = bool(initial_rollout_prompt and not single_clip_mode and len(curated_prompts) == 0) + initial_rollout_waiting_for_rewrite = generation_paused + rewrite_restart_pending = False + loop_iteration = 0 + project_active = True + project_stream_started = False + pending_project_end = False + + return True + + def drain_queue_nowait(queue: asyncio.Queue) -> int: + drained = 0 + while True: + try: + queue.get_nowait() + drained += 1 + except asyncio.QueueEmpty: + break + return drained + + async def enter_project_idle() -> None: + nonlocal seed_prompt_memory + nonlocal curated_prompts + nonlocal curated_idx + nonlocal segment_idx + nonlocal generation_paused + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal pending_simple_submission + nonlocal single_clip_waiting_for_request + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal initial_rollout_prompt + nonlocal project_active + nonlocal project_stream_started + nonlocal pending_project_end + + dropped_raw = drain_queue_nowait(raw_prompt_queue) + dropped_ready = drain_queue_nowait(ready_prompt_queue) + seed_prompt_memory = [] + curated_prompts = [] + curated_idx = 0 + segment_idx = 0 + generation_paused = False + pending_seed_reset = False + pending_seed_reset_reason = "" + pending_simple_submission = None + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + initial_rollout_waiting_for_rewrite = False + rewrite_restart_pending = False + initial_rollout_prompt = "" + project_active = False + pending_project_end = False + + if project_stream_started: + project_stream_started = False + await ws_send_json({"type": "ltx2_stream_complete"}) + + await ws_send_json({ + "type": "project_idle", + "dropped_user_raw_queue": dropped_raw, + "dropped_user_ready_queue": dropped_ready, + }) + + async def reader_loop() -> None: + nonlocal preset_id + nonlocal preset_label + nonlocal enhancement_enabled + nonlocal auto_extension_enabled + nonlocal loop_generation_enabled + nonlocal single_clip_mode + nonlocal generation_paused + nonlocal seed_prompt_memory + nonlocal curated_prompts + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal pending_simple_submission + nonlocal prompt_sources_blocked + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal rewrite_restart_pending + nonlocal project_active + nonlocal pending_project_end + while not stop_event.is_set(): + try: + data = await websocket.receive_json() + except WebSocketDisconnect: + stop_event.set() + return + except Exception: + stop_event.set() + return + + msg_type = data.get("type") + if msg_type == "leave": + stop_event.set() + return + + if msg_type == "end_project_keep_session": + loop_generation_enabled = False + generation_paused = False + rollout_waiting_for_rewrite = False + initial_rollout_waiting_for_rewrite = False + rewrite_restart_pending = False + pending_project_end = True + continue + + if msg_type == "project_init_v1": + if project_active or pending_project_end: + await ws_send_json({ + "type": + "error", + "message": ("Current project is still ending. " + "Wait for the reset to finish before starting another project."), + }) + continue + init_applied = await apply_project_init_payload(data) + if not init_applied: + continue + if initial_rollout_waiting_for_rewrite: + await ws_send_json({ + "type": "rewrite_seed_prompts_started", + "model": "mock-rewrite-model", + }) + seed_prompt_memory = [f"{initial_rollout_prompt} [segment {index}]" for index in range(1, 7)] + curated_prompts = list(seed_prompt_memory) + await ws_send_json({ + "type": + "seed_prompts_updated", + "prompts": + seed_prompt_memory, + "preset_id": + preset_id, + "preset_label": + preset_label, + "reason": + "rewrite", + "fallback_used": + False, + "error": + None, + "model": + "mock-rewrite-model", + "latency_ms": + float(min(LATENCY_MS, 150)), + "raw_llm_output": + json.dumps({ + "id": preset_id or "current_rollout", + "label": preset_label or "Current rollout", + "segment_prompts": seed_prompt_memory, + }), + }) + await ws_send_json({ + "type": "rewrite_seed_prompts_complete", + "fallback_used": False, + "error": None, + "preset_id": preset_id, + "preset_label": preset_label, + "model": "mock-rewrite-model", + "latency_ms": float(min(LATENCY_MS, 150)), + }) + pending_seed_reset = True + pending_seed_reset_reason = "initial_rewrite" + generation_paused = False + initial_rollout_waiting_for_rewrite = False + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + else: + pending_seed_reset = True + pending_seed_reset_reason = "project_init" + continue + + if not project_active: + if msg_type in { + "append_prompt", + "rewrite_seed_prompts", + "reset_to_seed_prompts", + "restart_generation", + }: + await ws_send_json({ + "type": "error", + "message": "Start a new project before sending prompts.", + }) + continue + + if msg_type == "append_prompt": + if ((rollout_waiting_for_rewrite or initial_rollout_waiting_for_rewrite) and not single_clip_mode): + await ws_send_json({ + "type": "error", + "message": ("Use Rewrite to create or restart the rollout."), + }) + continue + raw_prompt = str(data.get("prompt", "")).strip() + if not raw_prompt: + await ws_send_json({ + "type": "error", + "message": "append_prompt requires a non-empty prompt", + }) + continue + await promote_submission_to_ready( + PromptSubmission( + prompt_id=str(data.get("prompt_id") or uuid.uuid4()), + raw_prompt=raw_prompt, + )) + continue + + if msg_type == "simple_generate": + raw_prompt = str(data.get("prompt", "")).strip() + if not raw_prompt: + await ws_send_json({ + "type": "error", + "message": "simple_generate requires a non-empty prompt", + }) + continue + + try: + replace_session_image(data.get("initial_image")) + except ValueError as exc: + await ws_send_json({ + "type": "error", + "message": str(exc), + }) + continue + + next_preset_id = str(data.get("preset_id") or "").strip() + if next_preset_id: + preset_id = next_preset_id + next_preset_label = str(data.get("preset_label") or "").strip() + if next_preset_label: + preset_label = next_preset_label + + enhancement_enabled = bool(data.get("enhancement_enabled", True)) + auto_extension_enabled = False + loop_generation_enabled = False + single_clip_mode = True + generation_paused = False + pending_seed_reset = True + pending_seed_reset_reason = "simple_generate" + seed_prompt_memory = ([] if enhancement_enabled else [raw_prompt]) + curated_prompts = list(seed_prompt_memory) + pending_simple_submission = (PromptSubmission( + prompt_id=str(data.get("prompt_id") or uuid.uuid4()), + raw_prompt=raw_prompt, + ) if enhancement_enabled else None) + continue + + if msg_type == "rewrite_seed_prompts": + rewrite_instruction = str(data.get("rewrite_instruction", "")).strip() + if len(seed_prompt_memory) == 0 and not (initial_rollout_waiting_for_rewrite + and rewrite_instruction): + await ws_send_json({ + "type": "error", + "message": "No seed prompts available to rewrite.", + }) + continue + if (segment_generation_active and not single_clip_mode and not rollout_waiting_for_rewrite + and not initial_rollout_waiting_for_rewrite): + rewrite_restart_pending = True + await ws_send_json({ + "type": "rewrite_seed_prompts_started", + "model": "mock-rewrite-model", + }) + rollout_id = preset_id or "current_rollout" + rollout_label = preset_label or "Current rollout" + if seed_prompt_memory: + rewritten_prompts = [ + f"{prompt} [{rewrite_instruction or 'mock rewrite'}]" for prompt in seed_prompt_memory + ] + else: + base_prompt = rewrite_instruction or "mock rewrite" + rewritten_prompts = [f"{base_prompt} [segment {index}]" for index in range(1, 7)] + seed_prompt_memory = rewritten_prompts + curated_prompts = list(seed_prompt_memory) + preset_id = rollout_id + preset_label = rollout_label + if initial_rollout_waiting_for_rewrite: + pending_seed_reset = True + pending_seed_reset_reason = "initial_rewrite" + generation_paused = False + initial_rollout_waiting_for_rewrite = False + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + if rollout_waiting_for_rewrite: + pending_seed_reset = True + pending_seed_reset_reason = "rewrite_rollout" + generation_paused = False + elif rewrite_restart_pending: + pending_seed_reset = True + pending_seed_reset_reason = "rewrite_during_generation" + generation_paused = False + rewrite_restart_pending = False + await ws_send_json({ + "type": + "seed_prompts_updated", + "prompts": + seed_prompt_memory, + "preset_id": + preset_id, + "preset_label": + preset_label, + "reason": + "rewrite", + "fallback_used": + False, + "error": + None, + "model": + "mock-rewrite-model", + "latency_ms": + float(min(LATENCY_MS, 150)), + "raw_llm_output": + json.dumps({ + "id": preset_id, + "label": preset_label, + "segment_prompts": seed_prompt_memory, + }), + }) + await ws_send_json({ + "type": "rewrite_seed_prompts_complete", + "fallback_used": False, + "error": None, + "preset_id": preset_id, + "preset_label": preset_label, + "model": "mock-rewrite-model", + "latency_ms": float(min(LATENCY_MS, 150)), + }) + continue + + if msg_type == "reset_to_seed_prompts": + pending_seed_reset = True + pending_seed_reset_reason = "manual_reset" + continue + + if msg_type == "restart_generation": + pending_seed_reset = True + pending_seed_reset_reason = "cap_restart" + rollout_waiting_for_rewrite = False + await ws_send_json({ + "type": "generation_restarted", + "reason": "cap_restart", + "segment_cap": 0, + }) + continue + + if msg_type == "set_enhancement": + enhancement_enabled = bool(data.get("enabled", enhancement_enabled)) + await ws_send_json({ + "type": "enhancement_updated", + "enabled": enhancement_enabled, + }) + continue + + if msg_type == "set_auto_extension": + auto_extension_enabled = bool(data.get("enabled", auto_extension_enabled)) + await ws_send_json({ + "type": "auto_extension_updated", + "enabled": auto_extension_enabled, + }) + continue + + if msg_type == "set_loop_generation": + loop_generation_enabled = bool(data.get("enabled", loop_generation_enabled)) + await ws_send_json({ + "type": "loop_generation_updated", + "enabled": loop_generation_enabled, + }) + continue + + if msg_type == "set_generation_paused": + generation_paused = bool(data.get("paused", generation_paused)) + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + continue + + async def generation_loop() -> None: + nonlocal curated_idx + nonlocal curated_prompts + nonlocal segment_idx + nonlocal loop_iteration + nonlocal pending_seed_reset + nonlocal pending_seed_reset_reason + nonlocal pending_simple_submission + nonlocal prompt_sources_blocked + nonlocal single_clip_waiting_for_request + nonlocal rollout_waiting_for_rewrite + nonlocal initial_rollout_waiting_for_rewrite + nonlocal segment_generation_active + nonlocal rewrite_restart_pending + nonlocal project_active + nonlocal project_stream_started + nonlocal pending_project_end + while not stop_event.is_set(): + if pending_project_end: + await enter_project_idle() + await asyncio.sleep(0) + continue + + if not project_active: + await asyncio.sleep(0.05) + continue + + if pending_seed_reset: + curated_prompts = list(seed_prompt_memory) + curated_idx = 0 + segment_idx = 0 + pending_seed_reset = False + loop_iteration += 1 + prompt_sources_blocked = False + single_clip_waiting_for_request = False + rollout_waiting_for_rewrite = False + project_stream_started = True + await send_stream_start(pending_seed_reset_reason) + pending_seed_reset_reason = "" + if pending_simple_submission is not None: + submission = pending_simple_submission + pending_simple_submission = None + await promote_submission_to_ready(submission) + + if generation_paused: + await asyncio.sleep(0.05) + continue + + if rewrite_restart_pending: + await asyncio.sleep(0.05) + continue + + if (not single_clip_mode and not rollout_waiting_for_rewrite and GENERATION_SEGMENT_CAP > 0 + and segment_idx >= GENERATION_SEGMENT_CAP): + rollout_waiting_for_rewrite = True + loop_generation_enabled = False + project_stream_started = False + await ws_send_json({ + "type": "ltx2_stream_complete", + }) + await asyncio.sleep(0.05) + continue + + if rollout_waiting_for_rewrite: + await asyncio.sleep(0.05) + continue + + if initial_rollout_waiting_for_rewrite: + await asyncio.sleep(0.05) + continue + + selected = None + try: + selected = ready_prompt_queue.get_nowait() + except asyncio.QueueEmpty: + if curated_idx < len(curated_prompts): + prompt_idx = curated_idx + selected = ReadyPrompt( + prompt=curated_prompts[prompt_idx], + source="curated", + seed_prompt_index=prompt_idx, + loop_iteration=loop_iteration, + ) + curated_idx += 1 + + if selected is None: + if single_clip_mode: + await asyncio.sleep(0.05) + continue + if not prompt_sources_blocked: + prompt_sources_blocked = True + await ws_send_json({ + "type": "prompt_sources_blocked", + "segment_idx": segment_idx + 1, + }) + await asyncio.sleep(0.05) + continue + + if prompt_sources_blocked: + prompt_sources_blocked = False + await ws_send_json({ + "type": "prompt_sources_resumed", + "segment_idx": segment_idx + 1, + "source": selected.source, + "prompt_id": selected.prompt_id, + }) + + segment_idx += 1 + single_clip_waiting_for_request = False + total_segments_hint = max(segment_idx, len(curated_prompts)) + await ws_send_json({ + "type": "segment_prompt_source", + "segment_idx": segment_idx, + "source": selected.source, + "prompt_id": selected.prompt_id, + "fallback_used": False, + "seed_prompt_index": selected.seed_prompt_index, + "loop_iteration": selected.loop_iteration or loop_iteration, + }) + await ws_send_json({ + "type": "ltx2_segment_start", + "segment_idx": segment_idx, + "total_segments": total_segments_hint, + "prompt": selected.prompt, + "source": selected.source, + "seed_prompt_index": selected.seed_prompt_index, + "loop_iteration": selected.loop_iteration or loop_iteration, + }) + + segment_generation_active = True + await asyncio.sleep(LATENCY_MS / 1000.0) + await _stream_mock_segment( + ws_send_json, + ws_send_bytes, + segment_idx=segment_idx, + ) + await ws_send_json({ + "type": "step_complete", + "latency_ms": { + "total": round(LATENCY_MS + 120.0, 2), + "worker_e2e": round(float(LATENCY_MS), 2), + "main_user_step": round(LATENCY_MS + 120.0, 2), + "overhead": 120.0, + }, + }) + await ws_send_json({ + "type": "ltx2_segment_complete", + "segment_idx": segment_idx, + "total_segments": total_segments_hint, + }) + segment_generation_active = False + if single_clip_mode and not single_clip_waiting_for_request: + single_clip_waiting_for_request = True + project_stream_started = False + await ws_send_json({ + "type": "ltx2_stream_complete", + }) + + reader_task = asyncio.create_task(reader_loop()) + generation_task = asyncio.create_task(generation_loop()) + await asyncio.sleep(0) + + await ws_send_json({ + "type": "loop_generation_updated", + "enabled": loop_generation_enabled, + }) + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + if pending_project_end: + await enter_project_idle() + elif initial_rollout_waiting_for_rewrite: + await ws_send_json({ + "type": "rewrite_seed_prompts_started", + "model": "mock-rewrite-model", + }) + seed_prompt_memory = [f"{initial_rollout_prompt} [segment {index}]" for index in range(1, 7)] + curated_prompts = list(seed_prompt_memory) + await ws_send_json({ + "type": + "seed_prompts_updated", + "prompts": + seed_prompt_memory, + "preset_id": + preset_id, + "preset_label": + preset_label, + "reason": + "rewrite", + "fallback_used": + False, + "error": + None, + "model": + "mock-rewrite-model", + "latency_ms": + float(min(LATENCY_MS, 150)), + "raw_llm_output": + json.dumps({ + "id": preset_id or "current_rollout", + "label": preset_label or "Current rollout", + "segment_prompts": seed_prompt_memory, + }), + }) + await ws_send_json({ + "type": "rewrite_seed_prompts_complete", + "fallback_used": False, + "error": None, + "preset_id": preset_id, + "preset_label": preset_label, + "model": "mock-rewrite-model", + "latency_ms": float(min(LATENCY_MS, 150)), + }) + pending_seed_reset = True + pending_seed_reset_reason = "initial_rewrite" + generation_paused = False + initial_rollout_waiting_for_rewrite = False + await ws_send_json({ + "type": "generation_paused_updated", + "paused": generation_paused, + }) + else: + project_stream_started = True + await send_stream_start("init") + + done, pending = await asyncio.wait( + {reader_task, generation_task, timeout_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + for task in done: + try: + await task + except asyncio.CancelledError: + pass + except WebSocketDisconnect: + print(f"Client {client_id[:8]} disconnected") + except Exception as exc: + print(f"Client {client_id[:8]} error: {exc}") + try: + await websocket.send_json({ + "type": "error", + "message": f"Mock streaming failed: {exc}", + }) + except Exception: + pass + import traceback + traceback.print_exc() + finally: + stop_event.set() + cleanup_session_init_image(session_init_image) + + +for static_dir in FRONTEND_STATIC_DIR_CANDIDATES: + if os.path.isdir(static_dir): + app.mount("/", StaticFiles(directory=static_dir, html=True), name="static") + break + + +def cli() -> None: + global LATENCY_MS + + import argparse + import uvicorn + + require_dreamverse_runtime_deps() + + parser = argparse.ArgumentParser(description="Protocol-compatible mock server for the LTX2 UI") + parser.add_argument( + "--latency", + type=int, + default=200, + help="Simulated model latency in milliseconds", + ) + parser.add_argument( + "--port", + type=int, + default=8009, + help="Server port", + ) + args = parser.parse_args() + + LATENCY_MS = args.latency + print(f"Starting mock server with {LATENCY_MS}ms latency on port {args.port}") + + _install_heartbeat_log_filter() + uvicorn.run(app, host="0.0.0.0", port=args.port) + + +if __name__ == "__main__": + cli() diff --git a/apps/dreamverse/dreamverse/tests/test_gpu_pool.py b/apps/dreamverse/dreamverse/tests/test_gpu_pool.py new file mode 100644 index 000000000..e30814539 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_gpu_pool.py @@ -0,0 +1,121 @@ +# pyright: reportAttributeAccessIssue=false +from __future__ import annotations + +import asyncio +import multiprocessing as mp +from types import SimpleNamespace + +import pytest + + +import dreamverse.gpu_pool as gpu_pool + + +def _child_consume_and_exit(cmd_q, resp_q): + """Top-level so the spawn context can pickle it. + + Signals startup by putting "READY" on resp_q (so the parent can + wait out spawn-import latency separately from the actual test + assertion), then consumes one command from cmd_q and exits without + putting anything else on resp_q. Simulates a worker that dies + mid-command — e.g. SIGQUIT'd after a fatal pipeline error. + """ + resp_q.put("READY") + cmd_q.get() + + +def test_get_available_gpus_defaults_to_single_detected_gpu(monkeypatch): + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("FASTVIDEO_GPU_COUNT", raising=False) + + def fake_run(*args, **kwargs): + del args, kwargs + return SimpleNamespace(returncode=0, stdout="0\n1\n2\n") + + monkeypatch.setattr(gpu_pool.subprocess, "run", fake_run) + + assert gpu_pool.get_available_gpus() == [0] + + +def test_get_available_gpus_respects_explicit_gpu_count(monkeypatch): + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "2") + + def fake_run(*args, **kwargs): + del args, kwargs + return SimpleNamespace(returncode=0, stdout="0\n1\n2\n") + + monkeypatch.setattr(gpu_pool.subprocess, "run", fake_run) + + assert gpu_pool.get_available_gpus() == [0, 1] + + +def test_get_available_gpus_can_use_all_visible_devices(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,5,7") + monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "all") + + assert gpu_pool.get_available_gpus() == [3, 5, 7] + + +def test_get_available_gpus_defaults_to_first_visible_device(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,5,7") + monkeypatch.delenv("FASTVIDEO_GPU_COUNT", raising=False) + + assert gpu_pool.get_available_gpus() == [3] + + +def test_get_available_gpus_rejects_invalid_gpu_count(monkeypatch): + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.setenv("FASTVIDEO_GPU_COUNT", "zero") + + with pytest.raises(RuntimeError, match="Invalid FASTVIDEO_GPU_COUNT"): + gpu_pool.get_available_gpus() + + +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 + queue timeout. + + The child signals startup with a "READY" message; we drain that + first so the wait_for budget bounds only the sentinel-detection + time, not spawn-import latency. + """ + ctx = mp.get_context("spawn") + cmd_q = ctx.Queue() + resp_q = ctx.Queue() + + proc = ctx.Process( + target=_child_consume_and_exit, args=(cmd_q, resp_q) + ) + proc.start() + + # Wait for the spawn child to fully boot. Allow generous time — + # this isn't what we're measuring. + ready = resp_q.get(timeout=30.0) + assert ready == "READY" + + async def runner(): + slot = gpu_pool.GPUSlot(gpu_id=0, cuda_device="0") + slot.process = proc + slot.command_queue = cmd_q + slot.response_queue = resp_q + + # Now the child is blocked in cmd_q.get(). Sending the cmd + # makes it exit ~immediately; sentinel must fire well before + # the queue timeout. + await asyncio.wait_for( + slot._send_command( + gpu_pool.Command(gpu_pool.CommandType.SHUTDOWN), + timeout=10.0, + ), + timeout=5.0, + ) + + try: + with pytest.raises(RuntimeError, match="worker died"): + asyncio.run(runner()) + finally: + proc.join(timeout=5) + cmd_q.close() + resp_q.close() diff --git a/apps/dreamverse/dreamverse/tests/test_mock_server.py b/apps/dreamverse/dreamverse/tests/test_mock_server.py new file mode 100644 index 000000000..30730e074 --- /dev/null +++ b/apps/dreamverse/dreamverse/tests/test_mock_server.py @@ -0,0 +1,355 @@ +# pyright: reportArgumentType=false +from __future__ import annotations + +import asyncio +import os + +from fastapi import WebSocketDisconnect + + +os.environ.setdefault("CEREBRAS_API_KEY", "dummy") +os.environ.setdefault("GROQ_API_KEY", "dummy") + +import dreamverse.mock_server as mock_server + + +class _FakeWebSocket: + def __init__(self, messages: list[tuple[float, dict[str, object]]]): + self._messages = messages + self._index = 0 + self.sent_json: list[dict[str, object]] = [] + self.sent_bytes: list[bytes] = [] + + async def accept(self) -> None: + return None + + async def send_json(self, payload: dict[str, object]) -> None: + self.sent_json.append(payload) + + async def send_bytes(self, payload: bytes) -> None: + self.sent_bytes.append(payload) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + del code, reason + + async def receive_json(self) -> dict[str, object]: + if self._index >= len(self._messages): + raise WebSocketDisconnect() + delay_s, payload = self._messages[self._index] + self._index += 1 + if delay_s > 0: + await asyncio.sleep(delay_s) + return payload + + +def test_mock_server_matches_current_single5s_protocol(): + old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES + old_latency_ms = mock_server.LATENCY_MS + try: + mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" + mock_server.LATENCY_MS = 1 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "simple_prompt_1", + "curated_prompts": ["selected prompt"], + "single_clip_mode": True, + "enhancement_enabled": False, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.01, + { + "type": "simple_generate", + "preset_id": "simple_custom_prompt", + "prompt_id": "simple_custom_prompt", + "prompt": "custom prompt", + "enhancement_enabled": True, + "initial_image": None, + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(mock_server.websocket_endpoint(ws)) + + message_types = [payload["type"] for payload in ws.sent_json] + assert message_types[0] == "queue_status" + assert "gpu_assigned" in message_types + assert message_types.count("ltx2_stream_start") == 2 + assert "prompt_received" in message_types + assert "prompt_enhancing" in message_types + assert "prompt_ready" in message_types + assert "media_init" in message_types + assert "media_segment_complete" in message_types + assert message_types.count("ltx2_stream_complete") == 2 + assert "prompt_sources_blocked" not in message_types + + segment_start_events = [ + payload + for payload in ws.sent_json + if payload["type"] == "ltx2_segment_start" + ] + gpu_assigned_event = next( + payload for payload in ws.sent_json if payload["type"] == "gpu_assigned" + ) + assert gpu_assigned_event["session_timeout"] == mock_server.SESSION_TIMEOUT_SECONDS + assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1] + assert segment_start_events[0]["prompt"] == "selected prompt" + assert segment_start_events[1]["prompt"] == "custom prompt" + + step_complete_events = [ + payload + for payload in ws.sent_json + if payload["type"] == "step_complete" + ] + assert len(step_complete_events) == 2 + assert step_complete_events[0]["latency_ms"] == { + "total": 121.0, + "worker_e2e": 1.0, + "main_user_step": 121.0, + "overhead": 120.0, + } + + assert ws.sent_bytes + assert ws.sent_bytes[0] == b"mock-fmp4-bytes" + finally: + mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes + mock_server.LATENCY_MS = old_latency_ms + + +def test_mock_server_regular_cap_waits_for_rewrite_rollout(): + old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES + old_latency_ms = mock_server.LATENCY_MS + old_generation_segment_cap = mock_server.GENERATION_SEGMENT_CAP + try: + mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" + mock_server.LATENCY_MS = 1 + mock_server.GENERATION_SEGMENT_CAP = 1 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.02, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "start a new rollout", + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(mock_server.websocket_endpoint(ws)) + + message_types = [payload["type"] for payload in ws.sent_json] + assert message_types.count("ltx2_stream_start") == 2 + assert message_types.count("ltx2_stream_complete") == 2 + assert "generation_cap_reached" not in message_types + assert "prompt_sources_blocked" not in message_types + + segment_start_events = [ + payload + for payload in ws.sent_json + if payload["type"] == "ltx2_segment_start" + ] + assert [payload["segment_idx"] for payload in segment_start_events] == [1, 1] + assert segment_start_events[0]["prompt"] == "segment one" + assert segment_start_events[1]["prompt"] == "segment one [start a new rollout]" + finally: + mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes + mock_server.LATENCY_MS = old_latency_ms + mock_server.GENERATION_SEGMENT_CAP = old_generation_segment_cap + + +def test_mock_server_rewrite_during_active_segment_restarts_from_first_rewritten_prompt(): + old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES + old_latency_ms = mock_server.LATENCY_MS + try: + mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" + mock_server.LATENCY_MS = 100 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one", "segment two"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + ( + 0.02, + { + "type": "rewrite_seed_prompts", + "rewrite_instruction": "restart from rewrite", + }, + ), + (0.40, {"type": "leave"}), + ] + ) + + asyncio.run(mock_server.websocket_endpoint(ws)) + + segment_start_events = [ + payload + for payload in ws.sent_json + if payload["type"] == "ltx2_segment_start" + ] + assert [payload["prompt"] for payload in segment_start_events[:2]] == [ + "segment one", + "segment one [restart from rewrite]", + ] + assert all( + payload["prompt"] != "segment two" + for payload in segment_start_events[1:] + ) + reset_events = [ + payload + for payload in ws.sent_json + if payload.get("type") == "seed_prompts_reset_applied" + ] + assert any( + payload.get("reason") == "rewrite_during_generation" + for payload in reset_events + ) + finally: + mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes + mock_server.LATENCY_MS = old_latency_ms + + +def test_mock_server_supports_initial_custom_rollout_prompt(): + old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES + old_latency_ms = mock_server.LATENCY_MS + try: + mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" + mock_server.LATENCY_MS = 1 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "custom_editable", + "preset_label": "Custom rollout", + "curated_prompts": [], + "initial_rollout_prompt": "A moonbase corridor thriller with flooding", + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.20, {"type": "leave"}), + ] + ) + + asyncio.run(mock_server.websocket_endpoint(ws)) + + message_types = [payload["type"] for payload in ws.sent_json] + assert "rewrite_seed_prompts_started" in message_types + assert "seed_prompts_updated" in message_types + assert "rewrite_seed_prompts_complete" in message_types + assert "seed_prompts_reset_applied" in message_types + assert "ltx2_stream_start" in message_types + assert "prompt_sources_blocked" not in message_types + + segment_start_events = [ + payload + for payload in ws.sent_json + if payload["type"] == "ltx2_segment_start" + ] + assert segment_start_events + assert segment_start_events[0]["prompt"] == ( + "A moonbase corridor thriller with flooding [segment 1]" + ) + finally: + mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes + mock_server.LATENCY_MS = old_latency_ms + + +def test_mock_server_can_start_new_project_without_reconnecting(): + old_segment_bytes = mock_server.MOCK_SEGMENT_BYTES + old_latency_ms = mock_server.LATENCY_MS + try: + mock_server.MOCK_SEGMENT_BYTES = b"mock-fmp4-bytes" + mock_server.LATENCY_MS = 40 + + ws = _FakeWebSocket( + [ + ( + 0.0, + { + "type": "session_init_v2", + "preset_id": "test_preset", + "curated_prompts": ["segment one"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.02, {"type": "end_project_keep_session"}), + ( + 0.20, + { + "type": "project_init_v1", + "preset_id": "test_preset_2", + "preset_label": "Test Preset 2", + "curated_prompts": ["segment two"], + "enhancement_enabled": True, + "auto_extension_enabled": False, + "loop_generation_enabled": False, + }, + ), + (0.40, {"type": "leave"}), + ] + ) + + asyncio.run(mock_server.websocket_endpoint(ws)) + + message_types = [payload["type"] for payload in ws.sent_json] + assert message_types.count("gpu_assigned") == 1 + assert message_types.count("ltx2_stream_start") == 2 + assert "project_idle" in message_types + + project_idle_index = message_types.index("project_idle") + stream_start_indexes = [ + index for index, message_type in enumerate(message_types) + if message_type == "ltx2_stream_start" + ] + assert stream_start_indexes[0] < project_idle_index < stream_start_indexes[1] + + segment_start_events = [ + payload + for payload in ws.sent_json + if payload["type"] == "ltx2_segment_start" + ] + assert [payload["prompt"] for payload in segment_start_events[:2]] == [ + "segment one", + "segment two", + ] + finally: + mock_server.MOCK_SEGMENT_BYTES = old_segment_bytes + mock_server.LATENCY_MS = old_latency_ms diff --git a/apps/dreamverse/dreamverse/video_generation.py b/apps/dreamverse/dreamverse/video_generation.py new file mode 100644 index 000000000..b9e530217 --- /dev/null +++ b/apps/dreamverse/dreamverse/video_generation.py @@ -0,0 +1,535 @@ +"""LTX2 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.*`` +imports are deferred to method bodies so nothing touches CUDA at +module import time. +""" +# pyright: reportArgumentType=false, reportMissingImports=false, reportMissingTypeArgument=false, reportOptionalMemberAccess=false +# ruff: noqa: SIM105 +# mypy: ignore-errors +import gc +import os +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from dreamverse.config import ( + FRAME_HEIGHT, + FRAME_WIDTH, + MODEL_CONFIG, + NUM_FRAMES, + NUM_INFERENCE_STEPS, +) + +# Multi-frame decoded continuation defaults from +# examples/inference/basic/basic_ltx2_distilled_video_continuation.py. +# Overridable via environment variables. +LTX2_VIDEO_CONDITIONING_NUM_FRAMES = int(os.getenv("LTX2_VIDEO_CONDITIONING_NUM_FRAMES", "9")) +LTX2_VIDEO_CONDITIONING_END_OFFSET = int(os.getenv("LTX2_VIDEO_CONDITIONING_END_OFFSET", "0")) +LTX2_VIDEO_CONDITIONING_FRAME_IDX = int(os.getenv("LTX2_VIDEO_CONDITIONING_FRAME_IDX", "0")) +LTX2_VIDEO_CONDITIONING_STRENGTH = float(os.getenv("LTX2_VIDEO_CONDITIONING_STRENGTH", "1.0")) + +if (LTX2_VIDEO_CONDITIONING_NUM_FRAMES - 1) % 8 != 0: + raise ValueError("LTX2_VIDEO_CONDITIONING_NUM_FRAMES must satisfy " + "(frames - 1) % 8 == 0; got " + f"{LTX2_VIDEO_CONDITIONING_NUM_FRAMES}") + +# Audio conditioning: reuse denoised audio latents from the previous +# segment as initial latents for the next segment. +ENABLE_AUDIO_COND = os.getenv("ENABLE_AUDIO_COND", "1").lower() in ("1", "true", "yes") +# Number of decoded video frames worth of audio to condition on. +# Not subject to the (n-1)%8==0 constraint since this is audio-only. +AUDIO_CONDITIONING_NUM_FRAMES = int(os.getenv("AUDIO_CONDITIONING_NUM_FRAMES", '49')) +AUDIO_CONDITIONING_STRENGTH = float(os.getenv("AUDIO_CONDITIONING_STRENGTH", "1.0")) + +# Noise injected into conditioning context to prevent error +# accumulation across segments. 0 = no noise (default). +VIDEO_CONTEXT_NOISE = float(os.getenv("VIDEO_CONTEXT_NOISE", "0")) +AUDIO_CONTEXT_NOISE = float(os.getenv("AUDIO_CONTEXT_NOISE", "0")) + +# Re-encode audio latents through decode→encode round-trip to +# regularize noise accumulation (mirrors the video PIL→VAE path). +ENABLE_AUDIO_RE_ENCODE = os.getenv("ENABLE_AUDIO_RE_ENCODE", "").lower() in ("1", "true", "yes") + +DEFAULT_LTX2_AUDIO_SAMPLE_RATE = 16000 +DEFAULT_LTX2_AUDIO_HOP_LENGTH = 160 +DEFAULT_LTX2_AUDIO_DOWNSAMPLE = 4 + + +@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.""" + + def __init__(self): + self.video_images: list | None = None + self.audio_latents: torch.Tensor | None = None + + def clear(self) -> None: + if self.video_images: + for old_image in self.video_images: + try: + old_image.close() + except Exception: + pass + self.video_images = None + self.audio_latents = None + + def apply_video(self, request_kwargs: dict, segment_idx: int) -> None: + """Seed next-segment kwargs with the cached tail frames.""" + if segment_idx <= 1 or not self.video_images: + return + from PIL import Image + + cond_images = list(self.video_images) + if VIDEO_CONTEXT_NOISE > 0: + noisy = [] + for img in cond_images: + arr = np.array(img, dtype=np.float32) + arr += np.random.normal(0, VIDEO_CONTEXT_NOISE * 255, arr.shape) + arr = np.clip(arr, 0, 255).astype(np.uint8) + noisy.append(Image.fromarray(arr)) + cond_images = noisy + request_kwargs["ltx2_video_conditions"] = [( + cond_images, + LTX2_VIDEO_CONDITIONING_FRAME_IDX, + LTX2_VIDEO_CONDITIONING_STRENGTH, + )] + request_kwargs["ltx2_images"] = None + request_kwargs["image_path"] = None + + def apply_audio( + self, + request_kwargs: dict, + segment_idx: int, + audio_lps: float, + ) -> None: + """Seed next-segment kwargs with clean audio latents + denoise mask. + + When audio conditioning is longer than video, extend audio + generation and shift video RoPE forward so the audio prefix + sits before video t=0. ``audio_lps`` (audio latent frames per + second) is passed in so this class never imports fastvideo. + """ + if not (ENABLE_AUDIO_COND and segment_idx > 1 and self.audio_latents is not None): + return + cached = self.audio_latents # [B,C,T,mel] + cond_duration = float(AUDIO_CONDITIONING_NUM_FRAMES) / 24.0 + audio_cond_T = max(1, round(cond_duration * audio_lps)) + audio_cond_T = min(audio_cond_T, cached.shape[2]) + + audio_extra = max(0, AUDIO_CONDITIONING_NUM_FRAMES - LTX2_VIDEO_CONDITIONING_NUM_FRAMES) + if audio_extra > 0: + audio_num_frames = NUM_FRAMES + audio_extra + request_kwargs["audio_num_frames"] = (audio_num_frames) + prefix_sec = float(audio_extra) / 24.0 + request_kwargs["video_position_offset_sec"] = prefix_sec + + new_duration = float(NUM_FRAMES + audio_extra) / 24.0 + total_T = max( + round(new_duration * audio_lps), + audio_cond_T + 1, + ) + + B, C, _, mel = cached.shape + clean = torch.zeros((B, C, total_T, mel), dtype=cached.dtype) + clean[:, :, :audio_cond_T, :] = (cached[:, :, -audio_cond_T:, :]) + if AUDIO_CONTEXT_NOISE > 0: + noise = torch.randn_like(clean[:, :, :audio_cond_T, :]) + clean[:, :, :audio_cond_T, :] += (AUDIO_CONTEXT_NOISE * noise) + + mask = torch.ones((B, 1, total_T, 1), dtype=torch.float32) + mask[:, :, :audio_cond_T, :] = (1.0 - AUDIO_CONDITIONING_STRENGTH) + + request_kwargs["ltx2_audio_clean_latent"] = clean + request_kwargs["ltx2_audio_denoise_mask"] = mask + + def save_video(self, frames: list) -> None: + """Snapshot trailing N frames as PIL images for next-segment conditioning.""" + from PIL import Image + num_cond_frames = LTX2_VIDEO_CONDITIONING_NUM_FRAMES + end_offset = LTX2_VIDEO_CONDITIONING_END_OFFSET + if end_offset + num_cond_frames > len(frames): + raise RuntimeError(f"Cannot extract {num_cond_frames} conditioning frames with " + f"end_offset={end_offset} from {len(frames)} generated frames.") + start_idx = len(frames) - end_offset - num_cond_frames + self.video_images = [ + Image.fromarray(np.ascontiguousarray(frames[start_idx + i])) for i in range(num_cond_frames) + ] + + def save_audio_latents(self, latents: torch.Tensor | None) -> None: + if latents is None: + self.audio_latents = None + return + self.audio_latents = latents.detach().clone().cpu() + + +class VideoGenerationWorker: + """Single-GPU LTX2 generator with continuation state. + + Caller must set ``os.environ["CUDA_VISIBLE_DEVICES"]`` before + instantiating, and call ``initialize()`` before any + ``generate_step()`` / ``warmup()``. + """ + + def __init__(self, gpu_id: int): + self.gpu_id = gpu_id + self.generator = None + self.current_model_config: dict = dict(MODEL_CONFIG) + self.continuation = ContinuationState() + self.audio_encoder_module = None + self.audio_processor_module = None + + def _gpu_mem(self) -> str: + a = torch.cuda.memory_allocated() / 1024**3 + r = torch.cuda.memory_reserved() / 1024**3 + return f"alloc={a:.2f}GiB, reserved={r:.2f}GiB" + + @staticmethod + def _resolve_refine_upsampler_path(model_root: str) -> str: + candidates = ( + os.path.join(model_root, "spatial_upscaler"), + os.path.join(model_root, "spatial_upsampler"), + ) + for candidate in candidates: + config_path = os.path.join(candidate, "config.json") + if os.path.isfile(config_path): + return candidate + raise FileNotFoundError("Could not find an LTX2 refine upsampler directory under " + f"{model_root}. Checked: {', '.join(candidates)}") + + def initialize(self, model_config: dict | None = None) -> None: + """Load (or reload) the LTX2 generator on the visible GPU.""" + if model_config is not None: + self.current_model_config = model_config + + if self.generator is not None: + print(f"[GPU {self.gpu_id}] Freeing old model...") + try: + self.generator.shutdown() + except Exception: + pass + del self.generator + self.generator = None + gc.collect() + torch.cuda.empty_cache() + print(f"[GPU {self.gpu_id}] After cleanup: {self._gpu_mem()}") + + print(f"[GPU {self.gpu_id}] Loading model: " + f"{self.current_model_config['model_path']}") + print(f"[GPU {self.gpu_id}] Before model load: {self._gpu_mem()}") + + from fastvideo.api.schema import ( + CompileConfig, + ComponentConfig, + EngineConfig, + GeneratorConfig, + OffloadConfig, + PipelineSelection, + QuantizationConfig, + ) + from fastvideo.entrypoints.video_generator import VideoGenerator + from fastvideo.utils import maybe_download_model + + model_root = maybe_download_model(self.current_model_config["model_path"]) + refine_upsampler_path = self._resolve_refine_upsampler_path(model_root) + config_model_path = (self.current_model_config.get("config_model_path") + or self.current_model_config["model_path"]) + + enable_compile = os.getenv("ENABLE_TORCH_COMPILE", "1") == "1" + + components = ComponentConfig( + config_root=config_model_path, + upsampler_weights=refine_upsampler_path, + ) + init_weights = self.current_model_config.get("init_weights_from_safetensors") + if init_weights: + components.transformer_weights = init_weights + + generator_config = GeneratorConfig( + model_path=model_root, + engine=EngineConfig( + num_gpus=1, + offload=OffloadConfig( + dit=False, + dit_layerwise=False, + text_encoder=False, + vae=False, + pin_cpu_memory=True, + ), + compile=CompileConfig( + enabled=enable_compile, + text_encoder_enabled=enable_compile, + backend="inductor", + fullgraph=True, + mode="max-autotune-no-cudagraphs", + dynamic=False, + ), + use_fsdp_inference=False, + quantization=QuantizationConfig(transformer_quant="NVFP4"), + ), + pipeline=PipelineSelection( + components=components, + vae_tiling=False, + preset_overrides={ + "refine": { + "enabled": True, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "add_noise": True, + }, + }, + ), + ) + + self.generator = VideoGenerator.from_pretrained(config=generator_config) + print(f"[GPU {self.gpu_id}] After model load: {self._gpu_mem()}") + self._load_audio_encoder(model_root) + print(f"[GPU {self.gpu_id}] LTX2 model loaded (warmup pending)") + + def _load_audio_encoder(self, model_root: str) -> None: + if not ENABLE_AUDIO_RE_ENCODE: + return + from fastvideo.models.audio.ltx2_audio_processing import AudioProcessor + from fastvideo.models.loader.component_loader import ComponentLoader + + audio_vae_path = os.path.join(model_root, "audio_vae") + if not os.path.isdir(audio_vae_path): + print(f"[GPU {self.gpu_id}] audio_vae dir not found at " + f"{audio_vae_path}; disabling re-encode") + return + + loader = ComponentLoader.for_module_type("audio_encoder", "diffusers") + enc = loader.load(audio_vae_path, self.generator.fastvideo_args) + target = getattr(enc, "model", enc) + + proc = AudioProcessor( + sample_rate=target.sample_rate, + mel_bins=target.mel_bins, + mel_hop_length=target.mel_hop_length, + n_fft=target.n_fft, + ).to(torch.device("cuda")) + + self.audio_encoder_module = target + self.audio_processor_module = proc + print(f"[GPU {self.gpu_id}] Audio encoder loaded for " + f"re-encode conditioning ({self._gpu_mem()})") + + def _re_encode_audio( + self, + waveform: torch.Tensor, + sample_rate: int, + ) -> torch.Tensor | None: + """Waveform → mel → encoder → latents.""" + if (self.audio_encoder_module is None or self.audio_processor_module is None): + return None + device = torch.device("cuda") + if waveform.ndim == 1: + waveform = waveform.unsqueeze(0) # [channels, samples] + waveform = waveform.unsqueeze(0).to(device=device, dtype=torch.float32) # [1, ch, samples] + with torch.no_grad(): + mel = self.audio_processor_module.waveform_to_mel( + waveform, + waveform_sample_rate=sample_rate, + ).to(device=device, dtype=torch.float32) + latents = self.audio_encoder_module(mel) + return latents.detach() + + def shutdown(self) -> None: + if self.generator is not None: + try: + self.generator.shutdown() + except Exception: + pass + + def clear_conditioning(self) -> None: + self.continuation.clear() + + def generate_step( + self, + prompt: str, + segment_idx: int, + image_path: str | None, + reset_conditioning: bool, + ) -> StepResult: + """Execute one generation step; snapshot state for the next segment.""" + timings: dict = {} + + request_kwargs = dict( + prompt=prompt, + negative_prompt="", + save_video=False, + height=FRAME_HEIGHT, + width=FRAME_WIDTH, + num_frames=NUM_FRAMES, + fps=24, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=1.0, + seed=10, + ltx2_image_crf=0.0, + image_path=image_path if segment_idx == 1 else None, + return_continuation_state=False, + ) + + if reset_conditioning: + self.continuation.clear() + + audio_lps = (DEFAULT_LTX2_AUDIO_SAMPLE_RATE / DEFAULT_LTX2_AUDIO_HOP_LENGTH / DEFAULT_LTX2_AUDIO_DOWNSAMPLE) + + # Phase 1: seed kwargs with prior-segment conditioning. + self.continuation.apply_video(request_kwargs, segment_idx) + self.continuation.apply_audio(request_kwargs, segment_idx, audio_lps) + + # Phase 2: generate. + t0 = time.perf_counter() + result = self.generator.generate_video(**request_kwargs) + torch.cuda.synchronize() + timings["generation_ms"] = (time.perf_counter() - t0) * 1000 + + if not isinstance(result, dict): + raise RuntimeError("Expected dictionary output from generate_video.") + frames = result.get("frames") + if not isinstance(frames, list) or len(frames) == 0: + raise RuntimeError("Generation did not return frames.") + audio = result.get("audio") + audio_sample_rate = result.get("audio_sample_rate") + if audio is not None and audio_sample_rate is None: + # LTX2 audio decoding stage uses 24kHz output by default. + audio_sample_rate = 24000 + print(f"[GPU {self.gpu_id}] audio_sample_rate missing from result; " + f"defaulting to {audio_sample_rate}Hz") + + timings["generation_time_ms"] = result.get("generation_time", 0.0) * 1000 + + # Phase 3: snapshot continuation state for the next segment. + t_save_start = time.perf_counter() + self.continuation.clear() + self.continuation.save_video(frames) + next_audio_latents = self._derive_next_audio_latents(audio, audio_sample_rate, result, segment_idx) + self.continuation.save_audio_latents(next_audio_latents) + + timings["save_conditioning_ms"] = (time.perf_counter() - t_save_start) * 1000 + timings["e2e_latency_ms"] = (time.perf_counter() - t0) * 1000 + + print(f"[GPU {self.gpu_id}] LTX2 segment {segment_idx}: " + f"{len(frames)} frames, gen={timings['generation_ms']:.0f}ms, " + f"save_conditioning={timings['save_conditioning_ms']:.0f}ms, " + f"e2e={timings['e2e_latency_ms']:.0f}ms") + + # Head-trim values for downstream AV streaming — computed here so + # the streaming layer never needs to know conditioning constants. + is_continuation = segment_idx > 1 and not reset_conditioning + head_trim_frames = (LTX2_VIDEO_CONDITIONING_NUM_FRAMES if is_continuation else 0) + audio_extra = (max(0, AUDIO_CONDITIONING_NUM_FRAMES - + LTX2_VIDEO_CONDITIONING_NUM_FRAMES) if ENABLE_AUDIO_COND else 0) + head_trim_audio_frames = (head_trim_frames + audio_extra if is_continuation else 0) + + return StepResult( + frames=frames, + audio=audio, + audio_sample_rate=audio_sample_rate, + timings=timings, + head_trim_frames=head_trim_frames, + head_trim_audio_frames=head_trim_audio_frames, + ) + + def _derive_next_audio_latents( + self, + audio: object, + audio_sample_rate: int | None, + result: dict, + segment_idx: int, + ) -> torch.Tensor | None: + """Pick which tensor to cache for next-segment audio conditioning.""" + if not ENABLE_AUDIO_COND: + return None + if (ENABLE_AUDIO_RE_ENCODE and audio is not None and audio_sample_rate is not None): + re_encoded = self._re_encode_audio(audio, audio_sample_rate) + if re_encoded is not None: + print(f"[GPU {self.gpu_id}] Re-encoded audio " + f"latents shape=" + f"{tuple(re_encoded.shape)} " + f"for segment {segment_idx + 1}") + return re_encoded + return None + audio_latents = result.get("ltx2_audio_latents") + if audio_latents is not None: + print(f"[GPU {self.gpu_id}] Cached audio latents " + f"shape={tuple(audio_latents.shape)} " + f"for segment {segment_idx + 1}") + return audio_latents + return None + + def warmup(self, prompt: str) -> dict[str, float]: + warmup_prompt = (prompt or "").strip() + if not warmup_prompt: + raise RuntimeError("Startup warmup prompt must be non-empty.") + + print(f"[GPU {self.gpu_id}] Startup warmup starting " + "(synthetic segments: seg1, seg2, seg1-post-LoRA)") + warmup_t0 = time.perf_counter() + + r1 = self.generate_step( + warmup_prompt, + segment_idx=1, + image_path=None, + reset_conditioning=True, + ) + r2 = self.generate_step( + warmup_prompt, + segment_idx=2, + image_path=None, + reset_conditioning=False, + ) + # r1 stage 1 compiled BEFORE LoRA wrapping (which happens during + # r1 stage 2 via ltx2_refine_lora_stage), so the resulting graph + # is keyed off pre-LoRA module identity and is stale once r1 r2 + # finish. The first real user seg=1 then re-compiles stage 1 + # ("stage1-LoRA-nocont"), wasting ~90s on the user's first + # request. Run a 3rd warmup pass with seg=1 reset=True after r2 + # so this graph is compiled while no client is waiting. r3 + # stage 2 hits r1 stage 2's cache (shape match, both LoRA-nocont) + # so the only real work is the missing stage 1 graph. + self.continuation.clear() + r3 = self.generate_step( + warmup_prompt, + segment_idx=1, + image_path=None, + reset_conditioning=True, + ) + warmup_total_ms = (time.perf_counter() - warmup_t0) * 1000.0 + self.continuation.clear() + + segment1_ms = float(r1.timings.get("e2e_latency_ms", 0.0)) + segment2_ms = float(r2.timings.get("e2e_latency_ms", 0.0)) + segment3_ms = float(r3.timings.get("e2e_latency_ms", 0.0)) + print(f"[GPU {self.gpu_id}] Startup warmup complete: " + f"segment1={segment1_ms:.0f}ms, " + f"segment2={segment2_ms:.0f}ms, " + f"segment3={segment3_ms:.0f}ms, total={warmup_total_ms:.0f}ms") + return { + "warmup_segment1_ms": segment1_ms, + "warmup_segment2_ms": segment2_ms, + "warmup_segment3_ms": segment3_ms, + "warmup_total_ms": warmup_total_ms, + } diff --git a/apps/dreamverse/dreamverse/worker_ipc.py b/apps/dreamverse/dreamverse/worker_ipc.py new file mode 100644 index 000000000..e11478d0f --- /dev/null +++ b/apps/dreamverse/dreamverse/worker_ipc.py @@ -0,0 +1,157 @@ +# pyright: reportMissingTypeArgument=false +"""Typed worker IPC events: replaces the fat ``Response`` dataclass. + +Each message kind from GPU worker → main process is its own small +dataclass; ``WorkerEvent`` is the union of them. Consumers dispatch +via ``match``/``case`` or ``isinstance`` — mirroring the existing +``StreamEvent`` pattern in ``av_streaming.py``. + +Invalid states are unrepresentable: a ``MediaChunk`` simply has no +``frames`` field, a ``StepComplete`` has no ``chunk_offset`` field, +and a ``JoinAck`` can't accidentally default to ``kind="step_result"`` +because there is no ``kind`` string. +""" +from __future__ import annotations + +from dataclasses import dataclass + +# ---- User-scoped events (carry user_id) ------------------------------------ + + +@dataclass(frozen=True) +class StepComplete: + """Generation step finished. Frames/audio were already sent via ``MediaChunk``.""" + user_id: str + segment_idx: int + timings: dict[str, float] + + +@dataclass(frozen=True) +class WorkerError: + """Any worker failure. ``user_id`` is None for system-scoped failures.""" + user_id: str | None + message: str + + +@dataclass(frozen=True) +class JoinAck: + user_id: str + + +@dataclass(frozen=True) +class LeaveAck: + user_id: str + + +@dataclass(frozen=True) +class ReloadAck: + user_id: str + + +@dataclass(frozen=True) +class WarmupComplete: + user_id: str | None + timings: dict[str, float] + + +# ---- Streaming events (carry user_id + segment_idx for routing) ------------ + + +@dataclass(frozen=True) +class MediaInit: + user_id: str + segment_idx: int + stream_id: str + mime: str + uses_shared_buffer: bool + + +@dataclass(frozen=True) +class MediaChunk: + """One fMP4 chunk. + + Either ``chunk`` (raw bytes) is set, or both ``chunk_offset`` and + ``chunk_length`` are set (read from the shared buffer). The + invariant is enforced in ``__post_init__``. + """ + user_id: str + segment_idx: int + stream_id: str + chunk: bytes | None = None + chunk_offset: int | None = None + chunk_length: int | None = None + uses_shared_buffer: bool = False + + def __post_init__(self) -> None: + has_bytes = self.chunk is not None + has_offset = (self.chunk_offset is not None and self.chunk_length is not None) + if has_bytes == has_offset: + raise ValueError("MediaChunk must carry either chunk bytes or " + "(chunk_offset + chunk_length), not both or neither") + + +@dataclass(frozen=True) +class MediaComplete: + user_id: str + segment_idx: int + stream_id: str + chunks: int + + +# ---- System events (no user_id) -------------------------------------------- + + +@dataclass(frozen=True) +class InitAck: + success: bool + error: str | None = None + + +@dataclass(frozen=True) +class ShutdownAck: + pass + + +WorkerEvent = (StepComplete + | WorkerError + | JoinAck + | LeaveAck + | ReloadAck + | WarmupComplete + | MediaInit + | MediaChunk + | MediaComplete + | InitAck + | ShutdownAck) + +# ---- Command payloads (main process → worker) ------------------------------ +# +# The envelope (``Command`` + ``CommandType``) lives in ``gpu_pool.py`` +# alongside the enum; only the typed payloads live here so they sit +# next to the response-side types. Commands without a payload +# (``INIT``, ``SHUTDOWN``, ``USER_JOIN``, ``USER_LEAVE``) leave +# ``Command.payload`` as ``None``. + + +@dataclass(frozen=True) +class UserStepPayload: + prompt: str + segment_idx: int + image_path: str | None + reset_conditioning: bool + + +@dataclass(frozen=True) +class WarmupPayload: + prompt: str + + +@dataclass(frozen=True) +class ReloadModelPayload: + # ``model_config`` stays a dict because ``MODEL_REGISTRY`` values + # are dicts shaped by the external config module; typing that dict + # is a separate refactor. + model_config: dict + + +CommandPayload = UserStepPayload | WarmupPayload | ReloadModelPayload