upscaling poggers

package.json

@@ -21,6 +21,7 @@
     "@stripe/stripe-js": "^1.46.0",
     "@supabase/auth-helpers-react": "^0.3.1",
     "@supabase/supabase-js": "^2.1.3",
+    "@types/file-saver": "^2.0.5",
     "@types/node": "18.11.10",
     "@types/react": "18.0.26",
     "@types/react-dom": "18.0.9",
@@ -30,6 +31,7 @@
     "date-fns": "^2.29.3",
     "eslint": "8.29.0",
     "eslint-config-next": "13.0.6",
+    "file-saver": "^2.0.5",
     "framer-motion": "^7.6.19",
     "image-blob-reduce": "^4.1.0",
     "jsonwebtoken": "^8.5.1",

src/components/projects/CustomZoomContent.tsx

@@ -1,11 +1,23 @@
+import { getFullShotUrl } from "@/core/utils/bucketHelpers";
 import { ShotsPick } from "@/pages/api/projects";
 import { Box, Text } from "@chakra-ui/layout";
-import { Button, Portal, useDisclosure } from "@chakra-ui/react";
+import {
+  Button,
+  Popover,
+  PopoverArrow,
+  PopoverBody,
+  PopoverCloseButton,
+  PopoverContent,
+  PopoverFooter,
+  PopoverHeader,
+  PopoverTrigger,
+  Spinner,
+  useDisclosure,
+} from "@chakra-ui/react";
 import axios from "axios";
-import { useRouter } from "next/router";
-import { ReactElement, FC } from "react";
+import { ReactElement, FC, useState, useEffect } from "react";
 import { useMutation } from "react-query";
-import ConfirmationModal from "../shared/ConfirmationModal";
+import { saveAs } from "file-saver";
 
 declare const enum ModalState {
   LOADED = "LOADED",
@@ -20,6 +32,7 @@ export type ZoomContentProps = {
   modalState: ModalState;
   shot: ShotsPick;
   onUnzoom: () => void;
+  refetchShot: () => void;
 };
 
 type CustomZoomContentProps = {} & ZoomContentProps;
@@ -29,25 +42,37 @@ const CustomZoomContent: FC<CustomZoomContentProps> = ({
   modalState,
   img,
   shot,
+  refetchShot,
 }) => {
+  const [isLoadingUpscale, setIsLoadingUpscale] = useState(
+    shot.upscaleId && !shot.upscaledImageUrl
+  );
   // If we don't have this check, the unzoom button / description will render
   // before the modals transition is complete.
   const modalLoaded = modalState === "LOADED";
-  const { isOpen, onClose, onOpen } = useDisclosure();
+  const { isOpen, onClose, onOpen, onToggle } = useDisclosure();
 
-  const {
-    mutate: mutateUpscale,
-    isLoading: isLoadingUpscale,
-    isSuccess,
-  } = useMutation(
+  const { mutate: mutateUpscale, isLoading: isLoading } = useMutation(
     `upscale-shot-${shot.id}`,
-    (shot: ShotsPick) => axios.post(`/api/shots/upscale?id=${shot.id}`),
+    (shot: ShotsPick) =>
+      axios.post(`/api/shots/upscale?id=${shot.id}`).then((res) => res.data),
     {
       onSuccess: () => {
-        // refresh shots later
+        setIsLoadingUpscale(true);
+        refetchShot();
       },
     }
   );
+
+  useEffect(() => {
+    if (!isLoadingUpscale) return;
+    const interval = setInterval(() => {
+      refetchShot();
+    }, 5000);
+
+    return () => clearInterval(interval);
+  }, [isLoadingUpscale]);
+
   return (
     <>
       {modalLoaded && buttonUnzoom}
@@ -60,23 +85,95 @@ const CustomZoomContent: FC<CustomZoomContentProps> = ({
           </Text>
         )}
       </Box>
-      <Box display={"flex"} alignItems={"center"} justifyContent={"center"}>
-        {modalLoaded && (
-          <Button variant={"brand"} onClick={onOpen}>
-            Upscale
-          </Button>
-        )}
-      </Box>
+
       <Box
-        display={isOpen ? "flex" : "none"}
+        display={modalLoaded ? "flex" : "none"}
+        mt={4}
         alignItems="center"
         justifyContent="center"
       >
-        <ConfirmationModal
-          isOpen={isOpen}
-          onClose={onClose}
-          onConfirm={() => mutateUpscale(shot)}
-        />
+        {shot.imageUrl && (
+          <Button
+            variant="brand"
+            mr={4}
+            onClick={async () => {
+              saveAs(getFullShotUrl(shot, false), "image.png");
+            }}
+          >
+            Download SD
+          </Button>
+        )}
+        {shot.upscaledImageUrl && (
+          <Button
+            variant="brand"
+            mr={4}
+            onClick={async () => {
+              saveAs(getFullShotUrl(shot, true), "image.png");
+            }}
+          >
+            Download HD
+          </Button>
+        )}
+        {isLoadingUpscale && (
+          <Button variant="brand" disabled={true}>
+            <Spinner mr={2} /> <Text>Upscaling...</Text>
+          </Button>
+        )}
+        {!isLoadingUpscale && (
+          <Box
+            display={modalLoaded && !shot.upscaledImageUrl ? "flex" : "none"}
+            alignItems="center"
+            justifyContent="center"
+          >
+            <Popover
+              placement="bottom"
+              size="md"
+              isOpen={isOpen}
+              onClose={onClose}
+              arrowShadowColor="gray.300"
+            >
+              <PopoverTrigger>
+                <Button variant="brand" onClick={onToggle}>
+                  Upscale Image
+                </Button>
+              </PopoverTrigger>
+              <PopoverContent>
+                <PopoverArrow />
+                <PopoverHeader>Upscale Image</PopoverHeader>
+                <PopoverCloseButton />
+
+                <PopoverBody>
+                  <Text fontSize="sm" mb={4}>
+                    Are you sure you want to spend 1 credit to upscale this
+                    image?
+                  </Text>
+                </PopoverBody>
+                <PopoverFooter>
+                  <Box
+                    display="flex"
+                    justifyContent={"space-around"}
+                    width="100%"
+                  >
+                    <Button size="lg" variant="brand" onClick={() => onClose()}>
+                      Cancel
+                    </Button>
+                    <Button
+                      size="lg"
+                      variant="brand"
+                      onClick={() => {
+                        onClose();
+                        setIsLoadingUpscale(true);
+                        mutateUpscale(shot);
+                      }}
+                    >
+                      Yes
+                    </Button>
+                  </Box>
+                </PopoverFooter>
+              </PopoverContent>
+            </Popover>
+          </Box>
+        )}
       </Box>
     </>
   );

src/components/projects/ProjectCard.tsx

@@ -50,7 +50,6 @@ const ProjectCard = ({
       },
     }
   );
-  const [startTime, setStartTime] = useState<number>(0);
 
   const isWaitingPayment = !project.stripePaymentId;
   const isWaitingTraining =
@@ -66,8 +65,8 @@ const ProjectCard = ({
   };
 
   const checkBackTime = formatRelative(
-    new Date(new Date(startTime).getTime() + 30 * 60 * 1000),
-    new Date(startTime)
+    new Date(new Date(project.createdAt).getTime() + 30 * 60 * 1000),
+    new Date(project.createdAt)
   )
     .replace("at", "around")
     .replace("today", "");

src/components/projects/ShotCard.tsx

@@ -15,7 +15,15 @@ const ShotCard = ({
   shot: ShotsPick;
   projectId: string;
 }) => {
-  const { data } = useQuery(
+  // if shot has failed status, dont refetch,
+  // if shot has upscaleId and no upscaledImageUrl, refetch
+  // if shot has no imageUrl, refetch
+  const enableFetch =
+    initialShot.status !== "failed" &&
+    ((initialShot.upscaleId && !initialShot.upscaledImageUrl) ||
+      !initialShot.imageUrl);
+
+  const { data, refetch: refetchShot } = useQuery(
     `shot-${initialShot.id}`,
     () =>
       axios
@@ -24,22 +32,15 @@ const ShotCard = ({
         )
         .then((res) => res.data),
     {
-      refetchInterval: (data) =>
-        !data?.shot.imageUrl ||
-        (data?.shot.upscaleId && !data?.shot.upscaledImageUrl)
-          ? 5000
-          : false,
+      refetchInterval: (data) => (enableFetch ? 5000 : false),
       refetchOnWindowFocus: false,
-      enabled:
-        (initialShot.upscaleId &&
-          initialShot.status !== "failed" &&
-          !initialShot.upscaledImageUrl) ||
-        (!initialShot.imageUrl && initialShot.status !== "failed"),
+      enabled: enableFetch,
       initialData: { shot: initialShot },
     }
   );
 
   const shot = data!.shot;
+
   return (
     <Box key={shot.id} backgroundColor="gray.100" overflow="hidden">
       {shot.status === "failed" && (
@@ -56,6 +57,7 @@ const ShotCard = ({
               // Type script being annoying, i'll figure it out later.
               {...(zoomContentProps as unknown as ZoomContentProps)}
               shot={shot}
+              refetchShot={refetchShot}
             />
           )}
         >

src/core/clients/replicate.ts

@@ -347,7 +347,7 @@ export type PredictionResponse = {
   created_at: Date;
   started_at: Date | null;
   completed_at: Date | null;
-  status: "starting" | "pushing" | "succeeded";
+  status: "starting" | "pushing" | "succeeded" | "failed";
   input: any; // the input is free form it comes
   output: any;
   error: any;

src/pages/api/projects/[id]/predictions/[predictionId].tsx

@@ -18,6 +18,15 @@ const handler = async (req: NextApiRequest, res: NextApiResponse) => {
     const { data: prediction } = await replicateClient.get<PredictionResponse>(
       `https://api.replicate.com/v1/predictions/${fetchId}`
     );
+    if (prediction.status === "failed") {
+      await db.shot.update({
+        where: { id: shot.id },
+        data: {
+          status: prediction.status,
+          upscaleId: null,
+        },
+      });
+    }
     // If the initial shot status changes from the prediction, update the shot in database.
     if (shot.status !== prediction.status) {
       const outputUrl = shot.upscaleId

src/pages/dashboard.tsx

@@ -17,7 +17,6 @@ import { useQuery } from "react-query";
 import PageContainer from "@/components/layout/PageContainer";
 import { useRouter } from "next/router";
 import { useSession } from "next-auth/react";
-import { useEffect } from "react";
 import ExamplePictures from "@/components/dashboard/ExamplePictures";
 import { ProjectsGetResponse } from "./api/projects";