Support for returning multiple outputs

[andreasjansson]

This PR makes it possible to return more than one output from Cog, without having to wrap the outputs in a zip file.

Python API changes

The predict() function treats dict outputs as multiple outputs. This is a change in behavior from before, when dicts were returned as json.

You can mix pathlib.Paths and other non-Path types. The following is a valid response:

def predict(self):
    return {
        "file-output1": Path(path1),
        "file-output2": Path(path2),
        "text-output": "foobar",
        "list-output": [1, 2, 3],
    }

Server changes

The built-in HTTP server returns multiple outputs as multipart/form-data.

The Redis queue worker returns multiple outputs as a new multi_values field in the response struct. The response struct now is:

type InferResponse struct {
	Error       string        `json:"error"`
	File        *InferFile    `json:"file"`
	Value       string        `json:"value"`
	MultiValues []*MultiValue `json:"multi_values"`
	Status      string        `json:"status"`
}

type MultiValue struct {
	Name  string     `json:"name"`
	File  *InferFile `json:"file"`
	Value string     `json:"value"`
}

type InferFile struct {
	URL  string `json:"url"`
	Name string `json:"name"`
}

UI changes

cog predict can write multiple outputs to stdout if the outputs are non-Path or textual Path outputs.

For example, given the following predict.py:

    @cog.input("input", type=str)
    def predict(self, input):
        out_dir = Path(tempfile.mkdtemp())
        out_path1 = out_dir / "first.txt"
        out_path2 = out_dir / "second.txt"
        with open(out_path1, "w") as f:
            f.write(f"one {input}")
        with open(out_path2, "w") as f:
            f.write(f"two {input}")
        
        return {"first": out_path1, "second": out_path2, "third": "just text"}

cog predict -i input=hello outputs:

first:
one hello
second:
two hello
third:
just text

Non-textual Path outputs are written to files named after the key in the response dict.

If the user passes the -o flag to cog predict, the outputs get named <output-name>.<dict-key>.<extension>.

---

Signed-off-by: andreasjansson [email protected]

[bfirsh]

This is now being incorporated into #407.