Refactor route imports and enhance lazy loading for optional integrations to avoid CPU spikes on startup

- Updated route imports in `main.py` to use specific modules for better organization and lazy loading.
- Implemented lazy imports for `acestream` and `telegram` functionalities to avoid unnecessary loading when not enabled.
- Refactored the router initialization in `__init__.py` to support lazy loading of routes, improving startup performance.
- Enhanced the `acestream` and `telegram` routes to conditionally include their respective routers based on configuration settings.
- Improved the handling of transcode components across various routes to optimize resource usage.

Author: mhdzumair

mediaflow_proxy/main.py

@@ -12,23 +12,16 @@
 
 from mediaflow_proxy.configs import settings
 from mediaflow_proxy.middleware import UIAccessControlMiddleware
-from mediaflow_proxy.routes import (
-    proxy_router,
-    extractor_router,
-    speedtest_router,
-    playlist_builder_router,
-    xtream_root_router,
-    acestream_router,
-    telegram_router,
-)
+from mediaflow_proxy.routes.proxy import proxy_router
+from mediaflow_proxy.routes.extractor import extractor_router
+from mediaflow_proxy.routes.speedtest import speedtest_router
+from mediaflow_proxy.routes.playlist_builder import playlist_builder_router
+from mediaflow_proxy.routes.xtream import xtream_root_router
 from mediaflow_proxy.schemas import GenerateUrlRequest, GenerateMultiUrlRequest, MultiUrlRequestItem
 from mediaflow_proxy.utils.crypto_utils import EncryptionHandler, EncryptionMiddleware
 from mediaflow_proxy.utils import redis_utils
 from mediaflow_proxy.utils.http_utils import encode_mediaflow_proxy_url
 from mediaflow_proxy.utils.base64_utils import encode_url_to_base64, decode_base64_url, is_base64_url
-from mediaflow_proxy.utils.acestream import acestream_manager
-from mediaflow_proxy.remuxer.video_transcoder import get_hw_capability, HWAccelType
-from mediaflow_proxy.utils.telegram import telegram_manager
 
 logging.basicConfig(level=settings.log_level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
@@ -61,30 +54,22 @@ async def lifespan(app: FastAPI):
         # use redis-cli KEYS "mfp:*" | xargs redis-cli DEL
         logger.info("Cache clearing note: Redis entries will expire via TTL")
 
-    # Log transcoding capability
-    hw = get_hw_capability()
-    if hw.accel_type != HWAccelType.NONE and settings.transcode_prefer_gpu:
-        logger.info(
-            "Transcode ready: GPU %s (encoder=%s) | PyAV pipeline",
-            hw.accel_type.value,
-            hw.h264_encoder,
-        )
-    else:
-        logger.info(
-            "Transcode ready: CPU (%s) | PyAV pipeline",
-            hw.h264_encoder,
-        )
-
     yield
 
     # Shutdown
     logger.info("Shutting down...")
     # Close acestream sessions
-    await acestream_manager.close()
-    logger.info("Acestream manager closed")
+    if settings.enable_acestream:
+        from mediaflow_proxy.utils.acestream import acestream_manager
+
+        await acestream_manager.close()
+        logger.info("Acestream manager closed")
     # Close telegram session
-    await telegram_manager.close()
-    logger.info("Telegram manager closed")
+    if settings.enable_telegram:
+        from mediaflow_proxy.utils.telegram import telegram_manager
+
+        await telegram_manager.close()
+        logger.info("Telegram manager closed")
     # Close Redis connections
     await redis_utils.close_redis()
     logger.info("Redis connections closed")
@@ -318,8 +303,14 @@ async def check_base64_url(url: str):
 
 
 app.include_router(proxy_router, prefix="/proxy", tags=["proxy"], dependencies=[Depends(verify_api_key)])
-app.include_router(acestream_router, prefix="/proxy", tags=["acestream"], dependencies=[Depends(verify_api_key)])
-app.include_router(telegram_router, prefix="/proxy", tags=["telegram"], dependencies=[Depends(verify_api_key)])
+if settings.enable_acestream:
+    from mediaflow_proxy.routes.acestream import acestream_router
+
+    app.include_router(acestream_router, prefix="/proxy", tags=["acestream"], dependencies=[Depends(verify_api_key)])
+if settings.enable_telegram:
+    from mediaflow_proxy.routes.telegram import telegram_router
+
+    app.include_router(telegram_router, prefix="/proxy", tags=["telegram"], dependencies=[Depends(verify_api_key)])
 app.include_router(extractor_router, prefix="/extractor", tags=["extractors"], dependencies=[Depends(verify_api_key)])
 app.include_router(speedtest_router, prefix="/speedtest", tags=["speedtest"], dependencies=[Depends(verify_api_key)])
 app.include_router(playlist_builder_router, prefix="/playlist", tags=["playlist"])

mediaflow_proxy/remuxer/media_source.py

@@ -13,7 +13,6 @@
 from urllib.parse import urlparse, unquote
 
 from mediaflow_proxy.utils.http_client import create_aiohttp_session
-from mediaflow_proxy.utils.telegram import telegram_manager
 
 logger = logging.getLogger(__name__)
 
@@ -142,6 +141,9 @@ def filename_hint(self) -> str:
         return self._filename_hint
 
     async def stream(self, offset: int = 0, limit: int | None = None) -> AsyncIterator[bytes]:
+        # Lazy import to avoid loading Telegram dependencies for non-Telegram routes.
+        from mediaflow_proxy.utils.telegram import telegram_manager
+
         effective_limit = limit or self._file_size
         if self._use_single_client:
             async for chunk in telegram_manager.stream_media_single(

mediaflow_proxy/routes/__init__.py

@@ -1,11 +1,3 @@
-from .proxy import proxy_router
-from .extractor import extractor_router
-from .speedtest import speedtest_router
-from .playlist_builder import playlist_builder_router
-from .xtream import xtream_root_router
-from .acestream import acestream_router
-from .telegram import telegram_router
-
 __all__ = [
     "proxy_router",
     "extractor_router",
@@ -15,3 +7,37 @@
     "acestream_router",
     "telegram_router",
 ]
+
+
+def __getattr__(name: str):
+    # Lazy import routers so importing a single route module does not
+    # pull in optional integrations (telegram/acestream/transcode) at startup.
+    if name == "proxy_router":
+        from .proxy import proxy_router
+
+        return proxy_router
+    if name == "extractor_router":
+        from .extractor import extractor_router
+
+        return extractor_router
+    if name == "speedtest_router":
+        from .speedtest import speedtest_router
+
+        return speedtest_router
+    if name == "playlist_builder_router":
+        from .playlist_builder import playlist_builder_router
+
+        return playlist_builder_router
+    if name == "xtream_root_router":
+        from .xtream import xtream_root_router
+
+        return xtream_root_router
+    if name == "acestream_router":
+        from .acestream import acestream_router
+
+        return acestream_router
+    if name == "telegram_router":
+        from .telegram import telegram_router
+
+        return telegram_router
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

mediaflow_proxy/routes/acestream.py

@@ -9,16 +9,15 @@
 
 import asyncio
 import logging
-from typing import Annotated
+from functools import lru_cache
+from typing import Annotated, TYPE_CHECKING
 from urllib.parse import urlencode, urljoin, urlparse
 
 import aiohttp
 from fastapi import APIRouter, Query, Request, HTTPException, Response, Depends
 from starlette.background import BackgroundTask
 
 from mediaflow_proxy.configs import settings
-from mediaflow_proxy.remuxer.transcode_pipeline import stream_transcode_universal
-from mediaflow_proxy.utils.acestream import acestream_manager, AcestreamSession
 from mediaflow_proxy.utils.http_client import create_aiohttp_session
 from mediaflow_proxy.utils.http_utils import (
     get_original_scheme,
@@ -34,6 +33,22 @@
 logger = logging.getLogger(__name__)
 acestream_router = APIRouter()
 
+if TYPE_CHECKING:
+    from mediaflow_proxy.utils.acestream import AcestreamSession
+
+
+def _get_acestream_manager():
+    from mediaflow_proxy.utils.acestream import acestream_manager
+
+    return acestream_manager
+
+
+@lru_cache(maxsize=1)
+def _load_transcode_pipeline():
+    from mediaflow_proxy.remuxer.transcode_pipeline import stream_transcode_universal
+
+    return stream_transcode_universal
+
 
 class AcestreamM3U8Processor(M3U8Processor):
     """
@@ -46,7 +61,7 @@ class AcestreamM3U8Processor(M3U8Processor):
     def __init__(
         self,
         request: Request,
-        session: AcestreamSession,
+        session: "AcestreamSession",
         key_url: str = None,
         force_playlist_proxy: bool = True,
         key_only_proxy: bool = False,
@@ -140,6 +155,7 @@ async def acestream_hls_manifest(
     """
     if not settings.enable_acestream:
         raise HTTPException(status_code=503, detail="Acestream support is disabled")
+    acestream_manager = _get_acestream_manager()
 
     if not infohash and not id:
         raise HTTPException(status_code=400, detail="Either 'infohash' or 'id' parameter is required")
@@ -278,6 +294,7 @@ async def acestream_segment_proxy(
     """
     if not settings.enable_acestream:
         raise HTTPException(status_code=503, detail="Acestream support is disabled")
+    acestream_manager = _get_acestream_manager()
 
     # Use id or infohash for session lookup
     session_key = id or infohash
@@ -368,6 +385,7 @@ async def acestream_ts_stream(
     """
     if not settings.enable_acestream:
         raise HTTPException(status_code=503, detail="Acestream support is disabled")
+    acestream_manager = _get_acestream_manager()
 
     if not infohash and not id:
         raise HTTPException(status_code=400, detail="Either 'infohash' or 'id' parameter is required")
@@ -438,6 +456,7 @@ async def _acestream_ts_source():
             # Use our custom PyAV pipeline with forced video re-encoding
             # (live MPEG-TS sources often have corrupt H.264 bitstreams
             # that browsers reject; re-encoding produces a clean stream).
+            stream_transcode_universal = _load_transcode_pipeline()
             content = stream_transcode_universal(
                 _acestream_ts_source(),
                 force_video_reencode=True,
@@ -509,6 +528,7 @@ async def acestream_status(
     """
     if not settings.enable_acestream:
         raise HTTPException(status_code=503, detail="Acestream support is disabled")
+    acestream_manager = _get_acestream_manager()
 
     if infohash:
         session = acestream_manager.get_session(infohash)

mediaflow_proxy/routes/proxy.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 import re
+from functools import lru_cache
 from typing import Annotated
 from urllib.parse import quote, unquote
 
@@ -41,19 +42,31 @@
 from mediaflow_proxy.utils.http_client import create_aiohttp_session
 from mediaflow_proxy.utils.m3u8_processor import M3U8Processor
 from mediaflow_proxy.utils.stream_transformers import apply_transformer_to_bytes
-from mediaflow_proxy.remuxer.media_source import HTTPMediaSource
-from mediaflow_proxy.remuxer.transcode_handler import (
-    handle_transcode,
-    handle_transcode_hls_init,
-    handle_transcode_hls_playlist,
-    handle_transcode_hls_segment,
-)
 
 
 logger = logging.getLogger(__name__)
 proxy_router = APIRouter()
 
 
+@lru_cache(maxsize=1)
+def _load_transcode_components():
+    from mediaflow_proxy.remuxer.media_source import HTTPMediaSource
+    from mediaflow_proxy.remuxer.transcode_handler import (
+        handle_transcode,
+        handle_transcode_hls_init,
+        handle_transcode_hls_playlist,
+        handle_transcode_hls_segment,
+    )
+
+    return (
+        HTTPMediaSource,
+        handle_transcode,
+        handle_transcode_hls_init,
+        handle_transcode_hls_playlist,
+        handle_transcode_hls_segment,
+    )
+
+
 def sanitize_url(url: str) -> str:
     """
     Sanitize URL to fix common encoding issues and handle base64 encoded URLs.
@@ -499,6 +512,7 @@ async def transcode_hls_playlist(
     """
     if not settings.enable_transcode:
         raise HTTPException(status_code=503, detail="Transcoding support is disabled")
+    HTTPMediaSource, _, _, handle_transcode_hls_playlist, _ = _load_transcode_components()
     destination = sanitize_url(destination)
     source = HTTPMediaSource(url=destination, headers=dict(proxy_headers.request))
     await source.resolve_file_size()
@@ -540,6 +554,7 @@ async def transcode_hls_init(
     """
     if not settings.enable_transcode:
         raise HTTPException(status_code=503, detail="Transcoding support is disabled")
+    HTTPMediaSource, _, handle_transcode_hls_init, _, _ = _load_transcode_components()
     destination = sanitize_url(destination)
     source = HTTPMediaSource(url=destination, headers=dict(proxy_headers.request))
     await source.resolve_file_size()
@@ -572,6 +587,7 @@ async def transcode_hls_segment(
     """
     if not settings.enable_transcode:
         raise HTTPException(status_code=503, detail="Transcoding support is disabled")
+    HTTPMediaSource, _, _, _, handle_transcode_hls_segment = _load_transcode_components()
     destination = sanitize_url(destination)
     source = HTTPMediaSource(url=destination, headers=dict(proxy_headers.request))
     await source.resolve_file_size()
@@ -693,6 +709,7 @@ async def proxy_stream_endpoint(
     if transcode:
         if not settings.enable_transcode:
             raise HTTPException(status_code=503, detail="Transcoding support is disabled")
+        HTTPMediaSource, handle_transcode, _, _, _ = _load_transcode_components()
         transcode_headers = dict(proxy_headers.request)
         transcode_headers.pop("range", None)
         transcode_headers.pop("if-range", None)