feat: enhance thumbnail handling and add native thumbnail support in VirtualFS

Author: DrizzleTime

domain/adapters/providers/telegram.py

@@ -5,7 +5,7 @@
 import os
 import struct
 from models import StorageAdapter
-from telethon import TelegramClient
+from telethon import TelegramClient, utils
 from telethon.crypto import AuthKey
 from telethon.sessions import StringSession
 from telethon.tl import types
@@ -132,29 +132,42 @@ def _pick_photo_thumb(thumbs: list | None):
             return None
 
         cached = []
-        others = []
+        downloadable = []
         for t in thumbs:
             if isinstance(t, (types.PhotoCachedSize, types.PhotoStrippedSize)):
                 cached.append(t)
             elif isinstance(t, (types.PhotoSize, types.PhotoSizeProgressive)):
                 if not isinstance(t, types.PhotoSizeEmpty):
-                    others.append(t)
+                    downloadable.append(t)
 
-        if cached:
-            cached.sort(key=lambda x: len(getattr(x, "bytes", b"") or b""))
-            return cached[-1]
-
-        if others:
+        if downloadable:
             def _sz(x):
                 if isinstance(x, types.PhotoSizeProgressive):
                     return max(x.sizes or [0])
                 return int(getattr(x, "size", 0) or 0)
 
-            others.sort(key=_sz)
-            return others[-1]
+            downloadable.sort(key=_sz)
+            return downloadable[-1]
+
+        if cached:
+            cached.sort(key=lambda x: len(getattr(x, "bytes", b"") or b""))
+            return cached[-1]
 
         return None
 
+    @staticmethod
+    def _get_message_thumbs(message) -> list:
+        doc = message.document or message.video
+        if doc and getattr(doc, "thumbs", None):
+            return list(doc.thumbs or [])
+        if message.photo and getattr(message.photo, "sizes", None):
+            return list(message.photo.sizes or [])
+        return []
+
+    @classmethod
+    def _message_has_thumbnail(cls, message) -> bool:
+        return cls._pick_photo_thumb(cls._get_message_thumbs(message)) is not None
+
     def _build_session(self) -> StringSession:
         s = (self.session_string or "").strip()
         if not s:
@@ -229,6 +242,7 @@ async def list_dir(self, root: str, rel: str, page_num: int = 1, page_size: int
                     "size": size,
                     "mtime": int(message.date.timestamp()),
                     "type": "file",
+                    "has_thumbnail": self._message_has_thumbnail(message),
                 })
         finally:
             if client.is_connected():
@@ -386,17 +400,16 @@ async def get_thumbnail(self, root: str, rel: str, size: str = "medium"):
             if not message:
                 return None
 
-            doc = message.document or message.video
-            thumbs = None
-            if doc and getattr(doc, "thumbs", None):
-                thumbs = list(doc.thumbs or [])
-            elif message.photo and getattr(message.photo, "sizes", None):
-                thumbs = list(message.photo.sizes or [])
-
-            thumb = self._pick_photo_thumb(thumbs)
+            thumb = self._pick_photo_thumb(self._get_message_thumbs(message))
             if not thumb:
                 return None
 
+            embedded = getattr(thumb, "bytes", None)
+            if embedded and isinstance(thumb, types.PhotoCachedSize):
+                return bytes(embedded)
+            if embedded and isinstance(thumb, types.PhotoStrippedSize):
+                return utils.stripped_photo_to_jpg(bytes(embedded))
+
             result = await client.download_media(message, bytes, thumb=thumb)
             if isinstance(result, (bytes, bytearray)):
                 return bytes(result)
@@ -569,6 +582,7 @@ async def stat_file(self, root: str, rel: str):
                 "size": size,
                 "mtime": int(message.date.timestamp()),
                 "type": "file",
+                "has_thumbnail": self._message_has_thumbnail(message),
             }
         finally:
             if client.is_connected():

domain/virtual_fs/listing.py

@@ -89,6 +89,9 @@ def build_sort_key(item: Dict) -> Tuple:
 
         def annotate_entry(entry: Dict) -> None:
             if not entry.get("is_dir"):
+                if entry.get("has_thumbnail") is not None:
+                    entry["has_thumbnail"] = bool(entry.get("has_thumbnail"))
+                    return
                 name = entry.get("name", "")
                 entry["has_thumbnail"] = bool(is_image_filename(name) or is_video_filename(name))
             else:
@@ -273,7 +276,10 @@ async def stat_file(cls, path: str, verbose: bool = False):
                 is_dir = False
             rel_name = rel.rstrip("/").split("/")[-1] if rel else path.rstrip("/").split("/")[-1]
             name_hint = str(info.get("name") or rel_name or "")
-            info["has_thumbnail"] = bool(not is_dir and (is_image_filename(name_hint) or is_video_filename(name_hint)))
+            if not is_dir and info.get("has_thumbnail") is not None:
+                info["has_thumbnail"] = bool(info.get("has_thumbnail"))
+            else:
+                info["has_thumbnail"] = bool(not is_dir and (is_image_filename(name_hint) or is_video_filename(name_hint)))
             if verbose and not is_dir:
                 vector_index = await cls._gather_vector_index(path)
                 if vector_index is not None:

domain/virtual_fs/routes.py

@@ -89,8 +89,17 @@ async def get_thumbnail(cls, full_path: str, w: int, h: int, fit: str) -> Respon
         adapter, mount, root, rel = await cls.resolve_adapter_and_rel(full_path)
         if not rel or rel.endswith("/"):
             raise HTTPException(400, detail="Not a file")
-        if not (is_image_filename(rel) or is_video_filename(rel)):
-            raise HTTPException(404, detail="Not an image or video")
+        has_native_thumb = False
+        if callable(getattr(adapter, "get_thumbnail", None)):
+            stat_file = getattr(adapter, "stat_file", None)
+            if callable(stat_file):
+                try:
+                    stat = await stat_file(root, rel)
+                    has_native_thumb = bool(isinstance(stat, dict) and stat.get("has_thumbnail"))
+                except Exception:
+                    has_native_thumb = False
+        if not (is_image_filename(rel) or is_video_filename(rel) or has_native_thumb):
+            raise HTTPException(404, detail="Not an image, video, or native thumbnail file")
         data, mime, key = await get_or_create_thumb(adapter, mount.id, root, rel, w, h, fit)  # type: ignore
         headers = {
             "Cache-Control": "public, max-age=3600",

domain/virtual_fs/thumbnail.py

@@ -23,6 +23,7 @@
 VIDEO_THUMB_SEEK_SECONDS = (15, 10, 5, 3, 1, 0)
 VIDEO_BLACK_FRAME_MEAN_THRESHOLD = 12.0
 CACHE_ROOT = Path('data/.thumb_cache')
+THUMB_CACHE_VERSION = "v2"
 
 
 def is_image_filename(name: str) -> bool:
@@ -47,7 +48,7 @@ def is_video_filename(name: str) -> bool:
 
 
 def _cache_key(adapter_id: int, rel: str, size: int, mtime: int, w: int, h: int, fit: str) -> str:
-    raw = f"{adapter_id}|{rel}|{size}|{mtime}|{w}x{h}|{fit}".encode()
+    raw = f"{THUMB_CACHE_VERSION}|{adapter_id}|{rel}|{size}|{mtime}|{w}x{h}|{fit}".encode()
     return hashlib.sha1(raw).hexdigest()
 
 
@@ -385,8 +386,11 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w:
     stat = await adapter.stat_file(root, rel)
     size = int(stat.get('size') or 0)
     is_video = is_video_filename(rel)
-    if not is_video and size > MAX_IMAGE_SOURCE_SIZE:
-        raise HTTPException(400, detail="Image too large for thumbnail")
+    is_image = is_image_filename(rel)
+    get_thumb_impl = getattr(adapter, "get_thumbnail", None)
+    should_try_native_thumb = callable(get_thumb_impl) and (
+        is_image or is_video or bool(stat.get("has_thumbnail"))
+    )
 
     key = _cache_key(adapter_id, rel, size, int(
         stat.get('mtime', 0)), w, h, fit)
@@ -397,19 +401,15 @@ async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w:
     _ensure_cache_dir(path)
     thumb_bytes, mime = None, None
 
-    get_thumb_impl = getattr(adapter, "get_thumbnail", None)
-    if callable(get_thumb_impl):
+    if should_try_native_thumb:
         size_str = "large" if w > 400 else "medium" if w > 100 else "small"
         native_thumb_bytes = await get_thumb_impl(root, rel, size_str)
 
         if native_thumb_bytes:
             try:
                 from PIL import Image
                 im = Image.open(io.BytesIO(native_thumb_bytes))
-                buf = io.BytesIO()
-                im.save(buf, 'WEBP', quality=85)
-                thumb_bytes = buf.getvalue()
-                mime = 'image/webp'
+                thumb_bytes, mime = _image_to_webp(im, w, h, fit)
             except Exception as e:
                 print(
                     f"Failed to convert native thumbnail to WebP: {e}, falling back.")
@@ -493,7 +493,9 @@ async def _read_tail(limit: int) -> Tuple[bytes, int]:
                         thumb_bytes, mime = retry_thumb, retry_mime
                 except Exception:
                     pass
-        else:
+        elif is_image:
+            if size > MAX_IMAGE_SOURCE_SIZE:
+                raise HTTPException(400, detail="Image too large for thumbnail")
             read_data = await adapter.read_file(root, rel)
             try:
                 thumb_bytes, mime = generate_thumb(
@@ -502,6 +504,8 @@ async def _read_tail(limit: int) -> Tuple[bytes, int]:
                 print(e)
                 raise HTTPException(
                     500, detail=f"Thumbnail generation failed: {e}")
+        else:
+            raise HTTPException(500, detail="Native thumbnail unavailable")
 
     if thumb_bytes:
         path.write_bytes(thumb_bytes)