better snippet for image-to-image

Author: Wauplin

packages/tasks/src/snippets/inputs.ts

@@ -86,6 +86,11 @@ const inputsImageClassification = () => `"cats.jpg"`;
 
 const inputsImageToText = () => `"cats.jpg"`;
 
+const inputsImageToImage = () => `{
+	"image": "cat.png",
+	"prompt": "Turn the cat into a tiger."
+}`;
+
 const inputsImageSegmentation = () => `"cats.jpg"`;
 
 const inputsObjectDetection = () => `"cats.jpg"`;
@@ -118,6 +123,7 @@ const modelInputSnippets: {
 	"fill-mask": inputsFillMask,
 	"image-classification": inputsImageClassification,
 	"image-to-text": inputsImageToText,
+	"image-to-image": inputsImageToImage,
 	"image-segmentation": inputsImageSegmentation,
 	"object-detection": inputsObjectDetection,
 	"question-answering": inputsQuestionAnswering,

packages/tasks/src/snippets/python.spec.ts

@@ -175,6 +175,48 @@ output = query({
 })`);
 	});
 
+	it("image-to-image", async () => {
+		const model: ModelDataMinimal = {
+			id: "stabilityai/stable-diffusion-xl-refiner-1.0",
+			pipeline_tag: "image-to-image",
+			tags: [],
+			inference: "",
+		};
+		const snippets = getPythonInferenceSnippet(model, "api_token") as InferenceSnippet[];
+
+		expect(snippets.length).toEqual(2);
+
+		expect(snippets[0].client).toEqual("huggingface_hub");
+		expect(snippets[0].content).toEqual(`from huggingface_hub import InferenceClient
+client = InferenceClient("stabilityai/stable-diffusion-xl-refiner-1.0", token="api_token")
+
+# output is a PIL.Image object
+image = client.image_to_image("cat.png", prompt="Turn the cat into a tiger.")`);
+
+		expect(snippets[1].client).toEqual("requests");
+		expect(snippets[1].content).toEqual(`import requests
+
+API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-refiner-1.0"
+headers = {"Authorization": "Bearer api_token"}
+
+def query(payload):
+	with open(payload["inputs"], "rb") as f:
+		img = f.read()
+		payload["inputs"] = base64.b64encode(img).decode("utf-8")
+	response = requests.post(API_URL, headers=headers, json=payload)
+	return response.content
+
+image_bytes = query({
+	"inputs": "cat.png",
+	"parameters": {"prompt": "Turn the cat into a tiger."},
+})
+
+# You can access the image with PIL.Image for example
+import io
+from PIL import Image
+image = Image.open(io.BytesIO(image_bytes))`);
+	});
+
 	it("text-to-image", async () => {
 		const model: ModelDataMinimal = {
 			id: "black-forest-labs/FLUX.1-schnell",

packages/tasks/src/snippets/python.ts

@@ -182,6 +182,39 @@ output = query({
 	];
 };
 
+const snippetImageToImage = (model: ModelDataMinimal, accessToken: string): InferenceSnippet[] => {
+	const inputsAsStr = getModelInputSnippet(model) as string;
+	const inputsAsObj = JSON.parse(inputsAsStr);
+
+	return [
+		{
+			client: "huggingface_hub",
+			content: `${snippetImportInferenceClient(model, accessToken)}
+# output is a PIL.Image object
+image = client.image_to_image("${inputsAsObj.image}", prompt="${inputsAsObj.prompt}")`,
+		},
+		{
+			client: "requests",
+			content: `def query(payload):
+	with open(payload["inputs"], "rb") as f:
+		img = f.read()
+		payload["inputs"] = base64.b64encode(img).decode("utf-8")
+	response = requests.post(API_URL, headers=headers, json=payload)
+	return response.content
+
+image_bytes = query({
+	"inputs": "${inputsAsObj.image}",
+	"parameters": {"prompt": "${inputsAsObj.prompt}"},
+})
+
+# You can access the image with PIL.Image for example
+import io
+from PIL import Image
+image = Image.open(io.BytesIO(image_bytes))`,
+		},
+	];
+};
+
 const snippetTabular = (model: ModelDataMinimal): InferenceSnippet => ({
 	content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
@@ -312,6 +345,7 @@ const pythonSnippets: Partial<
 	"image-segmentation": snippetFile,
 	"document-question-answering": snippetDocumentQuestionAnswering,
 	"image-to-text": snippetFile,
+	"image-to-image": snippetImageToImage,
 	"zero-shot-image-classification": snippetZeroShotImageClassification,
 };