Image Studio

"use client"

import * as React from "react"

import { ImageStudio, sleep, type Generation, type ImageEvent, type ImageRunner } from "@/components/blocks/image-studio/image-studio"

type Job = {
  id: string
  images: {
    status: "queued" | "running" | "succeeded" | "failed"
    progress?: number
    url?: string
    width?: number
    height?: number
    seed?: number
    error?: string
  }[]
}

async function json<T>(response: Response): Promise<T> {
  if (!response.ok) throw new Error((await response.text()) || "The image service didn’t respond.")
  return response.json() as Promise<T>
}

const generate: ImageRunner = async function* (request, { signal }) {
  const job = await json<Job>(
    await fetch("/api/images", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        id: request.id,
        kind: request.kind,
        prompt: request.settings.prompt,
        negativePrompt: request.settings.negativePrompt,
        aspectRatio: request.settings.aspectRatio,
        model: request.settings.model,
        style: request.settings.preset,
        seeds: request.seeds,
        sourceUrl: request.source?.image.fullSrc ?? request.source?.image.src,
        strength: request.strength,
      }),
      signal,
    })
  )

  const finished = new Set<number>()
  while (finished.size < request.seeds.length) {
    await sleep(1000, signal)
    const current = await json<Job>(await fetch(`/api/images/${job.id}`, { signal }))
    for (const [index, image] of current.images.entries()) {
      if (finished.has(index)) continue
      let event: ImageEvent
      if (image.status === "succeeded" && image.url) {
        finished.add(index)
        event = {
          type: "done",
          index,
          image: { src: image.url, width: image.width, height: image.height, seed: image.seed },
        }
      } else if (image.status === "failed") {
        finished.add(index)
        event = { type: "error", index, message: image.error ?? "This image failed.", retryable: true }
      } else if (image.status === "queued") {
        event = { type: "queued", index }
      } else {
        event = { type: "progress", index, progress: image.progress }
      }
      yield event
    }
  }
}

export function Studio({ initial }: { initial: Generation[] }) {
  const [generations, setGenerations] = React.useState(initial)

  return (
    <ImageStudio
      generations={generations}
      models={[{ value: "image-2", label: "Image 2", credits: 2 }]}
      onGenerate={generate}
      onVary={generate}
      onUpscale={generate}
      onCancel={async (request) => {
        await fetch(`/api/images/${request.id}`, { method: "DELETE" })
      }}
      onComplete={(generation) => {
        setGenerations((list) => [generation, ...list.filter((item) => item.id !== generation.id)])
        void fetch("/api/generations", { method: "PUT", body: JSON.stringify(generation) })
      }}
      onFavorite={async (image, favorite) => {
        await json(await fetch(`/api/images/${image.id}/favorite`, { method: favorite ? "PUT" : "DELETE" }))
      }}
      onDelete={async (image) => {
        await json(await fetch(`/api/images/${image.id}`, { method: "DELETE" }))
      }}
    />
  )
}

"use client"

import * as React from "react"

import { Button } from "@/components/ui/button"

import {
  createImageEventQueue,
  ImageStudio,
  MaskEditor,
  type GeneratedImage,
  type Generation,
  type ImageRunner,
} from "@/components/blocks/image-studio/image-studio"

const edit: ImageRunner = (request, { signal }) => {
  const queue = createImageEventQueue()
  const body = new FormData()
  body.set("prompt", request.settings.prompt)
  body.set("image", request.source?.image.fullSrc ?? request.source?.image.src ?? "")
  body.set("seeds", request.seeds.join(","))
  if (request.mask) body.set("mask", request.mask, "mask.png")

  fetch("/api/images/edit", { method: "POST", body, signal })
    .then(async (response) => {
      if (!response.ok) throw new Error("The edit didn’t go through.")
      const { urls } = (await response.json()) as { urls: string[] }
      urls.forEach((src, index) => queue.push({ type: "done", index, image: { src } }))
      queue.end()
    })
    .catch((error) => queue.fail(error))

  return queue
}

export function StudioWithEdits({ generations }: { generations: Generation[] }) {
  return (
    <ImageStudio
      generations={generations}
      models={[{ value: "image-2", label: "Image 2" }]}
      onGenerate={edit}
      onEdit={edit}
    />
  )
}

export function EditButton({ image, generation }: { image: GeneratedImage; generation: Generation }) {
  const [open, setOpen] = React.useState(false)

  return (
    <>
      <Button variant="outline" onClick={() => setOpen(true)}>
        Edit region
      </Button>
      <MaskEditor
        open={open}
        onOpenChange={setOpen}
        image={image}
        generation={generation}
        onSubmit={async ({ prompt, mask }) => {
          const body = new FormData()
          body.set("prompt", prompt)
          body.set("image", image.fullSrc ?? image.src ?? "")
          body.set("mask", mask, "mask.png")
          const response = await fetch("/api/images/edit", { method: "POST", body })
          if (!response.ok) throw new Error("The edit didn’t go through. Try again.")
        }}
      />
    </>
  )
}

"use client"

import * as React from "react"

import {
  createImageEventQueue,
  defaultSettings,
  ImageComposer,
  ImageGeneration,
  useImageRuns,
  type ImageEvent,
  type ImageRunner,
  type ImageSettings,
} from "@/components/blocks/image-studio/image-studio"

const models = [{ value: "image-2", label: "Image 2", credits: 2 }]

const stream: ImageRunner = (request, { signal }) => {
  const queue = createImageEventQueue()
  const params = new URLSearchParams({
    prompt: request.settings.prompt,
    aspectRatio: request.settings.aspectRatio,
    seeds: request.seeds.join(","),
    kind: request.kind,
  })
  const source = new EventSource(`/api/images/stream?${params}`)
  source.onmessage = (message) => queue.push(JSON.parse(message.data) as ImageEvent)
  source.addEventListener("end", () => {
    source.close()
    queue.end()
  })
  source.onerror = () => {
    source.close()
    queue.fail(new Error("Lost the connection to the image service."))
  }
  signal.addEventListener("abort", () => source.close(), { once: true })
  return queue
}

export function ImageToolResult({ input }: { input: Partial<ImageSettings> & { prompt: string } }) {
  const [settings, setSettings] = React.useState<ImageSettings>({ ...defaultSettings, count: 2, ...input })
  const runs = useImageRuns()
  const started = React.useRef(false)

  React.useEffect(() => {
    if (started.current) return
    started.current = true
    runs.start(stream, { kind: "generate", settings })
  }, [runs, settings])

  return (
    <div className="flex flex-col gap-4 rounded-2xl border p-3">
      {runs.live.map((generation) => (
        <ImageGeneration
          key={generation.id}
          generation={generation}
          models={models}
          headingLevel={3}
          onCancel={() => runs.cancel(generation.id)}
          onRetry={(image) => runs.retry(stream, generation, image)}
          onVary={(image, from, strength) =>
            runs.start(stream, { kind: "variation", settings: from.settings, source: { generation: from, image }, strength })
          }
          onUpscale={(image, from) =>
            runs.start(stream, {
              kind: "upscale",
              settings: { ...from.settings, count: 1 },
              source: { generation: from, image },
            })
          }
        />
      ))}
      <ImageComposer
        settings={settings}
        onSettingsChange={setSettings}
        models={models}
        placeholder="Refine this image"
        onSubmit={(next, reference) => {
          runs.start(stream, { kind: "generate", settings: next, reference })
        }}
      />
    </div>
  )
}