add color palette generator

Monadical-SAS · Dec 1, 2023 · 5b8d0d5 · 5b8d0d5
1 parent 0893178
commit 5b8d0d5
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 82 deletions.
diff --git a/morpheus-client/components/ColorPaletteGenerator/ColorPaletteGenerator.tsx b/morpheus-client/components/ColorPaletteGenerator/ColorPaletteGenerator.tsx
@@ -0,0 +1,98 @@
+import { useRef, useState } from "react";
+import { CloseIcon } from "@/components/icons/close";
+import { CheckIcon } from "@/components/icons/check";
+import { useImagine } from "@/context/ImagineContext";
+import { useToastContext } from "@/context/ToastContext";
+
+type ColorPaletteGeneratorProps = {
+  closeModal: () => void;
+};
+
+const ColorPaletteGenerator = (props: ColorPaletteGeneratorProps) => {
+  const canvasRef = useRef<HTMLCanvasElement>(null);
+  const { setColorPaletteFile } = useImagine();
+  const { showInfoAlert } = useToastContext();
+  const [colors, setColors] = useState<string[]>(["#ffffff", "#ffffff"]);
+
+  const addColor = () => {
+    if (colors.length >= 5) {
+      showInfoAlert("You can only have 5 colors in your palette");
+      return;
+    }
+    setColors([...colors, "#ffffff"]);
+  };
+
+  const changeColor = (color: string, index: number) => {
+    const newColors = [...colors];
+    newColors[index] = color;
+    setColors(newColors);
+  };
+
+  const resetColors = () => {
+    setColors(["#ffffff"]);
+  };
+
+  const handleCompleted = () => {
+    if (!canvasRef) return;
+    const canvas = canvasRef.current;
+    if (!canvas) return;
+    const context = canvas.getContext("2d");
+    if (!context) return;
+
+    const gradient = context.createLinearGradient(0, 0, 0, canvas.height);
+    const step = 1 / (colors.length - 1);
+    for (let i = 0; i < colors.length; i++) {
+      gradient.addColorStop(step * i, colors[i]);
+    }
+    context.fillStyle = gradient;
+    context.fillRect(0, 0, canvas.width, canvas.height);
+
+    canvas.toBlob((blob: any) => {
+      const file = new File([blob as Blob], "color-palette.png", {
+        type: "image/png",
+      });
+      setColorPaletteFile(file);
+      props.closeModal();
+    }, "image/png");
+  };
+
+  return (
+    <div className="">
+      <h2 className="headline-3 white">Generate Color Palette</h2>
+
+      <canvas ref={canvasRef} width={500} height={500} hidden />
+      <div className="w-[500px] h-[500px] grid mt-5 p-1 bg-white">
+        {colors.map((color, index) => (
+          <div key={index}>
+            <input
+              type="color"
+              value={color}
+              className="w-full h-full"
+              onChange={(e) => changeColor(e.target.value, index)}
+            />
+          </div>
+        ))}
+      </div>
+
+      <div className="flex flex-row items-center gap-[16px] mt-10">
+        <button className="buttonSubmit w-full" onClick={addColor}>
+          Add Color
+        </button>
+
+        <span className="material-icons text-white" onClick={resetColors}>
+          replay
+        </span>
+
+        <span onClick={props.closeModal}>
+          <CloseIcon width={"24"} height={"24"} color={"#FFFFFF"} />
+        </span>
+
+        <span onClick={handleCompleted}>
+          <CheckIcon width={"24"} height={"24"} color={"#FFFFFF"} />
+        </span>
+      </div>
+    </div>
+  );
+};
+
+export default ColorPaletteGenerator;
diff --git a/morpheus-client/components/ImagineBase/ImagineBase.tsx b/morpheus-client/components/ImagineBase/ImagineBase.tsx
@@ -43,8 +43,8 @@ const ImagineBase = (props: MainContainerProps) => {
           label={"Base Image"}
           imageFile={img2imgFile}
           setImageFile={setImg2imgFile}
-          showEditImage={props.showMaskInput}
           showPaintImageLink={true}
+          showColorPalette={false}
         />
       )}
       {props.showMaskInput && (
@@ -54,14 +54,16 @@ const ImagineBase = (props: MainContainerProps) => {
           setImageFile={setMaskFile}
           icon={<UploadMaskIcon />}
           showPaintMask={img2imgFile !== null}
+          showColorPalette={false}
         />
       )}
       {props.showPaletteInput && (
         <ImagineImageInput
           label={"Palette image"}
           imageFile={colorPaletteFile}
           setImageFile={setColorPaletteFile}
-          showPaintMask={img2imgFile !== null}
+          showPaintMask={false}
+          showColorPalette={img2imgFile !== null}
         />
       )}
     </div>

diff --git a/morpheus-client/components/ImagineImageInput/ImagineImageInput.tsx b/morpheus-client/components/ImagineImageInput/ImagineImageInput.tsx
@@ -6,27 +6,27 @@ import React, {
   useRef,
   useState,
 } from "react";
-
+import Link from "next/link";
+import ColorPaletteGenerator from "@/components/ColorPaletteGenerator/ColorPaletteGenerator";
 import Modal from "../Modal/Modal";
 import MaskPaintingCanvas from "../MaskPaintingCanvas/MaskPaintingCanvas";
 import { CloseIcon } from "../icons/close";
 import { UploadImageIcon } from "../icons/uploadImage";
 import { PaintImageIcon } from "../icons/paintImage";
 import useWindowDimensions from "@/hooks/useWindowDimensions";
 import styles from "./ImagineImageInput.module.scss";
-import Link from "next/link";
 
 interface DragDropFileProps {
+  id?: string;
   imageFile: File | null;
   setImageFile: (file: File | null) => void;
   label?: string;
   buttonLabel?: string;
   icon?: ReactNode;
   styles?: CSSProperties;
-  showEditImage?: boolean;
   showPaintImageLink?: boolean;
   showPaintMask?: boolean;
-  id?: string;
+  showColorPalette?: boolean;
 }
 
 const ImagineImageInput = (props: DragDropFileProps) => {
@@ -136,7 +136,9 @@ const ImagineImageInput = (props: DragDropFileProps) => {
         )}
       </form>
 
-      {props.showPaintImageLink || props.showPaintMask ? (
+      {props.showPaintImageLink ||
+      props.showPaintMask ||
+      props.showColorPalette ? (
         <Fragment>
           <div className={styles.verticalSeparator} />
         </Fragment>
@@ -162,6 +164,18 @@ const ImagineImageInput = (props: DragDropFileProps) => {
           </a>
         </div>
       )}
+
+      {props.showColorPalette && (
+        <div className={styles.formInputItem}>
+          <PaintImageIcon />
+          <a
+            className="underline body-1 main pointer"
+            onClick={() => setShowEditModal(true)}
+          >
+            {isMobile ? "Generate" : "Generate a color palette"}
+          </a>
+        </div>
+      )}
     </div>
   );
 
@@ -193,11 +207,17 @@ const ImagineImageInput = (props: DragDropFileProps) => {
         isOpen={showEditModal}
         toggleModal={() => setShowEditModal(!showEditModal)}
       >
-        <MaskPaintingCanvas
-          width={isMobile && width ? width - 88 : 512}
-          height={isMobile && width ? width - 88 : 512}
-          closeModal={() => setShowEditModal(false)}
-        />
+        {props.showPaintMask && (
+          <MaskPaintingCanvas
+            width={isMobile && width ? width - 88 : 512}
+            height={isMobile && width ? width - 88 : 512}
+            closeModal={() => setShowEditModal(false)}
+          />
+        )}
+
+        {props.showColorPalette && (
+          <ColorPaletteGenerator closeModal={() => setShowEditModal(false)} />
+        )}
       </Modal>
     </Fragment>
   );

diff --git a/morpheus-client/components/ImagineSettings/ImagineSettings.tsx b/morpheus-client/components/ImagineSettings/ImagineSettings.tsx
@@ -15,8 +15,13 @@ import { useShowSettings } from "@/hooks/useShowSettings";
 import { CloseIcon } from "../icons/close";
 import { InfoIcon } from "../icons/info";
 import { SettingsIcon } from "../icons/settings";
-import styles from "./ImagineSettings.module.scss";
 import { useImagine } from "@/context/ImagineContext";
+import {
+  COLOR_PALETTES_CONTROLNET,
+  COLOR_PALETTES_IMAGE_TO_IMAGE,
+  IMAGE_SIZES,
+} from "@/utils/constants";
+import styles from "./ImagineSettings.module.scss";
 
 interface OptionState {
   title: string;
@@ -113,14 +118,7 @@ const ImagineSettings = () => {
           />
 
           <InputSelect
-            options={[
-              "480x480",
-              "512x512",
-              "640x640",
-              "720x720",
-              "768x768",
-              "1024x1024",
-            ]}
+            options={IMAGE_SIZES}
             selected={imageSize}
             setSelected={setImageSize}
           />
@@ -213,22 +211,7 @@ const ImagineSettings = () => {
               }
             />
             <InputSelect
-              options={[
-                "Quantization - Blend",
-                "Quantization - Contours",
-                "Quantization Gray - Blend",
-                "Quantization Gray - Contours",
-                "Random Polygons - Blend",
-                "Random Polygons - Contours",
-                "Random Color Blocks Small - Blend",
-                "Random Color Blocks Small - Contours",
-                "Random Color Blocks Large - Blend",
-                "Random Color Blocks Large - Contours",
-                "Color Matching - PCA",
-                "Color Matching - Cholesky",
-                "Color Matching - Symmetric",
-                "Linear Color Transfer",
-              ]}
+              options={COLOR_PALETTES_IMAGE_TO_IMAGE}
               selected={paletteTechnique}
               setSelected={setPaletteTechnique}
             />
@@ -245,17 +228,7 @@ const ImagineSettings = () => {
               }
             />
             <InputSelect
-              options={[
-                "Quantization",
-                "Quantization Gray",
-                "Random Polygons",
-                "Random Color Blocks Small",
-                "Random Color Blocks Large",
-                "Color Matching - PCA",
-                "Color Matching - Cholesky",
-                "Color Matching - Symmetric",
-                "Linear Color Transfer",
-              ]}
+              options={COLOR_PALETTES_CONTROLNET}
               selected={paletteTechnique}
               setSelected={setPaletteTechnique}
             />

diff --git a/morpheus-client/context/ImagineContext.tsx b/morpheus-client/context/ImagineContext.tsx
@@ -42,41 +42,26 @@ export interface ImagineContextProps {
   setImg2imgFile: (image: File | null) => void;
   maskFile: File | null;
   setMaskFile: (image: File | null) => void;
-  setColorPaletteURL: (image: string) => void;
   colorPaletteFile: File | null;
   setColorPaletteFile: (image: File | null) => void;
   generateImages: (option: ImagineOptions) => void;
   resultImages: Array<ImagineResult>;
   clearResults: () => void;
 }
 
-const defaultState = {
-  isLoading: false,
-  setImg2ImgURL: () => console.log("setImg2ImgURL"),
-  img2imgFile: null,
-  setImg2imgFile: () => console.log("setImg2imgFile"),
-  maskFile: null,
-  setMaskFile: () => console.log("setMaskFile"),
-  setColorPaletteURL: () => console.log("setColorPaletteURL"),
-  colorPaletteFile: null,
-  setColorPaletteFile: () => console.log("setColorPaletteFile"),
-  generateImages: () => console.log("generateImages"),
-  resultImages: [],
-  clearResults: () => console.log("clearResults"),
-};
-
-const ImagineContext = createContext<ImagineContextProps>(defaultState);
+const ImagineContext = createContext<ImagineContextProps>(
+  {} as ImagineContextProps
+);
 
 const ImagineProvider = (props: { children: ReactNode }) => {
   const { prompt, buildPrompt, restartSDSettings } = useDiffusion();
   const { buildControlNetPrompt } = useControlNet();
   const { showErrorAlert } = useToastContext();
 
-  // Images settings
+  // Image to image images
   const [img2ImgURL, setImg2ImgURL] = useState<string>("");
   const [img2imgFile, setImg2imgFile] = useState<File | null>(null);
   const [maskFile, setMaskFile] = useState<File | null>(null);
-  const [colorPaletteURL, setColorPaletteURL] = useState<string>("");
   const [colorPaletteFile, setColorPaletteFile] = useState<File | null>(null);
 
   // Image results
@@ -101,20 +86,6 @@ const ImagineProvider = (props: { children: ReactNode }) => {
     }
   }, [img2ImgURL]);
 
-  useEffect(() => {
-    if (colorPaletteURL) {
-      getFileBlobFromURL(colorPaletteURL)
-        .then((file: File) => {
-          setColorPaletteFile(file);
-          setColorPaletteURL("");
-        })
-        .catch((err) => {
-          showErrorAlert("Error getting file from URL");
-          console.log(err);
-        });
-    }
-  }, [colorPaletteURL]);
-
   useEffect(() => {
     maskFile && setMaskFile(null);
   }, [img2imgFile]);
@@ -222,7 +193,6 @@ const ImagineProvider = (props: { children: ReactNode }) => {
         setImg2imgFile,
         maskFile,
         setMaskFile,
-        setColorPaletteURL: setColorPaletteURL,
         colorPaletteFile,
         setColorPaletteFile,
         generateImages,