havvk-ComfyUI_AIIA/aiia_browser_node.py

# --- START OF FILE aiia_browser_node.py (V19 - 健壮性终极修复) ---

import os
import server
from aiohttp import web
import json
from pathlib import Path
import time
from PIL import Image, ImageOps
import asyncio
import shutil
import traceback
from io import BytesIO
import re
import tempfile
import folder_paths

try:
    # Attempt to find the ComfyUI root directory robustly
    output_dir = Path(folder_paths.get_output_directory())

except Exception as e:
    # Fallback path if auto-detection fails
    output_dir = Path("./ComfyUI/output")

# Define cache directories based on the final output_dir. This ensures consistency.
cache_main_dir = output_dir / ".aiia_cache"
image_thumb_dir = cache_main_dir / "thumbnails"
video_poster_dir = cache_main_dir / "posters"

# Create all necessary directories
try:
    output_dir.mkdir(parents=True, exist_ok=True)
    cache_main_dir.mkdir(exist_ok=True)
    image_thumb_dir.mkdir(exist_ok=True)
    video_poster_dir.mkdir(exist_ok=True)
except Exception as e:
    print(f"--- [AIIA] CRITICAL ERROR: Could not create output or cache directories at {output_dir}. Please check permissions. Error: {e}")
# --- END OF FIX: Path and Cache Logic Refinement ---



CONCURRENT_LIMIT = 16
semaphore = asyncio.Semaphore(CONCURRENT_LIMIT)

def get_safe_path(base_dir, relative_path_str):
    try:
        safe_base = base_dir.resolve()
        
        # If the relative path is empty or just a dot, return the base directory itself.
        if not relative_path_str or relative_path_str == '.':
            return safe_base

        clean_path = os.path.normpath(relative_path_str)
        if '..' in clean_path.split(os.sep):
             raise PermissionError("Path traversal attempt detected.")
        
        target_path = (safe_base / Path(clean_path)).resolve()
        
        # Ensure the resolved target path is still within the base directory.
        target_path.relative_to(safe_base)
        
        return target_path
    except ValueError:
        raise PermissionError("Forbidden path, not within the base directory.")
    except Exception as e:
        print(f"[AIIA Browser] Error in get_safe_path: {e}")
        raise


async def _get_file_metadata(file_path: Path):
    async with semaphore:
        metadata = {}
        try:
            if not file_path.exists() or file_path.stat().st_size == 0:
                return metadata
        except FileNotFoundError:
             return metadata 

        ext = file_path.suffix.lower()
        try:
            if ext in {'.png', '.jpg', '.jpeg', '.gif', '.webp'}:
                with Image.open(file_path) as img:
                    img.load() # Force loading of metadata chunks
                    metadata["width"], metadata["height"] = img.size
                    if 'prompt' in img.info or 'workflow' in img.info:
                        metadata['has_workflow'] = True
            elif ext in {'.mp4', '.mov', '.avi', '.wav', '.mp3', '.ogg'}:
                if not shutil.which("ffprobe"): return metadata
                command = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", str(file_path)]
                proc = await asyncio.create_subprocess_exec(*command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
                stdout, stderr = await proc.communicate()
                if proc.returncode == 0:
                    ffprobe_data = json.loads(stdout)
                    if 'format' in ffprobe_data:
                        if 'duration' in ffprobe_data['format']:
                            metadata['duration'] = float(ffprobe_data['format']['duration'])
                        # Check for workflow in format tags (e.g., 'comment' for MP4/MOV)
                        tags = ffprobe_data['format'].get('tags', {})
                        for key, value in tags.items():
                            # A more strict check for workflow JSON
                         if isinstance(value, str) and value.strip().startswith('{') and '"nodes"' in value and '"links"' in value:
                             metadata['has_workflow'] = True
                             break # Found it, no need to check other tags
                    
                    video_stream = next((s for s in ffprobe_data.get('streams', []) if s.get('codec_type') == 'video'), None)
                    if video_stream:
                        metadata['width'] = video_stream.get('width')
                        metadata['height'] = video_stream.get('height')
        except Exception as e:
            print(f"--- [AIIA] Warning: Could not get metadata for {file_path}. Reason: {e}")
            pass
        return metadata

async def list_items(request):
    relative_path_str = request.query.get("path", "")
    try:
        target_path = get_safe_path(output_dir, relative_path_str)
        if not target_path.is_dir(): 
            return web.Response(status=404, text=f"Directory not found: {target_path}")
        
        directories, files = [], []
        for item in os.scandir(target_path):
            # --- START OF FIX: Precise filtering ---
            # Only filter our specific cache directory, not all dotfiles.
            if item.name == '.aiia_cache':
                continue
            # --- END OF FIX: Precise filtering ---
            try:
                if item.is_dir():
                    try: item_count = len([name for name in os.listdir(item.path) if name != '.aiia_cache']) # Also exclude from count
                    except OSError: item_count = 0
                    directories.append({"name": item.name, "type": "directory", "mtime": item.stat().st_mtime, "item_count": item_count})
                elif item.is_file():
                    ext = Path(item.name).suffix.lower()
                    if ext in {'.png', '.jpg', '.jpeg', '.gif', '.webp', '.mp4', '.mov', '.avi', '.wav', '.mp3', '.ogg'}:
                        stat_info = item.stat()
                        file_data = {"name": item.name, "type": "file", "size": stat_info.st_size, "mtime": stat_info.st_mtime, "extension": ext}
                        files.append(file_data)
            except Exception:
                continue 
        
        def natural_sort_key(s, _re=re.compile(r'([0-9]+)')):
            return [int(text) if text.isdigit() else text.lower() for text in _re.split(s['name'])]

        directories.sort(key=natural_sort_key)
        files.sort(key=natural_sort_key)

        return web.json_response({"directories": directories, "files": files})
    except Exception as e:
        traceback.print_exc()
        return web.Response(status=500, text=str(e))

async def get_thumbnail(request):
    relative_path_str = request.query.get("path", "")
    filename = request.query.get("filename", "")
    if not filename: return web.Response(status=400, text="Filename is required.")
    
    try:
        original_file_path = get_safe_path(output_dir, os.path.join(relative_path_str, filename))
        if not original_file_path.is_file(): return web.Response(status=404, text="File not found")

        cache_sub_dir = get_safe_path(image_thumb_dir, relative_path_str)
        cache_sub_dir.mkdir(parents=True, exist_ok=True)
        cached_thumb_path = cache_sub_dir / f"{Path(filename).stem}.jpg"

        if cached_thumb_path.exists() and cached_thumb_path.stat().st_mtime > original_file_path.stat().st_mtime:
            return web.FileResponse(cached_thumb_path, headers={'Content-Type': 'image/jpeg', 'Cache-Control': 'public, must-revalidate'})

        with Image.open(original_file_path) as img:
            img = ImageOps.exif_transpose(img)
            if img.mode in ("RGBA", "P"): img = img.convert("RGB")
            img.thumbnail((256, 256), Image.Resampling.LANCZOS)
            
            buffer = BytesIO()
            img.save(buffer, format="JPEG", quality=85, optimize=True)
            buffer.seek(0)
            
            async def write_cache():
                try:
                    with open(cached_thumb_path, 'wb') as f: f.write(buffer.getvalue())
                except Exception as e: print(f"[AIIA] Error writing thumbnail cache for {filename}: {e}")
            
            asyncio.create_task(write_cache())
            
            return web.Response(body=buffer.getvalue(), content_type='image/jpeg', headers={'Cache-Control': 'public, must-revalidate'})

    except Exception as e:
        traceback.print_exc()
        return web.Response(status=500, text=str(e))

async def get_video_poster(request):
    relative_path_str = request.query.get("path", "")
    filename = request.query.get("filename", "")
    if not filename: return web.Response(status=400, text="Filename is required.")

    if not shutil.which("ffmpeg"):
        return web.Response(status=501, text="ffmpeg not found on server.")

    try:
        original_file_path = get_safe_path(output_dir, os.path.join(relative_path_str, filename))
        if not original_file_path.is_file() or original_file_path.stat().st_size == 0:
            return web.Response(status=404, text="File not found or is empty")

        cache_sub_dir = get_safe_path(video_poster_dir, relative_path_str)
        cache_sub_dir.mkdir(parents=True, exist_ok=True)
        cached_poster_path = cache_sub_dir / f"{Path(filename).stem}.jpg"
        
        if cached_poster_path.exists():
            # If the cached file is valid and up-to-date, serve it.
            if cached_poster_path.stat().st_size > 0 and cached_poster_path.stat().st_mtime > original_file_path.stat().st_mtime:
                return web.FileResponse(cached_poster_path, headers={'Content-Type': 'image/jpeg', 'Cache-Control': 'public, must-revalidate'})
            else:
                # Delete invalid (0-byte) or outdated cache files.
                print(f"--- [AIIA Browser Debug] Deleting invalid or outdated cache file: {cached_poster_path}")
                os.remove(cached_poster_path)
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmpfile:
            tmp_poster_path = tmpfile.name

        # Use the 'thumbnail' filter with the 'update' flag, which is robust for both
        # normal videos and special cases like single-frame videos.
        command = [
            "ffmpeg",
            "-i", str(original_file_path),
            "-vf", "thumbnail,scale=256:-1:force_original_aspect_ratio=decrease",
            "-frames:v", "1",
            "-update", "1",
            "-q:v", "3",
            "-y",
            str(tmp_poster_path)
        ]
        
        proc = await asyncio.create_subprocess_exec(*command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        _, stderr = await proc.communicate()

        if proc.returncode != 0 or not os.path.exists(tmp_poster_path) or os.path.getsize(tmp_poster_path) == 0:
            if os.path.exists(tmp_poster_path): os.remove(tmp_poster_path)
            print(f"[AIIA] FFmpeg failed to extract thumbnail for {filename}. Error: {stderr.decode()}")
            return web.Response(status=500, text="Failed to extract thumbnail from video.")

        shutil.move(tmp_poster_path, cached_poster_path)
        return web.FileResponse(cached_poster_path, headers={'Content-Type': 'image/jpeg', 'Cache-Control': 'public, must-revalidate'})

    except Exception as e:
        if 'tmp_poster_path' in locals() and os.path.exists(tmp_poster_path):
            os.remove(tmp_poster_path)
        traceback.print_exc()
        return web.Response(status=500, text=str(e))

async def delete_item(request):
    data = await request.json()
    relative_path_str = data.get("path", "")
    filename = data.get("filename", "")
    if not filename:
        return web.Response(status=400, text="Filename is required.")

    try:
        file_path = get_safe_path(output_dir, os.path.join(relative_path_str, filename))
        
        if file_path.is_file():
            os.remove(file_path)
            # Also cleanup potential cache files
            cache_thumb = get_safe_path(image_thumb_dir, os.path.join(relative_path_str, f"{file_path.stem}.jpg"))
            if cache_thumb.exists(): os.remove(cache_thumb)
            cache_poster = get_safe_path(video_poster_dir, os.path.join(relative_path_str, f"{file_path.stem}.jpg"))
            if cache_poster.exists(): os.remove(cache_poster)
            
            return web.json_response({"status": "success", "message": f"File {filename} deleted."})
        elif file_path.is_dir():
            shutil.rmtree(file_path)
            return web.json_response({"status": "success", "message": f"Directory {filename} deleted."})
        else:
            return web.Response(status=404, text="Item not found")

    except Exception as e:
        traceback.print_exc()
        return web.Response(status=500, text=str(e))

async def get_batch_metadata(request):
    data = await request.json()
    path = data.get("path", "")
    filenames = data.get("filenames", [])
    try:
        safe_base = output_dir.resolve()
        tasks = []
        valid_filenames = []
        for filename in filenames:
            file_path = (safe_base / Path(path) / filename).resolve()
            if safe_base in file_path.parents and file_path.is_file():
                tasks.append(_get_file_metadata(file_path))
                valid_filenames.append(filename)
        results = await asyncio.gather(*tasks)
        response_data = {valid_filenames[i]: results[i] for i in range(len(valid_filenames))}
        return web.json_response(response_data)
    except Exception as e:
        traceback.print_exc()
        return web.Response(status=500, text=str(e))

import math

def replace_nan_with_null(obj):
    """
    Recursively traverses a Python object (dicts, lists) and replaces
    float('nan') with None, which serializes to JSON 'null'.
    """
    if isinstance(obj, dict):
        return {k: replace_nan_with_null(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [replace_nan_with_null(elem) for elem in obj]
    # The core of the fix: only replace actual float NaN values
    elif isinstance(obj, float) and math.isnan(obj):
        return None
    else:
        return obj

async def get_workflow(request):
    relative_path_str = request.query.get("path", "")
    filename = request.query.get("filename", "")
    if not filename:
        return web.Response(status=400, text="Filename is required.")

    try:
        file_path = get_safe_path(output_dir, os.path.join(relative_path_str, filename))
        if not file_path.is_file():
            return web.Response(status=404, text="File not found")

        ext = file_path.suffix.lower()
        workflow_obj = None

        if ext in {'.png', '.jpg', '.jpeg', '.gif', '.webp'}:
            with Image.open(file_path) as img:
                workflow_data_str = img.info.get('workflow') or img.info.get('prompt')
                if workflow_data_str:
                    workflow_obj = json.loads(workflow_data_str)
        
        elif ext in {'.mp4', '.mov', '.avi', '.webm'}:
            if not shutil.which("ffprobe"):
                return web.Response(status=501, text="ffprobe is not installed on the server.")
            
            command = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", str(file_path)]
            proc = await asyncio.create_subprocess_exec(*command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
            stdout, stderr = await proc.communicate()

            if proc.returncode == 0:
                ffprobe_data = json.loads(stdout)
                tags = ffprobe_data.get("format", {}).get("tags", {})
                for key, value in tags.items():
                    if isinstance(value, str) and value.strip().startswith('{') and '"nodes"' in value:
                        workflow_obj = json.loads(value)
                        break
        
        if workflow_obj:
            sanitized_workflow_obj = replace_nan_with_null(workflow_obj)
            return web.json_response(sanitized_workflow_obj)
        else:
            return web.Response(status=404, text="No workflow data found in the file.")

    except Exception as e:
        traceback.print_exc()
        return web.Response(status=500, text=str(e))

server.PromptServer.instance.app.router.add_get('/api/aiia/v1/browser/list_items', list_items)
server.PromptServer.instance.app.router.add_post('/api/aiia/v1/browser/get_batch_metadata', get_batch_metadata)
server.PromptServer.instance.app.router.add_get('/api/aiia/v1/browser/thumbnail', get_thumbnail)
server.PromptServer.instance.app.router.add_get('/api/aiia/v1/browser/poster', get_video_poster)
server.PromptServer.instance.app.router.add_get('/api/aiia/v1/browser/get_workflow', get_workflow)
server.PromptServer.instance.app.router.add_post('/api/aiia/v1/browser/delete_item', delete_item)

NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}