From 9d8eedc77a0fdbd5d1651bec4ff8a06398715e3e Mon Sep 17 00:00:00 2001
From: BFlat <bombayapirat@gmail.com>
Date: Fri, 12 Sep 2025 14:12:35 +0700
Subject: [PATCH] Feature/dnd image prompts (#823)

- hide/show end frame input depending on the model selected
- image prompt drag and drop reordering if there's more than one image
added
---
 .../promptbox/src/lib/ImagePromptRow.tsx      | 189 +++++++++++++++---
 .../promptbox/src/lib/PromptBoxVideo.tsx      |  10 +-
 2 files changed, 171 insertions(+), 28 deletions(-)

diff --git a/frontend/libs/components/promptbox/src/lib/ImagePromptRow.tsx b/frontend/libs/components/promptbox/src/lib/ImagePromptRow.tsx
index 0881748237..4283180942 100644
--- a/frontend/libs/components/promptbox/src/lib/ImagePromptRow.tsx
+++ b/frontend/libs/components/promptbox/src/lib/ImagePromptRow.tsx
@@ -1,4 +1,10 @@
-import { useEffect, useMemo, useRef, useState } from "react";
+import {
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+  type CSSProperties,
+} from "react";
 import { Button } from "@storyteller/ui-button";
 import { Tooltip } from "@storyteller/ui-tooltip";
 import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
@@ -15,6 +21,23 @@ import { faImage } from "@fortawesome/pro-regular-svg-icons";
 import { RefImage } from "./promptStore";
 import { twMerge } from "tailwind-merge";
 import { UploaderState, UploaderStates } from "@storyteller/common";
+import {
+  DndContext,
+  closestCenter,
+  KeyboardSensor,
+  PointerSensor,
+  useSensor,
+  useSensors,
+  DragEndEvent,
+} from "@dnd-kit/core";
+import {
+  arrayMove,
+  SortableContext,
+  sortableKeyboardCoordinates,
+  horizontalListSortingStrategy,
+  useSortable,
+} from "@dnd-kit/sortable";
+import { CSS } from "@dnd-kit/utilities";
 
 export type UploadImageFn = ({
   title,
@@ -40,6 +63,7 @@ interface ImagePromptRowProps {
   endFrameImage?: RefImage;
   setEndFrameImage?: (image?: RefImage) => void;
   allowUploadEnd?: boolean;
+  showEndFrameSection?: boolean;
 }
 
 export const ImagePromptRow = ({
@@ -56,6 +80,7 @@ export const ImagePromptRow = ({
   endFrameImage,
   setEndFrameImage,
   allowUploadEnd,
+  showEndFrameSection = true,
 }: ImagePromptRowProps) => {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const [uploadingImages, setUploadingImages] = useState<
@@ -72,6 +97,89 @@ export const ImagePromptRow = ({
     file: File;
   } | null>(null);
 
+  const referenceImagesRef = useRef(referenceImages);
+  useEffect(() => {
+    referenceImagesRef.current = referenceImages;
+  }, [referenceImages]);
+
+  const allowReorder = useMemo(
+    () => maxImagePromptCount > 1 && referenceImages.length > 1,
+    [maxImagePromptCount, referenceImages.length]
+  );
+
+  const sensors = useSensors(
+    useSensor(PointerSensor, {
+      activationConstraint: { distance: 6 },
+    }),
+    useSensor(KeyboardSensor, {
+      coordinateGetter: sortableKeyboardCoordinates,
+    })
+  );
+
+  const SortableImage = ({ image }: { image: RefImage }) => {
+    const {
+      attributes,
+      listeners,
+      setNodeRef,
+      transform,
+      transition,
+      isDragging,
+    } = useSortable({ id: image.id });
+
+    const style: CSSProperties = {
+      transform: CSS.Transform.toString(transform),
+      transition,
+      zIndex: isDragging ? 9999 : undefined,
+    };
+
+    return (
+      <div
+        ref={setNodeRef}
+        style={style}
+        {...(allowReorder ? { ...attributes, ...listeners } : {})}
+        className={twMerge(
+          "glass relative aspect-square overflow-hidden rounded-lg w-14 border-2 border-white/30 transition-opacity group",
+          allowReorder
+            ? "cursor-move hover:border-white/80"
+            : "cursor-pointer hover:border-white/80",
+          isDragging ? "opacity-50 shadow-lg" : ""
+        )}
+        onClick={() => onImageClick?.(image)}
+      >
+        <img
+          src={image.url}
+          alt="Reference"
+          className="h-full w-full object-cover"
+        />
+        <button
+          onClick={(e) => {
+            e.stopPropagation();
+            handleRemoveReference(image.id);
+          }}
+          onMouseDown={(e) => {
+            e.stopPropagation();
+          }}
+          onPointerDown={(e) => {
+            e.stopPropagation();
+          }}
+          className="opacity-0 group-hover:opacity-100 absolute right-[2px] top-[2px] flex h-5 w-5 items-center justify-center rounded-full bg-black/50 hover:bg-red/70 text-white backdrop-blur-md transition-colors hover:bg-black cursor-pointer"
+        >
+          <FontAwesomeIcon icon={faXmark} className="h-2.5 w-2.5" />
+        </button>
+      </div>
+    );
+  };
+
+  const handleDragEnd = (event: DragEndEvent) => {
+    const { active, over } = event;
+    if (!over || active.id === over.id) return;
+    const oldIndex = referenceImages.findIndex((img) => img.id === active.id);
+    const newIndex = referenceImages.findIndex((img) => img.id === over.id);
+    if (oldIndex === -1 || newIndex === -1) return;
+    const newOrder = arrayMove(referenceImages, oldIndex, newIndex);
+    setReferenceImages(newOrder);
+  };
+
   const usedSlotsRender = useMemo(
     () =>
       Math.min(
@@ -167,7 +275,10 @@ export const ImagePromptRow = ({
                 if (uploadTarget === "end") {
                   setEndFrameImage?.(referenceImage);
                 } else {
-                  setReferenceImages([...referenceImages, referenceImage]);
+                  setReferenceImages([
+                    ...referenceImagesRef.current,
+                    referenceImage,
+                  ]);
                 }
               } else if (
                 newState.status === UploaderStates.assetError ||
@@ -200,7 +311,7 @@ export const ImagePromptRow = ({
           if (uploadTarget === "end") {
             setEndFrameImage?.(referenceImage);
           } else {
-            setReferenceImages([...referenceImages, referenceImage]);
+            setReferenceImages([...referenceImagesRef.current, referenceImage]);
           }
         }
 
@@ -286,7 +397,7 @@ export const ImagePromptRow = ({
         <div
           className={twMerge(
             "grow grid py-2 px-3 grid-cols-1",
-            isVideo && "grid-cols-2 gap-5"
+            isVideo && showEndFrameSection && "grid-cols-2 gap-5"
           )}
         >
           <div className="flex gap-2">
@@ -308,30 +419,54 @@ export const ImagePromptRow = ({
             </div>
 
             <div className="flex gap-2">
-              {referenceImages
-                .slice(0, Math.max(0, maxImagePromptCount))
-                .map((image) => (
-                  <div
-                    key={image.id}
-                    className="glass relative aspect-square overflow-hidden rounded-lg w-14 border-2 border-white/30 hover:border-white/80 transition-all group cursor-pointer hover:cursor-zoom-in"
-                    onClick={() => onImageClick?.(image)}
+              {allowReorder ? (
+                <DndContext
+                  sensors={sensors}
+                  collisionDetection={closestCenter}
+                  onDragEnd={handleDragEnd}
+                >
+                  <SortableContext
+                    items={referenceImages
+                      .slice(0, Math.max(0, maxImagePromptCount))
+                      .map((img) => img.id)}
+                    strategy={horizontalListSortingStrategy}
                   >
-                    <img
-                      src={image.url}
-                      alt="Reference"
-                      className="h-full w-full object-cover"
-                    />
-                    <button
-                      onClick={(e) => {
-                        e.stopPropagation();
-                        handleRemoveReference(image.id);
-                      }}
-                      className="opacity-0 group-hover:opacity-100 absolute right-[2px] top-[2px] flex h-5 w-5 items-center justify-center rounded-full bg-black/50 hover:bg-red/70 text-white backdrop-blur-md transition-colors hover:bg-black cursor-pointer"
+                    {referenceImages
+                      .slice(0, Math.max(0, maxImagePromptCount))
+                      .map((image) => (
+                        <SortableImage key={image.id} image={image} />
+                      ))}
+                  </SortableContext>
+                </DndContext>
+              ) : (
+                referenceImages
+                  .slice(0, Math.max(0, maxImagePromptCount))
+                  .map((image) => (
+                    <div
+                      key={image.id}
+                      className="glass relative aspect-square overflow-hidden rounded-lg w-14 border-2 border-white/30 hover:border-white/80 transition-all group cursor-pointer hover:cursor-zoom-in"
+                      onClick={() => onImageClick?.(image)}
                     >
-                      <FontAwesomeIcon icon={faXmark} className="h-2.5 w-2.5" />
-                    </button>
-                  </div>
-                ))}
+                      <img
+                        src={image.url}
+                        alt="Reference"
+                        className="h-full w-full object-cover"
+                      />
+                      <button
+                        onClick={(e) => {
+                          e.stopPropagation();
+                          handleRemoveReference(image.id);
+                        }}
+                        className="opacity-0 group-hover:opacity-100 absolute right-[2px] top-[2px] flex h-5 w-5 items-center justify-center rounded-full bg-black/50 hover:bg-red/70 text-white backdrop-blur-md transition-colors hover:bg-black cursor-pointer"
+                      >
+                        <FontAwesomeIcon
+                          icon={faXmark}
+                          className="h-2.5 w-2.5"
+                        />
+                      </button>
+                    </div>
+                  ))
+              )}
               {uploadingImages
                 .slice(
                   0,
@@ -414,7 +549,7 @@ export const ImagePromptRow = ({
               )}
             </div>
           </div>
-          {isVideo && (
+          {isVideo && showEndFrameSection && (
             <div className="flex gap-2">
               <div className="flex flex-col grow gap-1">
                 <div className="flex items-center gap-2 opacity-90">
diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxVideo.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxVideo.tsx
index c3dacdec48..e397918a14 100644
--- a/frontend/libs/components/promptbox/src/lib/PromptBoxVideo.tsx
+++ b/frontend/libs/components/promptbox/src/lib/PromptBoxVideo.tsx
@@ -222,6 +222,13 @@ export const PromptBoxVideo = ({
 
   const [isGalleryModalOpen, setIsGalleryModalOpen] = useState(false);
 
+  // Hide/clear ending frame if model doesn't support it
+  useEffect(() => {
+    if (selectedModel && !selectedModel.endFrame && endFrameImage) {
+      setEndFrameImage(undefined);
+    }
+  }, [selectedModel, endFrameImage, setEndFrameImage]);
+
   return (
     <>
       <Modal
@@ -255,7 +262,8 @@ export const PromptBoxVideo = ({
             uploadImage={uploadImage}
             endFrameImage={endFrameImage}
             setEndFrameImage={setEndFrameImage}
-            allowUploadEnd
+            allowUploadEnd={!!selectedModel?.endFrame}
+            showEndFrameSection={!!selectedModel?.endFrame}
           />
         )}
         <div