Add rolling LTX video extension

2026-09-14 20:19:08 +02:00
parent 719bc4ccf7
commit aaa96bbed5
4 changed files with 489 additions and 2 deletions
@@ -24,3 +24,50 @@ normal client browser through a local SSH tunnel to the container instead of
 requiring a browser or storing login credentials inside the VNC desktop.
 The container port is published only on Athena's `127.0.0.1`; it is not exposed
 on Athena's LAN or WireGuard interfaces.
+# VNC-Oberfläche
+
+Die Linux-Oberfläche unter noVNC kann im Athena-Dashboard schlafen gelegt und
+wieder aufgeweckt werden. Im Schlafzustand sind die Electron-Unterprozesse der
+Linux-Oberfläche angehalten. Der Electron-Hauptprozess, das LTX-Backend, laufende
+Generierungen und die native macOS-App bleiben aktiv. Der Status wird im
+Controller zwischengespeichert, damit die Dashboard-Abfrage keine fortlaufenden
+Docker-Exec-Prozesse erzeugt. Beim erneuten Start des LTX-Profils ist die
+VNC-Oberfläche zunächst aktiv.
+
+## Lokaler Prompt Enhancer auf der RTX 3060
+
+Der optionale Gemma-4-E2B-Prompt-Enhancer läuft getrennt von der
+Videogenerierung auf `cuda:1` (RTX 3060); LTX 2.5 bleibt auf `cuda:0`
+(RTX 5080). Das Modell liegt unter `models/gemma-4-E2B-it`, wird nur für einen
+Enhance-Aufruf geladen und danach wieder entladen. Gemessen wurden rund
+9.97 GiB VRAM und 19 Sekunden für Text sowie 9.95 GiB und 10 Sekunden mit
+Referenzbild.
+
+Die vorherige Einzel-GPU-Konfiguration liegt unter
+`/opt/mike-ai/rollback/ltx-gemma-e2b-3060-20260914/`. Rückbau:
+
+```sh
+/opt/mike-ai/rollback/ltx-gemma-e2b-3060-20260914/rollback.sh
+```
+
+## Rollierendes Extend-Fenster
+
+Lokales Extend verarbeitet auf der RTX 5080 höchstens 433 Frames (rund 18
+Sekunden bei 24 fps). Sobald Quelle plus Verlängerung größer werden, verwendet
+das Backend nur bis zu acht Sekunden vom Anfang oder Ende als Übergangskontext.
+Nach der Generierung wird ausschließlich der neue Abschnitt per FFmpeg an die
+vollständige Quelle gesetzt. Dadurch kann ein Ergebnis erneut verlängert
+werden, ohne dass der GPU-Speicherbedarf mit der gesamten Filmlänge wächst.
+
+Die Oberfläche wird unverändert verwendet: fertigen Clip auswählen, `Extend`
+wählen und bis zu zwölf Sekunden ergänzen. Bei zwölf Sekunden reduziert das
+Backend den Kontext automatisch auf ungefähr sechs Sekunden. Die abschließende
+Zusammensetzung läuft auf der CPU und kann nach der GPU-Generierung noch einige
+Sekunden dauern. Der Übergang kann bei sehr langen Ketten allmählich driften;
+für einen neuen Szenenabschnitt ist daher ein klarer Prompt sinnvoll.
+
+Rückbau:
+
+```sh
+/opt/mike-ai/rollback/ltx-rolling-extend-20260914/rollback.sh
+```
@@ -17,20 +17,28 @@ services:
            - driver: nvidia
              device_ids:
                - ${LTX2_GPU_UUID:-GPU-8ad38c6c-5a01-9d8e-1dfa-ed662ad78fbe}
+                - ${LTX2_PROMPT_ENHANCER_GPU_UUID:-GPU-4834d9d7-5b61-3004-1fb3-4ae49d482d4b}
              capabilities: [gpu]
    shm_size: 8g
    environment:
-      NVIDIA_VISIBLE_DEVICES: ${LTX2_GPU_UUID:-GPU-8ad38c6c-5a01-9d8e-1dfa-ed662ad78fbe}
+      NVIDIA_VISIBLE_DEVICES: ${LTX2_GPU_UUID:-GPU-8ad38c6c-5a01-9d8e-1dfa-ed662ad78fbe},${LTX2_PROMPT_ENHANCER_GPU_UUID:-GPU-4834d9d7-5b61-3004-1fb3-4ae49d482d4b}
      NVIDIA_DRIVER_CAPABILITIES: compute,utility,graphics
-      CUDA_VISIBLE_DEVICES: "0"
+      CUDA_VISIBLE_DEVICES: "0,1"
+      LTX_PROMPT_ENHANCER_DEVICE: cuda:1
      # Stable and intentionally different from the default local macOS app
      # port, so OAuth can return through an SSH tunnel without a collision.
      LTX_PORT: "41955"
      LTX_HOST: 0.0.0.0
+      LTX_REMOTE_AUTH_TOKEN_FILE: /run/secrets/ltx-remote-token
    ports:
      - 127.0.0.1:41955:41955
    volumes:
      - /data/video/ltx-desktop:/data
+      - /etc/mike-ai/secrets/ltx-remote-token:/run/secrets/ltx-remote-token:ro
+      - /opt/mike-ai/stack/platform/ltx2-studio/remote-app_factory.py:/opt/ltx-desktop/squashfs-root/resources/backend/app_factory.py:ro
+      - /opt/mike-ai/stack/platform/ltx2-studio/prompt_enhancement_handler.py:/opt/ltx-desktop/squashfs-root/resources/backend/handlers/prompt_enhancement_handler.py:ro
+      - /opt/mike-ai/stack/platform/ltx2-studio/extend_handler.py:/opt/ltx-desktop/squashfs-root/resources/backend/handlers/extend_handler.py:ro
+      - /opt/mike-ai/stack/platform/ltx2-studio/rolling_extend.py:/opt/ltx-desktop/squashfs-root/resources/backend/handlers/rolling_extend.py:ro
    networks:
      - frontend
    security_opt:
@@ -0,0 +1,289 @@
+"""Extend API orchestration handler.
+
+Mirrors ``RetakeHandler``'s dual API/local dispatch. Extend appends (``mode="end"``)
+or prepends (``mode="start"``) freshly generated frames to a source video. The cloud
+``/v1/extend`` endpoint and the local PyTorch wrapper share this entry point.
+"""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from threading import RLock
+
+from api_types import (
+    ExtendMode,
+    ExtendRequest,
+    ExtendResponse,
+    RetakeCancelledResponse,
+    RetakeExtendModel,
+    RetakePayloadResponse,
+    RetakeVideoResponse,
+    TargetResolution,
+)
+from _routes._errors import HTTPError
+from api_model_specs import FORCED_API_MODEL_MAP
+from handlers.base import StateHandlerBase
+from handlers.generation_handler import GenerationHandler
+from handlers.rolling_extend import assemble_rolling_result, extract_context, plan_rolling_extend
+from handlers.pipelines_handler import PipelinesHandler
+from handlers.text_handler import TextHandler
+from handlers.video_resolution import (
+    TIME_FACTOR,
+    correct_frame_count,
+    read_source_metadata,
+    resolve_target_resolution,
+    validate_source_video_path,
+)
+from runtime_config.ltx_capabilities import local_caps, supports
+from runtime_config.model_download_specs import resolve_active_ltx_model_id
+from runtime_config.runtime_config import RuntimeConfig
+from services.generation_interrupt import GenerationCancelledError, is_cancel_exception
+from services.ltx_api_client.ltx_api_client import LTXAPIClientError
+from services.interfaces import LTXAPIClient
+from state.app_state_types import AppState
+from state.app_settings import should_video_generate_with_ltx_api
+
+# Cloud caps a single extend at 20s; mirror it locally. Minimum is 2s (cloud min).
+_MIN_DURATION = 2.0
+_MAX_DURATION = 20.0
+
+
+class ExtendHandler(StateHandlerBase):
+    def __init__(
+        self,
+        state: AppState,
+        lock: RLock,
+        ltx_api_client: LTXAPIClient,
+        config: RuntimeConfig,
+        generation_handler: GenerationHandler,
+        pipelines_handler: PipelinesHandler,
+        text_handler: TextHandler,
+    ) -> None:
+        super().__init__(state, lock, config)
+        self._ltx_api_client = ltx_api_client
+        self._generation = generation_handler
+        self._pipelines = pipelines_handler
+        self._text = text_handler
+
+    def run(self, req: ExtendRequest) -> ExtendResponse:
+        video_path = req.video_path
+        duration = req.duration
+        prompt = req.prompt
+        mode = req.mode
+
+        if duration < _MIN_DURATION:
+            raise HTTPError(400, f"duration must be at least {int(_MIN_DURATION)} seconds")
+        if duration > _MAX_DURATION:
+            raise HTTPError(400, f"duration must be at most {int(_MAX_DURATION)} seconds")
+
+        video_file = validate_source_video_path(video_path)
+
+        if should_video_generate_with_ltx_api(
+            force_api_generations=self.config.force_api_generations,
+            settings=self.state.app_settings,
+        ):
+            # The cloud preserves source resolution (no resolution param); resolution
+            # selection is local-only.
+            return self._run_api_extend(
+                video_file=video_file, duration=duration, prompt=prompt, mode=mode, model=req.model,
+            )
+
+        model_id = resolve_active_ltx_model_id(
+            self.models_dir, self.state.app_settings.active_ltx_model_id
+        )
+        if model_id is None:
+            raise HTTPError(409, "NO_DOWNLOADED_LTX_MODEL")
+        if not supports(local_caps(model_id), "extend"):
+            raise HTTPError(
+                409,
+                "Extend is not supported for the active LTX model.",
+                code="UNSUPPORTED_EXTEND",
+            )
+
+        return self._run_local_extend(
+            video_file=video_file, duration=duration, prompt=prompt, mode=mode, resolution=req.resolution
+        )
+
+    def _run_api_extend(
+        self,
+        *,
+        video_file: Path,
+        duration: float,
+        prompt: str,
+        mode: ExtendMode,
+        model: RetakeExtendModel,
+    ) -> ExtendResponse:
+        api_key = self.state.app_settings.ltx_api_key
+        if not api_key:
+            raise HTTPError(400, "LTX API key not configured. Set it in Settings.")
+
+        with self._generation.reserved_generation_start():
+
+            # Drive the generation state machine so the result is recoverable via
+            # /api/generation/progress if the page unmounts mid-generation (mirrors retake).
+            try:
+                self._generation.start_api_generation(uuid.uuid4().hex[:8])
+            except RuntimeError as exc:
+                # Lost a race with a concurrent generation between the check above and here;
+                # surface the same 409 as the guard, not a bare 500 (the running generation
+                # owns the state now, so don't fail_generation).
+                raise HTTPError(409, str(exc)) from exc
+            try:
+                self._generation.update_progress("inference", 55, None, None)
+                result = self._ltx_api_client.extend(
+                    api_key=api_key,
+                    video_path=str(video_file),
+                    duration=duration,
+                    prompt=prompt,
+                    mode=mode,
+                    model=FORCED_API_MODEL_MAP[model],
+                )
+
+                if result.video_bytes is not None:
+                    self._generation.update_progress("downloading_output", 85, None, None)
+                    output = self.config.outputs_dir / f"extend_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.mp4"
+                    try:
+                        with open(output, "wb") as out:
+                            out.write(result.video_bytes)
+                        self._generation.update_progress("complete", 100, None, None)
+                        self._generation.complete_generation(str(output))
+                    except Exception:
+                        output.unlink(missing_ok=True)  # don't strand a half-written .mp4
+                        raise
+                    return RetakeVideoResponse(status="complete", video_path=str(output))
+
+                if result.result_payload is not None:
+                    self._generation.update_progress("complete", 100, None, None)
+                    self._generation.complete_generation(None)
+                    return RetakePayloadResponse(status="complete", result=result.result_payload)
+
+                raise HTTPError(500, "Extend API returned no result")
+            except LTXAPIClientError as exc:
+                self._generation.fail_generation(exc.detail)
+                raise HTTPError(exc.status_code, exc.detail) from exc
+            except HTTPError as exc:
+                self._generation.fail_generation(exc.detail)
+                raise
+            except Exception as exc:
+                self._generation.fail_generation(str(exc))
+                raise
+
+    def _run_local_extend(
+        self,
+        *,
+        video_file: Path,
+        duration: float,
+        prompt: str,
+        mode: ExtendMode,
+        resolution: TargetResolution | None,
+    ) -> ExtendResponse:
+        with self._generation.reserved_generation_start():
+
+            fps, source_width, source_height, source_frames = read_source_metadata(str(video_file))
+            target_frames = correct_frame_count(source_frames)
+            extend_frames = self._duration_to_extend_frames(duration, fps)
+            try:
+                rolling_plan = plan_rolling_extend(target_frames, extend_frames, fps)
+            except ValueError as exc:
+                raise HTTPError(400, str(exc)) from exc
+            target_width, target_height = resolve_target_resolution(resolution, source_width, source_height)
+
+            try:
+                self._text.prepare_text_encoding(prompt, enhance_prompt=False)
+            except RuntimeError as exc:
+                raise HTTPError(400, str(exc)) from exc
+
+            generation_id = uuid.uuid4().hex[:8]
+            seed = self._resolve_seed()
+            output_path = self.config.outputs_dir / f"extend_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{generation_id}.mp4"
+
+            try:
+                pipeline_state = self._pipelines.load_retake_pipeline(distilled=True)
+                self._generation.start_generation(generation_id)
+                self._generation.update_progress("loading_model", 5, 0, 1)
+                self._generation.update_progress("inference", 15, 0, 1)
+
+                if rolling_plan.rolling:
+                    with TemporaryDirectory(prefix="rolling-extend-", dir=self.config.outputs_dir) as temp:
+                        context_path = Path(temp) / "context.mp4"
+                        window_output = Path(temp) / "generated-window.mp4"
+                        extract_context(
+                            video_file,
+                            context_path,
+                            mode=mode,
+                            fps=fps,
+                            source_frames=target_frames,
+                            context_frames=rolling_plan.context_frames,
+                        )
+                        pipeline_state.pipeline.extend(
+                            video_path=str(context_path),
+                            prompt=prompt,
+                            extend_frames=extend_frames,
+                            mode=mode,
+                            seed=seed,
+                            output_path=str(window_output),
+                            negative_prompt=self.config.default_negative_prompt,
+                            regenerate_audio=True,
+                            enhance_prompt=False,
+                            distilled=True,
+                            target_width=target_width,
+                            target_height=target_height,
+                            target_frames=rolling_plan.context_frames,
+                        )
+                        self._generation.update_progress("assembling_video", 95, 0, 1)
+                        assemble_rolling_result(
+                            video_file,
+                            window_output,
+                            output_path,
+                            mode=mode,
+                            fps=fps,
+                            source_frames=target_frames,
+                            context_frames=rolling_plan.context_frames,
+                            extend_frames=extend_frames,
+                        )
+                else:
+                    pipeline_state.pipeline.extend(
+                        video_path=str(video_file),
+                        prompt=prompt,
+                        extend_frames=extend_frames,
+                        mode=mode,
+                        seed=seed,
+                        output_path=str(output_path),
+                        negative_prompt=self.config.default_negative_prompt,
+                        regenerate_audio=True,
+                        enhance_prompt=False,
+                        distilled=True,
+                        target_width=target_width,
+                        target_height=target_height,
+                        target_frames=target_frames,
+                    )
+
+                # Denoiser interrupt cannot abort VAE decode / ffmpeg; a Stop after the last
+                # denoise step still finishes encode, then this check drops the file.
+                if self._generation.is_generation_cancelled():
+                    output_path.unlink(missing_ok=True)
+                    raise GenerationCancelledError()
+
+                self._generation.update_progress("complete", 100, 1, 1)
+                self._generation.complete_generation(str(output_path))
+                return RetakeVideoResponse(status="complete", video_path=str(output_path))
+            except HTTPError:
+                self._generation.fail_generation("Extend generation failed")
+                raise
+            except Exception as exc:
+                self._generation.fail_generation(str(exc))
+                if is_cancel_exception(exc):
+                    return RetakeCancelledResponse(status="cancelled")
+                raise HTTPError(500, f"Generation error: {exc}") from exc
+            finally:
+                self._text.clear_api_embeddings()
+
+    @staticmethod
+    def _duration_to_extend_frames(duration: float, fps: float) -> int:
+        # Snap up to the nearest multiple of TIME_FACTOR so the output stays 8k+1: the
+        # source is 8k+1, so adding a multiple of 8 keeps the padded latent integer-sized.
+        frames = round(duration * fps)
+        return ((frames + TIME_FACTOR - 1) // TIME_FACTOR) * TIME_FACTOR
@@ -0,0 +1,143 @@
+"""Bounded-context local video extension and loss-minimizing assembly helpers."""
+
+from __future__ import annotations
+
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+
+import av
+from imageio_ffmpeg import get_ffmpeg_exe
+
+from api_types import ExtendMode
+
+# LTX video lengths must be 8n+1 frames. Keep this helper independent from the
+# handlers package so it can also be exercised by maintenance and migration
+# scripts without importing every request handler.
+TIME_FACTOR = 8
+MIN_FRAMES = 9
+
+
+def correct_frame_count(frames: int) -> int:
+    return max(MIN_FRAMES, (frames // TIME_FACTOR) * TIME_FACTOR + 1)
+
+_MAX_WINDOW_SECONDS = 18.0
+_MAX_CONTEXT_SECONDS = 8.0
+_MIN_CONTEXT_SECONDS = 2.0
+
+
+@dataclass(frozen=True)
+class RollingExtendPlan:
+    rolling: bool
+    context_frames: int
+    total_generation_frames: int
+
+
+def plan_rolling_extend(source_frames: int, extend_frames: int, fps: float) -> RollingExtendPlan:
+    """Keep the GPU job bounded while retaining as much source context as possible."""
+    max_total = ((round(_MAX_WINDOW_SECONDS * fps)) // TIME_FACTOR) * TIME_FACTOR + 1
+    if source_frames + extend_frames <= max_total:
+        return RollingExtendPlan(False, source_frames, source_frames + extend_frames)
+
+    max_context = correct_frame_count(round(_MAX_CONTEXT_SECONDS * fps) + 1)
+    available_context = max_total - extend_frames
+    if available_context < round(_MIN_CONTEXT_SECONDS * fps) + 1:
+        raise ValueError(
+            "The requested extension is too long for the rolling VRAM window. "
+            "Choose a shorter duration."
+        )
+    context_frames = min(source_frames, max_context, correct_frame_count(available_context))
+    return RollingExtendPlan(True, context_frames, context_frames + extend_frames)
+
+
+def extract_context(
+    source: Path,
+    destination: Path,
+    *,
+    mode: ExtendMode,
+    fps: float,
+    source_frames: int,
+    context_frames: int,
+) -> None:
+    start_frame = 0 if mode == "start" else source_frames - context_frames
+    start_seconds = max(0.0, start_frame / fps)
+    duration_seconds = context_frames / fps
+    _run_ffmpeg(
+        "-i", str(source),
+        "-ss", f"{start_seconds:.9f}",
+        "-t", f"{duration_seconds:.9f}",
+        "-map", "0:v:0", "-map", "0:a:0?",
+        "-c:v", "libx264", "-preset", "fast", "-crf", "16",
+        "-pix_fmt", "yuv420p", "-c:a", "aac", "-b:a", "192k",
+        "-movflags", "+faststart", "-y", str(destination),
+    )
+
+
+def assemble_rolling_result(
+    source: Path,
+    generated_window: Path,
+    destination: Path,
+    *,
+    mode: ExtendMode,
+    fps: float,
+    source_frames: int,
+    context_frames: int,
+    extend_frames: int,
+) -> None:
+    """Join only the newly generated window edge to the untouched logical source."""
+    if mode == "end":
+        window_video = f"trim=start_frame={context_frames}"
+        window_audio = f"atrim=start={context_frames / fps:.9f}"
+        video_inputs = "[v0][v1]"
+        audio_inputs = "[a0][a1]"
+    else:
+        window_video = f"trim=end_frame={extend_frames}"
+        window_audio = f"atrim=end={extend_frames / fps:.9f}"
+        video_inputs = "[v1][v0]"
+        audio_inputs = "[a1][a0]"
+
+    have_audio = _has_audio(source) and _has_audio(generated_window)
+    filters = [
+        f"[0:v]trim=end_frame={source_frames},fps={fps:.9f},format=yuv420p,setpts=PTS-STARTPTS[v0]",
+        f"[1:v]{window_video},fps={fps:.9f},format=yuv420p,setpts=PTS-STARTPTS[v1]",
+        f"{video_inputs}concat=n=2:v=1:a=0[v]",
+    ]
+    args = ["-i", str(source), "-i", str(generated_window)]
+    if have_audio:
+        filters.extend(
+            [
+                f"[0:a]atrim=end={source_frames / fps:.9f},aresample=48000,asetpts=PTS-STARTPTS[a0]",
+                f"[1:a]{window_audio},aresample=48000,asetpts=PTS-STARTPTS[a1]",
+                f"{audio_inputs}concat=n=2:v=0:a=1[a]",
+            ]
+        )
+    args.extend(["-filter_complex", ";".join(filters), "-map", "[v]"])
+    if have_audio:
+        args.extend(["-map", "[a]", "-c:a", "aac", "-b:a", "192k"])
+    else:
+        args.append("-an")
+    args.extend(
+        [
+            "-c:v", "libx264", "-preset", "fast", "-crf", "16",
+            "-pix_fmt", "yuv420p", "-r", f"{fps:.9f}",
+            "-movflags", "+faststart", "-y", str(destination),
+        ]
+    )
+    _run_ffmpeg(*args)
+
+
+def _has_audio(path: Path) -> bool:
+    with av.open(str(path)) as container:
+        return any(stream.type == "audio" for stream in container.streams)
+
+
+def _run_ffmpeg(*args: str) -> None:
+    process = subprocess.run(
+        [get_ffmpeg_exe(), "-hide_banner", "-loglevel", "error", *args],
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    if process.returncode != 0:
+        detail = process.stderr.strip()[-2000:] or f"ffmpeg exited with {process.returncode}"
+        raise RuntimeError(f"Rolling extend media assembly failed: {detail}")