fix: handle cloudflare ReadableStream image type

Author: monkeyWie

src/server/ai/provider/cloudflare.ts

@@ -1,5 +1,5 @@
 import { inCfWorker } from "@/server/lib/env";
-import { buildDataURI } from "@/server/lib/util";
+import { base64ToDataURI, readableStreamToDataURI } from "@/server/lib/util";
 import { getContext } from "@/server/service/context";
 import type { AiProvider, ApiProviderSettings, ApiProviderSettingsItem } from "../types/provider";
 import { type ProviderSettingsType, doParseSettings, getProviderSettingsSchema } from "../types/provider";
@@ -54,24 +54,24 @@ const Cloudflare: AiProvider = {
 			ability: "t2i",
 			enabledByDefault: true,
 		},
-		{
-			id: "@cf/bytedance/stable-diffusion-xl-lightning",
-			name: "Stable Diffusion XL Lightning",
-			ability: "t2i",
-			enabledByDefault: true,
-		},
 		{
 			id: "@cf/lykon/dreamshaper-8-lcm",
 			name: "DreamShaper 8 LCM",
 			ability: "t2i",
 			enabledByDefault: true,
 		},
 		{
-			id: "@cf/runwayml/stable-diffusion-v1-5-img2img",
-			name: "Stable Diffusion v1.5 Img2Img",
-			ability: "i2i",
+			id: "@cf/bytedance/stable-diffusion-xl-lightning",
+			name: "Stable Diffusion XL Lightning",
+			ability: "t2i",
 			enabledByDefault: true,
 		},
+		// {
+		// 	id: "@cf/runwayml/stable-diffusion-v1-5-img2img",
+		// 	name: "Stable Diffusion v1.5 Img2Img",
+		// 	ability: "i2i",
+		// 	enabledByDefault: true,
+		// },
 		{
 			id: "@cf/stabilityai/stable-diffusion-xl-base-1.0",
 			name: "Stable Diffusion XL Base 1.0",
@@ -92,8 +92,14 @@ const Cloudflare: AiProvider = {
 				prompt: request.prompt,
 			});
 
+			if (resp instanceof ReadableStream) {
+				return {
+					images: [await readableStreamToDataURI(resp)],
+				};
+			}
+
 			return {
-				images: [buildDataURI(resp.image)],
+				images: [base64ToDataURI(resp.image)],
 			};
 		}
 
@@ -116,13 +122,13 @@ const Cloudflare: AiProvider = {
 		if (contentType?.includes("image/png") === true) {
 			const imageBuffer = await resp.arrayBuffer();
 			return {
-				images: [buildDataURI(Buffer.from(imageBuffer).toString("base64"))],
+				images: [base64ToDataURI(Buffer.from(imageBuffer).toString("base64"))],
 			};
 		}
 
 		const result = (await resp.json()) as unknown as any;
 		return {
-			images: [buildDataURI(result.result.image)],
+			images: [base64ToDataURI(result.result.image)],
 		};
 	},
 };

src/server/lib/util.ts

@@ -1,12 +1,35 @@
-export function buildDataURI(base64: string, fmt = "png"): string {
+export function base64ToDataURI(base64: string, fmt = "png") {
 	return `data:image/${fmt};base64,${base64}`;
 }
 
-export function buildDataURIArray(base64Array: string[], fmt = "png"): string[] {
-	return base64Array.map((base64) => buildDataURI(base64, fmt));
+export async function readableStreamToDataURI(stream: ReadableStream<Uint8Array>, fmt = "png") {
+	const reader = stream.getReader();
+	const chunks: Uint8Array[] = [];
+
+	while (true) {
+		const result = await reader.read();
+		if (result.done) break;
+		chunks.push(result.value);
+	}
+
+	const totalLength = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
+	const mergedArray = new Uint8Array(totalLength);
+	let offset = 0;
+	for (const chunk of chunks) {
+		mergedArray.set(chunk, offset);
+		offset += chunk.length;
+	}
+
+	let binaryString = "";
+	for (let i = 0; i < mergedArray.length; i++) {
+		binaryString += String.fromCharCode(mergedArray[i]!);
+	}
+	const base64 = btoa(binaryString);
+
+	return base64ToDataURI(base64, fmt);
 }
 
-export async function fetchUrlToDataURI(url: string): Promise<string> {
+export async function fetchUrlToDataURI(url: string) {
 	const resp = await fetch(url);
 	if (!resp.ok) {
 		throw new Error(`Failed to fetch URL: ${url}, status: ${resp.status}`);
@@ -17,5 +40,5 @@ export async function fetchUrlToDataURI(url: string): Promise<string> {
 	const uint8Array = new Uint8Array(arrayBuffer);
 	const binaryString = Array.from(uint8Array, (byte) => String.fromCharCode(byte)).join("");
 	const base64 = btoa(binaryString);
-	return buildDataURI(base64);
+	return base64ToDataURI(base64);
 }