Add support for video streaming, video thumbnail generation, video search functionality, and video tagging

requirements.txt

@@ -2,4 +2,6 @@ fastapi
 uvicorn
 piexif
 python-dotenv
-Pillow
+Pillow
+imageio
+av

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)

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)

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,
+    )

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
 

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(() => {
       <a-menu-item key="openWithWalkMode">{{ $t('openWithWalkMode') }}</a-menu-item>
     </template>
     <template v-if="file.type === 'file'">
-      <template v-if="isImageFile(file.name)">
+      <template v-if="isMediaFile(file.name)">
         <a-menu-item key="viewGenInfo">{{ $t('viewGenerationInfo') }}</a-menu-item>
         <a-menu-divider />
         <template v-if="global.conf?.launch_mode !== 'server'">

vue/src/components/FileItem.vue

@@ -4,16 +4,17 @@ import { useGlobalStore } from '@/store/useGlobalStore'
 import { fallbackImage, ok } from 'vue3-ts-util'
 import type { FileNodeInfo } from '@/api/files'
 import { isImageFile, isVideoFile } from '@/util'
-import { toImageThumbnailUrl, toRawFileUrl } from '@/util/file'
+import { toImageThumbnailUrl, toVideoCoverUrl, toRawFileUrl } from '@/util/file'
 import type { MenuInfo } from 'ant-design-vue/lib/menu/src/interface'
 import { computed } from 'vue'
 import ContextMenu from './ContextMenu.vue'
 import ChangeIndicator from './ChangeIndicator.vue'
 import { useTagStore } from '@/store/useTagStore'
-import { CloseCircleOutlined, StarFilled, StarOutlined, PlayCircleFilled } from '@/icon'
+import { CloseCircleOutlined, StarFilled, StarOutlined } from '@/icon'
 import { Tag } from '@/api/db'
 import { openVideoModal } from './functionalCallableComp'
 import type { GenDiffInfo } from '@/api/files'
+import { play } from '@/icon'
 
 const global = useGlobalStore()
 const tagStore = useTagStore()
@@ -35,19 +36,23 @@ const props = withDefaults(
     enableChangeIndicator?: boolean
     extraTags?: Tag[]
   }>(),
-  { selected: false, enableRightClickMenu: true, enableCloseIcon: false, genDiffToNext: () => ({
+  {
-    empty: true,
+    selected: false, enableRightClickMenu: true, enableCloseIcon: false, genDiffToNext: () => ({
-    ownFile: "",
+      empty: true,
-    otherFile: "",
+      ownFile: "",
-    diff: "",    
+      otherFile: "",
-}), genDiffToPrevious: () => ({
+      diff: "",
-    empty: true,
+    }), genDiffToPrevious: () => ({
-    ownFile: "",
+      empty: true,
-    otherFile: "",
+      ownFile: "",
-    diff: "",    
+      otherFile: "",
-}) }
+      diff: "",
+    })
+  }
 )
 
+
+
 const emit = defineEmits<{
   'update:showMenuIdx': [v: number],
   'fileItemClick': [event: MouseEvent, file: FileNodeInfo, idx: number],
@@ -85,9 +90,9 @@ const taggleLikeTag = () => {
   <a-dropdown :trigger="['contextmenu']" :visible="!global.longPressOpenContextMenu ? undefined : typeof idx === 'number' && showMenuIdx === idx
     " @update:visible="(v: boolean) => typeof idx === 'number' && emit('update:showMenuIdx', v ? idx : -1)">
     <li class="file file-item-trigger grid" :class="{
-      clickable: file.type === 'dir',
+    clickable: file.type === 'dir',
-      selected
+    selected
-    }" :data-idx="idx" :key="file.name" draggable="true" @dragstart="emit('dragstart', $event, idx)"
+  }" :data-idx="idx" :key="file.name" draggable="true" @dragstart="emit('dragstart', $event, idx)"
       @dragend="emit('dragend', $event, idx)" @click.capture="emit('fileItemClick', $event, file, idx)">
 
       <div>
@@ -124,22 +129,25 @@ const taggleLikeTag = () => {
         <div :key="file.fullpath" :class="`idx-${idx} item-content`" v-if="isImageFile(file.name)">
 
           <!-- change indicators -->
-          <ChangeIndicator  v-if="enableChangeIndicator" :gen-diff-to-next="genDiffToNext" :gen-diff-to-previous="genDiffToPrevious"/>
+          <ChangeIndicator v-if="enableChangeIndicator" :gen-diff-to-next="genDiffToNext"
+            :gen-diff-to-previous="genDiffToPrevious" />
           <!-- change indicators END -->
 
           <a-image :src="imageSrc" :fallback="fallbackImage" :preview="{
-            src: fullScreenPreviewImageUrl,
+    src: fullScreenPreviewImageUrl,
-            onVisibleChange: (v: boolean, lv: boolean) => emit('previewVisibleChange', v, lv)
+    onVisibleChange: (v: boolean, lv: boolean) => emit('previewVisibleChange', v, lv)
-          }" />
+  }" />
           <div class="tags-container" v-if="customTags && cellWidth > 128">
             <a-tag v-for="tag in extraTags ?? customTags" :key="tag.id" :color="tagStore.getColor(tag.name)">
               {{ tag.name }}
             </a-tag>
           </div>
         </div>
-        <div :class="`idx-${idx} item-content video`" v-else-if="isVideoFile(file.name)" @click="openVideoModal(file)">
+        <div :class="`idx-${idx} item-content video`" :urld="toVideoCoverUrl(file)" :style="{ 'background-image': `url('${toVideoCoverUrl(file)}')` }"
+          v-else-if="isVideoFile(file.name)" @click="openVideoModal(file)">
+      
           <div class="play-icon">
-            <PlayCircleFilled />
+            <img :src="play" style="width: 40px;height: 40px;">
           </div>
           <div class="tags-container" v-if="customTags && cellWidth > 128">
             <a-tag v-for="tag in customTags" :key="tag.id" :color="tagStore.getColor(tag.name)">
@@ -174,7 +182,6 @@ const taggleLikeTag = () => {
   </a-dropdown>
 </template>
 <style lang="scss" scoped>
-
 .center {
   display: flex;
   justify-content: center;
@@ -190,6 +197,7 @@ const taggleLikeTag = () => {
     overflow: hidden;
     width: v-bind('$props.cellWidth + "px"');
     height: v-bind('$props.cellWidth + "px"');
+    background-size: cover;
     cursor: pointer;
   }
 
@@ -197,8 +205,9 @@ const taggleLikeTag = () => {
     position: absolute;
     top: 50%;
     left: 50%;
-    font-size: 3em;
+    transform: translate(-50%, -50%); 
-    transform: translate(-50%, -50%);
+    border-radius: 100%;
+    display: flex;
   }
 
   .tags-container {
@@ -345,4 +354,5 @@ const taggleLikeTag = () => {
     flex-direction: column;
     align-items: flex-end;
   }
-}</style>
+}
+</style>

vue/src/components/functionalCallableComp.tsx

@@ -3,7 +3,7 @@ import { ref } from 'vue'
 import * as Path from '@/util/path'
 import { FileNodeInfo, mkdirs } from '@/api/files'
 import { t } from '@/i18n'
-import { downloadFiles, globalEvents, toRawFileUrl } from '@/util'
+import { downloadFiles, globalEvents, toRawFileUrl, toStreamVideoUrl } from '@/util'
 import { DownloadOutlined } from '@/icon'
 import { isStandalone } from '@/util/env'
 import { rebuildImageIndex } from '@/api/db'
@@ -52,7 +52,7 @@ export const openVideoModal = (file: FileNodeInfo) => {
           flexDirection: 'column'
         }}
       >
-        <video style={{ maxHeight: isStandalone ? '80vh' : '60vh' }} src={toRawFileUrl(file)} controls autoplay></video>
+        <video style={{ maxHeight: isStandalone ? '80vh' : '60vh' }} src={toStreamVideoUrl(file)} controls autoplay></video>
         <div class="actions" style={{ marginTop: '16px' }}>
           <Button onClick={() => downloadFiles([toRawFileUrl(file, true)])}>
             {{

vue/src/icon/index.ts

@@ -1,5 +1,7 @@
 export * from '@ant-design/icons-vue'
 import regex from './regex.svg'
+import play from './play.svg'
 export {
-  regex
+  regex,
+  play
 }

vue/src/icon/play.svg

@@ -0,0 +1,4 @@
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 40 40" width="40" height="40">
+  <circle cx="20" cy="20" r="18" fill="#000000"/>
+  <polygon points="28,20 16,26 16,14" fill="#ffffff"/>
+</svg>

vue/src/util/file.ts

@@ -4,14 +4,19 @@ import { uniqBy } from 'lodash-es'
 
 const encode = encodeURIComponent
 export const toRawFileUrl = (file: FileNodeInfo, download = false) =>
-  `${apiBase.value}/file?path=${encode(file.fullpath)}&t=${encode(file.date)}${
+  `${apiBase.value}/file?path=${encode(file.fullpath)}&t=${encode(file.date)}${download ? `&disposition=${encode(file.name)}` : ''
-    download ? `&disposition=${encode(file.name)}` : ''
   }`
 export const toImageThumbnailUrl = (file: FileNodeInfo, size: string = '512x512') =>
   `${apiBase.value}/image-thumbnail?path=${encode(file.fullpath)}&size=${size}&t=${encode(
     file.date
   )}`
 
+export const toStreamVideoUrl = (file: FileNodeInfo) =>
+  `${apiBase.value}/stream_video?path=${encode(file.fullpath)}`
+
+  export const toVideoCoverUrl = (file: FileNodeInfo) =>
+  parent.document.location.origin+ `${apiBase.value}/video_cover?path=${encode(file.fullpath)}&t=${encode(file.date)}`
+
 export type FileTransferData = {
   path: string[]
   loc: string
@@ -30,7 +35,7 @@ export const getFileTransferDataFromDragEvent = (e: DragEvent) => {
 
 export const uniqueFile = (files: FileNodeInfo[]) => uniqBy(files, 'fullpath')
 
-export function isImageFile(filename: string): boolean {
+export function isImageFile (filename: string): boolean {
   if (typeof filename !== 'string') {
     return false
   }
@@ -39,7 +44,7 @@ export function isImageFile(filename: string): boolean {
   return extension !== undefined && exts.includes(`.${extension}`)
 }
 
-export function isVideoFile(filename: string): boolean {
+export function isVideoFile (filename: string): boolean {
   if (typeof filename !== 'string') {
     return false
   }
@@ -50,10 +55,10 @@ export function isVideoFile(filename: string): boolean {
 
 export const isMediaFile = (file: string) => isImageFile(file) || isVideoFile(file)
 
-export function downloadFiles(urls: string[]) {
+export function downloadFiles (urls: string[]) {
-  const link = document.createElement('a');
+  const link = document.createElement('a')
-  link.style.display = 'none';
+  link.style.display = 'none'
-  document.body.appendChild(link);
+  document.body.appendChild(link)
 
   urls.forEach((url) => {
     const urlObject = new URL(url, 'https://github.com/zanllp/sd-webui-infinite-image-browsing')
@@ -62,12 +67,12 @@ export function downloadFiles(urls: string[]) {
     if (disposition) {
       filename = disposition
     }
-    link.href = url;
+    link.href = url
-    link.download = filename;
+    link.download = filename
-    link.click();
+    link.click()
-  });
+  })
 
-  document.body.removeChild(link);
+  document.body.removeChild(link)
 }
 
 export const downloadFileInfoJSON = (files: FileNodeInfo[], name?: string) => {