model capabilities

aisbf/coderai_broker.py

@@ -143,8 +143,22 @@ class CoderAIBroker:
             result["gpus"] = normalized_gpus
             result["gpu_count"] = len(normalized_gpus)
             result["total_vram_mb"] = total_vram_mb
-            result["available_vram_mb"] = available_vram_mb
-            result["used_vram_mb"] = max(total_vram_mb - available_vram_mb, 0.0)
+            # Only overwrite the top-level available_vram_mb when the GPU loop
+            # actually produced data; otherwise keep the value that arrived at
+            # the top level of the metadata (e.g. coderai sends it there but
+            # does not repeat it inside each GPU dict).
+            if available_vram_mb > 0:
+                result["available_vram_mb"] = available_vram_mb
+            elif result.get("available_vram_mb") is None:
+                result["available_vram_mb"] = 0.0
+            # Propagate top-level free VRAM into the single GPU so the UI can
+            # read it from the per-GPU entry too.
+            if len(normalized_gpus) == 1 and normalized_gpus[0].get("available_vram_mb") is None:
+                top_free = result.get("available_vram_mb")
+                if top_free is not None:
+                    normalized_gpus[0]["available_vram_mb"] = float(top_free)
+            final_free = result.get("available_vram_mb") or 0.0
+            result["used_vram_mb"] = max(total_vram_mb - final_free, 0.0)
         return result
 
     @staticmethod

aisbf/models.py

@@ -76,6 +76,8 @@ class Model(BaseModel):
     top_provider: Optional[Dict] = None  # context_length, max_completion_tokens, is_moderated
     supported_parameters: Optional[List[str]] = None
     default_parameters: Optional[Dict] = None
+    type: Optional[str] = None  # e.g. "text", "image", "video", "audio"
+    capabilities: Optional[List[str]] = None  # provider-native capability list
 
 class Provider(BaseModel):
     id: str

aisbf/providers/coderai.py

@@ -417,6 +417,8 @@ class CoderAIProviderHandler(BaseProviderHandler):
                 supported_parameters=metadata.get("supported_parameters"),
                 default_parameters=metadata.get("default_parameters"),
                 description=metadata.get("description"),
+                type=metadata.get("type"),
+                capabilities=metadata.get("capabilities") or None,
             ))
         return result
 

aisbf/studio.py

@@ -64,6 +64,45 @@ STUDIO_CAPABILITY_MAP = {
     "segmentation": "segmentation",
     "3d_generation": "3d_generation",
     "animation": "animation",
+    # CoderAI / HuggingFace pipeline task names (long-form field names from ModelCapabilities)
+    "text_generation": "chat",
+    "text-generation": "chat",
+    "chat_completion": "chat",
+    "image_to_text": "vision",
+    "image_generation": "image_generation",
+    "text_to_image": "image_generation",
+    "image_to_image": "image_edit",
+    "inpainting": "image_edit",
+    "controlnet": "image_generation",
+    "video_generation": "video_generation",
+    "text_to_video": "video_generation",
+    "image_to_video": "video_generation",
+    "video_to_video": "video_generation",
+    "video_interpolation": "animation",
+    "video_upscaling": "video_generation",
+    "image_upscaling": "image_edit",
+    "face_restoration": "image_edit",
+    "depth_estimation": "object_detection",
+    "image_segmentation": "segmentation",
+    "image_classification": "classification",
+    "object_detection": "object_detection",
+    "style_transfer": "image_edit",
+    "automatic_speech_recognition": "transcription",
+    "speech_recognition": "transcription",
+    "speech_to_text": "transcription",
+    "subtitle_generation": "transcription",
+    "text_to_speech": "speech_generation",
+    "audio_generation": "audio_generation",
+    "lip_sync": "video_generation",
+    "video_dubbing": "video_generation",
+    "image_to_3d": "3d_generation",
+    "video_to_3d": "3d_generation",
+    "model_3d_generation": "3d_generation",
+    "model_3d_to_image": "3d_generation",
+    "feature_extraction": "embeddings",
+    "sentence_similarity": "embeddings",
+    "fill_mask": "chat",
+    "token_classification": "ner",
 }
 
 STUDIO_CAPABILITY_CHOICES = [
@@ -187,6 +226,16 @@ def infer_model_capabilities(
 
     provider_metadata = provider_metadata or {}
     capabilities = normalize_capabilities(provider_metadata.get("capabilities"))
+    if not capabilities and provider_metadata.get("type"):
+        _type_cap_map = {
+            "text": ["chat"],
+            "image": ["image_generation"],
+            "video": ["video_generation"],
+            "audio": ["audio_generation"],
+            "embedding": ["embeddings"],
+            "multimodal": ["multimodal", "chat"],
+        }
+        capabilities = _type_cap_map.get((provider_metadata["type"] or "").lower(), [])
     source = "provider_metadata" if capabilities else "heuristic"
     notes: List[str] = []