diff --git a/pkg/apis/ome/v1beta1/model.go b/pkg/apis/ome/v1beta1/model.go index d16a78aba..8dfc197b0 100644 --- a/pkg/apis/ome/v1beta1/model.go +++ b/pkg/apis/ome/v1beta1/model.go @@ -286,7 +286,7 @@ const ( // ModelCapability enum // TODO: Remove legacy capabilities // -// +kubebuilder:validation:Enum=TEXT_GENERATION;TEXT_SUMMARIZATION;TEXT_EMBEDDINGS;TEXT_RERANK;CHAT;VISION;EMBEDDING;RERANK;TEXT_TO_TEXT;TEXT_TO_AUDIO;TEXT_TO_IMAGE;TEXT_TO_VIDEO;IMAGE_TEXT_TO_TEXT;IMAGE_TEXT_TO_AUDIO;IMAGE_TEXT_TO_IMAGE;IMAGE_TEXT_TO_VIDEO;VIDEO_TEXT_TO_AUDIO;AUDIO_TO_TEXT;AUDIO_TO_AUDIO;AUDIO_TEXT_TO_TEXT;AUDIO_TRANSLATION +// +kubebuilder:validation:Enum=TEXT_GENERATION;TEXT_SUMMARIZATION;TEXT_EMBEDDINGS;TEXT_RERANK;CHAT;VISION;EMBEDDING;RERANK;TEXT_TO_TEXT;TEXT_TO_AUDIO;TEXT_TO_IMAGE;TEXT_TO_VIDEO;IMAGE_TEXT_TO_TEXT;IMAGE_TEXT_TO_AUDIO;IMAGE_TEXT_TO_IMAGE;IMAGE_TEXT_TO_VIDEO;VIDEO_TEXT_TO_AUDIO;VIDEO_TEXT_TO_TEXT;AUDIO_TO_TEXT;AUDIO_TO_AUDIO;AUDIO_TEXT_TO_TEXT;AUDIO_TRANSLATION type ModelCapability string const ( @@ -310,6 +310,7 @@ const ( ModelCapabilityImageTextToImage ModelCapability = "IMAGE_TEXT_TO_IMAGE" ModelCapabilityImageTextToVideo ModelCapability = "IMAGE_TEXT_TO_VIDEO" ModelCapabilityVideoTextToAudio ModelCapability = "VIDEO_TEXT_TO_AUDIO" + ModelCapabilityVideoTextToText ModelCapability = "VIDEO_TEXT_TO_TEXT" ModelCapabilityAudioToText ModelCapability = "AUDIO_TO_TEXT" ModelCapabilityAudioToAudio ModelCapability = "AUDIO_TO_AUDIO" ModelCapabilityAudioTextToText ModelCapability = "AUDIO_TEXT_TO_TEXT" diff --git a/pkg/modelconfig/capability.go b/pkg/modelconfig/capability.go index bf0bdff39..ca72ecf5b 100644 --- a/pkg/modelconfig/capability.go +++ b/pkg/modelconfig/capability.go @@ -21,6 +21,7 @@ const ( CapabilityTextToAudio Capability = "TEXT_TO_AUDIO" CapabilityImageTextToAudio Capability = "IMAGE_TEXT_TO_AUDIO" CapabilityVideoTextToAudio Capability = "VIDEO_TEXT_TO_AUDIO" + CapabilityVideoTextToText Capability = "VIDEO_TEXT_TO_TEXT" CapabilityAudioToText Capability = "AUDIO_TO_TEXT" CapabilityAudioToAudio Capability = "AUDIO_TO_AUDIO" CapabilityAudioTextToText Capability = "AUDIO_TEXT_TO_TEXT" @@ -166,14 +167,21 @@ func diffusionRule(hf HuggingFaceModel) []Capability { } func nemotronHNanoRule(hf HuggingFaceModel) []Capability { - if !strings.Contains(strings.ToLower(hf.GetModelType()), "nemotronh_nano") { + modelType := strings.ToLower(hf.GetModelType()) + if !strings.Contains(modelType, "nemotronh_nano") { return nil } - return []Capability{ + capabilities := []Capability{ CapabilityImageTextToText, CapabilityTextToText, - CapabilityAudioToText, } + if strings.Contains(modelType, "nemotronh_nano_omni") { + capabilities = append(capabilities, + CapabilityAudioTextToText, + CapabilityVideoTextToText, + ) + } + return capabilities } func omniRule(hf HuggingFaceModel) []Capability { diff --git a/pkg/modelconfig/capability_test.go b/pkg/modelconfig/capability_test.go index 34c20839d..9d6881fd5 100644 --- a/pkg/modelconfig/capability_test.go +++ b/pkg/modelconfig/capability_test.go @@ -198,7 +198,23 @@ func TestClassifyCapabilities_NemotronH_Nano(t *testing.T) { want := []Capability{ CapabilityImageTextToText, CapabilityTextToText, - CapabilityAudioToText, + CapabilityAudioTextToText, + CapabilityVideoTextToText, + } + if got := classifyCapabilities(m); !equalCaps(got, want) { + t.Errorf("got %v, want %v", got, want) + } +} + +func TestClassifyCapabilities_NemotronH_NanoNonOmni(t *testing.T) { + m := &stubModel{ + modelType: "NemotronH_Nano_VL_V2", + architecture: "NemotronH_Nano_VL_V2", + hasVision: true, + } + want := []Capability{ + CapabilityImageTextToText, + CapabilityTextToText, } if got := classifyCapabilities(m); !equalCaps(got, want) { t.Errorf("got %v, want %v", got, want) @@ -413,7 +429,8 @@ func TestClassifyCapabilities_RegressionSet(t *testing.T) { want: []Capability{ CapabilityImageTextToText, CapabilityTextToText, - CapabilityAudioToText, + CapabilityAudioTextToText, + CapabilityVideoTextToText, }, }, { diff --git a/pkg/modelparser/capability.go b/pkg/modelparser/capability.go index 2ed774599..98b2c9ac0 100644 --- a/pkg/modelparser/capability.go +++ b/pkg/modelparser/capability.go @@ -23,6 +23,7 @@ var capabilityToOME = map[modelconfig.Capability]v1beta1.ModelCapability{ modelconfig.CapabilityTextToAudio: v1beta1.ModelCapabilityTextToAudio, modelconfig.CapabilityImageTextToAudio: v1beta1.ModelCapabilityImageTextToAudio, modelconfig.CapabilityVideoTextToAudio: v1beta1.ModelCapabilityVideoTextToAudio, + modelconfig.CapabilityVideoTextToText: v1beta1.ModelCapabilityVideoTextToText, modelconfig.CapabilityAudioToText: v1beta1.ModelCapabilityAudioToText, modelconfig.CapabilityAudioToAudio: v1beta1.ModelCapabilityAudioToAudio, modelconfig.CapabilityAudioTextToText: v1beta1.ModelCapabilityAudioTextToText, diff --git a/pkg/modelparser/capability_test.go b/pkg/modelparser/capability_test.go index 8b1eb64ef..250b1e0dd 100644 --- a/pkg/modelparser/capability_test.go +++ b/pkg/modelparser/capability_test.go @@ -29,6 +29,7 @@ func TestCapabilityToOME(t *testing.T) { {modelconfig.CapabilityTextToAudio, v1beta1.ModelCapabilityTextToAudio}, {modelconfig.CapabilityImageTextToAudio, v1beta1.ModelCapabilityImageTextToAudio}, {modelconfig.CapabilityVideoTextToAudio, v1beta1.ModelCapabilityVideoTextToAudio}, + {modelconfig.CapabilityVideoTextToText, v1beta1.ModelCapabilityVideoTextToText}, {modelconfig.CapabilityAudioToText, v1beta1.ModelCapabilityAudioToText}, {modelconfig.CapabilityAudioToAudio, v1beta1.ModelCapabilityAudioToAudio}, {modelconfig.CapabilityAudioTextToText, v1beta1.ModelCapabilityAudioTextToText},