IDGallagher-MotionVideoSearch/nodes.py

# 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
from .database import get_connection  # Ensure accessible to ComfyUI
from torchvision import transforms

logger = logging.getLogger(__name__)

# Define URLs for data.bin and embeddings.db on HuggingFace
DATA_BIN_URL = "https://huggingface.co/your-username/your-repo/raw/main/data.bin"
EMBEDDINGS_DB_URL = "https://huggingface.co/your-username/your-repo/raw/main/embeddings.db"

# Directory to store downloaded files
DATA_DIR = Path(__file__).parent / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)

INDEX_PATH = DATA_DIR / "data.bin"
EMBEDDINGS_DB_PATH = DATA_DIR / "embeddings.db"
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; downloads it from the URL if it does not.
    """
    if not file_path.exists():
        logger.info(f"{file_path.name} not found. Downloading from HuggingFace...")
        download_file(url, file_path)
    else:
        logger.info(f"{file_path.name} already exists.")

def load_model_and_index():
    """
    Loads DINOv2 model and the FAISS index (if not already loaded).
    Downloads required files if they are missing.
    """
    global _dinov2_vitb14_reg, _faiss_index

    # Ensure that data.bin and embeddings.db are present
    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 the top 5 search results from the FAISS index.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),  # ComfyUI "IMAGE" type
                "top_k": ("INT", {"default": 5, "min": 1, "max": 20, "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 = "Motion Video DB"  # Appears in ComfyUI under this category in the node menu

    def search(self, image, top_k):
        """
        Perform the search using the loaded FAISS index and DINOv2 model.

        :param image: A torch.Tensor, shape [batch_size, C, H, W]
        :param top_k: Number of top results to retrieve
        :return: 5 separate URLs for the search results and a status string with scores
        """
        # Log input details for debugging
        logger.debug(f"Image type: {type(image)}")
        logger.debug(f"Image shape: {image.shape}")
        logger.debug(f"Image dtype: {image.dtype}")

        # 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  # Directly use the tensor without indexing
        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)  # Remove the batch dimension
        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 as your main.py does (multiple of 14, etc.)
        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")

        # 5. Search in FAISS
        distances, ids = index.search(embedding, top_k)

        # Handle edge cases
        if ids.size == 0 or (ids.size == 1 and ids[0][0] == -1):
            return ("No embeddings found in the FAISS index.", "", "", "", "", "No scores available.")

        # 6. Retrieve metadata from SQLite
        conn = get_connection()
        cursor = conn.cursor()

        urls = [""] * 5  # Initialize list of 5 URL strings
        results_str = []

        for rank, (dist, uid) in enumerate(zip(distances[0], ids[0]), start=1):
            if rank > 5:
                break
            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[rank - 1] = 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)