diff --git a/requirements.txt b/requirements.txt index f87addc..cf61d96 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,6 @@ fastapi uvicorn piexif python-dotenv -Pillow \ No newline at end of file +Pillow +imageio +av \ No newline at end of file diff --git a/scripts/iib/api.py b/scripts/iib/api.py index adeeef0..d1a3ca4 100644 --- a/scripts/iib/api.py +++ b/scripts/iib/api.py @@ -1,14 +1,16 @@ from datetime import datetime, timedelta import os +from pathlib import Path import shutil import sqlite3 + from scripts.iib.tool import ( comfyui_exif_data_to_str, get_comfyui_exif_data, human_readable_size, is_img_created_by_comfyui, is_img_created_by_comfyui_with_webui_gen_info, - is_valid_image_path, + is_valid_media_path, temp_path, read_sd_webui_gen_info_from_image, get_formatted_date, @@ -28,12 +30,12 @@ from scripts.iib.tool import ( is_secret_key_required, open_file_with_default_app ) -from fastapi import FastAPI, HTTPException, Response +from fastapi import FastAPI, HTTPException, Header, Response from fastapi.staticfiles import StaticFiles import asyncio from typing import List, Optional from pydantic import BaseModel -from fastapi.responses import FileResponse, JSONResponse +from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from PIL import Image from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware @@ -53,6 +55,7 @@ from scripts.iib.db.update_image_data import update_image_data, rebuild_image_in from scripts.iib.logger import logger from scripts.iib.seq import seq import urllib.parse +from scripts.iib.fastapi_video import range_requests_response index_html_path = os.path.join(cwd, "vue/dist/index.html") # 在app.py也被使用 @@ -493,7 +496,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs): encoded_filename = urllib.parse.quote(disposition.encode('utf-8')) headers['Content-Disposition'] = f"attachment; filename*=UTF-8''{encoded_filename}" - if is_path_under_parents(filename) and is_valid_image_path( + if is_path_under_parents(filename) and is_valid_media_path( filename ): # 认为永远不变,不要协商缓存了试试 headers[ @@ -508,6 +511,54 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs): media_type=media_type, headers=headers, ) + + @app.get(api_base + "/stream_video", dependencies=[Depends(verify_secret)]) + async def stream_video(path: str, request: Request): + import mimetypes + print(path) + print(request) + media_type, _ = mimetypes.guess_type(path) + return range_requests_response( + request, file_path=path, content_type=media_type + ) + + @app.get(api_base + "/video_cover", dependencies=[Depends(verify_secret)]) + async def video_cover(path: str, t: str): + check_path_trust(path) + if not temp_path: + return + # 生成缓存文件的路径 + hash_dir = hashlib.md5((path + t).encode("utf-8")).hexdigest() + hash = hash_dir + cache_dir = os.path.join(temp_path, "iib_cache", "video_cover", hash_dir) + cache_path = os.path.join(cache_dir, "cover.webp") + + # 如果缓存文件存在，则直接返回该文件 + if os.path.exists(cache_path): + print("命中缓存") + return FileResponse( + cache_path, + media_type="image/webp", + headers={"Cache-Control": "max-age=31536000", "ETag": hash}, + ) + # 如果缓存文件不存在，则生成缩略图并保存 + + import imageio.v3 as iio + frame = iio.imread( + path, + index=16, + plugin="pyav", + ) + + os.makedirs(cache_dir, exist_ok=True) + iio.imwrite(cache_path,frame, extension=".webp") + + # 返回缓存文件 + return FileResponse( + cache_path, + media_type="image/webp", + headers={"Cache-Control": "max-age=31536000", "ETag": hash}, + ) @app.post(api_base + "/send_img_path", dependencies=[Depends(verify_secret)]) async def api_set_send_img_path(path: str): @@ -728,7 +779,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs): @app.get(db_api_base + "/img_selected_custom_tag", dependencies=[Depends(verify_secret)]) async def get_img_selected_custom_tag(path: str): path = os.path.normpath(path) - if not is_valid_image_path(path): + if not is_valid_media_path(path): return [] conn = DataBase.get_conn() update_extra_paths(conn) diff --git a/scripts/iib/db/update_image_data.py b/scripts/iib/db/update_image_data.py index d91fc1c..ecbf4fd 100644 --- a/scripts/iib/db/update_image_data.py +++ b/scripts/iib/db/update_image_data.py @@ -6,8 +6,9 @@ from PIL import Image from scripts.iib.tool import ( read_sd_webui_gen_info_from_image, parse_generation_parameters, - is_valid_image_path, + is_valid_media_path, get_modified_date, + get_video_type, is_dev, get_comfyui_exif_data, comfyui_exif_data_to_str, @@ -22,6 +23,8 @@ from scripts.iib.logger import logger def get_exif_data(file_path): info = '' params = None + if get_video_type(file_path): + return params, info try: with Image.open(file_path) as img: if is_img_created_by_comfyui(img): @@ -62,7 +65,7 @@ def update_image_data(search_dirs: List[str], is_rebuild = False): if os.path.isdir(file_path): process_folder(file_path) - elif is_valid_image_path(file_path): + elif is_valid_media_path(file_path): img = DbImg.get(conn, file_path) if is_rebuild: parsed_params, info = get_exif_data(file_path) diff --git a/scripts/iib/fastapi_video.py b/scripts/iib/fastapi_video.py new file mode 100644 index 0000000..bdfef6f --- /dev/null +++ b/scripts/iib/fastapi_video.py @@ -0,0 +1,74 @@ +import os +from typing import BinaryIO + +from fastapi import FastAPI, HTTPException, Request, status +from fastapi.responses import StreamingResponse + + +def send_bytes_range_requests( + file_obj: BinaryIO, start: int, end: int, chunk_size: int = 10_000 +): + """Send a file in chunks using Range Requests specification RFC7233 + + `start` and `end` parameters are inclusive due to specification + """ + with file_obj as f: + f.seek(start) + while (pos := f.tell()) <= end: + read_size = min(chunk_size, end + 1 - pos) + yield f.read(read_size) + + +def _get_range_header(range_header: str, file_size: int) -> tuple[int, int]: + def _invalid_range(): + return HTTPException( + status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE, + detail=f"Invalid request range (Range:{range_header!r})", + ) + + try: + h = range_header.replace("bytes=", "").split("-") + start = int(h[0]) if h[0] != "" else 0 + end = int(h[1]) if h[1] != "" else file_size - 1 + except ValueError: + raise _invalid_range() + + if start > end or start < 0 or end > file_size - 1: + raise _invalid_range() + return start, end + + +def range_requests_response( + request: Request, file_path: str, content_type: str +): + """Returns StreamingResponse using Range Requests of a given file""" + + file_size = os.stat(file_path).st_size + range_header = request.headers.get("range") + + headers = { + "content-type": content_type, + "accept-ranges": "bytes", + "content-encoding": "identity", + "content-length": str(file_size), + "access-control-expose-headers": ( + "content-type, accept-ranges, content-length, " + "content-range, content-encoding" + ), + } + start = 0 + end = file_size - 1 + status_code = status.HTTP_200_OK + + if range_header is not None: + start, end = _get_range_header(range_header, file_size) + size = end - start + 1 + headers["content-length"] = str(size) + headers["content-range"] = f"bytes {start}-{end}/{file_size}" + status_code = status.HTTP_206_PARTIAL_CONTENT + + return StreamingResponse( + send_bytes_range_requests(open(file_path, mode="rb"), start, end), + headers=headers, + status_code=status_code, + ) diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py index f91abf6..d79b78b 100644 --- a/scripts/iib/tool.py +++ b/scripts/iib/tool.py @@ -167,9 +167,17 @@ def convert_to_bytes(file_size_str): return int(size) else: raise ValueError(f"Invalid file size string '{file_size_str}'") + +def get_video_type(file_path): + video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv'] + file_extension = file_path[file_path.rfind('.'):].lower() + if file_extension in video_extensions: + return file_extension[1:] + else: + return None -def is_valid_image_path(path): +def is_valid_media_path(path): """ 判断给定的路径是否是图像文件 """ @@ -178,7 +186,7 @@ def is_valid_image_path(path): return False if not os.path.isfile(abs_path): # 判断是否是文件 return False - if not imghdr.what(abs_path): # 判断是否是图像文件 + if not imghdr.what(abs_path) and not get_video_type(abs_path): # 判断是否是图像文件 return False return True diff --git a/vue/src/components/ContextMenu.vue b/vue/src/components/ContextMenu.vue index b1cc676..c35f853 100644 --- a/vue/src/components/ContextMenu.vue +++ b/vue/src/components/ContextMenu.vue @@ -2,7 +2,7 @@ import type { Tag } from '@/api/db' import type { FileNodeInfo } from '@/api/files' import type { MenuInfo } from 'ant-design-vue/lib/menu/src/interface' -import { isImageFile } from '@/util' +import { isMediaFile } from '@/util' import { StarFilled, StarOutlined } from '@/icon' import { useGlobalStore } from '@/store/useGlobalStore' import { computed } from 'vue' @@ -35,7 +35,7 @@ const tags = computed(() => { {{ $t('openWithWalkMode') }}