Add job queueing
This commit is contained in:
@@ -16,6 +16,8 @@ The UI currently supports:
|
||||
```
|
||||
apps/fastvideo_studio/
|
||||
├── server.py / job_runner.py / database.py # FastAPI backend + job lifecycle
|
||||
├── job_queue.py # Queue scheduling rules + last-frame references (GPU-free)
|
||||
├── merge.py # Join a scene's clips into one video (ffmpeg stream copy, GPU-free)
|
||||
├── mock_server.py # In-memory API mock for e2e tests
|
||||
├── models/ # Pydantic request models (shared with the mock)
|
||||
├── training_config.py # Studio workloads → fastvideo/train YAML configs
|
||||
@@ -88,6 +90,42 @@ Running it this way allows you to pass command line parameters.
|
||||
The API server defaults to port 8189. To configure this, you can edit `.env.local` file.
|
||||
Refer to `.env.example` for reference.
|
||||
|
||||
### Job queue
|
||||
|
||||
`POST /api/jobs/queue` puts jobs in a queue; the server runs them oldest-first,
|
||||
`--max-concurrent-jobs` at a time (default 1). Start on a job while that many
|
||||
are already running is refused; queue it instead. Queued jobs survive a server
|
||||
restart and carry on.
|
||||
|
||||
Inference jobs run on that many long-lived worker threads. The model's GPU
|
||||
processes are killed when the thread that started them exits, so a worker that
|
||||
never exits is what keeps the model loaded from one job to the next (it loads
|
||||
once, not per job). A worker holds one model at a time: a job with a different
|
||||
model or settings replaces it, and a failed or stopped job unloads it so the next
|
||||
one starts clean. Each extra worker loads its own copy of the model.
|
||||
|
||||
A job can start from another job's last frame by holding a reference whose
|
||||
source is `job-last-frame:<job id>`. The frame is extracted when the job runs,
|
||||
so the reference can be added before that job has finished, and the queue holds
|
||||
the job until it has. If that job fails, is stopped, or is deleted, the waiting
|
||||
job fails with a message saying so (and so does anything waiting on *it*); fix
|
||||
the earlier job and queue them again. Starting such a job directly while its
|
||||
source is unfinished is refused with a 409.
|
||||
|
||||
### Merging a scene
|
||||
|
||||
**Merge scene** on the Scenes page joins a scene's clips, in the order shown (the
|
||||
take chosen for each), into one video under `<output-dir>/merged/`, then plays it
|
||||
and offers a download. It is enabled once every clip has finished, and refuses
|
||||
otherwise, so a scene missing a clip can't pass for the whole thing.
|
||||
|
||||
The clips are copied, not re-encoded (about a second for 15 clips, no quality
|
||||
loss). That needs them to share size, frame rate and audio, which clips from one
|
||||
scene do; if one differs, the merge says which. Each clip is cut at the end of its
|
||||
last video frame, because H3's audio runs a few milliseconds longer and would
|
||||
otherwise leave a small gap at every join. It uses `ffmpeg` from `PATH`,
|
||||
`$FASTVIDEO_FFMPEG_BIN`, or the copy bundled with `imageio-ffmpeg`.
|
||||
|
||||
### API Endpoints
|
||||
|
||||
| Method | Path | Description |
|
||||
@@ -98,6 +136,11 @@ Refer to `.env.example` for reference.
|
||||
| `POST` | `/api/jobs` | Create a new job |
|
||||
| `POST` | `/api/jobs/{id}/start` | Start a pending/stopped/failed job |
|
||||
| `POST` | `/api/jobs/{id}/stop` | Request a running job to stop |
|
||||
| `POST` | `/api/jobs/queue` | Queue jobs (`{"job_ids": [...]}`), in order |
|
||||
| `POST` | `/api/jobs/{id}/queue` | Queue one job |
|
||||
| `POST` | `/api/jobs/{id}/dequeue` | Take a queued job back out of the queue |
|
||||
| `POST` | `/api/scenes/merge` | Join finished jobs' videos (`{"job_ids": [...], "name": ""}`) in order |
|
||||
| `GET` | `/api/merged/{filename}` | Stream a merged scene video |
|
||||
| `DELETE` | `/api/jobs/{id}` | Delete a job |
|
||||
| `GET` | `/api/jobs/{id}/video` | Stream the generated video/image |
|
||||
| `GET` | `/api/jobs/{id}/logs` | Get job logs (polling, supports `?after=`) |
|
||||
|
||||
@@ -110,6 +110,7 @@ def _migrate_db(conn: sqlite3.Connection) -> None:
|
||||
_add_column_if_missing(conn, "jobs", "generator_update_interval", "INTEGER", "5")
|
||||
_add_column_if_missing(conn, "jobs", "real_score_model_path", "TEXT", "''")
|
||||
_add_column_if_missing(conn, "jobs", "fake_score_model_path", "TEXT", "''")
|
||||
_add_column_if_missing(conn, "jobs", "queued_at", "REAL", "NULL")
|
||||
# Settings table
|
||||
_add_column_if_missing(conn, "settings", "vae_cpu_offload", "INTEGER", "0")
|
||||
_add_column_if_missing(conn, "settings", "image_encoder_cpu_offload", "INTEGER", "0")
|
||||
@@ -326,6 +327,7 @@ class Database:
|
||||
return
|
||||
allowed = {
|
||||
"status",
|
||||
"queued_at",
|
||||
"started_at",
|
||||
"finished_at",
|
||||
"error",
|
||||
@@ -345,6 +347,42 @@ class Database:
|
||||
self._execute(sql, tuple(vals))
|
||||
self._commit()
|
||||
|
||||
#: Job fields stored under a different column name, and how they are encoded.
|
||||
_CONFIG_COLUMNS = {"references": "references_json"}
|
||||
_BOOL_CONFIG_FIELDS = {
|
||||
"dit_cpu_offload",
|
||||
"text_encoder_cpu_offload",
|
||||
"vae_cpu_offload",
|
||||
"image_encoder_cpu_offload",
|
||||
"use_fsdp_inference",
|
||||
"enable_torch_compile",
|
||||
"dmd_use_vsa",
|
||||
}
|
||||
|
||||
def update_job_config(self, job_id: str, updates: dict[str, Any]) -> None:
|
||||
"""Persist edits to a job's configuration (prompt, references, sizes...).
|
||||
|
||||
Fields the table has no column for are skipped, as they are on insert.
|
||||
"""
|
||||
columns = _get_table_columns(self._conn(), "jobs")
|
||||
cols = []
|
||||
vals: list[Any] = []
|
||||
for key, value in updates.items():
|
||||
column = self._CONFIG_COLUMNS.get(key, key)
|
||||
if column not in columns or column == "id":
|
||||
continue
|
||||
if key == "references":
|
||||
value = json.dumps(value or [])
|
||||
elif key in self._BOOL_CONFIG_FIELDS:
|
||||
value = 1 if value else 0
|
||||
cols.append(f"{column} = ?")
|
||||
vals.append(value)
|
||||
if not cols:
|
||||
return
|
||||
vals.append(job_id)
|
||||
self._execute(f"UPDATE jobs SET {', '.join(cols)} WHERE id = ?", tuple(vals))
|
||||
self._commit()
|
||||
|
||||
def delete_job(self, job_id: str) -> bool:
|
||||
"""Delete a job. Returns True if a row was deleted."""
|
||||
cur = self._execute("DELETE FROM jobs WHERE id = ?", (job_id, ))
|
||||
@@ -537,8 +575,31 @@ def _row_to_dataset(row: sqlite3.Row) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _row_to_job(row: sqlite3.Row) -> dict[str, Any]:
|
||||
class _ColumnRow:
|
||||
"""A sqlite3.Row whose ``in`` tests column names.
|
||||
|
||||
sqlite3.Row's own ``in`` tests *values*, so ``"fps" in row`` is always False
|
||||
and a job restored from the database silently lost fps, negative_prompt and
|
||||
the other optional columns.
|
||||
"""
|
||||
|
||||
def __init__(self, row: sqlite3.Row):
|
||||
self._row = row
|
||||
self._columns = set(row.keys())
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
return self._row[key]
|
||||
|
||||
def __contains__(self, key: object) -> bool:
|
||||
return key in self._columns
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
return self._row.keys()
|
||||
|
||||
|
||||
def _row_to_job(sqlite_row: sqlite3.Row) -> dict[str, Any]:
|
||||
"""Convert a DB row to job dict (snake_case, matching Job.to_dict)."""
|
||||
row = _ColumnRow(sqlite_row)
|
||||
int_defaults = {
|
||||
"max_train_steps": 1000,
|
||||
"train_batch_size": 1,
|
||||
@@ -562,6 +623,7 @@ def _row_to_job(row: sqlite3.Row) -> dict[str, Any]:
|
||||
"job_type": _sqlite_row_get(row, "job_type", "inference"),
|
||||
"status": row["status"],
|
||||
"created_at": row["created_at"],
|
||||
"queued_at": row["queued_at"] if "queued_at" in row else None,
|
||||
"started_at": row["started_at"],
|
||||
"finished_at": row["finished_at"],
|
||||
"error": row["error"],
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Save the final frame of a finished clip, so the next clip can start from it.
|
||||
|
||||
Kept apart from ``server.py`` because it needs nothing from FastVideo -- only
|
||||
PyAV -- and so can be exercised without a GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
VIDEO_EXTENSIONS = {".mp4", ".webm", ".avi", ".mov", ".mkv"}
|
||||
|
||||
|
||||
class FrameError(Exception):
|
||||
"""A request that can't be satisfied, carrying the HTTP status to report."""
|
||||
|
||||
def __init__(self, status_code: int, detail: str):
|
||||
super().__init__(detail)
|
||||
self.status_code = status_code
|
||||
self.detail = detail
|
||||
|
||||
|
||||
def save_last_frame(video_path: str, out_path: str) -> None:
|
||||
"""Decode ``video_path`` and write its final frame to ``out_path`` as a PNG.
|
||||
|
||||
Clips are a few hundred frames, so decoding through to the end is simpler
|
||||
and more reliable than seeking (which lands on keyframes, not the last frame).
|
||||
"""
|
||||
import av
|
||||
|
||||
tmp_path = f"{out_path}.tmp.png"
|
||||
with av.open(video_path) as container:
|
||||
last = None
|
||||
for frame in container.decode(container.streams.video[0]):
|
||||
last = frame
|
||||
if last is None:
|
||||
raise ValueError("The video has no frames.")
|
||||
last.to_image().save(tmp_path)
|
||||
os.replace(tmp_path, out_path) # atomic, so a reader never sees a half-written file
|
||||
|
||||
|
||||
def last_frame_for_job(job: Any, upload_dir: str) -> str:
|
||||
"""Path of a PNG of the last frame of ``job``'s video, extracting it if needed.
|
||||
|
||||
Raises :class:`FrameError` with the status the API should return.
|
||||
"""
|
||||
status = getattr(job.status, "value", job.status)
|
||||
if status != "completed" or not job.output_path:
|
||||
raise FrameError(404, "No output available for this job")
|
||||
if not os.path.isfile(job.output_path):
|
||||
raise FrameError(404, "Output file not found on disk")
|
||||
if Path(job.output_path).suffix.lower() not in VIDEO_EXTENSIONS:
|
||||
raise FrameError(400, "This job's output is not a video")
|
||||
if not upload_dir:
|
||||
raise FrameError(503, "Upload directory not configured")
|
||||
|
||||
out_dir = os.path.join(upload_dir, "last_frames")
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
out_path = os.path.abspath(os.path.join(out_dir, f"last_frame_{job.id}.png"))
|
||||
# Reuse an earlier extraction unless the video has been rewritten since.
|
||||
if not (os.path.isfile(out_path) and os.path.getmtime(out_path) >= os.path.getmtime(job.output_path)):
|
||||
try:
|
||||
save_last_frame(job.output_path, out_path)
|
||||
except Exception as e:
|
||||
raise FrameError(500, f"Could not read the last frame: {e}") from e
|
||||
return out_path
|
||||
@@ -0,0 +1,185 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Scheduling rules for queued jobs, and last-frame references between them.
|
||||
|
||||
Kept apart from ``job_runner.py`` because it is pure -- it works on plain
|
||||
snapshots and needs neither FastVideo nor a GPU -- so it can be tested anywhere.
|
||||
|
||||
A job can start from the last frame of another job's video. That frame doesn't
|
||||
exist until the other job finishes, so the reference is *deferred*: its source
|
||||
is ``job-last-frame:<job id>`` and it is resolved to a real image only when the
|
||||
job runs. That link is also the queue's dependency graph: a queued job waits
|
||||
until every job it takes a frame from has completed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
DEFERRED_LAST_FRAME_PREFIX = "job-last-frame:"
|
||||
|
||||
QUEUED = "queued"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
PENDING = "pending"
|
||||
FAILED = "failed"
|
||||
STOPPED = "stopped"
|
||||
|
||||
|
||||
def deferred_last_frame_source(job_id: str) -> str:
|
||||
return f"{DEFERRED_LAST_FRAME_PREFIX}{job_id}"
|
||||
|
||||
|
||||
def dependency_of(reference: Mapping[str, Any] | None) -> str | None:
|
||||
"""The job a deferred reference takes its frame from, or None for any other."""
|
||||
source = (reference or {}).get("source")
|
||||
if isinstance(source, str) and source.startswith(DEFERRED_LAST_FRAME_PREFIX):
|
||||
return source[len(DEFERRED_LAST_FRAME_PREFIX):] or None
|
||||
return None
|
||||
|
||||
|
||||
def dependencies(references: Iterable[Mapping[str, Any]] | None) -> list[str]:
|
||||
"""Job ids a job needs finished first, in reference order, without repeats."""
|
||||
found: list[str] = []
|
||||
for ref in references or []:
|
||||
dep = dependency_of(ref)
|
||||
if dep and dep not in found:
|
||||
found.append(dep)
|
||||
return found
|
||||
|
||||
|
||||
class JobView(Protocol):
|
||||
"""The parts of a job the scheduler looks at."""
|
||||
|
||||
id: str
|
||||
references: Any
|
||||
queued_at: float | None
|
||||
|
||||
@property
|
||||
def status(self) -> Any: # a JobStatus (str enum) or plain str
|
||||
...
|
||||
|
||||
|
||||
def _status(job: Any) -> str:
|
||||
return str(getattr(job.status, "value", job.status))
|
||||
|
||||
|
||||
def circular_dependency(jobs: Mapping[str, JobView], job_id: str) -> list[str] | None:
|
||||
"""The chain of job ids if ``job_id`` (transitively) waits on itself, else None."""
|
||||
|
||||
def walk(current: str, path: list[str]) -> list[str] | None:
|
||||
job = jobs.get(current)
|
||||
if job is None:
|
||||
return None
|
||||
for dep in dependencies(job.references):
|
||||
if dep == job_id:
|
||||
return [*path, current, dep]
|
||||
if dep in path or dep == current:
|
||||
continue
|
||||
found = walk(dep, [*path, current])
|
||||
if found:
|
||||
return found
|
||||
return None
|
||||
|
||||
return walk(job_id, [])
|
||||
|
||||
|
||||
def _label(job: Any) -> str:
|
||||
name = getattr(job, "name", "") or ""
|
||||
return f"'{name}'" if name else job.id
|
||||
|
||||
|
||||
@dataclass
|
||||
class Plan:
|
||||
"""What to do now: jobs to launch, and jobs that can never run."""
|
||||
|
||||
start: list[str] = field(default_factory=list)
|
||||
fail: dict[str, str] = field(default_factory=dict) # job id -> reason
|
||||
|
||||
|
||||
def plan_dispatch(jobs: Mapping[str, JobView], max_concurrent: int) -> Plan:
|
||||
"""Decide which queued jobs launch now and which have become impossible.
|
||||
|
||||
* Queued jobs run oldest-queued first, at most ``max_concurrent`` at a time
|
||||
(counting jobs already running).
|
||||
* A job whose frame source hasn't completed *waits*, and doesn't hold up
|
||||
jobs behind it that don't need it.
|
||||
* A job whose frame source failed, was stopped, or was deleted can never
|
||||
run, so it fails -- and so does anything waiting on *it*.
|
||||
* A job whose frame source is idle (neither queued nor running) keeps
|
||||
waiting: the user may still queue it.
|
||||
"""
|
||||
plan = Plan()
|
||||
status = {jid: _status(job) for jid, job in jobs.items()}
|
||||
queued = sorted(
|
||||
(job for job in jobs.values() if status[job.id] == QUEUED),
|
||||
key=lambda job: (job.queued_at or 0.0, job.id),
|
||||
)
|
||||
|
||||
# Failures cascade down a chain, so repeat until nothing new fails.
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for job in queued:
|
||||
if job.id in plan.fail:
|
||||
continue
|
||||
for dep_id in dependencies(job.references):
|
||||
dep = jobs.get(dep_id)
|
||||
if dep is None:
|
||||
plan.fail[job.id] = "Waiting on the last frame of a job that was deleted."
|
||||
elif dep_id in plan.fail or status[dep_id] in (FAILED, STOPPED):
|
||||
reason = "failed" if dep_id in plan.fail or status[dep_id] == FAILED else "was stopped"
|
||||
plan.fail[job.id] = f"Needs the last frame of {_label(dep)}, which {reason}."
|
||||
else:
|
||||
continue
|
||||
status[job.id] = FAILED
|
||||
changed = True
|
||||
break
|
||||
|
||||
slots = max(max_concurrent, 1) - sum(1 for s in status.values() if s == RUNNING)
|
||||
for job in queued:
|
||||
if slots <= 0:
|
||||
break
|
||||
if job.id in plan.fail:
|
||||
continue
|
||||
if all(status.get(dep) == COMPLETED for dep in dependencies(job.references)):
|
||||
plan.start.append(job.id)
|
||||
slots -= 1
|
||||
return plan
|
||||
|
||||
|
||||
class UnresolvedReference(ValueError):
|
||||
"""A deferred reference can't be turned into an image."""
|
||||
|
||||
|
||||
def resolve_references(
|
||||
references: list[dict[str, Any]] | None,
|
||||
get_job: Callable[[str], Any],
|
||||
last_frame: Callable[[Any], str],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Copy of ``references`` with each deferred last-frame reference made real.
|
||||
|
||||
``last_frame(job)`` returns the path of an extracted frame and may raise
|
||||
(see ``frames.FrameError``); its message is passed on. Other references are
|
||||
returned unchanged, and the stored list is never modified, so a re-run
|
||||
resolves again against the source job's current output.
|
||||
"""
|
||||
resolved: list[dict[str, Any]] = []
|
||||
for ref in references or []:
|
||||
dep_id = dependency_of(ref)
|
||||
if dep_id is None:
|
||||
resolved.append(ref)
|
||||
continue
|
||||
dep = get_job(dep_id)
|
||||
if dep is None:
|
||||
raise UnresolvedReference("A reference needs the last frame of a job that no longer exists.")
|
||||
if _status(dep) != COMPLETED:
|
||||
raise UnresolvedReference(f"A reference needs the last frame of {_label(dep)}, which hasn't finished "
|
||||
f"(it is {_status(dep)}).")
|
||||
try:
|
||||
path = last_frame(dep)
|
||||
except Exception as exc: # FrameError carries a readable .detail
|
||||
raise UnresolvedReference(str(getattr(exc, "detail", None) or exc)) from exc
|
||||
resolved.append({**ref, "source": path})
|
||||
return resolved
|
||||
@@ -17,6 +17,7 @@ import logging
|
||||
import logging.handlers
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -29,6 +30,14 @@ import yaml
|
||||
|
||||
from fastvideo.utils import get_mp_context
|
||||
from fastvideo_studio.database import Database
|
||||
from fastvideo_studio.frames import last_frame_for_job
|
||||
from fastvideo_studio.job_queue import (
|
||||
UnresolvedReference,
|
||||
circular_dependency,
|
||||
dependencies,
|
||||
plan_dispatch,
|
||||
resolve_references,
|
||||
)
|
||||
from fastvideo_studio.training_config import (
|
||||
build_training_config,
|
||||
get_training_env,
|
||||
@@ -46,6 +55,7 @@ _MAX_LOG_LINES = 2000 # ring-buffer cap per job
|
||||
|
||||
class JobStatus(str, enum.Enum):
|
||||
PENDING = "pending"
|
||||
QUEUED = "queued"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
@@ -138,6 +148,7 @@ class Job:
|
||||
finished_at: float | None = None
|
||||
error: str | None = None
|
||||
output_path: str | None = None
|
||||
queued_at: float | None = None
|
||||
log_file_path: str | None = None # Path to the job's log file
|
||||
num_inference_steps: int = 50
|
||||
num_frames: int = 81
|
||||
@@ -195,6 +206,7 @@ class Job:
|
||||
"references": self.references,
|
||||
"status": self.status.value,
|
||||
"created_at": self.created_at,
|
||||
"queued_at": self.queued_at,
|
||||
"started_at": self.started_at,
|
||||
"finished_at": self.finished_at,
|
||||
"error": self.error,
|
||||
@@ -320,6 +332,8 @@ class JobRunner:
|
||||
log_dir: str,
|
||||
database: Database,
|
||||
verbose: bool = False,
|
||||
max_concurrent_jobs: int = 1,
|
||||
upload_dir: str = "",
|
||||
):
|
||||
"""Initialize the job runner.
|
||||
|
||||
@@ -328,19 +342,38 @@ class JobRunner:
|
||||
log_dir: Directory where job log files are saved
|
||||
verbose: Whether to print full tracebacks in error messages
|
||||
database: Optional SQLite database for job persistence
|
||||
max_concurrent_jobs: How many *queued* jobs may run at once. Jobs
|
||||
started directly with start_job() run immediately regardless
|
||||
(but count toward the limit, so the queue waits for them).
|
||||
upload_dir: Where extracted last frames are saved.
|
||||
"""
|
||||
self.output_dir = output_dir
|
||||
self.log_dir = log_dir
|
||||
self.verbose = verbose
|
||||
self._db = database
|
||||
self.max_concurrent_jobs = max(max_concurrent_jobs, 1)
|
||||
self.upload_dir = upload_dir or os.path.join(output_dir, "uploads")
|
||||
|
||||
self._jobs: dict[str, Job] = {}
|
||||
self._jobs_lock = threading.Lock()
|
||||
# Serializes queue decisions (enqueue / dequeue / dispatch) so two
|
||||
# finishing jobs can't both take the same free slot.
|
||||
self._queue_lock = threading.RLock()
|
||||
self._last_queued_at = 0.0
|
||||
self._load_jobs()
|
||||
|
||||
# Cache of loaded generators keyed by model config so that we only pay
|
||||
# the model-loading cost once per model configuration.
|
||||
self._generators: dict[tuple, Any] = {}
|
||||
# Cached generators are keyed by (owning thread, config). The model's GPU
|
||||
# worker processes are killed (PR_SET_PDEATHSIG) when the *thread* that
|
||||
# started them exits, so a generator is only safe for that thread to reuse
|
||||
# -- which is why inference jobs run on long-lived worker threads (below)
|
||||
# rather than a new thread each.
|
||||
self._generator_threads: dict[tuple, threading.Thread] = {}
|
||||
# Inference jobs are handed to these threads through one shared queue.
|
||||
self._work: queue.Queue[Job | None] = queue.Queue()
|
||||
self._workers: list[threading.Thread] = []
|
||||
self._generators_lock = threading.Lock()
|
||||
|
||||
# Shared Manager for log queues (avoids spawning a new process per job)
|
||||
@@ -351,6 +384,14 @@ class JobRunner:
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
os.makedirs(self.log_dir, exist_ok=True)
|
||||
|
||||
for i in range(self.max_concurrent_jobs):
|
||||
worker = threading.Thread(target=self._worker_loop, name=f"inference-worker-{i}", daemon=True)
|
||||
worker.start()
|
||||
self._workers.append(worker)
|
||||
|
||||
# Jobs that were queued when the server last stopped carry on.
|
||||
self._dispatch()
|
||||
|
||||
def _load_logs(self, job: Job) -> None:
|
||||
"""Populate job's log buffer from its log file if it exists."""
|
||||
path = job.log_file_path
|
||||
@@ -416,6 +457,7 @@ class JobRunner:
|
||||
fake_score_model_path=row.get("fake_score_model_path", "") or "",
|
||||
status=JobStatus(status),
|
||||
created_at=row["created_at"],
|
||||
queued_at=row.get("queued_at"),
|
||||
started_at=row.get("started_at"),
|
||||
finished_at=row.get("finished_at"),
|
||||
error=row.get("error"),
|
||||
@@ -456,6 +498,7 @@ class JobRunner:
|
||||
self._db.update_job(
|
||||
job.id, {
|
||||
"status": job.status.value,
|
||||
"queued_at": job.queued_at,
|
||||
"started_at": job.started_at,
|
||||
"finished_at": job.finished_at,
|
||||
"error": job.error,
|
||||
@@ -466,7 +509,9 @@ class JobRunner:
|
||||
logger.warning("Failed to persist job %s: %s", job.id, exc)
|
||||
|
||||
def _shutdown(self) -> None:
|
||||
"""Shutdown the shared multiprocessing manager on exit."""
|
||||
"""Stop the inference workers and the shared multiprocessing manager on exit."""
|
||||
for _ in self._workers:
|
||||
self._work.put(None)
|
||||
try:
|
||||
self._mp_manager.shutdown()
|
||||
except Exception as exc:
|
||||
@@ -610,6 +655,9 @@ class JobRunner:
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to delete job %s from database: %s", job_id, exc)
|
||||
logger.info("Deleted job %s", job.id)
|
||||
# Anything queued behind this job's last frame can no longer run, and a
|
||||
# freed queue slot may let the next job start.
|
||||
self._dispatch()
|
||||
return True
|
||||
|
||||
CONFIG_FIELDS: tuple[str, ...] = (
|
||||
@@ -688,24 +736,30 @@ class JobRunner:
|
||||
for field_name, value in updates.items():
|
||||
setattr(job, field_name, value)
|
||||
self._save_job(job)
|
||||
try:
|
||||
self._db.update_job_config(job.id, updates)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to persist edits to job %s: %s", job.id, exc)
|
||||
return job
|
||||
|
||||
def start_job(self, job_id: str) -> Job:
|
||||
"""Start (or restart) a pending / stopped / failed job.
|
||||
|
||||
Raises:
|
||||
ValueError: If job not found or cannot be started
|
||||
"""
|
||||
#: Statuses a job can be started or queued from.
|
||||
STARTABLE_STATUSES = (JobStatus.PENDING, JobStatus.FAILED, JobStatus.STOPPED)
|
||||
|
||||
def _startable_job(self, job_id: str) -> Job:
|
||||
with self._jobs_lock:
|
||||
job = self._jobs.get(job_id)
|
||||
if job is None:
|
||||
raise ValueError(f"Job {job_id} not found")
|
||||
if job.status == JobStatus.RUNNING:
|
||||
raise ValueError("Job is already running")
|
||||
if job.status == JobStatus.QUEUED:
|
||||
raise ValueError("Job is already queued")
|
||||
if job.status == JobStatus.COMPLETED:
|
||||
raise ValueError("Job already completed. Delete and re-create to run again.")
|
||||
# Reset state for re-run
|
||||
job.status = JobStatus.PENDING
|
||||
return job
|
||||
|
||||
def _reset_for_run(self, job: Job) -> None:
|
||||
"""Clear the previous attempt's state so the job starts clean."""
|
||||
job.error = None
|
||||
job.output_path = None
|
||||
job.log_file_path = None # Will be set when job starts
|
||||
@@ -716,26 +770,189 @@ class JobRunner:
|
||||
job.log_file_handler = None # Reset file handler
|
||||
job._process = None # Reset subprocess handle
|
||||
|
||||
# Wrap _run_job in an additional safety layer to catch any exceptions
|
||||
# that might escape (though they shouldn't with our comprehensive handling)
|
||||
def safe_run_job(job: Job):
|
||||
"""Wrapper to ensure _run_job never raises an unhandled exception."""
|
||||
try:
|
||||
self._run_job(job)
|
||||
except BaseException as exc:
|
||||
# This should never happen, but if it does, we catch it here
|
||||
logger.critical("Unhandled exception escaped from _run_job for job %s: %s", job.id, exc, exc_info=True)
|
||||
with contextlib.suppress(Exception):
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = (f"Unhandled exception: {type(exc).__name__}: {str(exc)}")
|
||||
job.finished_at = time.time()
|
||||
def _launch(self, job: Job) -> None:
|
||||
"""Start ``job``; when it ends, look for more queued work.
|
||||
|
||||
thread = threading.Thread(target=safe_run_job, args=(job, ), daemon=True)
|
||||
Inference jobs go to a long-lived worker thread so the loaded model
|
||||
survives from one job to the next. Training jobs only babysit a
|
||||
subprocess, so they get a thread of their own.
|
||||
"""
|
||||
# Marked running now, not when the job gets going, so the queue counts it
|
||||
# immediately and can't start one job too many.
|
||||
job.status = JobStatus.RUNNING
|
||||
job.started_at = time.time()
|
||||
self._save_job(job)
|
||||
|
||||
if job.job_type == "inference":
|
||||
self._work.put(job)
|
||||
logger.info("Sent job %s to an inference worker", job.id)
|
||||
return
|
||||
|
||||
def run_training_job(job: Job):
|
||||
self._run_job_safely(job)
|
||||
self._dispatch_after_exit(threading.current_thread())
|
||||
|
||||
thread = threading.Thread(target=run_training_job, args=(job, ), daemon=True)
|
||||
job._thread = thread
|
||||
thread.start()
|
||||
logger.info("Started job %s", job.id)
|
||||
|
||||
def _run_job_safely(self, job: Job) -> None:
|
||||
"""Run a job, turning anything that escapes into a failed job."""
|
||||
try:
|
||||
self._run_job(job)
|
||||
except BaseException as exc:
|
||||
# This should never happen, but if it does, we catch it here
|
||||
logger.critical("Unhandled exception escaped from _run_job for job %s: %s", job.id, exc, exc_info=True)
|
||||
with contextlib.suppress(Exception):
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = (f"Unhandled exception: {type(exc).__name__}: {str(exc)}")
|
||||
job.finished_at = time.time()
|
||||
self._save_job(job)
|
||||
|
||||
def _worker_loop(self) -> None:
|
||||
"""Run inference jobs one after another on this thread, for the life of the server.
|
||||
|
||||
The thread never exits between jobs. That matters: it started the model's
|
||||
GPU processes, which are killed when it exits, so staying alive is what
|
||||
lets the next job find the model already loaded.
|
||||
"""
|
||||
while True:
|
||||
job = self._work.get()
|
||||
if job is None:
|
||||
return
|
||||
self._run_job_safely(job)
|
||||
# The slot is free; dependents may now be runnable (or impossible).
|
||||
with contextlib.suppress(Exception):
|
||||
self._dispatch()
|
||||
|
||||
def start_job(self, job_id: str) -> Job:
|
||||
"""Start (or restart) a pending / stopped / failed job right now.
|
||||
|
||||
Refused if every inference worker is already busy; use enqueue_job() to
|
||||
wait for a free one. A job that takes another job's last frame can only
|
||||
start once that job has completed.
|
||||
|
||||
Raises:
|
||||
ValueError: If job not found or cannot be started
|
||||
"""
|
||||
with self._queue_lock:
|
||||
job = self._startable_job(job_id)
|
||||
# Fail here, before touching the job, rather than as a failed run.
|
||||
self._check_dependencies_finished(job)
|
||||
if job.job_type == "inference":
|
||||
busy = self._running_inference_jobs()
|
||||
if busy >= self.max_concurrent_jobs:
|
||||
raise ValueError(f"{busy} inference job(s) already running and the model workers are busy. "
|
||||
"Queue this job to run next.")
|
||||
self._reset_for_run(job)
|
||||
job.queued_at = None
|
||||
self._launch(job)
|
||||
return job
|
||||
|
||||
def _running_inference_jobs(self) -> int:
|
||||
with self._jobs_lock:
|
||||
return sum(1 for j in self._jobs.values() if j.job_type == "inference" and j.status == JobStatus.RUNNING)
|
||||
|
||||
def _check_dependencies_finished(self, job: Job) -> None:
|
||||
for dep_id in dependencies(job.references):
|
||||
dep = self.get_job(dep_id)
|
||||
if dep is None:
|
||||
raise ValueError("This job starts from the last frame of a job that no longer exists.")
|
||||
if dep.status != JobStatus.COMPLETED:
|
||||
label = f"'{dep.name}'" if dep.name else dep.id
|
||||
raise ValueError(f"This job starts from the last frame of {label}, which hasn't finished "
|
||||
f"(it is {dep.status.value}). Queue it to run after that job.")
|
||||
|
||||
# -- queue ---------------------------------------------------------------
|
||||
|
||||
def _next_queued_at(self) -> float:
|
||||
"""A strictly increasing timestamp, so jobs queued together keep their order."""
|
||||
self._last_queued_at = max(time.time(), self._last_queued_at + 1e-6)
|
||||
return self._last_queued_at
|
||||
|
||||
def enqueue_jobs(self, job_ids: list[str]) -> list[Job]:
|
||||
"""Queue jobs, in the given order, to run as slots free up.
|
||||
|
||||
All-or-nothing: if any job can't be queued, none are.
|
||||
|
||||
Raises:
|
||||
ValueError: If a job is missing, not startable, or waits on itself.
|
||||
"""
|
||||
if len(set(job_ids)) != len(job_ids):
|
||||
raise ValueError("A job was listed more than once.")
|
||||
with self._queue_lock:
|
||||
jobs = [self._startable_job(job_id) for job_id in job_ids]
|
||||
with self._jobs_lock:
|
||||
snapshot = dict(self._jobs)
|
||||
for job in jobs:
|
||||
loop = circular_dependency(snapshot, job.id)
|
||||
if loop:
|
||||
raise ValueError("Jobs would wait on each other's last frames in a loop: " + " -> ".join(loop))
|
||||
for job in jobs:
|
||||
self._reset_for_run(job)
|
||||
job.status = JobStatus.QUEUED
|
||||
job.queued_at = self._next_queued_at()
|
||||
self._save_job(job)
|
||||
logger.info("Queued job %s", job.id)
|
||||
self._dispatch()
|
||||
return jobs
|
||||
|
||||
def enqueue_job(self, job_id: str) -> Job:
|
||||
return self.enqueue_jobs([job_id])[0]
|
||||
|
||||
def dequeue_job(self, job_id: str) -> Job:
|
||||
"""Take a queued job out of the queue, back to pending.
|
||||
|
||||
Raises:
|
||||
ValueError: If the job is missing or isn't queued.
|
||||
"""
|
||||
with self._queue_lock:
|
||||
job = self.get_job(job_id)
|
||||
if job is None:
|
||||
raise ValueError(f"Job {job_id} not found")
|
||||
if job.status != JobStatus.QUEUED:
|
||||
raise ValueError(f"Job is not queued (status={job.status.value})")
|
||||
job.status = JobStatus.PENDING
|
||||
job.queued_at = None
|
||||
self._save_job(job)
|
||||
logger.info("Dequeued job %s", job.id)
|
||||
# Anything queued behind this one's frame keeps waiting; a slot it held
|
||||
# in line may now go to another job.
|
||||
self._dispatch()
|
||||
return job
|
||||
|
||||
def _dispatch_after_exit(self, finished: threading.Thread) -> None:
|
||||
"""Wait for a job's thread to end, then start whatever is next in the queue."""
|
||||
finished.join()
|
||||
with contextlib.suppress(Exception):
|
||||
self._dispatch()
|
||||
|
||||
def _dispatch(self) -> None:
|
||||
"""Launch queued jobs that are ready, and fail those that can never be."""
|
||||
with self._queue_lock:
|
||||
with self._jobs_lock:
|
||||
snapshot = dict(self._jobs)
|
||||
plan = plan_dispatch(snapshot, self.max_concurrent_jobs)
|
||||
for job_id, reason in plan.fail.items():
|
||||
job = snapshot[job_id]
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = reason
|
||||
job.finished_at = time.time()
|
||||
self._save_job(job)
|
||||
logger.warning("Queued job %s can't run: %s", job_id, reason)
|
||||
for job_id in plan.start:
|
||||
job = snapshot[job_id]
|
||||
job.queued_at = None
|
||||
try:
|
||||
self._launch(job)
|
||||
except Exception as exc: # e.g. the OS refused a thread
|
||||
logger.exception("Could not launch queued job %s", job_id)
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = f"Could not start: {exc}"
|
||||
job.finished_at = time.time()
|
||||
self._save_job(job)
|
||||
|
||||
def stop_job(self, job_id: str) -> Job:
|
||||
"""Request a running job to stop.
|
||||
|
||||
@@ -749,6 +966,8 @@ class JobRunner:
|
||||
job = self._jobs.get(job_id)
|
||||
if job is None:
|
||||
raise ValueError(f"Job {job_id} not found")
|
||||
if job.status == JobStatus.QUEUED:
|
||||
return self.dequeue_job(job_id)
|
||||
if job.status != JobStatus.RUNNING:
|
||||
raise ValueError(f"Job is not running (status={job.status.value})")
|
||||
|
||||
@@ -786,6 +1005,22 @@ class JobRunner:
|
||||
"phase": job._log_buf.phase,
|
||||
}
|
||||
|
||||
def _discard_generators(self) -> None:
|
||||
"""Unload this thread's models so the next job on it starts from a clean load."""
|
||||
with self._generators_lock:
|
||||
for key in [k for k in self._generators if k[0] == threading.current_thread().ident]:
|
||||
self._drop_generator(key)
|
||||
|
||||
def _drop_generator(self, key: tuple) -> None:
|
||||
"""Forget a cached generator and shut it down. Call with ``_generators_lock`` held."""
|
||||
gen = self._generators.pop(key, None)
|
||||
self._generator_threads.pop(key, None)
|
||||
if gen is not None:
|
||||
try:
|
||||
gen.shutdown()
|
||||
except Exception:
|
||||
logger.debug("Shutdown of a cached generator failed", exc_info=True)
|
||||
|
||||
def _get_or_create_generator(
|
||||
self,
|
||||
model_id: str,
|
||||
@@ -823,23 +1058,27 @@ class JobRunner:
|
||||
sp_size,
|
||||
)
|
||||
|
||||
# Generators are cached by model_id and configuration parameters
|
||||
# Generators are cached per thread and configuration. Only the thread that
|
||||
# built one may reuse it: its GPU processes die with that thread.
|
||||
thread = threading.current_thread()
|
||||
entry_key = (thread.ident, cache_key)
|
||||
with self._generators_lock:
|
||||
cached = self._generators.get(cache_key)
|
||||
# Generators whose thread has exited have dead workers already.
|
||||
for key, owner in list(self._generator_threads.items()):
|
||||
if not owner.is_alive():
|
||||
self._drop_generator(key)
|
||||
cached = self._generators.get(entry_key)
|
||||
if cached is not None:
|
||||
if _generator_is_alive(cached):
|
||||
return cached
|
||||
# Workers can exit while a generator sits idle in the cache;
|
||||
# reusing it fails every later job with the same config.
|
||||
logger.warning(
|
||||
"Cached generator for %s has dead workers; reloading.",
|
||||
model_id,
|
||||
)
|
||||
self._generators.pop(cache_key, None)
|
||||
try:
|
||||
cached.shutdown()
|
||||
except Exception:
|
||||
logger.debug("Shutdown of the dead generator failed", exc_info=True)
|
||||
# Workers can exit while a generator sits idle in the cache.
|
||||
logger.warning("Cached generator for %s has dead workers; reloading.", model_id)
|
||||
self._drop_generator(entry_key)
|
||||
# A worker keeps one model resident. A different configuration replaces
|
||||
# it, freeing its GPU memory first, rather than loading beside it.
|
||||
for key in [k for k in self._generators if k[0] == thread.ident]:
|
||||
logger.info("Unloading the previous model to load %s.", model_id)
|
||||
self._drop_generator(key)
|
||||
|
||||
# Import lazily so starting the server is fast even without a GPU.
|
||||
from fastvideo import VideoGenerator
|
||||
@@ -883,11 +1122,8 @@ class JobRunner:
|
||||
)
|
||||
|
||||
with self._generators_lock:
|
||||
if cache_key not in self._generators:
|
||||
self._generators[cache_key] = gen
|
||||
else: # Another thread may have created it while we were loading.
|
||||
gen.shutdown()
|
||||
gen = self._generators[cache_key]
|
||||
self._generators[entry_key] = gen
|
||||
self._generator_threads[entry_key] = thread
|
||||
return gen
|
||||
|
||||
def _run_job(self, job: Job):
|
||||
@@ -1064,6 +1300,14 @@ class JobRunner:
|
||||
logger.warning("Job stopped before execution started")
|
||||
return
|
||||
|
||||
# Frames borrowed from other clips are extracted now, when they exist.
|
||||
# Done before loading the model so a missing frame fails fast.
|
||||
resolved_references = resolve_references(
|
||||
job.references,
|
||||
self.get_job,
|
||||
lambda source: last_frame_for_job(source, self.upload_dir),
|
||||
)
|
||||
|
||||
buf.phase = "loading model"
|
||||
logger.info("Loading model...")
|
||||
|
||||
@@ -1121,7 +1365,7 @@ class JobRunner:
|
||||
if job.image_path:
|
||||
gen_kwargs["image_path"] = job.image_path
|
||||
if job.references:
|
||||
gen_kwargs["references"] = _build_h3_references(job.references)
|
||||
gen_kwargs["references"] = _build_h3_references(resolved_references)
|
||||
if job.last_image_path:
|
||||
# _prepare_fl2va requires a PIL image, not a path.
|
||||
from PIL import Image as _PILImage
|
||||
@@ -1148,6 +1392,7 @@ class JobRunner:
|
||||
if job._stop_event.is_set():
|
||||
job.status = JobStatus.STOPPED
|
||||
logger.warning("Job was stopped during execution")
|
||||
self._discard_generators()
|
||||
else:
|
||||
job.status = JobStatus.COMPLETED
|
||||
buf.progress = 100.0
|
||||
@@ -1156,9 +1401,19 @@ class JobRunner:
|
||||
self._save_job(job)
|
||||
buf.phase = "done"
|
||||
|
||||
except UnresolvedReference as exception:
|
||||
logger.error("Job %s can't get its reference frame: %s", job.id, exception)
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = str(exception)
|
||||
job.finished_at = time.time()
|
||||
self._save_job(job)
|
||||
buf.phase = "failed"
|
||||
|
||||
except Exception as exception:
|
||||
error_msg = str(exception)
|
||||
logger.exception("Critical error in job thread: %s", error_msg)
|
||||
# The model may be mid-request. Reload rather than trust it for the next job.
|
||||
self._discard_generators()
|
||||
job.status = JobStatus.FAILED
|
||||
job.error = f"Critical error ({type(exception).__name__}): {error_msg}"
|
||||
job.finished_at = time.time()
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Join a scene's clips, in order, into one video.
|
||||
|
||||
Kept apart from ``server.py`` because it needs nothing from FastVideo -- only
|
||||
PyAV (to read each clip's format) and an ``ffmpeg`` binary -- so it can be
|
||||
exercised without a GPU.
|
||||
|
||||
Clips from one scene share codec, size and frame rate, so they are joined by
|
||||
copying the streams rather than re-encoding: it takes about a second, loses no
|
||||
quality, and keeps picture and sound in step. Each clip is cut off at the end of
|
||||
its last video frame, because the audio track runs a few milliseconds past it and
|
||||
would otherwise open a gap between clips.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
VIDEO_EXTENSIONS = {".mp4", ".mov", ".mkv", ".webm", ".avi"}
|
||||
MERGED_DIRNAME = "merged"
|
||||
_FILENAME_RE = re.compile(r"^[\w.-]+\.mp4$")
|
||||
|
||||
|
||||
class MergeError(Exception):
|
||||
"""A merge that can't be done, carrying the HTTP status to report."""
|
||||
|
||||
def __init__(self, status_code: int, detail: str):
|
||||
super().__init__(detail)
|
||||
self.status_code = status_code
|
||||
self.detail = detail
|
||||
|
||||
|
||||
def find_ffmpeg() -> str | None:
|
||||
"""An ``ffmpeg`` to run: $FASTVIDEO_FFMPEG_BIN, then PATH, then imageio-ffmpeg's bundled one."""
|
||||
configured = os.getenv("FASTVIDEO_FFMPEG_BIN")
|
||||
if configured and (shutil.which(configured) or os.path.isfile(configured)):
|
||||
return shutil.which(configured) or configured
|
||||
on_path = shutil.which("ffmpeg")
|
||||
if on_path:
|
||||
return on_path
|
||||
try:
|
||||
import imageio_ffmpeg
|
||||
|
||||
return imageio_ffmpeg.get_ffmpeg_exe()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClipInfo:
|
||||
path: str
|
||||
video_codec: str
|
||||
width: int
|
||||
height: int
|
||||
fps: float
|
||||
pix_fmt: str
|
||||
video_seconds: float
|
||||
audio: tuple[str, int, int] | None # codec, sample rate, channels
|
||||
|
||||
def describe(self) -> str:
|
||||
audio = f"{self.audio[0]} {self.audio[1]} Hz x{self.audio[2]}" if self.audio else "no audio"
|
||||
return f"{self.width}x{self.height} {self.fps:.2f} fps {self.video_codec}/{self.pix_fmt}, {audio}"
|
||||
|
||||
def signature(self) -> tuple[Any, ...]:
|
||||
return (self.video_codec, self.width, self.height, round(self.fps, 3), self.pix_fmt, self.audio)
|
||||
|
||||
|
||||
def probe_clip(path: str) -> ClipInfo:
|
||||
"""Read what ``check_compatible`` needs from a video file."""
|
||||
import av
|
||||
|
||||
try:
|
||||
with av.open(path) as container:
|
||||
video = next((s for s in container.streams if s.type == "video"), None)
|
||||
if video is None:
|
||||
raise MergeError(409, f"{os.path.basename(path)} has no video stream.")
|
||||
audio = next((s for s in container.streams if s.type == "audio"), None)
|
||||
fps = float(video.average_rate or video.guessed_rate or 0)
|
||||
if video.duration is not None:
|
||||
seconds = float(video.duration * video.time_base)
|
||||
elif container.duration is not None:
|
||||
seconds = container.duration / 1_000_000
|
||||
else:
|
||||
seconds = (video.frames or 0) / fps if fps else 0.0
|
||||
if seconds <= 0:
|
||||
raise MergeError(409, f"{os.path.basename(path)} has no playable video.")
|
||||
return ClipInfo(
|
||||
path=path,
|
||||
video_codec=video.codec_context.name,
|
||||
width=video.codec_context.width,
|
||||
height=video.codec_context.height,
|
||||
fps=fps,
|
||||
pix_fmt=video.codec_context.pix_fmt or "",
|
||||
video_seconds=seconds,
|
||||
audio=((audio.codec_context.name, audio.codec_context.sample_rate, audio.codec_context.channels)
|
||||
if audio else None),
|
||||
)
|
||||
except MergeError:
|
||||
raise
|
||||
except Exception as exc: # av.error.* on a corrupt or unreadable file
|
||||
raise MergeError(409, f"Could not read {os.path.basename(path)}: {exc}") from exc
|
||||
|
||||
|
||||
def check_compatible(infos: list[ClipInfo], labels: list[str]) -> None:
|
||||
"""Refuse clips that can't be joined by copying, naming the first that differs."""
|
||||
first = infos[0]
|
||||
for info, label in zip(infos[1:], labels[1:], strict=True):
|
||||
if info.signature() != first.signature():
|
||||
raise MergeError(
|
||||
409,
|
||||
f"{label} ({info.describe()}) doesn't match {labels[0]} ({first.describe()}). "
|
||||
"Clips need the same size, frame rate and audio to be joined.",
|
||||
)
|
||||
|
||||
|
||||
def _concat_list(infos: list[ClipInfo]) -> str:
|
||||
lines = ["ffconcat version 1.0"]
|
||||
for info in infos:
|
||||
escaped = info.path.replace("'", "'\\''")
|
||||
lines += [f"file '{escaped}'", f"outpoint {info.video_seconds:.6f}"]
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def merge_clips(paths: list[str], out_path: str, labels: list[str] | None = None, ffmpeg: str | None = None) -> float:
|
||||
"""Join ``paths`` in order into ``out_path``. Returns the merged length in seconds."""
|
||||
if not paths:
|
||||
raise MergeError(400, "No clips to merge.")
|
||||
labels = labels or [os.path.basename(p) for p in paths]
|
||||
ffmpeg = ffmpeg or find_ffmpeg()
|
||||
if not ffmpeg:
|
||||
raise MergeError(503, "ffmpeg isn't available on the server. Install it or set FASTVIDEO_FFMPEG_BIN.")
|
||||
|
||||
infos = [probe_clip(p) for p in paths]
|
||||
check_compatible(infos, labels)
|
||||
|
||||
os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
|
||||
tmp_out = f"{out_path}.tmp.mp4"
|
||||
list_fd, list_path = tempfile.mkstemp(suffix=".ffconcat", text=True)
|
||||
try:
|
||||
with os.fdopen(list_fd, "w", encoding="utf-8") as f:
|
||||
f.write(_concat_list(infos))
|
||||
result = subprocess.run(
|
||||
[ffmpeg, "-y", "-loglevel", "error", "-f", "concat", "-safe", "0", "-i", list_path, "-c", "copy",
|
||||
"-movflags", "+faststart", tmp_out],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=600,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
tail = " ".join(result.stderr.strip().splitlines()[-3:])
|
||||
raise MergeError(500, f"ffmpeg failed: {tail or 'unknown error'}")
|
||||
os.replace(tmp_out, out_path) # atomic, so a reader never sees a half-written file
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
raise MergeError(500, "ffmpeg took too long.") from exc
|
||||
finally:
|
||||
for leftover in (list_path, tmp_out):
|
||||
try:
|
||||
os.remove(leftover)
|
||||
except OSError:
|
||||
pass
|
||||
return sum(i.video_seconds for i in infos)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MergeResult:
|
||||
filename: str
|
||||
path: str
|
||||
clips: int
|
||||
seconds: float
|
||||
|
||||
|
||||
def _safe_stem(name: str) -> str:
|
||||
return re.sub(r"[^\w.-]+", "-", name).strip("-.")[:60] or "scene"
|
||||
|
||||
|
||||
def _label(job: Any) -> str:
|
||||
name = (getattr(job, "name", "") or "").strip()
|
||||
return f"'{name}'" if name else str(job.id)
|
||||
|
||||
|
||||
def merge_jobs(jobs: list[Any], out_dir: str, name: str = "") -> MergeResult:
|
||||
"""Join the finished videos of ``jobs`` (in the order given) into ``<out_dir>/merged``.
|
||||
|
||||
Every job must be completed with a video on disk; otherwise nothing is made,
|
||||
so a scene missing a clip can't be mistaken for the whole thing.
|
||||
"""
|
||||
if not jobs:
|
||||
raise MergeError(400, "No clips to merge.")
|
||||
paths: list[str] = []
|
||||
for job in jobs:
|
||||
status = getattr(job.status, "value", job.status)
|
||||
if status != "completed" or not job.output_path:
|
||||
raise MergeError(409, f"{_label(job)} hasn't finished ({status}), so the scene can't be merged yet.")
|
||||
if not os.path.isfile(job.output_path):
|
||||
raise MergeError(409, f"{_label(job)}'s video is missing from disk.")
|
||||
if Path(job.output_path).suffix.lower() not in VIDEO_EXTENSIONS:
|
||||
raise MergeError(409, f"{_label(job)}'s output isn't a video.")
|
||||
paths.append(job.output_path)
|
||||
|
||||
filename = f"{_safe_stem(name)}-{time.strftime('%Y%m%d-%H%M%S')}.mp4"
|
||||
merged_dir = os.path.join(out_dir, MERGED_DIRNAME)
|
||||
out_path = os.path.join(merged_dir, filename)
|
||||
seconds = merge_clips(paths, out_path, labels=[_label(j) for j in jobs])
|
||||
return MergeResult(filename=filename, path=out_path, clips=len(paths), seconds=seconds)
|
||||
|
||||
|
||||
def merged_path(out_dir: str, filename: str) -> str:
|
||||
"""Path of a merged video by filename, refusing anything that isn't one."""
|
||||
if not _FILENAME_RE.match(filename):
|
||||
raise MergeError(404, "Merged video not found.")
|
||||
path = os.path.join(out_dir, MERGED_DIRNAME, filename)
|
||||
if not os.path.isfile(path):
|
||||
raise MergeError(404, "Merged video not found.")
|
||||
return path
|
||||
@@ -40,6 +40,7 @@ import threading
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Annotated, Any
|
||||
|
||||
from fastapi import FastAPI, File, HTTPException, UploadFile
|
||||
@@ -47,7 +48,10 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, PlainTextResponse
|
||||
|
||||
from fastvideo_studio.database import Database, default_settings_dict
|
||||
from fastvideo_studio.models import (CreateDatasetRequest, CreateJobRequest, SettingsUpdate, UpdateCaptionRequest,
|
||||
from fastvideo_studio.job_queue import circular_dependency, dependencies, plan_dispatch
|
||||
from fastvideo_studio.merge import MergeError, merge_jobs, merged_path
|
||||
from fastvideo_studio.models import (CreateDatasetRequest, CreateJobRequest, MergeSceneRequest, QueueJobsRequest,
|
||||
SettingsUpdate, UpdateCaptionRequest,
|
||||
model_label)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
|
||||
@@ -268,6 +272,7 @@ def _new_job_dict(req: CreateJobRequest) -> dict[str, Any]:
|
||||
"data_path": req.data_path or "",
|
||||
"status": "pending",
|
||||
"created_at": time.time(),
|
||||
"queued_at": None,
|
||||
"started_at": None,
|
||||
"finished_at": None,
|
||||
"error": None,
|
||||
@@ -337,8 +342,8 @@ def _compute_logs(job: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return JobLogs-shaped data, growing the visible lines as time passes."""
|
||||
seq = _log_sequence(job)
|
||||
status = job["status"]
|
||||
if status == "pending":
|
||||
return {"lines": [], "progress": 0.0, "progress_msg": "", "phase": "pending"}
|
||||
if status in ("pending", "queued"):
|
||||
return {"lines": [], "progress": 0.0, "progress_msg": "", "phase": status}
|
||||
if status == "completed":
|
||||
return {"lines": seq, "progress": 100.0, "progress_msg": "50/50 steps", "phase": "done"}
|
||||
if status in ("failed", "stopped"):
|
||||
@@ -461,6 +466,11 @@ def _job_from_snapshot(row: dict[str, Any]) -> dict[str, Any]:
|
||||
status = "failed"
|
||||
job["error"] = job.get("error") or "Server restarted (job was running)"
|
||||
job["finished_at"] = job.get("finished_at") or time.time()
|
||||
elif status == "queued":
|
||||
# The real server resumes its queue on restart; the mock never runs
|
||||
# snapshot jobs, so leave it as something the user can queue again.
|
||||
status = "pending"
|
||||
job["queued_at"] = None
|
||||
job["status"] = status
|
||||
job["references"] = _decode_references(row.get("references"))
|
||||
job["progress"] = 100.0 if status == "completed" else 0.0
|
||||
@@ -633,6 +643,7 @@ def list_jobs(job_type: str | None = None) -> list[dict[str, Any]]:
|
||||
# _public_job mutates the shared job dict via _advance_job, so it must run
|
||||
# under the lock (matching every other job route) to avoid racing start/stop.
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
jobs = list(_jobs.values())
|
||||
if job_type:
|
||||
jobs = [j for j in jobs if j.get("job_type") == job_type]
|
||||
@@ -643,6 +654,7 @@ def list_jobs(job_type: str | None = None) -> list[dict[str, Any]]:
|
||||
@app.get("/api/jobs/{job_id}")
|
||||
def get_job(job_id: str) -> dict[str, Any]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
@@ -653,18 +665,20 @@ def get_job(job_id: str) -> dict[str, Any]:
|
||||
def create_job(req: CreateJobRequest) -> dict[str, Any]:
|
||||
job = _new_job_dict(req)
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
_jobs[job["id"]] = job
|
||||
if _settings.get("autoStartJob"):
|
||||
_start(job)
|
||||
return _public_job(job)
|
||||
|
||||
|
||||
def _start(job: dict[str, Any]) -> None:
|
||||
def _start(job: dict[str, Any], at: float | None = None) -> None:
|
||||
# Once the mock is actually driving this job, it's no longer a frozen
|
||||
# historical record -- let _advance_job simulate it like any other.
|
||||
job.pop(_FROM_SNAPSHOT_KEY, None)
|
||||
job["status"] = "running"
|
||||
job["started_at"] = time.time()
|
||||
job["queued_at"] = None
|
||||
job["started_at"] = at if at is not None else time.time()
|
||||
job["finished_at"] = None
|
||||
job["error"] = None
|
||||
job["output_path"] = None
|
||||
@@ -676,17 +690,185 @@ def _start(job: dict[str, Any]) -> None:
|
||||
job["phase"] = "starting"
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/start")
|
||||
def start_job(job_id: str) -> dict[str, Any]:
|
||||
class _QueueView:
|
||||
"""The scheduler's view of a mock job dict (it reads attributes, not keys)."""
|
||||
|
||||
def __init__(self, job: dict[str, Any]):
|
||||
self.id = job["id"]
|
||||
self.name = job.get("name", "")
|
||||
self.status = job["status"]
|
||||
self.references = job.get("references")
|
||||
self.queued_at = job.get("queued_at")
|
||||
|
||||
|
||||
def _drain_queue() -> None:
|
||||
"""Run the queue up to now. Call with ``_state_lock`` held.
|
||||
|
||||
Like everything else in the mock this is computed on read: each pass
|
||||
finishes jobs whose time is up, then starts whatever the queue allows.
|
||||
A job takes over the slot at the moment the previous job *finished* (not at
|
||||
the moment someone happened to poll), so a queue that has been sitting for
|
||||
a minute shows every clip already done rather than advancing one per poll.
|
||||
Same rules as the real runner: one at a time, oldest first, and a clip that
|
||||
takes another's last frame waits for it -- see job_queue.plan_dispatch.
|
||||
"""
|
||||
while True:
|
||||
for job in _jobs.values():
|
||||
_advance_job(job)
|
||||
plan = plan_dispatch({jid: _QueueView(j) for jid, j in _jobs.items()}, max_concurrent=1)
|
||||
now = time.time()
|
||||
for job_id, reason in plan.fail.items():
|
||||
job = _jobs[job_id]
|
||||
job["status"] = "failed"
|
||||
job["error"] = reason
|
||||
job["finished_at"] = now
|
||||
job["queued_at"] = None
|
||||
job["phase"] = "failed"
|
||||
if not plan.start:
|
||||
if plan.fail:
|
||||
continue # a failure can free a slot or fail more jobs
|
||||
return
|
||||
for job_id in plan.start:
|
||||
job = _jobs[job_id]
|
||||
finished = [j["finished_at"] for j in _jobs.values() if j.get("finished_at")]
|
||||
# The slot came free when the last job finished, if that was after the job was queued.
|
||||
slot_free = max([job.get("queued_at") or 0.0, *[f for f in finished if f <= now]])
|
||||
_start(job, at=min(now, slot_free))
|
||||
|
||||
|
||||
def _check_dependencies_finished(job: dict[str, Any]) -> None:
|
||||
for dep_id in dependencies(job.get("references")):
|
||||
dep = _jobs.get(dep_id)
|
||||
if dep is None:
|
||||
raise HTTPException(status_code=409,
|
||||
detail="This job starts from the last frame of a job that no longer exists.")
|
||||
if dep["status"] != "completed":
|
||||
label = f"'{dep['name']}'" if dep.get("name") else dep["id"]
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=(f"This job starts from the last frame of {label}, which hasn't finished "
|
||||
f"(it is {dep['status']}). Queue it to run after that job."),
|
||||
)
|
||||
|
||||
|
||||
# Mirrors JobRunner.EDITABLE_STATUSES: exactly the jobs that can still be started.
|
||||
_EDITABLE_STATUSES = ("pending", "failed", "stopped")
|
||||
|
||||
|
||||
@app.patch("/api/jobs/{job_id}")
|
||||
def update_job(job_id: str, updates: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Edit a job's configuration; only jobs that have not produced a result."""
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
_advance_job(job)
|
||||
if job["status"] == "running":
|
||||
raise HTTPException(status_code=409, detail="Job is already running")
|
||||
if job["status"] == "completed":
|
||||
raise HTTPException(status_code=409, detail="Job already completed. Delete and re-create to run again.")
|
||||
if job["status"] not in _EDITABLE_STATUSES:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(f"Job is {job['status']}; only {', '.join(_EDITABLE_STATUSES)} jobs can be edited. "
|
||||
"Duplicate it instead."),
|
||||
)
|
||||
unknown = set(updates) - set(CreateJobRequest.model_fields)
|
||||
if unknown:
|
||||
raise HTTPException(status_code=400, detail=f"Not editable: {', '.join(sorted(unknown))}")
|
||||
job.update(updates)
|
||||
return _public_job(job)
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/last-frame")
|
||||
def extract_last_frame(job_id: str) -> dict[str, str]:
|
||||
"""Stand-in for the real route: there is no video to read, so hand back a made-up path."""
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
_advance_job(job)
|
||||
if job["status"] != "completed" or not job.get("output_path"):
|
||||
raise HTTPException(status_code=404, detail="No output available for this job")
|
||||
return {"path": f"/mock/last_frames/last_frame_{job_id}.png", "media_type": "image"}
|
||||
|
||||
|
||||
def _check_startable(job: dict[str, Any]) -> None:
|
||||
if job["status"] == "running":
|
||||
raise HTTPException(status_code=409, detail="Job is already running")
|
||||
if job["status"] == "queued":
|
||||
raise HTTPException(status_code=409, detail="Job is already queued")
|
||||
if job["status"] == "completed":
|
||||
raise HTTPException(status_code=409, detail="Job already completed. Delete and re-create to run again.")
|
||||
|
||||
|
||||
def _enqueue(job_ids: list[str]) -> list[dict[str, Any]]:
|
||||
"""Queue jobs in order, all-or-nothing (mirrors JobRunner.enqueue_jobs). Lock held."""
|
||||
if len(set(job_ids)) != len(job_ids):
|
||||
raise HTTPException(status_code=409, detail="A job was listed more than once.")
|
||||
jobs = []
|
||||
for job_id in job_ids:
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
_check_startable(job)
|
||||
jobs.append(job)
|
||||
views = {jid: _QueueView(j) for jid, j in _jobs.items()}
|
||||
for job in jobs:
|
||||
loop = circular_dependency(views, job["id"])
|
||||
if loop:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Jobs would wait on each other's last frames in a loop: " + " -> ".join(loop),
|
||||
)
|
||||
for job in jobs:
|
||||
job.pop(_FROM_SNAPSHOT_KEY, None)
|
||||
job.update(status="queued", queued_at=time.time(), error=None, output_path=None, log_file_path=None,
|
||||
started_at=None, finished_at=None, progress=0.0, progress_msg="", phase="queued")
|
||||
# Strictly increasing, so jobs queued together keep their order.
|
||||
others = [j["queued_at"] for j in _jobs.values() if j.get("queued_at") and j is not job]
|
||||
job["queued_at"] = max([job["queued_at"], *[o + 1e-6 for o in others]])
|
||||
_drain_queue()
|
||||
return [_public_job(j) for j in jobs]
|
||||
|
||||
|
||||
@app.post("/api/jobs/queue")
|
||||
def queue_jobs(req: QueueJobsRequest) -> list[dict[str, Any]]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
return _enqueue(req.job_ids)
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/queue")
|
||||
def queue_job(job_id: str) -> dict[str, Any]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
return _enqueue([job_id])[0]
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/dequeue")
|
||||
def dequeue_job(job_id: str) -> dict[str, Any]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
if job["status"] != "queued":
|
||||
raise HTTPException(status_code=409, detail=f"Job is not queued (status={job['status']})")
|
||||
job["status"] = "pending"
|
||||
job["queued_at"] = None
|
||||
_drain_queue()
|
||||
return _public_job(job)
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/start")
|
||||
def start_job(job_id: str) -> dict[str, Any]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
_advance_job(job)
|
||||
_check_startable(job)
|
||||
_check_dependencies_finished(job)
|
||||
_start(job)
|
||||
return _public_job(job)
|
||||
|
||||
@@ -694,10 +876,16 @@ def start_job(job_id: str) -> dict[str, Any]:
|
||||
@app.post("/api/jobs/{job_id}/stop")
|
||||
def stop_job(job_id: str) -> dict[str, Any]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
_advance_job(job)
|
||||
if job["status"] == "queued":
|
||||
job["status"] = "pending"
|
||||
job["queued_at"] = None
|
||||
_drain_queue()
|
||||
return _public_job(job)
|
||||
if job["status"] != "running":
|
||||
raise HTTPException(status_code=409, detail=f"Job is not running (status={job['status']})")
|
||||
job["status"] = "stopped"
|
||||
@@ -709,14 +897,17 @@ def stop_job(job_id: str) -> dict[str, Any]:
|
||||
@app.delete("/api/jobs/{job_id}")
|
||||
def delete_job(job_id: str) -> dict[str, str]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
if _jobs.pop(job_id, None) is None:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
_drain_queue() # clips waiting on this one's last frame can no longer run
|
||||
return {"detail": f"Job {job_id} deleted"}
|
||||
|
||||
|
||||
@app.get("/api/jobs/{job_id}/logs")
|
||||
def get_job_logs(job_id: str, after: int = 0) -> dict[str, Any]:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
@@ -735,6 +926,7 @@ def get_job_logs(job_id: str, after: int = 0) -> dict[str, Any]:
|
||||
@app.get("/api/jobs/{job_id}/video")
|
||||
def get_video(job_id: str) -> FileResponse:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
@@ -755,9 +947,59 @@ def get_video(job_id: str) -> FileResponse:
|
||||
return FileResponse(_build_mock_media("mp4"), media_type="video/mp4", filename=f"job_{job_id}.mp4")
|
||||
|
||||
|
||||
_MERGED_DIR = tempfile.mkdtemp(prefix="fvstudio_merged_")
|
||||
|
||||
|
||||
@app.post("/api/scenes/merge")
|
||||
def merge_scene(req: MergeSceneRequest) -> dict[str, Any]:
|
||||
"""Real merge (same code as the real server) of whatever videos the jobs have.
|
||||
|
||||
Jobs whose video only exists in the mock get the stand-in clip, so the
|
||||
flow can be tried end to end without generating anything.
|
||||
"""
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
picked = []
|
||||
for job_id in req.job_ids:
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
picked.append(dict(job))
|
||||
all_real = all(j.get("output_path") and os.path.isfile(j["output_path"]) for j in picked)
|
||||
stand_in = None if all_real else _build_mock_media("mp4")
|
||||
jobs = [
|
||||
SimpleNamespace(
|
||||
id=j["id"],
|
||||
name=j.get("name", ""),
|
||||
status=j["status"],
|
||||
output_path=(j["output_path"] if all_real else (stand_in if j["status"] == "completed" else None)),
|
||||
) for j in picked
|
||||
]
|
||||
try:
|
||||
result = merge_jobs(jobs, _MERGED_DIR, req.name)
|
||||
except MergeError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.detail) from e
|
||||
return {
|
||||
"filename": result.filename,
|
||||
"url": f"/api/merged/{result.filename}",
|
||||
"clips": result.clips,
|
||||
"seconds": result.seconds,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/merged/{filename}")
|
||||
def get_merged_video(filename: str) -> FileResponse:
|
||||
try:
|
||||
path = merged_path(_MERGED_DIR, filename)
|
||||
except MergeError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.detail) from e
|
||||
return FileResponse(path, media_type="video/mp4", filename=filename)
|
||||
|
||||
|
||||
@app.get("/api/jobs/{job_id}/download_log")
|
||||
def download_log(job_id: str) -> PlainTextResponse:
|
||||
with _state_lock:
|
||||
_drain_queue()
|
||||
job = _jobs.get(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
|
||||
@@ -6,6 +6,8 @@ from fastvideo_studio.models.create_job_request import CreateJobRequest
|
||||
from fastvideo_studio.models.settings_update import SettingsUpdate
|
||||
from fastvideo_studio.models.create_dataset_request import CreateDatasetRequest
|
||||
from fastvideo_studio.models.update_caption_request import UpdateCaptionRequest
|
||||
from fastvideo_studio.models.queue_jobs_request import QueueJobsRequest
|
||||
from fastvideo_studio.models.merge_scene_request import MergeSceneRequest
|
||||
|
||||
|
||||
def model_label(model_path: str) -> str:
|
||||
@@ -18,5 +20,7 @@ __all__ = [
|
||||
"SettingsUpdate",
|
||||
"CreateDatasetRequest",
|
||||
"UpdateCaptionRequest",
|
||||
"QueueJobsRequest",
|
||||
"MergeSceneRequest",
|
||||
"model_label",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Request model for merging a scene's clips into one video."""
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class MergeSceneRequest(BaseModel):
|
||||
#: Completed jobs, in the order their videos should play.
|
||||
job_ids: list[str]
|
||||
#: Used to name the merged file.
|
||||
name: str = ""
|
||||
@@ -0,0 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Request model for queueing jobs."""
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class QueueJobsRequest(BaseModel):
|
||||
job_ids: list[str]
|
||||
@@ -33,10 +33,12 @@ from fastapi.responses import FileResponse
|
||||
|
||||
from fastvideo.registry import (get_registered_model_paths, get_registered_models_with_workloads)
|
||||
from fastvideo_studio.database import Database, _get_db_path
|
||||
from fastvideo_studio.frames import FrameError, last_frame_for_job
|
||||
from fastvideo_studio.gpu import get_gpu_snapshot
|
||||
from fastvideo_studio.job_runner import JobRunner, JobStatus
|
||||
from fastvideo_studio.models import (CreateDatasetRequest, CreateJobRequest, SettingsUpdate, UpdateCaptionRequest,
|
||||
model_label)
|
||||
from fastvideo_studio.merge import MergeError, merge_jobs, merged_path
|
||||
from fastvideo_studio.models import (CreateDatasetRequest, CreateJobRequest, MergeSceneRequest, QueueJobsRequest,
|
||||
SettingsUpdate, UpdateCaptionRequest, model_label)
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
@@ -439,6 +441,41 @@ def start_job(job_id: str) -> dict[str, Any]:
|
||||
) from e
|
||||
|
||||
|
||||
def _queue_error(e: ValueError) -> HTTPException:
|
||||
return HTTPException(status_code=404 if "not found" in str(e) else 409, detail=str(e))
|
||||
|
||||
|
||||
@app.post("/api/jobs/queue")
|
||||
def queue_jobs(req: QueueJobsRequest) -> list[dict[str, Any]]:
|
||||
"""Queue jobs, in the order given, to run as GPU slots free up.
|
||||
|
||||
A job that starts from another job's last frame waits for that job to
|
||||
complete. All-or-nothing: if any job can't be queued, none are.
|
||||
"""
|
||||
try:
|
||||
return [j.to_dict() for j in job_runner.enqueue_jobs(req.job_ids)]
|
||||
except ValueError as e:
|
||||
raise _queue_error(e) from e
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/queue")
|
||||
def queue_job(job_id: str) -> dict[str, Any]:
|
||||
"""Queue one job to run when a slot is free."""
|
||||
try:
|
||||
return job_runner.enqueue_job(job_id).to_dict()
|
||||
except ValueError as e:
|
||||
raise _queue_error(e) from e
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/dequeue")
|
||||
def dequeue_job(job_id: str) -> dict[str, Any]:
|
||||
"""Take a queued job out of the queue (back to pending)."""
|
||||
try:
|
||||
return job_runner.dequeue_job(job_id).to_dict()
|
||||
except ValueError as e:
|
||||
raise _queue_error(e) from e
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/stop")
|
||||
def stop_job(job_id: str) -> dict[str, Any]:
|
||||
"""Request a running job to stop.
|
||||
@@ -665,6 +702,62 @@ def get_video(job_id: str) -> FileResponse:
|
||||
return FileResponse(job.output_path, media_type=media_type)
|
||||
|
||||
|
||||
@app.post("/api/jobs/{job_id}/last-frame")
|
||||
def extract_last_frame(job_id: str) -> dict[str, str]:
|
||||
"""Save the last frame of a completed job's video and return its path.
|
||||
|
||||
The path is meant to be attached to another job as an image reference, so
|
||||
the next clip of a scene can start from where this one ended.
|
||||
"""
|
||||
job = job_runner.get_job(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
try:
|
||||
path = last_frame_for_job(job, upload_dir)
|
||||
except FrameError as e:
|
||||
if e.status_code >= 500:
|
||||
logger.warning("Last-frame extraction failed for job %s: %s", job_id, e.detail)
|
||||
raise HTTPException(status_code=e.status_code, detail=e.detail) from e
|
||||
return {"path": path, "media_type": "image"}
|
||||
|
||||
|
||||
@app.post("/api/scenes/merge")
|
||||
def merge_scene(req: MergeSceneRequest) -> dict[str, Any]:
|
||||
"""Join the finished clips of a scene, in the order given, into one video.
|
||||
|
||||
The clips are copied, not re-encoded, so this takes about a second.
|
||||
"""
|
||||
jobs = []
|
||||
for job_id in req.job_ids:
|
||||
job = job_runner.get_job(job_id)
|
||||
if job is None:
|
||||
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
||||
jobs.append(job)
|
||||
try:
|
||||
result = merge_jobs(jobs, job_runner.output_dir, req.name)
|
||||
except MergeError as e:
|
||||
if e.status_code >= 500:
|
||||
logger.warning("Scene merge failed: %s", e.detail)
|
||||
raise HTTPException(status_code=e.status_code, detail=e.detail) from e
|
||||
logger.info("Merged %d clips into %s", result.clips, result.path)
|
||||
return {
|
||||
"filename": result.filename,
|
||||
"url": f"/api/merged/{result.filename}",
|
||||
"clips": result.clips,
|
||||
"seconds": result.seconds,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/merged/{filename}")
|
||||
def get_merged_video(filename: str) -> FileResponse:
|
||||
"""Stream a merged scene video."""
|
||||
try:
|
||||
path = merged_path(job_runner.output_dir, filename)
|
||||
except MergeError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.detail) from e
|
||||
return FileResponse(path, media_type="video/mp4", filename=filename)
|
||||
|
||||
|
||||
@app.get("/api/jobs/{job_id}/download_log")
|
||||
def get_job_log_file(job_id: str) -> FileResponse:
|
||||
"""Download the log file for a job."""
|
||||
@@ -750,6 +843,14 @@ def main() -> None:
|
||||
help=("Directory for SQLite database (jobs + settings persistence) "
|
||||
f"(default: {default_data_dir})"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrent-jobs",
|
||||
type=int,
|
||||
default=1,
|
||||
help=("How many inference jobs may run at once (default: 1). Each gets its own "
|
||||
"long-lived worker that keeps its model loaded between jobs, so raising "
|
||||
"this loads one copy of the model per worker."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
@@ -775,6 +876,8 @@ def main() -> None:
|
||||
log_dir=log_dir,
|
||||
verbose=args.verbose,
|
||||
database=database,
|
||||
max_concurrent_jobs=args.max_concurrent_jobs,
|
||||
upload_dir=upload_dir,
|
||||
)
|
||||
|
||||
logger.info("Output directory: %s", output_dir)
|
||||
|
||||
@@ -3,14 +3,29 @@ import userEvent from '@testing-library/user-event';
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import ScenesPage from './page';
|
||||
import { getJobsList } from '@/lib/api';
|
||||
import { toast } from 'sonner';
|
||||
import { dequeueJob, downloadMergedVideo, getJobsList, mergeScene, queueJobs, updateJob } from '@/lib/api';
|
||||
import { downloadBlob } from '@/lib/utils';
|
||||
import { makeJob } from '@/test/factories';
|
||||
|
||||
vi.mock('sonner', () => ({ toast: { success: vi.fn(), error: vi.fn() } }));
|
||||
|
||||
vi.mock('@/lib/api', () => ({
|
||||
getJobsList: vi.fn(),
|
||||
queueJobs: vi.fn(),
|
||||
dequeueJob: vi.fn(),
|
||||
mergeScene: vi.fn(),
|
||||
downloadMergedVideo: vi.fn(),
|
||||
getMergedVideoUrl: (name: string) => `http://test.local/api/merged/${name}`,
|
||||
updateJob: vi.fn(),
|
||||
getJobVideoUrl: (id: string) => `http://test.local/api/jobs/${id}/video`,
|
||||
}));
|
||||
|
||||
vi.mock('@/lib/utils', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof import('@/lib/utils')>()),
|
||||
downloadBlob: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('@/components/jobs/CreateJobModal', () => ({
|
||||
default: (props: { readOnly?: boolean; editingJob?: { name?: string } | null; onClose: () => void }) => (
|
||||
<div role="dialog" data-readonly={String(!!props.readOnly)}>
|
||||
@@ -35,6 +50,40 @@ const wolf = (n: number, over: Parameters<typeof makeJob>[0] = {}) =>
|
||||
...over,
|
||||
});
|
||||
|
||||
import { attachLastFrame } from '@/lib/continuation';
|
||||
|
||||
const sixSection = (n: number) =>
|
||||
[
|
||||
'subject_definitions:',
|
||||
'<Subject 1> is a man.',
|
||||
'',
|
||||
'summary:',
|
||||
'[reference generation] A man.',
|
||||
'',
|
||||
'retention_analysis:',
|
||||
'<Subject 1> (appears in [Shot 1]): fully_preserved - suit.',
|
||||
'',
|
||||
'detailed_description:',
|
||||
`[Shot 1] Wide shot. Clip ${n}.`,
|
||||
'',
|
||||
'overall_soundscape:',
|
||||
'Quiet.',
|
||||
'',
|
||||
'non_diegetic_music:',
|
||||
'N/A',
|
||||
].join('\n');
|
||||
|
||||
/** Clip n's prompt after 'Use last frame' rewrote it (built by the same function the page uses). */
|
||||
const attachedPrompt = (n: number) => {
|
||||
const images = [
|
||||
{ source: '/refs/veteran.png', media_type: 'image' },
|
||||
{ source: '/refs/young.png', media_type: 'image' },
|
||||
];
|
||||
const r = attachLastFrame(sixSection(n), images, 'job-last-frame:wolf-1');
|
||||
if (!r.ok) throw new Error(r.error);
|
||||
return r.prompt;
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.mocked(getJobsList).mockReset();
|
||||
});
|
||||
@@ -241,6 +290,8 @@ describe('ScenesPage', () => {
|
||||
afterEach(() => vi.unstubAllGlobals());
|
||||
|
||||
const header = () => screen.getByRole('heading', { name: 'wolf-lunch' }).closest('[data-stuck]') as HTMLElement;
|
||||
// The observer is created in an effect, which can land just after the heading does.
|
||||
const observerReady = () => waitFor(() => expect(callback).not.toBeNull());
|
||||
|
||||
it('is sticky to the top of the scrolling area', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([wolf(1), wolf(2)]);
|
||||
@@ -257,6 +308,7 @@ describe('ScenesPage', () => {
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await observerReady();
|
||||
expect(header()).toHaveAttribute('data-stuck', 'false');
|
||||
expect(header()).not.toHaveClass('shadow-md');
|
||||
|
||||
@@ -275,8 +327,600 @@ describe('ScenesPage', () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([wolf(1)]);
|
||||
const { unmount } = render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await observerReady();
|
||||
unmount();
|
||||
expect(disconnect).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('using the last frame of the previous clip', () => {
|
||||
const FRAME = 'job-last-frame:wolf-1';
|
||||
const IMAGES = [
|
||||
{ source: '/refs/veteran.png', media_type: 'image' },
|
||||
{ source: '/refs/young.png', media_type: 'image' },
|
||||
];
|
||||
const done = (over: Parameters<typeof wolf>[1] = {}) =>
|
||||
wolf(1, { status: 'completed', output_path: '/out/a.mp4', prompt: sixSection(1), ...over });
|
||||
const next = (over: Parameters<typeof wolf>[1] = {}) =>
|
||||
wolf(2, { status: 'pending', prompt: sixSection(2), references: IMAGES, ...over });
|
||||
const button = () => screen.getByRole('checkbox', { name: /Use last frame of clip 1/ });
|
||||
|
||||
beforeEach(() => {
|
||||
vi.mocked(updateJob).mockReset().mockResolvedValue({});
|
||||
vi.mocked(toast.success).mockClear();
|
||||
vi.mocked(toast.error).mockClear();
|
||||
});
|
||||
|
||||
it('is offered on a clip that can still be edited', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
expect(button()).toBeEnabled();
|
||||
// the first clip has nothing before it
|
||||
expect(screen.getAllByRole('checkbox', { name: /Use last frame/ })).toHaveLength(1);
|
||||
expect(button()).not.toBeChecked();
|
||||
});
|
||||
|
||||
it('is not offered on a clip that has already run', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next({ status: 'completed', output_path: '/out/b.mp4' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(screen.queryByRole('checkbox', { name: /Use last frame/ })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("attaches the frame as a reference the server resolves, rewrites the prompt, and reloads", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(button());
|
||||
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalledTimes(1));
|
||||
const [jobId, updates] = vi.mocked(updateJob).mock.calls[0] as [string, { references: unknown[]; prompt: string }];
|
||||
expect(jobId).toBe('wolf-2');
|
||||
expect(updates.references).toEqual([...IMAGES, { source: FRAME, media_type: 'image' }]);
|
||||
expect(updates.prompt).toContain('[reference generation + keyframe completion]');
|
||||
expect(updates.prompt).toContain('<Picture 3> is the first frame of [Shot 1]');
|
||||
expect(updates.prompt).toContain('[Shot 1] Wide shot. The shot begins from <Picture 3>. Clip 2.');
|
||||
|
||||
await waitFor(() => expect(toast.success).toHaveBeenCalled());
|
||||
expect(vi.mocked(toast.success).mock.calls[0][0]).toContain('last frame of clip 1');
|
||||
await waitFor(() => expect(getJobsList).toHaveBeenCalledTimes(2));
|
||||
});
|
||||
|
||||
it('works before the previous clip has finished', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done({ status: 'pending', output_path: null }), next()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
expect(button()).toBeEnabled();
|
||||
expect(button().closest('label')?.getAttribute('title')).toMatch(/after clip 1 has finished/);
|
||||
await user.click(button());
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalled());
|
||||
});
|
||||
|
||||
it('shows progress while the job is being updated', async () => {
|
||||
const user = userEvent.setup();
|
||||
let release: (v: unknown) => void = () => {};
|
||||
vi.mocked(updateJob).mockReturnValue(new Promise((r) => (release = r)));
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(button());
|
||||
expect(screen.getByRole('checkbox', { name: 'Attaching…' })).toBeDisabled();
|
||||
|
||||
release({});
|
||||
await waitFor(() => expect(toast.success).toHaveBeenCalled());
|
||||
});
|
||||
|
||||
it("explains why it is unavailable for a prompt it can't rewrite", async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next({ prompt: 'A plain prompt.' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(button()).toBeDisabled();
|
||||
expect(button().closest('label')?.getAttribute('title')).toMatch(/six-section/);
|
||||
});
|
||||
|
||||
it.each([
|
||||
['a reference the server resolves', FRAME],
|
||||
['a frame already saved to disk', '/data/uploads/last_frames/last_frame_wolf-1.png'],
|
||||
])('shows the box checked when the frame is attached (%s)', async (_label, source) => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next({ references: [...IMAGES, { source, media_type: 'image' }] })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(button()).toBeChecked();
|
||||
});
|
||||
|
||||
it("shows a frame from another take as attached instead of stacking a second one", async () => {
|
||||
const other = 'job-last-frame:some-other-take';
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next({ references: [...IMAGES, { source: other, media_type: 'image' }] })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(screen.getByRole('checkbox', { name: "Starts from another take's last frame" })).toBeChecked();
|
||||
});
|
||||
|
||||
describe('unchecking it', () => {
|
||||
const attached = (source = FRAME, extra: { source: string; media_type: string }[] = []) =>
|
||||
next({
|
||||
prompt: attachedPrompt(2),
|
||||
references: [...IMAGES, { source, media_type: 'image' }, ...extra],
|
||||
});
|
||||
const updates = () => vi.mocked(updateJob).mock.calls[0][1] as { references: unknown[]; prompt: string };
|
||||
|
||||
it('removes the frame and restores the prompt', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), attached()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(button());
|
||||
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalledTimes(1));
|
||||
expect(vi.mocked(updateJob).mock.calls[0][0]).toBe('wolf-2');
|
||||
expect(updates().references).toEqual(IMAGES);
|
||||
expect(updates().prompt).toBe(sixSection(2));
|
||||
await waitFor(() => expect(toast.success).toHaveBeenCalled());
|
||||
await waitFor(() => expect(getJobsList).toHaveBeenCalledTimes(2));
|
||||
});
|
||||
|
||||
it('also removes a frame saved to disk', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), attached('/data/uploads/last_frames/last_frame_wolf-1.png')]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await user.click(button());
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalled());
|
||||
expect(updates().references).toEqual(IMAGES);
|
||||
});
|
||||
|
||||
it("removes another take's frame too", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), attached('job-last-frame:some-other-take')]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await user.click(screen.getByRole('checkbox', { name: "Starts from another take's last frame" }));
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalled());
|
||||
expect(updates().references).toEqual(IMAGES);
|
||||
});
|
||||
|
||||
it('is undone exactly by checking it and then unchecking it', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
// Once saved, the reload shows the clip as the first click left it.
|
||||
const attachedClip = attached();
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), attachedClip]);
|
||||
await user.click(button());
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalledTimes(1));
|
||||
expect(updates().prompt).toBe(attachedClip.prompt); // same rewrite the page's own helper produces
|
||||
|
||||
await user.click(await screen.findByRole('checkbox', { name: /Use last frame of clip 1/, checked: true }));
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalledTimes(2));
|
||||
const restored = vi.mocked(updateJob).mock.calls[1][1] as { prompt: string; references: unknown[] };
|
||||
expect(restored.prompt).toBe(sixSection(2));
|
||||
expect(restored.references).toEqual(IMAGES);
|
||||
});
|
||||
|
||||
it('says why when other images were added after the frame, and changes nothing', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
done(),
|
||||
attached(FRAME, [{ source: '/refs/extra.png', media_type: 'image' }]),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await user.click(button());
|
||||
await waitFor(() => expect(toast.error).toHaveBeenCalled());
|
||||
expect(vi.mocked(toast.error).mock.calls[0][0]).toMatch(/Other images were added/);
|
||||
expect(updateJob).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
it('reports a failure to update the job', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(updateJob).mockRejectedValue(new Error('Job is running; only pending jobs can be edited.'));
|
||||
vi.mocked(getJobsList).mockResolvedValue([done(), next()]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(button());
|
||||
|
||||
await waitFor(() => expect(toast.error).toHaveBeenCalledWith('Job is running; only pending jobs can be edited.'));
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('uses the clip that is shown when a different take is chosen', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
done({ id: 'take-a', created_at: 1 }),
|
||||
done({ id: 'take-b', created_at: 2 }),
|
||||
next(),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.selectOptions(screen.getByLabelText('Take for clip 1'), 'take-a');
|
||||
await user.click(button());
|
||||
await waitFor(() => expect(updateJob).toHaveBeenCalled());
|
||||
const updates = vi.mocked(updateJob).mock.calls[0][1] as { references: { source: string }[] };
|
||||
expect(updates.references.at(-1)?.source).toBe('job-last-frame:take-a');
|
||||
});
|
||||
});
|
||||
|
||||
describe('queueing', () => {
|
||||
const clip = (n: number, over: Parameters<typeof wolf>[1] = {}) =>
|
||||
wolf(n, { status: 'pending', prompt: sixSection(n), ...over });
|
||||
const openDialog = async (user: ReturnType<typeof userEvent.setup>) => {
|
||||
await user.click(screen.getByRole('button', { name: 'Queue scene' }));
|
||||
return screen.findByRole('dialog');
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.mocked(updateJob).mockReset().mockResolvedValue({});
|
||||
vi.mocked(queueJobs).mockReset().mockResolvedValue([]);
|
||||
vi.mocked(dequeueJob).mockReset().mockResolvedValue(makeJob());
|
||||
vi.mocked(toast.success).mockClear();
|
||||
vi.mocked(toast.error).mockClear();
|
||||
});
|
||||
|
||||
it('queues the whole scene in order and leaves the clips as they are', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(3), clip(1), clip(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
expect(within(dialog).getByText('3 clips will be queued and run one at a time, in order.')).toBeInTheDocument();
|
||||
expect(within(dialog).getByRole('checkbox')).not.toBeChecked();
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 3 clips' }));
|
||||
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['wolf-1', 'wolf-2', 'wolf-3']));
|
||||
expect(updateJob).not.toHaveBeenCalled();
|
||||
expect(toast.success).toHaveBeenCalledWith('Queued 3 clips.');
|
||||
await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument());
|
||||
});
|
||||
|
||||
it('links each clip to the frame before it when asked to, before queueing', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(3), clip(1), clip(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
await user.click(within(dialog).getByRole('checkbox'));
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 3 clips' }));
|
||||
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['wolf-1', 'wolf-2', 'wolf-3']));
|
||||
expect(updateJob).toHaveBeenCalledTimes(2);
|
||||
const refs = (n: number) =>
|
||||
(vi.mocked(updateJob).mock.calls.find((c) => c[0] === `wolf-${n}`)?.[1] as { references: { source: string }[] })
|
||||
.references.map((r) => r.source);
|
||||
expect(refs(2)).toEqual(['job-last-frame:wolf-1']);
|
||||
expect(refs(3)).toEqual(['job-last-frame:wolf-2']);
|
||||
// linked before queued, so the server never sees a queued clip without its link
|
||||
expect(vi.mocked(updateJob).mock.invocationCallOrder.at(-1)!).toBeLessThan(
|
||||
vi.mocked(queueJobs).mock.invocationCallOrder[0],
|
||||
);
|
||||
});
|
||||
|
||||
it('keeps the links you set on individual clips, and adds no others unless asked', async () => {
|
||||
const user = userEvent.setup();
|
||||
const link = [{ source: 'job-last-frame:wolf-1', media_type: 'image' }];
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1), clip(2, { references: link }), clip(3)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
expect(within(dialog).getByText(/Clip 2 already starts from a previous frame/)).toBeInTheDocument();
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 3 clips' }));
|
||||
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['wolf-1', 'wolf-2', 'wolf-3']));
|
||||
expect(updateJob).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('skips clips that already ran and says so', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
clip(1, { status: 'completed', output_path: '/o/1.mp4' }),
|
||||
clip(2),
|
||||
clip(3, { status: 'running' }),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
await user.click(within(dialog).getByRole('checkbox'));
|
||||
expect(within(dialog).getByText(/Skipping clip 1 \(already finished\), clip 3 \(already running\)/)).toBeInTheDocument();
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 1 clip' }));
|
||||
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['wolf-2']));
|
||||
// clip 2 still follows the finished clip 1's last frame
|
||||
expect(vi.mocked(updateJob).mock.calls[0][0]).toBe('wolf-2');
|
||||
});
|
||||
|
||||
it("mentions clips whose prompt can't be linked once linking is on, and still queues them", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1), clip(2), clip(3, { prompt: 'A plain prompt.' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
expect(within(dialog).queryByText(/can't be linked automatically/)).not.toBeInTheDocument();
|
||||
await user.click(within(dialog).getByRole('checkbox'));
|
||||
expect(within(dialog).getByText(/Clip 3 can't be linked automatically/)).toBeInTheDocument();
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 3 clips' }));
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['wolf-1', 'wolf-2', 'wolf-3']));
|
||||
expect(updateJob).toHaveBeenCalledTimes(1); // only clip 2
|
||||
});
|
||||
|
||||
it('queues nothing, and says which clip, when a link cannot be saved', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(updateJob).mockRejectedValue(new Error('Job is running'));
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1), clip(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
await user.click(within(dialog).getByRole('checkbox'));
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 2 clips' }));
|
||||
|
||||
await waitFor(() => expect(toast.error).toHaveBeenCalled());
|
||||
expect(vi.mocked(toast.error).mock.calls[0][0]).toContain('Clip 2: Job is running. Nothing was queued.');
|
||||
expect(queueJobs).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("reports the server's refusal", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(queueJobs).mockRejectedValue(new Error('Job already completed. Delete and re-create to run again.'));
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Queue 1 clip' }));
|
||||
await waitFor(() =>
|
||||
expect(toast.error).toHaveBeenCalledWith('Job already completed. Delete and re-create to run again.'),
|
||||
);
|
||||
});
|
||||
|
||||
it('changes nothing when cancelled', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1), clip(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const dialog = await openDialog(user);
|
||||
await user.click(within(dialog).getByRole('button', { name: 'Cancel' }));
|
||||
|
||||
await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument());
|
||||
expect(queueJobs).not.toHaveBeenCalled();
|
||||
expect(updateJob).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('is unavailable when every clip has already run or is queued', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
clip(1, { status: 'completed', output_path: '/o/1.mp4' }),
|
||||
clip(2, { status: 'queued' }),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(screen.getByRole('button', { name: 'Queue scene' })).toBeDisabled();
|
||||
});
|
||||
|
||||
it('queues a single clip', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1), clip(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(within(document.getElementById('clip-wolf-2') as HTMLElement).getByRole('button', { name: 'Queue' }));
|
||||
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['wolf-2']));
|
||||
await waitFor(() => expect(getJobsList).toHaveBeenCalledTimes(2));
|
||||
});
|
||||
|
||||
it('takes a queued clip back out of the queue', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1, { status: 'running' }), clip(2, { status: 'queued' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const row = document.getElementById('clip-wolf-2') as HTMLElement;
|
||||
expect(within(row).queryByRole('button', { name: 'Queue' })).not.toBeInTheDocument();
|
||||
await user.click(within(row).getByRole('button', { name: 'Remove from queue' }));
|
||||
await waitFor(() => expect(dequeueJob).toHaveBeenCalledWith('wolf-2'));
|
||||
});
|
||||
|
||||
it('shows what a queued clip is waiting for', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
clip(1, { status: 'running' }),
|
||||
clip(2, { status: 'queued', references: [{ source: 'job-last-frame:wolf-1', media_type: 'image' }] }),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(screen.getByText('waiting for clip 1')).toBeInTheDocument();
|
||||
expect(screen.getByText("starts from clip 1's last frame")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it('stops waiting once the previous clip has finished', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
clip(1, { status: 'completed', output_path: '/o/1.mp4' }),
|
||||
clip(2, { status: 'queued', references: [{ source: 'job-last-frame:wolf-1', media_type: 'image' }] }),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(screen.queryByText('waiting for clip 1')).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it('refreshes on its own while something is queued or running', async () => {
|
||||
const setIntervalSpy = vi.spyOn(globalThis, 'setInterval');
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1, { status: 'running' }), clip(2, { status: 'queued' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
const poll = setIntervalSpy.mock.calls.find(([, ms]) => ms === 3000)?.[0] as () => void;
|
||||
expect(poll).toBeTypeOf('function');
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
clip(1, { status: 'completed', output_path: '/o/1.mp4' }),
|
||||
clip(2, { status: 'running' }),
|
||||
]);
|
||||
await act(async () => poll());
|
||||
await waitFor(() => expect(within(document.getElementById('clip-wolf-2') as HTMLElement).getByText('running')).toBeInTheDocument());
|
||||
setIntervalSpy.mockRestore();
|
||||
});
|
||||
|
||||
it("doesn't poll when nothing is in progress", async () => {
|
||||
const setIntervalSpy = vi.spyOn(globalThis, 'setInterval');
|
||||
vi.mocked(getJobsList).mockResolvedValue([clip(1), clip(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(setIntervalSpy.mock.calls.some(([, ms]) => ms === 3000)).toBe(false);
|
||||
setIntervalSpy.mockRestore();
|
||||
});
|
||||
});
|
||||
|
||||
describe('merging a scene', () => {
|
||||
const finished = (n: number, over: Parameters<typeof wolf>[1] = {}) =>
|
||||
wolf(n, { status: 'completed', output_path: `/o/${n}.mp4`, ...over });
|
||||
const mergeButton = () => screen.getByRole('button', { name: /Merge scene|Merging/ });
|
||||
|
||||
beforeEach(() => {
|
||||
vi.mocked(mergeScene).mockReset().mockResolvedValue({ filename: 'wolf-lunch-20260927-100000.mp4', clips: 3, seconds: 33 });
|
||||
vi.mocked(downloadMergedVideo).mockReset().mockResolvedValue(new Blob(['x']));
|
||||
vi.mocked(downloadBlob).mockReset();
|
||||
vi.mocked(toast.error).mockClear();
|
||||
});
|
||||
|
||||
it('joins the clips in scene order and shows the result', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(3), finished(1), finished(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
expect(mergeButton()).toBeEnabled();
|
||||
await user.click(mergeButton());
|
||||
|
||||
await waitFor(() => expect(mergeScene).toHaveBeenCalledWith(['wolf-1', 'wolf-2', 'wolf-3'], 'wolf-lunch'));
|
||||
const dialog = await screen.findByRole('dialog');
|
||||
expect(within(dialog).getByText('wolf-lunch merged')).toBeInTheDocument();
|
||||
expect(within(dialog).getByText('3 clips joined in order · 0:33')).toBeInTheDocument();
|
||||
expect(within(dialog).getByLabelText('Merged video of wolf-lunch')).toHaveAttribute(
|
||||
'src',
|
||||
'http://test.local/api/merged/wolf-lunch-20260927-100000.mp4',
|
||||
);
|
||||
});
|
||||
|
||||
it('merges the take that is shown for a clip', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([
|
||||
finished(1, { id: 'take-a', created_at: 1 }),
|
||||
finished(1, { id: 'take-b', created_at: 2 }),
|
||||
finished(2),
|
||||
]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.selectOptions(screen.getByLabelText('Take for clip 1'), 'take-a');
|
||||
await user.click(mergeButton());
|
||||
await waitFor(() => expect(mergeScene).toHaveBeenCalledWith(['take-a', 'wolf-2'], 'wolf-lunch'));
|
||||
});
|
||||
|
||||
it('waits until every clip has finished, and says how many have not', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), wolf(2, { status: 'running' }), wolf(3, { status: 'queued' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(mergeButton()).toBeDisabled();
|
||||
expect(mergeButton()).toHaveAttribute('title', "2 clips haven't finished yet");
|
||||
});
|
||||
|
||||
it('says so for a single unfinished clip', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), wolf(2, { status: 'failed' })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(mergeButton()).toHaveAttribute('title', "1 clip hasn't finished yet");
|
||||
});
|
||||
|
||||
it('needs a video on every clip, not just a completed status', async () => {
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), finished(2, { output_path: null })]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
expect(mergeButton()).toBeDisabled();
|
||||
});
|
||||
|
||||
it("reports the server's refusal and shows nothing", async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(mergeScene).mockRejectedValue(new Error("'wolf-lunch-clip-02' hasn't finished (running), so the scene can't be merged yet."));
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), finished(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(mergeButton());
|
||||
await waitFor(() => expect(toast.error).toHaveBeenCalledWith(expect.stringContaining("hasn't finished (running)")));
|
||||
expect(screen.queryByRole('dialog')).not.toBeInTheDocument();
|
||||
expect(mergeButton()).toBeEnabled(); // can try again
|
||||
});
|
||||
|
||||
it('shows progress while merging', async () => {
|
||||
const user = userEvent.setup();
|
||||
let release: (v: { filename: string; clips: number; seconds: number }) => void = () => {};
|
||||
vi.mocked(mergeScene).mockReturnValue(new Promise((r) => (release = r)));
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), finished(2)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
|
||||
await user.click(mergeButton());
|
||||
expect(screen.getByRole('button', { name: 'Merging…' })).toBeDisabled();
|
||||
release({ filename: 'a.mp4', clips: 2, seconds: 22 });
|
||||
expect(await screen.findByRole('dialog')).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it('downloads the merged video', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), finished(2), finished(3)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await user.click(mergeButton());
|
||||
|
||||
await user.click(within(await screen.findByRole('dialog')).getByRole('button', { name: 'Download' }));
|
||||
await waitFor(() => expect(downloadBlob).toHaveBeenCalled());
|
||||
expect(downloadMergedVideo).toHaveBeenCalledWith('wolf-lunch-20260927-100000.mp4');
|
||||
expect(vi.mocked(downloadBlob).mock.calls[0][1]).toBe('wolf-lunch-20260927-100000.mp4');
|
||||
});
|
||||
|
||||
it('reports a failed download', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(downloadMergedVideo).mockRejectedValue(new Error('Failed to download the merged video'));
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), finished(2), finished(3)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await user.click(mergeButton());
|
||||
|
||||
await user.click(within(await screen.findByRole('dialog')).getByRole('button', { name: 'Download' }));
|
||||
await waitFor(() => expect(toast.error).toHaveBeenCalledWith('Failed to download the merged video'));
|
||||
expect(downloadBlob).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('closes', async () => {
|
||||
const user = userEvent.setup();
|
||||
vi.mocked(getJobsList).mockResolvedValue([finished(1), finished(2), finished(3)]);
|
||||
render(<ScenesPage />);
|
||||
await screen.findByRole('heading', { name: 'wolf-lunch' });
|
||||
await user.click(mergeButton());
|
||||
|
||||
// the footer button and the corner X both close it
|
||||
const dialog = await screen.findByRole('dialog');
|
||||
await user.click(within(dialog).getAllByRole('button', { name: 'Close' })[0]);
|
||||
await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument());
|
||||
|
||||
await user.click(mergeButton());
|
||||
await user.click(within(await screen.findByRole('dialog')).getAllByRole('button', { name: 'Close' }).at(-1)!);
|
||||
await waitFor(() => expect(screen.queryByRole('dialog')).not.toBeInTheDocument());
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -5,13 +5,20 @@ import * as React from 'react';
|
||||
|
||||
import CreateJobModal from '@/components/jobs/CreateJobModal';
|
||||
import { EDITABLE_STATUSES, SceneBoard } from '@/components/scenes/SceneBoard';
|
||||
import { MergedSceneDialog } from '@/components/scenes/MergedSceneDialog';
|
||||
import { SceneQueueDialog } from '@/components/scenes/SceneQueueDialog';
|
||||
import { Button } from '@/components/ui/button';
|
||||
import { Card } from '@/components/ui/card';
|
||||
import { NativeSelect } from '@/components/ui/native-select';
|
||||
import { getJobsList } from '@/lib/api';
|
||||
import { buildScenes } from '@/lib/scenes';
|
||||
import { toast } from 'sonner';
|
||||
import { dequeueJob, getJobsList, mergeScene, queueJobs, updateJob, type MergedScene } from '@/lib/api';
|
||||
import { attachedLastFrames, attachLastFrame, deferredLastFrameSource, detachLastFrame } from '@/lib/continuation';
|
||||
import type { SceneQueuePlan } from '@/lib/sceneQueue';
|
||||
import { buildScenes, type Scene, type SceneClip } from '@/lib/scenes';
|
||||
import type { Job } from '@/lib/types';
|
||||
|
||||
const POLL_MS = 3000;
|
||||
|
||||
export default function ScenesPage() {
|
||||
const [jobs, setJobs] = React.useState<Job[]>([]);
|
||||
const [isLoading, setIsLoading] = React.useState(true);
|
||||
@@ -21,6 +28,8 @@ export default function ScenesPage() {
|
||||
const [openJob, setOpenJob] = React.useState<Job | null>(null);
|
||||
// which re-run of a clip number the user picked, by clip
|
||||
const [choices, setChoices] = React.useState<Record<string, string>>({});
|
||||
const [queueTarget, setQueueTarget] = React.useState<Scene | null>(null);
|
||||
const [merged, setMerged] = React.useState<(MergedScene & { title: string }) | null>(null);
|
||||
|
||||
React.useEffect(() => {
|
||||
let cancelled = false;
|
||||
@@ -47,6 +56,110 @@ export default function ScenesPage() {
|
||||
setReloadKey((k) => k + 1);
|
||||
}
|
||||
|
||||
// While anything is queued or running, keep the board current without the
|
||||
// loading state (which would blank or flicker it).
|
||||
const hasActive = jobs.some((j) => j.status === 'queued' || j.status === 'running');
|
||||
React.useEffect(() => {
|
||||
if (!hasActive) return;
|
||||
const interval = setInterval(() => {
|
||||
getJobsList('inference')
|
||||
.then(setJobs)
|
||||
.catch(() => {});
|
||||
}, POLL_MS);
|
||||
return () => clearInterval(interval);
|
||||
}, [hasActive]);
|
||||
|
||||
// Attach the previous clip's last frame to `clip` and rewrite its prompt to open on
|
||||
// it. The frame is a reference the server resolves when the clip runs, so the
|
||||
// previous clip doesn't have to be finished yet. Returns the picture label.
|
||||
async function attachPreviousFrame(clip: SceneClip, previous: SceneClip): Promise<string> {
|
||||
const result = attachLastFrame(
|
||||
clip.job.prompt,
|
||||
clip.job.references ?? [],
|
||||
deferredLastFrameSource(previous.job.id),
|
||||
);
|
||||
if (!result.ok) throw new Error(result.error);
|
||||
await updateJob(clip.job.id, { references: result.references, prompt: result.prompt });
|
||||
return result.label;
|
||||
}
|
||||
|
||||
async function useLastFrame(clip: SceneClip, previous: SceneClip) {
|
||||
try {
|
||||
const label = await attachPreviousFrame(clip, previous);
|
||||
toast.success(`Clip ${clip.position} now starts from the last frame of clip ${previous.position} (${label}).`);
|
||||
reload();
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't use the last frame.");
|
||||
}
|
||||
}
|
||||
|
||||
// Undo the above: take every attached last frame back off the clip.
|
||||
async function stopUsingLastFrame(clip: SceneClip) {
|
||||
try {
|
||||
let prompt = clip.job.prompt;
|
||||
let references = clip.job.references ?? [];
|
||||
// Newest first: only the most recently added image can be removed cleanly.
|
||||
for (const { source } of attachedLastFrames(references).reverse()) {
|
||||
const result = detachLastFrame(prompt, references, source);
|
||||
if (!result.ok) throw new Error(result.error);
|
||||
prompt = result.prompt;
|
||||
references = result.references;
|
||||
}
|
||||
await updateJob(clip.job.id, { references, prompt });
|
||||
toast.success(`Clip ${clip.position} no longer starts from a previous clip's last frame.`);
|
||||
reload();
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't remove the last frame.");
|
||||
}
|
||||
}
|
||||
|
||||
// Join the scene's clips, in the order shown (the take chosen for each), into one video.
|
||||
async function mergeSceneClips(scene: Scene) {
|
||||
try {
|
||||
const result = await mergeScene(scene.clips.map((c) => c.job.id), scene.title);
|
||||
setMerged({ ...result, title: scene.title });
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't merge the scene.");
|
||||
}
|
||||
}
|
||||
|
||||
async function queueClip(clip: SceneClip) {
|
||||
try {
|
||||
await queueJobs([clip.job.id]);
|
||||
toast.success(`Queued clip ${clip.position}.`);
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't queue the clip.");
|
||||
}
|
||||
reload();
|
||||
}
|
||||
|
||||
async function dequeueClip(clip: SceneClip) {
|
||||
try {
|
||||
await dequeueJob(clip.job.id);
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't remove the clip from the queue.");
|
||||
}
|
||||
reload();
|
||||
}
|
||||
|
||||
// Link clips to the frames before them (if asked), then queue them all in scene order.
|
||||
async function queueScene(plan: SceneQueuePlan) {
|
||||
try {
|
||||
for (const { clip, previous } of plan.chain) {
|
||||
await attachPreviousFrame(clip, previous).catch((e) => {
|
||||
throw new Error(`Clip ${clip.position}: ${e instanceof Error ? e.message : e}. Nothing was queued.`);
|
||||
});
|
||||
}
|
||||
await queueJobs(plan.queue.map((c) => c.job.id));
|
||||
const n = plan.queue.length;
|
||||
toast.success(`Queued ${n} clip${n === 1 ? '' : 's'}.`);
|
||||
setQueueTarget(null);
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't queue the scene.");
|
||||
}
|
||||
reload();
|
||||
}
|
||||
|
||||
const scenes = React.useMemo(() => buildScenes(jobs, choices), [jobs, choices]);
|
||||
const scene = scenes.find((s) => s.key === selectedKey) ?? scenes[0];
|
||||
|
||||
@@ -107,12 +220,21 @@ export default function ScenesPage() {
|
||||
<SceneBoard
|
||||
scene={scene}
|
||||
onOpenJob={setOpenJob}
|
||||
onUseLastFrame={useLastFrame}
|
||||
onDetachLastFrame={stopUsingLastFrame}
|
||||
onQueueScene={setQueueTarget}
|
||||
onMergeScene={mergeSceneClips}
|
||||
onQueueClip={queueClip}
|
||||
onDequeueClip={dequeueClip}
|
||||
onChooseTake={(takeKey, jobId) => setChoices((c) => ({ ...c, [takeKey]: jobId }))}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</Card>
|
||||
|
||||
<MergedSceneDialog merged={merged} onClose={() => setMerged(null)} />
|
||||
<SceneQueueDialog scene={queueTarget} onConfirm={queueScene} onClose={() => setQueueTarget(null)} />
|
||||
|
||||
{openJob && (
|
||||
<CreateJobModal
|
||||
isOpen
|
||||
|
||||
@@ -5,7 +5,9 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import JobCard from '@/components/jobs/JobCard';
|
||||
import {
|
||||
deleteJob,
|
||||
dequeueJob,
|
||||
downloadJobVideo,
|
||||
queueJobs,
|
||||
startJob,
|
||||
stopJob,
|
||||
} from '@/lib/api';
|
||||
@@ -18,6 +20,8 @@ vi.mock('@/lib/api', () => ({
|
||||
stopJob: vi.fn(),
|
||||
deleteJob: vi.fn(),
|
||||
downloadJobVideo: vi.fn(),
|
||||
queueJobs: vi.fn(),
|
||||
dequeueJob: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('@/lib/utils', async (importOriginal) => {
|
||||
@@ -47,6 +51,8 @@ beforeEach(() => {
|
||||
setActiveJobId(null);
|
||||
vi.mocked(startJob).mockResolvedValue({} as Job);
|
||||
vi.mocked(stopJob).mockResolvedValue({} as Job);
|
||||
vi.mocked(queueJobs).mockResolvedValue([]);
|
||||
vi.mocked(dequeueJob).mockResolvedValue({} as Job);
|
||||
vi.mocked(deleteJob).mockResolvedValue(undefined);
|
||||
vi.mocked(downloadJobVideo).mockResolvedValue(new Blob());
|
||||
vi.spyOn(window, 'confirm').mockReturnValue(true);
|
||||
@@ -91,6 +97,46 @@ describe('JobCard', () => {
|
||||
expect(onJobUpdated).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it.each(['pending', 'failed', 'stopped'])('queues a %s job and notifies the parent', async (status) => {
|
||||
const onJobUpdated = vi.fn();
|
||||
render(<JobCard job={makeJob({ status })} onJobUpdated={onJobUpdated} />);
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Queue' }));
|
||||
await waitFor(() => expect(queueJobs).toHaveBeenCalledWith(['job-1']));
|
||||
expect(onJobUpdated).toHaveBeenCalled();
|
||||
expect(startJob).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('reports a job that cannot be queued', async () => {
|
||||
vi.mocked(queueJobs).mockRejectedValue(new Error('Jobs would wait on each other'));
|
||||
render(<JobCard job={makeJob({ status: 'pending' })} />);
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Queue' }));
|
||||
await waitFor(() => expect(window.alert).toHaveBeenCalledWith('Jobs would wait on each other'));
|
||||
});
|
||||
|
||||
it('shows a queued job as queued, and takes it back out of the queue', async () => {
|
||||
const onJobUpdated = vi.fn();
|
||||
render(<JobCard job={makeJob({ status: 'queued' })} onJobUpdated={onJobUpdated} />);
|
||||
|
||||
expect(screen.getByText('queued')).toBeInTheDocument();
|
||||
// nothing to start or queue again; it can't be edited either
|
||||
for (const name of ['Start', 'Restart', 'Queue', 'Stop', 'Edit']) {
|
||||
expect(screen.queryByRole('button', { name })).not.toBeInTheDocument();
|
||||
}
|
||||
expect(screen.getByRole('button', { name: 'View' })).toBeInTheDocument();
|
||||
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Remove from queue' }));
|
||||
await waitFor(() => expect(dequeueJob).toHaveBeenCalledWith('job-1'));
|
||||
expect(onJobUpdated).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('does not offer to queue a running or completed job', () => {
|
||||
for (const status of ['running', 'completed']) {
|
||||
const { unmount } = render(<JobCard job={makeJob({ status, started_at: Date.now() })} />);
|
||||
expect(screen.queryByRole('button', { name: 'Queue' })).not.toBeInTheDocument();
|
||||
unmount();
|
||||
}
|
||||
});
|
||||
|
||||
it('stops a running job', async () => {
|
||||
render(<JobCard job={makeJob({ status: 'running', started_at: Date.now() })} />);
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Stop' }));
|
||||
|
||||
@@ -10,7 +10,9 @@ import { useStore } from '@/hooks/useStore';
|
||||
import {
|
||||
deleteJob,
|
||||
duplicateJob,
|
||||
dequeueJob,
|
||||
downloadJobVideo,
|
||||
queueJobs,
|
||||
startJob,
|
||||
stopJob,
|
||||
} from '@/lib/api';
|
||||
@@ -51,6 +53,7 @@ function computeElapsed(job: Job, currentTime: number): string | null {
|
||||
|
||||
const BADGE_VARIANTS: Record<string, BadgeProps['variant']> = {
|
||||
pending: 'secondary',
|
||||
queued: 'default',
|
||||
running: 'warning',
|
||||
completed: 'success',
|
||||
ready: 'success',
|
||||
@@ -92,6 +95,36 @@ export default function JobCard({ job, onJobUpdated }: JobCardProps) {
|
||||
}
|
||||
}
|
||||
|
||||
async function handleQueue(e: React.MouseEvent) {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
if (isLoading) return;
|
||||
setIsLoading(true);
|
||||
try {
|
||||
await queueJobs([job.id]);
|
||||
onJobUpdated?.();
|
||||
} catch (err) {
|
||||
alert(err instanceof Error ? err.message : 'Failed to queue job');
|
||||
} finally {
|
||||
setIsLoading(false);
|
||||
}
|
||||
}
|
||||
|
||||
async function handleDequeue(e: React.MouseEvent) {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
if (isLoading) return;
|
||||
setIsLoading(true);
|
||||
try {
|
||||
await dequeueJob(job.id);
|
||||
onJobUpdated?.();
|
||||
} catch (err) {
|
||||
alert(err instanceof Error ? err.message : 'Failed to remove job from the queue');
|
||||
} finally {
|
||||
setIsLoading(false);
|
||||
}
|
||||
}
|
||||
|
||||
async function handleStop(e: React.MouseEvent) {
|
||||
e.preventDefault();
|
||||
e.stopPropagation();
|
||||
@@ -238,7 +271,30 @@ export default function JobCard({ job, onJobUpdated }: JobCardProps) {
|
||||
>
|
||||
Start
|
||||
</Button>
|
||||
) : job.status === 'queued' ? (
|
||||
<Button
|
||||
size="sm"
|
||||
variant="outline"
|
||||
onClick={handleDequeue}
|
||||
disabled={isLoading}
|
||||
title="Take this job out of the queue"
|
||||
>
|
||||
Remove from queue
|
||||
</Button>
|
||||
) : null}
|
||||
{(job.status === 'pending' ||
|
||||
job.status === 'failed' ||
|
||||
job.status === 'stopped') && (
|
||||
<Button
|
||||
size="sm"
|
||||
variant="outline"
|
||||
onClick={handleQueue}
|
||||
disabled={isLoading}
|
||||
title="Run this job when the ones queued before it have finished"
|
||||
>
|
||||
Queue
|
||||
</Button>
|
||||
)}
|
||||
{job.status === 'completed' &&
|
||||
job.output_path &&
|
||||
job.job_type === 'inference' && (
|
||||
|
||||
@@ -77,7 +77,7 @@ export default function JobDetailsSidebar({
|
||||
const isRestarting =
|
||||
previousJobId.current === job.id &&
|
||||
wasTerminal &&
|
||||
(job.status === 'pending' || job.status === 'running');
|
||||
(job.status === 'pending' || job.status === 'queued' || job.status === 'running');
|
||||
|
||||
if (previousJobId.current !== job.id || isRestarting) {
|
||||
stateRef.current.text = '';
|
||||
@@ -89,7 +89,10 @@ export default function JobDetailsSidebar({
|
||||
}, [job.id, job.status]);
|
||||
|
||||
React.useEffect(() => {
|
||||
const shouldPoll = job.status === 'running' || job.status === 'pending';
|
||||
const shouldPoll =
|
||||
job.status === 'running' ||
|
||||
job.status === 'pending' ||
|
||||
job.status === 'queued';
|
||||
let pollInterval: ReturnType<typeof setInterval> | null = null;
|
||||
let mounted = true;
|
||||
// Effect-local lock so this job's first fetch is never blocked by a
|
||||
|
||||
@@ -108,7 +108,8 @@ export default function JobQueue({ jobType, jobTypesForList }: JobQueueProps) {
|
||||
|
||||
// Poll every second while any job is running/pending; stop otherwise.
|
||||
const hasActive = jobs.some(
|
||||
(j) => j.status === 'running' || j.status === 'pending',
|
||||
(j) =>
|
||||
j.status === 'running' || j.status === 'pending' || j.status === 'queued',
|
||||
);
|
||||
React.useEffect(() => {
|
||||
if (!hasActive) return;
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
'use client';
|
||||
|
||||
import * as React from 'react';
|
||||
import { toast } from 'sonner';
|
||||
|
||||
import { Button } from '@/components/ui/button';
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from '@/components/ui/dialog';
|
||||
import { downloadMergedVideo, getMergedVideoUrl, type MergedScene } from '@/lib/api';
|
||||
import { formatClock } from '@/lib/scenes';
|
||||
import { downloadBlob } from '@/lib/utils';
|
||||
|
||||
export interface MergedSceneDialogProps {
|
||||
/** the merged video to show, or null when closed */
|
||||
merged: (MergedScene & { title: string }) | null;
|
||||
onClose: () => void;
|
||||
}
|
||||
|
||||
/** The finished, merged scene: play it, and save it. */
|
||||
export function MergedSceneDialog({ merged, onClose }: MergedSceneDialogProps) {
|
||||
const [downloading, setDownloading] = React.useState(false);
|
||||
|
||||
async function download() {
|
||||
if (!merged) return;
|
||||
setDownloading(true);
|
||||
try {
|
||||
downloadBlob(await downloadMergedVideo(merged.filename), merged.filename);
|
||||
} catch (e) {
|
||||
toast.error(e instanceof Error ? e.message : "Couldn't download the video.");
|
||||
} finally {
|
||||
setDownloading(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<Dialog open={!!merged} onOpenChange={(open) => !open && onClose()}>
|
||||
<DialogContent className="max-w-3xl">
|
||||
<DialogHeader>
|
||||
<DialogTitle>{merged?.title} merged</DialogTitle>
|
||||
<DialogDescription>
|
||||
{merged &&
|
||||
`${merged.clips} clip${merged.clips === 1 ? '' : 's'} joined in order · ${formatClock(merged.seconds)}`}
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
{merged && (
|
||||
<video
|
||||
src={getMergedVideoUrl(merged.filename)}
|
||||
aria-label={`Merged video of ${merged.title}`}
|
||||
className="max-h-[60vh] w-full rounded-md border border-border bg-black"
|
||||
controls
|
||||
playsInline
|
||||
preload="metadata"
|
||||
/>
|
||||
)}
|
||||
<DialogFooter>
|
||||
<Button type="button" variant="outline" onClick={onClose}>
|
||||
Close
|
||||
</Button>
|
||||
<Button type="button" onClick={download} disabled={downloading}>
|
||||
{downloading ? 'Downloading…' : 'Download'}
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -7,13 +7,16 @@ import { Badge, type BadgeProps } from '@/components/ui/badge';
|
||||
import { Button } from '@/components/ui/button';
|
||||
import { NativeSelect } from '@/components/ui/native-select';
|
||||
import { getJobVideoUrl } from '@/lib/api';
|
||||
import { attachedLastFrames, whyCannotAttach } from '@/lib/continuation';
|
||||
import { formatCutTime } from '@/lib/h3Shots';
|
||||
import { QUEUEABLE_STATUSES } from '@/lib/sceneQueue';
|
||||
import { formatClock, type Scene, type SceneClip } from '@/lib/scenes';
|
||||
import type { Job } from '@/lib/types';
|
||||
import { cn } from '@/lib/utils';
|
||||
|
||||
const BADGE_VARIANTS: Record<string, BadgeProps['variant']> = {
|
||||
pending: 'secondary',
|
||||
queued: 'default',
|
||||
running: 'warning',
|
||||
completed: 'success',
|
||||
failed: 'destructive',
|
||||
@@ -23,6 +26,7 @@ const BADGE_VARIANTS: Record<string, BadgeProps['variant']> = {
|
||||
const SEGMENT_TONE: Record<string, string> = {
|
||||
completed: 'bg-emerald-500/70',
|
||||
running: 'bg-amber-500/70',
|
||||
queued: 'bg-sky-500/60',
|
||||
failed: 'bg-rose-500/70',
|
||||
};
|
||||
|
||||
@@ -91,20 +95,66 @@ function ClipPreview({ job }: { job: Job }) {
|
||||
|
||||
function ClipRow({
|
||||
clip,
|
||||
previous,
|
||||
firstShotNumber,
|
||||
onOpenJob,
|
||||
onChooseTake,
|
||||
onUseLastFrame,
|
||||
onDetachLastFrame,
|
||||
onQueueClip,
|
||||
onDequeueClip,
|
||||
}: {
|
||||
clip: SceneClip;
|
||||
/** the clip before this one in the scene, if any */
|
||||
previous: SceneClip | undefined;
|
||||
firstShotNumber: number;
|
||||
onOpenJob: (job: Job) => void;
|
||||
onChooseTake: (takeKey: string, jobId: string) => void;
|
||||
onUseLastFrame: (clip: SceneClip, previous: SceneClip) => Promise<void>;
|
||||
onDetachLastFrame: (clip: SceneClip) => Promise<void>;
|
||||
onQueueClip: (clip: SceneClip) => Promise<void>;
|
||||
onDequeueClip: (clip: SceneClip) => Promise<void>;
|
||||
}) {
|
||||
const [previewing, setPreviewing] = React.useState(false);
|
||||
const [attaching, setAttaching] = React.useState(false);
|
||||
const [queueing, setQueueing] = React.useState(false);
|
||||
const { job } = clip;
|
||||
const canPreview = job.status === 'completed' && !!job.output_path;
|
||||
const editable = EDITABLE_STATUSES.includes(job.status);
|
||||
|
||||
const queued = job.status === 'queued';
|
||||
const canQueue = QUEUEABLE_STATUSES.includes(job.status);
|
||||
|
||||
// Starting from the previous clip's last frame: offered on clips that can still be
|
||||
// edited. The frame is taken when this clip runs, so the previous clip needn't be done yet.
|
||||
const frames = attachedLastFrames(job.references);
|
||||
const fromPrevious = previous && frames.some((f) => f.jobId === previous.job.id);
|
||||
const fromAnotherTake = frames.some((f) => f.jobId !== previous?.job.id);
|
||||
const previousDone = !!previous && previous.job.status === 'completed';
|
||||
const cannotUseFrame = previous ? whyCannotAttach(job.prompt, job.references) : null;
|
||||
|
||||
const usingLastFrame = frames.length > 0;
|
||||
const lastFrameBlocked = attaching || (!usingLastFrame && !!cannotUseFrame);
|
||||
|
||||
async function toggleLastFrame() {
|
||||
if (!previous) return;
|
||||
setAttaching(true);
|
||||
try {
|
||||
await (usingLastFrame ? onDetachLastFrame(clip) : onUseLastFrame(clip, previous));
|
||||
} finally {
|
||||
setAttaching(false);
|
||||
}
|
||||
}
|
||||
|
||||
async function queueOrDequeue(action: (clip: SceneClip) => Promise<void>) {
|
||||
setQueueing(true);
|
||||
try {
|
||||
await action(clip);
|
||||
} finally {
|
||||
setQueueing(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<li id={`clip-${job.id}`} className="scroll-mt-32 rounded-xl border border-border bg-card p-3">
|
||||
<div className="mb-2 flex flex-wrap items-center gap-2">
|
||||
@@ -121,7 +171,48 @@ function ClipRow({
|
||||
continues previous shot
|
||||
</Badge>
|
||||
)}
|
||||
<span className="ml-auto flex items-center gap-2">
|
||||
<span className="ml-auto flex flex-wrap items-center justify-end gap-2">
|
||||
{previous && queued && fromPrevious && (
|
||||
<Badge variant="outline" className="normal-case tracking-normal">
|
||||
starts from clip {previous.position}'s last frame
|
||||
</Badge>
|
||||
)}
|
||||
{previous && queued && fromPrevious && !previousDone && (
|
||||
<Badge variant="outline" className="normal-case tracking-normal">
|
||||
waiting for clip {previous.position}
|
||||
</Badge>
|
||||
)}
|
||||
{previous && editable && (
|
||||
<label
|
||||
className={cn(
|
||||
'flex items-center gap-1.5 text-xs text-foreground',
|
||||
lastFrameBlocked && 'text-muted-foreground',
|
||||
)}
|
||||
title={
|
||||
usingLastFrame
|
||||
? "Stop starting this clip from the previous clip's last frame"
|
||||
: (cannotUseFrame ??
|
||||
`Attach the last frame of clip ${previous.position} as this clip's opening frame${
|
||||
previousDone ? '' : ` (taken when this clip runs, after clip ${previous.position} has finished)`
|
||||
}`)
|
||||
}
|
||||
>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={usingLastFrame}
|
||||
onChange={toggleLastFrame}
|
||||
disabled={lastFrameBlocked}
|
||||
className="size-4"
|
||||
/>
|
||||
{attaching
|
||||
? usingLastFrame
|
||||
? 'Removing…'
|
||||
: 'Attaching…'
|
||||
: fromPrevious || !usingLastFrame
|
||||
? `Use last frame of clip ${previous.position}`
|
||||
: "Starts from another take's last frame"}
|
||||
</label>
|
||||
)}
|
||||
{clip.takeKey && clip.takes.length > 1 && (
|
||||
<NativeSelect
|
||||
aria-label={`Take for clip ${clip.position}`}
|
||||
@@ -141,6 +232,29 @@ function ClipRow({
|
||||
{previewing ? 'Hide video' : 'Preview'}
|
||||
</Button>
|
||||
)}
|
||||
{canQueue && (
|
||||
<Button
|
||||
type="button"
|
||||
size="sm"
|
||||
variant="outline"
|
||||
onClick={() => queueOrDequeue(onQueueClip)}
|
||||
disabled={queueing}
|
||||
title="Run this clip when the ones queued before it have finished"
|
||||
>
|
||||
Queue
|
||||
</Button>
|
||||
)}
|
||||
{queued && (
|
||||
<Button
|
||||
type="button"
|
||||
size="sm"
|
||||
variant="outline"
|
||||
onClick={() => queueOrDequeue(onDequeueClip)}
|
||||
disabled={queueing}
|
||||
>
|
||||
Remove from queue
|
||||
</Button>
|
||||
)}
|
||||
<Button type="button" size="sm" variant="outline" onClick={() => onOpenJob(job)}>
|
||||
{editable ? 'Edit' : 'View'}
|
||||
</Button>
|
||||
@@ -180,14 +294,48 @@ function ClipRow({
|
||||
export interface SceneBoardProps {
|
||||
scene: Scene;
|
||||
onOpenJob: (job: Job) => void;
|
||||
/** start a clip from the last frame of the one before it; resolves when done (errors are reported by the caller) */
|
||||
onUseLastFrame: (clip: SceneClip, previous: SceneClip) => Promise<void>;
|
||||
/** pick which re-run of a clip number to show */
|
||||
onChooseTake: (takeKey: string, jobId: string) => void;
|
||||
/** open the confirmation for queueing every runnable clip of the scene */
|
||||
onQueueScene: (scene: Scene) => void;
|
||||
/** join the scene's finished clips into one video; resolves when done (errors are reported by the caller) */
|
||||
onMergeScene: (scene: Scene) => Promise<void>;
|
||||
/** stop starting a clip from a previous clip's last frame; resolves when done */
|
||||
onDetachLastFrame: (clip: SceneClip) => Promise<void>;
|
||||
/** queue / unqueue one clip; resolve when done (errors are reported by the caller) */
|
||||
onQueueClip: (clip: SceneClip) => Promise<void>;
|
||||
onDequeueClip: (clip: SceneClip) => Promise<void>;
|
||||
}
|
||||
|
||||
/** One scene: a timeline of its clips, then every clip's shots in order, numbered straight through. */
|
||||
export function SceneBoard({ scene, onOpenJob, onChooseTake }: SceneBoardProps) {
|
||||
export function SceneBoard({
|
||||
scene,
|
||||
onOpenJob,
|
||||
onChooseTake,
|
||||
onUseLastFrame,
|
||||
onDetachLastFrame,
|
||||
onQueueScene,
|
||||
onMergeScene,
|
||||
onQueueClip,
|
||||
onDequeueClip,
|
||||
}: SceneBoardProps) {
|
||||
let shotNumber = 1;
|
||||
const statuses = Object.entries(scene.statusCounts);
|
||||
const anythingToQueue = scene.clips.some((c) => QUEUEABLE_STATUSES.includes(c.job.status));
|
||||
|
||||
// Merging needs every clip's video, so a scene missing one can't pass for the whole thing.
|
||||
const unfinished = scene.clips.filter((c) => c.job.status !== 'completed' || !c.job.output_path).length;
|
||||
const [merging, setMerging] = React.useState(false);
|
||||
async function merge() {
|
||||
setMerging(true);
|
||||
try {
|
||||
await onMergeScene(scene);
|
||||
} finally {
|
||||
setMerging(false);
|
||||
}
|
||||
}
|
||||
|
||||
// The header stays where it is until it reaches the top of the scrolling area,
|
||||
// then sticks there. A 1px sentinel sitting just above it tells us which: once
|
||||
@@ -225,21 +373,51 @@ export function SceneBoard({ scene, onOpenJob, onChooseTake }: SceneBoardProps)
|
||||
</Badge>
|
||||
))}
|
||||
</span>
|
||||
<span className="ml-auto flex flex-wrap gap-2">
|
||||
<Button
|
||||
type="button"
|
||||
size="sm"
|
||||
variant="outline"
|
||||
onClick={merge}
|
||||
disabled={unfinished > 0 || merging}
|
||||
title={
|
||||
unfinished > 0
|
||||
? `${unfinished} clip${unfinished === 1 ? " hasn't" : "s haven't"} finished yet`
|
||||
: 'Join all the clips, in order, into one video'
|
||||
}
|
||||
>
|
||||
{merging ? 'Merging…' : 'Merge scene'}
|
||||
</Button>
|
||||
<Button
|
||||
type="button"
|
||||
size="sm"
|
||||
onClick={() => onQueueScene(scene)}
|
||||
disabled={!anythingToQueue}
|
||||
title={anythingToQueue ? 'Queue every clip that has not run yet' : 'Every clip has already run or is queued'}
|
||||
>
|
||||
Queue scene
|
||||
</Button>
|
||||
</span>
|
||||
</div>
|
||||
<Timeline scene={scene} />
|
||||
</div>
|
||||
|
||||
<ol className="flex list-none flex-col gap-3 p-0">
|
||||
{scene.clips.map((clip) => {
|
||||
{scene.clips.map((clip, index) => {
|
||||
const first = shotNumber;
|
||||
shotNumber += clip.shots.length;
|
||||
return (
|
||||
<ClipRow
|
||||
key={clip.takeKey ?? clip.job.id}
|
||||
clip={clip}
|
||||
previous={scene.clips[index - 1]}
|
||||
firstShotNumber={first}
|
||||
onOpenJob={onOpenJob}
|
||||
onChooseTake={onChooseTake}
|
||||
onUseLastFrame={onUseLastFrame}
|
||||
onDetachLastFrame={onDetachLastFrame}
|
||||
onQueueClip={onQueueClip}
|
||||
onDequeueClip={onDequeueClip}
|
||||
/>
|
||||
);
|
||||
})}
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
'use client';
|
||||
|
||||
import * as React from 'react';
|
||||
|
||||
import { Button } from '@/components/ui/button';
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from '@/components/ui/dialog';
|
||||
import { attachedLastFrames } from '@/lib/continuation';
|
||||
import { planSceneQueue, QUEUEABLE_STATUSES, type SceneQueuePlan } from '@/lib/sceneQueue';
|
||||
import type { Scene, SceneClip } from '@/lib/scenes';
|
||||
|
||||
const plural = (n: number, word: string) => `${n} ${word}${n === 1 ? '' : 's'}`;
|
||||
const clipNumbers = (clips: SceneClip[]) => clips.map((c) => c.position).join(', ');
|
||||
|
||||
export interface SceneQueueDialogProps {
|
||||
scene: Scene | null;
|
||||
onConfirm: (plan: SceneQueuePlan) => Promise<void>;
|
||||
onClose: () => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Confirm queueing a whole scene: which clips are queued, and whether each
|
||||
* starts from the previous clip's last frame. Nothing is sent until "Queue".
|
||||
*/
|
||||
export function SceneQueueDialog({ scene, onConfirm, onClose }: SceneQueueDialogProps) {
|
||||
// Off by default: the per-clip "Use last frame" boxes are the source of truth, and
|
||||
// this would add links to every clip that doesn't have one.
|
||||
const [chain, setChain] = React.useState(false);
|
||||
const [busy, setBusy] = React.useState(false);
|
||||
|
||||
// Each opening starts from the default choice.
|
||||
React.useEffect(() => {
|
||||
if (scene) setChain(false);
|
||||
}, [scene]);
|
||||
|
||||
const plan = React.useMemo(() => (scene ? planSceneQueue(scene.clips, { chain }) : null), [scene, chain]);
|
||||
const alreadyLinked = React.useMemo(
|
||||
() =>
|
||||
scene
|
||||
? scene.clips.filter(
|
||||
(c) => QUEUEABLE_STATUSES.includes(c.job.status) && attachedLastFrames(c.job.references).length > 0,
|
||||
)
|
||||
: [],
|
||||
[scene],
|
||||
);
|
||||
// Whether the option would do anything is independent of the checkbox.
|
||||
const linkable = React.useMemo(
|
||||
() => (scene ? planSceneQueue(scene.clips, { chain: true }) : null),
|
||||
[scene],
|
||||
);
|
||||
|
||||
async function confirm() {
|
||||
if (!plan) return;
|
||||
setBusy(true);
|
||||
try {
|
||||
await onConfirm(plan);
|
||||
} finally {
|
||||
setBusy(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<Dialog open={!!scene} onOpenChange={(open) => !open && !busy && onClose()}>
|
||||
<DialogContent>
|
||||
<DialogHeader>
|
||||
<DialogTitle>Queue {scene?.title}</DialogTitle>
|
||||
<DialogDescription>
|
||||
{plan && plan.queue.length > 0
|
||||
? `${plural(plan.queue.length, 'clip')} will be queued and run one at a time, in order.`
|
||||
: 'There is nothing in this scene to queue.'}
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
|
||||
{plan && linkable && (
|
||||
<div className="flex flex-col gap-3 text-sm">
|
||||
{alreadyLinked.length > 0 && (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Clip {clipNumbers(alreadyLinked)} already start{alreadyLinked.length === 1 ? 's' : ''} from a previous
|
||||
frame (as you set on the clip{alreadyLinked.length === 1 ? '' : 's'}). Other clips are queued as they are.
|
||||
</p>
|
||||
)}
|
||||
{linkable.chain.length > 0 && (
|
||||
<label className="flex items-start gap-2">
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={chain}
|
||||
onChange={(e) => setChain(e.target.checked)}
|
||||
className="mt-1 size-4"
|
||||
/>
|
||||
<span>
|
||||
Start each clip from the previous clip's last frame
|
||||
<span className="block text-xs text-muted-foreground">
|
||||
Also links {plural(linkable.chain.length, 'clip')} that don't have it (
|
||||
{clipNumbers(linkable.chain.map((c) => c.clip))}), rewriting their prompts. Each waits for the clip
|
||||
before it to finish.
|
||||
</span>
|
||||
</span>
|
||||
</label>
|
||||
)}
|
||||
{chain && plan.cannotChain.length > 0 && (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Clip {clipNumbers(plan.cannotChain.map((c) => c.clip))} can't be linked automatically (its prompt
|
||||
isn't in the six-section format) and will start without the previous frame.
|
||||
</p>
|
||||
)}
|
||||
{plan.skipped.length > 0 && (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Skipping{' '}
|
||||
{plan.skipped.map((s) => `clip ${s.clip.position} (${s.reason})`).join(', ')}.
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<DialogFooter>
|
||||
<Button type="button" variant="outline" onClick={onClose} disabled={busy}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button type="button" onClick={confirm} disabled={busy || !plan || plan.queue.length === 0}>
|
||||
{busy ? 'Queueing…' : plan ? `Queue ${plural(plan.queue.length, 'clip')}` : 'Queue'}
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -194,6 +194,20 @@ export async function updateJob(
|
||||
return response.json();
|
||||
}
|
||||
|
||||
/**
|
||||
* Ask the server to save the last frame of a finished clip and return its path,
|
||||
* ready to attach to another job as an image reference.
|
||||
*/
|
||||
export async function extractLastFrame(jobId: string): Promise<{ path: string; media_type: MediaType }> {
|
||||
const baseApiUrl = getApiBaseUrl();
|
||||
const response = await fetch(`${baseApiUrl}/jobs/${jobId}/last-frame`, { method: "POST" });
|
||||
if (!response.ok) {
|
||||
const err = await response.json().catch(() => ({ detail: "Could not get the last frame" }));
|
||||
throw new Error(err.detail || "Could not get the last frame");
|
||||
}
|
||||
return response.json();
|
||||
}
|
||||
|
||||
export async function uploadMedia(
|
||||
file: File,
|
||||
): Promise<{ path: string; media_type: MediaType }> {
|
||||
@@ -313,6 +327,80 @@ export async function startJob(id: string): Promise<Job> {
|
||||
return response.json();
|
||||
}
|
||||
|
||||
/**
|
||||
* Queue jobs, in the order given, to run one after another as the server has room.
|
||||
* A job that starts from another job's last frame waits for that job to finish.
|
||||
* All-or-nothing: if any job can't be queued, none are.
|
||||
*/
|
||||
export async function queueJobs(ids: string[]): Promise<Job[]> {
|
||||
const baseApiUrl = getApiBaseUrl();
|
||||
const response = await fetch(`${baseApiUrl}/jobs/queue`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ job_ids: ids }),
|
||||
});
|
||||
if (!response.ok) {
|
||||
const error = await response
|
||||
.json()
|
||||
.catch(() => ({ detail: "Failed to queue jobs" }));
|
||||
throw new Error(error.detail || "Failed to queue jobs");
|
||||
}
|
||||
return response.json();
|
||||
}
|
||||
|
||||
/** Take a queued job out of the queue; it goes back to pending. */
|
||||
export async function dequeueJob(id: string): Promise<Job> {
|
||||
const baseApiUrl = getApiBaseUrl();
|
||||
const response = await fetch(`${baseApiUrl}/jobs/${id}/dequeue`, {
|
||||
method: "POST",
|
||||
});
|
||||
if (!response.ok) {
|
||||
const error = await response
|
||||
.json()
|
||||
.catch(() => ({ detail: "Failed to remove job from the queue" }));
|
||||
throw new Error(error.detail || "Failed to remove job from the queue");
|
||||
}
|
||||
return response.json();
|
||||
}
|
||||
|
||||
export interface MergedScene {
|
||||
/** file name on the server; pass to getMergedVideoUrl / downloadMergedVideo */
|
||||
filename: string;
|
||||
clips: number;
|
||||
/** length of the merged video */
|
||||
seconds: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Join the finished clips of a scene, in the order given, into one video. Fails,
|
||||
* making nothing, if any clip hasn't finished.
|
||||
*/
|
||||
export async function mergeScene(jobIds: string[], name: string): Promise<MergedScene> {
|
||||
const baseApiUrl = getApiBaseUrl();
|
||||
const response = await fetch(`${baseApiUrl}/scenes/merge`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ job_ids: jobIds, name }),
|
||||
});
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({ detail: "Failed to merge the scene" }));
|
||||
throw new Error(error.detail || "Failed to merge the scene");
|
||||
}
|
||||
return response.json();
|
||||
}
|
||||
|
||||
export function getMergedVideoUrl(filename: string): string {
|
||||
return `${getApiBaseUrl()}/merged/${encodeURIComponent(filename)}`;
|
||||
}
|
||||
|
||||
export async function downloadMergedVideo(filename: string): Promise<Blob> {
|
||||
const response = await fetch(getMergedVideoUrl(filename));
|
||||
if (!response.ok) {
|
||||
throw new Error("Failed to download the merged video");
|
||||
}
|
||||
return response.blob();
|
||||
}
|
||||
|
||||
export async function stopJob(id: string): Promise<Job> {
|
||||
const baseApiUrl = getApiBaseUrl();
|
||||
const response = await fetch(`${baseApiUrl}/jobs/${id}/stop`, {
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
attachLastFrame,
|
||||
attachedLastFrames,
|
||||
deferredLastFrameSource,
|
||||
detachLastFrame,
|
||||
lastFrameSourceJobId,
|
||||
whyCannotAttach,
|
||||
} from "@/lib/continuation";
|
||||
import { parseH3Prompt } from "@/lib/h3Prompt";
|
||||
import { readField } from "@/lib/h3PromptFields";
|
||||
|
||||
const FRAME = "/data/uploads/last_frames/last_frame_job-5.png";
|
||||
const REFS = [
|
||||
{ source: "/refs/veteran.png", media_type: "image" },
|
||||
{ source: "/refs/young.png", media_type: "image" },
|
||||
];
|
||||
|
||||
const CLIP = [
|
||||
"subject_definitions:",
|
||||
"<Subject 1> is the veteran broker from <Picture 1>: a confident man.",
|
||||
"<Subject 2> is the young broker from <Picture 2>: eager.",
|
||||
"",
|
||||
"summary:",
|
||||
"[reference generation] The target video shows <Subject 1> and <Subject 2> at a lunch.",
|
||||
"",
|
||||
"retention_analysis:",
|
||||
"<Subject 1> (appears in [Shot 1]): fully_preserved - his suit.",
|
||||
"<Subject 2> (appears in [Shot 1]): fully_preserved - his suit.",
|
||||
"",
|
||||
"detailed_description:",
|
||||
"[Shot 1] Close-up shot, dolly. Continuing the same shot on <Subject 1>, still speaking. (S1) <d>[English] And then, well.</d>",
|
||||
"[Shot 2] At 00:04.000, the shot cuts to a wide shot. Both men at the table.",
|
||||
"",
|
||||
"overall_soundscape:",
|
||||
"Muted restaurant ambience.",
|
||||
"",
|
||||
"non_diegetic_music:",
|
||||
"N/A",
|
||||
].join("\n");
|
||||
|
||||
function attach(prompt = CLIP, refs = REFS, frame = FRAME) {
|
||||
const r = attachLastFrame(prompt, refs, frame);
|
||||
if (!r.ok) throw new Error(r.error);
|
||||
return r;
|
||||
}
|
||||
|
||||
describe("attachLastFrame", () => {
|
||||
it("adds the frame as the next image reference and names it by its picture number", () => {
|
||||
const r = attach();
|
||||
expect(r.label).toBe("<Picture 3>");
|
||||
expect(r.references).toEqual([...REFS, { source: FRAME, media_type: "image" }]);
|
||||
});
|
||||
|
||||
it("defines the frame as the first frame of Shot 1, after the existing definitions", () => {
|
||||
const defs = readField(attach().prompt, "subject_definitions") as string;
|
||||
expect(defs.split("\n")).toEqual([
|
||||
"<Subject 1> is the veteran broker from <Picture 1>: a confident man.",
|
||||
"<Subject 2> is the young broker from <Picture 2>: eager.",
|
||||
"<Picture 3> is the first frame of [Shot 1]: the last frame of the previous clip, showing the scene, positions, and lighting exactly as that clip ended.",
|
||||
]);
|
||||
});
|
||||
|
||||
it("tags the summary as keyframe completion alongside its existing task type, and names the frame", () => {
|
||||
const summary = readField(attach().prompt, "summary") as string;
|
||||
expect(summary.startsWith("[reference generation + keyframe completion] The target video shows")).toBe(true);
|
||||
expect(summary.endsWith("The target video starts from <Picture 3>, the last frame of the previous clip.")).toBe(true);
|
||||
});
|
||||
|
||||
it("adds a first-frame line to the retention analysis", () => {
|
||||
const lines = (readField(attach().prompt, "retention_analysis") as string).split("\n");
|
||||
expect(lines).toHaveLength(3);
|
||||
expect(lines[2]).toBe(
|
||||
"<Picture 3> ([Shot 1] first frame): fully_preserved - the scene, positions, and lighting of the previous clip's last frame are retained.",
|
||||
);
|
||||
});
|
||||
|
||||
it("starts the first shot from the frame and leaves the other shots alone", () => {
|
||||
const shots = (readField(attach().prompt, "detailed_description") as string).split("\n");
|
||||
expect(shots[0]).toBe(
|
||||
"[Shot 1] Close-up shot, dolly. The shot begins from <Picture 3>. Continuing the same shot on <Subject 1>, still speaking. (S1) <d>[English] And then, well.</d>",
|
||||
);
|
||||
expect(shots[1]).toBe("[Shot 2] At 00:04.000, the shot cuts to a wide shot. Both men at the table.");
|
||||
});
|
||||
|
||||
it("leaves every other section untouched", () => {
|
||||
const out = attach().prompt;
|
||||
expect(readField(out, "overall_soundscape")).toBe("Muted restaurant ambience.");
|
||||
expect(readField(out, "non_diegetic_music")).toBe("N/A");
|
||||
});
|
||||
|
||||
it("still reads as a six-section prompt", () => {
|
||||
expect(parseH3Prompt(attach().prompt)).not.toBeNull();
|
||||
});
|
||||
|
||||
it("counts only images when numbering, so a video reference doesn't shift the label", () => {
|
||||
const refs = [REFS[0], { source: "/refs/clip.mp4", media_type: "video" }, REFS[1]];
|
||||
expect(attach(CLIP, refs).label).toBe("<Picture 3>");
|
||||
expect(attach(CLIP, [{ source: "/refs/clip.mp4", media_type: "video" }]).label).toBe("<Picture 1>");
|
||||
});
|
||||
|
||||
it("does not mutate what it was given", () => {
|
||||
const refs = [...REFS];
|
||||
attach(CLIP, refs);
|
||||
expect(refs).toEqual(REFS);
|
||||
});
|
||||
});
|
||||
|
||||
describe("attachLastFrame: summary wording", () => {
|
||||
const summaryOf = (summary: string) =>
|
||||
readField(attach(CLIP.replace("[reference generation] The target video shows <Subject 1> and <Subject 2> at a lunch.", summary)).prompt, "summary") as string;
|
||||
|
||||
it("adds the tag when the summary has none", () => {
|
||||
expect(summaryOf("Two men at a lunch.").startsWith("[keyframe completion] Two men at a lunch.")).toBe(true);
|
||||
});
|
||||
|
||||
it("extends a combined tag", () => {
|
||||
expect(summaryOf("[video editing + reference generation] Edit.").startsWith("[video editing + reference generation + keyframe completion] Edit.")).toBe(true);
|
||||
});
|
||||
|
||||
it("does not repeat a tag that is already there", () => {
|
||||
const s = summaryOf("[keyframe completion + reference generation] Both.");
|
||||
expect(s.startsWith("[keyframe completion + reference generation] Both.")).toBe(true);
|
||||
expect(s.match(/keyframe completion/g)).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe("attachLastFrame: other prompt shapes", () => {
|
||||
it("replaces an N/A subject_definitions instead of appending under it", () => {
|
||||
const out = attach(CLIP.replace(/subject_definitions:\n[\s\S]*?\n\nsummary:/, "subject_definitions:\nN/A\n\nsummary:")).prompt;
|
||||
expect(readField(out, "subject_definitions")).toBe(
|
||||
"<Picture 3> is the first frame of [Shot 1]: the last frame of the previous clip, showing the scene, positions, and lighting exactly as that clip ended.",
|
||||
);
|
||||
});
|
||||
|
||||
it("keeps hand-written shots verbatim, only prefixing the first", () => {
|
||||
const hand = CLIP.replace(
|
||||
/\[Shot 1\][^\n]*\n\[Shot 2\][^\n]*/,
|
||||
"[Shot 1] A handheld medium shot holds <Subject 1> at frame left.\n[Shot 2] At 00:03.800, the shot cuts to a low tracking shot behind him.",
|
||||
);
|
||||
const shots = (readField(attach(hand).prompt, "detailed_description") as string).split("\n");
|
||||
expect(shots).toEqual([
|
||||
"[Shot 1] The shot begins from <Picture 3>. A handheld medium shot holds <Subject 1> at frame left.",
|
||||
"[Shot 2] At 00:03.800, the shot cuts to a low tracking shot behind him.",
|
||||
]);
|
||||
});
|
||||
|
||||
it("adds the sentence to a first shot that is only dialogue", () => {
|
||||
const dialogueOnly = CLIP.replace(
|
||||
"[Shot 1] Close-up shot, dolly. Continuing the same shot on <Subject 1>, still speaking. (S1) <d>[English] And then, well.</d>",
|
||||
"[Shot 1] Wide shot. (S1) <d>[English] Hi.</d>",
|
||||
);
|
||||
const first = (readField(attach(dialogueOnly).prompt, "detailed_description") as string).split("\n")[0];
|
||||
expect(first).toBe("[Shot 1] Wide shot. The shot begins from <Picture 3>. (S1) <d>[English] Hi.</d>");
|
||||
});
|
||||
});
|
||||
|
||||
describe("attachLastFrame: refusals", () => {
|
||||
it("won't attach the same frame twice", () => {
|
||||
const r = attachLastFrame(CLIP, [...REFS, { source: FRAME, media_type: "image" }], FRAME);
|
||||
expect(r).toEqual({ ok: false, error: "That frame is already attached to this clip." });
|
||||
});
|
||||
|
||||
it("won't rewrite a prompt that isn't in the six-section format", () => {
|
||||
const r = attachLastFrame("A toy car drives into a plush dog.", REFS, FRAME);
|
||||
expect(r.ok).toBe(false);
|
||||
expect(!r.ok && r.error).toMatch(/six-section/);
|
||||
});
|
||||
|
||||
it("won't rewrite a base-format prompt either", () => {
|
||||
const base = "For the target video, <Picture 1> is used.\n\nintegrated_multimodal_description: [Shot 1] A.\n\noverall_soundscape: Quiet.";
|
||||
expect(whyCannotAttach(base, REFS)).toMatch(/six-section/);
|
||||
});
|
||||
|
||||
it("stops at nine image references", () => {
|
||||
const nine = Array.from({ length: 9 }, (_, i) => ({ source: `/r/${i}.png`, media_type: "image" }));
|
||||
const r = attachLastFrame(CLIP, nine, FRAME);
|
||||
expect(!r.ok && r.error).toMatch(/at most 9 image/);
|
||||
});
|
||||
|
||||
it("reports a prompt with no shots", () => {
|
||||
const empty = CLIP.replace(/detailed_description:\n[\s\S]*?\n\noverall_soundscape:/, "detailed_description:\n\n\noverall_soundscape:");
|
||||
const r = attachLastFrame(empty, REFS, FRAME);
|
||||
expect(r.ok).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("recognizing an attached last frame", () => {
|
||||
it("reads the job id out of the saved path", () => {
|
||||
expect(lastFrameSourceJobId(FRAME)).toBe("job-5");
|
||||
expect(lastFrameSourceJobId("last_frame_abc-123.png")).toBe("abc-123");
|
||||
expect(lastFrameSourceJobId("/refs/veteran.png")).toBeNull();
|
||||
expect(lastFrameSourceJobId("/x/not_last_frame_job-5.png")).toBeNull();
|
||||
});
|
||||
|
||||
it("finds last-frame references among a job's references", () => {
|
||||
expect(attachedLastFrames([...REFS, { source: FRAME, media_type: "image" }])).toEqual([{ source: FRAME, jobId: "job-5" }]);
|
||||
expect(attachedLastFrames(REFS)).toEqual([]);
|
||||
expect(attachedLastFrames(null)).toEqual([]);
|
||||
});
|
||||
|
||||
it("reads the job id out of a reference the server resolves when the job runs", () => {
|
||||
expect(deferredLastFrameSource("job-5")).toBe("job-last-frame:job-5");
|
||||
expect(lastFrameSourceJobId(deferredLastFrameSource("job-5"))).toBe("job-5");
|
||||
expect(lastFrameSourceJobId("job-last-frame:")).toBeNull();
|
||||
expect(attachedLastFrames([...REFS, { source: deferredLastFrameSource("job-5"), media_type: "image" }])).toEqual([
|
||||
{ source: "job-last-frame:job-5", jobId: "job-5" },
|
||||
]);
|
||||
});
|
||||
|
||||
it("attaches a deferred frame exactly like a saved one", () => {
|
||||
const result = attachLastFrame(CLIP, REFS, deferredLastFrameSource("job-5"));
|
||||
expect(result.ok).toBe(true);
|
||||
if (result.ok) {
|
||||
expect(result.references.at(-1)).toEqual({ source: "job-last-frame:job-5", media_type: "image" });
|
||||
expect(result.label).toBe("<Picture 3>");
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("detachLastFrame", () => {
|
||||
const roundTrip = (prompt: string, frame = FRAME) => {
|
||||
const attached = attach(prompt, REFS, frame);
|
||||
const r = detachLastFrame(attached.prompt, attached.references, frame);
|
||||
if (!r.ok) throw new Error(r.error);
|
||||
return r;
|
||||
};
|
||||
|
||||
it("takes back exactly what attaching added", () => {
|
||||
const r = roundTrip(CLIP);
|
||||
expect(r.references).toEqual(REFS);
|
||||
expect(r.prompt).toBe(CLIP);
|
||||
});
|
||||
|
||||
it("works for a frame the server resolves when the job runs", () => {
|
||||
expect(roundTrip(CLIP, deferredLastFrameSource("job-5")).prompt).toBe(CLIP);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["a summary with no tag", "Two men at a lunch."],
|
||||
["a combined tag", "[video editing + reference generation] Edit."],
|
||||
])("restores %s", (_name, summary) => {
|
||||
const prompt = CLIP.replace("[reference generation] The target video shows <Subject 1> and <Subject 2> at a lunch.", summary);
|
||||
expect(roundTrip(prompt).prompt).toBe(prompt);
|
||||
});
|
||||
|
||||
it("restores hand-written shots and an N/A section", () => {
|
||||
const hand = CLIP.replace(/\[Shot 1\][^\n]*\n\[Shot 2\][^\n]*/, "[Shot 1] A handheld medium shot.\n[Shot 2] At 00:03.800, the shot cuts to a low tracking shot.").replace(
|
||||
/subject_definitions:\n[\s\S]*?\n\nsummary:/,
|
||||
"subject_definitions:\nN/A\n\nsummary:",
|
||||
);
|
||||
expect(roundTrip(hand).prompt).toBe(hand);
|
||||
});
|
||||
|
||||
it("removes the sentence from a first shot that is only dialogue", () => {
|
||||
const dialogueOnly = CLIP.replace(
|
||||
"[Shot 1] Close-up shot, dolly. Continuing the same shot on <Subject 1>, still speaking. (S1) <d>[English] And then, well.</d>",
|
||||
"[Shot 1] Wide shot. (S1) <d>[English] Hi.</d>",
|
||||
);
|
||||
expect(roundTrip(dialogueOnly).prompt).toBe(dialogueOnly);
|
||||
});
|
||||
|
||||
it("only removes the reference from a prompt that was never rewritten", () => {
|
||||
const r = detachLastFrame("A plain prompt.", [...REFS, { source: FRAME, media_type: "image" }], FRAME);
|
||||
expect(r).toEqual({ ok: true, prompt: "A plain prompt.", references: REFS });
|
||||
});
|
||||
|
||||
it("refuses when other images were added after the frame", () => {
|
||||
const attached = attach();
|
||||
const r = detachLastFrame(attached.prompt, [...attached.references, { source: "/refs/extra.png", media_type: "image" }], FRAME);
|
||||
expect(r.ok).toBe(false);
|
||||
});
|
||||
|
||||
it("refuses a frame that isn't attached", () => {
|
||||
expect(detachLastFrame(CLIP, REFS, FRAME)).toEqual({ ok: false, error: "That frame isn't attached to this clip." });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,191 @@
|
||||
/**
|
||||
* Start a clip from the last frame of the one before it.
|
||||
*
|
||||
* A scene is generated as many short clips. To make one pick up where the
|
||||
* previous ended, the previous clip's last frame is attached as an image
|
||||
* reference, and the prompt is rewritten the way the prompt guide describes for
|
||||
* a picture that anchors the opening frame: a `[keyframe completion]` task tag,
|
||||
* a standalone `<Picture N>` definition, a retention line, and "the shot begins
|
||||
* from <Picture N>" in the first shot. The wording is kept together here so it
|
||||
* is easy to change if the guide is read differently.
|
||||
*/
|
||||
import { readField, writeField } from "@/lib/h3PromptFields";
|
||||
import { H3_REFERENCE_LIMITS } from "@/lib/h3References";
|
||||
import {
|
||||
detailedDescriptionFromShots,
|
||||
newTextPart,
|
||||
parseDetailedDescription,
|
||||
} from "@/lib/h3Shots";
|
||||
|
||||
export interface Reference {
|
||||
source: string;
|
||||
media_type: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* A reference that means "the last frame of that job's video", resolved by the
|
||||
* server when this job runs -- so it can be attached before that job has
|
||||
* finished, and the server's queue knows this job must wait for it.
|
||||
*/
|
||||
export const DEFERRED_LAST_FRAME_PREFIX = "job-last-frame:";
|
||||
|
||||
export const deferredLastFrameSource = (jobId: string) => `${DEFERRED_LAST_FRAME_PREFIX}${jobId}`;
|
||||
|
||||
/** Saved by the server as `.../last_frames/last_frame_<job id>.png`. */
|
||||
const LAST_FRAME_RE = /(?:^|\/)last_frame_([^/]+)\.png$/;
|
||||
|
||||
const REQUIRED_SECTIONS = [
|
||||
"subject_definitions",
|
||||
"summary",
|
||||
"retention_analysis",
|
||||
"detailed_description",
|
||||
] as const;
|
||||
|
||||
/** The job a last-frame reference came from, or null if `source` isn't one. */
|
||||
export function lastFrameSourceJobId(source: string): string | null {
|
||||
if (source.startsWith(DEFERRED_LAST_FRAME_PREFIX)) return source.slice(DEFERRED_LAST_FRAME_PREFIX.length) || null;
|
||||
return source.match(LAST_FRAME_RE)?.[1] ?? null;
|
||||
}
|
||||
|
||||
/** Last-frame references already on a job, with the job each came from. */
|
||||
export function attachedLastFrames(references: Reference[] | null | undefined): { source: string; jobId: string }[] {
|
||||
return (references ?? []).flatMap((r) => {
|
||||
const jobId = lastFrameSourceJobId(r.source);
|
||||
return jobId ? [{ source: r.source, jobId }] : [];
|
||||
});
|
||||
}
|
||||
|
||||
/** Why a last frame can't be attached to this job, or null if it can. */
|
||||
export function whyCannotAttach(prompt: string, references: Reference[] | null | undefined): string | null {
|
||||
const missing = REQUIRED_SECTIONS.filter((s) => readField(prompt, s) === null);
|
||||
if (missing.length > 0) {
|
||||
return "This prompt isn't in the six-section format, so it can't be rewritten automatically. Attach the frame in Edit instead.";
|
||||
}
|
||||
const images = (references ?? []).filter((r) => r.media_type === "image").length;
|
||||
if (images >= H3_REFERENCE_LIMITS.image) {
|
||||
return `A job can have at most ${H3_REFERENCE_LIMITS.image} image references.`;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** "[reference generation]" -> "[reference generation + keyframe completion]"; no tag -> "[keyframe completion] ...". */
|
||||
function withKeyframeCompletion(summary: string): string {
|
||||
const tag = summary.match(/^\[([^\]]*)\]\s*/);
|
||||
if (!tag) return `[keyframe completion] ${summary}`.trim();
|
||||
const types = tag[1].split("+").map((t) => t.trim()).filter(Boolean);
|
||||
if (types.includes("keyframe completion")) return summary;
|
||||
return `[${[...types, "keyframe completion"].join(" + ")}] ${summary.slice(tag[0].length)}`.trim();
|
||||
}
|
||||
|
||||
const isBlank = (text: string) => !text.trim() || text.trim().toUpperCase() === "N/A";
|
||||
const appendLine = (body: string, line: string) => (isBlank(body) ? line : `${body}\n${line}`);
|
||||
|
||||
export type AttachResult =
|
||||
| { ok: true; prompt: string; references: Reference[]; label: string }
|
||||
| { ok: false; error: string };
|
||||
|
||||
/**
|
||||
* Add `framePath` as the next image reference and rewrite the prompt to use it
|
||||
* as the opening frame of Shot 1. Fails, changing nothing, if the frame is
|
||||
* already attached or the prompt isn't in a form that can be rewritten safely.
|
||||
*/
|
||||
export function attachLastFrame(prompt: string, references: Reference[], framePath: string): AttachResult {
|
||||
if (references.some((r) => r.source === framePath)) {
|
||||
return { ok: false, error: "That frame is already attached to this clip." };
|
||||
}
|
||||
const blocked = whyCannotAttach(prompt, references);
|
||||
if (blocked) return { ok: false, error: blocked };
|
||||
|
||||
// Labels count per media type in reference order, so the new image is the next <Picture N>.
|
||||
const label = `<Picture ${references.filter((r) => r.media_type === "image").length + 1}>`;
|
||||
|
||||
const { preamble, shots } = parseDetailedDescription(readField(prompt, "detailed_description") ?? "", []);
|
||||
if (shots.length === 0) return { ok: false, error: "This prompt has no shots to start from the frame." };
|
||||
const first = shots[0];
|
||||
const firstText = first.body.find((p) => p.kind === "text");
|
||||
if (firstText && firstText.kind === "text") firstText.text = `The shot begins from ${label}. ${firstText.text}`;
|
||||
else first.body.unshift(newTextPart(`The shot begins from ${label}.`));
|
||||
|
||||
let next = prompt;
|
||||
next = writeField(
|
||||
next,
|
||||
"subject_definitions",
|
||||
appendLine(
|
||||
readField(next, "subject_definitions") ?? "",
|
||||
`${label} is the first frame of [Shot 1]: the last frame of the previous clip, showing the scene, positions, and lighting exactly as that clip ended.`,
|
||||
),
|
||||
);
|
||||
next = writeField(
|
||||
next,
|
||||
"summary",
|
||||
`${withKeyframeCompletion(readField(next, "summary") ?? "")} The target video starts from ${label}, the last frame of the previous clip.`,
|
||||
);
|
||||
next = writeField(
|
||||
next,
|
||||
"retention_analysis",
|
||||
appendLine(
|
||||
readField(next, "retention_analysis") ?? "",
|
||||
`${label} ([Shot 1] first frame): fully_preserved - the scene, positions, and lighting of the previous clip's last frame are retained.`,
|
||||
),
|
||||
);
|
||||
next = writeField(next, "detailed_description", detailedDescriptionFromShots(shots, preamble));
|
||||
|
||||
return { ok: true, prompt: next, references: [...references, { source: framePath, media_type: "image" }], label };
|
||||
}
|
||||
|
||||
/** Drop the lines `attachLastFrame` added, leaving "N/A" if nothing else was there. */
|
||||
function removeLine(body: string, startsWith: string): string {
|
||||
const kept = body.split("\n").filter((l) => !l.trim().startsWith(startsWith));
|
||||
const text = kept.join("\n").trim();
|
||||
return text || "N/A";
|
||||
}
|
||||
|
||||
export type DetachResult = { ok: true; prompt: string; references: Reference[] } | { ok: false; error: string };
|
||||
|
||||
/**
|
||||
* Undo `attachLastFrame`: remove the reference at `source` and take back what was
|
||||
* added to the prompt for it. Only the most recently added image can be removed --
|
||||
* removing an earlier one would renumber the pictures after it.
|
||||
*/
|
||||
export function detachLastFrame(prompt: string, references: Reference[], source: string): DetachResult {
|
||||
const images = references.filter((r) => r.media_type === "image");
|
||||
const index = images.findIndex((r) => r.source === source);
|
||||
if (index < 0) return { ok: false, error: "That frame isn't attached to this clip." };
|
||||
if (index !== images.length - 1) {
|
||||
return { ok: false, error: "Other images were added after that frame. Remove it in Edit instead." };
|
||||
}
|
||||
const remaining = references.filter((r) => r.source !== source);
|
||||
|
||||
// A prompt that isn't in the six-section form was never rewritten; only the reference goes.
|
||||
if (REQUIRED_SECTIONS.some((s) => readField(prompt, s) === null)) {
|
||||
return { ok: true, prompt, references: remaining };
|
||||
}
|
||||
|
||||
const label = `<Picture ${index + 1}>`;
|
||||
let next = prompt;
|
||||
next = writeField(next, "subject_definitions", removeLine(readField(next, "subject_definitions") ?? "", `${label} is the first frame of [Shot 1]`));
|
||||
next = writeField(next, "retention_analysis", removeLine(readField(next, "retention_analysis") ?? "", `${label} ([Shot 1] first frame)`));
|
||||
|
||||
const summary = readField(next, "summary") ?? "";
|
||||
const sentence = ` The target video starts from ${label}, the last frame of the previous clip.`;
|
||||
if (summary.includes(sentence)) {
|
||||
let restored = summary.replace(sentence, "");
|
||||
const tag = restored.match(/^\[([^\]]*)\]\s*/);
|
||||
if (tag) {
|
||||
const types = tag[1].split("+").map((t) => t.trim()).filter((t) => t && t !== "keyframe completion");
|
||||
restored = types.length > 0 ? `[${types.join(" + ")}] ${restored.slice(tag[0].length)}` : restored.slice(tag[0].length);
|
||||
}
|
||||
next = writeField(next, "summary", restored.trim());
|
||||
}
|
||||
|
||||
const { preamble, shots } = parseDetailedDescription(readField(next, "detailed_description") ?? "", []);
|
||||
const lead = `The shot begins from ${label}.`;
|
||||
const first = shots[0];
|
||||
const part = first?.body.find((p) => p.kind === "text" && p.text.startsWith(lead));
|
||||
if (first && part && part.kind === "text") {
|
||||
part.text = part.text.slice(lead.length).trimStart();
|
||||
if (!part.text) first.body = first.body.filter((p) => p !== part);
|
||||
next = writeField(next, "detailed_description", detailedDescriptionFromShots(shots, preamble));
|
||||
}
|
||||
return { ok: true, prompt: next, references: remaining };
|
||||
}
|
||||
@@ -94,6 +94,21 @@ describe("buildClips", () => {
|
||||
expect(first.subjectIds).toEqual(["<Subject 1>", "<Subject 2>"]);
|
||||
});
|
||||
|
||||
it("never adds camera movement, however long the shot runs", () => {
|
||||
const long = "We keep the clients on the wheel and the wheel keeps turning, every day of the week. ".repeat(4).trim();
|
||||
const clips = buildClips(
|
||||
[
|
||||
{ speaker: "CAPTAIN", text: long },
|
||||
{ speaker: "CLERK", text: long },
|
||||
{ speaker: "CAPTAIN", text: "Yes." },
|
||||
],
|
||||
{ speakers: SPEAKERS, targetSeconds: 60 },
|
||||
);
|
||||
const shots = clips.flatMap((c) => c.shots);
|
||||
expect(shots.length).toBeGreaterThan(2);
|
||||
expect(shots.every((s) => s.movement === "static")).toBe(true);
|
||||
});
|
||||
|
||||
it("uses a close-up when only one person speaks in a shot", () => {
|
||||
const [clip] = buildClips([{ speaker: "CAPTAIN", text: "Nobody leaves until the count is done." }], {
|
||||
speakers: SPEAKERS,
|
||||
|
||||
@@ -205,13 +205,7 @@ function joinNatural(items: string[]): string {
|
||||
return items.length <= 2 ? items.join(" and ") : `${items.slice(0, -1).join(", ")}, and ${items[items.length - 1]}`;
|
||||
}
|
||||
|
||||
function shotFromBeat(
|
||||
beat: Beat,
|
||||
opts: ImportOptions,
|
||||
totalSeconds: number,
|
||||
continued = false,
|
||||
truncated = false,
|
||||
): Shot {
|
||||
function shotFromBeat(beat: Beat, opts: ImportOptions, continued = false, truncated = false): Shot {
|
||||
const speakers = [...new Set(beat.lines.map((l) => l.speaker))];
|
||||
const subjectIds = [
|
||||
...new Set(speakers.map((s) => opts.speakers[s]?.subject).filter((s): s is string => Boolean(s))),
|
||||
@@ -231,7 +225,9 @@ function shotFromBeat(
|
||||
|
||||
const shot = makeShot({
|
||||
type: speakers.length > 1 ? "medium" : "close-up",
|
||||
movement: totalSeconds >= 5 ? "dolly" : "static",
|
||||
// A static camera: added movement (a dolly on every long shot) showed up as
|
||||
// unwanted panning. Movement can still be set per shot in the shot editor.
|
||||
movement: "static",
|
||||
subjectIds,
|
||||
description,
|
||||
lines: beat.lines.map((l) => newLine({ speaker: opts.speakers[l.speaker]?.tag ?? "", text: l.text })),
|
||||
@@ -308,7 +304,7 @@ export function buildClips(turns: ScriptTurn[], opts: ImportOptions): ImportClip
|
||||
|
||||
if (total <= opts.targetSeconds) {
|
||||
if (hasDialogue() && seconds + total > opts.targetSeconds) close(false);
|
||||
current.shots.push(shotFromBeat(beat, opts, total));
|
||||
current.shots.push(shotFromBeat(beat, opts));
|
||||
seconds += total;
|
||||
continue;
|
||||
}
|
||||
@@ -321,7 +317,7 @@ export function buildClips(turns: ScriptTurn[], opts: ImportOptions): ImportClip
|
||||
if (k > 0) close(true);
|
||||
const lines = [{ ...main, text: part }, ...(k === parts.length - 1 ? reactions : [])];
|
||||
const pieceBeat: Beat = { lines, directions: k === 0 ? beat.directions : [] };
|
||||
current.shots.push(shotFromBeat(pieceBeat, opts, total, k > 0, k < parts.length - 1));
|
||||
current.shots.push(shotFromBeat(pieceBeat, opts, k > 0, k < parts.length - 1));
|
||||
seconds += beatSeconds(pieceBeat);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -71,4 +71,8 @@ describe("referenceFileName", () => {
|
||||
expect(referenceFileName("/a/b/c.mp4")).toBe("c.mp4");
|
||||
expect(referenceFileName("c.mp4")).toBe("c.mp4");
|
||||
});
|
||||
|
||||
it("names a last-frame reference by the job it comes from, not its internal id", () => {
|
||||
expect(referenceFileName("job-last-frame:0a1b2c3d-1111-2222-3333-444455556666")).toBe("Last frame of job 0a1b2c3d");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -72,6 +72,7 @@ export interface JobFormFields {
|
||||
|
||||
/** Uploads keep their original basename, so this is the display name. */
|
||||
export function referenceFileName(source: string): string {
|
||||
if (source.startsWith("job-last-frame:")) return `Last frame of job ${source.slice(15, 23)}`;
|
||||
return source.split("/").filter(Boolean).pop() ?? source;
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { deferredLastFrameSource } from '@/lib/continuation';
|
||||
import { planSceneQueue } from '@/lib/sceneQueue';
|
||||
import { buildScenes, type SceneClip } from '@/lib/scenes';
|
||||
import { makeJob } from '@/test/factories';
|
||||
|
||||
const SIX_SECTION = [
|
||||
'subject_definitions:',
|
||||
'<Subject 1> is a man.',
|
||||
'',
|
||||
'summary:',
|
||||
'[reference generation] A man.',
|
||||
'',
|
||||
'retention_analysis:',
|
||||
'<Subject 1> (appears in [Shot 1]): fully_preserved - suit.',
|
||||
'',
|
||||
'detailed_description:',
|
||||
'[Shot 1] Wide shot. A man.',
|
||||
'',
|
||||
'overall_soundscape:',
|
||||
'Quiet.',
|
||||
'',
|
||||
'non_diegetic_music:',
|
||||
'N/A',
|
||||
].join('\n');
|
||||
|
||||
const clips = (...overrides: Parameters<typeof makeJob>[0][]): SceneClip[] => {
|
||||
const jobs = overrides.map((over, i) =>
|
||||
makeJob({
|
||||
id: `c${i + 1}`,
|
||||
name: `s-clip-${i + 1}`,
|
||||
created_at: i + 1,
|
||||
status: 'pending',
|
||||
prompt: SIX_SECTION,
|
||||
...over,
|
||||
}),
|
||||
);
|
||||
return buildScenes(jobs)[0].clips;
|
||||
};
|
||||
|
||||
const ids = (list: { job: { id: string } }[]) => list.map((c) => c.job.id);
|
||||
|
||||
describe('planSceneQueue', () => {
|
||||
it('queues every clip that can run, in scene order', () => {
|
||||
const plan = planSceneQueue(clips({}, {}, {}), { chain: false });
|
||||
expect(ids(plan.queue)).toEqual(['c1', 'c2', 'c3']);
|
||||
expect(plan.skipped).toEqual([]);
|
||||
});
|
||||
|
||||
it.each([
|
||||
['completed', 'already finished'],
|
||||
['running', 'already running'],
|
||||
['queued', 'already queued'],
|
||||
])('leaves %s clips alone', (status, reason) => {
|
||||
const plan = planSceneQueue(clips({}, { status }, {}), { chain: false });
|
||||
expect(ids(plan.queue)).toEqual(['c1', 'c3']);
|
||||
expect(plan.skipped.map((s) => [s.clip.job.id, s.reason])).toEqual([['c2', reason]]);
|
||||
});
|
||||
|
||||
it('requeues failed and stopped clips', () => {
|
||||
const plan = planSceneQueue(clips({ status: 'failed' }, { status: 'stopped' }), { chain: false });
|
||||
expect(ids(plan.queue)).toEqual(['c1', 'c2']);
|
||||
});
|
||||
|
||||
describe('chaining clips through last frames', () => {
|
||||
it('links each clip after the first to the one before it', () => {
|
||||
const plan = planSceneQueue(clips({}, {}, {}), { chain: true });
|
||||
expect(plan.chain.map((c) => [c.clip.job.id, c.previous.job.id])).toEqual([
|
||||
['c2', 'c1'],
|
||||
['c3', 'c2'],
|
||||
]);
|
||||
});
|
||||
|
||||
it('links nothing when not asked to', () => {
|
||||
expect(planSceneQueue(clips({}, {}), { chain: false }).chain).toEqual([]);
|
||||
});
|
||||
|
||||
it('links a clip to a previous clip that has already finished', () => {
|
||||
const plan = planSceneQueue(clips({ status: 'completed' }, {}), { chain: true });
|
||||
expect(plan.chain.map((c) => c.previous.job.id)).toEqual(['c1']);
|
||||
expect(ids(plan.queue)).toEqual(['c2']);
|
||||
});
|
||||
|
||||
it('does not link clips that are not being queued', () => {
|
||||
const plan = planSceneQueue(clips({}, { status: 'completed' }, {}), { chain: true });
|
||||
expect(plan.chain.map((c) => c.clip.job.id)).toEqual(['c3']);
|
||||
});
|
||||
|
||||
it('keeps a frame the clip already opens on', () => {
|
||||
const own = [{ source: deferredLastFrameSource('c1'), media_type: 'image' }];
|
||||
const stale = [{ source: '/uploads/last_frames/last_frame_old-take.png', media_type: 'image' }];
|
||||
const plan = planSceneQueue(clips({}, { references: own }, { references: stale }), { chain: true });
|
||||
expect(plan.chain).toEqual([]);
|
||||
expect(plan.cannotChain).toEqual([]);
|
||||
});
|
||||
|
||||
it("reports clips whose prompt can't be rewritten instead of linking them", () => {
|
||||
const plan = planSceneQueue(clips({}, { prompt: 'A plain prompt.' }), { chain: true });
|
||||
expect(plan.chain).toEqual([]);
|
||||
expect(plan.cannotChain).toHaveLength(1);
|
||||
expect(plan.cannotChain[0].clip.job.id).toBe('c2');
|
||||
expect(plan.cannotChain[0].reason).toMatch(/six-section/);
|
||||
expect(ids(plan.queue)).toEqual(['c1', 'c2']); // still queued, just not linked
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* Work out what "queue this scene" would do, so the user can confirm it first.
|
||||
*
|
||||
* Clips are queued in scene order. The server runs them one after another and
|
||||
* holds a clip that starts from another clip's last frame until that clip has
|
||||
* finished -- so a whole scene can be queued at once and each clip still opens
|
||||
* on the frame the one before it ended on. This decides which clips are queued,
|
||||
* and (optionally) which get that link attached first.
|
||||
*/
|
||||
import { attachedLastFrames, whyCannotAttach } from "@/lib/continuation";
|
||||
import type { SceneClip } from "@/lib/scenes";
|
||||
|
||||
/** Statuses a job can be queued from (the same ones it can be started or edited in). */
|
||||
export const QUEUEABLE_STATUSES = ["pending", "failed", "stopped"];
|
||||
|
||||
export interface SceneQueuePlan {
|
||||
/** clips to queue, in order */
|
||||
queue: SceneClip[];
|
||||
/** clips that will get the previous clip's last frame attached first */
|
||||
chain: { clip: SceneClip; previous: SceneClip }[];
|
||||
/** clips that would need the link but can't have it attached automatically */
|
||||
cannotChain: { clip: SceneClip; reason: string }[];
|
||||
/** clips left alone, and why */
|
||||
skipped: { clip: SceneClip; reason: string }[];
|
||||
}
|
||||
|
||||
function skipReason(status: string): string {
|
||||
if (status === "completed") return "already finished";
|
||||
if (status === "running") return "already running";
|
||||
if (status === "queued") return "already queued";
|
||||
return `is ${status}`;
|
||||
}
|
||||
|
||||
export function planSceneQueue(clips: SceneClip[], { chain: wantChain }: { chain: boolean }): SceneQueuePlan {
|
||||
const plan: SceneQueuePlan = { queue: [], chain: [], cannotChain: [], skipped: [] };
|
||||
clips.forEach((clip, i) => {
|
||||
if (!QUEUEABLE_STATUSES.includes(clip.job.status)) {
|
||||
plan.skipped.push({ clip, reason: skipReason(clip.job.status) });
|
||||
return;
|
||||
}
|
||||
plan.queue.push(clip);
|
||||
|
||||
const previous = clips[i - 1];
|
||||
if (!wantChain || !previous) return;
|
||||
// A clip that already opens on a frame (of this or another take) keeps it.
|
||||
if (attachedLastFrames(clip.job.references).length > 0) return;
|
||||
const blocked = whyCannotAttach(clip.job.prompt, clip.job.references);
|
||||
if (blocked) plan.cannotChain.push({ clip, reason: blocked });
|
||||
else plan.chain.push({ clip, previous });
|
||||
});
|
||||
return plan;
|
||||
}
|
||||
@@ -11,6 +11,8 @@ export interface Job {
|
||||
workload_type?: string;
|
||||
status: string;
|
||||
created_at: number;
|
||||
/** when the job joined the queue; null unless it is queued */
|
||||
queued_at?: number | null;
|
||||
started_at: number | null;
|
||||
finished_at: number | null;
|
||||
error: string | null;
|
||||
@@ -20,6 +22,8 @@ export interface Job {
|
||||
num_frames: number;
|
||||
/** frames per second; the API sends it, older cached objects may not have it */
|
||||
fps?: number;
|
||||
/** reference media (image/video/audio) attached to the job, as saved on the server */
|
||||
references?: { source: string; media_type: string }[] | null;
|
||||
height: number;
|
||||
width: number;
|
||||
guidance_scale: number;
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Queued jobs run in order, one at a time, and a clip waits for the clip whose last frame it uses."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo_studio.job_queue import (
|
||||
circular_dependency,
|
||||
deferred_last_frame_source,
|
||||
dependencies,
|
||||
dependency_of,
|
||||
plan_dispatch,
|
||||
resolve_references,
|
||||
UnresolvedReference,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class J:
|
||||
id: str
|
||||
status: str = "pending"
|
||||
queued_at: float | None = None
|
||||
references: list[dict[str, Any]] = field(default_factory=list)
|
||||
name: str = ""
|
||||
|
||||
|
||||
def after(job_id: str) -> list[dict[str, str]]:
|
||||
return [{"source": deferred_last_frame_source(job_id), "media_type": "image"}]
|
||||
|
||||
|
||||
def queue(*jobs: J) -> dict[str, J]:
|
||||
for i, job in enumerate(jobs):
|
||||
if job.status == "queued" and job.queued_at is None:
|
||||
job.queued_at = float(i)
|
||||
return {j.id: j for j in jobs}
|
||||
|
||||
|
||||
def test_deferred_source_round_trips():
|
||||
assert dependency_of({"source": deferred_last_frame_source("abc")}) == "abc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ref", [None, {}, {"source": "/tmp/a.png"}, {"source": "job-last-frame:"}, {"source": 3}])
|
||||
def test_ordinary_references_have_no_dependency(ref):
|
||||
assert dependency_of(ref) is None
|
||||
|
||||
|
||||
def test_dependencies_keep_order_and_drop_repeats():
|
||||
refs = [*after("b"), {"source": "/x.png"}, *after("a"), *after("b")]
|
||||
assert dependencies(refs) == ["b", "a"]
|
||||
|
||||
|
||||
class TestDispatchOrder:
|
||||
def test_starts_the_oldest_queued_job_first(self):
|
||||
jobs = queue(J("late", "queued", 5.0), J("early", "queued", 1.0))
|
||||
assert plan_dispatch(jobs, 1).start == ["early"]
|
||||
|
||||
def test_runs_only_max_concurrent_including_running_jobs(self):
|
||||
jobs = queue(J("r", "running"), J("a", "queued"), J("b", "queued"))
|
||||
assert plan_dispatch(jobs, 1).start == []
|
||||
assert plan_dispatch(jobs, 2).start == ["a"]
|
||||
assert plan_dispatch(jobs, 3).start == ["a", "b"]
|
||||
|
||||
def test_ignores_jobs_that_are_not_queued(self):
|
||||
jobs = queue(J("p"), J("c", "completed"), J("f", "failed"), J("s", "stopped"))
|
||||
assert plan_dispatch(jobs, 4).start == []
|
||||
|
||||
def test_a_limit_below_one_still_runs_one_job(self):
|
||||
assert plan_dispatch(queue(J("a", "queued")), 0).start == ["a"]
|
||||
|
||||
def test_ties_are_broken_by_id_for_a_stable_order(self):
|
||||
jobs = {"b": J("b", "queued", 1.0), "a": J("a", "queued", 1.0)}
|
||||
assert plan_dispatch(jobs, 2).start == ["a", "b"]
|
||||
|
||||
|
||||
class TestWaitingOnAFrame:
|
||||
def test_waits_while_the_source_clip_is_running_or_queued(self):
|
||||
for source_status in ("running", "queued"):
|
||||
jobs = queue(J("one", source_status), J("two", "queued", references=after("one")))
|
||||
plan = plan_dispatch(jobs, 2)
|
||||
assert "two" not in plan.start and not plan.fail
|
||||
|
||||
def test_starts_once_the_source_clip_completed(self):
|
||||
jobs = queue(J("one", "completed"), J("two", "queued", references=after("one")))
|
||||
assert plan_dispatch(jobs, 1).start == ["two"]
|
||||
|
||||
def test_waiting_does_not_hold_up_later_jobs(self):
|
||||
jobs = queue(
|
||||
J("one", "running"),
|
||||
J("two", "queued", references=after("one")),
|
||||
J("free", "queued"),
|
||||
)
|
||||
assert plan_dispatch(jobs, 2).start == ["free"]
|
||||
|
||||
def test_a_source_that_is_only_pending_keeps_waiting(self):
|
||||
# The user may still queue it.
|
||||
jobs = queue(J("one", "pending"), J("two", "queued", references=after("one")))
|
||||
plan = plan_dispatch(jobs, 1)
|
||||
assert plan.start == [] and plan.fail == {}
|
||||
|
||||
def test_needs_every_source(self):
|
||||
refs = [*after("a"), *after("b")]
|
||||
jobs = queue(J("a", "completed"), J("b", "running"), J("c", "queued", references=refs))
|
||||
assert plan_dispatch(jobs, 2).start == []
|
||||
jobs["b"].status = "completed"
|
||||
assert plan_dispatch(jobs, 2).start == ["c"]
|
||||
|
||||
def test_a_chain_advances_one_clip_at_a_time(self):
|
||||
jobs = queue(
|
||||
J("c1", "queued"),
|
||||
J("c2", "queued", references=after("c1")),
|
||||
J("c3", "queued", references=after("c2")),
|
||||
)
|
||||
assert plan_dispatch(jobs, 5).start == ["c1"]
|
||||
jobs["c1"].status = "completed"
|
||||
assert plan_dispatch(jobs, 5).start == ["c2"]
|
||||
jobs["c2"].status = "completed"
|
||||
assert plan_dispatch(jobs, 5).start == ["c3"]
|
||||
|
||||
|
||||
class TestImpossibleJobs:
|
||||
@pytest.mark.parametrize(("status", "word"), [("failed", "failed"), ("stopped", "stopped")])
|
||||
def test_fails_when_the_source_failed_or_was_stopped(self, status, word):
|
||||
jobs = queue(J("one", status, name="Clip 1"), J("two", "queued", references=after("one")))
|
||||
plan = plan_dispatch(jobs, 1)
|
||||
assert plan.start == []
|
||||
assert "Clip 1" in plan.fail["two"] and word in plan.fail["two"]
|
||||
|
||||
def test_fails_when_the_source_was_deleted(self):
|
||||
jobs = queue(J("two", "queued", references=after("gone")))
|
||||
assert "deleted" in plan_dispatch(jobs, 1).fail["two"]
|
||||
|
||||
def test_failure_cascades_down_the_chain(self):
|
||||
jobs = queue(
|
||||
J("c1", "failed"),
|
||||
J("c2", "queued", references=after("c1")),
|
||||
J("c3", "queued", references=after("c2")),
|
||||
J("c4", "queued", references=after("c3")),
|
||||
J("other", "queued"),
|
||||
)
|
||||
plan = plan_dispatch(jobs, 3)
|
||||
assert set(plan.fail) == {"c2", "c3", "c4"}
|
||||
assert plan.start == ["other"]
|
||||
|
||||
def test_a_failed_job_frees_its_slot_for_the_next(self):
|
||||
jobs = queue(
|
||||
J("c1", "failed"),
|
||||
J("c2", "queued", references=after("c1")),
|
||||
J("next", "queued"),
|
||||
)
|
||||
assert plan_dispatch(jobs, 1).start == ["next"]
|
||||
|
||||
def test_planning_does_not_change_the_jobs(self):
|
||||
jobs = queue(J("c1", "failed"), J("c2", "queued", references=after("c1")))
|
||||
plan_dispatch(jobs, 1)
|
||||
assert jobs["c2"].status == "queued"
|
||||
|
||||
|
||||
class TestCircularDependency:
|
||||
def test_none_for_a_plain_chain(self):
|
||||
jobs = queue(J("a"), J("b", references=after("a")), J("c", references=after("b")))
|
||||
assert circular_dependency(jobs, "c") is None
|
||||
|
||||
def test_finds_a_loop_through_the_job(self):
|
||||
jobs = queue(J("a", references=after("c")), J("b", references=after("a")), J("c", references=after("b")))
|
||||
assert circular_dependency(jobs, "a") is not None
|
||||
|
||||
def test_a_job_that_waits_on_itself(self):
|
||||
assert circular_dependency(queue(J("a", references=after("a"))), "a") is not None
|
||||
|
||||
def test_a_loop_elsewhere_is_not_this_jobs_problem(self):
|
||||
jobs = queue(
|
||||
J("a", references=after("b")),
|
||||
J("b", references=after("a")),
|
||||
J("c", references=after("a")),
|
||||
)
|
||||
assert circular_dependency(jobs, "c") is None
|
||||
|
||||
|
||||
class TestResolveReferences:
|
||||
def _jobs(self, **statuses: str) -> dict[str, J]:
|
||||
return {k: J(k, v, name=f"Clip {k}") for k, v in statuses.items()}
|
||||
|
||||
def test_replaces_a_deferred_reference_with_the_extracted_frame(self):
|
||||
jobs = self._jobs(one="completed")
|
||||
stored = [{"source": "/ref.png", "media_type": "image"}, *after("one")]
|
||||
resolved = resolve_references(stored, jobs.get, lambda job: f"/frames/{job.id}.png")
|
||||
assert resolved == [
|
||||
{"source": "/ref.png", "media_type": "image"},
|
||||
{"source": "/frames/one.png", "media_type": "image"},
|
||||
]
|
||||
|
||||
def test_leaves_the_stored_references_untouched(self):
|
||||
jobs = self._jobs(one="completed")
|
||||
stored = after("one")
|
||||
resolve_references(stored, jobs.get, lambda job: "/f.png")
|
||||
assert stored == after("one")
|
||||
|
||||
def test_no_references(self):
|
||||
assert resolve_references(None, {}.get, lambda job: "") == []
|
||||
|
||||
def test_refuses_an_unfinished_source(self):
|
||||
jobs = self._jobs(one="running")
|
||||
with pytest.raises(UnresolvedReference, match="hasn't finished.*running"):
|
||||
resolve_references(after("one"), jobs.get, lambda job: "/f.png")
|
||||
|
||||
def test_refuses_a_deleted_source(self):
|
||||
with pytest.raises(UnresolvedReference, match="no longer exists"):
|
||||
resolve_references(after("gone"), {}.get, lambda job: "/f.png")
|
||||
|
||||
def test_passes_on_the_extraction_error(self):
|
||||
class Boom(Exception):
|
||||
detail = "The video has no frames."
|
||||
|
||||
def fail(job):
|
||||
raise Boom()
|
||||
|
||||
with pytest.raises(UnresolvedReference, match="no frames"):
|
||||
resolve_references(after("one"), self._jobs(one="completed").get, fail)
|
||||
@@ -0,0 +1,560 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""JobRunner's queue: order, concurrency, last-frame dependencies, and surviving a restart.
|
||||
|
||||
Runs without a GPU: `fastvideo` is stubbed (importing it needs a working triton
|
||||
driver) and jobs "run" by waiting on an event instead of generating video.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import importlib.util
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from fastvideo_studio.database import Database
|
||||
from fastvideo_studio.job_queue import deferred_last_frame_source
|
||||
|
||||
RUNNER_PATH = Path(__file__).resolve().parents[1] / "job_runner.py"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def job_runner_module(monkeypatch):
|
||||
class Manager:
|
||||
def shutdown(self):
|
||||
pass
|
||||
|
||||
stub_utils = types.ModuleType("fastvideo.utils")
|
||||
stub_utils.get_mp_context = lambda: types.SimpleNamespace(Manager=Manager)
|
||||
stub_pkg = types.ModuleType("fastvideo")
|
||||
stub_pkg.utils = stub_utils
|
||||
monkeypatch.setitem(sys.modules, "fastvideo", stub_pkg)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo.utils", stub_utils)
|
||||
|
||||
spec = importlib.util.spec_from_file_location("job_runner_under_test", RUNNER_PATH)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, module) # @dataclass looks its module up here
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class Harness:
|
||||
"""Builds runners whose jobs finish only when the test says so."""
|
||||
|
||||
def __init__(self, module, tmp_path, monkeypatch):
|
||||
self.module = module
|
||||
self.tmp_path = tmp_path
|
||||
self.gates: dict[str, threading.Event] = collections.defaultdict(threading.Event)
|
||||
self.outcomes: dict[str, object] = {}
|
||||
self.ran: list[str] = []
|
||||
self.runners = []
|
||||
self.thread_of: dict[str, threading.Thread] = {}
|
||||
self.closed = False
|
||||
|
||||
harness = self
|
||||
|
||||
def fake_run(runner_self, job):
|
||||
harness.thread_of[job.id] = threading.current_thread()
|
||||
harness.ran.append(job.id)
|
||||
if harness.closed:
|
||||
harness.gates[job.id].set()
|
||||
harness.gates[job.id].wait(10)
|
||||
job.status = harness.outcomes.get(job.id, module.JobStatus.COMPLETED)
|
||||
if job.status == module.JobStatus.FAILED:
|
||||
job.error = "boom"
|
||||
job.finished_at = time.time()
|
||||
runner_self._save_job(job)
|
||||
|
||||
monkeypatch.setattr(module.JobRunner, "_run_job", fake_run)
|
||||
|
||||
def runner(self, **kwargs):
|
||||
runner = self.module.JobRunner(
|
||||
output_dir=str(self.tmp_path / "out"),
|
||||
log_dir=str(self.tmp_path / "logs"),
|
||||
database=Database(self.tmp_path / "t.db"),
|
||||
upload_dir=str(self.tmp_path / "uploads"),
|
||||
**kwargs,
|
||||
)
|
||||
self.runners.append(runner)
|
||||
return runner
|
||||
|
||||
def make(self, runner, name, **over):
|
||||
kwargs = dict(job_id=str(uuid.uuid4()), model_id="m", name=name, prompt="p")
|
||||
kwargs.update(over)
|
||||
return runner.create_job(**kwargs)
|
||||
|
||||
def finish(self, job, status=None):
|
||||
if status is not None:
|
||||
self.outcomes[job.id] = status
|
||||
self.gates[job.id].set()
|
||||
|
||||
def release_all(self):
|
||||
for gate in self.gates.values():
|
||||
gate.set()
|
||||
|
||||
|
||||
def wait_for(condition, timeout=5.0):
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if condition():
|
||||
return True
|
||||
time.sleep(0.01)
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def h(job_runner_module, tmp_path, monkeypatch):
|
||||
harness = Harness(job_runner_module, tmp_path, monkeypatch)
|
||||
yield harness
|
||||
# Let every job (and the queue's follow-up launches) finish while `_run_job` is
|
||||
# still patched, so no stray thread starts a real job after the test is over.
|
||||
harness.closed = True
|
||||
active = ("queued", "running")
|
||||
for runner in harness.runners:
|
||||
for _ in range(100):
|
||||
harness.release_all()
|
||||
if not any(j.status.value in active for j in runner.list_jobs()):
|
||||
break
|
||||
time.sleep(0.05)
|
||||
runner._shutdown() # stops the worker threads
|
||||
|
||||
|
||||
def statuses(runner, *jobs):
|
||||
return [runner.get_job(j.id).status.value for j in jobs]
|
||||
|
||||
|
||||
def frame_of(job):
|
||||
return [{"source": deferred_last_frame_source(job.id), "media_type": "image"}]
|
||||
|
||||
|
||||
class TestOrderAndConcurrency:
|
||||
def test_one_job_runs_at_a_time_in_the_order_queued(self, h):
|
||||
r = h.runner()
|
||||
a, b, c = (h.make(r, n) for n in "abc")
|
||||
r.enqueue_jobs([b.id, c.id, a.id])
|
||||
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["queued", "running", "queued"])
|
||||
h.finish(b)
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["queued", "completed", "running"])
|
||||
h.finish(c)
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["running", "completed", "completed"])
|
||||
h.finish(a)
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["completed"] * 3)
|
||||
assert h.ran == [b.id, c.id, a.id]
|
||||
|
||||
def test_runs_as_many_at_once_as_allowed(self, h):
|
||||
r = h.runner(max_concurrent_jobs=2)
|
||||
a, b, c = (h.make(r, n) for n in "abc")
|
||||
r.enqueue_jobs([a.id, b.id, c.id])
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["running", "running", "queued"])
|
||||
h.finish(a)
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["completed", "running", "running"])
|
||||
|
||||
def test_a_failed_job_does_not_stop_the_queue(self, h):
|
||||
r = h.runner()
|
||||
a, b = h.make(r, "a"), h.make(r, "b")
|
||||
r.enqueue_jobs([a.id, b.id])
|
||||
h.finish(a, h.module.JobStatus.FAILED)
|
||||
assert wait_for(lambda: statuses(r, a, b) == ["failed", "running"])
|
||||
|
||||
def test_queued_jobs_have_a_queue_time_and_running_ones_do_not(self, h):
|
||||
r = h.runner()
|
||||
a, b = h.make(r, "a"), h.make(r, "b")
|
||||
r.enqueue_jobs([a.id, b.id])
|
||||
assert wait_for(lambda: r.get_job(a.id).status.value == "running")
|
||||
assert r.get_job(a.id).queued_at is None
|
||||
assert r.get_job(b.id).queued_at is not None
|
||||
assert r.get_job(b.id).to_dict()["status"] == "queued"
|
||||
|
||||
|
||||
class TestLongLivedWorkers:
|
||||
"""Model workers die with the thread that started them, so inference jobs share long-lived threads."""
|
||||
|
||||
def _run_all(self, h, r, n):
|
||||
jobs = [h.make(r, f"j{i}") for i in range(n)]
|
||||
for job in jobs:
|
||||
h.gates[job.id].set() # each finishes as soon as it starts
|
||||
r.enqueue_jobs([j.id for j in jobs])
|
||||
assert wait_for(lambda: statuses(r, *jobs) == ["completed"] * n)
|
||||
return jobs
|
||||
|
||||
def test_consecutive_jobs_run_on_the_same_thread_and_it_stays_alive(self, h):
|
||||
r = h.runner()
|
||||
jobs = self._run_all(h, r, 4)
|
||||
threads = {h.thread_of[j.id] for j in jobs}
|
||||
assert len(threads) == 1
|
||||
assert threads.pop().is_alive()
|
||||
|
||||
def test_the_thread_is_not_one_made_for_the_job(self, h):
|
||||
r = h.runner()
|
||||
(job,) = self._run_all(h, r, 1)
|
||||
assert h.thread_of[job.id].name.startswith("inference-worker")
|
||||
|
||||
def test_as_many_workers_as_jobs_allowed_at_once(self, h):
|
||||
r = h.runner(max_concurrent_jobs=2)
|
||||
jobs = self._run_all(h, r, 6)
|
||||
assert len({h.thread_of[j.id] for j in jobs}) <= 2
|
||||
|
||||
def test_two_jobs_at_once_use_two_different_workers(self, h):
|
||||
r = h.runner(max_concurrent_jobs=2)
|
||||
a, b = h.make(r, "a"), h.make(r, "b")
|
||||
r.enqueue_jobs([a.id, b.id])
|
||||
assert wait_for(lambda: statuses(r, a, b) == ["running", "running"])
|
||||
assert wait_for(lambda: len(h.thread_of) == 2)
|
||||
assert h.thread_of[a.id] is not h.thread_of[b.id]
|
||||
|
||||
def test_a_job_that_blows_up_does_not_kill_its_worker(self, h):
|
||||
r = h.runner()
|
||||
bad, good = h.make(r, "bad"), h.make(r, "good")
|
||||
original = h.module.JobRunner._run_job
|
||||
|
||||
def explode_once(self, job):
|
||||
if job.id == bad.id:
|
||||
raise RuntimeError("boom")
|
||||
return original(self, job)
|
||||
|
||||
h.module.JobRunner._run_job = explode_once
|
||||
h.gates[good.id].set()
|
||||
r.enqueue_jobs([bad.id, good.id])
|
||||
assert wait_for(lambda: statuses(r, bad, good) == ["failed", "completed"])
|
||||
assert "boom" in r.get_job(bad.id).error
|
||||
|
||||
def test_a_direct_start_is_refused_while_the_workers_are_busy(self, h):
|
||||
r = h.runner()
|
||||
first, second = h.make(r, "first"), h.make(r, "second")
|
||||
r.start_job(first.id)
|
||||
with pytest.raises(ValueError, match="already running.*Queue this job"):
|
||||
r.start_job(second.id)
|
||||
assert r.get_job(second.id).status.value == "pending"
|
||||
|
||||
def test_a_chain_runs_clean_when_every_clip_finishes_instantly(self, h):
|
||||
r = h.runner()
|
||||
clips = []
|
||||
for i in range(6):
|
||||
clips.append(h.make(r, f"clip {i}", references=frame_of(clips[-1]) if clips else []))
|
||||
for clip in clips:
|
||||
h.gates[clip.id].set()
|
||||
r.enqueue_jobs([c.id for c in clips])
|
||||
assert wait_for(lambda: statuses(r, *clips) == ["completed"] * 6)
|
||||
assert h.ran == [c.id for c in clips]
|
||||
|
||||
|
||||
class TestGeneratorCache:
|
||||
@pytest.fixture
|
||||
def generators(self, job_runner_module, monkeypatch):
|
||||
made = []
|
||||
|
||||
class FakeGenerator:
|
||||
def __init__(self, model_id):
|
||||
self.model_id = model_id
|
||||
self.shut_down = False
|
||||
made.append(self)
|
||||
|
||||
def shutdown(self):
|
||||
self.shut_down = True
|
||||
|
||||
stub = types.SimpleNamespace(from_pretrained=lambda model_id, *a, **k: FakeGenerator(model_id))
|
||||
monkeypatch.setattr(sys.modules["fastvideo"], "VideoGenerator", stub, raising=False)
|
||||
return made
|
||||
|
||||
def _get(self, runner, model="m"):
|
||||
return runner._get_or_create_generator(model, "i2v", 1)
|
||||
|
||||
def _on_thread(self, fn):
|
||||
box = []
|
||||
t = threading.Thread(target=lambda: box.append(fn()))
|
||||
t.start()
|
||||
t.join()
|
||||
return box[0]
|
||||
|
||||
def test_the_thread_that_built_it_keeps_reusing_it(self, h, generators):
|
||||
r = h.runner()
|
||||
first = self._get(r)
|
||||
assert self._get(r) is first
|
||||
assert len(generators) == 1
|
||||
|
||||
def test_a_generator_from_an_exited_thread_is_dropped_and_replaced(self, h, generators):
|
||||
r = h.runner()
|
||||
old = self._on_thread(lambda: self._get(r)) # that thread has exited: its workers are gone
|
||||
new = self._get(r)
|
||||
assert new is not old
|
||||
assert old.shut_down
|
||||
assert len(generators) == 2
|
||||
|
||||
def test_threads_do_not_share_generators(self, h, generators):
|
||||
# Sharing would break the second the moment the first one's thread ends.
|
||||
r = h.runner()
|
||||
mine = self._get(r)
|
||||
other = self._on_thread(lambda: self._get(r))
|
||||
assert other is not mine
|
||||
assert not mine.shut_down # still owned by a live thread
|
||||
|
||||
def test_a_different_model_replaces_the_loaded_one_instead_of_sitting_beside_it(self, h, generators):
|
||||
r = h.runner()
|
||||
first = self._get(r, "model-a")
|
||||
second = self._get(r, "model-b")
|
||||
assert first.shut_down # its GPU memory is freed before the next load
|
||||
assert not second.shut_down
|
||||
assert self._get(r, "model-b") is second
|
||||
|
||||
def test_switching_back_reloads(self, h, generators):
|
||||
r = h.runner()
|
||||
a1 = self._get(r, "model-a")
|
||||
self._get(r, "model-b")
|
||||
a2 = self._get(r, "model-a")
|
||||
assert a2 is not a1
|
||||
|
||||
def test_discarding_unloads_this_threads_models_only(self, h, generators):
|
||||
r = h.runner()
|
||||
mine = self._get(r)
|
||||
gate, done = threading.Event(), threading.Event()
|
||||
theirs = []
|
||||
|
||||
def other_thread():
|
||||
theirs.append(self._get(r, "other"))
|
||||
done.set()
|
||||
gate.wait(5)
|
||||
|
||||
t = threading.Thread(target=other_thread)
|
||||
t.start()
|
||||
done.wait(5)
|
||||
r._discard_generators()
|
||||
gate.set()
|
||||
t.join()
|
||||
assert mine.shut_down and not theirs[0].shut_down
|
||||
assert self._get(r) is not mine # reloads fresh next time
|
||||
|
||||
|
||||
class TestLastFrameChains:
|
||||
def _chain(self, h, r, n=3):
|
||||
clips = []
|
||||
for i in range(n):
|
||||
refs = frame_of(clips[-1]) if clips else []
|
||||
clips.append(h.make(r, f"clip {i + 1}", references=refs))
|
||||
return clips
|
||||
|
||||
def test_a_scene_runs_clip_by_clip_however_many_slots_there_are(self, h):
|
||||
r = h.runner(max_concurrent_jobs=4)
|
||||
c1, c2, c3 = self._chain(h, r)
|
||||
r.enqueue_jobs([c1.id, c2.id, c3.id])
|
||||
|
||||
assert wait_for(lambda: statuses(r, c1, c2, c3) == ["running", "queued", "queued"])
|
||||
time.sleep(0.05)
|
||||
assert statuses(r, c1, c2, c3) == ["running", "queued", "queued"] # spare slots stay unused
|
||||
h.finish(c1)
|
||||
assert wait_for(lambda: statuses(r, c1, c2, c3) == ["completed", "running", "queued"])
|
||||
h.finish(c2)
|
||||
assert wait_for(lambda: statuses(r, c1, c2, c3) == ["completed", "completed", "running"])
|
||||
|
||||
def test_a_clip_does_not_block_unrelated_jobs_behind_it(self, h):
|
||||
r = h.runner(max_concurrent_jobs=2)
|
||||
c1, c2 = self._chain(h, r, 2)
|
||||
other = h.make(r, "other")
|
||||
r.enqueue_jobs([c1.id, c2.id, other.id])
|
||||
assert wait_for(lambda: statuses(r, c1, c2, other) == ["running", "queued", "running"])
|
||||
|
||||
def test_a_failed_clip_fails_the_clips_after_it_and_says_why(self, h):
|
||||
r = h.runner()
|
||||
c1, c2, c3 = self._chain(h, r)
|
||||
r.enqueue_jobs([c1.id, c2.id, c3.id])
|
||||
h.finish(c1, h.module.JobStatus.FAILED)
|
||||
|
||||
assert wait_for(lambda: statuses(r, c1, c2, c3) == ["failed"] * 3)
|
||||
assert "clip 1" in r.get_job(c2.id).error and "failed" in r.get_job(c2.id).error
|
||||
assert "clip 2" in r.get_job(c3.id).error
|
||||
assert h.ran == [c1.id]
|
||||
|
||||
def test_a_stopped_clip_fails_the_clips_after_it(self, h):
|
||||
r = h.runner()
|
||||
c1, c2 = self._chain(h, r, 2)
|
||||
r.enqueue_jobs([c1.id, c2.id])
|
||||
h.finish(c1, h.module.JobStatus.STOPPED)
|
||||
assert wait_for(lambda: statuses(r, c1, c2) == ["stopped", "failed"])
|
||||
assert "stopped" in r.get_job(c2.id).error
|
||||
|
||||
def test_deleting_a_queued_clip_fails_the_clips_after_it(self, h):
|
||||
r = h.runner()
|
||||
blocker = h.make(r, "blocker")
|
||||
c1, c2 = self._chain(h, r, 2)
|
||||
r.enqueue_jobs([blocker.id, c1.id, c2.id])
|
||||
assert wait_for(lambda: statuses(r, blocker, c1, c2) == ["running", "queued", "queued"])
|
||||
assert r.delete_job(c1.id)
|
||||
assert wait_for(lambda: r.get_job(c2.id).status.value == "failed")
|
||||
assert "deleted" in r.get_job(c2.id).error
|
||||
|
||||
def test_a_failed_clip_can_be_fixed_and_the_scene_requeued(self, h):
|
||||
r = h.runner()
|
||||
c1, c2 = self._chain(h, r, 2)
|
||||
r.enqueue_jobs([c1.id, c2.id])
|
||||
h.finish(c1, h.module.JobStatus.FAILED)
|
||||
assert wait_for(lambda: statuses(r, c1, c2) == ["failed", "failed"])
|
||||
|
||||
h.outcomes.pop(c1.id)
|
||||
h.gates[c1.id].clear()
|
||||
r.enqueue_jobs([c1.id, c2.id])
|
||||
assert wait_for(lambda: statuses(r, c1, c2) == ["running", "queued"])
|
||||
h.finish(c1)
|
||||
assert wait_for(lambda: statuses(r, c1, c2) == ["completed", "running"])
|
||||
|
||||
def test_the_stored_reference_stays_deferred(self, h):
|
||||
r = h.runner()
|
||||
c1, c2 = self._chain(h, r, 2)
|
||||
r.enqueue_jobs([c1.id, c2.id])
|
||||
h.finish(c1)
|
||||
assert wait_for(lambda: r.get_job(c2.id).status.value == "running")
|
||||
assert r.get_job(c2.id).references == frame_of(c1)
|
||||
|
||||
|
||||
class TestStartingDirectly:
|
||||
def test_refuses_while_the_frame_source_is_unfinished(self, h):
|
||||
r = h.runner()
|
||||
c1, c2 = h.make(r, "clip 1"), None
|
||||
c2 = h.make(r, "clip 2", references=frame_of(c1))
|
||||
with pytest.raises(ValueError, match="clip 1.*hasn't finished.*pending"):
|
||||
r.start_job(c2.id)
|
||||
assert r.get_job(c2.id).status.value == "pending"
|
||||
assert h.ran == []
|
||||
|
||||
def test_starts_once_the_frame_source_completed(self, h):
|
||||
r = h.runner()
|
||||
c1 = h.make(r, "clip 1")
|
||||
c1.status = h.module.JobStatus.COMPLETED
|
||||
c2 = h.make(r, "clip 2", references=frame_of(c1))
|
||||
r.start_job(c2.id)
|
||||
assert r.get_job(c2.id).status.value == "running"
|
||||
|
||||
def test_refuses_when_the_frame_source_was_deleted(self, h):
|
||||
r = h.runner()
|
||||
c1 = h.make(r, "clip 1")
|
||||
c2 = h.make(r, "clip 2", references=frame_of(c1))
|
||||
r.delete_job(c1.id)
|
||||
with pytest.raises(ValueError, match="no longer exists"):
|
||||
r.start_job(c2.id)
|
||||
|
||||
def test_a_direct_start_takes_a_queue_slot(self, h):
|
||||
r = h.runner()
|
||||
running, queued = h.make(r, "running"), h.make(r, "queued")
|
||||
r.start_job(running.id)
|
||||
r.enqueue_job(queued.id)
|
||||
time.sleep(0.05)
|
||||
assert statuses(r, running, queued) == ["running", "queued"]
|
||||
h.finish(running)
|
||||
assert wait_for(lambda: r.get_job(queued.id).status.value == "running")
|
||||
|
||||
def test_a_queued_job_cannot_also_be_started(self, h):
|
||||
r = h.runner()
|
||||
blocker, waiting = h.make(r, "blocker"), h.make(r, "waiting")
|
||||
r.enqueue_jobs([blocker.id, waiting.id])
|
||||
with pytest.raises(ValueError, match="already queued"):
|
||||
r.start_job(waiting.id)
|
||||
|
||||
|
||||
class TestQueueManagement:
|
||||
def test_dequeue_returns_a_job_to_pending_and_lets_others_go(self, h):
|
||||
r = h.runner()
|
||||
a, b, c = (h.make(r, n) for n in "abc")
|
||||
r.enqueue_jobs([a.id, b.id, c.id])
|
||||
assert wait_for(lambda: r.get_job(a.id).status.value == "running")
|
||||
r.dequeue_job(b.id)
|
||||
assert r.get_job(b.id).status.value == "pending"
|
||||
assert r.get_job(b.id).queued_at is None
|
||||
h.finish(a)
|
||||
assert wait_for(lambda: statuses(r, a, b, c) == ["completed", "pending", "running"])
|
||||
|
||||
def test_stopping_a_queued_job_dequeues_it(self, h):
|
||||
r = h.runner()
|
||||
a, b = h.make(r, "a"), h.make(r, "b")
|
||||
r.enqueue_jobs([a.id, b.id])
|
||||
assert wait_for(lambda: r.get_job(a.id).status.value == "running")
|
||||
assert r.stop_job(b.id).status.value == "pending"
|
||||
|
||||
def test_dequeue_of_a_job_not_queued_is_refused(self, h):
|
||||
r = h.runner()
|
||||
a = h.make(r, "a")
|
||||
with pytest.raises(ValueError, match="not queued"):
|
||||
r.dequeue_job(a.id)
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
r.dequeue_job("nope")
|
||||
|
||||
def test_enqueue_is_all_or_nothing(self, h):
|
||||
r = h.runner()
|
||||
a, done = h.make(r, "a"), h.make(r, "done")
|
||||
done.status = h.module.JobStatus.COMPLETED
|
||||
with pytest.raises(ValueError, match="already completed"):
|
||||
r.enqueue_jobs([a.id, done.id])
|
||||
assert r.get_job(a.id).status.value == "pending"
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
r.enqueue_jobs([a.id, "nope"])
|
||||
assert h.ran == []
|
||||
|
||||
def test_a_job_cannot_be_queued_twice(self, h):
|
||||
r = h.runner()
|
||||
blocker, a = h.make(r, "blocker"), h.make(r, "a")
|
||||
r.enqueue_jobs([blocker.id, a.id])
|
||||
with pytest.raises(ValueError, match="already queued"):
|
||||
r.enqueue_job(a.id)
|
||||
with pytest.raises(ValueError, match="more than once"):
|
||||
r.enqueue_jobs([a.id, a.id])
|
||||
|
||||
def test_loops_of_last_frames_are_refused(self, h):
|
||||
r = h.runner()
|
||||
a = h.make(r, "a")
|
||||
b = h.make(r, "b", references=frame_of(a))
|
||||
r.update_job_config(a.id, {"references": frame_of(b)})
|
||||
with pytest.raises(ValueError, match="loop"):
|
||||
r.enqueue_jobs([a.id, b.id])
|
||||
assert statuses(r, a, b) == ["pending", "pending"]
|
||||
|
||||
def test_queued_jobs_cannot_be_edited(self, h):
|
||||
r = h.runner()
|
||||
blocker, a = h.make(r, "blocker"), h.make(r, "a")
|
||||
r.enqueue_jobs([blocker.id, a.id])
|
||||
with pytest.raises(ValueError, match="queued"):
|
||||
r.update_job_config(a.id, {"prompt": "changed"})
|
||||
|
||||
def test_requeueing_clears_the_previous_error(self, h):
|
||||
r = h.runner()
|
||||
a = h.make(r, "a")
|
||||
h.finish(a, h.module.JobStatus.FAILED)
|
||||
r.enqueue_job(a.id)
|
||||
assert wait_for(lambda: r.get_job(a.id).status.value == "failed")
|
||||
h.outcomes.pop(a.id)
|
||||
h.gates[a.id].clear()
|
||||
r.enqueue_job(a.id)
|
||||
assert wait_for(lambda: r.get_job(a.id).status.value == "running")
|
||||
assert r.get_job(a.id).error is None
|
||||
|
||||
|
||||
class TestRestart:
|
||||
def test_queued_jobs_resume_in_order_after_a_restart(self, h, job_runner_module):
|
||||
first = h.runner()
|
||||
blocker = h.make(first, "blocker")
|
||||
a, b = h.make(first, "a"), h.make(first, "b")
|
||||
first.enqueue_jobs([blocker.id, b.id, a.id])
|
||||
assert wait_for(lambda: h.ran == [blocker.id]) # actually started, not just handed to a worker
|
||||
|
||||
h.ran.clear()
|
||||
second = h.runner() # same database
|
||||
# The job that was running is failed, exactly as before queues existed...
|
||||
assert second.get_job(blocker.id).status.value == "failed"
|
||||
assert "restarted" in second.get_job(blocker.id).error
|
||||
# ...and the queue carries on where it was.
|
||||
assert wait_for(lambda: h.ran[:1] == [b.id])
|
||||
assert second.get_job(a.id).status.value == "queued"
|
||||
|
||||
def test_edits_survive_a_restart(self, h):
|
||||
first = h.runner()
|
||||
job = h.make(first, "a", fps=24, negative_prompt="neg")
|
||||
first.update_job_config(job.id, {"prompt": "edited", "references": frame_of(job), "fps": 12})
|
||||
|
||||
second = h.runner()
|
||||
restored = second.get_job(job.id)
|
||||
assert restored.prompt == "edited"
|
||||
assert restored.references == frame_of(job)
|
||||
assert restored.fps == 12
|
||||
assert restored.negative_prompt == "neg"
|
||||
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""A finished clip's final frame is saved as an image the next clip can use as a reference."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from fastvideo_studio import frames
|
||||
from fastvideo_studio.frames import FrameError, last_frame_for_job
|
||||
|
||||
WIDTH, HEIGHT, FRAMES = 128, 96, 12
|
||||
|
||||
|
||||
class Status(Enum): # mirrors JobStatus without importing job_runner (which needs fastvideo)
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
STOPPED = "stopped"
|
||||
|
||||
|
||||
def _write_video(path: str, *, last_rgb=(0, 0, 255), other_rgb=(255, 0, 0)) -> None:
|
||||
"""A clip whose frames are all `other_rgb` except the last, which is `last_rgb`."""
|
||||
import av
|
||||
|
||||
container = av.open(path, mode="w")
|
||||
try:
|
||||
stream = container.add_stream("h264", rate=24)
|
||||
stream.width, stream.height, stream.pix_fmt = WIDTH, HEIGHT, "yuv420p"
|
||||
for i in range(FRAMES):
|
||||
rgb = last_rgb if i == FRAMES - 1 else other_rgb
|
||||
array = np.full((HEIGHT, WIDTH, 3), rgb, dtype=np.uint8)
|
||||
frame = av.VideoFrame.from_ndarray(array, format="rgb24")
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def _job(job_id="job-1", status=Status.COMPLETED, output_path=None):
|
||||
return SimpleNamespace(id=job_id, status=status, output_path=output_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def video(tmp_path):
|
||||
path = str(tmp_path / "clip.mp4")
|
||||
_write_video(path)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def uploads(tmp_path):
|
||||
return str(tmp_path / "uploads")
|
||||
|
||||
|
||||
def test_saves_the_final_frame_not_the_first(video, uploads, tmp_path):
|
||||
path = last_frame_for_job(_job(output_path=video), uploads)
|
||||
|
||||
assert path == str(tmp_path / "uploads" / "last_frames" / "last_frame_job-1.png")
|
||||
assert os.path.isabs(path)
|
||||
image = Image.open(path).convert("RGB")
|
||||
assert image.size == (WIDTH, HEIGHT)
|
||||
r, _g, b = image.getpixel((WIDTH // 2, HEIGHT // 2))
|
||||
# the last frame is blue and every other frame is red -- a wrong frame would be red
|
||||
assert b > 150 and r < 100, (r, b)
|
||||
|
||||
|
||||
def test_accepts_a_plain_string_status(video, uploads):
|
||||
assert last_frame_for_job(_job(status="completed", output_path=video), uploads).endswith(".png")
|
||||
|
||||
|
||||
def test_extraction_is_reused_until_the_video_changes(video, uploads, monkeypatch):
|
||||
calls = []
|
||||
real = frames.save_last_frame
|
||||
monkeypatch.setattr(frames, "save_last_frame", lambda v, o: (calls.append(o), real(v, o))[1])
|
||||
|
||||
first = last_frame_for_job(_job(output_path=video), uploads)
|
||||
second = last_frame_for_job(_job(output_path=video), uploads)
|
||||
assert first == second
|
||||
assert len(calls) == 1
|
||||
|
||||
# the clip is regenerated -> the saved frame is stale and gets rebuilt
|
||||
later = os.path.getmtime(first) + 10
|
||||
os.utime(video, (later, later))
|
||||
last_frame_for_job(_job(output_path=video), uploads)
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
def test_each_job_gets_its_own_file(video, uploads):
|
||||
a = last_frame_for_job(_job("a", output_path=video), uploads)
|
||||
b = last_frame_for_job(_job("b", output_path=video), uploads)
|
||||
assert a != b and os.path.isfile(a) and os.path.isfile(b)
|
||||
|
||||
|
||||
def test_leaves_no_temp_file_behind(video, uploads):
|
||||
last_frame_for_job(_job(output_path=video), uploads)
|
||||
assert os.listdir(os.path.join(uploads, "last_frames")) == ["last_frame_job-1.png"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [Status.PENDING, Status.RUNNING, Status.FAILED, Status.STOPPED])
|
||||
def test_a_job_that_has_not_finished_is_404(video, uploads, status):
|
||||
with pytest.raises(FrameError) as e:
|
||||
last_frame_for_job(_job(status=status, output_path=video), uploads)
|
||||
assert e.value.status_code == 404
|
||||
assert "No output" in e.value.detail
|
||||
|
||||
|
||||
def test_a_completed_job_with_no_output_path_is_404(uploads):
|
||||
with pytest.raises(FrameError) as e:
|
||||
last_frame_for_job(_job(output_path=None), uploads)
|
||||
assert e.value.status_code == 404
|
||||
|
||||
|
||||
def test_missing_output_file_is_404(uploads, tmp_path):
|
||||
with pytest.raises(FrameError) as e:
|
||||
last_frame_for_job(_job(output_path=str(tmp_path / "gone.mp4")), uploads)
|
||||
assert (e.value.status_code, "not found on disk" in e.value.detail) == (404, True)
|
||||
|
||||
|
||||
def test_an_image_output_is_rejected(uploads, tmp_path):
|
||||
png = tmp_path / "out.png"
|
||||
Image.new("RGB", (8, 8)).save(png)
|
||||
with pytest.raises(FrameError) as e:
|
||||
last_frame_for_job(_job(output_path=str(png)), uploads)
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
def test_an_unreadable_video_is_a_500_with_a_reason(uploads, tmp_path):
|
||||
bad = tmp_path / "bad.mp4"
|
||||
bad.write_bytes(b"not a video")
|
||||
with pytest.raises(FrameError) as e:
|
||||
last_frame_for_job(_job(output_path=str(bad)), uploads)
|
||||
assert e.value.status_code == 500
|
||||
assert "Could not read the last frame" in e.value.detail
|
||||
|
||||
|
||||
def test_a_failed_extraction_does_not_leave_a_partial_file(uploads, tmp_path):
|
||||
bad = tmp_path / "bad.mp4"
|
||||
bad.write_bytes(b"not a video")
|
||||
with pytest.raises(FrameError):
|
||||
last_frame_for_job(_job(output_path=str(bad)), uploads)
|
||||
assert os.listdir(os.path.join(uploads, "last_frames")) == []
|
||||
|
||||
|
||||
def test_needs_an_upload_directory(video):
|
||||
with pytest.raises(FrameError) as e:
|
||||
last_frame_for_job(_job(output_path=video), "")
|
||||
assert e.value.status_code == 503
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""A scene's clips are joined, in order, into one video without re-encoding."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from enum import Enum
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from fastvideo_studio import merge
|
||||
from fastvideo_studio.merge import MergeError, find_ffmpeg, merge_clips, merge_jobs, merged_path, probe_clip
|
||||
|
||||
pytestmark = pytest.mark.skipif(find_ffmpeg() is None, reason="needs an ffmpeg binary")
|
||||
|
||||
FPS = 24
|
||||
COLORS = {"red": (255, 0, 0), "green": (0, 255, 0), "blue": (0, 0, 255)}
|
||||
|
||||
|
||||
class Status(Enum): # mirrors JobStatus without importing job_runner (which needs fastvideo)
|
||||
PENDING = "pending"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
def write_clip(path, *, frames=24, color="red", size=(128, 96), fps=FPS, audio=True, audio_extra=0.042):
|
||||
"""A clip of one flat colour. Like real H3 output, its audio runs a little past the video."""
|
||||
import av
|
||||
|
||||
container = av.open(path, mode="w")
|
||||
try:
|
||||
video = container.add_stream("h264", rate=fps)
|
||||
video.width, video.height, video.pix_fmt = size[0], size[1], "yuv420p"
|
||||
sound = None
|
||||
if audio:
|
||||
sound = container.add_stream("aac", rate=32000)
|
||||
sound.layout = "stereo"
|
||||
for _ in range(frames):
|
||||
array = np.full((size[1], size[0], 3), COLORS[color], dtype=np.uint8)
|
||||
for packet in video.encode(av.VideoFrame.from_ndarray(array, format="rgb24")):
|
||||
container.mux(packet)
|
||||
for packet in video.encode():
|
||||
container.mux(packet)
|
||||
if sound is not None:
|
||||
total = int((frames / fps + audio_extra) * 32000)
|
||||
tone = (np.sin(np.arange(total) * 2 * np.pi * 440 / 32000) * 0.2).astype(np.float32)
|
||||
pts = 0
|
||||
for start in range(0, total, 1024):
|
||||
chunk = tone[start:start + 1024]
|
||||
chunk = np.pad(chunk, (0, 1024 - len(chunk)))
|
||||
frame = av.AudioFrame.from_ndarray(np.stack([chunk, chunk]), format="fltp", layout="stereo")
|
||||
frame.sample_rate, frame.pts = 32000, pts
|
||||
pts += 1024
|
||||
for packet in sound.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in sound.encode():
|
||||
container.mux(packet)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
|
||||
def read_back(path):
|
||||
"""(frame colours in order, audio seconds, video seconds) of a finished video."""
|
||||
import av
|
||||
|
||||
with av.open(path) as c:
|
||||
v = c.streams.video[0]
|
||||
colours = []
|
||||
for frame in c.decode(v):
|
||||
r, g, b = frame.to_ndarray(format="rgb24")[48, 64]
|
||||
colours.append("red" if r > 150 else "green" if g > 150 else "blue")
|
||||
audio = next((s for s in c.streams if s.type == "audio"), None)
|
||||
a_secs = float(audio.duration * audio.time_base) if audio and audio.duration else None
|
||||
return colours, a_secs, float(v.duration * v.time_base)
|
||||
|
||||
|
||||
def job(path, name="", status=Status.COMPLETED, job_id="j"):
|
||||
return SimpleNamespace(id=job_id, name=name, status=status, output_path=path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clips(tmp_path):
|
||||
paths = []
|
||||
for i, (color, frames) in enumerate([("red", 24), ("green", 36), ("blue", 12)]):
|
||||
p = str(tmp_path / f"clip{i}.mp4")
|
||||
write_clip(p, frames=frames, color=color)
|
||||
paths.append(p)
|
||||
return paths
|
||||
|
||||
|
||||
class TestMergeClips:
|
||||
def test_joins_clips_in_the_order_given(self, clips, tmp_path):
|
||||
out = str(tmp_path / "out.mp4")
|
||||
merge_clips(clips, out)
|
||||
colours, _, _ = read_back(out)
|
||||
assert colours == ["red"] * 24 + ["green"] * 36 + ["blue"] * 12
|
||||
|
||||
def test_the_order_is_the_callers_not_the_filenames(self, clips, tmp_path):
|
||||
out = str(tmp_path / "out.mp4")
|
||||
merge_clips([clips[2], clips[0], clips[1]], out)
|
||||
colours, _, _ = read_back(out)
|
||||
assert colours == ["blue"] * 12 + ["red"] * 24 + ["green"] * 36
|
||||
|
||||
def test_keeps_every_frame_and_the_full_length(self, clips, tmp_path):
|
||||
out = str(tmp_path / "out.mp4")
|
||||
seconds = merge_clips(clips, out)
|
||||
colours, _, video_secs = read_back(out)
|
||||
assert len(colours) == 72
|
||||
assert seconds == pytest.approx(72 / FPS, abs=0.02)
|
||||
assert video_secs == pytest.approx(72 / FPS, abs=0.05)
|
||||
|
||||
def test_audio_stays_in_step_with_the_picture(self, clips, tmp_path):
|
||||
# Each clip's audio runs ~40 ms past its video; joined naively that adds up to a gap or drift.
|
||||
out = str(tmp_path / "out.mp4")
|
||||
merge_clips(clips, out)
|
||||
_, audio_secs, video_secs = read_back(out)
|
||||
assert audio_secs is not None
|
||||
assert abs(audio_secs - video_secs) < 0.08 # not 3 x 0.042 and growing
|
||||
|
||||
def test_a_single_clip_is_fine(self, clips, tmp_path):
|
||||
out = str(tmp_path / "out.mp4")
|
||||
merge_clips(clips[:1], out)
|
||||
assert len(read_back(out)[0]) == 24
|
||||
|
||||
def test_clips_without_audio_join_too(self, tmp_path):
|
||||
paths = []
|
||||
for i, color in enumerate(("red", "green")):
|
||||
p = str(tmp_path / f"s{i}.mp4")
|
||||
write_clip(p, frames=12, color=color, audio=False)
|
||||
paths.append(p)
|
||||
out = str(tmp_path / "out.mp4")
|
||||
merge_clips(paths, out)
|
||||
colours, audio_secs, _ = read_back(out)
|
||||
assert colours == ["red"] * 12 + ["green"] * 12 and audio_secs is None
|
||||
|
||||
def test_creates_the_output_folder_and_leaves_no_temp_files(self, clips, tmp_path):
|
||||
out = str(tmp_path / "new" / "dir" / "out.mp4")
|
||||
merge_clips(clips, out)
|
||||
assert os.listdir(os.path.dirname(out)) == ["out.mp4"]
|
||||
|
||||
def test_paths_with_quotes_and_spaces(self, tmp_path):
|
||||
paths = []
|
||||
for i, color in enumerate(("red", "blue")):
|
||||
p = str(tmp_path / f"it's clip {i}.mp4")
|
||||
write_clip(p, frames=6, color=color)
|
||||
paths.append(p)
|
||||
out = str(tmp_path / "out.mp4")
|
||||
merge_clips(paths, out)
|
||||
assert read_back(out)[0] == ["red"] * 6 + ["blue"] * 6
|
||||
|
||||
|
||||
class TestRefusals:
|
||||
def test_clips_of_different_sizes(self, clips, tmp_path):
|
||||
odd = str(tmp_path / "odd.mp4")
|
||||
write_clip(odd, size=(160, 96))
|
||||
with pytest.raises(MergeError, match=r"'clip 3' \(160x96.*doesn't match 'clip 1' \(128x96") as e:
|
||||
merge_clips([clips[0], clips[1], odd], str(tmp_path / "o.mp4"), labels=["'clip 1'", "'clip 2'", "'clip 3'"])
|
||||
assert e.value.status_code == 409
|
||||
assert not (tmp_path / "o.mp4").exists()
|
||||
|
||||
def test_clips_of_different_frame_rates(self, clips, tmp_path):
|
||||
odd = str(tmp_path / "odd.mp4")
|
||||
write_clip(odd, fps=30)
|
||||
with pytest.raises(MergeError, match="doesn't match"):
|
||||
merge_clips([clips[0], odd], str(tmp_path / "o.mp4"))
|
||||
|
||||
def test_one_clip_with_audio_and_one_without(self, clips, tmp_path):
|
||||
silent = str(tmp_path / "silent.mp4")
|
||||
write_clip(silent, audio=False)
|
||||
with pytest.raises(MergeError, match="no audio"):
|
||||
merge_clips([clips[0], silent], str(tmp_path / "o.mp4"))
|
||||
|
||||
def test_nothing_to_merge(self, tmp_path):
|
||||
with pytest.raises(MergeError) as e:
|
||||
merge_clips([], str(tmp_path / "o.mp4"))
|
||||
assert e.value.status_code == 400
|
||||
|
||||
def test_a_file_that_is_not_a_video(self, tmp_path):
|
||||
bad = tmp_path / "bad.mp4"
|
||||
bad.write_bytes(b"not a video")
|
||||
with pytest.raises(MergeError, match="Could not read bad.mp4"):
|
||||
merge_clips([str(bad)], str(tmp_path / "o.mp4"))
|
||||
|
||||
def test_no_ffmpeg(self, clips, tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(merge, "find_ffmpeg", lambda: None)
|
||||
with pytest.raises(MergeError) as e:
|
||||
merge_clips(clips, str(tmp_path / "o.mp4"))
|
||||
assert e.value.status_code == 503 and "ffmpeg" in e.value.detail
|
||||
|
||||
|
||||
class TestMergeJobs:
|
||||
def test_merges_the_jobs_videos_into_a_merged_folder(self, clips, tmp_path):
|
||||
jobs = [job(p, name=f"clip {i + 1}", job_id=f"j{i}") for i, p in enumerate(clips)]
|
||||
result = merge_jobs(jobs, str(tmp_path / "out"), name="Wolf Lunch")
|
||||
assert result.clips == 3 and result.seconds == pytest.approx(3.0, abs=0.02)
|
||||
assert result.filename.startswith("Wolf-Lunch-") and result.filename.endswith(".mp4")
|
||||
assert result.path == str(tmp_path / "out" / "merged" / result.filename)
|
||||
assert len(read_back(result.path)[0]) == 72
|
||||
|
||||
def test_a_name_cannot_escape_the_folder(self, clips, tmp_path):
|
||||
result = merge_jobs([job(clips[0])], str(tmp_path / "out"), name="../../etc/passwd")
|
||||
assert os.path.dirname(result.path) == str(tmp_path / "out" / "merged")
|
||||
assert "/" not in result.filename
|
||||
|
||||
def test_an_empty_name_still_gets_a_filename(self, clips, tmp_path):
|
||||
assert merge_jobs([job(clips[0])], str(tmp_path / "out"), name=" ").filename.startswith("scene-")
|
||||
|
||||
def test_refuses_a_clip_that_has_not_finished_and_names_it(self, clips, tmp_path):
|
||||
jobs = [job(clips[0], name="clip 1"), job(None, name="clip 2", status=Status.PENDING)]
|
||||
with pytest.raises(MergeError, match=r"'clip 2' hasn't finished \(pending\)") as e:
|
||||
merge_jobs(jobs, str(tmp_path / "out"))
|
||||
assert e.value.status_code == 409
|
||||
assert not (tmp_path / "out").exists() # nothing partial
|
||||
|
||||
def test_refuses_a_video_missing_from_disk(self, clips, tmp_path):
|
||||
jobs = [job(clips[0], name="a"), job(str(tmp_path / "gone.mp4"), name="b")]
|
||||
with pytest.raises(MergeError, match="'b'.*missing from disk"):
|
||||
merge_jobs(jobs, str(tmp_path / "out"))
|
||||
|
||||
def test_refuses_an_output_that_is_not_a_video(self, tmp_path):
|
||||
img = tmp_path / "a.png"
|
||||
img.write_bytes(b"x")
|
||||
with pytest.raises(MergeError, match="isn't a video"):
|
||||
merge_jobs([job(str(img), name="a")], str(tmp_path / "out"))
|
||||
|
||||
def test_refuses_no_jobs(self, tmp_path):
|
||||
with pytest.raises(MergeError) as e:
|
||||
merge_jobs([], str(tmp_path))
|
||||
assert e.value.status_code == 400
|
||||
|
||||
|
||||
class TestMergedPath:
|
||||
def test_finds_a_merged_video(self, clips, tmp_path):
|
||||
result = merge_jobs([job(clips[0])], str(tmp_path))
|
||||
assert merged_path(str(tmp_path), result.filename) == result.path
|
||||
|
||||
@pytest.mark.parametrize("name", ["../secret.mp4", "a/b.mp4", "..", "x.txt", "", "nope.mp4", "a b.mp4"])
|
||||
def test_refuses_anything_else(self, tmp_path, name):
|
||||
(tmp_path / "merged").mkdir()
|
||||
(tmp_path / "secret.mp4").write_bytes(b"x")
|
||||
with pytest.raises(MergeError) as e:
|
||||
merged_path(str(tmp_path), name)
|
||||
assert e.value.status_code == 404
|
||||
|
||||
|
||||
def test_probe_reads_real_clip_settings(clips):
|
||||
info = probe_clip(clips[1])
|
||||
assert (info.width, info.height, info.video_codec, info.pix_fmt) == (128, 96, "h264", "yuv420p")
|
||||
assert info.fps == pytest.approx(24) and info.video_seconds == pytest.approx(1.5, abs=0.03)
|
||||
assert info.audio == ("aac", 32000, 2)
|
||||
@@ -0,0 +1,238 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""The mock API's job edit and last-frame routes behave like the real server's."""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from fastvideo_studio import mock_server
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_state():
|
||||
mock_server._jobs.clear()
|
||||
yield
|
||||
mock_server._jobs.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
return TestClient(mock_server.app)
|
||||
|
||||
|
||||
def _create(client, **over):
|
||||
body = {"model_id": "MiniMaxAI/MiniMax-H3", "prompt": "p", "workload_type": "i2v", **over}
|
||||
res = client.post("/api/jobs", json=body)
|
||||
assert res.status_code == 201
|
||||
return res.json()
|
||||
|
||||
|
||||
def test_edits_a_pending_jobs_prompt_and_references(client):
|
||||
job = _create(client, name="clip-2", references=[{"source": "/a.png", "media_type": "image"}])
|
||||
refs = job["references"] + [{"source": "/b.png", "media_type": "image"}]
|
||||
|
||||
res = client.patch(f"/api/jobs/{job['id']}", json={"prompt": "new prompt", "references": refs})
|
||||
|
||||
assert res.status_code == 200
|
||||
body = res.json()
|
||||
assert body["prompt"] == "new prompt"
|
||||
assert [r["source"] for r in body["references"]] == ["/a.png", "/b.png"]
|
||||
assert body["name"] == "clip-2" # untouched fields stay
|
||||
assert client.get(f"/api/jobs/{job['id']}").json()["prompt"] == "new prompt"
|
||||
|
||||
|
||||
def test_a_started_job_cannot_be_edited(client):
|
||||
job = _create(client)
|
||||
mock_server._jobs[job["id"]]["status"] = "completed"
|
||||
res = client.patch(f"/api/jobs/{job['id']}", json={"prompt": "x"})
|
||||
assert res.status_code == 400
|
||||
assert "only pending, failed, stopped jobs can be edited" in res.json()["detail"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["failed", "stopped"])
|
||||
def test_failed_and_stopped_jobs_can_be_edited(client, status):
|
||||
job = _create(client)
|
||||
mock_server._jobs[job["id"]]["status"] = status
|
||||
assert client.patch(f"/api/jobs/{job['id']}", json={"prompt": "again"}).status_code == 200
|
||||
|
||||
|
||||
def test_unknown_fields_are_rejected(client):
|
||||
job = _create(client)
|
||||
res = client.patch(f"/api/jobs/{job['id']}", json={"status": "completed", "nonsense": 1})
|
||||
assert res.status_code == 400
|
||||
assert "Not editable: nonsense, status" in res.json()["detail"]
|
||||
|
||||
|
||||
def test_editing_an_unknown_job_is_404(client):
|
||||
assert client.patch("/api/jobs/nope", json={"prompt": "x"}).status_code == 404
|
||||
|
||||
|
||||
def test_last_frame_of_a_completed_job(client):
|
||||
job = _create(client)
|
||||
mock_server._jobs[job["id"]].update(status="completed", output_path="/out/a.mp4")
|
||||
res = client.post(f"/api/jobs/{job['id']}/last-frame")
|
||||
assert res.status_code == 200
|
||||
assert res.json() == {"path": f"/mock/last_frames/last_frame_{job['id']}.png", "media_type": "image"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", ["pending", "failed", "stopped"])
|
||||
def test_last_frame_of_an_unfinished_job_is_404(client, status):
|
||||
job = _create(client)
|
||||
mock_server._jobs[job["id"]]["status"] = status
|
||||
res = client.post(f"/api/jobs/{job['id']}/last-frame")
|
||||
assert res.status_code == 404
|
||||
assert "No output" in res.json()["detail"]
|
||||
|
||||
|
||||
def test_last_frame_of_an_unknown_job_is_404(client):
|
||||
assert client.post("/api/jobs/nope/last-frame").status_code == 404
|
||||
|
||||
|
||||
# --- queue -------------------------------------------------------------------
|
||||
|
||||
|
||||
def _frame_of(job):
|
||||
return [{"source": f"job-last-frame:{job['id']}", "media_type": "image"}]
|
||||
|
||||
|
||||
def _statuses(client, *jobs):
|
||||
return [client.get(f"/api/jobs/{j['id']}").json()["status"] for j in jobs]
|
||||
|
||||
|
||||
def _age(job, seconds):
|
||||
"""Pretend `seconds` have passed since `job` started."""
|
||||
record = mock_server._jobs[job["id"]]
|
||||
record["started_at"] -= seconds
|
||||
|
||||
|
||||
def test_queued_jobs_run_one_at_a_time_in_order(client):
|
||||
a, b, c = (_create(client, name=n) for n in "abc")
|
||||
res = client.post("/api/jobs/queue", json={"job_ids": [b["id"], c["id"], a["id"]]})
|
||||
|
||||
assert res.status_code == 200
|
||||
assert [j["name"] for j in res.json()] == ["b", "c", "a"]
|
||||
assert _statuses(client, a, b, c) == ["queued", "running", "queued"]
|
||||
_age(b, 10)
|
||||
assert _statuses(client, a, b, c) == ["queued", "completed", "running"]
|
||||
|
||||
|
||||
def test_a_long_idle_queue_is_all_done_on_the_next_poll(client):
|
||||
clips = [_create(client, name=f"c{i}") for i in range(3)]
|
||||
client.post("/api/jobs/queue", json={"job_ids": [c["id"] for c in clips]})
|
||||
for record in mock_server._jobs.values(): # everything happened a minute ago
|
||||
if record.get("queued_at"):
|
||||
record["queued_at"] -= 60
|
||||
if record.get("started_at"):
|
||||
record["started_at"] -= 60
|
||||
assert _statuses(client, *clips) == ["completed"] * 3
|
||||
starts = [mock_server._jobs[c["id"]]["started_at"] for c in clips]
|
||||
assert starts == sorted(starts) # each took over when the previous finished
|
||||
|
||||
|
||||
def test_a_clip_waits_for_the_clip_whose_last_frame_it_uses(client):
|
||||
one = _create(client, name="one")
|
||||
two = _create(client, name="two", references=_frame_of(one))
|
||||
client.post("/api/jobs/queue", json={"job_ids": [one["id"], two["id"]]})
|
||||
assert _statuses(client, one, two) == ["running", "queued"]
|
||||
_age(one, 10)
|
||||
assert _statuses(client, one, two) == ["completed", "running"]
|
||||
|
||||
|
||||
def test_failure_of_a_frame_source_fails_the_clips_after_it(client):
|
||||
one = _create(client, name="Clip 1")
|
||||
two = _create(client, name="Clip 2", references=_frame_of(one))
|
||||
three = _create(client, name="Clip 3", references=_frame_of(two))
|
||||
mock_server._jobs[one["id"]]["status"] = "failed"
|
||||
client.post("/api/jobs/queue", json={"job_ids": [two["id"], three["id"]]})
|
||||
|
||||
assert _statuses(client, two, three) == ["failed", "failed"]
|
||||
error = client.get(f"/api/jobs/{two['id']}").json()["error"]
|
||||
assert "Clip 1" in error and "failed" in error
|
||||
|
||||
|
||||
def test_deleting_a_frame_source_fails_the_clips_waiting_on_it(client):
|
||||
blocker = _create(client)
|
||||
one = _create(client, name="one")
|
||||
two = _create(client, name="two", references=_frame_of(one))
|
||||
client.post("/api/jobs/queue", json={"job_ids": [blocker["id"], one["id"], two["id"]]})
|
||||
assert client.delete(f"/api/jobs/{one['id']}").status_code == 200
|
||||
assert _statuses(client, two) == ["failed"]
|
||||
assert "deleted" in client.get(f"/api/jobs/{two['id']}").json()["error"]
|
||||
|
||||
|
||||
def test_dequeue_and_stop_return_a_queued_job_to_pending(client):
|
||||
blocker, a, b = _create(client), _create(client), _create(client)
|
||||
client.post("/api/jobs/queue", json={"job_ids": [blocker["id"], a["id"], b["id"]]})
|
||||
|
||||
assert client.post(f"/api/jobs/{a['id']}/dequeue").json()["status"] == "pending"
|
||||
assert client.post(f"/api/jobs/{b['id']}/stop").json()["status"] == "pending"
|
||||
assert client.post(f"/api/jobs/{a['id']}/dequeue").status_code == 409
|
||||
assert client.post("/api/jobs/nope/dequeue").status_code == 404
|
||||
|
||||
|
||||
def test_queueing_is_all_or_nothing_and_refuses_started_jobs(client):
|
||||
a, done = _create(client), _create(client)
|
||||
mock_server._jobs[done["id"]]["status"] = "completed"
|
||||
res = client.post("/api/jobs/queue", json={"job_ids": [a["id"], done["id"]]})
|
||||
assert res.status_code == 409 and "already completed" in res.json()["detail"]
|
||||
assert _statuses(client, a) == ["pending"]
|
||||
assert client.post("/api/jobs/queue", json={"job_ids": [a["id"], "nope"]}).status_code == 404
|
||||
assert client.post("/api/jobs/queue", json={"job_ids": [a["id"], a["id"]]}).status_code == 409
|
||||
|
||||
assert client.post(f"/api/jobs/{a['id']}/queue").status_code == 200
|
||||
assert client.post(f"/api/jobs/{a['id']}/queue").status_code == 409 # running now
|
||||
|
||||
|
||||
def test_starting_a_clip_before_its_frame_source_finished_is_refused(client):
|
||||
one = _create(client, name="Clip 1")
|
||||
two = _create(client, references=_frame_of(one))
|
||||
res = client.post(f"/api/jobs/{two['id']}/start")
|
||||
assert res.status_code == 409
|
||||
assert "Clip 1" in res.json()["detail"] and "hasn't finished" in res.json()["detail"]
|
||||
|
||||
mock_server._jobs[one["id"]]["status"] = "completed"
|
||||
assert client.post(f"/api/jobs/{two['id']}/start").status_code == 200
|
||||
|
||||
|
||||
def test_queued_jobs_cannot_be_edited(client):
|
||||
blocker, a = _create(client), _create(client)
|
||||
client.post("/api/jobs/queue", json={"job_ids": [blocker["id"], a["id"]]})
|
||||
assert client.patch(f"/api/jobs/{a['id']}", json={"prompt": "x"}).status_code == 400
|
||||
|
||||
|
||||
# --- merging a scene -----------------------------------------------------------
|
||||
|
||||
|
||||
def _completed(client, name):
|
||||
job = _create(client, name=name)
|
||||
record = mock_server._jobs[job["id"]]
|
||||
record.update(status="completed", output_path=f"/mock/outputs/{job['id']}/output.mp4", finished_at=1.0)
|
||||
return job
|
||||
|
||||
|
||||
def test_merges_finished_clips_into_one_downloadable_video(client):
|
||||
clips = [_completed(client, f"wolf-{i}") for i in range(3)]
|
||||
res = client.post("/api/scenes/merge", json={"job_ids": [c["id"] for c in clips], "name": "wolf"})
|
||||
|
||||
assert res.status_code == 200
|
||||
body = res.json()
|
||||
assert body["clips"] == 3 and body["filename"].startswith("wolf-") and body["seconds"] == pytest.approx(3.0, abs=0.1)
|
||||
video = client.get(body["url"])
|
||||
assert video.status_code == 200 and video.headers["content-type"] == "video/mp4"
|
||||
assert len(video.content) > 1000
|
||||
|
||||
|
||||
def test_refuses_to_merge_a_scene_with_an_unfinished_clip(client):
|
||||
done, pending = _completed(client, "a"), _create(client, name="b")
|
||||
res = client.post("/api/scenes/merge", json={"job_ids": [done["id"], pending["id"]]})
|
||||
assert res.status_code == 409 and "'b' hasn't finished" in res.json()["detail"]
|
||||
|
||||
|
||||
def test_merge_of_an_unknown_job_and_of_nothing(client):
|
||||
assert client.post("/api/scenes/merge", json={"job_ids": ["nope"]}).status_code == 404
|
||||
assert client.post("/api/scenes/merge", json={"job_ids": []}).status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filename", ["nope.mp4", "..%2Fetc%2Fpasswd", "x.txt"])
|
||||
def test_an_unknown_merged_file_is_not_found(client, filename):
|
||||
assert client.get(f"/api/merged/{filename}").status_code == 404
|
||||
Reference in New Issue
Block a user