|
| 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.env import ( |
| 7 | + COSMOS3_ENABLED, |
| 8 | + HOSTED_CORE_MODEL_URL, |
| 9 | + LOCAL_INFERENCE_API_URL, |
| 10 | + WORKFLOWS_REMOTE_API_TARGET, |
| 11 | +) |
| 12 | +from inference.core.managers.base import ModelManager |
| 13 | +from inference.core.workflows.core_steps.common.entities import StepExecutionMode |
| 14 | +from inference.core.workflows.execution_engine.entities.base import ( |
| 15 | + Batch, |
| 16 | + OutputDefinition, |
| 17 | + WorkflowImageData, |
| 18 | +) |
| 19 | +from inference.core.workflows.execution_engine.entities.types import ( |
| 20 | + IMAGE_KIND, |
| 21 | + LANGUAGE_MODEL_OUTPUT_KIND, |
| 22 | + ROBOFLOW_MODEL_ID_KIND, |
| 23 | + STRING_KIND, |
| 24 | + ImageInputField, |
| 25 | + Selector, |
| 26 | +) |
| 27 | +from inference.core.workflows.prototypes.block import ( |
| 28 | + BlockResult, |
| 29 | + Runtime, |
| 30 | + RuntimeRestriction, |
| 31 | + Severity, |
| 32 | + WorkflowBlock, |
| 33 | + WorkflowBlockManifest, |
| 34 | +) |
| 35 | +from inference_sdk import InferenceHTTPClient |
| 36 | + |
| 37 | +DEFAULT_PROMPT = "Describe what's in this image." |
| 38 | +DEFAULT_SYSTEM_PROMPT = ( |
| 39 | + "You are Cosmos, a helpful assistant that understands physical scenes " |
| 40 | + "and answers questions about images and videos." |
| 41 | +) |
| 42 | + |
| 43 | + |
| 44 | +class BlockManifest(WorkflowBlockManifest): |
| 45 | + model_config = ConfigDict( |
| 46 | + json_schema_extra={ |
| 47 | + "name": "Cosmos 3", |
| 48 | + "version": "v1", |
| 49 | + "short_description": "Reason about physical scenes with NVIDIA Cosmos 3.", |
| 50 | + "long_description": ( |
| 51 | + "This workflow block runs the NVIDIA Cosmos 3 reasoner — a world model " |
| 52 | + "tuned for physical scene understanding—on an image with an optional text " |
| 53 | + "prompt, and returns a text answer. The model is well suited to questions " |
| 54 | + "about spatial relations, safety, and what will happen next in a scene." |
| 55 | + ), |
| 56 | + "license": "OpenMDW-1.1", |
| 57 | + "block_type": "model", |
| 58 | + "search_keywords": [ |
| 59 | + "Cosmos", |
| 60 | + "cosmos3", |
| 61 | + "NVIDIA", |
| 62 | + "world model", |
| 63 | + "vision language model", |
| 64 | + "VLM", |
| 65 | + "robotics", |
| 66 | + ], |
| 67 | + "is_vlm_block": True, |
| 68 | + "ui_manifest": { |
| 69 | + "section": "model", |
| 70 | + "icon": "fal fa-atom", |
| 71 | + "blockPriority": 5.5, |
| 72 | + }, |
| 73 | + }, |
| 74 | + protected_namespaces=(), |
| 75 | + ) |
| 76 | + type: Literal["roboflow_core/cosmos3_edge@v1"] |
| 77 | + |
| 78 | + images: Selector(kind=[IMAGE_KIND]) = ImageInputField |
| 79 | + prompt: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( |
| 80 | + default=None, |
| 81 | + description="Optional text prompt to pass to Cosmos 3 Edge. Otherwise a default " |
| 82 | + "scene-description prompt is used, which may affect the desired model behavior.", |
| 83 | + examples=["Is the walkway free of obstacles?", "$inputs.prompt"], |
| 84 | + ) |
| 85 | + model_version: Union[ |
| 86 | + Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), Literal["nvidia/cosmos-3-edge"] |
| 87 | + ] = Field( |
| 88 | + default="nvidia/cosmos-3-edge", |
| 89 | + description="The Cosmos 3 Edge model to be used for inference.", |
| 90 | + examples=["nvidia/cosmos-3-edge"], |
| 91 | + ) |
| 92 | + system_prompt: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( |
| 93 | + default=None, |
| 94 | + description="Optional system prompt to provide additional context to Cosmos 3 Edge.", |
| 95 | + examples=["You are a safety inspector.", "$inputs.system_prompt"], |
| 96 | + ) |
| 97 | + |
| 98 | + @classmethod |
| 99 | + def describe_outputs(cls) -> List[OutputDefinition]: |
| 100 | + return [ |
| 101 | + OutputDefinition( |
| 102 | + name="output", |
| 103 | + kind=[STRING_KIND, LANGUAGE_MODEL_OUTPUT_KIND], |
| 104 | + description="The model's text answer for each input image.", |
| 105 | + ), |
| 106 | + ] |
| 107 | + |
| 108 | + @classmethod |
| 109 | + def get_parameters_accepting_batches(cls) -> List[str]: |
| 110 | + return ["images"] |
| 111 | + |
| 112 | + @classmethod |
| 113 | + def get_execution_engine_compatibility(cls) -> Optional[str]: |
| 114 | + return ">=1.3.0,<2.0.0" |
| 115 | + |
| 116 | + @classmethod |
| 117 | + def get_restrictions(cls) -> List[RuntimeRestriction]: |
| 118 | + restrictions = [ |
| 119 | + RuntimeRestriction( |
| 120 | + severity=Severity.HARD, |
| 121 | + note="Requires a GPU; run_locally() loads a model that needs CUDA.", |
| 122 | + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], |
| 123 | + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], |
| 124 | + ), |
| 125 | + ] |
| 126 | + if not COSMOS3_ENABLED: |
| 127 | + restrictions.append( |
| 128 | + RuntimeRestriction( |
| 129 | + severity=Severity.HARD, |
| 130 | + note=( |
| 131 | + "COSMOS3_ENABLED=False: the Cosmos 3 Edge endpoint is not " |
| 132 | + "registered, so run_remotely() returns 404." |
| 133 | + ), |
| 134 | + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], |
| 135 | + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], |
| 136 | + ) |
| 137 | + ) |
| 138 | + return restrictions |
| 139 | + |
| 140 | + @classmethod |
| 141 | + def get_supported_model_variants(cls) -> Optional[List[str]]: |
| 142 | + """Return list of model_id variants that can satisfy this block.""" |
| 143 | + return ["nvidia/cosmos-3-edge"] |
| 144 | + |
| 145 | + |
| 146 | +class Cosmos3EdgeBlockV1(WorkflowBlock): |
| 147 | + def __init__( |
| 148 | + self, |
| 149 | + model_manager: ModelManager, |
| 150 | + api_key: Optional[str], |
| 151 | + step_execution_mode: StepExecutionMode, |
| 152 | + ): |
| 153 | + self._model_manager = model_manager |
| 154 | + self._api_key = api_key |
| 155 | + self._step_execution_mode = step_execution_mode |
| 156 | + |
| 157 | + @classmethod |
| 158 | + def get_init_parameters(cls) -> List[str]: |
| 159 | + return ["model_manager", "api_key", "step_execution_mode"] |
| 160 | + |
| 161 | + @classmethod |
| 162 | + def get_manifest(cls) -> Type[WorkflowBlockManifest]: |
| 163 | + return BlockManifest |
| 164 | + |
| 165 | + def run( |
| 166 | + self, |
| 167 | + images: Batch[WorkflowImageData], |
| 168 | + model_version: str, |
| 169 | + prompt: Optional[str], |
| 170 | + system_prompt: Optional[str], |
| 171 | + ) -> BlockResult: |
| 172 | + if self._step_execution_mode == StepExecutionMode.LOCAL: |
| 173 | + return self.run_locally( |
| 174 | + images=images, |
| 175 | + model_version=model_version, |
| 176 | + prompt=prompt, |
| 177 | + system_prompt=system_prompt, |
| 178 | + ) |
| 179 | + elif self._step_execution_mode == StepExecutionMode.REMOTE: |
| 180 | + return self.run_remotely( |
| 181 | + images=images, |
| 182 | + model_version=model_version, |
| 183 | + prompt=prompt, |
| 184 | + system_prompt=system_prompt, |
| 185 | + ) |
| 186 | + else: |
| 187 | + raise ValueError( |
| 188 | + f"Unknown step execution mode: {self._step_execution_mode}" |
| 189 | + ) |
| 190 | + |
| 191 | + def run_remotely( |
| 192 | + self, |
| 193 | + images: Batch[WorkflowImageData], |
| 194 | + model_version: str, |
| 195 | + prompt: Optional[str], |
| 196 | + system_prompt: Optional[str], |
| 197 | + ) -> BlockResult: |
| 198 | + api_url = ( |
| 199 | + LOCAL_INFERENCE_API_URL |
| 200 | + if WORKFLOWS_REMOTE_API_TARGET != "hosted" |
| 201 | + else HOSTED_CORE_MODEL_URL |
| 202 | + ) |
| 203 | + client = InferenceHTTPClient( |
| 204 | + api_url=api_url, |
| 205 | + api_key=self._api_key, |
| 206 | + ) |
| 207 | + if WORKFLOWS_REMOTE_API_TARGET == "hosted": |
| 208 | + client.select_api_v0() |
| 209 | + |
| 210 | + combined_prompt = _combine_prompt(prompt=prompt, system_prompt=system_prompt) |
| 211 | + predictions = [] |
| 212 | + for image in images: |
| 213 | + result = client.infer_lmm( |
| 214 | + inference_input=image.base64_image, |
| 215 | + model_id=model_version, |
| 216 | + prompt=combined_prompt, |
| 217 | + model_id_in_path=True, |
| 218 | + ) |
| 219 | + predictions.append({"output": result.get("response", result)}) |
| 220 | + return predictions |
| 221 | + |
| 222 | + def run_locally( |
| 223 | + self, |
| 224 | + images: Batch[WorkflowImageData], |
| 225 | + model_version: str, |
| 226 | + prompt: Optional[str], |
| 227 | + system_prompt: Optional[str], |
| 228 | + ) -> BlockResult: |
| 229 | + inference_images = [ |
| 230 | + i.to_inference_format(numpy_preferred=False) for i in images |
| 231 | + ] |
| 232 | + combined_prompt = _combine_prompt(prompt=prompt, system_prompt=system_prompt) |
| 233 | + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) |
| 234 | + |
| 235 | + predictions = [] |
| 236 | + for image in inference_images: |
| 237 | + request = LMMInferenceRequest( |
| 238 | + api_key=self._api_key, |
| 239 | + model_id=model_version, |
| 240 | + image=image, |
| 241 | + source="workflow-execution", |
| 242 | + prompt=combined_prompt, |
| 243 | + ) |
| 244 | + prediction = self._model_manager.infer_from_request_sync( |
| 245 | + model_id=model_version, request=request |
| 246 | + ) |
| 247 | + predictions.append({"output": prediction.response}) |
| 248 | + return predictions |
| 249 | + |
| 250 | + |
| 251 | +def _combine_prompt(prompt: Optional[str], system_prompt: Optional[str]) -> str: |
| 252 | + prompt = prompt or DEFAULT_PROMPT |
| 253 | + system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT |
| 254 | + return prompt + "<system_prompt>" + system_prompt |
0 commit comments