From 3798ce193ee74c79c62730f607ba8b6bdb9ebb8e Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 11 Jun 2026 15:20:43 +0300 Subject: [PATCH] Stream partial chain scene previews while multi-scene runs are in flight. After each scene finishes sampling, decode a downscaled temp mp4 and publish progress so the Movie Editor can update timeline and preview before the full chain completes. --- movie_editor/backend/chain_progress.py | 66 ++++++++++++++++++ movie_editor/frontend/store.js | 21 ++++-- movie_editor/server.py | 18 ++++- movie_editor/tests/test_chain_progress.py | 22 ++++++ samplers.py | 81 +++++++++++++++++++++++ 5 files changed, 203 insertions(+), 5 deletions(-) create mode 100644 movie_editor/backend/chain_progress.py create mode 100644 movie_editor/tests/test_chain_progress.py diff --git a/movie_editor/backend/chain_progress.py b/movie_editor/backend/chain_progress.py new file mode 100644 index 00000000..f80e60d7 --- /dev/null +++ b/movie_editor/backend/chain_progress.py @@ -0,0 +1,66 @@ +"""In-flight chain sampler progress for the Movie Editor. + +While a multi-scene FunPackLTXAVSceneChainSampler run is executing, partial decodes +are published here so the editor can show finished scenes before the full run ends. +""" +from __future__ import annotations + +import threading +import time +from typing import Any, Optional + +_lock = threading.Lock() +_by_key: dict[str, dict[str, Any]] = {} +_prompt_to_key: dict[str, str] = {} + + +def begin(prompt_id: str, progress_key: str, scene_count: int) -> None: + if not prompt_id or not progress_key or scene_count <= 1: + return + with _lock: + _prompt_to_key[prompt_id] = progress_key + _by_key[progress_key] = { + "scene_count": int(scene_count), + "completed_scenes": 0, + "media": None, + "ts": time.time(), + } + + +def update(progress_key: str, completed_scenes: int, media: dict) -> None: + if not progress_key or not media: + return + with _lock: + slot = _by_key.get(progress_key) + if not slot: + return + slot["completed_scenes"] = max(int(completed_scenes), int(slot.get("completed_scenes") or 0)) + slot["media"] = dict(media) + slot["ts"] = time.time() + + +def read_for_prompt(prompt_id: str) -> Optional[dict]: + with _lock: + key = _prompt_to_key.get(prompt_id) + if not key: + return None + slot = _by_key.get(key) + if not slot or not slot.get("media"): + return None + completed = int(slot.get("completed_scenes") or 0) + total = int(slot.get("scene_count") or 0) + if completed <= 0 or total <= 1 or completed >= total: + return None + return { + "partial": True, + "completed_scenes": completed, + "scene_count": total, + "media": [dict(slot["media"])], + } + + +def finish(prompt_id: str) -> None: + with _lock: + key = _prompt_to_key.pop(prompt_id, None) + if key: + _by_key.pop(key, None) diff --git a/movie_editor/frontend/store.js b/movie_editor/frontend/store.js index 4b43fdb6..498b3893 100644 --- a/movie_editor/frontend/store.js +++ b/movie_editor/frontend/store.js @@ -1338,16 +1338,19 @@ } } - function _recordSegment(mediaList, targetSceneIds) { + function _recordSegment(mediaList, targetSceneIds, opts) { if (!mediaList || !mediaList.length || !targetSceneIds || !targetSceneIds.length) return; + const completedScenes = opts?.completedScenes ?? targetSceneIds.length; + const ids = targetSceneIds.slice(0, Math.max(0, completedScenes)); + if (!ids.length) return; const primary = mediaList.find((m) => m.kind === "videos" || m.kind === "gifs") || mediaList[0]; - const layout = _chainRunLayout(targetSceneIds); + const layout = _chainRunLayout(ids); let sourceEnd = 0; let lastChain = null; let clearedRating = false; const recordedSceneIds = []; - for (let i = 0; i < targetSceneIds.length; i++) { - const id = targetSceneIds[i]; + for (let i = 0; i < ids.length; i++) { + const id = ids[i]; if (_removedDuringGen.has(id)) { const ghosts = state.sceneGhosts || []; const gi = ghosts.findIndex((g) => g.id === id); @@ -1384,6 +1387,7 @@ // Poll a single queued prompt to completion. Resolves true on success, false on error. function _pollPromise(promptId, targetSceneIds, prefix) { prefix = prefix || "Generating…"; + let lastPartialScenes = 0; return new Promise((resolve) => { _clearGenTimers(); let pendingStreak = 0; @@ -1424,6 +1428,15 @@ } else { // "pending" after "running" means the job left the queue without a history // entry — it likely crashed or was interrupted by ComfyUI. + if (s.partial?.media?.length && s.partial.completed_scenes > lastPartialScenes) { + lastPartialScenes = s.partial.completed_scenes; + _recordSegment(s.partial.media, targetSceneIds, { completedScenes: lastPartialScenes }); + notify(); + const step = (state.gen.maxStep > 0) ? ` · sampling ${state.gen.step}/${state.gen.maxStep}` : ""; + updateGenProgress({ + msg: `${prefix} ${_elapsed()} · ${lastPartialScenes}/${targetSceneIds.length} scene(s) ready${step}`, + }); + } if (s.state === "pending") pendingStreak++; else pendingStreak = 0; if (pendingStreak >= 3) { _clearGenTimers(); diff --git a/movie_editor/server.py b/movie_editor/server.py index 8006fe20..a20a1da7 100644 --- a/movie_editor/server.py +++ b/movie_editor/server.py @@ -855,6 +855,11 @@ if web is not None and PromptServer is not None: media_pack = _prepare_media(target, continuity_media_refs(p, target)) sampler_inputs = _run_sampler_inputs(target, active_scene_count, full=p) _attach_scene_anchors(sampler_inputs, media_pack, target) + progress_key = "" + if active_scene_count > 1: + import uuid + progress_key = uuid.uuid4().hex + sampler_inputs["chain_progress_key"] = progress_key graph, report = builder.build(oi, _project_models(target), { "prompt": prompt, "seed": _resolve_run_seed(target), "num_frames_per_scene": effective_frames, @@ -873,6 +878,10 @@ if web is not None and PromptServer is not None: result = await bridge.queue_prompt(graph) except Exception as e: # noqa: BLE001 return web.json_response({"detail": f"Failed to queue with ComfyUI: {e}"}, status=502) + prompt_id = result.get("prompt_id") + if progress_key and prompt_id: + from .backend import chain_progress + chain_progress.begin(prompt_id, progress_key, active_scene_count) p.generation_meta = { **(p.generation_meta or {}), "prompt_hash": validation["prompt_hash"], @@ -907,6 +916,8 @@ if web is not None and PromptServer is not None: exec_error = f"{node_type}: {exc}" break is_error = raw_status.get("status_str") == "error" + from .backend import chain_progress + chain_progress.finish(prompt_id) return web.json_response({ "state": "error" if (is_error and not media) else "completed", "media": media, @@ -917,7 +928,12 @@ if web is not None and PromptServer is not None: running = await bridge.is_running(prompt_id) except Exception: # noqa: BLE001 running = True - return web.json_response({"state": "running" if running else "pending", "media": []}) + from .backend import chain_progress + partial = chain_progress.read_for_prompt(prompt_id) + payload = {"state": "running" if running else "pending", "media": []} + if partial: + payload["partial"] = partial + return web.json_response(payload) @routes.get(UI_PREFIX + "/api/rating-labels") async def _rating_labels(_req): diff --git a/movie_editor/tests/test_chain_progress.py b/movie_editor/tests/test_chain_progress.py new file mode 100644 index 00000000..468f50e6 --- /dev/null +++ b/movie_editor/tests/test_chain_progress.py @@ -0,0 +1,22 @@ +from movie_editor.backend.chain_progress import begin, finish, read_for_prompt, update + + +def test_chain_progress_partial_read(): + begin("pid1", "run-a", 3) + update("run-a", 1, {"kind": "videos", "filename": "partial_s1.mp4", "subfolder": "funpack_chain_progress/run-a", "type": "temp"}) + snap = read_for_prompt("pid1") + assert snap is not None + assert snap["partial"] is True + assert snap["completed_scenes"] == 1 + assert snap["scene_count"] == 3 + assert snap["media"][0]["filename"] == "partial_s1.mp4" + update("run-a", 3, {"kind": "videos", "filename": "partial_s3.mp4", "subfolder": "x", "type": "temp"}) + assert read_for_prompt("pid1") is None + finish("pid1") + assert read_for_prompt("pid1") is None + + +def test_chain_progress_single_scene_skipped(): + begin("pid2", "run-b", 1) + assert read_for_prompt("pid2") is None + finish("pid2") diff --git a/samplers.py b/samplers.py index 9792a97b..709da2cf 100644 --- a/samplers.py +++ b/samplers.py @@ -2095,6 +2095,11 @@ class FunPackLTXAVSceneChainSampler: "multiline": True, "tooltip": "Optional media_ref → filename map for image-type i2v guides in custom guide stacks.", }), + "chain_progress_key": ("STRING", { + "default": "", + "multiline": False, + "tooltip": "Movie Editor only: token for publishing partial scene previews while a multi-scene chain run is in flight.", + }), }, "hidden": { "unique_id": "UNIQUE_ID", @@ -3293,6 +3298,7 @@ class FunPackLTXAVSceneChainSampler: refinement_key_input="", funpack_scene_guides="", funpack_scene_anchors="", funpack_scene_media_refs="", + chain_progress_key="", unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -3520,6 +3526,12 @@ class FunPackLTXAVSceneChainSampler: f"Scene {scene_index + 1}: seed={scene_seed}, text={scene_meta['text']}" + (f" | encode≠text" if scene_meta["encode_text"] != scene_meta["text"] else "") ) + if chain_progress_key and scene_count > 1: + self._save_chain_progress_preview( + vae, output, chain_progress_key, scene_index, scene_count, + decode_tile_size=decode_tile_size, + frame_rate=24, + ) del scene_cond, scene_positive, scene_negative, scene_conditionings, chunk, sampled @@ -3630,6 +3642,75 @@ class FunPackLTXAVSceneChainSampler: return None return [groups[k] for k in sorted(groups)] + def _save_chain_progress_preview(self, vae, output, progress_key, scene_index, scene_count, + decode_tile_size=0, frame_rate=24): + """Decode cumulative chain output to a small temp mp4 for Movie Editor preview.""" + if not progress_key or scene_count <= 1 or scene_index >= scene_count - 1: + return + try: + try: + from movie_editor.backend import chain_progress as _cp + except ImportError: + from .movie_editor.backend import chain_progress as _cp # type: ignore + except Exception: + return + import os + import re as _re + import shutil + import subprocess + import tempfile + try: + import folder_paths + base = folder_paths.get_temp_directory() + except Exception: + return + rel_dir = os.path.join("funpack_chain_progress", _re.sub(r"[^A-Za-z0-9_.-]", "_", progress_key)) + out_dir = os.path.join(base, rel_dir) + os.makedirs(out_dir, exist_ok=True) + fn = f"partial_s{scene_index + 1}.mp4" + out_path = os.path.join(out_dir, fn) + try: + decoded = self._decode_for_preview(vae, output, decode_tile_size=decode_tile_size or 512) + if decoded is None: + return + t = decoded.detach().float().clamp(0, 1).cpu() + if t.dim() == 4 and t.shape[-1] not in (1, 3): + t = t.permute(0, 2, 3, 1) + frames = t.numpy() + tmp = tempfile.mkdtemp(prefix="funpack_chain_prev_") + try: + from PIL import Image + import numpy as _np + for i, f in enumerate(frames): + arr = (_np.asarray(f) * 255).astype(_np.uint8) + if arr.shape[-1] == 1: + arr = arr.repeat(3, axis=-1) + Image.fromarray(arr).save(os.path.join(tmp, f"{i:04d}.png")) + ff = shutil.which("ffmpeg") + if not ff: + return + subprocess.run( + [ + ff, "-y", "-hide_banner", "-loglevel", "error", + "-framerate", str(max(1, int(frame_rate))), + "-i", os.path.join(tmp, "%04d.png"), + "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart", + out_path, + ], + check=True, + capture_output=True, + ) + finally: + shutil.rmtree(tmp, ignore_errors=True) + _cp.update(progress_key, scene_index + 1, { + "kind": "videos", + "filename": fn, + "subfolder": rel_dir.replace("\\", "/"), + "type": "temp", + }) + except Exception as e: + print(f"[FunPackSceneChain] chain progress preview failed (scene {scene_index + 1}): {e}") + def _save_batch_preview(self, decoded, path, max_frames=16, width=256): """Save a decoded video tensor [T,H,W,C] in 0..1 as a downscaled animated webp.""" try: