diff --git a/apps/dreamverse/server/av_streaming.py b/apps/dreamverse/server/av_streaming.py new file mode 100644 index 000000000..505dc722b --- /dev/null +++ b/apps/dreamverse/server/av_streaming.py @@ -0,0 +1,403 @@ +# 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 + 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()) + 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") + wav_path = _write_audio_wav(audio_int16, num_channels, sample_rate) + + 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) + + try: + 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 + fcntl.fcntl(proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, 1048576) + fcntl.fcntl(proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, 1048576) + + 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 + while True: + chunk = proc.stdout.read(AV_CHUNK_SIZE_BYTES) + if not chunk: + break + chunk_count += 1 + total_bytes += len(chunk) + 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 + 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, + )) + + 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)) + 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/server/gpu_pool.py b/apps/dreamverse/server/gpu_pool.py new file mode 100644 index 000000000..eb9424eab --- /dev/null +++ b/apps/dreamverse/server/gpu_pool.py @@ -0,0 +1,952 @@ +# 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 config import ( + DEFAULT_MODEL_ID, + MODEL_REGISTRY, + STARTUP_WARMUP_ENABLED, + STARTUP_WARMUP_PROMPT, + STARTUP_WARMUP_TIMEOUT_SECONDS, +) +from av_streaming import ( + SHARED_STREAM_BUFFER_BYTES, + USE_SHARED_STREAM_BUFFER, + StreamChunk, + StreamComplete, + StreamEvent, + StreamInit, + generate_stream_id, + stream_fmp4, +) +from 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 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") + 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/server/main.py b/apps/dreamverse/server/main.py index 878cc2663..48fe4cdb4 100644 --- a/apps/dreamverse/server/main.py +++ b/apps/dreamverse/server/main.py @@ -9,9 +9,8 @@ from pathlib import Path from fastapi import FastAPI, WebSocket from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles -from fastvideo.entrypoints.streaming.gpu_pool import GpuPool as GPUPool -from fastvideo.entrypoints.streaming.gpu_pool import get_available_gpus -from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger +from gpu_pool import GPUPool, get_available_gpus +from session_logger import SessionEventLogger from config import ( DEVTOOLS_ENABLED, diff --git a/apps/dreamverse/server/mock_server.py b/apps/dreamverse/server/mock_server.py new file mode 100644 index 000000000..42bf4ccd1 --- /dev/null +++ b/apps/dreamverse/server/mock_server.py @@ -0,0 +1,1161 @@ +# 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 config import FRONTEND_STATIC_DIR_CANDIDATES, GENERATION_SEGMENT_CAP +from 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 + + 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/server/runtime.py b/apps/dreamverse/server/runtime.py index e2c6bfbfa..74fcb0f27 100644 --- a/apps/dreamverse/server/runtime.py +++ b/apps/dreamverse/server/runtime.py @@ -11,8 +11,8 @@ from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: - from fastvideo.entrypoints.streaming.gpu_pool import GpuPool as GPUPool - from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger + from gpu_pool import GPUPool + from session_logger import SessionEventLogger from prompt_enhancer import PromptEnhancer from prompt_safety import PromptSafetyFilter diff --git a/apps/dreamverse/server/server_entry.py b/apps/dreamverse/server/server_entry.py new file mode 100644 index 000000000..4369f5e91 --- /dev/null +++ b/apps/dreamverse/server/server_entry.py @@ -0,0 +1,14 @@ +# pyright: reportMissingTypeArgument=false +from __future__ import annotations + + +def cli() -> None: + try: + from main import cli as main_cli + except ModuleNotFoundError as exc: + if exc.name in {"fastvideo", "torch", "safetensors"}: + raise SystemExit("dreamverse-server requires the Python backend extra. " + "Run `uv sync --extra server` first.") from exc + raise + + main_cli() diff --git a/apps/dreamverse/server/session/controller.py b/apps/dreamverse/server/session/controller.py index 6cdddac1c..cbe2eb7d5 100644 --- a/apps/dreamverse/server/session/controller.py +++ b/apps/dreamverse/server/session/controller.py @@ -25,9 +25,9 @@ import uuid from typing import TYPE_CHECKING from fastapi import WebSocket, WebSocketDisconnect -from fastvideo.entrypoints.streaming.gpu_pool import GPUSlot -from fastvideo.entrypoints.streaming.protocol import MediaChunk, MediaComplete, MediaInit -from fastvideo.entrypoints.streaming.session_init_image import cleanup_session_init_image, persist_session_init_image +from gpu_pool import GPUSlot +from session_init_image import cleanup_session_init_image, persist_session_init_image +from worker_ipc import MediaChunk, MediaComplete, MediaInit from config import ( DEFAULT_MODEL_ID, @@ -47,8 +47,8 @@ from utils import ( ) if TYPE_CHECKING: - from fastvideo.entrypoints.streaming.gpu_pool import GpuPool as GPUPool - from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger + from gpu_pool import GPUPool + from session_logger import SessionEventLogger from prompt_enhancer import PromptEnhancer from prompt_safety import PromptSafetyFilter diff --git a/apps/dreamverse/server/session_init_image.py b/apps/dreamverse/server/session_init_image.py new file mode 100644 index 000000000..5ad4ae8f1 --- /dev/null +++ b/apps/dreamverse/server/session_init_image.py @@ -0,0 +1,101 @@ +# pyright: reportMissingTypeArgument=false, reportArgumentType=false +from __future__ import annotations + +import base64 +import io +import re +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image + +MAX_SESSION_INIT_IMAGE_BYTES = 15 * 1024 * 1024 +SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES = { + "image/jpeg", + "image/png", + "image/webp", +} + +_DATA_URL_RE = re.compile(r"^data:(?P[-\w.+/]+);base64,(?P[A-Za-z0-9+/=\s]+)$") + + +@dataclass(frozen=True) +class SessionInitImage: + file_path: Path + temp_dir: Path + display_name: str + mime_type: str + + +def _sanitize_display_name(raw_name: object) -> str: + text = str(raw_name or "").strip() + if not text: + return "uploaded-image" + return Path(text).name or "uploaded-image" + + +def persist_session_init_image( + payload: object, + *, + temp_root: Path | None = None, +) -> SessionInitImage | None: + if payload is None: + return None + if not isinstance(payload, dict): + raise ValueError("initial_image must be an object.") + + data_url = str(payload.get("data_url") or "").strip() + if not data_url: + return None + + data_url_match = _DATA_URL_RE.match(data_url) + if data_url_match is None: + raise ValueError("initial_image.data_url must be a base64 data URL.") + + data_url_mime_type = data_url_match.group("mime").strip().lower() + mime_type = str(payload.get("mime_type") or data_url_mime_type).strip().lower() + if mime_type != data_url_mime_type: + raise ValueError("initial_image.mime_type must match the data URL mime type.") + if mime_type not in SUPPORTED_SESSION_INIT_IMAGE_MIME_TYPES: + raise ValueError("initial_image must be a PNG, JPEG, or WebP image.") + + encoded_bytes = re.sub(r"\s+", "", data_url_match.group("data")) + try: + raw_bytes = base64.b64decode(encoded_bytes, validate=True) + except Exception as exc: + raise ValueError("initial_image.data_url is not valid base64 data.") from exc + + if len(raw_bytes) == 0: + raise ValueError("initial_image.data_url did not contain image bytes.") + if len(raw_bytes) > MAX_SESSION_INIT_IMAGE_BYTES: + raise ValueError("initial_image must be 15 MB or smaller.") + + try: + with Image.open(io.BytesIO(raw_bytes)) as image: + image.load() + normalized = image.convert("RGBA" if "A" in image.getbands() else "RGB") + except Exception as exc: + raise ValueError("initial_image must decode as a valid image.") from exc + + temp_dir = Path( + tempfile.mkdtemp( + prefix="ltx2_session_init_", + dir=str(temp_root) if temp_root is not None else None, + )) + file_path = temp_dir / "initial_frame.png" + normalized.save(file_path, format="PNG") + + return SessionInitImage( + file_path=file_path, + temp_dir=temp_dir, + display_name=_sanitize_display_name(payload.get("name")), + mime_type=mime_type, + ) + + +def cleanup_session_init_image(session_image: SessionInitImage | None) -> None: + if session_image is None: + return + shutil.rmtree(session_image.temp_dir, ignore_errors=True) diff --git a/apps/dreamverse/server/session_logger.py b/apps/dreamverse/server/session_logger.py new file mode 100644 index 000000000..f537fdacc --- /dev/null +++ b/apps/dreamverse/server/session_logger.py @@ -0,0 +1,46 @@ +# pyright: reportMissingTypeArgument=false +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone +import json +import socket +from pathlib import Path +from typing import Any + + +def _utc_now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +class SessionEventLogger: + + def __init__(self, root_dir: Path): + self.hostname = socket.gethostname() + timestamp = datetime.now(timezone.utc).strftime("%y%m%d_%H%M%S") + self.directory = root_dir / self.hostname + self.path = self.directory / f"{timestamp}.jsonl" + self._lock = asyncio.Lock() + + self.directory.mkdir(parents=True, exist_ok=True) + self.path.touch(exist_ok=False) + + async def write_event( + self, + *, + event: str, + client_id: str, + payload: dict[str, Any] | None = None, + ) -> None: + entry = { + "ts": _utc_now_iso(), + "event": event, + "hostname": self.hostname, + "client_id": client_id, + } + if payload: + entry.update(payload) + + async with self._lock: + with self.path.open("a", encoding="utf-8") as fp: + fp.write(json.dumps(entry, ensure_ascii=False) + "\n") diff --git a/apps/dreamverse/server/tests/test_entrypoints.py b/apps/dreamverse/server/tests/test_entrypoints.py index ab951e863..8e6d67490 100644 --- a/apps/dreamverse/server/tests/test_entrypoints.py +++ b/apps/dreamverse/server/tests/test_entrypoints.py @@ -5,17 +5,16 @@ import sys from pathlib import Path import types -import fastvideo.entrypoints.streaming.mock_server as mock_server - SERVER_DIR = Path(__file__).resolve().parent if str(SERVER_DIR) not in sys.path: sys.path.insert(0, str(SERVER_DIR)) import main as server_main # noqa: E402 +import mock_server # noqa: E402 -def _run_cli(module, monkeypatch, argv: list[str], entrypoint: str = "cli") -> list[dict[str, object]]: +def _run_cli(module, monkeypatch, argv: list[str]) -> list[dict[str, object]]: calls: list[dict[str, object]] = [] uvicorn_stub = types.ModuleType("uvicorn") @@ -32,7 +31,7 @@ def _run_cli(module, monkeypatch, argv: list[str], entrypoint: str = "cli") -> l monkeypatch.setitem(sys.modules, "uvicorn", uvicorn_stub) monkeypatch.setattr(sys, "argv", argv) - getattr(module, entrypoint)() + module.cli() return calls @@ -65,52 +64,37 @@ def test_server_cli_allows_explicit_host_and_port(monkeypatch): def test_mock_server_cli_defaults_to_local_web_port(monkeypatch): - mock_app = object() - build_calls: list[dict[str, object]] = [] - - def build_mock_app(*, sleep_ms: float = 0.0): - build_calls.append({"sleep_ms": sleep_ms}) - return mock_app - - monkeypatch.setattr(mock_server, "build_mock_app", build_mock_app) calls = _run_cli( mock_server, monkeypatch, ["dreamverse-mock-server"], - entrypoint="main", ) - assert build_calls == [{"sleep_ms": 0.0}] assert calls == [ { - "app": mock_app, - "host": "127.0.0.1", - "port": 8000, + "app": mock_server.app, + "host": "0.0.0.0", + "port": 8009, } ] def test_mock_server_cli_updates_latency(monkeypatch): - mock_app = object() - build_calls: list[dict[str, object]] = [] + old_latency_ms = mock_server.LATENCY_MS + try: + calls = _run_cli( + mock_server, + monkeypatch, + ["dreamverse-mock-server", "--latency", "321", "--port", "8111"], + ) - def build_mock_app(*, sleep_ms: float = 0.0): - build_calls.append({"sleep_ms": sleep_ms}) - return mock_app - - monkeypatch.setattr(mock_server, "build_mock_app", build_mock_app) - calls = _run_cli( - mock_server, - monkeypatch, - ["dreamverse-mock-server", "--sleep-ms", "321", "--port", "8111"], - entrypoint="main", - ) - - assert build_calls == [{"sleep_ms": 321.0}] - assert calls == [ - { - "app": mock_app, - "host": "127.0.0.1", - "port": 8111, - } - ] + assert calls == [ + { + "app": mock_server.app, + "host": "0.0.0.0", + "port": 8111, + } + ] + assert mock_server.LATENCY_MS == 321 + finally: + mock_server.LATENCY_MS = old_latency_ms diff --git a/apps/dreamverse/server/tests/test_gpu_pool.py b/apps/dreamverse/server/tests/test_gpu_pool.py new file mode 100644 index 000000000..6a38f9050 --- /dev/null +++ b/apps/dreamverse/server/tests/test_gpu_pool.py @@ -0,0 +1,127 @@ +# pyright: reportAttributeAccessIssue=false +from __future__ import annotations + +import asyncio +import multiprocessing as mp +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +SERVER_DIR = Path(__file__).resolve().parent +if str(SERVER_DIR) not in sys.path: + sys.path.insert(0, str(SERVER_DIR)) + +import 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/server/tests/test_mock_server.py b/apps/dreamverse/server/tests/test_mock_server.py new file mode 100644 index 000000000..ea8ba00ba --- /dev/null +++ b/apps/dreamverse/server/tests/test_mock_server.py @@ -0,0 +1,361 @@ +# pyright: reportArgumentType=false +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +from fastapi import WebSocketDisconnect + + +SERVER_DIR = Path(__file__).resolve().parent +if str(SERVER_DIR) not in sys.path: + sys.path.insert(0, str(SERVER_DIR)) + +os.environ.setdefault("CEREBRAS_API_KEY", "dummy") +os.environ.setdefault("GROQ_API_KEY", "dummy") + +import 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/server/tests/test_prompt_enhancer.py b/apps/dreamverse/server/tests/test_prompt_enhancer.py new file mode 100644 index 000000000..306afdcb6 --- /dev/null +++ b/apps/dreamverse/server/tests/test_prompt_enhancer.py @@ -0,0 +1,1177 @@ +# pyright: reportMissingTypeArgument=false, reportArgumentType=false, reportOptionalSubscript=false, reportOperatorIssue=false +from __future__ import annotations + +import asyncio +import os +import re +import sys +import time +from pathlib import Path + + +SERVER_DIR = Path(__file__).resolve().parent +if str(SERVER_DIR) not in sys.path: + sys.path.insert(0, str(SERVER_DIR)) + +os.environ.setdefault("CEREBRAS_API_KEY", "dummy") +os.environ.setdefault("GROQ_API_KEY", "dummy") + +import prompt_enhancer as prompt_enhancer_module +from prompt_enhancer import ( + ProviderRuntime, + PromptEnhancer, + _build_prompt_backup_path, + _load_prompt_required, + _parse_json_response, + _resolve_prompt_save_path, +) + + +class _FakeResponse: + def __init__(self, payload: dict): + self._payload = payload + + def model_dump(self, mode: str = "json"): # noqa: ARG002 + return self._payload + + +class _FakeSyncCompletions: + def __init__(self, payload: dict): + self._payload = payload + + def create(self, **kwargs): # noqa: ARG002 + return _FakeResponse(self._payload) + + +class _FakeSyncClient: + def __init__(self, payload: dict): + self.chat = type( + "_FakeChat", + (), + {"completions": _FakeSyncCompletions(payload)}, + )() + + +class _DelayedSyncCompletions: + def __init__(self, payload: dict, delay_s: float = 0.0, exc: Exception | None = None): + self._payload = payload + self._delay_s = delay_s + self._exc = exc + + def create(self, **kwargs): # noqa: ARG002 + if self._delay_s > 0: + time.sleep(self._delay_s) + if self._exc is not None: + raise self._exc + return _FakeResponse(self._payload) + + +class _DelayedSyncClient: + def __init__(self, payload: dict, delay_s: float = 0.0, exc: Exception | None = None): + self.chat = type( + "_FakeChat", + (), + { + "completions": _DelayedSyncCompletions( + payload, + delay_s=delay_s, + exc=exc, + ) + }, + )() + + +def _chat_payload_with_content(content: str) -> dict: + return { + "choices": [ + { + "message": { + "content": content, + } + } + ] + } + + +def _build_test_enhancer(payload: dict) -> PromptEnhancer: + enhancer = PromptEnhancer.__new__(PromptEnhancer) + enhancer.provider = "cerebras" + enhancer.provider_label = "cerebras" + enhancer.model = "gpt-test" + enhancer.rewrite_default_model = "gpt-test" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.default_timeout_ms = 1200 + enhancer.http_timeout_ms = 1200 + enhancer.temperature = 0.4 + enhancer.rewrite_default_temperature = 0.4 + enhancer.max_completion_tokens = 512 + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_user_system_prompt_path = "/tmp/rewrite_user_system_prompt.md" + enhancer.rewrite_user_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt = "system prompt" + enhancer.rewrite_user_system_prompt = "system prompt" + enhancer.provider_request_models = { + "cerebras": "gpt-test", + "groq": "openai/gpt-test", + } + enhancer.provider_success_counts = { + "cerebras": 0, + "groq": 0, + } + enhancer.client = _FakeSyncClient(payload) + return enhancer + + +def _build_staged_enhancer( + *, + cerebras_payload: dict, + groq_payload: dict, + cerebras_request_model: str = "gpt-test", + groq_request_model: str = "openai/gpt-test", + cerebras_delay_s: float = 0.0, + groq_delay_s: float = 0.0, + cerebras_exc: Exception | None = None, + groq_exc: Exception | None = None, +) -> PromptEnhancer: + enhancer = _build_test_enhancer(cerebras_payload) + enhancer.provider_request_models = { + "cerebras": cerebras_request_model, + "groq": groq_request_model, + } + enhancer.enhance_system_prompt = "enhance prompt" + enhancer.auto_system_prompt = "auto prompt" + enhancer.provider_runtimes = [ + ProviderRuntime( + name="cerebras", + api_key="cerebras-key", + api_base_url=None, + request_model=cerebras_request_model, + client=_DelayedSyncClient( + cerebras_payload, + delay_s=cerebras_delay_s, + exc=cerebras_exc, + ), + ), + ProviderRuntime( + name="groq", + api_key="groq-key", + api_base_url="https://api.groq.com/openai/v1", + request_model=groq_request_model, + client=_DelayedSyncClient( + groq_payload, + delay_s=groq_delay_s, + exc=groq_exc, + ), + ), + ] + enhancer.provider_runtime_stages = [ + list(enhancer.provider_runtimes), + ] + enhancer.provider_success_counts = { + "cerebras": 0, + "groq": 0, + } + enhancer.client = enhancer.provider_runtimes[0].client + return enhancer + + +class _FakeOpenAIClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.chat = type( + "_FakeChat", + (), + {"completions": _FakeSyncCompletions(_chat_payload_with_content("{}"))}, + )() + + +class _FakeCerebrasClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.chat = type( + "_FakeChat", + (), + {"completions": _FakeSyncCompletions(_chat_payload_with_content("{}"))}, + )() + + +def test_parse_json_response_accepts_fenced_json_with_prose(): + parsed = _parse_json_response( + "Here is the rewrite:\n```json\n{\"segment_prompts\":[\"A\",\"B\"]}\n```\nThanks." + ) + assert parsed == {"segment_prompts": ["A", "B"]} + + +def test_parse_json_response_extracts_first_embedded_object(): + parsed = _parse_json_response( + "Model output:\n{\"segment_prompts\":[\"A\",\"B\"]}\n(complete)" + ) + assert parsed == {"segment_prompts": ["A", "B"]} + + +def test_load_prompt_required_falls_back_to_default_path(tmp_path): + fallback_path = tmp_path / "next_segment_system_prompt.md" + fallback_path.write_text("fallback prompt\n", encoding="utf-8") + + loaded_prompt = _load_prompt_required( + str(tmp_path / "prompts.local" / "next_segment_system_prompt.md"), + "next-segment", + str(fallback_path), + ) + + assert loaded_prompt == "fallback prompt" + + +def test_load_prompt_required_unwraps_python_style_system_prompt_assignment(tmp_path): + prompt_path = tmp_path / "rewrite_window_system_prompt.md" + prompt_path.write_text( + 'SYSTEM_PROMPT = """\nfirst line\nsecond line\n"""\n', + encoding="utf-8", + ) + + loaded_prompt = _load_prompt_required( + str(prompt_path), + "rewrite-window", + ) + + assert loaded_prompt == "first line\nsecond line" + + +def test_build_client_supports_cerebras_provider(monkeypatch): + monkeypatch.setattr(prompt_enhancer_module, "Cerebras", _FakeCerebrasClient) + + enhancer = PromptEnhancer.__new__(PromptEnhancer) + enhancer.provider = "cerebras" + enhancer.api_key = "cerebras-key" + enhancer.api_base_url = None + + client = PromptEnhancer._build_client(enhancer) + + assert isinstance(client, _FakeCerebrasClient) + assert client.kwargs == { + "api_key": "cerebras-key", + } + + +def test_build_client_supports_groq_provider(monkeypatch): + monkeypatch.setattr(prompt_enhancer_module, "OpenAI", _FakeOpenAIClient) + + enhancer = PromptEnhancer.__new__(PromptEnhancer) + enhancer.provider = "groq" + enhancer.api_key = "groq-key" + enhancer.api_base_url = "https://api.groq.com/openai/v1" + + client = PromptEnhancer._build_client(enhancer) + + assert isinstance(client, _FakeOpenAIClient) + assert client.kwargs == { + "api_key": "groq-key", + "base_url": "https://api.groq.com/openai/v1", + } + + +def test_rewrite_prompt_sequence_accepts_segment_prompts_output(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "preset_a" + assert result.rollout_label == "Preset A" + assert result.prompts == ["A", "B"] + + +def test_rewrite_prompt_sequence_accepts_legacy_rewritten_prompts_output(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"rewritten_prompts":["A","B"]}') + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "current_rollout" + assert result.rollout_label == "Current rollout" + assert result.prompts == ["A", "B"] + + +def test_rewrite_prompt_sequence_accepts_segment_dicts_without_top_level_rollout_metadata(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"segments":[{"prompt":"A"},{"text":"B"}]}' + ) + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + preset_id="preset_a", + preset_label="Preset A", + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "preset_a" + assert result.rollout_label == "Preset A" + assert result.prompts == ["A", "B"] + + +def test_rewrite_prompt_sequence_accepts_numbered_prose_output(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + "The user is asking for a cinematic rewrite.\n\n" + "1. A dog bounds across the moon's dusty surface, kicking up silver regolith as it chases a rabbit beneath the black sky.\n" + "2. The rabbit darts around a crater rim while the dog lunges after it, Earth glowing blue in the distance.\n" + ) + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.rollout_id == "current_rollout" + assert result.rollout_label == "Current rollout" + assert result.prompts == [ + "A dog bounds across the moon's dusty surface, kicking up silver regolith as it chases a rabbit beneath the black sky.", + "The rabbit darts around a crater rim while the dog lunges after it, Earth glowing blue in the distance.", + ] + + +def test_enhance_prompt_prefers_cerebras_before_groq_fallback(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content('{"prompt":"Cerebras prompt"}'), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_delay_s=0.01, + groq_delay_s=0.01, + ) + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "cerebras" + assert result.model == "gpt-test" + assert result.prompt == "Cerebras prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 1, + "groq": 0, + } + + +def test_enhance_prompt_uses_groq_when_cerebras_fails(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content("{}"), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_exc=RuntimeError("cerebras overloaded"), + groq_delay_s=0.01, + ) + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "groq" + assert result.prompt == "Groq prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 0, + "groq": 1, + } + + +def test_enhance_prompt_can_use_groq_when_cerebras_times_out(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content('{"prompt":"Slow Cerebras prompt"}'), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_delay_s=0.4, + groq_delay_s=0.01, + ) + enhancer.http_timeout_ms = 50 + enhancer.default_timeout_ms = 50 + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + timeout_ms=50, + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "groq" + assert result.prompt == "Groq prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 0, + "groq": 1, + } + + +def test_enhance_prompt_can_use_cerebras_when_it_returns_first(): + enhancer = _build_staged_enhancer( + cerebras_payload=_chat_payload_with_content('{"prompt":"Primary Cerebras prompt"}'), + groq_payload=_chat_payload_with_content('{"prompt":"Groq prompt"}'), + cerebras_delay_s=0.01, + groq_delay_s=0.08, + ) + + result = asyncio.run( + enhancer.enhance_prompt( + "A rainy alley at night", + mode="single_clip", + ) + ) + + assert result.fallback_used is False + assert result.error is None + assert result.provider == "cerebras" + assert result.model == "gpt-test" + assert result.prompt == "Primary Cerebras prompt" + assert enhancer.get_provider_success_counts() == { + "cerebras": 1, + "groq": 0, + } + + +def test_rewrite_prompt_sequence_keeps_raw_output_on_parse_error(): + enhancer = _build_test_enhancer( + _chat_payload_with_content("I cannot comply with JSON right now.") + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is True + assert "No JSON object found in assistant response." in (result.error or "") + assert result.raw_response_text == "I cannot comply with JSON right now." + assert result.rollout_id == "current_rollout" + assert result.rollout_label == "Current rollout" + assert result.prompts == ["prompt one", "prompt two"] + + +def test_rewrite_prompt_sequence_uses_current_rollout_payload_shape(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}' + ) + ) + captured = { + "body": None, + "timeout_seconds": None, + } + + async def _fake_request_content_with_body(*, body: dict, timeout_seconds: float): + captured["body"] = body + captured["timeout_seconds"] = timeout_seconds + return ( + _chat_payload_with_content( + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}' + ), + '{"id":"rewritten_rollout","label":"Rewritten Rollout","segment_prompts":["A","B"]}', + ) + + enhancer._request_content_with_body = _fake_request_content_with_body # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + preset_id="preset_a", + preset_label="Preset A", + rewrite_instruction="make it cinematic", + rewrite_model="gpt-test", + rewrite_temperature=0.2, + timeout_ms=800, + ) + ) + + assert result.fallback_used is False + assert captured["body"]["messages"][0] == { + "role": "system", + "content": "system prompt", + } + assert captured["body"]["messages"][1]["role"] == "user" + assert prompt_enhancer_module.json.loads(captured["body"]["messages"][1]["content"]) == { + "mode": "edit_existing_rollout", + "request": ( + "Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical." + ), + "user_instruction": "make it cinematic", + "current_rollout": { + "id": "preset_a", + "label": "Preset A", + "segment_prompts": ["prompt one", "prompt two"], + }, + } + + +def test_rewrite_prompt_sequence_supports_new_rollout_mode(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}' + ) + ) + captured = { + "body": None, + } + + async def _fake_request_content_with_body(*, body: dict, timeout_seconds: float): + del timeout_seconds + captured["body"] = body + return ( + _chat_payload_with_content( + '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}' + ), + '{"id":"custom_editable","label":"Custom rollout","segment_prompts":[' + '"A","B","C","D","E","F"]}', + ) + + enhancer._request_content_with_body = _fake_request_content_with_body # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + [], + preset_id="custom_editable", + preset_label="Custom rollout", + rewrite_instruction="A moonbase corridor thriller with flooding and red alarms", + rewrite_model="gpt-test", + rewrite_temperature=0.2, + timeout_ms=800, + ) + ) + + assert result.fallback_used is False + assert result.prompts == ["A", "B", "C", "D", "E", "F"] + assert prompt_enhancer_module.json.loads(captured["body"]["messages"][1]["content"]) == { + "mode": "new_rollout", + "request": ( + "Rewrite all segment prompts with improved continuity and cinematic detail. " + "Keep count and ordering identical." + ), + "user_instruction": "A moonbase corridor thriller with flooding and red alarms", + "desired_segment_count": 6, + "rollout_id_hint": "custom_editable", + "rollout_label_hint": "Custom rollout", + } + + +def test_rewrite_prompt_sequence_uses_session_override_system_prompt(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_all_system_prompt = "shared system prompt" + captured = { + "body": None, + } + + async def _fake_request_content_with_body(*, body: dict, timeout_seconds: float): + del timeout_seconds + captured["body"] = body + return ( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ), + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}', + ) + + enhancer._request_content_with_body = _fake_request_content_with_body # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + preset_id="preset_a", + preset_label="Preset A", + rewrite_instruction="make it cinematic", + rewrite_model="gpt-test", + system_prompt_override="session specific system prompt", + ) + ) + + assert result.fallback_used is False + assert captured["body"]["messages"][0] == { + "role": "system", + "content": "session specific system prompt", + } + + +def test_resolve_rewrite_new_rollout_system_prompt_uses_dedicated_prompt(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_all_system_prompt = "shared rewrite system prompt" + enhancer.rewrite_user_system_prompt = "new rollout rewrite system prompt" + + resolved = enhancer.resolve_rewrite_new_rollout_system_prompt(None) + + assert resolved == "new rollout rewrite system prompt" + + +def test_resolve_rewrite_new_rollout_system_prompt_prefers_override(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_all_system_prompt = "shared rewrite system prompt" + enhancer.rewrite_user_system_prompt = "new rollout rewrite system prompt" + + resolved = enhancer.resolve_rewrite_new_rollout_system_prompt( + "session specific system prompt" + ) + + assert resolved == "session specific system prompt" + + +def test_generate_auto_prompt_uses_selected_model(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"next_prompt":"Auto next"}') + ) + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + enhancer.rewrite_default_model = "gpt-test" + captured_model = {"value": None} + + async def _fake_request_content( + *, + system_prompt: str, + user_payload: dict, + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ): + captured_model["value"] = model + return ( + _chat_payload_with_content('{"next_prompt":"Auto next"}'), + '{"next_prompt":"Auto next"}', + ) + + enhancer._request_content = _fake_request_content # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.generate_auto_prompt( + locked_segments=["One"], + next_segment_idx=2, + model="gpt-alt", + timeout_ms=800, + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.prompt == "Auto next" + assert captured_model["value"] == "gpt-alt" + assert result.model == "gpt-alt" + + +def test_enhance_prompt_uses_selected_model(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"next_prompt":"Enhanced next"}') + ) + enhancer.enhance_system_prompt = "enhance system prompt" + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + enhancer.rewrite_default_model = "gpt-test" + captured_model = {"value": None} + + async def _fake_request_content( + *, + system_prompt: str, + user_payload: dict, + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ): + captured_model["value"] = model + return ( + _chat_payload_with_content('{"next_prompt":"Enhanced next"}'), + '{"next_prompt":"Enhanced next"}', + ) + + enhancer._request_content = _fake_request_content # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.enhance_prompt( + "A user live prompt", + locked_segments=["Segment 1"], + next_segment_idx=2, + model="gpt-alt", + timeout_ms=800, + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.prompt == "Enhanced next" + assert captured_model["value"] == "gpt-alt" + assert result.model == "gpt-alt" + + +def test_enhance_prompt_single_clip_uses_auto_extension_prompt_and_prompt_field(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"prompt":"Extended single clip"}') + ) + enhancer.enhance_system_prompt = "enhance system prompt" + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + enhancer.rewrite_default_model = "gpt-test" + captured = { + "model": None, + "system_prompt": None, + "user_payload": None, + } + + async def _fake_request_content( + *, + system_prompt: str, + user_payload: dict, + timeout_ms: int, + model: str | None = None, + temperature: float | None = None, + ): + del timeout_ms, temperature + captured["model"] = model + captured["system_prompt"] = system_prompt + captured["user_payload"] = user_payload + return ( + _chat_payload_with_content('{"prompt":"Extended single clip"}'), + '{"prompt":"Extended single clip"}', + ) + + enhancer._request_content = _fake_request_content # type: ignore[attr-defined] + + result = asyncio.run( + enhancer.enhance_prompt( + "short 5s idea", + mode="single_clip", + model="gpt-alt", + timeout_ms=800, + ) + ) + assert result.fallback_used is False + assert result.error is None + assert result.prompt == "Extended single clip" + assert result.model == "gpt-alt" + assert captured["model"] == "gpt-alt" + assert captured["system_prompt"] == "auto system prompt" + assert captured["user_payload"] == { + "request": ( + "Expand the user prompt into one detailed prompt for a " + "single 5-second LTX-2.3 video clip. Respond with " + 'valid JSON only as {"prompt": "..."}.' # noqa: E501 + ), + "user_prompt": "short 5s idea", + } + + +def test_enhance_prompt_single_clip_rejects_plain_text_response(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + "Medium shot of a woman by a rainy cafe window as she lifts her " + "phone, exhales softly, and the camera makes a slow push in." + ) + ) + enhancer.auto_system_prompt = "auto system prompt" + + result = asyncio.run( + enhancer.enhance_prompt( + "woman in a cafe looking at her phone", + mode="single_clip", + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert "No JSON object found in assistant response." in result.error + assert result.prompt == "" + + +def test_enhance_prompt_single_clip_rejects_segment_prompts_json(): + enhancer = _build_test_enhancer( + _chat_payload_with_content('{"segment_prompts":["A","B"]}') + ) + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.rewrite_default_model = "gpt-test" + + result = asyncio.run( + enhancer.enhance_prompt( + "short 5s idea", + mode="single_clip", + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert result.prompt == "" + assert "Missing prompt string." in (result.error or "") + + +def test_enhance_prompt_requires_json_and_does_not_fallback_to_raw_text(): + enhancer = _build_test_enhancer( + _chat_payload_with_content("A cinematic continuation with slow dolly movement.") + ) + enhancer.enhance_system_prompt = "enhance system prompt" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.rewrite_default_model = "gpt-test" + + result = asyncio.run( + enhancer.enhance_prompt( + "raw input prompt", + locked_segments=["Segment 1"], + next_segment_idx=2, + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert result.prompt == "" + assert "No JSON object found in assistant response." in (result.error or "") + + +def test_generate_auto_prompt_requires_json_and_does_not_fallback_to_raw_text(): + enhancer = _build_test_enhancer( + _chat_payload_with_content("A calm, grounded continuation with subtle motion.") + ) + enhancer.auto_system_prompt = "auto system prompt" + enhancer.rewrite_model_options = ["gpt-test"] + enhancer.rewrite_default_model = "gpt-test" + + result = asyncio.run( + enhancer.generate_auto_prompt( + locked_segments=["Segment 1"], + next_segment_idx=2, + model="gpt-test", + timeout_ms=800, + ) + ) + assert result.fallback_used is True + assert result.prompt == "" + assert "No JSON object found in assistant response." in (result.error or "") + + +def test_rewrite_prompt_sequence_includes_raw_json_when_content_empty(): + enhancer = _build_test_enhancer( + { + "choices": [ + { + "finish_reason": "length", + "message": { + "content": [], + "refusal": None, + }, + } + ], + "usage": {"completion_tokens": 0}, + } + ) + result = asyncio.run( + enhancer.rewrite_prompt_sequence( + ["prompt one", "prompt two"], + rewrite_instruction="make it cinematic", + ) + ) + assert result.fallback_used is True + assert "No rewrite segment prompts found in assistant response." in (result.error or "") + assert isinstance(result.raw_response_text, str) + assert '"finish_reason": "length"' in result.raw_response_text + + +def test_get_rewrite_model_config_returns_fixed_defaults(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.rewrite_default_model = "gpt-oss-120b" + enhancer.rewrite_model_options = ["gpt-oss-120b"] + + config = enhancer.get_rewrite_model_config() + + assert config["default_model"] == "gpt-oss-120b" + assert config["options"] == ["gpt-oss-120b"] + + +def test_get_prompt_config_includes_auto_extension_prompt(): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + enhancer.enhance_system_prompt_path = "/tmp/next.md" + enhancer.auto_system_prompt_path = "/tmp/auto.md" + enhancer.rewrite_all_system_prompt_path = "/tmp/rewrite.md" + enhancer.rewrite_user_system_prompt_path = "/tmp/rewrite_user.md" + enhancer.enhance_system_prompt = "next prompt" + enhancer.auto_system_prompt = "auto prompt" + enhancer.rewrite_all_system_prompt = "rewrite prompt" + enhancer.rewrite_user_system_prompt = "rewrite user prompt" + + config = enhancer.get_prompt_config() + + assert config["next_segment_system_prompt_path"] == "/tmp/next.md" + assert config["auto_extension_system_prompt_path"] == "/tmp/auto.md" + assert config["rewrite_window_system_prompt_path"] == "/tmp/rewrite.md" + assert config["rewrite_user_system_prompt_path"] == "/tmp/rewrite_user.md" + assert config["next_segment_system_prompt"] == "next prompt" + assert config["auto_extension_system_prompt"] == "auto prompt" + assert config["rewrite_window_system_prompt"] == "rewrite prompt" + assert config["rewrite_user_system_prompt"] == "rewrite user prompt" + assert config["rewrite_model"] == "gpt-test" + assert config["rewrite_model_options"] == ["gpt-test"] + assert config["rewrite_temperature"] == 0.4 + + +def test_get_prompt_config_reports_loaded_fallback_prompt_path(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + rewrite_fallback_path = tmp_path / "rewrite_window_system_prompt.md" + rewrite_fallback_path.write_text("rewrite prompt\n", encoding="utf-8") + next_path = tmp_path / "next.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path = tmp_path / "auto.md" + auto_path.write_text("auto prompt\n", encoding="utf-8") + enhancer.rewrite_all_system_prompt_path = str( + tmp_path / "prompts.local" / "rewrite_window_system_prompt.md" + ) + enhancer.rewrite_all_system_prompt_fallback_path = str(rewrite_fallback_path) + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.reload_system_prompts() + + config = enhancer.get_prompt_config() + + assert config["rewrite_window_system_prompt_path"] == str(rewrite_fallback_path) + + +def test_reload_system_prompts_falls_back_to_rewrite_window_when_user_prompt_empty( + tmp_path, +): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite_window_system_prompt.md" + rewrite_user_path = tmp_path / "rewrite_user_system_prompt.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path.write_text("auto prompt\n", encoding="utf-8") + rewrite_path.write_text("rewrite window prompt\n", encoding="utf-8") + rewrite_user_path.write_text("\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.rewrite_user_system_prompt_path = str(rewrite_user_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_user_system_prompt_fallback_path = None + + enhancer.reload_system_prompts() + + assert enhancer.rewrite_all_system_prompt == "rewrite window prompt" + assert enhancer.rewrite_user_system_prompt == "rewrite window prompt" + assert enhancer.rewrite_user_system_prompt_source_path == str(rewrite_path) + + +def test_save_prompt_config_updates_auto_extension_prompt(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + next_path.write_text("next original\n", encoding="utf-8") + auto_path.write_text("auto original\n", encoding="utf-8") + rewrite_path.write_text("rewrite original\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + + config = enhancer.save_prompt_config( + auto_extension_system_prompt="auto updated", + ) + + assert auto_path.read_text(encoding="utf-8").strip() == "auto updated" + assert config["auto_extension_system_prompt"] == "auto updated" + + +def test_save_prompt_config_updates_rewrite_user_prompt(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + rewrite_user_path = tmp_path / "rewrite_user.md" + next_path.write_text("next original\n", encoding="utf-8") + auto_path.write_text("auto original\n", encoding="utf-8") + rewrite_path.write_text("rewrite original\n", encoding="utf-8") + rewrite_user_path.write_text("rewrite user original\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.rewrite_user_system_prompt_path = str(rewrite_user_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_user_system_prompt_fallback_path = None + + config = enhancer.save_prompt_config( + rewrite_user_system_prompt="rewrite user updated", + ) + + assert rewrite_user_path.read_text(encoding="utf-8").strip() == "rewrite user updated" + assert config["rewrite_user_system_prompt"] == "rewrite user updated" + + +def test_save_prompt_config_updates_rewrite_model(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path.write_text("auto prompt\n", encoding="utf-8") + rewrite_path.write_text("rewrite prompt\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + enhancer.rewrite_default_model = "gpt-test" + enhancer.rewrite_model_options = ["gpt-test", "gpt-alt"] + + config = enhancer.save_prompt_config( + rewrite_model="gpt-alt", + ) + + assert enhancer.rewrite_default_model == "gpt-alt" + assert config["rewrite_model"] == "gpt-alt" + assert config["rewrite_model_options"] == ["gpt-test", "gpt-alt"] + + +def test_save_prompt_config_updates_rewrite_temperature(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite.md" + next_path.write_text("next prompt\n", encoding="utf-8") + auto_path.write_text("auto prompt\n", encoding="utf-8") + rewrite_path.write_text("rewrite prompt\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + + config = enhancer.save_prompt_config( + rewrite_temperature=1.3, + ) + + assert enhancer.rewrite_default_temperature == 1.3 + assert config["rewrite_temperature"] == 1.3 + + +def test_save_prompt_config_creates_versioned_backup_for_existing_prompt(tmp_path): + enhancer = _build_test_enhancer( + _chat_payload_with_content( + '{"id":"preset_a","label":"Preset A","segment_prompts":["A","B"]}' + ) + ) + next_path = tmp_path / "next.md" + auto_path = tmp_path / "auto.md" + rewrite_path = tmp_path / "rewrite_window_system_prompt.md" + next_path.write_text("next original\n", encoding="utf-8") + auto_path.write_text("auto original\n", encoding="utf-8") + rewrite_path.write_text("rewrite original\n", encoding="utf-8") + enhancer.enhance_system_prompt_path = str(next_path) + enhancer.auto_system_prompt_path = str(auto_path) + enhancer.rewrite_all_system_prompt_path = str(rewrite_path) + enhancer.enhance_system_prompt_fallback_path = None + enhancer.auto_system_prompt_fallback_path = None + enhancer.rewrite_all_system_prompt_fallback_path = None + + enhancer.save_prompt_config( + rewrite_window_system_prompt="rewrite updated", + ) + + backup_paths = sorted( + tmp_path.glob("rewrite_window_system_prompt.*.bak.md") + ) + + assert rewrite_path.read_text(encoding="utf-8").strip() == "rewrite updated" + assert len(backup_paths) == 1 + assert backup_paths[0].read_text(encoding="utf-8").strip() == "rewrite original" + + +def test_resolve_prompt_save_path_prefers_existing_fallback_file(tmp_path): + primary_path = tmp_path / "prompts.local" / "rewrite.md" + fallback_path = tmp_path / "rewrite.md" + fallback_path.write_text("existing\n", encoding="utf-8") + + resolved = _resolve_prompt_save_path(str(primary_path), str(fallback_path)) + + assert resolved == str(fallback_path) + + +def test_build_prompt_backup_path_preserves_extension(tmp_path): + prompt_path = tmp_path / "rewrite_window_system_prompt.md" + + backup_path = _build_prompt_backup_path(prompt_path) + + assert backup_path.parent == tmp_path + assert re.match( + r"rewrite_window_system_prompt\.\d{8}_\d{6}\.bak\.md$", + backup_path.name, + ) diff --git a/apps/dreamverse/server/tests/test_session_init_image.py b/apps/dreamverse/server/tests/test_session_init_image.py new file mode 100644 index 000000000..2f33b2ea9 --- /dev/null +++ b/apps/dreamverse/server/tests/test_session_init_image.py @@ -0,0 +1,75 @@ +# pyright: reportMissingTypeArgument=false +import base64 +import io +from pathlib import Path + +from PIL import Image +import pytest + +from session_init_image import ( + MAX_SESSION_INIT_IMAGE_BYTES, + cleanup_session_init_image, + persist_session_init_image, +) + + +def make_data_url(format_name: str = "PNG") -> str: + image = Image.new("RGB", (2, 2), color=(16, 32, 64)) + buffer = io.BytesIO() + image.save(buffer, format=format_name) + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + mime_type = "image/png" if format_name == "PNG" else "image/jpeg" + return f"data:{mime_type};base64,{encoded}" + + +def test_persist_session_init_image_saves_normalized_file(tmp_path: Path): + session_image = persist_session_init_image( + { + "name": "frame.png", + "mime_type": "image/png", + "data_url": make_data_url("PNG"), + }, + temp_root=tmp_path, + ) + + assert session_image is not None + assert session_image.display_name == "frame.png" + assert session_image.file_path.is_file() + + cleanup_session_init_image(session_image) + assert not session_image.temp_dir.exists() + + +def test_persist_session_init_image_returns_none_when_missing_data(): + assert persist_session_init_image(None) is None + assert persist_session_init_image({"name": "frame.png", "data_url": ""}) is None + + +def test_persist_session_init_image_rejects_unsupported_mime(): + with pytest.raises(ValueError, match="PNG, JPEG, or WebP"): + persist_session_init_image( + { + "name": "frame.gif", + "mime_type": "image/gif", + "data_url": "data:image/gif;base64,R0lGODlhAQABAAAAACw=", + } + ) + + +def test_persist_session_init_image_rejects_large_payload(monkeypatch): + data_url = make_data_url("PNG") + oversized = "a" * (MAX_SESSION_INIT_IMAGE_BYTES + 1) + + def fake_b64decode(value: str, validate: bool = True): + return oversized.encode("ascii") + + monkeypatch.setattr(base64, "b64decode", fake_b64decode) + + with pytest.raises(ValueError, match="15 MB or smaller"): + persist_session_init_image( + { + "name": "frame.png", + "mime_type": "image/png", + "data_url": data_url, + } + ) diff --git a/apps/dreamverse/server/tests/test_session_logging.py b/apps/dreamverse/server/tests/test_session_logging.py index a77ff0de2..1a74ae17b 100644 --- a/apps/dreamverse/server/tests/test_session_logging.py +++ b/apps/dreamverse/server/tests/test_session_logging.py @@ -13,8 +13,6 @@ from pathlib import Path from types import SimpleNamespace from fastapi import WebSocketDisconnect -from fastvideo.entrypoints.streaming.protocol import MediaChunk, MediaComplete, MediaInit -from fastvideo.entrypoints.streaming.session_logger import SessionEventLogger from PIL import Image import pytest @@ -28,6 +26,8 @@ os.environ.setdefault("GROQ_API_KEY", "dummy") import main as server_main # noqa: E402 import runtime # noqa: E402 +from session_logger import SessionEventLogger # noqa: E402 +from worker_ipc import MediaChunk, MediaComplete, MediaInit # noqa: E402 from session import controller as session_controller # noqa: E402 pytestmark = pytest.mark.gpu diff --git a/apps/dreamverse/server/worker_ipc.py b/apps/dreamverse/server/worker_ipc.py new file mode 100644 index 000000000..e11478d0f --- /dev/null +++ b/apps/dreamverse/server/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 diff --git a/fastvideo/entrypoints/streaming/gpu_pool.py b/fastvideo/entrypoints/streaming/gpu_pool.py index f6c881a61..1da2f6f38 100644 --- a/fastvideo/entrypoints/streaming/gpu_pool.py +++ b/fastvideo/entrypoints/streaming/gpu_pool.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 -# pyright: reportMissingTypeArgument=false """GPU pool manager for the streaming server. Replaces the single-generator path in PR 7.5 with a typed pool @@ -27,9 +26,7 @@ from __future__ import annotations import asyncio import multiprocessing as mp -import os import queue -import subprocess import threading import time import uuid @@ -74,13 +71,6 @@ class PoolAssignment: pinned_at: float = field(default_factory=time.monotonic) -# Backwards-compatible name for the Dreamverse product controller during -# the monorepo migration. The controller only uses this name for runtime -# annotations; Phase 4 replaces the product-local controller path with the -# generic streaming runtime. -GPUSlot = PoolAssignment - - class GpuPool(ABC): """Abstract GPU pool. @@ -540,32 +530,7 @@ def _default_worker_factory( ) -def get_available_gpus() -> list[int]: - """Return CUDA-visible GPU ids, falling back to ``nvidia-smi`` then ``[0]``.""" - - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") - if cuda_visible: - return [int(item.strip()) for item in cuda_visible.split(",") if item.strip()] - - try: - result = subprocess.run( - ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], - capture_output=True, - text=True, - check=False, - ) - except OSError: - return [0] - - if result.returncode != 0: - return [0] - gpu_ids = [int(line.strip()) for line in result.stdout.splitlines() if line.strip()] - return gpu_ids or [0] - - __all__ = [ - "get_available_gpus", - "GPUSlot", "GpuPool", "InProcessGpuPool", "PoolAcquireTimeout", diff --git a/fastvideo/entrypoints/streaming/protocol.py b/fastvideo/entrypoints/streaming/protocol.py index abc8884a7..e5912570d 100644 --- a/fastvideo/entrypoints/streaming/protocol.py +++ b/fastvideo/entrypoints/streaming/protocol.py @@ -11,7 +11,6 @@ both places in the same PR. """ from __future__ import annotations -from dataclasses import dataclass from typing import Annotated, Any, Literal, Union from pydantic import BaseModel, ConfigDict, Field @@ -143,8 +142,6 @@ class StepComplete(BaseModel): class MediaInit(BaseModel): """Descriptor for the fMP4 initialization segment that follows.""" - model_config = ConfigDict(extra="allow") - type: Literal["media_init"] = "media_init" segment_idx: int mime: str = "video/mp4; codecs=\"avc1.64001f, mp4a.40.2\"" @@ -152,36 +149,6 @@ class MediaInit(BaseModel): mode: Literal["av_fmp4"] = "av_fmp4" -@dataclass(frozen=True) -class MediaChunk: - """One fMP4 chunk emitted by the legacy Dreamverse worker bridge.""" - - 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: - """Legacy Dreamverse fMP4 segment-complete event.""" - - user_id: str - segment_idx: int - stream_id: str - chunks: int - - class MediaSegmentComplete(BaseModel): type: Literal["media_segment_complete"] = "media_segment_complete" segment_idx: int @@ -272,8 +239,6 @@ __all__ = [ "Ltx2StreamStart", "LoopGenerationUpdated", "MediaInit", - "MediaChunk", - "MediaComplete", "MediaSegmentComplete", "QueueStatus", "SeedPromptsUpdated", diff --git a/fastvideo/entrypoints/streaming/session_init_image.py b/fastvideo/entrypoints/streaming/session_init_image.py index 157feb060..cf4672a21 100644 --- a/fastvideo/entrypoints/streaming/session_init_image.py +++ b/fastvideo/entrypoints/streaming/session_init_image.py @@ -6,9 +6,6 @@ import base64 import binascii import contextlib import os -from pathlib import Path -import re -import shutil import tempfile from dataclasses import dataclass from typing import Any @@ -21,7 +18,6 @@ _ACCEPTED_MIMES = { } _MAX_IMAGE_BYTES = 32 * 1024 * 1024 # 32 MiB cap -_DATA_URL_RE = re.compile(r"^data:(?P[-\w.+/]+);base64,(?P[A-Za-z0-9+/=\s]+)$") @dataclass(frozen=True) @@ -36,24 +32,6 @@ class SessionInitImage: display_name: str mime: str - @property - def file_path(self) -> Path: - """Compatibility alias used by the migrated Dreamverse controller.""" - - return Path(self.path) - - @property - def temp_dir(self) -> Path: - """Compatibility alias for cleanup of the temporary image directory.""" - - return Path(self.path).parent - - @property - def mime_type(self) -> str: - """Compatibility alias for the UI payload's MIME field.""" - - return self.mime - def persist_session_init_image( payload: Any, @@ -79,19 +57,11 @@ def persist_session_init_image( if not isinstance(payload, dict): raise ValueError("session init image must be an object") - data_url = payload.get("data_url") - if isinstance(data_url, str) and data_url.strip(): - match = _DATA_URL_RE.match(data_url.strip()) - if match is None: - raise ValueError("session init image data URL is not valid base64") - mime = payload.get("mime_type") or match.group("mime").strip().lower() - data_b64 = match.group("data") - else: - mime = payload.get("mime") - data_b64 = payload.get("data") + mime = payload.get("mime") if mime not in _ACCEPTED_MIMES: raise ValueError(f"session init image mime {mime!r} is not one of " f"{sorted(_ACCEPTED_MIMES)}") + data_b64 = payload.get("data") if not isinstance(data_b64, str): raise ValueError("session init image data must be a base64 string") try: @@ -117,14 +87,6 @@ def persist_session_init_image( return SessionInitImage(path=path, display_name=display_name, mime=mime) -def cleanup_session_init_image(session_image: SessionInitImage | None) -> None: - """Remove a persisted session init image and its temporary directory.""" - - if session_image is None: - return - shutil.rmtree(session_image.temp_dir, ignore_errors=True) - - def _sanitize_display_name(name: Any) -> str | None: if not isinstance(name, str): return None @@ -137,6 +99,5 @@ def _sanitize_display_name(name: Any) -> str | None: __all__ = [ "SessionInitImage", - "cleanup_session_init_image", "persist_session_init_image", ] diff --git a/fastvideo/entrypoints/streaming/session_logger.py b/fastvideo/entrypoints/streaming/session_logger.py index cac0e66d1..1a88e4d57 100644 --- a/fastvideo/entrypoints/streaming/session_logger.py +++ b/fastvideo/entrypoints/streaming/session_logger.py @@ -9,16 +9,12 @@ format; keeping the same shape makes log tooling portable. from __future__ import annotations import contextlib -from datetime import datetime, timezone -import asyncio import json import os import re -import socket import threading import time from dataclasses import dataclass, field -from pathlib import Path from typing import Any, TextIO _FILENAME_SANITIZE_RE = re.compile(r"[^A-Za-z0-9._-]") @@ -111,50 +107,7 @@ class SessionLogger: return handle, lock -def _utc_now_iso() -> str: - return datetime.now(timezone.utc).isoformat() - - -class SessionEventLogger: - """Dreamverse-compatible async session event logger. - - This keeps the product server importable during the monorepo migration while - the generic streaming logger remains :class:`SessionLogger`. - """ - - def __init__(self, root_dir: Path): - self.hostname = socket.gethostname() - timestamp = datetime.now(timezone.utc).strftime("%y%m%d_%H%M%S") - self.directory = root_dir / self.hostname - self.path = self.directory / f"{timestamp}.jsonl" - self._lock = asyncio.Lock() - - self.directory.mkdir(parents=True, exist_ok=True) - self.path.touch(exist_ok=False) - - async def write_event( - self, - *, - event: str, - client_id: str, - payload: dict[str, Any] | None = None, - ) -> None: - entry = { - "ts": _utc_now_iso(), - "event": event, - "hostname": self.hostname, - "client_id": client_id, - } - if payload: - entry.update(payload) - - async with self._lock: - with self.path.open("a", encoding="utf-8") as fp: - fp.write(json.dumps(entry, ensure_ascii=False) + "\n") - - __all__ = [ "SessionLogEvent", - "SessionEventLogger", "SessionLogger", ]