Skip to content

Commit

Permalink
chore: regeneration
Browse files Browse the repository at this point in the history
  • Loading branch information
@tomwisecodes
tomwisecodes committed Dec 5, 2024
1 parent 60fc633 commit a39930e
Show file tree
Hide file tree
Showing 4 changed files with 406 additions and 47 deletions.
155 changes: 109 additions & 46 deletions apps/nextjs/src/app/image-spike/[slug]/images.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import Link from "next/link";
import type { ImageResponse } from "types/imageTypes";

import LoadingWheel from "@/components/LoadingWheel";
import { RegenerationForm } from "@/components/RegenerationForm";

export type Cycle = {
title: string;
Expand Down Expand Up @@ -70,8 +71,51 @@ const ImagesPage = ({ pageData }) => {
},
]);

// In your ImagesPage component, add this state
const [isRegenerating, setIsRegenerating] = useState(false);

const { bestImage, findBestImage } = useBestImage({ pageData });
const { fetchImages, availableSources } = useImageSearch({ pageData });
const { fetchImages, availableSources, regenerateImageWithAnalysis } =
useImageSearch({ pageData });

// Add this function to handle regeneration
const handleRegeneration = async (columnId: string, feedback: string) => {
const column = comparisonColumns.find((col) => col.id === columnId);
if (!column?.imageSearchBatch?.[0]) return;

const image = column.imageSearchBatch[0];

updateColumn(columnId, { isLoading: true });

if (!image.imageSource) {
throw new Error("Image source is required to regenerate image.");
}

try {
setIsRegenerating(true);
const regeneratedImage = await regenerateImageWithAnalysis(
image.url,
selectedImagePrompt,
feedback,
image.imageSource.toLowerCase().includes("Stability AI")
? "stability"
: "openai",
);

setIsRegenerating(false);
updateColumn(columnId, {
imageSearchBatch: [regeneratedImage],
});
} catch (error) {
setIsRegenerating(false);
console.error("Error regenerating image:", error);
updateColumn(columnId, {
error: "Failed to regenerate image",
});
} finally {
updateColumn(columnId, { isLoading: false });
}
};

if (
!pageData?.lessonPlan?.cycle1?.explanation?.imagePrompt ||
Expand Down Expand Up @@ -196,8 +240,8 @@ const ImagesPage = ({ pageData }) => {
{promptConstructor(
prompt,
pageData.title,
pageData.keyStage,
pageData.subject,
pageData.keyStage,
pageData.lessonPlan,
)}
</p>
Expand Down Expand Up @@ -348,58 +392,77 @@ const ImagesPage = ({ pageData }) => {

{column.imageSearchBatch && (
<div className="mt-6 space-y-6">
{column.imageSearchBatch?.map((image) => (
<div
key={image.id}
className="rounded-lg border bg-gray-50 p-4"
>
<div className="relative mb-4 h-48 w-full overflow-hidden rounded-lg">
<Image
src={image.url}
alt={image.alt || "Image"}
fill
className="object-cover"
/>
</div>
<div className="space-y-2 text-sm text-gray-600">
<p>
<span className="font-medium">License:</span>{" "}
{image.license}
</p>
<p>
<span className="font-medium">Score:</span>{" "}
{image.appropriatenessScore}
</p>
<p>
<span className="font-medium">Prompt used:</span>{" "}
{image.imagePrompt}
</p>
<p>
<span className="font-medium">Reasoning:</span>{" "}
{image.appropriatenessReasoning}
</p>
{image.photographer && (
{column.imageSearchBatch?.map((image) => {
console.log("iumage", image);
return (
<div
key={image.id}
className="rounded-lg border bg-gray-50 p-4"
>
<div className="relative mb-4 h-48 w-full rounded-lg object-contain">
{isRegenerating ? (
<LoadingWheel />
) : (
<Image
src={image.url}
alt={image.alt || "Image"}
fill
className="object-cover"
/>
)}
</div>
{(image.imageSource === "DAL-E" ||
image.imageSource?.includes(
"Stable Diffusion",
)) && (
<RegenerationForm
onSubmit={(feedback) =>
handleRegeneration(column.id, feedback)
}
imageSource={image.imageSource}
/>
)}
<div className="space-y-2 text-sm text-gray-600">
<p>
<span className="font-medium">By:</span>{" "}
{image.photographer}
<span className="font-medium">License:</span>{" "}
{image.license}
</p>
)}
{image.title && (
<p>
<span className="font-medium">Title:</span>{" "}
{image.title}
<span className="font-medium">Score:</span>{" "}
{image.appropriatenessScore}
</p>
<p>
<span className="font-medium">Prompt used:</span>{" "}
{image.imagePrompt}
</p>
)}
<div className="mt-4 border-t pt-4 text-xs">
<p>Fetch: {image.timing.fetch.toFixed(2)}ms</p>
<p>
Validation: {image.timing.validation.toFixed(2)}ms
<span className="font-medium">Reasoning:</span>{" "}
{image.appropriatenessReasoning}
</p>
<p>Total: {image.timing.total.toFixed(2)}ms</p>
{image.photographer && (
<p>
<span className="font-medium">By:</span>{" "}
{image.photographer}
</p>
)}
{image.title && (
<p>
<span className="font-medium">Title:</span>{" "}
{image.title}
</p>
)}
<div className="mt-4 border-t pt-4 text-xs">
<p>Fetch: {image.timing.fetch.toFixed(2)}ms</p>
<p>
Validation: {image.timing.validation.toFixed(2)}
ms
</p>
<p>Total: {image.timing.total.toFixed(2)}ms</p>
</div>
</div>
</div>
</div>
))}
);
})}
</div>
)}
</div>
Expand Down
96 changes: 96 additions & 0 deletions apps/nextjs/src/components/RegenerationForm.tsx
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
import React, { useState } from "react";

import { OakIcon } from "@oaknational/oak-components";
import * as Dialog from "@radix-ui/react-dialog";

import LoadingWheel from "@/components/LoadingWheel";

import { Icon } from "./Icon";

interface RegenerationFormProps {
onSubmit: (feedback: string) => Promise<void>;
imageSource: string;
}

export const RegenerationForm: React.FC<RegenerationFormProps> = ({
onSubmit,
imageSource,
}) => {
const [isOpen, setIsOpen] = useState(false);
const [feedback, setFeedback] = useState("");
const [isLoading, setIsLoading] = useState(false);

const handleSubmit = async (e: React.FormEvent) => {
e.preventDefault();
setIsLoading(true);
try {
await onSubmit(feedback);
setIsOpen(false);
setFeedback("");
} catch (error) {
console.error("Error regenerating image:", error);
} finally {
setIsLoading(false);
}
};

return (
<Dialog.Root open={isOpen} onOpenChange={setIsOpen}>
<Dialog.Trigger asChild>
<button className="bg-blue-600 hover:bg-blue-700 focus:ring-blue-500 mt-4 flex w-full items-center justify-center gap-2 rounded-lg px-4 py-2 text-sm font-medium text-white focus:outline-none focus:ring-2 focus:ring-offset-2">
<Icon icon="reload" size="sm" color="white" />
Regenerate with {imageSource}
</button>
</Dialog.Trigger>

<Dialog.Portal>
<Dialog.Overlay className="fixed inset-0 bg-black/50" />
<Dialog.Content className="fixed left-1/2 top-1/2 w-full max-w-md -translate-x-1/2 -translate-y-1/2 rounded-lg bg-white p-6 shadow-xl focus:outline-none">
<Dialog.Title className="text-lg font-medium text-gray-900">
Regenerate Image
</Dialog.Title>

<form onSubmit={handleSubmit} className="mt-4">
<div className="space-y-2">
<label
htmlFor="feedback"
className="block text-sm font-medium text-gray-700"
>
What would you like to change about this image?
</label>
<textarea
id="feedback"
placeholder="e.g., Make the colors more vibrant, add more detail to..."
value={feedback}
onChange={(e) => setFeedback(e.target.value)}
className="focus:border-blue-500 focus:ring-blue-500 min-h-[100px] w-full rounded-md border border-gray-300 px-3 py-2 text-sm placeholder:text-gray-400 focus:outline-none focus:ring-1"
required
/>
</div>

<div className="mt-4 flex justify-end gap-3">
<button
type="button"
onClick={() => setIsOpen(false)}
className="focus:ring-blue-500 rounded-md border border-gray-300 bg-white px-4 py-2 text-sm font-medium text-gray-700 hover:bg-gray-50 focus:outline-none focus:ring-2 focus:ring-offset-2"
>
Cancel
</button>
<button
type="submit"
disabled={isLoading || !feedback.trim()}
className="bg-blue-600 hover:bg-blue-700 focus:ring-blue-500 rounded-md px-4 py-2 text-sm font-medium text-white focus:outline-none focus:ring-2 focus:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50"
>
{isLoading ? <LoadingWheel /> : "Submit"}
</button>
</div>
</form>

<Dialog.Close className="focus:ring-blue-500 absolute right-4 top-4 text-gray-400 hover:text-gray-500 focus:outline-none focus:ring-2">
<OakIcon iconName="cross" />
</Dialog.Close>
</Dialog.Content>
</Dialog.Portal>
</Dialog.Root>
);
};
67 changes: 67 additions & 0 deletions apps/nextjs/src/hooks/useImageSearch.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@ export const useImageSearch = ({ pageData }: ImageSearchHookProps) => {
const flickrMutation = trpc.imageSearch.getImagesFromFlickr.useMutation();
const unsplashMutation = trpc.imageSearch.getImagesFromUnsplash.useMutation();
const stableDiffusionCoreMutation = trpc.imageGen.stableDifCore.useMutation();
const analyzeAndRegenerateMutation =
trpc.imageGen.analyzeAndRegenerate.useMutation();

const daleMutation = trpc.imageGen.openAi.useMutation();
const validateImageMutation = trpc.imageGen.validateImage.useMutation();
const validateImagesInParallel =
Expand All @@ -34,6 +37,69 @@ export const useImageSearch = ({ pageData }: ImageSearchHookProps) => {
Cloudinary: cloudinaryMutation,
};

const regenerateImageWithAnalysis = async (
originalImageUrl: string,
originalPrompt: string,
feedback: string,
provider: "openai" | "stability",
): Promise<ImageResponse> => {
try {
if (
!pageData?.title ||
!pageData?.keyStage ||
!pageData?.subject ||
!pageData?.lessonPlan
) {
throw new Error("Missing required page data");
}

const startTime = performance.now();
const fetchStart = performance.now();

const response = await analyzeAndRegenerateMutation.mutateAsync({
originalImageUrl,
originalPrompt,
feedback,
lessonTitle: pageData.title,
subject: pageData.subject,
keyStage: pageData.keyStage,
lessonPlan: pageData.lessonPlan,
provider,
});

const fetchEnd = performance.now();

// Validate the regenerated image
const validationStart = performance.now();
const validationResult = await validateSingleImage(
response,
originalPrompt,
true,
);
const endTime = performance.now();
// generate random string for id
const id = Math.random().toString(36).substr(2, 9);

return {
id: id,
url: response,
license: provider === "openai" ? "OpenAI DALL-E 3" : "Stability AI",
imagePrompt: validationResult.metadata.promptUsed ?? originalPrompt,
appropriatenessScore: validationResult.metadata.appropriatenessScore,
appropriatenessReasoning: validationResult.metadata.validationReasoning,
imageSource: provider === "openai" ? "OpenAI DALL-E 3" : "Stability AI",
timing: {
total: endTime - startTime,
fetch: fetchEnd - fetchStart,
validation: endTime - validationStart,
},
};
} catch (error) {
console.error("Error regenerating image:", error);
throw error;
}
};

const validateSingleImage = async (
imageUrl: string,
prompt: string,
Expand Down Expand Up @@ -303,5 +369,6 @@ export const useImageSearch = ({ pageData }: ImageSearchHookProps) => {
return {
fetchImages,
availableSources: trpcMutations,
regenerateImageWithAnalysis,
};
};
Loading

0 comments on commit a39930e

Please sign in to comment.