diff --git a/inference_gateway/__init__.py b/inference_gateway/__init__.py
index 3895afd..ddc9934 100644
--- a/inference_gateway/__init__.py
+++ b/inference_gateway/__init__.py
@@ -58,7 +58,9 @@
     MessagesToolResultBlock,
     MessagesToolUseBlock,
     MessagesUsage,
+    Modality,
     Model,
+    ModelModalities,
     Pricing,
     Provider,
     Response,
@@ -108,6 +110,8 @@
     "MessageContent",
     "MessageRole",
     "ListModelsResponse",
+    "Modality",
+    "ModelModalities",
     "CreateChatCompletionRequest",
     "CreateChatCompletionResponse",
     "CreateChatCompletionStreamResponse",
diff --git a/inference_gateway/models.py b/inference_gateway/models.py
index 856df78..dc3715b 100644
--- a/inference_gateway/models.py
+++ b/inference_gateway/models.py
@@ -367,44 +367,23 @@ class Pricing(BaseModel):
     """
 
 
-class Model(BaseModel):
-    """
-    Common model information
-    """
-
-    model_config = ConfigDict(
-        populate_by_name=True,
-    )
-    id: str
-    object: str
-    created: int
-    owned_by: str
-    served_by: Provider
-    modalities: Sequence[Literal["text", "image", "audio", "video"]] | None = None
-    """
-    The modalities the model supports natively (included when `include=modalities`)
-    """
-    context_window: ContextWindow | None = None
+class Modality(RootModel[Literal["text", "image", "audio", "video"]]):
+    root: Literal["text", "image", "audio", "video"]
     """
-    Context window information for the model (included when `include=context_window`)
-    """
-    pricing: Pricing | None = None
-    """
-    Pricing information for the model (included when `include=pricing`)
+    A single input or output modality
     """
 
 
-class ListModelsResponse(BaseModel):
+class ModelModalities(BaseModel):
     """
-    Response structure for listing models
+    The input and output modalities of a model, mirroring the models.dev dataset shape. Vision models accept `image` in `input`; image-generation models list `image` in `output` — when `output` carries `image` but not `text`, the model only generates images and cannot chat.
     """
 
     model_config = ConfigDict(
         populate_by_name=True,
     )
-    provider: Provider | None = None
-    object: str
-    data: Annotated[Sequence[Model], Field(validate_default=True)] = []
+    input: Sequence[Modality]
+    output: Sequence[Modality]
 
 
 class MCPTool(BaseModel):
@@ -1736,6 +1715,46 @@ class ImageContentPart(BaseModel):
     image_url: ImageURL
 
 
+class Model(BaseModel):
+    """
+    Common model information
+    """
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    id: str
+    object: str
+    created: int
+    owned_by: str
+    served_by: Provider
+    modalities: ModelModalities | None = None
+    """
+    The input and output modalities of the model (included when `include=modalities`)
+    """
+    context_window: ContextWindow | None = None
+    """
+    Context window information for the model (included when `include=context_window`)
+    """
+    pricing: Pricing | None = None
+    """
+    Pricing information for the model (included when `include=pricing`)
+    """
+
+
+class ListModelsResponse(BaseModel):
+    """
+    Response structure for listing models
+    """
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+    )
+    provider: Provider | None = None
+    object: str
+    data: Annotated[Sequence[Model], Field(validate_default=True)] = []
+
+
 class ListToolsResponse(BaseModel):
     """
     Response structure for listing MCP tools
diff --git a/openapi.yaml b/openapi.yaml
index f5c083f..32e0397 100644
--- a/openapi.yaml
+++ b/openapi.yaml
@@ -256,16 +256,22 @@ paths:
                         owned_by: 'openai'
                         served_by: 'openai'
                         modalities:
-                          - text
-                          - image
-                      - id: 'openai/gpt-4-turbo'
+                          input:
+                            - text
+                            - image
+                          output:
+                            - text
+                      - id: 'openai/gpt-image-2'
                         object: 'model'
                         created: 1687882410
                         owned_by: 'openai'
                         served_by: 'openai'
                         modalities:
-                          - text
-                          - image
+                          input:
+                            - text
+                            - image
+                          output:
+                            - image
         '400':
           description: Bad request - unsupported include value
           content:
@@ -1630,16 +1636,9 @@ components:
           $ref: '#/components/schemas/Provider'
         modalities:
           oneOf:
-            - type: array
-              items:
-                type: string
-                enum:
-                  - text
-                  - image
-                  - audio
-                  - video
+            - $ref: '#/components/schemas/ModelModalities'
             - type: 'null'
-          description: The modalities the model supports natively (included when `include=modalities`)
+          description: The input and output modalities of the model (included when `include=modalities`)
         context_window:
           oneOf:
             - $ref: '#/components/schemas/ContextWindow'
@@ -1656,6 +1655,33 @@ components:
         - created
         - owned_by
         - served_by
+    Modality:
+      type: string
+      description: A single input or output modality
+      enum:
+        - text
+        - image
+        - audio
+        - video
+    ModelModalities:
+      type: object
+      description: >-
+        The input and output modalities of a model, mirroring the models.dev dataset shape.
+        Vision models accept `image` in `input`; image-generation models list `image` in
+        `output` — when `output` carries `image` but not `text`, the model only generates
+        images and cannot chat.
+      properties:
+        input:
+          type: array
+          items:
+            $ref: '#/components/schemas/Modality'
+        output:
+          type: array
+          items:
+            $ref: '#/components/schemas/Modality'
+      required:
+        - input
+        - output
     ListModelsResponse:
       type: object
       description: Response structure for listing models