From 2fa4b59c201494b0aed3b75899b396c8bd317c9f Mon Sep 17 00:00:00 2001 From: Lukas Schmidt Date: Mon, 9 Nov 2026 19:39:16 +0100 Subject: [PATCH] Add DA3 depth preparation, automatic stereo passes and directory queue --- DA3_ARBEITSAUFTRAG.md | 38 +++ README.md | 113 +++++++++ TODO.md | 6 + __init__.py | 5 + da3_sequence.py | 413 ++++++++++++++++++++++++++++++++ directory_queue.py | 106 ++++++++ duke_stereo_streaming.py | 17 +- examples/da3_directory_api.json | 123 ++++++++++ stereoimage_generation.py | 6 +- tests/test_da3_sequence.py | 116 +++++++++ tests/test_directory_queue.py | 146 +++++++++++ 11 files changed, 1079 insertions(+), 10 deletions(-) create mode 100644 DA3_ARBEITSAUFTRAG.md create mode 100644 TODO.md create mode 100644 da3_sequence.py create mode 100644 directory_queue.py create mode 100644 examples/da3_directory_api.json create mode 100644 tests/test_da3_sequence.py create mode 100644 tests/test_directory_queue.py diff --git a/DA3_ARBEITSAUFTRAG.md b/DA3_ARBEITSAUFTRAG.md new file mode 100644 index 0000000..c657b41 --- /dev/null +++ b/DA3_ARBEITSAUFTRAG.md @@ -0,0 +1,38 @@ +# Arbeitsauftrag: zeitlich konsistente DA3-Tiefe für Stereo-Videos + +## Ziel + +Den bestehenden Stereo-Workflow um eine speichersparende Tiefenvorberechnung mit Depth Anything 3 und einen separaten Stereo-Node für vorbereitete Tiefenfolgen ergänzen. Die vorhandenen Nodes und Workflows sollen weiter funktionieren. Neue Nodes gehören in dieses Repository; die installierte `ComfyUI-DepthAnythingV3`-Erweiterung und VideoHelperSuite (VHS) sollen nach Möglichkeit wiederverwendet werden. + +## Gewünschter Ablauf + +1. Das Originalvideo über VHS-Meta-Batches abschnittsweise einlesen und DA3-Tiefe vorab für den ganzen Clip berechnen. Benachbarte Abschnitte sollen sich überlappen und anhand der gemeinsamen Frames ausgerichtet und im Übergang verblendet werden. +2. Zwischenstände auf Platte ablegen, sodass der RAM-Bedarf nicht mit dem gesamten Video und seinen vollständigen Ergebnisstapeln wächst. Nach dem letzten Batch eine gemeinsame, für Stereo geeignete Tiefenskalierung bestimmen und eine Tiefenfolge speichern. +3. In einem zweiten Durchlauf die Originalframes über den bestehenden Meta-Batch- und Upscaler-Pfad verarbeiten. Ein neuer, separater Stereo-Node liest dazu die passenden vorbereiteten Tiefenframes und schreibt das SBS-Video. Audio und bisherige Ausgabeparameter sollen erhalten bleiben. Der bestehende Stereo-Node wird nicht überladen oder ersetzt. + +Der neue Stereo-Pfad muss die gemeinsame Skalierung der Tiefenfolge erhalten: `apply_stereo_divergence` normalisiert bisher jede Tiefenkarte anhand ihres eigenen Minimums und Maximums. Für die vorbereitete Folge darf diese frameweise Neunormalisierung die zeitliche Konsistenz nicht wieder zerstören. Nähe/Ferne, Konvergenz und Tiefenbereich anhand eines tatsächlichen Ergebnisses überprüfen. + +## Technische Prüfung und Umsetzung + +- VHS-Meta-Batch-Lebenszyklus, Requeue, Abbruch, letzter unvollständiger Batch und Aufräumen temporärer Dateien berücksichtigen. Die Zuordnung von Originalvideo, Tiefenfolge und Audio muss bei mehreren Videos stimmen. +- 64 Frames **insgesamt** pro DA3-Abschnitt mit beispielsweise 8 Frames Überlappung als Ausgangspunkt untersuchen; nach dem ersten Abschnitt entsprechen das 56 neuen Frames pro Schritt. VHS liefert regulär feste Batchgrößen: Pufferung und Abschluss so umsetzen, dass keine Frames fehlen, doppelt ausgegeben oder falsch zugeordnet werden. Abschnittsgröße und Überlappung sollen zur verfügbaren Hardware passen. +- Bestehende DA3-Logik für Multi-View-Inferenz, Überlappung, Sim(3)-Ausrichtung und Verblendung untersuchen und wenn sinnvoll wiederverwenden. Die heutige DA3-Streaming-Node lädt das gesamte VIDEO und hält mehrere große Ergebnisstapel im RAM; bloßes Aufrufen dieser Node pro VHS-Batch erhält keine globale Konsistenz. Öffentliche und private Schnittstellen und ihre Abhängigkeiten vor der Integration prüfen. +- Die aktuelle DA3-Node nutzt gemeinsame Normalisierung über alle ausgerichteten Frames. Eine äquivalente speicherbegrenzte Lösung entwerfen; die Skalierung darf nicht je Batch unabhängig erfolgen. Tiefenausgabe und temporäres Format mit ausreichender Präzision und ohne unnötige verlustbehaftete Kompression wählen. +- Frameanzahl, Reihenfolge, Bildrate, Startposition, Größenanpassung sowie die letzte Teilmenge beider Durchläufe prüfen. Getestetes Video: 723 Frames, 48 FPS, 1024 × 1024. Für Tests vorhandenes Videomaterial aus `homemade/tmp` verwenden; keine persönlichen Inhalte oder Dateinamen in Repo, Logs oder Testdaten übernehmen. +- Den vorhandenen Importpfad zur VHS-Hilfe im Stereo-Node prüfen: Ein Fallback verweist auf `custom_nodes/comfyui-videohelpersuite`, die ältere Kopie wurde deaktiviert; aktiv ist `custom_nodes/ComfyUI-VideoHelperSuite`. +- Keine neuen Downloads ohne konkrete Zustimmung. Vorhandene Small- und Base-Gewichte liegen in `models/depthanything3` der ComfyUI-Installation. + +## Bisherige Messungen + +GPU: AMD Radeon RX 7900 XTX, 24 GiB VRAM; Rechner: 31 GiB RAM. DA3 Streaming wurde mit 72 Frames bei 1024 × 1024, 64 Frames pro Abschnitt, 8 Frames Überlappung, `bf16` und Torch-Ausrichtung getestet. Die Zeiten enthalten Modellaufbau, Ausrichtung und Schreiben der NPZ-Dateien: + +| Modell | Zeit | Spitzen-VRAM | Spitzen-RAM des Prozesses | +| --- | ---: | ---: | ---: | +| Small | 67,9 s | 6,94 GiB | 6,60 GiB | +| Base | 108,0 s | 12,35 GiB | 6,63 GiB | + +Large mit 32er-Abschnitten wurde wegen sehr langer Laufzeit auf Wunsch abgebrochen, bevor der erste Abschnitt fertig war. Keine belastbare Large-Zeit oder VRAM-Spitze vorhanden. Small oder Base sind für die Umsetzung vorzuziehen; keine Qualitätsbewertung der Tiefenbilder vorgenommen. Für den Einzelbildpfad vorher gemessene Zeiten sind nicht auf Multi-View-Streaming übertragbar. + +## Abnahme + +Einen vollständigen Clip aus dem vorhandenen Material durch beide Phasen verarbeiten. Überlappung und Batchgrenzen, Anzahl und Synchronität der Frames, Audio, normalisierte Tiefenwerte und sichtbar ruhiges Stereo prüfen. RAM/VRAM, Laufzeit und Plattenbedarf erfassen, einschließlich Abbruch und Wiederanlauf. Frühere Workflows mit dem bestehenden Stereo-Node müssen unverändert ausführbar bleiben. diff --git a/README.md b/README.md index 64d507a..34f5908 100644 --- a/README.md +++ b/README.md @@ -72,3 +72,116 @@ restart ComfyUI before starting a new meta-batch run. ## License MIT + +## DA3: automatische Tiefenvorbereitung und Stereo + +Nach dem Update ComfyUI neu starten und die Webseite neu laden. Beide Nodes +bleiben aktiv (Modus Always). Ein Queue-Start führt die beiden Durchläufe +automatisch hintereinander aus: + +1. VHS Batch Manager an VHS Load Video, **Duke DA3 Prepare Sequence** und + **Duke Stereo Prepared Sequence** anschließen. Nur dieses Node-Paar an + diesem Manager betreiben. +2. Die Originalbilder des Loaders direkt an Prepare anschließen. Parallel die + Bilder über den Upscaler an den IMAGE-Eingang des Stereo-Nodes führen. + AUDIO des Loaders direkt an Stereo anschließen. +3. **(Down)Load Depth Anything V3 Model** mit Prepare verbinden (Typ + `DA3MODEL`). Small oder Base mit bf16 und sdpa verwenden. Die STRING-Ausgabe + `depth_sequence` von Prepare direkt mit Stereo verbinden. +4. Bei Prepare `workflow_mode=prepare_and_stereo` verwenden (Standard). + Mit VHS-Batches von 16, DA3 `chunk_size=64` und `overlap=8` beginnen. + `fps` auf die effektive VHS-Bildrate stellen; die Werte für Quelle und + Frameauswahl in beiden Nodes identisch setzen. Ausgabe- und Cache-Pfad + müssen beschreibbar sein. +5. Einmal starten. Prepare blockiert den Stereo-Ausgang während der gesamten + Tiefenvorbereitung. Nach dem letzten Batch wird VHS zurückgesetzt und liest + dieselbe Quelle mit denselben Einstellungen erneut. Erst dieser zweite + Durchlauf skaliert die Bilder hoch und schreibt Stereo samt Audio. + +Der IMAGE-Eingang am neuen Stereo-Node ist lazy: Der Upscaler wird für diesen +Node erst im zweiten Durchlauf benötigt. Zusätzliche aktive Preview-Ausgaben +am Upscaler können ihn unabhängig davon auch im ersten Durchlauf ausführen. +Für den automatischen Einzelvideo-Ablauf keine Quelle wechseln, bis beide +Durchläufe abgeschlossen sind. + +### Verzeichnis automatisch abarbeiten + +Für DA3 den **Duke DA3 Directory Queue** anstelle von Loopy Dir einsetzen. +Ein vorhandener GIR-Verzeichnispicker kann weiterhin den Verzeichnispfad liefern. +Die übrigen VHS-/Upscaler-Verbindungen bleiben erhalten: + +- Derselbe VHS Batch Manager wird mit Directory Queue, VHS Load Video (Path), + Prepare und Stereo verbunden. +- `video_path` der Directory Queue wird an `video` von VHS Load Video (Path) + angeschlossen. Beide Bild-/Audiozweige müssen diesen Loader verwenden. +- Prepare bleibt auf `prepare_and_stereo`; Prepare und Stereo bleiben aktiv. +- `directory` bezeichnet den Eingabeordner, `start_index` ist nullbasiert, + `max_videos=0` verarbeitet alle verbleibenden Videos. Die Liste wird zu Beginn + aufgenommen, nach Dateiname sortiert und endet nach dem letzten Video. + +Die Queue hält die aktuelle Datei während beider Durchläufe fest. Erst nach +fertigem Stereo und erfolgreichem Audio-Mux folgt das nächste Video. Die +Quellenkennung wird automatisch aus Dateipfad, Größe und Änderungszeit gehasht; +Prepare und Stereo übernehmen sie aus dem gemeinsamen Manager. Persönliche +Quelldateinamen werden nicht im Tiefenmanifest gespeichert. Änderungen an +Quelldateien während des Laufs werden abgewiesen. + +Im Verzeichnismodus wird der Tiefencache nach erfolgreicher Ausgabe gelöscht. +Bei Fehlern beim Stereo-Schreiben oder Audio-Mux bleibt der Cache erhalten und +die Queue schaltet nicht weiter. Neue manuelle Läufe starten bei `start_index`; +es gibt keine automatische Wiederaufnahme eines abgebrochenen Prozesses. +Weitere aktive Video-Loader oder Ausgaben benötigen einen eigenen Manager. + +Ein vollständiges API-Prompt-Beispiel liegt in +[`examples/da3_directory_api.json`](examples/da3_directory_api.json). +Eingabeverzeichnis, vorhandenes Modell, Bildrate und Ausgabepfade vor Verwendung +anpassen. Das Beispiel verwendet den Duke NCNN Upscaler und enthält keine +Vorschau, die Upscaling im ersten Durchlauf auslösen könnte. `%source_key%` im +Ausgabepfad wird durch die automatische Quellenkennung ersetzt. + +Für reine Tiefenvorbereitung gibt es weiterhin `workflow_mode=prepare_only`. +Dabei den Stereo-Node deaktivieren. Eine fertige Folge kann anschließend über +einen STRING-/Primitive-Node am `depth_sequence`-Eingang erneut verwendet werden, +ohne DA3 neu auszuführen. + +`source_key` muss für jedes Video eindeutig und in beiden Durchläufen gleich +sein. Ebenso müssen `fps`, `start_frame` und `frame_stride` übereinstimmen. +Die letzten beiden Felder dokumentieren Startframe und Auswahlabstand des +Loaders; sie verändern dessen Auswahl nicht. Alle weiteren Loader-Einstellungen +(z. B. Frame-Limit und erzwungene Bildrate) ebenfalls identisch halten. +Die Kennung ist eine manuelle Zuordnung, keine automatische Inhaltsprüfung. +Im automatischen Modus teilen sich genau ein Prepare- und ein Stereo-Node +denselben Batch Manager; weitere Video-Ausgaben benötigen einen eigenen Ablauf. + +Die Datei `manifest.json` markiert eine fertige Folge. Pro Frame wird eine +verlustfreie float32-NPY-Tiefenkarte gespeichert; die gemeinsame globale +Min-/Max-Skalierung erfolgt beim Lesen. Das entspricht DA3 **Standard** ohne +Invertierung, nicht dessen V2-Style-Normalisierung. Nah liegt bei 0, fern bei 1; +`convergence_point` bezeichnet die Ebene ohne Parallaxe. Es findet keine erneute +Min-/Max-Skalierung pro Frame statt. Upscaling darf die räumliche Auflösung +ändern, aber keine Frames entfernen, hinzufügen oder umsortieren. + +Bei 723 Frames und etwa 1024² Pixeln sind rund 2,8 GiB für die Tiefe einzuplanen. +Einzelvideo-Folgen bleiben mit `cache_policy=keep` (Standard) für weitere +Stereo-Versuche erhalten. `delete_after_stereo` entfernt sie nach erfolgreichem +Schreiben und Audio-Mux. Im Verzeichnismodus gilt automatisch diese Löschregel. +Es werden ausschließlich entsprechend markierte, von DukeStereo erzeugte +Cache-Verzeichnisse gelöscht; die Ausgabedateien bleiben erhalten. Fehler während der +Vorbereitung entfernen den unvollständigen Cache; VHS-Reset räumt ihn ebenfalls +auf. Nach Prozessabsturz können unvollständige Verzeichnisse ohne Manifest +zurückbleiben. Ein Neustart berechnet neu; eine Fortsetzung ab Checkpoint ist +noch nicht implementiert. + +Der Prepare-Fortschritt zeigt Pufferung, Inferenzabschnitte, Ausrichtung und +Ermittlung des gemeinsamen Tiefenbereichs. Der Stereo-Fortschritt zählt die +Frames im zweiten Durchlauf; im Verzeichnismodus wird zusätzlich die +Videoposition protokolliert. VHS selbst zählt Requeues weiterhin über beide +Phasen; dessen Batch-Zähler ist kein Gesamtfortschritt für den DA3-Ablauf. + +Pufferung und Überlappung wurden synthetisch auch mit 723 Frames und wechselnden +Batchgrößen geprüft. Der automatische Einzelvideo-Ablauf wurde in ComfyUI mit +Small und Base bei bf16, 64/8, 1024² und 48 FPS durchlaufen: Beide Ausgaben +enthielten 723 decodierbare Frames in 4096 × 2048 SBS mit Audio. Die +Directory-Queue ist mit simulierten Zwei-Video-Läufen inklusive Audiofehler, +Cache-Aufräumen und Quellenwechsel getestet; ein vollständiger Verzeichnislauf +im ComfyUI-Frontend steht noch aus. diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000..8edf18b --- /dev/null +++ b/TODO.md @@ -0,0 +1,6 @@ +# TODO + +- [x] **Depth-Eingang oben am Stereo-Node:** fester STRING-Anschluss direkt hinter IMAGE; keine Tiefenpfad-Widgets zwischen den übrigen Einstellungen. +- [x] **Fortschritt für Prepare Sequence:** Pufferfüllstand, Inferenzabschnitte, Ausrichtung und globaler Tiefenbereich; Stereo zählt Frames separat. Gesamtframezahl wird von VHS übernommen, wenn verfügbar. +- [x] **Verzeichnislauf mit temporärem Tiefencache:** eigener Duke DA3 Directory Queue statt Loopy Dir; Quelldatei über beide Durchläufe festhalten, Kennung ableiten, erst nach erfolgreichem Stereo samt Audio Cache löschen und weiterschalten. Fehler halten die Queue an. +- [ ] **Abnahme im ComfyUI-Frontend:** Verzeichnislauf mit mehreren Videos sowie neue Fortschrittsanzeige prüfen. Einzelvideo-Durchläufe mit Small und Base wurden bereits erfolgreich geprüft. diff --git a/__init__.py b/__init__.py index 679e742..7186912 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,8 @@ from .duke_stereo_streaming import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +from .da3_sequence import NODE_CLASS_MAPPINGS as DA3_NODES +from .da3_sequence import NODE_DISPLAY_NAME_MAPPINGS as DA3_NAMES +NODE_CLASS_MAPPINGS.update(DA3_NODES) +NODE_DISPLAY_NAME_MAPPINGS.update(DA3_NAMES) diff --git a/da3_sequence.py b/da3_sequence.py new file mode 100644 index 0000000..38a2ce2 --- /dev/null +++ b/da3_sequence.py @@ -0,0 +1,413 @@ +"""Bounded DA3 processing. Uses the installed DA3 extension's private pipeline API.""" +import importlib +import logging +import math +import json +import sys +import shutil +import tempfile +from pathlib import Path +from dataclasses import dataclass + +import numpy as np +import torch + +from .duke_stereo_streaming import DukeStereoVideoBatch +from .directory_queue import DIRECTORY_KEY, DukeDA3DirectoryQueue, directory_run, rewind, finish_video + +logger = logging.getLogger(__name__) + +def report_progress(stage, current, total, detail=""): + logger.info("[Duke DA3] %s: %s/%s %s", stage, current, total or "?", detail) + if total: + try: + from comfy.utils import ProgressBar + except ImportError: + return + ProgressBar(total).update_absolute(min(current, total)) + +def remove_temporary_sequence(path, metadata): + # Only delete a cache explicitly created and marked by this implementation. + if (metadata.get("temporary_cache") and metadata.get("producer") == "DukeStereo-DA3-v1" + and path.name.startswith("duke-depth-") and not path.is_symlink()): + shutil.rmtree(path) + + + +def da3_module(name): + # ComfyUI may load custom nodes under a path-derived module name. + for key, module in list(sys.modules.items()): + if key.endswith('.nodes.streaming.pipeline') and 'DepthAnything' in str(getattr(module, '__file__', '')): + root = key.rsplit('.nodes.', 1)[0] + return importlib.import_module(root + '.nodes.' + name) + import folder_paths + for base in folder_paths.get_folder_paths('custom_nodes'): + if (Path(base) / 'ComfyUI-DepthAnythingV3').is_dir(): + if base not in sys.path: + sys.path.insert(0, base) + return importlib.import_module('ComfyUI-DepthAnythingV3.nodes.' + name) + raise RuntimeError('Installed ComfyUI-DepthAnythingV3 extension not found') + + +def requeue(): + try: + from videohelpersuite.utils import requeue_workflow_unchecked + except ModuleNotFoundError: + import folder_paths + for base in folder_paths.get_folder_paths('custom_nodes'): + candidate = str(Path(base) / 'ComfyUI-VideoHelperSuite') + if candidate not in sys.path: + sys.path.insert(0, candidate) + from videohelpersuite.utils import requeue_workflow_unchecked + requeue_workflow_unchecked() + + +class SequenceBuilder: + def __init__(self, parent, chunk_size, overlap, metadata): + if not 0 < overlap <= chunk_size // 2: + raise ValueError('Require 0 < overlap <= chunk_size / 2') + Path(parent).mkdir(parents=True, exist_ok=True) + self.path = Path(tempfile.mkdtemp(prefix='duke-depth-', dir=parent)) + self.size, self.overlap = chunk_size, overlap + self.meta = metadata + self.count = self.start = 0 + self.previous = None + self.scale = 1.0 + self.pending = [] + self.progress = None + self.sections = 0 + + def append(self, images, infer, align, final=False): + # Only chunk_size source frames and two inference results are resident. + for image in images: + self.pending.append(image.detach().cpu().clone()) + self.count += 1 + if len(self.pending) == self.size: + self.process(infer, align) + if final: + if self.pending and (self.previous is None or len(self.pending) > self.overlap): + self.process(infer, align) + self.finish() + + def process(self, infer, align): + self.sections += 1 + if self.progress: + self.progress("Inference", self.start, f"section {self.sections}, {len(self.pending)} frames") + result = infer(torch.stack(self.pending)) + if self.previous is not None: + if self.progress: + self.progress("Alignment", self.start, f"section {self.sections}") + scale, _, _ = align(self.previous, result) + if not np.isfinite(scale) or scale <= 0: + raise ValueError('Invalid Sim(3) alignment scale') + self.scale *= float(scale) + depth = result.depth.astype(np.float32) * self.scale + if not np.isfinite(depth).all(): + raise ValueError('DA3 produced non-finite depth') + for i, frame in enumerate(depth): + target = self.path / f'{self.start+i:08d}.npy' + if self.previous is not None and i < self.overlap: + weight = (i + 1) / (self.overlap + 1) + frame = np.load(target) * (1 - weight) + frame * weight + np.save(target, frame.astype(np.float32), allow_pickle=False) + # Retain raw overlap for pairwise local Sim(3), cumulative scale separately. + from dataclasses import replace + self.previous = replace(result, **{k: getattr(result, k)[-self.overlap:].copy() + for k in ('depth', 'conf', 'sky', 'extrinsics', 'intrinsics')}) + self.start += len(self.pending) - self.overlap + self.pending = self.pending[-self.overlap:] + + def finish(self): + if not self.count: + raise ValueError('Empty sequence') + lo, hi = float('inf'), -float('inf') + for i in range(self.count): + if self.progress and (i % 16 == 0 or i + 1 == self.count): + self.progress('Global depth range', i + 1, f'{self.count} frames') + depth = np.load(self.path / f'{i:08d}.npy', mmap_mode='r') + lo, hi = min(lo, float(depth.min())), max(hi, float(depth.max())) + # Store lossless float32 raw depth and one global Standard normalization. + manifest = dict(self.meta, version=1, frames=self.count, minimum=lo, maximum=hi, + normalization='global_standard_near_zero', chunk_size=self.size, + overlap=self.overlap) + tmp = self.path / 'manifest.tmp' + tmp.write_text(json.dumps(manifest, indent=2)) + tmp.replace(self.path / 'manifest.json') + self.pending.clear() + self.previous = None + + +@dataclass +class SequenceReplay: + path: str + + +def idle_guard(): + # VHS recognizes and closes generator outputs when a run is reset. + yield + + +class DukeDA3Prepare: + @classmethod + def INPUT_TYPES(cls): + return {'required': {'images': ('IMAGE',), 'da3_model': ('DA3MODEL',), + 'cache_directory': ('STRING', {'default': 'output/duke-depth'}), + 'source_key': ('STRING', {'default': 'clip-1'}), + 'fps': ('FLOAT', {'default': 48.0, 'min': 0.01}), + 'start_frame': ('INT', {'default': 0, 'min': 0}), + 'frame_stride': ('INT', {'default': 1, 'min': 1}), + 'chunk_size': ('INT', {'default': 64, 'min': 2, 'max': 128}), + 'overlap': ('INT', {'default': 8, 'min': 1, 'max': 64})}, + 'optional': {'meta_batch': ('VHS_BatchManager',), + 'workflow_mode': (['prepare_and_stereo', 'prepare_only'], + {'default': 'prepare_and_stereo'}), + 'cache_policy': (['keep', 'delete_after_stereo'], {'default': 'keep'})}, + 'hidden': {'unique_id': 'UNIQUE_ID'}} + RETURN_TYPES = ('STRING',) + RETURN_NAMES = ('depth_sequence',) + FUNCTION = 'execute' + CATEGORY = 'Stereo/DA3' + OUTPUT_NODE = True + + def execute(self, images, da3_model, cache_directory, source_key, fps, + start_frame, frame_stride, chunk_size, overlap, meta_batch=None, unique_id=None, + workflow_mode="prepare_and_stereo", cache_policy="keep"): + import comfy.model_management as mm + state_key = str(unique_id) + run = directory_run(meta_batch) + if run is not None: + if workflow_mode != 'prepare_and_stereo': + raise ValueError('Directory queue requires prepare_and_stereo mode') + source_key = run.source_key + total = getattr(meta_batch, 'total_frames', 0) + total = int(total) if total and math.isfinite(total) else 0 + clip = f'video {run.index + 1}/{len(run.files)}' if run else '' + if meta_batch is not None and state_key in meta_batch.outputs: + state = meta_batch.outputs[state_key][0] + if isinstance(state, SequenceReplay): + if meta_batch.has_closed_inputs: + _, guard = meta_batch.outputs.pop(state_key) + guard.close() + report_progress('Stereo pass', 0, 0, clip) + return {'ui': {'text': ['Stereo pass: prepared depth loaded']}, + 'result': (state.path,)} + else: + state = SequenceBuilder(cache_directory, chunk_size, overlap, + dict(source_key=source_key, fps=fps, start_frame=start_frame, + frame_stride=frame_stride, producer='DukeStereo-DA3-v1', + temporary_cache=run is not None or cache_policy == 'delete_after_stereo')) + if meta_batch is not None: + # Generator protocol is recognized by VHS reset on cancellation/restart. + def cleanup(): + try: + yield + finally: + state.pending.clear() + state.previous = None + if state.path.exists() and not (state.path / 'manifest.json').exists(): + shutil.rmtree(state.path) + guard = cleanup() + next(guard) + meta_batch.outputs[state_key] = (state, guard) + state.progress = lambda stage, current, detail: report_progress( + stage, current, state.count if stage == "Global depth range" else total, f"{clip} {detail}") + report_progress("Prepare buffer", state.count, total, + f"{clip} {len(state.pending)}/{chunk_size} buffered") + def infer(frames): + pipeline_mod = da3_module('streaming.pipeline') + utils = da3_module('utils') + mm.throw_exception_if_processing_interrupted() + if not Path(da3_model['model_path']).is_file(): + raise ValueError('DA3 weights must already exist locally') + patcher = da3_module('load_model')._get_or_build_da3_model(da3_model) + if not hasattr(patcher.model, '__call__'): + raise RuntimeError('Incompatible DA3 model') + dtype = patcher.model_options.get('da3_dtype', torch.float16) + mm.load_models_gpu([patcher]) + pipeline = pipeline_mod.StreamingPipeline(patcher.model, + pipeline_mod.StreamingConfig(chunk_size=chunk_size, overlap=overlap, + align_lib='torch'), mm.get_torch_device(), dtype) + self.pipeline = pipeline + frames, _, _ = utils.resize_to_patch_multiple(frames.permute(0, 3, 1, 2)) + frames = utils.imagenet_normalize(frames).to(dtype=dtype).unsqueeze(0) + with torch.inference_mode(), torch.autocast(device_type=mm.get_torch_device().type, dtype=dtype): + return pipeline._process_chunk(frames, 0, frames.shape[1]) + final = meta_batch is None or meta_batch.has_closed_inputs + try: + state.append(images, infer, lambda a, b: self.pipeline._align_chunks(a, b), final) + except BaseException: + state.pending.clear() + state.previous = None + self.pipeline = None + shutil.rmtree(state.path) + if meta_batch is not None: + entry = meta_batch.outputs.pop(state_key, None) + if entry: + entry[-1].close() + raise + self.pipeline = None + path = str(state.path) + if not final: + from comfy_execution.graph import ExecutionBlocker + requeue() + return {'ui': {'unfinished_batch': [True]}, + 'result': (ExecutionBlocker(None),)} + if meta_batch is not None: + entry = meta_batch.outputs.pop(state_key, None) + if entry: + entry[-1].close() + if workflow_mode == 'prepare_and_stereo': + if set(meta_batch.outputs) - {DIRECTORY_KEY}: + raise ValueError('Automatic two-pass mode requires its own VHS Batch Manager.') + # VHS reopens its source generators on the next queued batch. + # Keep the completed sequence in manager state so manual restart/reset + # clears the phase as well; no stale state lives on the node instance. + rewind(meta_batch) + report_progress("Depth complete; starting Stereo", state.count, state.count, clip) + guard = idle_guard() + next(guard) + meta_batch.outputs[state_key] = (SequenceReplay(path), guard) + from comfy_execution.graph import ExecutionBlocker + requeue() + return {'ui': {'text': ['Depth complete; restarting source for Stereo'], + 'unfinished_batch': [True]}, + 'result': (ExecutionBlocker(None),)} + if not meta_batch.outputs: + meta_batch.reset() + return {'ui': {'text': [path]}, 'result': (path,)} + + +class DukeStereoPrepared(DukeStereoVideoBatch): + normalize_depth = False + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + del inputs['required']['depth_maps'] + inputs['required']['fps'] = ('FLOAT', {'default': 48.0, 'min': 0.01}) + # A socket-only input stays with IMAGE/AUDIO instead of below widgets. + inputs['required'].pop('images') + inputs['required'] = { + 'images': ('IMAGE', {'lazy': True}), + 'depth_sequence': ('STRING', {'forceInput': True}), + **inputs['required'], + } + inputs['required'].update(source_key=('STRING', {'default': 'clip-1'}), + start_frame=('INT', {'default': 0, 'min': 0}), + frame_stride=('INT', {'default': 1, 'min': 1})) + return inputs + + def check_lazy_status(self, images=None, **kwargs): + # The prepared path is blocked during pass one, so upstream upscaling + # is requested only once the complete depth sequence is available. + return ['images'] if images is None else [] + + def execute(self, images, depth_sequence, source_key, start_frame, frame_stride, **kwargs): + path = Path(depth_sequence) + manifest = path / 'manifest.json' + if not manifest.is_file(): + raise ValueError( + 'Depth sequence is incomplete or the directory is incorrect. ' + 'Use prepare_and_stereo on Prepare with a shared VHS Batch Manager, ' + 'or select an already completed sequence directory.') + meta = json.loads(manifest.read_text()) + run = directory_run(kwargs.get('meta_batch')) + if run is not None: + source_key = run.source_key + for key, value in dict(source_key=source_key, start_frame=start_frame, + frame_stride=frame_stride, fps=kwargs['fps']).items(): + if meta[key] != value: + raise ValueError(f'Depth sequence mismatch: {key}') + manager, uid = kwargs.get('meta_batch'), kwargs.get('unique_id') + import subprocess + from .duke_stereo_streaming import resolve_output_path, mux_audio + from .stereoimage_generation import apply_stereo_divergence + import comfy.model_management as mm + h, w = images.shape[1:3] + if manager is not None and uid in manager.outputs: + state = manager.outputs[uid][0] + if state['sequence'] != str(path.resolve()) or state['shape'] != (h, w): + raise ValueError('Sequence or image dimensions changed during batch') + else: + output = resolve_output_path(kwargs['output_path'], dict(kwargs, source_key=source_key, width=w, height=h, frames=meta['frames'])) + proc = subprocess.Popen(['ffmpeg', '-v', 'error', '-y', '-f', 'rawvideo', + '-pix_fmt', 'rgb24', '-s', f'{w*2}x{h}', '-r', str(kwargs['fps']), + '-i', '-', '-c:v', 'libx264', '-pix_fmt', 'yuv420p', '-crf', + str(kwargs['crf']), '-preset', 'medium', output], stdin=subprocess.PIPE) + state = dict(cursor=0, proc=proc, output=output, audio=kwargs.get('audio'), + sequence=str(path.resolve()), shape=(h, w)) + def cleanup(): + try: + yield + finally: + try: + if not proc.stdin.closed: + proc.stdin.close() + except BrokenPipeError: + pass + finally: + if proc.poll() is None: + proc.wait() + guard = cleanup() + next(guard) + state['guard'] = guard + if manager is not None: + manager.outputs[uid] = (state, guard) + end = state['cursor'] + len(images) + final = manager is None or manager.has_closed_inputs + try: + if end > meta['frames'] or (final and end != meta['frames']): + raise ValueError('Source and prepared depth frame counts differ') + for i, image in enumerate(images, state['cursor']): + mm.throw_exception_if_processing_interrupted() + depth = np.load(path / f'{i:08d}.npy', allow_pickle=False) + depth = (depth - meta['minimum']) / max(meta['maximum'] - meta['minimum'], 1e-8) + depth = torch.nn.functional.interpolate(torch.from_numpy(depth.copy())[None, None], + size=(h, w), mode='bilinear', align_corners=True)[0, 0].numpy() + if kwargs['depth_blur'] > 0: + from scipy.ndimage import gaussian_filter + blurred = gaussian_filter(depth, sigma=kwargs['depth_blur']) + threshold = kwargs['depth_blur_edge_threshold'] / 255 + if threshold > 0: + gy, gx = np.gradient(depth) + depth = np.where(np.hypot(gx, gy) > threshold, depth, blurred) + else: + depth = blurred + rgb = (image.detach().cpu().numpy().clip(0, 1) * 255).astype(np.uint8) + views = [apply_stereo_divergence(rgb, depth, sign * kwargs['divergence'], + -sign * kwargs['separation'], kwargs['stereo_offset_exponent'], + kwargs['fill_technique'], kwargs['convergence_point'], normalize_depth=False) + for sign in (1, -1)] + state['proc'].stdin.write(np.concatenate(views, axis=1).tobytes()) + if (i + 1) % 8 == 0 or i + 1 == meta['frames']: + from comfy.utils import ProgressBar + ProgressBar(meta['frames']).update_absolute(i + 1) + state['cursor'] = end + report_progress('Stereo frames', end, meta['frames'], + f'video {run.index + 1}/{len(run.files)}' if run else '') + if not final: + requeue() + return {'ui': {'unfinished_batch': [True]}, 'result': (state['output'],)} + state['guard'].close() + if state['proc'].returncode: + raise RuntimeError('Stereo ffmpeg failed') + report_progress('Finalize video and audio', end, meta['frames']) + mux_audio(state['output'], state['audio']) + remove_temporary_sequence(path, meta) + except BaseException: + state['guard'].close() + if manager is not None: + manager.outputs.pop(uid, None) + raise + if manager is not None: + manager.outputs.pop(uid, None) + if finish_video(manager): + requeue() + return {'ui': {'unfinished_batch': [True]}, 'result': (state['output'],)} + return (state['output'],) + + +NODE_CLASS_MAPPINGS = {'Duke DA3 Directory Queue': DukeDA3DirectoryQueue, + 'Duke DA3 Prepare Sequence': DukeDA3Prepare, + 'Duke Stereo Prepared Sequence': DukeStereoPrepared} +NODE_DISPLAY_NAME_MAPPINGS = {key: key for key in NODE_CLASS_MAPPINGS} diff --git a/directory_queue.py b/directory_queue.py new file mode 100644 index 0000000..c380ada --- /dev/null +++ b/directory_queue.py @@ -0,0 +1,106 @@ +"""Directory sequencing tied to completion of both DA3 passes, not prompt count.""" +import hashlib +from dataclasses import dataclass +from pathlib import Path + +DIRECTORY_KEY = '__duke_da3_directory__' +EXTENSIONS = {'.mp4', '.mkv', '.mov', '.webm', '.avi', '.m4v'} + + +def idle_guard(): + yield + + +@dataclass +class DirectoryRun: + files: tuple + directory: str + owner: str + index: int = 0 + identities: tuple = () + + def __post_init__(self): + self.identities = tuple(self.identity(p) for p in self.files) + + @staticmethod + def identity(p): + stat = p.stat() + identity = f'{p.resolve()}:{stat.st_size}:{stat.st_mtime_ns}' + return hashlib.sha256(identity.encode()).hexdigest()[:24] + + @property + def source_key(self): + key = self.identity(self.files[self.index]) + if key != self.identities[self.index]: + raise ValueError('Source video changed during the directory run') + return key + + +def directory_run(manager): + entry = manager.outputs.get(DIRECTORY_KEY) if manager is not None else None + return entry[0] if entry else None + + +def rewind(manager): + """Reopen VHS readers while preserving the pinned directory selection.""" + entry = manager.outputs.pop(DIRECTORY_KEY, None) + manager.reset() + if entry: + manager.outputs[DIRECTORY_KEY] = entry + + +def finish_video(manager): + """Called only after video closure, successful audio mux and cache cleanup.""" + run = directory_run(manager) + if run is None: + if not manager.outputs: + manager.reset() + return False + unexpected = set(manager.outputs) - {DIRECTORY_KEY} + if unexpected: + raise ValueError('DA3 directory mode needs a dedicated VHS Batch Manager') + entry = manager.outputs.pop(DIRECTORY_KEY) + manager.reset() + run.index += 1 + if run.index == len(run.files): + entry[-1].close() + return False + manager.outputs[DIRECTORY_KEY] = entry + return True + + +class DukeDA3DirectoryQueue: + @classmethod + def INPUT_TYPES(cls): + return {'required': { + 'directory': ('STRING', {'default': '', 'tooltip': 'Input directory; source list is fixed for this run.'}), + 'meta_batch': ('VHS_BatchManager',), + 'start_index': ('INT', {'default': 0, 'min': 0}), + 'max_videos': ('INT', {'default': 0, 'min': 0, 'tooltip': '0 = all remaining videos'})}, + 'hidden': {'unique_id': 'UNIQUE_ID'}} + + RETURN_TYPES = ('STRING', 'STRING', 'INT', 'INT') + RETURN_NAMES = ('video_path', 'source_key', 'current_index', 'video_count') + FUNCTION = 'execute' + CATEGORY = 'Stereo/DA3' + + def execute(self, directory, meta_batch, start_index=0, max_videos=0, unique_id=None): + root = Path(directory).expanduser().resolve() + run = directory_run(meta_batch) + if run is None: + if not root.is_dir(): + raise ValueError('DA3 input directory does not exist') + files = sorted((p for p in root.iterdir() if p.is_file() and p.suffix.lower() in EXTENSIONS), + key=lambda p: p.name) + files = files[start_index:] + if max_videos: + files = files[:max_videos] + if not files: + raise ValueError('No videos at the selected directory start index') + run = DirectoryRun(tuple(files), str(root), str(unique_id)) + guard = idle_guard() + next(guard) + meta_batch.outputs[DIRECTORY_KEY] = (run, guard) + elif run.directory != str(root) or run.owner != str(unique_id): + raise ValueError('Directory selection changed during a DA3 run') + return (str(run.files[run.index]), run.source_key, run.index + 1, len(run.files)) diff --git a/duke_stereo_streaming.py b/duke_stereo_streaming.py index b58a9c3..a125e80 100644 --- a/duke_stereo_streaming.py +++ b/duke_stereo_streaming.py @@ -123,6 +123,7 @@ def mux_audio(output_path, audio): class DukeStereoVideo: + normalize_depth = True """ ComfyUI Node für Stereo SBS Video mit Streaming (RAM-effizient). Verwendet die bewährte stereoimage_generation Logik. @@ -220,7 +221,7 @@ class DukeStereoVideo: depth_np = depth_t.squeeze().numpy() # Falls Depth 0-1 ist, auf 0-255 skalieren - if depth_np.max() <= 1.0: + if self.normalize_depth and depth_np.max() <= 1.0: depth_np = (depth_np * 255).astype(np.float32) # Edge-aware depth blur @@ -239,11 +240,11 @@ class DukeStereoVideo: # Stereo-Warp mit originaler Funktion left = apply_stereo_divergence( img_np, depth_np, divergence, -separation, - stereo_offset_exponent, fill_technique, convergence_point + stereo_offset_exponent, fill_technique, convergence_point, normalize_depth=self.normalize_depth ) right = apply_stereo_divergence( img_np, depth_np, -divergence, separation, - stereo_offset_exponent, fill_technique, convergence_point + stereo_offset_exponent, fill_technique, convergence_point, normalize_depth=self.normalize_depth ) # SBS zusammenfügen @@ -274,7 +275,7 @@ class DukeStereoVideoBatch(DukeStereoVideo): try: from videohelpersuite.utils import requeue_workflow_unchecked except ModuleNotFoundError: import sys - sys.path.insert(0, str(Path(__file__).parents[1] / 'comfyui-videohelpersuite')) + sys.path.insert(0, str(Path(__file__).parents[1] / 'ComfyUI-VideoHelperSuite')) from videohelpersuite.utils import requeue_workflow_unchecked n_frames = images.shape[0] @@ -308,19 +309,19 @@ class DukeStereoVideoBatch(DukeStereoVideo): if depth_np.shape[:2] != (h, w): depth_t = torch.from_numpy(depth_np).unsqueeze(0).unsqueeze(0).float() depth_np = torch.nn.functional.interpolate(depth_t, size=(h, w), mode="bilinear", align_corners=True).squeeze().numpy() - if depth_np.max() <= 1.0: depth_np = (depth_np * 255).astype(np.float32) + if self.normalize_depth and depth_np.max() <= 1.0: depth_np = (depth_np * 255).astype(np.float32) if depth_blur > 0: from scipy.ndimage import gaussian_filter blurred = gaussian_filter(depth_np, sigma=depth_blur) if depth_blur_edge_threshold > 0: grad_y, grad_x = np.gradient(depth_np) - edge_mask = np.sqrt(grad_x**2 + grad_y**2) > depth_blur_edge_threshold + edge_mask = np.sqrt(grad_x**2 + grad_y**2) > (depth_blur_edge_threshold if self.normalize_depth else depth_blur_edge_threshold / 255) depth_np = np.where(edge_mask, depth_np, blurred) else: depth_np = blurred left = apply_stereo_divergence(img_np, depth_np, divergence, -separation, - stereo_offset_exponent, fill_technique, convergence_point) + stereo_offset_exponent, fill_technique, convergence_point, normalize_depth=self.normalize_depth) right = apply_stereo_divergence(img_np, depth_np, -divergence, separation, - stereo_offset_exponent, fill_technique, convergence_point) + stereo_offset_exponent, fill_technique, convergence_point, normalize_depth=self.normalize_depth) write_queue.put(np.concatenate([left, right], axis=1)) pbar.update(1) diff --git a/examples/da3_directory_api.json b/examples/da3_directory_api.json new file mode 100644 index 0000000..83c0add --- /dev/null +++ b/examples/da3_directory_api.json @@ -0,0 +1,123 @@ +{ + "1": { + "class_type": "VHS_BatchManager", + "inputs": { + "frames_per_batch": 64 + } + }, + "2": { + "class_type": "Duke DA3 Directory Queue", + "inputs": { + "directory": "/path/to/input-videos", + "meta_batch": [ + "1", + 0 + ], + "start_index": 0, + "max_videos": 0 + } + }, + "3": { + "class_type": "VHS_LoadVideoPath", + "inputs": { + "video": [ + "2", + 0 + ], + "meta_batch": [ + "1", + 0 + ], + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "AnimateDiff" + } + }, + "4": { + "class_type": "DownloadAndLoadDepthAnythingV3Model", + "inputs": { + "model": "da3_base.safetensors", + "precision": "bf16", + "attention": "sdpa" + } + }, + "5": { + "class_type": "Duke DA3 Prepare Sequence", + "inputs": { + "images": [ + "3", + 0 + ], + "da3_model": [ + "4", + 0 + ], + "meta_batch": [ + "1", + 0 + ], + "cache_directory": "output/duke-depth", + "source_key": "auto", + "fps": 48.0, + "start_frame": 0, + "frame_stride": 1, + "chunk_size": 64, + "overlap": 8, + "workflow_mode": "prepare_and_stereo", + "cache_policy": "delete_after_stereo" + } + }, + "6": { + "class_type": "DukeUpscalerNCNN", + "inputs": { + "images": [ + "3", + 0 + ], + "model": "realesr-animevideov3", + "scale": 2, + "tile_size": 1024, + "gpu_id": -1, + "threads": "8:4:8", + "keep_temp": false + } + }, + "7": { + "class_type": "Duke Stereo Prepared Sequence", + "inputs": { + "images": [ + "6", + 0 + ], + "audio": [ + "3", + 2 + ], + "meta_batch": [ + "1", + 0 + ], + "depth_sequence": [ + "5", + 0 + ], + "output_path": "output/stereo_%source_key%.mp4", + "fps": 48.0, + "source_key": "auto", + "start_frame": 0, + "frame_stride": 1, + "divergence": -2.5, + "separation": 0.0, + "convergence_point": 1.0, + "stereo_offset_exponent": 2.0, + "fill_technique": "polylines_soft", + "depth_blur": 6, + "depth_blur_edge_threshold": 6.0, + "crf": 19 + } + } +} diff --git a/stereoimage_generation.py b/stereoimage_generation.py index 3dbe0b5..5107670 100644 --- a/stereoimage_generation.py +++ b/stereoimage_generation.py @@ -725,7 +725,7 @@ def create_stereoimages(original_image, depthmap, divergence, separation=0.0, mo else: return stereo_images -def apply_stereo_divergence(original_image, depth, divergence, separation, stereo_offset_exponent, fill_technique, convergence_point=0.5): +def apply_stereo_divergence(original_image, depth, divergence, separation, stereo_offset_exponent, fill_technique, convergence_point=0.5, normalize_depth=True): """ Dispatches to the desired stereo mapping algorithm. @@ -740,7 +740,9 @@ def apply_stereo_divergence(original_image, depth, divergence, separation, stere depth_max = depth.max() # Protect against division by zero if depth map is flat (all pixels same value) - if depth_max == depth_min: + if not normalize_depth: + normalized_depth = np.clip(depth, 0.0, 1.0) + elif depth_max == depth_min: normalized_depth = np.zeros_like(depth) else: normalized_depth = (depth - depth_min) / (depth_max - depth_min) diff --git a/tests/test_da3_sequence.py b/tests/test_da3_sequence.py new file mode 100644 index 0000000..046a1eb --- /dev/null +++ b/tests/test_da3_sequence.py @@ -0,0 +1,116 @@ +"""Run with the ComfyUI Python environment: python -m unittest discover -s tests.""" +import importlib.util +import json +import sys +import tempfile +import types +import unittest +from unittest.mock import patch +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +ROOT = Path(__file__).resolve().parents[1] +package = types.ModuleType('duke_test') +package.__path__ = [str(ROOT)] +sys.modules['duke_test'] = package +from duke_test.da3_sequence import SequenceBuilder, DukeDA3Prepare, SequenceReplay, DukeStereoPrepared + + +@dataclass +class Result: + depth: np.ndarray + conf: np.ndarray + sky: np.ndarray + extrinsics: np.ndarray + intrinsics: np.ndarray + + +class SequenceTests(unittest.TestCase): + def test_arbitrary_batches_and_tail(self): + for count in (1, 8, 64, 65, 120, 121, 723): + for batch in (7, 16, 64, 80): + with self.subTest(count=count, batch=batch), tempfile.TemporaryDirectory() as tmp: + state = SequenceBuilder(tmp, 64, 8, {}) + calls = [] + def infer(images): + depth = images[..., 0].numpy() + calls.append(len(images)) + return Result(depth, depth, depth, depth, depth) + for start in range(0, count, batch): + end = min(start + batch, count) + images = torch.arange(start, end).float()[:, None, None, None].expand(-1, 2, 2, 3) + state.append(images, infer, lambda a, b: (1., None, None), end == count) + self.assertLessEqual(max(calls), 64) + self.assertEqual(len(list(state.path.glob('*.npy'))), count) + for i in range(count): + np.testing.assert_allclose(np.load(state.path / f'{i:08d}.npy'), i, atol=1e-4) + manifest = json.loads((state.path / 'manifest.json').read_text()) + self.assertEqual((manifest['minimum'], manifest['maximum']), (0, count-1)) + + def test_bad_overlap(self): + with tempfile.TemporaryDirectory() as tmp: + for overlap in (0, 64, 33): + with self.assertRaises(ValueError): + SequenceBuilder(tmp, 64, overlap, {}) + + +class AutomaticPassTests(unittest.TestCase): + def test_complete_prepare_rewinds_and_replays_then_clears_state(self): + class Manager: + def __init__(self): + self.outputs = {} + self.has_closed_inputs = False + self.resets = 0 + def reset(self): + self.outputs.clear() + self.has_closed_inputs = False + self.resets += 1 + class Blocker: + def __init__(self, message): + self.message = message + comfy = types.ModuleType('comfy') + mm = types.ModuleType('comfy.model_management') + comfy.model_management = mm + graph = types.ModuleType('comfy_execution.graph') + graph.ExecutionBlocker = Blocker + def append(builder, images, infer, align, final=False): + builder.count += len(images) + if final: + (builder.path / 'manifest.json').write_text('{}') + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, { + 'comfy': comfy, 'comfy.model_management': mm, + 'comfy_execution.graph': graph}), patch.object(SequenceBuilder, 'append', append), \ + patch('duke_test.da3_sequence.requeue') as queue: + manager = Manager() + node = DukeDA3Prepare() + args = dict(images=torch.zeros(2, 2, 2, 3), da3_model={}, cache_directory=tmp, + source_key='test', fps=48, start_frame=0, frame_stride=1, + chunk_size=64, overlap=8, meta_batch=manager, unique_id='prepare') + first = node.execute(**args) + self.assertIsInstance(first['result'][0], Blocker) + manager.has_closed_inputs = True + last = node.execute(**args) + self.assertIsInstance(last['result'][0], Blocker) + self.assertEqual(manager.resets, 1) + replay = manager.outputs['prepare'][0] + self.assertIsInstance(replay, SequenceReplay) + # A new node instance still finds phase state on the manager. + replay_node = DukeDA3Prepare() + self.assertEqual(replay_node.execute(**args)['result'], (replay.path,)) + manager.has_closed_inputs = True + self.assertEqual(replay_node.execute(**args)['result'], (replay.path,)) + self.assertNotIn('prepare', manager.outputs) + self.assertEqual(queue.call_count, 2) + + def test_stereo_images_are_lazy(self): + node = DukeStereoPrepared() + self.assertTrue(node.INPUT_TYPES()['required']['images'][1]['lazy']) + self.assertEqual(node.check_lazy_status(), ['images']) + self.assertEqual(node.check_lazy_status(images=torch.zeros(1)), []) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_directory_queue.py b/tests/test_directory_queue.py new file mode 100644 index 0000000..9cba491 --- /dev/null +++ b/tests/test_directory_queue.py @@ -0,0 +1,146 @@ +"""Exercise directory pinning, both passes, successful cleanup, and failed mux.""" +import json +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import torch + +# Reuse the lightweight package loader without importing ComfyUI's GPU runtime. +from test_da3_sequence import SequenceBuilder, DukeDA3Prepare, DukeStereoPrepared +from duke_test.directory_queue import DukeDA3DirectoryQueue, directory_run, DIRECTORY_KEY + + +class Manager: + def __init__(self): + self.outputs = {} + self.inputs = {} + self.has_closed_inputs = False + self.total_frames = 3 + self.resets = 0 + + def reset(self): + for entry in list(self.outputs.values()): + entry[-1].close() + self.outputs = {} + self.inputs = {} + self.has_closed_inputs = False + self.total_frames = float('inf') + self.resets += 1 + + +class Process: + def __init__(self, *args, **kwargs): + import io + self.stdin = io.BytesIO() + self.returncode = None + + def poll(self): + return self.returncode + + def wait(self): + self.returncode = 0 + return 0 + + +class DirectoryTests(unittest.TestCase): + def test_two_videos_two_passes_and_temporary_cleanup(self): + self.run_directory(False) + + def test_mux_failure_keeps_cache_and_does_not_advance(self): + self.run_directory(True) + + def run_directory(self, mux_failure): + comfy = types.ModuleType('comfy') + mm = types.ModuleType('comfy.model_management') + mm.throw_exception_if_processing_interrupted = lambda: None + utils = types.ModuleType('comfy.utils') + class Progress: + def __init__(self, total): pass + def update_absolute(self, value): pass + utils.ProgressBar = Progress + comfy.model_management = mm + comfy.utils = utils + graph = types.ModuleType('comfy_execution.graph') + class Blocker: + def __init__(self, message): pass + graph.ExecutionBlocker = Blocker + def append(builder, images, infer, align, final=False): + for _ in images: + np.save(builder.path / f'{builder.count:08d}.npy', np.ones((2, 2), dtype=np.float32)) + builder.count += 1 + if final: + builder.finish() + modules = {'comfy': comfy, 'comfy.model_management': mm, 'comfy.utils': utils, + 'comfy_execution.graph': graph} + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, modules), \ + patch.object(SequenceBuilder, 'append', append), \ + patch('duke_test.da3_sequence.requeue') as requeue, \ + patch('duke_test.duke_stereo_streaming.mux_audio', + side_effect=RuntimeError('mux failed') if mux_failure else None) as mux, \ + patch('subprocess.Popen', Process), \ + patch('duke_test.stereoimage_generation.apply_stereo_divergence', + side_effect=lambda rgb, *a, **k: rgb): + root = Path(tmp) / 'inputs'; root.mkdir() + for name in ('a.mp4', 'b.mp4'): (root / name).touch() + manager = Manager(); queue = DukeDA3DirectoryQueue() + prepare = DukeDA3Prepare(); stereo = DukeStereoPrepared() + keys = [] + for video_index in range(2): + selection = queue.execute(str(root), manager, unique_id='directory') + keys.append(selection[1]) + self.assertEqual(selection[2:], (video_index + 1, 2)) + for final, count in ((False, 2), (True, 1)): + manager.has_closed_inputs = final + result = prepare.execute(torch.zeros(count, 2, 2, 3), {}, tmp, 'ignored', + 48, 0, 1, 64, 8, manager, 'prepare') + self.assertIsInstance(result['result'][0], Blocker) + self.assertEqual(directory_run(manager).index, video_index) + path = Path(manager.outputs['prepare'][0].path) + meta = json.loads((path / 'manifest.json').read_text()) + self.assertEqual(meta['source_key'], selection[1]) + self.assertTrue(meta['temporary_cache']) + for final, count in ((False, 2), (True, 1)): + manager.has_closed_inputs = final + self.assertEqual(queue.execute(str(root), manager, unique_id='directory')[:2], selection[:2]) + prepared = prepare.execute(torch.zeros(count, 2, 2, 3), {}, tmp, 'ignored', + 48, 0, 1, 64, 8, manager, 'prepare') + self.assertEqual(prepared['result'], (str(path),)) + args = dict(images=torch.zeros(count, 2, 2, 3), depth_sequence=str(path), + source_key='ignored', start_frame=0, frame_stride=1, + output_path=str(Path(tmp) / 'stereo.mp4'), fps=48, divergence=-2.5, + separation=0, convergence_point=1, stereo_offset_exponent=2, + fill_technique='none', depth_blur=0, depth_blur_edge_threshold=6, + crf=19, meta_batch=manager, unique_id='stereo') + if mux_failure and final: + with self.assertRaisesRegex(RuntimeError, 'mux failed'): + stereo.execute(**args) + self.assertTrue(path.exists()) + self.assertEqual(directory_run(manager).index, 0) + return + stereo.execute(**args) + self.assertFalse(path.exists()) + self.assertNotEqual(keys[0], keys[1]) + self.assertFalse(manager.outputs) + self.assertEqual(mux.call_count, 2) + # Three continuations per clip plus one to move to the second clip. + self.assertEqual(requeue.call_count, 7) + + def test_source_changes_are_rejected(self): + with tempfile.TemporaryDirectory() as tmp: + source = Path(tmp) / 'a.mp4'; source.touch() + queue = DukeDA3DirectoryQueue(); manager = Manager() + queue.execute(tmp, manager, unique_id='q') + source.write_bytes(b'changed') + with self.assertRaisesRegex(ValueError, 'Source video changed'): + queue.execute(tmp, manager, unique_id='q') + manager.reset() + self.assertNotIn(DIRECTORY_KEY, manager.outputs) + + +if __name__ == '__main__': + unittest.main()