model-compose 0.4.85__tar.gz → 0.4.87__tar.gz

{model_compose-0.4.85 → model_compose-0.4.87}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-compose
-Version: 0.4.85
+Version: 0.4.87
 Summary: model-compose: Declarative AI Workflow Orchestrator
 Author-email: Hanyeol Cho <hanyeol.cho@gmail.com>
 License-Expression: MIT

{model_compose-0.4.85 → model_compose-0.4.87}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "model-compose"
-version = "0.4.85"
+version = "0.4.87"
 description = "model-compose: Declarative AI Workflow Orchestrator"
 authors = [
     { name = "Hanyeol Cho", email = "hanyeol.cho@gmail.com" }

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/base.py

@@ -9,10 +9,13 @@ from mindor.dsl.schema.gateway import GatewayConfig
 from mindor.dsl.schema.workflow import WorkflowConfig
 from mindor.dsl.schema.runtime import RuntimeType
 from mindor.core.foundation import AsyncService
+from mindor.core.foundation.streaming.iterators import StreamIterator
 from mindor.core.utils.work_queue import WorkQueue
 from mindor.core.utils.active_counter import ActiveCounter
 from mindor.core.logger import logging
+from collections.abc import AsyncIterator
 from .context import ComponentActionContext
+from .streaming import ComponentOutputStreamIterator, StreamTerminatedEvent
 import asyncio
 
 class ActionResolver:
@@ -136,6 +139,17 @@ class ComponentService(AsyncService):
             await context.event_notifier.notify("failed", error=str(e))
             raise
 
+        if isinstance(output, (StreamIterator, AsyncIterator)):
+            async def _on_terminated(event: StreamTerminatedEvent, error: Optional[str]) -> None:
+                if event == "completed":
+                    await context.event_notifier.notify("completed", output=None)
+                elif event == "cancelled":
+                    await context.event_notifier.notify("cancelled")
+                else:
+                    await context.event_notifier.notify("failed", error=error)
+
+            return ComponentOutputStreamIterator(output, _on_terminated)
+
         await context.event_notifier.notify("completed", output=output)
 
         return output

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/context.py

@@ -97,11 +97,11 @@ class ComponentActionContext:
     async def render_audio(self, value: Any) -> Union[MediaSource, List[MediaSource]]:
         return await AudioValueRenderer().render(await self.render_variable(value))
 
-    async def render_audio_buffer(self, value: Any) -> Union[AudioBuffer, List[AudioBuffer], AsyncIterator[AudioBuffer]]:
-        return await AudioBufferValueRenderer().render(await self.render_variable(value))
+    async def render_audio_buffer(self, value: Any, sample_rate: Optional[int] = None, channel: Optional[int] = None) -> Union[AudioBuffer, List[AudioBuffer], AsyncIterator[AudioBuffer]]:
+        return await AudioBufferValueRenderer(sample_rate, channel).render(await self.render_variable(value))
 
-    async def render_audio_buffer_array(self, value: Any) -> Union[AudioBufferArrayValue, List[AudioBufferArrayValue], AsyncIterator[AudioBufferArrayValue]]:
-        return await AudioBufferValueRenderer().render_array(await self.render_variable(value))
+    async def render_audio_buffer_array(self, value: Any, sample_rate: Optional[int] = None, channel: Optional[int] = None) -> Union[AudioBufferArrayValue, List[AudioBufferArrayValue], AsyncIterator[AudioBufferArrayValue]]:
+        return await AudioBufferValueRenderer(sample_rate, channel).render_array(await self.render_variable(value))
 
     async def render_video(self, value: Any) -> Union[MediaSource, List[MediaSource]]:
         return await VideoValueRenderer().render(await self.render_variable(value))

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/runtime/base/ipc_proxy.py

@@ -119,7 +119,7 @@ class IpcRuntimeProxy(ABC):
             payload = message.payload or {}
 
             if message.type == IpcMessageType.ERROR:
-                raise RuntimeError(f"Worker '{self.worker_id}' failed to start: {payload.get('error', 'Unknown error')}")
+                raise RuntimeError(f"Worker '{self.worker_id}' failed to start: {self._format_error(payload)}")
 
             if message.type == IpcMessageType.STATUS:
                 if payload.get("status") == "ready":
@@ -164,7 +164,7 @@ class IpcRuntimeProxy(ABC):
                     continue
 
                 if message.type == IpcMessageType.ERROR:
-                    future.set_exception(Exception(payload.get("error", "Unknown error")))
+                    future.set_exception(Exception(self._format_error(payload)))
                     continue
         except asyncio.CancelledError:
             pass
@@ -331,6 +331,16 @@ class IpcRuntimeProxy(ABC):
             payload={ "stream_id": stream_id, "error": error },
         ).serialize())
 
+    @staticmethod
+    def _format_error(payload: Dict[str, Any]) -> str:
+        message = payload.get("error") or "Unknown error"
+        traceback = payload.get("traceback")
+
+        if traceback:
+            return f"{message}\n\nWorker traceback:\n{traceback}"
+
+        return message
+
     @abstractmethod
     async def _start(self) -> None:
         pass

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/runtime/base/ipc_worker.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
 from abc import ABC, abstractmethod
 from mindor.core.foundation.variable.codec import StreamKind, VariableCodec
 from .ipc_message import IpcMessage, IpcMessageType
 from .ipc_stream import IpcInboundStream, IpcOutboundStream, IpcStreamReader
-import asyncio
+import asyncio, traceback
 
 class IpcRuntimeWorker(ABC):
     """
@@ -75,9 +75,9 @@ class IpcRuntimeWorker(ABC):
                 try:
                     await self._dispatch_message(message)
                 except Exception as e:
-                    await self._send_error(message.request_id, str(e))
+                    await self._send_error(message.request_id, e)
         except Exception as e:
-            await self._notify_error(str(e))
+            await self._notify_error(e)
         finally:
             await self._abort_all_streams()
             for task in list(self._run_tasks.values()):
@@ -95,7 +95,7 @@ class IpcRuntimeWorker(ABC):
             await self._send_error(message.request_id, "Request was cancelled")
             raise
         except Exception as e:
-            await self._send_error(message.request_id, str(e))
+            await self._send_error(message.request_id, e)
 
     async def _dispatch_message(self, message: IpcMessage) -> Dict[str, Any]:
         if message.type == IpcMessageType.RUN:
@@ -273,11 +273,11 @@ class IpcRuntimeWorker(ABC):
             payload=payload,
         ).serialize())
 
-    async def _send_error(self, request_id: str, error: str) -> None:
+    async def _send_error(self, request_id: str, error: Union[BaseException, str]) -> None:
         await self._send_message(IpcMessage(
             type=IpcMessageType.ERROR,
             request_id=request_id,
-            payload={ "error": error },
+            payload=self._build_error_payload(error),
         ).serialize())
 
     async def _notify_status(self, status: str) -> None:
@@ -286,15 +286,23 @@ class IpcRuntimeWorker(ABC):
             payload={ "status": status },
         ).serialize())
 
-    async def _notify_error(self, error: str) -> None:
+    async def _notify_error(self, error: Union[BaseException, str]) -> None:
         try:
             await self._send_message(IpcMessage(
                 type=IpcMessageType.ERROR,
-                payload={ "error": error },
+                payload=self._build_error_payload(error),
             ).serialize())
         except Exception:
             pass
 
+    def _build_error_payload(self, error: Union[BaseException, str]) -> Dict[str, Any]:
+        if isinstance(error, BaseException):
+            return {
+                "error":     str(error) or type(error).__name__,
+                "traceback": "".join(traceback.format_exception(type(error), error, error.__traceback__)),
+            }
+        return { "error": error }
+
     @abstractmethod
     async def _start(self) -> None:
         pass

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/services/model/tasks/music_generation/custom/ace_step.py

@@ -99,10 +99,10 @@ class AceStepMusicGenerationTaskService(MusicGenerationTaskService):
 
     def get_setup_requirements(self) -> Optional[List[str]]:
         return [
-            "torch==2.8.0+cu128@https://download.pytorch.org/whl/cu128",
-            "torchaudio==2.8.0+cu128@https://download.pytorch.org/whl/cu128",
-            "torchvision==0.23.0+cu128@https://download.pytorch.org/whl/cu128",
-            "flash-attn==2.8.3.post1@https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3.post1+cu12torch2.8cxx11abiFALSE-cp310-cp310-linux_x86_64.whl",
+            "torch==2.10.0+cu128@https://download.pytorch.org/whl/cu128",
+            "torchaudio==2.10.0+cu128@https://download.pytorch.org/whl/cu128",
+            "torchvision==0.25.0+cu128@https://download.pytorch.org/whl/cu128",
+            "flash-attn==2.8.1@https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.10cxx11abiTRUE-cp312-cp312-linux_x86_64.whl",
             "nano-vllm@git+https://github.com/GeeeekExplorer/nano-vllm.git",
             "ace-step@git+https://github.com/ace-step/ACE-Step-1.5.git",
             "numpy",

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/services/model/tasks/text_to_speech/custom/cosyvoice.py

@@ -3,15 +3,15 @@ from typing import TYPE_CHECKING
 
 from typing import Union, Dict, Optional, List, Tuple, Any
 from abc import abstractmethod
-from mindor.dsl.schema.component import ModelComponentConfig, HuggingfaceModelConfig
+from mindor.dsl.schema.component import ModelComponentConfig
 from mindor.dsl.schema.action import ModelActionConfig, TextToSpeechActionMethod
 from mindor.dsl.schema.action import CommonTextToSpeechModelActionConfig
 from mindor.dsl.schema.action import CosyvoiceTextToSpeechModelGenerateActionConfig
 from mindor.dsl.schema.action import CosyvoiceTextToSpeechModelCloneActionConfig
 from mindor.dsl.schema.action import CosyvoiceTextToSpeechModelDesignActionConfig
 from mindor.core.foundation.cancellation import CancellationToken
-from mindor.core.foundation.streaming.audio import PcmStreamResource
-from mindor.core.foundation.streaming.resources import StreamResource
+from mindor.core.foundation.streaming.audio import PcmStreamResource, WavStreamResource
+from mindor.core.foundation.streaming.resources import StreamResource, save_stream_to_temporary_file
 from mindor.core.utils.audio import encode_waveform_to_pcm
 from ......base import ComponentActionContext
 from ..common import TextToSpeechTaskService, TextToSpeechTaskAction
@@ -67,6 +67,20 @@ class CosyvoiceTextToSpeechTaskAction(TextToSpeechTaskAction):
     def _invoke(self, text: str, params: Dict[str, Any]) -> Any:
         pass
 
+    async def _render_reference_audio(self, context: ComponentActionContext, value: Any) -> str:
+        """Decode the reference audio and re-emit it as a 16 kHz mono WAV
+        tempfile path. CosyVoice's frontend re-reads the prompt with its own
+        load_wav (which in newer torchaudio routes through torchcodec and
+        breaks on tensor inputs), so we hand it a plain on-disk WAV path."""
+        buffer = await context.render_audio_buffer(value, sample_rate=16000)
+        pcm_bytes, channels = encode_waveform_to_pcm(buffer.waveform)
+        stream = WavStreamResource(PcmStreamResource(pcm_bytes, {
+            "sample_rate": str(buffer.sample_rate),
+            "channels":    str(channels),
+            "bit_depth":   "16",
+        }))
+        return await save_stream_to_temporary_file(stream, "wav")
+
 class CosyvoiceTextToSpeechGenerateTaskAction(CosyvoiceTextToSpeechTaskAction):
     config: CosyvoiceTextToSpeechModelGenerateActionConfig
 
@@ -116,34 +130,29 @@ class CosyvoiceTextToSpeechCloneTaskAction(CosyvoiceTextToSpeechTaskAction):
     async def _resolve_params(self, context: ComponentActionContext) -> Dict[str, Any]:
         params = await super()._resolve_params(context)
 
-        reference_audio = await context.render_file(self.config.reference_audio)
-        reference_text  = await context.render_variable(self.config.reference_text)
-        speed           = await context.render_variable(self.config.speed)
-        text_frontend   = await context.render_variable(self.config.text_frontend)
+        prompt_wav     = await self._render_prompt_wav_file(context, self.config.reference_audio)
+        prompt_text    = await context.render_variable(self.config.reference_text)
+        speed          = await context.render_variable(self.config.speed)
+        text_frontend  = await context.render_variable(self.config.text_frontend)
 
         params.update({
-            "reference_audio": reference_audio,
-            "reference_text":  reference_text,
-            "speed":           float(speed) if speed is not None else 1.0,
-            "text_frontend":   bool(text_frontend) if text_frontend is not None else True,
+            "prompt_wav":    prompt_wav,
+            "prompt_text":   prompt_text,
+            "speed":         float(speed) if speed is not None else 1.0,
+            "text_frontend": bool(text_frontend) if text_frontend is not None else True,
         })
 
         return params
 
     def _invoke(self, text: str, params: Dict[str, Any]) -> Any:
-        from cosyvoice.utils.file_utils import load_wav
-
-        # CosyVoice's frontend expects a 16kHz mono tensor for the prompt wav.
-        prompt_wav = load_wav(str(params["reference_audio"]), 16000)
-
         # With a transcript, use zero_shot (better speaker similarity). Without
         # one, fall back to cross_lingual — which accepts any language and does
         # not require reference text.
-        if params["reference_text"]:
+        if params["prompt_text"]:
             return self.model.inference_zero_shot(
                 tts_text=text,
-                prompt_text=params["reference_text"],
-                prompt_wav=prompt_wav,
+                prompt_text=params["prompt_text"],
+                prompt_wav=params["prompt_wav"],
                 stream=False,
                 speed=params["speed"],
                 text_frontend=params["text_frontend"],
@@ -151,7 +160,7 @@ class CosyvoiceTextToSpeechCloneTaskAction(CosyvoiceTextToSpeechTaskAction):
 
         return self.model.inference_cross_lingual(
             tts_text=text,
-            prompt_wav=prompt_wav,
+            prompt_wav=params["prompt_wav"],
             stream=False,
             speed=params["speed"],
             text_frontend=params["text_frontend"],
@@ -164,7 +173,7 @@ class CosyvoiceTextToSpeechDesignTaskAction(CosyvoiceTextToSpeechTaskAction):
         params = await super()._resolve_params(context)
 
         instructions    = await context.render_variable(self.config.instructions)
-        reference_audio = await context.render_file(self.config.reference_audio)
+        reference_audio = await self._render_reference_audio(context, self.config.reference_audio)
         speed           = await context.render_variable(self.config.speed)
         text_frontend   = await context.render_variable(self.config.text_frontend)
 
@@ -188,14 +197,10 @@ class CosyvoiceTextToSpeechDesignTaskAction(CosyvoiceTextToSpeechTaskAction):
                 "Use family=cosyvoice with a CosyVoice2/3 checkpoint (e.g. FunAudioLLM/CosyVoice2-0.5B)."
             )
 
-        from cosyvoice.utils.file_utils import load_wav
-
-        prompt_wav = load_wav(str(params["reference_audio"]), 16000)
-
         return self.model.inference_instruct2(
             tts_text=text,
             instruct_text=params["instructions"],
-            prompt_wav=prompt_wav,
+            prompt_wav=params["reference_audio"],
             stream=False,
             speed=params["speed"],
             text_frontend=params["text_frontend"],
@@ -210,10 +215,45 @@ class CosyvoiceTextToSpeechTaskService(TextToSpeechTaskService):
         self.device: Optional[torch.device] = None
 
     def get_setup_requirements(self) -> Optional[List[str]]:
-        # cosyvoice itself is installed from source (see the repo README). We
-        # still declare its runtime deps so the environment matches what the
+        # cosyvoice has no pip distribution — clone it from source and add both the
+        # repo root and its Matcha-TTS submodule to sys.path (or a venv .pth file).
+        # We still declare the runtime deps so the environment matches what the
         # inference paths expect.
-        return [ "cosyvoice", "torch", "torchaudio", "modelscope", "hyperpyyaml", "numpy", "soundfile" ]
+        return [
+            "torch==2.10.0+cu128@https://download.pytorch.org/whl/cu128",
+            "torchaudio==2.10.0+cu128@https://download.pytorch.org/whl/cu128",
+            # torchaudio >=2.9 delegates audio decoding to torchcodec.
+            # PyPI default targets CUDA 13; pull the cu128 variant from
+            # PyTorch's index to match torch above.
+            "torchcodec@https://download.pytorch.org/whl/cu128",
+            "accelerate",
+            "diffusers",
+            "gdown",
+            "transformers",
+            "librosa",
+            "matplotlib",
+            "numpy",
+            "soundfile",
+            "conformer",
+            "hydra-core",
+            "hyperpyyaml",
+            "lightning",
+            "modelscope",
+            "omegaconf",
+            # onnxruntime-gpu >=1.20 links CUDA 13. torch above is cu128, so
+            # pin to the last CUDA 12 compatible release to keep everything
+             # on the same CUDA major version.
+            "onnxruntime-gpu<1.20",
+            "openai-whisper",
+            "pyarrow",
+            "pyworld",
+            # setuptools >=81 no longer ships pkg_resources by default; CosyVoice
+            # (via its transitive deps) still does `import pkg_resources`.
+            "setuptools<81",
+            "wget",
+            "inflect",
+            "wetext",
+        ]
 
     async def _load_model(self) -> None:
         self.model, self.sample_rate, self.device = self._load_pretrained_model()
@@ -224,14 +264,22 @@ class CosyvoiceTextToSpeechTaskService(TextToSpeechTaskService):
 
     def _load_pretrained_model(self) -> Tuple[Any, int, Any]:
         # CosyVoice ships an AutoModel factory that picks CosyVoice / CosyVoice2 /
-        # CosyVoice3 by looking for cosyvoice{,2,3}.yaml inside model_dir, and
-        # calls snapshot_download() itself if model_dir isn't a local path.
-        from cosyvoice.cli.cosyvoice import AutoModel
-
-        if isinstance(self.config.model, HuggingfaceModelConfig):
-            model_dir = self.config.model.repository
-        else:
-            model_dir = self._get_model_path()
+        # CosyVoice3 by looking for cosyvoice{,2,3}.yaml inside model_dir. Its
+        # internal snapshot_download only supports ModelScope, so we always
+        # resolve to a local dir via mindor's HF-aware _get_model_path().
+        try:
+            from cosyvoice.cli.cosyvoice import AutoModel
+        except ImportError as e:
+            raise RuntimeError(
+                f"Failed to import cosyvoice.cli.cosyvoice: {e}\n"
+                "If the cosyvoice package itself is missing, clone the repo and add it to the venv:\n"
+                "  git clone --recursive https://github.com/FunAudioLLM/CosyVoice.git\n"
+                "then add both the repo root and its third_party/Matcha-TTS to the venv "
+                "(e.g. via a .pth file in the venv's site-packages).\n"
+                "If cosyvoice is present but a dependency is missing, install it into the venv."
+            ) from e
+
+        model_dir = self._get_model_path()
 
         device = self._resolve_device(self.config.device)
 

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/services/video_frame_extractor/drivers/common.py

@@ -83,6 +83,7 @@ class VideoFrameExtractorAction(ComponentAction):
             "start_time":      parse_time(start_time) if start_time is not None else None,
             "end_time":        parse_time(end_time) if end_time is not None else None,
             "max_frame_count": max_frame_count,
+            "filename_format": self.config.filename_format,
         }
 
     async def _process_batch(
@@ -114,6 +115,7 @@ class VideoFrameExtractorAction(ComponentAction):
             params["start_time"],
             params["end_time"],
             params["max_frame_count"],
+            params["filename_format"],
             streaming,
             cancellation_token,
         )
@@ -126,6 +128,7 @@ class VideoFrameExtractorAction(ComponentAction):
         start_time: Optional[float],
         end_time: Optional[float],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         streaming: bool,
         cancellation_token: Optional[CancellationToken] = None,
     ) -> Union[List[Dict[str, Any]], AsyncIterable[Dict[str, Any]]]:

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/services/video_frame_extractor/drivers/ffmpeg.py

@@ -13,6 +13,7 @@ from mindor.core.logger import logging
 from ..base import VideoFrameExtractorService, VideoFrameExtractorDriver, register_video_frame_extractor_service
 from ..base import ComponentActionContext
 from .common import VideoFrameExtractorAction
+from ..filename import format_filename
 from PIL import Image as PILImage
 from io import BytesIO
 import asyncio, os, re
@@ -35,6 +36,7 @@ class FFmpegVideoFrameExtractorAction(VideoFrameExtractorAction):
         start_time: Optional[float],
         end_time: Optional[float],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         streaming: bool,
         cancellation_token: Optional[CancellationToken] = None,
     ) -> Union[List[Dict[str, Any]], AsyncIterator[Dict[str, Any]]]:
@@ -85,9 +87,25 @@ class FFmpegVideoFrameExtractorAction(VideoFrameExtractorAction):
                     pass
 
         if streaming:
-            return self._stream_frames(command, video, input_path, max_frame_count, _cleanup, cancellation_token)
-
-        return await self._collect_frames(command, video, input_path, max_frame_count, _cleanup, cancellation_token)
+            return self._stream_frames(
+                command,
+                video,
+                input_path,
+                max_frame_count,
+                filename_format,
+                _cleanup,
+                cancellation_token,
+            )
+
+        return await self._collect_frames(
+            command,
+            video,
+            input_path,
+            max_frame_count,
+            filename_format,
+            _cleanup,
+            cancellation_token,
+        )
 
     async def _collect_frames(
         self,
@@ -95,6 +113,7 @@ class FFmpegVideoFrameExtractorAction(VideoFrameExtractorAction):
         video: MediaSource,
         input_path: Optional[str],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         cleanup: Callable[[], None],
         cancellation_token: Optional[CancellationToken] = None,
     ) -> List[Dict[str, Any]]:
@@ -176,10 +195,17 @@ class FFmpegVideoFrameExtractorAction(VideoFrameExtractorAction):
             frames = []
 
             for index, image in enumerate(images):
-                frames.append({
+                frame_count = index + 1
+                frame: Dict[str, Any] = {
+                    "number": frame_count,
                     "image": image,
-                    "timestamp": timestamps[index] if index < len(timestamps) else 0.0
-                })
+                    "timestamp": timestamps[index] if index < len(timestamps) else 0.0,
+                }
+
+                if filename_format is not None:
+                    frame["filename"] = format_filename(filename_format, frame_count)
+
+                frames.append(frame)
 
             return frames
         except asyncio.CancelledError:
@@ -201,6 +227,7 @@ class FFmpegVideoFrameExtractorAction(VideoFrameExtractorAction):
         video: MediaSource,
         input_path: Optional[str],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         cleanup: Callable[[], None],
         cancellation_token: Optional[CancellationToken] = None,
     ) -> AsyncIterator[Dict[str, Any]]:
@@ -268,11 +295,18 @@ class FFmpegVideoFrameExtractorAction(VideoFrameExtractorAction):
 
                 async for image in images:
                     timestamp = await timestamps.get()
-                    yield {
+                    frame_count += 1
+
+                    frame: Dict[str, Any] = {
+                        "number": frame_count,
                         "image": image,
                         "timestamp": timestamp if timestamp is not None else 0.0,
                     }
-                    frame_count += 1
+
+                    if filename_format is not None:
+                        frame["filename"] = format_filename(filename_format, frame_count)
+
+                    yield frame
 
                     if max_frame_count and frame_count >= max_frame_count:
                         return

{model_compose-0.4.85 → model_compose-0.4.87}/src/mindor/core/component/services/video_frame_extractor/drivers/opencv.py

@@ -13,6 +13,7 @@ from mindor.core.utils.streamer import SyncGeneratorStreamer
 from ..base import VideoFrameExtractorService, VideoFrameExtractorDriver, register_video_frame_extractor_service
 from ..base import ComponentActionContext
 from .common import VideoFrameExtractorAction
+from ..filename import format_filename
 from PIL import Image as PILImage
 import asyncio, os, threading
 
@@ -26,6 +27,7 @@ class OpenCVVideoFrameExtractorAction(VideoFrameExtractorAction):
         start_time: Optional[float],
         end_time: Optional[float],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         streaming: bool,
         cancellation_token: Optional[CancellationToken] = None,
     ) -> Union[List[Dict[str, Any]], AsyncIterator[Dict[str, Any]]]:
@@ -41,9 +43,27 @@ class OpenCVVideoFrameExtractorAction(VideoFrameExtractorAction):
                 pass
 
         if streaming:
-            return self._stream_frames(input_path, frame_interval, start_time, end_time, max_frame_count, _cleanup, cancellation_token)
-
-        return await self._collect_frames(input_path, frame_interval, start_time, end_time, max_frame_count, _cleanup, cancellation_token)
+            return self._stream_frames(
+                input_path,
+                frame_interval,
+                start_time,
+                end_time,
+                max_frame_count,
+                filename_format,
+                _cleanup,
+                cancellation_token,
+            )
+
+        return await self._collect_frames(
+            input_path,
+            frame_interval,
+            start_time,
+            end_time,
+            max_frame_count,
+            filename_format,
+            _cleanup,
+            cancellation_token,
+        )
 
     async def _collect_frames(
         self,
@@ -52,12 +72,21 @@ class OpenCVVideoFrameExtractorAction(VideoFrameExtractorAction):
         start_time: Optional[float],
         end_time: Optional[float],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         cleanup: Callable[[], None],
         cancellation_token: Optional[CancellationToken] = None,
     ) -> List[Dict[str, Any]]:
         def _collect() -> List[Dict[str, Any]]:
             try:
-                return list(self._extract_frames(input_path, frame_interval, start_time, end_time, max_frame_count, None))
+                return list(self._extract_frames(
+                    input_path,
+                    frame_interval,
+                    start_time,
+                    end_time,
+                    max_frame_count,
+                    filename_format,
+                    None,
+                ))
             finally:
                 cleanup()
 
@@ -70,12 +99,25 @@ class OpenCVVideoFrameExtractorAction(VideoFrameExtractorAction):
         start_time: Optional[float],
         end_time: Optional[float],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         cleanup: Callable[[], None],
         cancellation_token: Optional[CancellationToken] = None,
     ) -> AsyncIterator[Dict[str, Any]]:
         stop_event = threading.Event()
-        generator = self._extract_frames(input_path, frame_interval, start_time, end_time, max_frame_count, stop_event)
-        streamer = SyncGeneratorStreamer(generator, asyncio.get_running_loop(), maxsize=_FRAME_QUEUE_MAXSIZE, stop_event=stop_event)
+        streamer = SyncGeneratorStreamer(
+            self._extract_frames(
+                input_path,
+                frame_interval,
+                start_time,
+                end_time,
+                max_frame_count,
+                filename_format,
+                stop_event,
+            ),
+            asyncio.get_running_loop(),
+            maxsize=_FRAME_QUEUE_MAXSIZE,
+            stop_event=stop_event,
+        )
 
         try:
             async for frame in streamer:
@@ -91,6 +133,7 @@ class OpenCVVideoFrameExtractorAction(VideoFrameExtractorAction):
         start_time: Optional[float],
         end_time: Optional[float],
         max_frame_count: Optional[int],
+        filename_format: Optional[str],
         stop_event: Optional[threading.Event],
     ) -> Iterator[Dict[str, Any]]:
         import cv2
@@ -123,11 +166,17 @@ class OpenCVVideoFrameExtractorAction(VideoFrameExtractorAction):
 
                 if (current_frame - start_frame) % frame_interval == 0:
                     rgb_frame = cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2RGB)
-                    yield {
+                    frame_count += 1
+                    frame: Dict[str, Any] = {
+                        "number": frame_count,
                         "image": PILImage.fromarray(rgb_frame),
                         "timestamp": (current_frame / fps) if fps > 0 else 0.0,
                     }
-                    frame_count += 1
+
+                    if filename_format is not None:
+                        frame["filename"] = format_filename(filename_format, frame_count)
+
+                    yield frame
 
                     if max_frame_count and frame_count >= max_frame_count:
                         break