Resolve DOO-88 "Stable Diffusion"

Pipfile

@@ -0,0 +1,12 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+
+[dev-packages]
+ipykernel = "*"
+
+[requires]
+python_version = "3.11"

client/package.json

@@ -18,6 +18,7 @@
     "@radix-ui/react-avatar": "^1.0.4",
     "@radix-ui/react-checkbox": "^1.0.4",
     "@radix-ui/react-context-menu": "^2.1.5",
+    "@radix-ui/react-dialog": "^1.0.5",
     "@radix-ui/react-dropdown-menu": "^2.0.6",
     "@radix-ui/react-icons": "^1.3.0",
     "@radix-ui/react-label": "^2.0.2",

client/src/api/index.ts

@@ -0,0 +1,4 @@
+import { sfu } from './sfu';
+import { stableDiffusion } from './stableDiffusion';
+
+export { sfu, stableDiffusion };

client/src/api/sfu.ts

@@ -0,0 +1,25 @@
+import { REST } from '@/constants';
+import axios from 'axios';
+
+/**
+ * Defines REST methods for the SFU API.
+ * @author Yousef Yassin
+ */
+
+/**
+ * Function to poll for an ongoing stream in a room and initialize a consumer if available.
+ * @param roomID String, the ID of the room to poll.
+ * @param initConsumer Function, callback to initialize the consumer when a stream is available.
+ */
+const pollOngoingStream = async (roomID: string, initConsumer: () => void) => {
+  try {
+    const { data } = await axios.put(REST.sfu.poll, {
+      roomId: roomID,
+    });
+    data.roomHasProducer && initConsumer();
+  } catch (e) {
+    console.error('Failed to poll for ongoing stream');
+  }
+};
+
+export const sfu = { poll: pollOngoingStream };

client/src/api/stableDiffusion.ts

@@ -0,0 +1,28 @@
+import axios from 'axios';
+
+/**
+ * Defines REST methods for the Stable Diffusion API.
+ * @author Yousef Yassin
+ */
+
+/**
+ * Sends a textual prompt and an image, as a dataURL, to the server
+ * to perform stable diffusion. Return the generated images.
+ * TODO: Add more options here. / more methods to set models.
+ * @param prompt The textual prompt.
+ * @param dataURL The image to condition on.
+ * @returns The generated images { image_data_urls: string[] }.
+ */
+const diffusion = async (prompt: string, dataURL: string) => {
+  try {
+    const { data } = await axios.post('http://localhost:5000/diffusion', {
+      prompt,
+      im_dataurl: dataURL,
+    });
+    return data;
+  } catch (e) {
+    console.error('Failed to diffuse.');
+  }
+};
+
+export const stableDiffusion = { diffusion };

client/src/components/lib/ContextMenu.tsx

@@ -5,6 +5,7 @@ import { useCanvasElementStore } from '@/stores/CanvasElementsStore';
 import ExportSelectedPNGContextItem from './ExportSelectedPNGContextItem';
 import ContextMenuItem from './ContextMenuItem';
 import { useWebSocketStore } from '@/stores/WebSocketStore';
+import StableDiffusionContextItem from './StableDiffusion/StableDiffusionContextItem';
 
 /**
  * Defines a context menu, with options, that is revealed
@@ -55,6 +56,7 @@ const ContextMenu = () => {
               </div>
             </ContextMenuItem>
             <ExportSelectedPNGContextItem />
+            <StableDiffusionContextItem />
           </>
         ) : null}
       </RadixContextMenu.Content>

client/src/components/lib/StableDiffusion/StableDiffusionContextItem.tsx

@@ -0,0 +1,30 @@
+import React from 'react';
+import { ImageIcon } from '@radix-ui/react-icons';
+import { useAppStore } from '@/stores/AppStore';
+import ContextMenuItem from '../ContextMenuItem';
+
+/**
+ * Context Menu option that opens up the Stable Diffusion Sheet.
+ * @author Yousef Yassin
+ */
+
+const StableDiffusionContextItem = () => {
+  const { setIsUsingStableDiffusion } = useAppStore([
+    'setIsUsingStableDiffusion',
+  ]);
+  return (
+    <ContextMenuItem
+      className="text-violet-500"
+      onClick={() => {
+        setIsUsingStableDiffusion(true);
+      }}
+    >
+      Stable Diffusion
+      <div className="ml-auto pl-5 text-violet-500 group-data-[highlighted]:text-white group-data-[disabled]:text-mauve8">
+        <ImageIcon />
+      </div>
+    </ContextMenuItem>
+  );
+};
+
+export default StableDiffusionContextItem;

client/src/components/lib/StableDiffusion/StableDiffusionSheet.tsx

@@ -0,0 +1,238 @@
+import React, { useEffect } from 'react';
+import {
+  Sheet,
+  SheetClose,
+  SheetContent,
+  SheetDescription,
+  SheetFooter,
+  SheetHeader,
+  SheetTitle,
+} from '@/components/ui/sheet';
+import { Button } from '../../ui/button';
+import { Label } from '../../ui/label';
+import { Input } from '../../ui/input';
+import { useAppStore } from '@/stores/AppStore';
+import { useCanvasElementStore } from '@/stores/CanvasElementsStore';
+import { createElement } from '@/lib/canvasElements/canvasElementUtils';
+import { renderElementsOnOffscreenCanvas } from '@/lib/export';
+import { stableDiffusion } from '@/api';
+import {
+  dataURLToFile,
+  generateRandId,
+  getDataURL,
+  resizeImageFile,
+} from '@/lib/bytes';
+import { injectImageElement } from '@/lib/image';
+
+/**
+ * Sidebar sheet that allows the user to generate images using Stable Diffusion.
+ * @author Yousef Yassin
+ */
+
+const StableDiffusionSheet = () => {
+  const { setPendingImageElement, addCanvasShape, editCanvasElement } =
+    useCanvasElementStore([
+      'setPendingImageElement',
+      'addCanvasShape',
+      'editCanvasElement',
+    ]);
+  const { isUsingStableDiffusion, setIsUsingStableDiffusion, zoom, appHeight } =
+    useAppStore([
+      'isUsingStableDiffusion',
+      'setIsUsingStableDiffusion',
+      'zoom',
+      'appHeight',
+    ]);
+  const { canvasColor, setTool } = useAppStore(['canvasColor', 'setTool']);
+  const {
+    selectedElementIds,
+    p1,
+    p2,
+    types,
+    freehandPoints,
+    freehandBounds,
+    textStrings,
+    fontFamilies,
+    fontSizes,
+    fillColors,
+    fileIds,
+    isImagePlaceds,
+    angles,
+    roughElements,
+    opacities,
+    strokeColors,
+    strokeWidths,
+  } = useCanvasElementStore([
+    'selectedElementIds',
+    'p1',
+    'p2',
+    'types',
+    'freehandPoints',
+    'freehandBounds',
+    'textStrings',
+    'fontFamilies',
+    'fontSizes',
+    'fillColors',
+    'fileIds',
+    'isImagePlaceds',
+    'angles',
+    'roughElements',
+    'opacities',
+    'strokeColors',
+    'strokeWidths',
+  ]);
+  const [dataURL, setDataURL] = React.useState<string>('');
+  const [prompt, setPrompt] = React.useState<string>('');
+  const [diffusionImages, setDiffusionImages] = React.useState<string[]>([]);
+  const [selectedIdx, setSelectedIdx] = React.useState<number>(0);
+  const [isLoading, setIsLoading] = React.useState<boolean>(false);
+
+  const populateDataURL = () => {
+    const canvas = renderElementsOnOffscreenCanvas(
+      selectedElementIds,
+      {
+        p1,
+        p2,
+        angles,
+        types,
+        freehandPoints,
+        freehandBounds,
+        textStrings,
+        fontFamilies,
+        fontSizes,
+        fillColors,
+        isImagePlaceds,
+        fileIds,
+        roughElements,
+        opacities,
+        strokeColors,
+        strokeWidths,
+      },
+      {
+        margin: 20,
+        canvasColor,
+      },
+    );
+    canvas && setDataURL(canvas.toDataURL('image/png'));
+  };
+
+  useEffect(() => {
+    populateDataURL();
+  }, [isUsingStableDiffusion]);
+  return (
+    <Sheet
+      modal={false}
+      open={isUsingStableDiffusion}
+      onOpenChange={setIsUsingStableDiffusion}
+    >
+      {/* Prevent closing when clicking outside */}
+      <SheetContent onInteractOutside={(e) => e.preventDefault()}>
+        <SheetHeader>
+          <SheetTitle>Stable Diffusion</SheetTitle>
+          <SheetDescription>
+            Enter a prompt, and watch the magic happen!
+          </SheetDescription>
+        </SheetHeader>
+        <div className="flex flex-col gap-2 pt-[1rem]">
+          <div>
+            <Label htmlFor="prompt">Prompt</Label>
+          </div>
+
+          <Input
+            id="promp"
+            placeholder="Something Creative"
+            value={prompt}
+            onChange={(e) => setPrompt(e.target.value)}
+          />
+        </div>
+        <div className="flex items-center justify-center p-[2rem]">
+          <div className={`w-[10rem] overflow-hidden rounded-md`}>
+            <img className="h-full w-full object-cover" src={dataURL} />
+          </div>
+        </div>
+
+        {diffusionImages.length > 0 && (
+          <div className="flex flex-wrap justify-center gap-4 overflow-y-scroll max-h-[35rem] p-[1rem]">
+            {diffusionImages.map((imageDataURL, i) => (
+              <div
+                key={`SD_IM_${i}`}
+                className={`flex items-center justify-center ${
+                  i === selectedIdx
+                    ? 'ring-4 ring-offset-2 ring-[#818cf8]'
+                    : 'hover:ring-2 hover:ring-offset-2 hover:ring-[#818cf8]'
+                }`}
+              >
+                <div
+                  className={`w-[16rem] overflow-hidden rounded-md`}
+                  onClick={() => setSelectedIdx(i)}
+                >
+                  <img
+                    className="h-full w-full object-cover"
+                    src={imageDataURL}
+                  />
+                </div>
+              </div>
+            ))}
+          </div>
+        )}
+        <SheetFooter className="sm:justify-start pt-4">
+          <Button onClick={populateDataURL}>Refresh</Button>
+          <Button
+            disabled={isLoading}
+            onClick={async () => {
+              // TODO: resize on set
+              const imSize = 96;
+              const resizedFile = dataURL && dataURLToFile(dataURL);
+              if (!resizedFile) return;
+              const resizedImage = await resizeImageFile(resizedFile, {
+                maxWidthOrHeight: imSize,
+              });
+              const resizedDataURL = await getDataURL(resizedImage);
+              setIsLoading(true);
+              const { image_data_urls: imageDataURLS } =
+                await stableDiffusion.diffusion(prompt, resizedDataURL);
+              setIsLoading(false);
+              setDiffusionImages(imageDataURLS);
+            }}
+          >
+            {isLoading ? 'Generating...' : 'Generate'}
+          </Button>
+          <SheetClose asChild>
+            <Button
+              onClick={async () => {
+                const selectedDataURL = diffusionImages[selectedIdx];
+                const imageFile = dataURLToFile(selectedDataURL);
+                // Create a proxy element
+                const id = generateRandId();
+                const placeholderElement = createElement(
+                  id,
+                  0,
+                  0,
+                  0,
+                  0,
+                  'image',
+                );
+                setPendingImageElement(id);
+                // Inject the image into the proxy
+                await injectImageElement(
+                  placeholderElement,
+                  imageFile,
+                  addCanvasShape,
+                  editCanvasElement,
+                  { zoom, appHeight },
+                  true,
+                );
+                // And let the user place the image
+                setTool('image');
+              }}
+            >
+              Done
+            </Button>
+          </SheetClose>
+        </SheetFooter>
+      </SheetContent>
+    </Sheet>
+  );
+};
+
+export default StableDiffusionSheet;