loader.ts

import { Ai } from "@cloudflare/ai";
import { LoaderFunctionArgs } from "@remix-run/cloudflare";

export async function loader({ request, context }: LoaderFunctionArgs) {
  const url = new URL(request.url);
  const query = url.searchParams.get("query");

  if (!query) {
    return new Response('Please provide a query parameter', {
      status: 400,
    });
  }

  const { env } = context.cloudflare;
  const ai = new Ai(env.AI);

  const stream = (await ai.run('@cf/stabilityai/stable-diffusion-xl-base-1.0', {
    prompt: query ?? '',
  })) as unknown as ReadableStream; // Oops

  async function streamToArrayBuffer(stream: ReadableStream) {
    const reader = stream.getReader();
    const chunks = [];
  
    // eslint-disable-next-line no-constant-condition
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }
  
    const totalLength = chunks.reduce((acc, chunk) => acc + chunk.length, 0);
    const result = new Uint8Array(totalLength);
  
    let position = 0;
    for (const chunk of chunks) {
      result.set(chunk, position);
      position += chunk.length;
    }
  
    return result;
  }

  const arrayBuffer = await streamToArrayBuffer(stream);

  const bucket = env.BUCKET;
  const key = `${Date.now()}.png`;
  const r2Object = await bucket.put(key, arrayBuffer, {
    httpMetadata: {
      contentType: 'image/png',
    },
  });

  console.log('Object created:', {
    key: r2Object.key,
    size: r2Object.size,
    version: r2Object.version,
  });

  const object = await bucket.get(r2Object.key);

  if (!object) {
    return new Response('Object not found', {
      status: 404,
    });
  }

  return new Response(object.body, {
    headers: {
      'Content-Type': 'image/png',
    },
  })
}