|
| 1 | +from typing import List, Literal, Optional, Type, Union |
| 2 | + |
| 3 | +from pydantic import ConfigDict, Field |
| 4 | + |
| 5 | +from inference.core.entities.requests.inference import LMMInferenceRequest |
| 6 | +from inference.core.managers.base import ModelManager |
| 7 | +from inference.core.workflows.core_steps.common.entities import StepExecutionMode |
| 8 | +from inference.core.workflows.execution_engine.entities.base import ( |
| 9 | + Batch, |
| 10 | + OutputDefinition, |
| 11 | + WorkflowImageData, |
| 12 | +) |
| 13 | +from inference.core.workflows.execution_engine.entities.types import ( |
| 14 | + DICTIONARY_KIND, |
| 15 | + IMAGE_KIND, |
| 16 | + ROBOFLOW_MODEL_ID_KIND, |
| 17 | + ImageInputField, |
| 18 | + Selector, |
| 19 | +) |
| 20 | +from inference.core.workflows.prototypes.block import ( |
| 21 | + BlockResult, |
| 22 | + WorkflowBlock, |
| 23 | + WorkflowBlockManifest, |
| 24 | +) |
| 25 | + |
| 26 | + |
| 27 | +class BlockManifest(WorkflowBlockManifest): |
| 28 | + # SmolVLM needs an image and a text prompt. |
| 29 | + images: Selector(kind=[IMAGE_KIND]) = ImageInputField |
| 30 | + prompt: Optional[str] = Field( |
| 31 | + default=None, |
| 32 | + description="Optional text prompt to provide additional context to SmolVLM2. Otherwise it will just be None", |
| 33 | + examples=["What is in this image?"], |
| 34 | + ) |
| 35 | + |
| 36 | + # Standard model configuration for UI, schema, etc. |
| 37 | + model_config = ConfigDict( |
| 38 | + json_schema_extra={ |
| 39 | + "name": "SmolVLM2", |
| 40 | + "version": "v1", |
| 41 | + "short_description": "Run SmolVLM2 on an image.", |
| 42 | + "long_description": ( |
| 43 | + "This workflow block runs SmolVLM2, a multimodal vision-language model. You can ask questions about images" |
| 44 | + " -- including documents and photos -- and get answers in natural language." |
| 45 | + ), |
| 46 | + "license": "Apache-2.0", |
| 47 | + "block_type": "model", |
| 48 | + "search_keywords": [ |
| 49 | + "SmolVLM2", |
| 50 | + "smolvlm", |
| 51 | + "vision language model", |
| 52 | + "VLM", |
| 53 | + ], |
| 54 | + "is_vlm_block": True, |
| 55 | + "ui_manifest": { |
| 56 | + "section": "model", |
| 57 | + "icon": "fal fa-atom", |
| 58 | + "blockPriority": 5.5, |
| 59 | + }, |
| 60 | + }, |
| 61 | + protected_namespaces=(), |
| 62 | + ) |
| 63 | + type: Literal["roboflow_core/smolvlm2@v1"] |
| 64 | + |
| 65 | + model_version: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = Field( |
| 66 | + default="smolvlm2/smolvlm-2.2b-instruct", |
| 67 | + description="The SmolVLM2 model to be used for inference.", |
| 68 | + examples=["smolvlm2/smolvlm-2.2b-instruct"], |
| 69 | + ) |
| 70 | + |
| 71 | + @classmethod |
| 72 | + def describe_outputs(cls) -> List[OutputDefinition]: |
| 73 | + return [ |
| 74 | + OutputDefinition( |
| 75 | + name="parsed_output", |
| 76 | + kind=[DICTIONARY_KIND], |
| 77 | + description="A parsed version of the output, provided as a dictionary containing the text.", |
| 78 | + ), |
| 79 | + ] |
| 80 | + |
| 81 | + @classmethod |
| 82 | + def get_parameters_accepting_batches(cls) -> List[str]: |
| 83 | + # Only images can be passed in as a list/batch |
| 84 | + return ["images"] |
| 85 | + |
| 86 | + @classmethod |
| 87 | + def get_execution_engine_compatibility(cls) -> Optional[str]: |
| 88 | + return ">=1.3.0,<2.0.0" |
| 89 | + |
| 90 | + |
| 91 | +class SmolVLM2BlockV1(WorkflowBlock): |
| 92 | + def __init__( |
| 93 | + self, |
| 94 | + model_manager: ModelManager, |
| 95 | + api_key: Optional[str], |
| 96 | + step_execution_mode: StepExecutionMode, |
| 97 | + ): |
| 98 | + self._model_manager = model_manager |
| 99 | + self._api_key = api_key |
| 100 | + self._step_execution_mode = step_execution_mode |
| 101 | + |
| 102 | + @classmethod |
| 103 | + def get_init_parameters(cls) -> List[str]: |
| 104 | + return ["model_manager", "api_key", "step_execution_mode"] |
| 105 | + |
| 106 | + @classmethod |
| 107 | + def get_manifest(cls) -> Type[WorkflowBlockManifest]: |
| 108 | + return BlockManifest |
| 109 | + |
| 110 | + def run( |
| 111 | + self, |
| 112 | + images: Batch[WorkflowImageData], |
| 113 | + model_version: str, |
| 114 | + prompt: Optional[str], |
| 115 | + ) -> BlockResult: |
| 116 | + if self._step_execution_mode == StepExecutionMode.LOCAL: |
| 117 | + return self.run_locally( |
| 118 | + images=images, |
| 119 | + model_version=model_version, |
| 120 | + prompt=prompt, |
| 121 | + ) |
| 122 | + elif self._step_execution_mode == StepExecutionMode.REMOTE: |
| 123 | + raise NotImplementedError( |
| 124 | + "Remote execution is not supported for SmolVLM2. Please use a local or dedicated inference server." |
| 125 | + ) |
| 126 | + else: |
| 127 | + raise ValueError( |
| 128 | + f"Unknown step execution mode: {self._step_execution_mode}" |
| 129 | + ) |
| 130 | + |
| 131 | + def run_locally( |
| 132 | + self, |
| 133 | + images: Batch[WorkflowImageData], |
| 134 | + model_version: str, |
| 135 | + prompt: Optional[str], |
| 136 | + ) -> BlockResult: |
| 137 | + # Convert each image to the format required by the model. |
| 138 | + inference_images = [ |
| 139 | + i.to_inference_format(numpy_preferred=False) for i in images |
| 140 | + ] |
| 141 | + # Use the provided prompt (or an empty string if None) for every image. |
| 142 | + prompt = prompt or "" |
| 143 | + prompts = [prompt] * len(inference_images) |
| 144 | + |
| 145 | + # Register SmolVLM2 with the model manager. |
| 146 | + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) |
| 147 | + |
| 148 | + predictions = [] |
| 149 | + for image, single_prompt in zip(inference_images, prompts): |
| 150 | + # Build an LMMInferenceRequest with both prompt and image. |
| 151 | + request = LMMInferenceRequest( |
| 152 | + api_key=self._api_key, |
| 153 | + model_id=model_version, |
| 154 | + image=image, |
| 155 | + source="workflow-execution", |
| 156 | + prompt=single_prompt, |
| 157 | + ) |
| 158 | + # Run inference. |
| 159 | + prediction = self._model_manager.infer_from_request_sync( |
| 160 | + model_id=model_version, request=request |
| 161 | + ) |
| 162 | + response_text = prediction.response |
| 163 | + predictions.append( |
| 164 | + { |
| 165 | + "parsed_output": response_text, |
| 166 | + } |
| 167 | + ) |
| 168 | + return predictions |
0 commit comments