Add x-cog-accept MIME type annotation for Path/File input fields (#2863)

* Add x-cog-accept MIME type annotation for Path/File input fields

Add an 'accept' parameter to Input() that specifies allowed MIME types
or file extensions for Path/File inputs. The Go static schema generator
extracts this and emits it as 'x-cog-accept' in the OpenAPI schema,
giving schema consumers (UIs, validators, API clients) visibility into
what file types an input expects.

Usage: Input(accept="image/*"), Input(accept="audio/wav,audio/mp3"),
or Input(accept=".safetensors,.bin"). Using accept on non-Path/File
types is a hard build error.

* Set COG_STATIC_SCHEMA=1 in accept integration tests

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Anish Sahoo <anishsahoo2005@gmail.com>

Author: 3 people

integration-tests/tests/accept_mime_type.txtar

@@ -0,0 +1,57 @@
+# Test that the accept parameter on Path/File inputs produces the
+# x-cog-accept annotation in the generated OpenAPI schema.
+#
+# Verifies:
+# - accept="image/*" on a Path input emits x-cog-accept in the schema
+# - accept with multiple MIME types works
+# - accept with file extensions works
+# - Fields without accept do not have x-cog-accept
+# - Prediction still works end-to-end
+
+env COG_STATIC_SCHEMA=1
+
+cog build -t $TEST_IMAGE
+
+# Extract the schema from the image label
+exec docker inspect $TEST_IMAGE --format '{{index .Config.Labels "run.cog.openapi_schema"}}'
+
+# x-cog-accept annotations are present
+stdout '"x-cog-accept":"image/\*"'
+stdout '"x-cog-accept":"audio/wav,audio/mp3"'
+stdout '"x-cog-accept":".safetensors,.bin"'
+
+# The prompt field (str) should NOT have x-cog-accept
+# (we check the schema has prompt but confirm no extra x-cog-accept entries)
+stdout '"prompt":'
+
+# Path fields still have uri format
+stdout '"format":"uri"'
+
+# Prediction works end-to-end
+cog predict $TEST_IMAGE -i prompt=hello -i image=@test.png
+stdout 'hello-png'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from typing import Optional
+
+from cog import BasePredictor, Input, Path
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self,
+        prompt: str = Input(description="Text prompt", default="test"),
+        image: Path = Input(description="Input image", accept="image/*"),
+        audio: Optional[Path] = Input(description="Audio clip", accept="audio/wav,audio/mp3", default=None),
+        weights: Optional[Path] = Input(description="Model weights", accept=".safetensors,.bin", default=None),
+    ) -> str:
+        ext = str(image).split(".")[-1]
+        return f"{prompt}-{ext}"
+
+-- test.png --
+fake image content

integration-tests/tests/accept_mime_type_error.txtar

@@ -0,0 +1,25 @@
+# Test that using accept on a non-Path/File input type causes a build error.
+#
+# The accept parameter is only valid on Path or File inputs. Using it on
+# str, int, float, etc. should produce a clear error at build time.
+
+env COG_STATIC_SCHEMA=1
+
+! cog build -t $TEST_IMAGE
+stderr 'accept is only valid on Path or File inputs'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self,
+        name: str = Input(description="User name", accept="text/plain"),
+    ) -> str:
+        return f"hello {name}"

pkg/schema/errors.go

@@ -29,6 +29,7 @@ const (
 	ErrChoicesNotResolvable
 	ErrDefaultNotResolvable
 	ErrUnresolvableType
+	ErrAcceptOnNonFileType
 	ErrOther
 )
 

pkg/schema/openapi.go

@@ -438,6 +438,11 @@ func buildInputSchema(info *PredictorInfo) (map[string]any, []enumSchema) {
 			prop["deprecated"] = true
 		}
 
+		// MIME type constraint for Path/File inputs
+		if field.Accept != nil {
+			prop["x-cog-accept"] = *field.Accept
+		}
+
 		properties.Set(name, prop)
 	})
 

pkg/schema/openapi_test.go

@@ -1007,6 +1007,80 @@ func TestMultipleInputTypes(t *testing.T) {
 	assert.NotContains(t, required, "secret_key")
 }
 
+// ---------------------------------------------------------------------------
+// Tests: Accept (MIME type) annotation
+// ---------------------------------------------------------------------------
+
+func TestAcceptAnnotation(t *testing.T) {
+	accept := "image/*"
+	inputs := NewOrderedMap[string, InputField]()
+	inputs.Set("image", InputField{
+		Name:      "image",
+		Order:     0,
+		FieldType: FieldType{Primitive: TypePath, Repetition: Required},
+		Accept:    &accept,
+	})
+
+	info := &PredictorInfo{
+		Inputs: inputs,
+		Output: SchemaPrim(TypeString),
+		Mode:   ModePredict,
+	}
+
+	spec := parseSpec(t, info)
+	props := getPath(spec, "components", "schemas", "Input", "properties").(map[string]any)
+
+	imageField := props["image"].(map[string]any)
+	assert.Equal(t, "string", imageField["type"])
+	assert.Equal(t, "uri", imageField["format"])
+	assert.Equal(t, "image/*", imageField["x-cog-accept"])
+}
+
+func TestAcceptAnnotationMultipleMimeTypes(t *testing.T) {
+	accept := "audio/wav,audio/mp3,audio/flac"
+	inputs := NewOrderedMap[string, InputField]()
+	inputs.Set("audio", InputField{
+		Name:      "audio",
+		Order:     0,
+		FieldType: FieldType{Primitive: TypePath, Repetition: Required},
+		Accept:    &accept,
+	})
+
+	info := &PredictorInfo{
+		Inputs: inputs,
+		Output: SchemaPrim(TypeString),
+		Mode:   ModePredict,
+	}
+
+	spec := parseSpec(t, info)
+	props := getPath(spec, "components", "schemas", "Input", "properties").(map[string]any)
+
+	audioField := props["audio"].(map[string]any)
+	assert.Equal(t, "audio/wav,audio/mp3,audio/flac", audioField["x-cog-accept"])
+}
+
+func TestAcceptAnnotationNotPresentWhenNil(t *testing.T) {
+	inputs := NewOrderedMap[string, InputField]()
+	inputs.Set("image", InputField{
+		Name:      "image",
+		Order:     0,
+		FieldType: FieldType{Primitive: TypePath, Repetition: Required},
+	})
+
+	info := &PredictorInfo{
+		Inputs: inputs,
+		Output: SchemaPrim(TypeString),
+		Mode:   ModePredict,
+	}
+
+	spec := parseSpec(t, info)
+	props := getPath(spec, "components", "schemas", "Input", "properties").(map[string]any)
+
+	imageField := props["image"].(map[string]any)
+	_, hasAccept := imageField["x-cog-accept"]
+	assert.False(t, hasAccept, "x-cog-accept should not be present when Accept is nil")
+}
+
 // ---------------------------------------------------------------------------
 // Tests: Edge cases
 // ---------------------------------------------------------------------------

pkg/schema/python/inputs.go

@@ -18,6 +18,7 @@ type inputCallInfo struct {
 	Regex       *string
 	Choices     []schema.DefaultValue
 	Deprecated  *bool
+	Accept      *string
 }
 
 type inputMethodInfo struct {
@@ -271,7 +272,11 @@ func inputField(name string, order int, inputType schema.InputType, fieldType sc
 	}
 }
 
-func inputFieldWithInfo(name string, order int, inputType schema.InputType, fieldType schema.FieldType, info inputCallInfo) schema.InputField {
+func inputFieldWithInfo(name string, order int, inputType schema.InputType, fieldType schema.FieldType, info inputCallInfo) (schema.InputField, error) {
+	if info.Accept != nil && fieldType.Primitive != schema.TypePath && fieldType.Primitive != schema.TypeFile {
+		return schema.InputField{}, schema.NewError(schema.ErrAcceptOnNonFileType,
+			fmt.Sprintf("accept is only valid on Path or File inputs (parameter '%s')", name))
+	}
 	field := inputField(name, order, inputType, fieldType)
 	field.Default = info.Default
 	field.Description = info.Description
@@ -282,7 +287,8 @@ func inputFieldWithInfo(name string, order int, inputType schema.InputType, fiel
 	field.Regex = info.Regex
 	field.Choices = info.Choices
 	field.Deprecated = info.Deprecated
-	return field
+	field.Accept = info.Accept
+	return field, nil
 }
 
 func firstParamIsSelf(params *sitter.Node, source []byte) bool {
@@ -398,13 +404,19 @@ func parseTypedDefaultParameter(node *sitter.Node, source []byte, order int, ctx
 			if err != nil {
 				return schema.InputField{}, err
 			}
-			field := inputFieldWithInfo(name, order, inputType, fieldType, info)
+			field, err := inputFieldWithInfo(name, order, inputType, fieldType, info)
+			if err != nil {
+				return schema.InputField{}, err
+			}
 			return field, schema.ValidateInputField(field)
 		}
 
 		// 2. Reference to Input() via class attribute or static method
 		if info, ok := resolveInputReference(valNode, source, ctx.registry); ok {
-			field := inputFieldWithInfo(name, order, inputType, fieldType, info)
+			field, err := inputFieldWithInfo(name, order, inputType, fieldType, info)
+			if err != nil {
+				return schema.InputField{}, err
+			}
 			return field, schema.ValidateInputField(field)
 		}
 
@@ -511,6 +523,10 @@ func parseInputCall(node *sitter.Node, source []byte, paramName string, scope mo
 			if b, ok := parseBoolLiteral(valNode, source); ok {
 				info.Deprecated = &b
 			}
+		case "accept":
+			if s, ok := parseStringLiteral(valNode, source); ok {
+				info.Accept = &s
+			}
 		}
 	}
 

pkg/schema/python/parser_test.go

@@ -1653,6 +1653,97 @@ class Predictor(BasePredictor):
 	require.True(t, *old.Deprecated)
 }
 
+// ---------------------------------------------------------------------------
+// Accept MIME type
+// ---------------------------------------------------------------------------
+
+func TestAcceptMimeType(t *testing.T) {
+	source := `
+from cog import BasePredictor, Input, Path
+
+class Predictor(BasePredictor):
+    def predict(self, image: Path = Input(description="An image", accept="image/*")) -> str:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	image, ok := info.Inputs.Get("image")
+	require.True(t, ok)
+	require.NotNil(t, image.Accept)
+	require.Equal(t, "image/*", *image.Accept)
+}
+
+func TestAcceptMultipleMimeTypes(t *testing.T) {
+	source := `
+from cog import BasePredictor, Input, Path
+
+class Predictor(BasePredictor):
+    def predict(self, audio: Path = Input(accept="audio/wav,audio/mp3")) -> str:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	audio, ok := info.Inputs.Get("audio")
+	require.True(t, ok)
+	require.NotNil(t, audio.Accept)
+	require.Equal(t, "audio/wav,audio/mp3", *audio.Accept)
+}
+
+func TestAcceptFileExtensions(t *testing.T) {
+	source := `
+from cog import BasePredictor, Input, Path
+
+class Predictor(BasePredictor):
+    def predict(self, weights: Path = Input(accept=".safetensors,.bin")) -> str:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	weights, ok := info.Inputs.Get("weights")
+	require.True(t, ok)
+	require.NotNil(t, weights.Accept)
+	require.Equal(t, ".safetensors,.bin", *weights.Accept)
+}
+
+func TestAcceptOnFileType(t *testing.T) {
+	source := `
+from cog import BasePredictor, Input, File
+
+class Predictor(BasePredictor):
+    def predict(self, f: File = Input(accept="image/png")) -> str:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	f, ok := info.Inputs.Get("f")
+	require.True(t, ok)
+	require.NotNil(t, f.Accept)
+	require.Equal(t, "image/png", *f.Accept)
+}
+
+func TestAcceptOnNonFileTypeErrors(t *testing.T) {
+	source := `
+from cog import BasePredictor, Input
+
+class Predictor(BasePredictor):
+    def predict(self, name: str = Input(accept="image/*")) -> str:
+        pass
+`
+	se := parseErr(t, source, "Predictor", schema.ModePredict)
+	require.Equal(t, schema.ErrAcceptOnNonFileType, se.Kind)
+	require.Contains(t, se.Error(), "name")
+}
+
+func TestAcceptNotSetWhenOmitted(t *testing.T) {
+	source := `
+from cog import BasePredictor, Input, Path
+
+class Predictor(BasePredictor):
+    def predict(self, image: Path = Input(description="An image")) -> str:
+        pass
+`
+	info := parse(t, source, "Predictor")
+	image, ok := info.Inputs.Get("image")
+	require.True(t, ok)
+	require.Nil(t, image.Accept)
+}
+
 // ---------------------------------------------------------------------------
 // File type (deprecated alias for Path)
 // ---------------------------------------------------------------------------

pkg/schema/types.go

@@ -231,6 +231,7 @@ type InputField struct {
 	Regex       *string
 	Choices     []DefaultValue
 	Deprecated  *bool
+	Accept      *string // MIME types / file extensions for Path/File inputs (e.g. "image/*")
 }
 
 // IsRequired returns true if this field is required in the schema.

python/cog/input.py

@@ -27,6 +27,7 @@ class FieldInfo:
     regex: Optional[str] = None
     choices: Optional[List[Union[str, int]]] = None
     deprecated: Optional[bool] = None
+    accept: Optional[str] = None
 
 
 def Input(
@@ -41,6 +42,7 @@ def Input(
     regex: Optional[str] = None,
     choices: Optional[List[Union[str, int]]] = None,
     deprecated: Optional[bool] = None,
+    accept: Optional[str] = None,
 ) -> Any:
     """
     Create an input field specification for a predictor parameter.
@@ -70,6 +72,9 @@ def predict(
         regex: Regular expression pattern for string inputs.
         choices: List of allowed values.
         deprecated: Whether the input is deprecated.
+        accept: Allowed MIME types or file extensions for Path/File inputs,
+            using the same format as the HTML accept attribute (e.g.
+            ``"image/*"``, ``"image/png,image/jpeg"``, ``".safetensors,.bin"``).
 
     Returns:
         A FieldInfo instance containing the field metadata.
@@ -92,4 +97,5 @@ def predict(
         regex=regex,
         choices=choices,
         deprecated=deprecated,
+        accept=accept,
     )