263 lines
12 KiB
Python
263 lines
12 KiB
Python
# --- 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
|
|
|
|
print("--- [AIIA] Loading Media Browser API Endpoints (V19 - Robustness Fixes) ---")
|
|
|
|
# --- START OF FIX: Path and Cache Logic Refinement ---
|
|
try:
|
|
# Attempt to find the ComfyUI root directory robustly
|
|
comfyui_root_path = Path(__file__).resolve().parents[2]
|
|
if not (comfyui_root_path / "main.py").exists():
|
|
current_dir = Path(__file__).resolve().parent
|
|
while not (current_dir / "main.py").exists() and current_dir != current_dir.parent:
|
|
current_dir = current_dir.parent
|
|
if (current_dir / "main.py").exists():
|
|
comfyui_root_path = current_dir
|
|
else:
|
|
raise FileNotFoundError("Could not find ComfyUI root directory.")
|
|
output_dir = comfyui_root_path / "output"
|
|
print(f"--- [AIIA] Successfully located ComfyUI root. Output directory set to: {output_dir}")
|
|
|
|
except Exception as e:
|
|
# Fallback path if auto-detection fails
|
|
print(f"--- [AIIA] Warning: Could not auto-detect ComfyUI root path ({e}). Falling back to default './ComfyUI/output'.")
|
|
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)
|
|
print(f"--- [AIIA] Image thumbnail cache: {image_thumb_dir}")
|
|
print(f"--- [AIIA] Video poster cache: {video_poster_dir}")
|
|
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()
|
|
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()
|
|
target_path.relative_to(safe_base)
|
|
return target_path
|
|
except ValueError:
|
|
raise PermissionError("Forbidden path, not within the base directory.")
|
|
|
|
|
|
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:
|
|
metadata["width"], metadata["height"] = img.size
|
|
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 and 'duration' in ffprobe_data['format']:
|
|
metadata['duration'] = float(ffprobe_data['format']['duration'])
|
|
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:
|
|
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="Directory not found")
|
|
|
|
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'})
|
|
|
|
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')
|
|
|
|
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(): return web.Response(status=404, text="File not found")
|
|
|
|
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() and cached_poster_path.stat().st_mtime > original_file_path.stat().st_mtime:
|
|
return web.FileResponse(cached_poster_path, headers={'Content-Type': 'image/jpeg'})
|
|
|
|
# 使用临时文件确保 ffmpeg 操作的原子性
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmpfile:
|
|
tmp_poster_path = tmpfile.name
|
|
|
|
command = [
|
|
"ffmpeg",
|
|
"-ss", "00:00:01.00",
|
|
"-i", str(original_file_path),
|
|
"-frames:v", "1",
|
|
"-vf", "thumbnail,scale=256:-1",
|
|
"-y", # Overwrite output file if it exists
|
|
str(tmp_poster_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:
|
|
if os.path.exists(tmp_poster_path): os.remove(tmp_poster_path)
|
|
print(f"[AIIA] FFmpeg error for {filename}: {stderr.decode()}")
|
|
return web.Response(status=500, text="Failed to extract frame from video.")
|
|
|
|
# 检查生成的文件是否有效
|
|
if 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)
|
|
return web.Response(status=500, text="FFmpeg produced an empty file.")
|
|
|
|
# 原子性地移动/重命名临时文件到缓存位置
|
|
shutil.move(tmp_poster_path, cached_poster_path)
|
|
|
|
return web.FileResponse(cached_poster_path, headers={'Content-Type': 'image/jpeg'})
|
|
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 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))
|
|
|
|
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)
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
print("--- [AIIA] Media Browser API Endpoints (V19) loaded successfully. ---") |