diff --git a/src/lib/nodes/ComfyActionNodes.ts b/src/lib/nodes/ComfyActionNodes.ts
index 09bc68e..2da838a 100644
--- a/src/lib/nodes/ComfyActionNodes.ts
+++ b/src/lib/nodes/ComfyActionNodes.ts
@@ -8,7 +8,7 @@ import ComfyGraphNode, { type ComfyGraphNodeProperties } from "./ComfyGraphNode"
 import type { ComfyWidgetNode } from "$lib/nodes/widgets";
 import type { NotifyOptions } from "$lib/notify";
 import type { FileData as GradioFileData } from "@gradio/upload";
-import { type ComfyExecutionResult, type ComfyImageLocation, convertComfyOutputToGradio, type ComfyUploadImageAPIResponse } from "$lib/utils";
+import { type ComfyExecutionResult, type ComfyImageLocation, convertComfyOutputToGradio, type ComfyUploadImageAPIResponse, parseWhateverIntoComfyImageLocations } from "$lib/utils";
 
 export class ComfyQueueEvents extends ComfyGraphNode {
     static slotLayout: SlotLayout = {
@@ -598,27 +598,7 @@ export class ComfySetPromptThumbnailsAction extends ComfyGraphNode {
 
     override getPromptThumbnails(): ComfyImageLocation[] | null {
         const data = this.getInputData(0)
-
-        const folderType = this.properties.folderType || "input";
-
-        const convertString = (s: string): ComfyImageLocation => {
-            return { filename: data, subfolder: "", type: folderType }
-        }
-
-        if (typeof data === "string") {
-            return [convertString(data)]
-        }
-        else if (data != null && typeof data === "object") {
-            if ("filename" in data && "type" in data)
-                return [data as ComfyImageLocation];
-        }
-        else if (Array.isArray(data) && data.length > 0) {
-            if (typeof data[0] === "string")
-                return data.map(convertString)
-            else if (typeof data[0] === "object" && "filename" in data[0] && "type" in data[0])
-                return data as ComfyImageLocation[]
-        }
-        return null;
+        return parseWhateverIntoComfyImageLocations(data);
     }
 }
 
diff --git a/src/lib/nodes/ComfyPickImageNode.ts b/src/lib/nodes/ComfyPickImageNode.ts
index 862fef5..2253a28 100644
--- a/src/lib/nodes/ComfyPickImageNode.ts
+++ b/src/lib/nodes/ComfyPickImageNode.ts
@@ -63,7 +63,7 @@ export default class ComfyPickImageNode extends ComfyGraphNode {
 
     override onExecute() {
         const data = this.getInputData(0)
-        const index = this.getInputData(1)
+        const index = this.getInputData(1) || 0
         this.setValue(data, index);
 
         if (this._image == null) {
diff --git a/src/lib/utils.ts b/src/lib/utils.ts
index dc53262..4509436 100644
--- a/src/lib/utils.ts
+++ b/src/lib/utils.ts
@@ -419,6 +419,12 @@ export function executionResultToImageMetadata(result: ComfyExecutionResult): Co
     return result.images.map(comfyFileToComfyBoxMetadata)
 }
 
+export function isComfyImageLocation(param: any): param is ComfyImageLocation {
+    return param != null && typeof param === "object"
+        && typeof param.filename === "string"
+        && typeof param.type === "string"
+}
+
 export function parseWhateverIntoImageMetadata(param: any): ComfyBoxImageMetadata[] | null {
     let meta: ComfyBoxImageMetadata[] | null = null
 
@@ -429,12 +435,26 @@ export function parseWhateverIntoImageMetadata(param: any): ComfyBoxImageMetadat
         meta = param
     }
     else if (isComfyExecutionResult(param)) {
-        meta = executionResultToImageMetadata(param)
+        meta = executionResultToImageMetadata(param);
+    }
+    else if (isComfyImageLocation(param)) {
+        meta = [comfyFileToComfyBoxMetadata(param)]
+    }
+    else if (Array.isArray(param) && param.every(isComfyImageLocation)) {
+        meta = param.map(comfyFileToComfyBoxMetadata)
     }
 
     return meta;
 }
 
+export function parseWhateverIntoComfyImageLocations(param: any): ComfyImageLocation[] | null {
+    const meta = parseWhateverIntoImageMetadata(param);
+    if (!Array.isArray(meta))
+        return null
+
+    return meta.map(m => m.comfyUIFile);
+}
+
 export function comfyBoxImageToComfyFile(image: ComfyBoxImageMetadata): ComfyImageLocation {
     return image.comfyUIFile
 }