Handle badly formatted data URIs in cog predict (#2409)

* Currently cog returns data:None;base64
* The data URI parser does not recognise None, and
requires data:;base64
* Handle receiving these URIs by removing the None
format specified

Author: 8W9aG

pkg/cli/predict.go

@@ -19,7 +19,6 @@ import (
 	"github.com/getkin/kin-openapi/openapi3"
 	"github.com/mitchellh/go-homedir"
 	"github.com/spf13/cobra"
-	"github.com/vincent-petithory/dataurl"
 	"golang.org/x/sys/unix"
 
 	"github.com/replicate/cog/pkg/config"
@@ -30,6 +29,7 @@ import (
 	"github.com/replicate/cog/pkg/predict"
 	"github.com/replicate/cog/pkg/registry"
 	"github.com/replicate/cog/pkg/util/console"
+	"github.com/replicate/cog/pkg/util/files"
 	"github.com/replicate/cog/pkg/util/mime"
 )
 
@@ -443,7 +443,7 @@ func runPrediction(predictor predict.Predictor, inputs predict.Inputs, outputPat
 		}
 
 		if writeOutputToDisk {
-			path, err := writeFile(indentedJSON.Bytes(), outputPath)
+			path, err := files.WriteFile(indentedJSON.Bytes(), outputPath)
 			if err != nil {
 				return fmt.Errorf("Failed to write output: %w", err)
 			}
@@ -485,7 +485,7 @@ func runPrediction(predictor predict.Predictor, inputs predict.Inputs, outputPat
 		}
 
 		if writeOutputToDisk {
-			path, err := writeFile([]byte(s), outputPath)
+			path, err := files.WriteFile([]byte(s), outputPath)
 			if err != nil {
 				return fmt.Errorf("Failed to write output: %w", err)
 			}
@@ -505,7 +505,7 @@ func runPrediction(predictor predict.Predictor, inputs predict.Inputs, outputPat
 
 		// No special handling for needsJSON here.
 		if writeOutputToDisk {
-			path, err := writeFile(output, outputPath)
+			path, err := files.WriteFile(output, outputPath)
 			if err != nil {
 				return fmt.Errorf("Failed to write output: %w", err)
 			}
@@ -585,7 +585,7 @@ func processFileOutputs(output any, schema *openapi3.Schema, destination string)
 			return nil, fmt.Errorf("Failed to convert prediction output to string: %v", output)
 		}
 
-		path, err := writeDataURLToFile(outputStr, destination)
+		path, err := files.WriteDataURLToFile(outputStr, destination)
 		if err != nil {
 			return nil, fmt.Errorf("Failed to write output: %w", err)
 		}
@@ -615,58 +615,6 @@ func processFileOutputs(output any, schema *openapi3.Schema, destination string)
 	return output, nil
 }
 
-func writeFile(output []byte, outputPath string) (string, error) {
-	outputPath, err := homedir.Expand(outputPath)
-	if err != nil {
-		return "", err
-	}
-
-	// Write to file
-	outFile, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
-	if err != nil {
-		return "", err
-	}
-
-	if _, err := outFile.Write(output); err != nil {
-		return "", err
-	}
-	if err := outFile.Close(); err != nil {
-		return "", err
-	}
-	return outputPath, nil
-}
-
-// Writes a data URL to the destination. If no file extension is provided then it
-// will be inferred from the data URL mime type and appended.
-func writeDataURLToFile(url string, destination string) (string, error) {
-	dataurlObj, err := dataurl.DecodeString(url)
-	if err != nil {
-		return "", fmt.Errorf("Failed to decode data URL: %w", err)
-	}
-	output := dataurlObj.Data
-
-	ext := path.Ext(destination)
-	dir := path.Dir(destination)
-	name := r8_path.TrimExt(path.Base(destination))
-
-	// Check if ext is an integer, in which case ignore it...
-	if r8_path.IsExtInteger(ext) {
-		ext = ""
-		name = path.Base(destination)
-	}
-
-	if ext == "" {
-		ext = mime.ExtensionByType(dataurlObj.ContentType())
-	}
-
-	path, err := writeFile(output, path.Join(dir, name+ext))
-	if err != nil {
-		return "", err
-	}
-
-	return path, nil
-}
-
 func parseInputFlags(inputs []string, schema *openapi3.T) (predict.Inputs, error) {
 	keyVals := map[string][]string{}
 	for _, input := range inputs {

pkg/util/files/files.go

@@ -5,8 +5,15 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"path"
+	"strings"
 
+	"github.com/mitchellh/go-homedir"
+	"github.com/vincent-petithory/dataurl"
 	"golang.org/x/sys/unix"
+
+	r8_path "github.com/replicate/cog/pkg/path"
+	"github.com/replicate/cog/pkg/util/mime"
 )
 
 func Exists(path string) (bool, error) {
@@ -82,3 +89,56 @@ func WriteIfDifferent(file, content string) error {
 	}
 	return nil
 }
+
+func WriteDataURLToFile(url string, destination string) (string, error) {
+	if strings.HasPrefix(url, "data:None;base64") {
+		url = strings.Replace(url, "data:None;base64", "data:;base64", 1)
+	}
+	dataurlObj, err := dataurl.DecodeString(url)
+	if err != nil {
+		return "", fmt.Errorf("Failed to decode data URL: %w", err)
+	}
+	output := dataurlObj.Data
+
+	ext := path.Ext(destination)
+	dir := path.Dir(destination)
+	name := r8_path.TrimExt(path.Base(destination))
+
+	// Check if ext is an integer, in which case ignore it...
+	if r8_path.IsExtInteger(ext) {
+		ext = ""
+		name = path.Base(destination)
+	}
+
+	if ext == "" {
+		ext = mime.ExtensionByType(dataurlObj.ContentType())
+	}
+
+	path, err := WriteFile(output, path.Join(dir, name+ext))
+	if err != nil {
+		return "", err
+	}
+
+	return path, nil
+}
+
+func WriteFile(output []byte, outputPath string) (string, error) {
+	outputPath, err := homedir.Expand(outputPath)
+	if err != nil {
+		return "", err
+	}
+
+	// Write to file
+	outFile, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
+	if err != nil {
+		return "", err
+	}
+
+	if _, err := outFile.Write(output); err != nil {
+		return "", err
+	}
+	if err := outFile.Close(); err != nil {
+		return "", err
+	}
+	return outputPath, nil
+}

pkg/util/files/files_test.go

@@ -18,3 +18,10 @@ func TestIsExecutable(t *testing.T) {
 	require.NoError(t, os.Chmod(path, 0o744))
 	require.True(t, IsExecutable(path))
 }
+
+func TestWriteBadlyFormattedBase64DataURI(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test-file")
+	_, err := WriteDataURLToFile("data:None;base64,SGVsbG8gVGhlcmU=", path)
+	require.NoError(t, err)
+}

test-integration/test_integration/fixtures/glb-project/cog.yaml

@@ -0,0 +1,4 @@
+build:
+  python_version: "3.13"
+
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/glb-project/mesh.glb

@@ -0,0 +1,17 @@
+# Prediction interface for Cog ⚙️
+# https://github.com/replicate/cog/blob/main/docs/python.md
+
+from cog import BasePredictor, Path
+
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+        if not Path("mesh.glb").exists():
+            raise ValueError("Example file mesh.glb does not exist")
+
+    def predict(
+        self,
+    ) -> Path:
+        """Run a single prediction on the model"""
+        return Path("mesh.glb")

test-integration/test_integration/fixtures/glb-project/output.txt

@@ -769,3 +769,21 @@ def test_predict_json_input_stdin_dash(cog_binary):
 }
 """
     )
+
+
+def test_predict_glb_file(cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/glb-project"
+
+    result = subprocess.run(
+        [
+            cog_binary,
+            "predict",
+            "--debug",
+        ],
+        cwd=project_dir,
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=120.0,
+    )
+    assert result.returncode == 0