Add DA3 depth preparation, automatic stereo passes and directory queue
This commit is contained in:
Lukas Schmidt
@@ -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.
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
@@ -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)
|
||||
|
||||
+413
@@ -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}
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user