model-compose 0.4.5__tar.gz → 0.4.6__tar.gz

{model_compose-0.4.5 → model_compose-0.4.6}/PKG-INFO

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

{model_compose-0.4.5 → model_compose-0.4.6}/pyproject.toml

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

model_compose-0.4.6/src/mindor/core/component/services/model/base/__init__.py

@@ -0,0 +1,3 @@
+from .common import *
+from .huggingface import *
+from .unsloth import *

model_compose-0.4.6/src/mindor/core/component/services/model/base/common.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Callable, Mapping, Any
+from abc import ABC, abstractmethod
+from mindor.dsl.schema.component import ModelComponentConfig, ModelTaskType, ModelDriver, HuggingfaceModelConfig, LocalModelConfig
+from mindor.dsl.schema.action import ModelActionConfig
+from mindor.core.services import AsyncService
+from mindor.core.logger import logging
+from ....context import ComponentActionContext
+import asyncio
+
+if TYPE_CHECKING:
+    import torch
+
+class ModelTaskService(AsyncService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(daemon)
+
+        self.id: str = id
+        self.config: ModelComponentConfig = config
+
+    def get_setup_requirements(self) -> Optional[List[str]]:
+        return [ "torch" ]
+
+    async def run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+        loop: asyncio.AbstractEventLoop = asyncio.get_running_loop()
+
+        async def _run():
+            return await self._run(action, context, loop)
+
+        return await self.run_in_thread(_run)
+
+    @abstractmethod
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        pass
+
+    def _get_model_path(self) -> str:
+        if isinstance(self.config.model, HuggingfaceModelConfig):
+            from .huggingface_hub import get_model_path
+            return get_model_path(self.config.model)
+
+        if isinstance(self.config.model, LocalModelConfig):
+            return self.config.model.path
+        
+        if isinstance(self.config.model, str):
+            return self.config.model
+
+        raise ValueError(f"Unknown model config type: {type(self.config.model)}")
+
+    def _resolve_device(self) -> torch.device:
+        import torch
+
+        try:
+            return torch.device(self.config.device)
+        except:
+            logging.warning(f"Invalid device '{self.config.device}', falling back to 'cpu'")
+        
+        return torch.device("cpu")
+
+    def _load_model_checkpoint(self, model: torch.nn.Module, model_path: str) -> None:
+        checkpoint = torch.load(model_path, map_location="cpu")
+        state_dict = self._get_state_dict_from_checkpoint(checkpoint)
+
+        model.load_state_dict(state_dict, strict=True)
+
+    def _get_state_dict_from_checkpoint(self, checkpoint: Any) -> Mapping[str, Any]:
+        for key in [ "params", "state_dict" ]:
+            if key in checkpoint:
+                return checkpoint[key]
+        return checkpoint
+
+def register_model_task_service(type: ModelTaskType, driver: ModelDriver):
+    def decorator(cls: Type[ModelTaskService]) -> Type[ModelTaskService]:
+        if type not in ModelTaskServiceRegistry:
+            ModelTaskServiceRegistry[type] = {}
+        ModelTaskServiceRegistry[type][driver] = cls
+        return cls
+    return decorator
+
+ModelTaskServiceRegistry: Dict[ModelTaskType, Dict[ModelDriver, Type[ModelTaskService]]] = {}

model_compose-0.4.6/src/mindor/core/component/services/model/base/huggingface.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Callable, Any
+from mindor.dsl.schema.component import ModelComponentConfig, HuggingfaceModelConfig, DeviceMode
+from mindor.core.logger import logging
+from .common import ModelTaskService
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
+    import torch
+
+class HuggingfaceModelTaskService(ModelTaskService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(id, config, daemon)
+
+    def _load_pretrained_model(self) -> PreTrainedModel:
+        model_cls = self._get_model_class()
+        model = model_cls.from_pretrained(self.config.model, **self._get_model_params())
+
+        if self.config.device_mode == DeviceMode.SINGLE:
+            model = model.to(torch.device(self.config.device))
+
+        return model
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        raise NotImplementedError("Model class loader not implemented.")
+
+    def _get_model_params(self) -> Dict[str, Any]:
+        params: Dict[str, Any] = {}
+
+        if isinstance(self.config.model, HuggingfaceModelConfig):
+            if self.config.model.revision:
+                params["revision"] = self.config.model.revision
+            
+            if self.config.model.cache_dir:
+                params["cache_dir"] = self.config.model.cache_dir
+
+            if self.config.model.local_files_only:
+                params["local_files_only"] = True
+
+        if self.config.device_mode != DeviceMode.SINGLE:
+            params["device_map"] = self.config.device_mode.value
+    
+        if self.config.precision is not None:
+            params["torch_dtype"] = getattr(torch, self.config.precision.value)
+    
+        if self.config.low_cpu_mem_usage:
+            params["low_cpu_mem_usage"] = True
+
+        return params
+
+    def _get_model_device(self, model: PreTrainedModel) -> torch.device:
+        return next(model.parameters()).device
+
+class HuggingfaceLanguageModelTaskService(HuggingfaceModelTaskService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(id, config, daemon)
+
+        self.model: Optional[PreTrainedModel] = None
+        self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
+
+    def get_setup_requirements(self) -> Optional[List[str]]:
+        return [ 
+            "transformers>=4.21.0",
+            "torch",
+            "sentencepiece",
+            "accelerate"
+        ]
+
+    async def _serve(self) -> None:
+        try:
+            self.model = self._load_pretrained_model()
+            self.tokenizer = self._load_pretrained_tokenizer()
+            self.device = self._get_model_device(self.model)
+            logging.info(f"Model and tokenizer loaded successfully on device '{self.device}': {self.config.model}")
+        except Exception as e:
+            logging.error(f"Failed to load model '{self.config.model}': {e}")
+            raise
+
+    async def _shutdown(self) -> None:
+        self.model = None
+        self.tokenizer = None
+        self.device = None
+
+    def _load_pretrained_tokenizer(self) -> Optional[PreTrainedTokenizer]:
+        tokenizer_cls = self._get_tokenizer_class()
+
+        if not tokenizer_cls:
+            return None
+
+        return tokenizer_cls.from_pretrained(self.config.model, **self._get_tokenizer_params())
+
+    def _get_tokenizer_class(self) -> Optional[Type[PreTrainedTokenizer]]:
+        return None
+
+    def _get_tokenizer_params(self) -> Dict[str, Any]:
+        params: Dict[str, Any] = {}
+
+        if isinstance(self.config.model, HuggingfaceModelConfig):
+            if self.config.model.revision:
+                params["revision"] = self.config.model.revision
+
+            if self.config.model.cache_dir:
+                params["cache_dir"] = self.config.model.cache_dir
+
+            if self.config.model.local_files_only:
+                params["local_files_only"] = True
+
+        if not self.config.fast_tokenizer:
+            params["use_fast"] = False
+
+        return params
+
+class HuggingfaceMultimodalModelTaskService(HuggingfaceModelTaskService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(id, config, daemon)
+
+        self.model: Optional[PreTrainedModel] = None
+        self.processor: Optional[ProcessorMixin] = None
+        self.device: Optional[torch.device] = None
+
+    def get_setup_requirements(self) -> Optional[List[str]]:
+        return [ 
+            "transformers>=4.21.0",
+            "torch",
+            "sentencepiece",
+            "accelerate"
+        ]
+
+    async def _serve(self) -> None:
+        try:
+            self.model = self._load_pretrained_model()
+            self.processor = self._load_pretrained_processor()
+            self.device = self._get_model_device(self.model)
+            logging.info(f"Model and processor loaded successfully on device '{self.device}': {self.config.model}")
+        except Exception as e:
+            logging.error(f"Failed to load model '{self.config.model}': {e}")
+            raise
+
+    async def _shutdown(self) -> None:
+        self.model = None
+        self.processor = None
+        self.device = None
+
+    def _load_pretrained_processor(self) -> Optional[ProcessorMixin]:
+        processor_cls = self._get_processor_class()
+
+        if not processor_cls:
+            return None
+
+        return processor_cls.from_pretrained(self.config.model, **self._get_processor_params())
+
+    def _get_processor_class(self) -> Optional[Type[ProcessorMixin]]:
+        return None
+
+    def _get_processor_params(self) -> Dict[str, Any]:
+        params: Dict[str, Any] = {}
+
+        if isinstance(self.config.model, HuggingfaceModelConfig):
+            if self.config.model.revision:
+                params["revision"] = self.config.model.revision
+ 
+            if self.config.model.cache_dir:
+                params["cache_dir"] = self.config.model.cache_dir
+
+            if self.config.model.local_files_only:
+                params["local_files_only"] = True
+
+        return params

model_compose-0.4.6/src/mindor/core/component/services/model/base/huggingface_hub.py

@@ -0,0 +1,23 @@
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Callable, Any
+from mindor.dsl.schema.component import HuggingfaceModelConfig
+
+def get_model_path(model: HuggingfaceModelConfig) -> Optional[str]:
+    if model.filename:
+        from huggingface_hub import hf_hub_download
+        return hf_hub_download(
+            repo_id=model.repository,
+            filename=model.filename,
+            revision=model.revision,
+            cache_dir=model.cache_dir,
+            token=model.token,
+            local_files_only=model.local_files_only
+        )
+    else:
+        from huggingface_hub import snapshot_download
+        return snapshot_download(
+            repo_id=model.repository,
+            revision=model.revision,
+            cache_dir=model.cache_dir,
+            token=model.token,
+            local_files_only=model.local_files_only
+        )

model_compose-0.4.6/src/mindor/core/component/services/model/base/unsloth.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Callable, Any
+from mindor.dsl.schema.component import ModelComponentConfig, DeviceMode
+from mindor.core.logger import logging
+from .common import ModelTaskService
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+
+class UnslothModelTaskService(ModelTaskService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(id, config, daemon)
+
+        self.model: Optional[PreTrainedModel] = None
+        self.tokenizer: Optional[PreTrainedTokenizer] = None
+
+    def get_setup_requirements(self) -> Optional[List[str]]:
+        return [ 
+            "unsloth",
+            "torch"
+        ]
+
+    async def _serve(self) -> None:
+        try:
+            self.model, self.tokenizer = self._load_pretrained_model()
+            logging.info(f"Model and tokenizer loaded successfully: {self.config.model}")
+        except Exception as e:
+            logging.error(f"Failed to load model '{self.config.model}': {e}")
+            raise
+
+    async def _shutdown(self) -> None:
+        self.model = None
+        self.tokenizer = None
+        self.device = None
+
+    def _load_pretrained_model(self) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+        from unsloth import FastLanguageModel
+        
+        model = FastLanguageModel.from_pretrained(self.config.model, **self._get_model_params())
+
+        return model
+
+    def _get_model_params(self) -> Dict[str, Any]:
+        params: Dict[str, Any] = {}
+
+        if self.config.device_mode != DeviceMode.SINGLE:
+            params["device_map"] = self.config.device_mode.value
+    
+        if self.config.precision is not None:
+            params["dtype"] = self.config.precision.value
+
+        return params

{model_compose-0.4.5 → model_compose-0.4.6}/src/mindor/core/component/services/model/model.py

@@ -1,5 +1,5 @@
 from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, AsyncIterator, Any
-from mindor.dsl.schema.component import ModelComponentConfig, ModelTaskType
+from mindor.dsl.schema.component import ModelComponentConfig, ModelTaskType, ModelDriver
 from mindor.dsl.schema.action import ActionConfig, ModelActionConfig
 from ...base import ComponentService, ComponentType, ComponentGlobalConfigs, register_component
 from ...context import ComponentActionContext
@@ -17,26 +17,18 @@ class ModelComponent(ComponentService):
     def __init__(self, id: str, config: ModelComponentConfig, global_configs: ComponentGlobalConfigs, daemon: bool):
         super().__init__(id, config, global_configs, daemon)
 
-        self.task_service: ModelTaskService = self._create_task_service(self.config.task)
+        self.task_service: ModelTaskService = self._create_task_service(self.config.task, self.config.driver)
 
-    def _create_task_service(self, type: ModelTaskType) -> ModelTaskService:
+    def _create_task_service(self, type: ModelTaskType, driver: ModelDriver) -> ModelTaskService:
         try:
             if not ModelTaskServiceRegistry:
                 from . import tasks
-            return ModelTaskServiceRegistry[type](self.id, self.config, self.daemon)
+            return ModelTaskServiceRegistry[type][driver](self.id, self.config, self.daemon)
         except KeyError:
-            raise ValueError(f"Unsupported model task type: {type}")
+            raise ValueError(f"Unsupported model task type: {type} on {driver}")
 
     def _get_setup_requirements(self) -> Optional[List[str]]:
-        task_dependencies = self.task_service.get_setup_requirements()
-
-        return [ 
-            "transformers>=4.21.0",
-            "torch",
-            "sentencepiece",
-            "accelerate",
-            *(task_dependencies or [])
-        ]
+        return self.task_service.get_setup_requirements()
 
     async def _serve(self) -> None:
         await self.task_service.start()

{model_compose-0.4.5 → model_compose-0.4.6}/src/mindor/core/component/services/model/tasks/__init__.py

@@ -1,6 +1,6 @@
 from .text_generation import *
-from .summarization import *
-from .translation import *
+from .chat_completion import *
 from .text_classification import *
 from .text_embedding import *
 from .image_to_text import *
+from .image_upscale import *

model_compose-0.4.6/src/mindor/core/component/services/model/tasks/chat_completion/__init__.py

@@ -0,0 +1 @@
+from .huggingface import *

model_compose-0.4.6/src/mindor/core/component/services/model/tasks/chat_completion/huggingface.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Any
+from mindor.dsl.schema.action import ModelActionConfig, ChatCompletionModelActionConfig
+from ...base import ModelTaskType, ModelDriver, register_model_task_service
+from ...base import HuggingfaceLanguageModelTaskService, ComponentActionContext
+from ..text_generation import HuggingfaceTextGenerationTaskAction
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+    import torch
+
+class HuggingfaceChatCompletionTaskAction(HuggingfaceTextGenerationTaskAction):
+    def __init__(self, config: ChatCompletionModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
+        super().__init__(config, model, tokenizer, device)
+
+        self.config: ChatCompletionModelActionConfig = config # For type only
+
+    async def _prepare_input(self, context: ComponentActionContext) -> Union[str, List[str]]:
+        messages = await context.render_variable(self.config.messages)
+
+        return self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+@register_model_task_service(ModelTaskType.CHAT_COMPLETION, ModelDriver.HUGGINGFACE)
+class HuggingfaceChatCompletionTaskService(HuggingfaceLanguageModelTaskService):
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        return await HuggingfaceChatCompletionTaskAction(action, self.model, self.tokenizer, self.device).run(context, loop)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        from transformers import AutoModelForCausalLM
+        return AutoModelForCausalLM
+
+    def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        from transformers import AutoTokenizer
+        return AutoTokenizer

model_compose-0.4.6/src/mindor/core/component/services/model/tasks/image_to_text/__init__.py

@@ -0,0 +1 @@
+from .huggingface import *

model_compose-0.4.5/src/mindor/core/component/services/model/tasks/image_to_text.py → model_compose-0.4.6/src/mindor/core/component/services/model/tasks/image_to_text/huggingface.py

@@ -2,12 +2,12 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Protocol, Any
-from mindor.dsl.schema.component import ModelComponentConfig, ImageToTextModelArchitecture
+from mindor.dsl.schema.component import ImageToTextModelArchitecture
 from mindor.dsl.schema.action import ModelActionConfig, ImageToTextModelActionConfig
 from mindor.core.utils.streamer import AsyncStreamer
 from mindor.core.logger import logging
-from ..base import ModelTaskService, ModelTaskType, register_model_task_service
-from ..base import ComponentActionContext
+from ...base import ModelTaskType, ModelDriver, register_model_task_service
+from ...base import HuggingfaceMultimodalModelTaskService, ComponentActionContext
 from PIL import Image as PILImage
 from threading import Thread
 import asyncio
@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 class WithTokenizer(Protocol):
     tokenizer: PreTrainedTokenizer
 
-class ImageToTextTaskAction:
+class HuggingfaceImageToTextTaskAction:
     def __init__(self, config: ImageToTextModelActionConfig, model: PreTrainedModel, processor: ProcessorMixin, device: torch.device):
         self.config: ImageToTextModelActionConfig = config
         self.model: Union[PreTrainedModel, GenerationMixin] = model
@@ -28,60 +28,39 @@ class ImageToTextTaskAction:
         self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
-        from transformers import TextIteratorStreamer, StopStringCriteria
+        from transformers import TextIteratorStreamer, StopStringCriteria, GenerationConfig
         import torch
 
-        image: Union[PILImage.Image, List[PILImage.Image]] = await context.render_image(self.config.image)
-        prompt: Optional[Union[str, List[str]]] = await context.render_variable(self.config.prompt)
-
-        max_input_length     = await context.render_variable(self.config.params.max_input_length)
-        max_output_length    = await context.render_variable(self.config.params.max_output_length)
-        min_output_length    = await context.render_variable(self.config.params.min_output_length)
-        num_return_sequences = await context.render_variable(self.config.params.num_return_sequences)
-        do_sample            = await context.render_variable(self.config.params.do_sample)
-        temperature          = await context.render_variable(self.config.params.temperature) if do_sample else None
-        top_k                = await context.render_variable(self.config.params.top_k) if do_sample else None
-        top_p                = await context.render_variable(self.config.params.top_p) if do_sample else None
-        num_beams            = await context.render_variable(self.config.params.num_beams)
-        length_penalty       = await context.render_variable(self.config.params.length_penalty) if num_beams > 1 else None
-        early_stopping       = await context.render_variable(self.config.params.early_stopping) if num_beams > 1 else False
-        stop_sequences       = await context.render_variable(self.config.params.stop_sequences)
-        batch_size           = await context.render_variable(self.config.params.batch_size)
-        stream               = await context.render_variable(self.config.stream)
-
-        is_single_input: bool = bool(not isinstance(images, list))
+        image, text = await self._prepare_input(context)
+        is_single_input: bool = bool(not isinstance(image, list))
         images: List[PILImage.Image] = [ image ] if is_single_input else image
-        prompts: Optional[List[str]] = [ prompt ] if is_single_input else prompt
+        texts: Optional[List[str]] = [ text ] if is_single_input else text
         results = []
 
+        batch_size        = await context.render_variable(self.config.batch_size)
+        stream            = await context.render_variable(self.config.stream)
+        stop_sequences    = await context.render_variable(self.config.stop_sequences)
+        processor_params  = await self._resolve_processor_params(context)
+        generation_params = await self._resolve_generation_params(context)
+
         if stream and (batch_size != 1 or len(images) != 1):
             raise ValueError("Streaming mode only supports a single input image with batch size of 1.")
 
         streamer = TextIteratorStreamer(self.processor.tokenizer, skip_prompt=True, skip_special_tokens=True) if stream else None
         stopping_criteria = [ StopStringCriteria(self.processor.tokenizer, stop_sequences) ] if stop_sequences else None
+
         for index in range(0, len(images), batch_size):
             batch_images = images[index:index + batch_size]
-            batch_prompts = prompts[index:index + batch_size] if prompts else None
-            
-            inputs: Tensor = self.processor(images=batch_images, text=batch_prompts, max_length=max_input_length, return_tensors="pt")
+            batch_texts = texts[index:index + batch_size] if texts else None
+
+            inputs: Tensor = self.processor(images=batch_images, texts=batch_texts, **processor_params)
             inputs = inputs.to(self.device)
 
             def _generate():
                 with torch.inference_mode():
                     outputs = self.model.generate(
-                        **inputs,
-                        max_new_tokens=max_output_length,
-                        min_length=min_output_length,
-                        num_return_sequences=num_return_sequences,
-                        do_sample=do_sample,
-                        temperature=temperature,
-                        top_k=top_k,
-                        top_p=top_p,
-                        num_beams=num_beams,
-                        length_penalty=length_penalty,
-                        early_stopping=early_stopping,
-                        pad_token_id=getattr(self.processor.tokenizer, "pad_token_id", None),
-                        eos_token_id=getattr(self.processor.tokenizer, "eos_token_id", None),
+                        **inputs, 
+                        generation_config=GenerationConfig(**generation_params),
                         stopping_criteria=stopping_criteria,
                         streamer=streamer
                     )
@@ -109,32 +88,68 @@ class ImageToTextTaskAction:
 
             return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
 
-@register_model_task_service(ModelTaskType.IMAGE_TO_TEXT)
-class ImageToTextTaskService(ModelTaskService):
-    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
-        super().__init__(id, config, daemon)
-
-        self.model: Optional[Union[PreTrainedModel, GenerationMixin]] = None
-        self.processor: Optional[ProcessorMixin] = None
-        self.device: Optional[torch.device] = None
-
-    async def _serve(self) -> None:
-        try:
-            self.model = self._load_pretrained_model()
-            self.processor = self._load_pretrained_processor()
-            self.device = self._get_model_device(self.model)
-            logging.info(f"Model and processor loaded successfully on device '{self.device}': {self.config.model}")
-        except Exception as e:
-            logging.error(f"Failed to load model '{self.config.model}': {e}")
-            raise
-
-    async def _shutdown(self) -> None:
-        self.model = None
-        self.processor = None
-        self.device = None
+    async def _prepare_input(self, context: ComponentActionContext) -> Tuple[Union[PILImage.Image, List[PILImage.Image]], Optional[Union[str, List[str]]]]:
+        image = await context.render_image(self.config.image)
+        text  = await context.render_variable(self.config.text)
+        
+        return image, text
+
+    async def _resolve_processor_params(self, context: ComponentActionContext) -> Dict[str, Any]:
+        max_input_length = await context.render_variable(self.config.max_input_length)
+
+        params = {
+            "return_tensors": "pt",
+            "padding": True,
+            "truncation": False
+        }
+
+        if max_input_length is not None:
+            params["max_length"] = max_input_length
+            params["truncation"] = True
+
+        return params
+
+    async def _resolve_generation_params(self, context: ComponentActionContext) -> Dict[str, Any]:
+        max_output_length    = await context.render_variable(self.config.params.max_output_length)
+        min_output_length    = await context.render_variable(self.config.params.min_output_length)
+        num_return_sequences = await context.render_variable(self.config.params.num_return_sequences)
+        do_sample            = await context.render_variable(self.config.params.do_sample)
+        temperature          = await context.render_variable(self.config.params.temperature) if do_sample else None
+        top_k                = await context.render_variable(self.config.params.top_k) if do_sample else None
+        top_p                = await context.render_variable(self.config.params.top_p) if do_sample else None
+        num_beams            = await context.render_variable(self.config.params.num_beams)
+        length_penalty       = await context.render_variable(self.config.params.length_penalty) if num_beams > 1 else None
+        early_stopping       = await context.render_variable(self.config.params.early_stopping) if num_beams > 1 else False
 
+        params = {
+            "max_new_tokens": max_output_length,
+            "min_length": min_output_length,
+            "num_return_sequences": num_return_sequences,
+            "do_sample": do_sample,
+            "num_beams": num_beams,
+            "pad_token_id": getattr(self.processor.tokenizer, "pad_token_id", None),
+            "eos_token_id": getattr(self.processor.tokenizer, "eos_token_id", None),
+        }
+
+        if do_sample:
+            if temperature is not None:
+                params["temperature"] = temperature
+            if top_k is not None:
+                params["top_k"] = top_k
+            if top_p is not None:
+                params["top_p"] = top_p
+
+        if num_beams > 1:
+            if length_penalty is not None:
+                params["length_penalty"] = length_penalty
+            params["early_stopping"] = early_stopping
+
+        return params
+
+@register_model_task_service(ModelTaskType.IMAGE_TO_TEXT, ModelDriver.HUGGINGFACE)
+class HuggingfaceImageToTextTaskService(HuggingfaceMultimodalModelTaskService):
     async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
-        return await ImageToTextTaskAction(action, self.model, self.processor, self.device).run(context, loop)
+        return await HuggingfaceImageToTextTaskAction(action, self.model, self.processor, self.device).run(context, loop)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
         if self.config.architecture == ImageToTextModelArchitecture.BLIP:

model_compose-0.4.6/src/mindor/core/component/services/model/tasks/image_upscale/__init__.py

@@ -0,0 +1 @@
+from .custom import *