SolitaryThinker
aef113a861
Source: Dreamverse repo at will/integrate-public-fastvideo @ ec8ef92.\n\nCarry the generic-merged backend modules as Dreamverse product-local code because the public streaming APIs are not drop-in compatible with the product runtime yet. Revert the Phase 2 public compatibility shims so FastVideo's public streaming surface returns to the pre-migration state.\n\nCo-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com>\nCo-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com>\nCo-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com>\nCo-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
1162 lines
44 KiB
Python
1162 lines
44 KiB
Python
# 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()
|