easydiffusion/ui/frontend/build_src/src/stores/imageQueueStore.ts

import create from "zustand";
import produce from "immer";
import { useRandomSeed } from "../utils";

import { ImageRequest } from "./imageCreateStore";

interface ImageQueueState {
  images: ImageRequest[];
  completedImageIds: string[];
  addNewImage: (id: string, imgRec: ImageRequest) => void;
  hasQueuedImages: () => boolean;
  firstInQueue: () => ImageRequest | [];
  removeFirstInQueue: () => void;
  clearCachedIds: () => void;
}

export const useImageQueue = create<ImageQueueState>((set, get) => ({
  images: [],
  completedImageIds: [],
  // use produce to make sure we don't mutate state
  addNewImage: (id: string, imgRec: ImageRequest, isRandom = false) => {
    set(
      produce((state) => {
        let { seed } = imgRec;
        if (isRandom) {
          seed = useRandomSeed();
        }
        state.images.push({ id, options: { ...imgRec, seed } });
      })
    );
  },

  hasQueuedImages: () => {
    return get().images.length > 0;
  },
  firstInQueue: () => {
    return (get().images[0] as ImageRequest) || [];
  },

  removeFirstInQueue: () => {
    set(
      produce((state) => {
        const image = state.images.shift();
        state.completedImageIds.push(image.id);
      })
    );
  },
  clearCachedIds: () => {
    set(
      produce((state) => {
        state.completedImageIds = [];
      })
    );
  },
}));