356 lines
12 KiB
Python
356 lines
12 KiB
Python
# nodes.py
|
|
import os
|
|
import torch
|
|
import faiss
|
|
import logging
|
|
import numpy as np
|
|
import requests
|
|
from tqdm import tqdm
|
|
from PIL import Image
|
|
from pathlib import Path
|
|
import torch.nn.functional as F
|
|
|
|
# Make sure this import can find the database.py in your package
|
|
from .database import get_connection
|
|
from torchvision import transforms
|
|
|
|
# Import your get_dot_frame function
|
|
from .dot_functions import get_dot_frame
|
|
from .functions import download_checkpoints
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Define URLs for data.bin and embeddings.db on HuggingFace
|
|
DATA_BIN_URL = "https://huggingface.co/iggy101/MotionVideoSearch/resolve/main/index.faiss"
|
|
EMBEDDINGS_DB_URL = "https://huggingface.co/iggy101/MotionVideoSearch/resolve/main/data.sqlite"
|
|
|
|
# Directory to store downloaded files
|
|
DATA_DIR = Path(__file__).parent / "data"
|
|
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
INDEX_PATH = DATA_DIR / "index.faiss"
|
|
EMBEDDINGS_DB_PATH = DATA_DIR / "data.sqlite"
|
|
EMBEDDING_DIM = 768
|
|
|
|
_dinov2_vitb14_reg = None
|
|
_faiss_index = None
|
|
|
|
_transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
])
|
|
|
|
|
|
def download_file(url, dest_path):
|
|
"""
|
|
Downloads a file from the specified URL to the destination path with a progress bar.
|
|
"""
|
|
try:
|
|
response = requests.get(url, stream=True)
|
|
response.raise_for_status()
|
|
total_size = int(response.headers.get('content-length', 0))
|
|
with open(dest_path, 'wb') as f, tqdm(
|
|
desc=f"Downloading {dest_path.name}",
|
|
total=total_size,
|
|
unit='iB',
|
|
unit_scale=True,
|
|
unit_divisor=1024,
|
|
) as bar:
|
|
for data in response.iter_content(chunk_size=1024):
|
|
size = f.write(data)
|
|
bar.update(size)
|
|
logger.info(f"Downloaded {dest_path.name} successfully.")
|
|
except requests.exceptions.RequestException as e:
|
|
logger.error(f"Failed to download {url}: {e}")
|
|
raise
|
|
|
|
|
|
def ensure_file_exists(file_path, url):
|
|
"""
|
|
Ensures that the file exists locally, and uses ETag to detect if it needs updating.
|
|
If ETag indicates the remote file changed, download the new version.
|
|
"""
|
|
local_etag_path = file_path.with_suffix(file_path.suffix + ".etag")
|
|
|
|
# Attempt to retrieve the remote ETag
|
|
remote_etag = None
|
|
try:
|
|
head_resp = requests.head(url, allow_redirects=True)
|
|
head_resp.raise_for_status()
|
|
remote_etag = head_resp.headers.get("ETag")
|
|
except requests.exceptions.RequestException as e:
|
|
logger.warning(f"Failed to retrieve ETag for {url}: {e}")
|
|
|
|
# If the file does not exist locally, always download
|
|
if not file_path.exists():
|
|
logger.info(f"{file_path.name} not found. Downloading from HuggingFace...")
|
|
download_file(url, file_path)
|
|
if remote_etag:
|
|
local_etag_path.write_text(remote_etag)
|
|
return
|
|
|
|
# If we have a remote ETag, compare to local ETag
|
|
if remote_etag:
|
|
if local_etag_path.exists():
|
|
local_etag = local_etag_path.read_text().strip()
|
|
# If they match, do nothing
|
|
if local_etag == remote_etag:
|
|
logger.info(f"{file_path.name} is already up to date (ETag match).")
|
|
return
|
|
# Otherwise, re-download
|
|
logger.info(f"{file_path.name} is outdated. Downloading new version from HuggingFace...")
|
|
download_file(url, file_path)
|
|
local_etag_path.write_text(remote_etag)
|
|
else:
|
|
# If no remote ETag is available, default to the old logic
|
|
logger.info(f"{file_path.name} already exists (no ETag to compare).")
|
|
|
|
|
|
def load_model_and_index():
|
|
"""
|
|
Loads DINOv2 model and the FAISS index (if not already loaded).
|
|
Downloads required files if they are missing or outdated.
|
|
"""
|
|
global _dinov2_vitb14_reg, _faiss_index
|
|
|
|
# Ensure that data.bin and embeddings.db are present/updated
|
|
ensure_file_exists(INDEX_PATH, DATA_BIN_URL)
|
|
ensure_file_exists(EMBEDDINGS_DB_PATH, EMBEDDINGS_DB_URL)
|
|
|
|
# Load the model lazily
|
|
if _dinov2_vitb14_reg is None:
|
|
logger.info("Loading DINOv2 model...")
|
|
_dinov2_vitb14_reg = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14_reg')
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
_dinov2_vitb14_reg.to(device)
|
|
_dinov2_vitb14_reg.eval()
|
|
|
|
# Load or initialize the FAISS index
|
|
if _faiss_index is None:
|
|
if not INDEX_PATH.exists():
|
|
raise FileNotFoundError(f"FAISS index not found at {INDEX_PATH}. Please run the 'store' command first.")
|
|
logger.info(f"Loading FAISS index from {INDEX_PATH}...")
|
|
_faiss_index = faiss.read_index(str(INDEX_PATH))
|
|
|
|
return _dinov2_vitb14_reg, _faiss_index
|
|
|
|
|
|
class IG_MotionVideoSearch:
|
|
"""
|
|
A ComfyUI node that accepts a ComfyUI image and
|
|
returns 5 ranked search results from the FAISS index based on a given starting rank.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # ComfyUI "IMAGE" type
|
|
"starting_rank": ("INT", {"default": 1, "min": 1, "max": 9999999, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("url_1", "url_2", "url_3", "url_4", "url_5", "status")
|
|
FUNCTION = "search"
|
|
|
|
CATEGORY = "🐓 IG Motion Search Nodes"
|
|
|
|
def search(self, image, starting_rank):
|
|
"""
|
|
Perform the search using the loaded FAISS index and DINOv2 model.
|
|
|
|
:param image: A torch.Tensor, shape [batch_size, C, H, W]
|
|
:param starting_rank: The starting rank for the search results (defaults to 1).
|
|
The node will return this rank plus the next 4 lower-ranked results.
|
|
:return: 5 separate URLs for the search results and a status string with scores
|
|
"""
|
|
logger.debug(f"Image type: {type(image)}")
|
|
logger.debug(f"Image shape: {image.shape}")
|
|
logger.debug(f"Image dtype: {image.dtype}")
|
|
logger.debug(f"Starting rank: {starting_rank}")
|
|
|
|
# 1. Load model and index if needed
|
|
model, index = load_model_and_index()
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# 2. Convert ComfyUI image (torch tensor) into a PIL Image
|
|
c_img = image
|
|
if c_img.ndim == 4:
|
|
# Assume shape is [batch_size, C, H, W]
|
|
if c_img.size(0) > 1:
|
|
logger.warning("Received batch size > 1. Only processing the first image in the batch.")
|
|
c_img = c_img[0]
|
|
else:
|
|
c_img = c_img.squeeze(0)
|
|
elif c_img.ndim != 3:
|
|
raise ValueError(f"Expected image tensor to have 3 or 4 dimensions, got {c_img.ndim}")
|
|
|
|
# Ensure values are in [0, 1] range
|
|
c_img = c_img.clamp(0, 1)
|
|
|
|
# Convert to uint8 and numpy array
|
|
np_img = (c_img * 255.0).byte().cpu().numpy()
|
|
# Convert to PIL image
|
|
pil_img = Image.fromarray(np_img, mode='RGB')
|
|
|
|
# 3. Apply the same resizing logic if desired
|
|
with torch.no_grad():
|
|
tensor_img = _transform(pil_img).unsqueeze(0).to(device) # shape [1, C, H, W]
|
|
_, _, h, w = tensor_img.shape
|
|
new_h = (h // 14) * 14
|
|
new_w = (w // 14) * 14
|
|
h_start = (h - new_h) // 2
|
|
w_start = (w - new_w) // 2
|
|
tensor_img = tensor_img[:, :, h_start: h_start + new_h, w_start: w_start + new_w]
|
|
|
|
# 4. Get the embedding
|
|
embedding = model(tensor_img).cpu().numpy().astype("float32")
|
|
|
|
# We want top_k = starting_rank + 4
|
|
top_k = starting_rank + 4
|
|
distances, ids = index.search(embedding, top_k)
|
|
|
|
if ids.size == 0 or (ids.size == 1 and ids[0][0] == -1):
|
|
return (
|
|
"No embeddings found in the FAISS index.",
|
|
"",
|
|
"",
|
|
"",
|
|
"",
|
|
"No scores available."
|
|
)
|
|
|
|
# Slice out the 5 results we actually want
|
|
selected_distances = distances[0][starting_rank - 1 : starting_rank - 1 + 5]
|
|
selected_ids = ids[0][starting_rank - 1 : starting_rank - 1 + 5]
|
|
|
|
# 6. Retrieve metadata from SQLite
|
|
conn = get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
urls = [""] * 5
|
|
results_str = []
|
|
|
|
for offset, (dist, uid) in enumerate(zip(selected_distances, selected_ids)):
|
|
rank = starting_rank + offset
|
|
if uid == -1:
|
|
results_str.append(f"{rank}. [No valid ID] - Distance: {dist:.4f}")
|
|
continue
|
|
|
|
cursor.execute(
|
|
"""
|
|
SELECT videos.url
|
|
FROM embeddings
|
|
JOIN videos ON embeddings.video_id = videos.id
|
|
WHERE embeddings.id = ?
|
|
""",
|
|
(int(uid),),
|
|
)
|
|
row = cursor.fetchone()
|
|
if row:
|
|
urls[offset] = row[0]
|
|
results_str.append(f"{rank}. Distance: {dist:.4f}")
|
|
else:
|
|
results_str.append(f"{rank}. [Missing DB row for ID {uid}], Distance: {dist:.4f}")
|
|
|
|
conn.close()
|
|
status = "\n".join(results_str)
|
|
|
|
return (*urls, status)
|
|
|
|
|
|
###############################################################################
|
|
# New Node: IG_MotionVideoDotFrame
|
|
###############################################################################
|
|
class IG_MotionVideoFrame:
|
|
"""
|
|
A ComfyUI node that takes a batch of frames from a video, ensures we have at least
|
|
24 frames, trims (cuts off) the video to exactly 24 frames, then calls get_dot_frame
|
|
to produce a colorized motion image. Finally, it returns that as a ComfyUI image.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video_frames": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("colored_motion_image",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "🐓 IG Motion Search Nodes"
|
|
|
|
def apply(self, video_frames):
|
|
"""
|
|
:param video_frames: A 4D torch.Tensor of shape [N, C, H, W], with N >= 24
|
|
:return: A single colorized motion image (torch.Tensor of shape [1, C, H, W])
|
|
"""
|
|
|
|
if not isinstance(video_frames, torch.Tensor):
|
|
raise TypeError("video_frames must be a torch.Tensor.")
|
|
|
|
if video_frames.ndim != 4:
|
|
raise ValueError(
|
|
f"Expected a 4D tensor [N, C, H, W], but got shape {video_frames.shape}."
|
|
)
|
|
|
|
# AI is convinced that Comfy images are B, C, H, W but they're actually B, H, W, C
|
|
video_frames = video_frames.permute(0, 3, 1, 2)
|
|
|
|
num_frames, channels, height, width = video_frames.shape
|
|
if num_frames < 24:
|
|
raise ValueError(
|
|
f"Video must have at least 24 frames, but got {num_frames}."
|
|
)
|
|
|
|
# 1) Trim the video to 24 frames
|
|
video_frames = video_frames[:24] # shape: [24, C, H, W]
|
|
|
|
# 2) Scale so the shorter side is 'fit_to':
|
|
fit_to = 336
|
|
_, _, H, W = video_frames.shape
|
|
|
|
# If width < height, we set width to 336 and scale height accordingly.
|
|
# Otherwise, we set height to 336 and scale width accordingly.
|
|
if W < H:
|
|
new_W = fit_to
|
|
new_H = int(round(H * fit_to / W))
|
|
else:
|
|
new_H = fit_to
|
|
new_W = int(round(W * fit_to / H))
|
|
|
|
# Interpolate all frames to this new size
|
|
video_frames = F.interpolate(
|
|
video_frames,
|
|
size=(new_H, new_W),
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
|
|
# 3) Crop so the height and width are multiples of 8
|
|
new_H_aligned = (new_H // 8) * 8
|
|
new_W_aligned = (new_W // 8) * 8
|
|
|
|
h_start = (new_H - new_H_aligned) // 2
|
|
w_start = (new_W - new_W_aligned) // 2
|
|
|
|
video_frames = video_frames[
|
|
:,
|
|
:,
|
|
h_start : h_start + new_H_aligned,
|
|
w_start : w_start + new_W_aligned
|
|
]
|
|
|
|
# get_dot_frame typically expects the shape [N, C, H, W].
|
|
# This function returns a single frame of shape [C, H, W].
|
|
download_checkpoints()
|
|
frame = get_dot_frame(video_frames)
|
|
|
|
# Convert to ComfyUI's "IMAGE" format: [1, H, W, C]
|
|
frame = frame.unsqueeze(0) # shape = [1, C, H, W]
|
|
frame = frame.permute(0, 2, 3, 1)
|
|
|
|
return (frame,)
|