From a3d64ccc02859739b1a934b9b81ba1d388615ddb Mon Sep 17 00:00:00 2001 From: Baptiste Adrien Date: Sat, 14 Jan 2023 18:22:46 +0100 Subject: [PATCH] feat: add 4k support (#36) --- .env.example | 1 + README.md | 1 + .../20230114145652_hd_status/migration.sql | 5 + .../migration.sql | 2 + .../20230114160322_url_hd/migration.sql | 2 + prisma/schema.prisma | 11 +- src/components/home/Pricing.tsx | 2 +- src/components/projects/FormPayment.tsx | 4 +- src/components/projects/shot/ShotCard.tsx | 112 +++++++++++++++--- src/components/projects/shot/ShotImage.tsx | 4 +- .../[id]/predictions/[predictionId]/hd.ts | 75 ++++++++++++ .../index.ts} | 0 12 files changed, 196 insertions(+), 23 deletions(-) create mode 100644 prisma/migrations/20230114145652_hd_status/migration.sql create mode 100644 prisma/migrations/20230114151121_hd_prediction_id/migration.sql create mode 100644 prisma/migrations/20230114160322_url_hd/migration.sql create mode 100644 src/pages/api/projects/[id]/predictions/[predictionId]/hd.ts rename src/pages/api/projects/[id]/predictions/{[predictionId].tsx => [predictionId]/index.ts} (100%) diff --git a/.env.example b/.env.example index e87aa3d..f856849 100644 --- a/.env.example +++ b/.env.example @@ -8,6 +8,7 @@ REPLICATE_API_TOKEN= REPLICATE_USERNAME= REPLICATE_MAX_TRAIN_STEPS=3000 REPLICATE_NEGATIVE_PROMPT="cropped face, cover face, cover visage, mutated hands" +REPLICATE_HD_VERSION_MODEL_ID= NEXT_PUBLIC_REPLICATE_INSTANCE_TOKEN= SECRET= EMAIL_FROM= diff --git a/README.md b/README.md index e7fbbca..de5c12f 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,7 @@ REPLICATE_API_TOKEN= REPLICATE_USERNAME= REPLICATE_MAX_TRAIN_STEPS=3000 REPLICATE_NEGATIVE_PROMPT= +REPLICATE_HD_VERSION_MODEL_ID= // Replicate instance token (should be rare) NEXT_PUBLIC_REPLICATE_INSTANCE_TOKEN= diff --git a/prisma/migrations/20230114145652_hd_status/migration.sql b/prisma/migrations/20230114145652_hd_status/migration.sql new file mode 100644 index 0000000..d2c012e --- /dev/null +++ b/prisma/migrations/20230114145652_hd_status/migration.sql @@ -0,0 +1,5 @@ +-- CreateEnum +CREATE TYPE "HdStatus" AS ENUM ('NO', 'PENDING', 'PROCESSED'); + +-- AlterTable +ALTER TABLE "Shot" ADD COLUMN "hdStatus" "HdStatus" NOT NULL DEFAULT 'NO'; diff --git a/prisma/migrations/20230114151121_hd_prediction_id/migration.sql b/prisma/migrations/20230114151121_hd_prediction_id/migration.sql new file mode 100644 index 0000000..3a97109 --- /dev/null +++ b/prisma/migrations/20230114151121_hd_prediction_id/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "Shot" ADD COLUMN "hdPredictionId" TEXT; diff --git a/prisma/migrations/20230114160322_url_hd/migration.sql b/prisma/migrations/20230114160322_url_hd/migration.sql new file mode 100644 index 0000000..b7e3fb1 --- /dev/null +++ b/prisma/migrations/20230114160322_url_hd/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "Shot" ADD COLUMN "hdOutputUrl" TEXT; diff --git a/prisma/schema.prisma b/prisma/schema.prisma index aa8eff3..386502b 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -7,6 +7,12 @@ generator client { provider = "prisma-client-js" } +enum HdStatus { + NO + PENDING + PROCESSED +} + model Account { id String @id @default(cuid()) userId String @map("user_id") @@ -76,7 +82,7 @@ model Project { userId String? shots Shot[] credits Int @default(100) - promptWizardCredits Int @default(30) + promptWizardCredits Int @default(20) Payment Payment[] } @@ -93,6 +99,9 @@ model Shot { bookmarked Boolean? @default(false) blurhash String? seed Int? + hdStatus HdStatus @default(NO) + hdPredictionId String? + hdOutputUrl String? } model Payment { diff --git a/src/components/home/Pricing.tsx b/src/components/home/Pricing.tsx index d3a4579..225efcc 100644 --- a/src/components/home/Pricing.tsx +++ b/src/components/home/Pricing.tsx @@ -74,7 +74,7 @@ const Pricing = () => { 1 Studio with a custom trained model - {process.env.NEXT_PUBLIC_STUDIO_SHOT_AMOUNT} avatars + {process.env.NEXT_PUBLIC_STUDIO_SHOT_AMOUNT} avatars 4K generation diff --git a/src/components/projects/FormPayment.tsx b/src/components/projects/FormPayment.tsx index 1a1c8d0..e4a50f5 100644 --- a/src/components/projects/FormPayment.tsx +++ b/src/components/projects/FormPayment.tsx @@ -76,8 +76,8 @@ const FormPayment = ({ 1 Studio with a custom trained model - {process.env.NEXT_PUBLIC_STUDIO_SHOT_AMOUNT} avatars - generation (512x512 resolution) + {process.env.NEXT_PUBLIC_STUDIO_SHOT_AMOUNT} avatars 4K + generation 30 AI prompt assists diff --git a/src/components/projects/shot/ShotCard.tsx b/src/components/projects/shot/ShotCard.tsx index b9eada3..6588c56 100644 --- a/src/components/projects/shot/ShotCard.tsx +++ b/src/components/projects/shot/ShotCard.tsx @@ -23,10 +23,27 @@ import { BsHeart, BsHeartFill } from "react-icons/bs"; import { HiDownload } from "react-icons/hi"; import { IoMdCheckmarkCircleOutline } from "react-icons/io"; import { MdOutlineModelTraining } from "react-icons/md"; +import { Ri4KFill } from "react-icons/ri"; import { useMutation, useQuery } from "react-query"; import ShotImage from "./ShotImage"; import { TbFaceIdError } from "react-icons/tb"; +const getHdLabel = (shot: Shot, isHd: boolean) => { + if (shot.hdStatus === "NO") { + return "Generate in 4K"; + } + + if (shot.hdStatus === "PENDING") { + return "4K in progress"; + } + + if (shot.hdStatus === "PROCESSED" && isHd) { + return "Show standard resolution"; + } + + return "Show 4K"; +}; + const ShotCard = ({ shot: initialShot, handleSeed, @@ -37,6 +54,7 @@ const ShotCard = ({ const { onCopy, hasCopied } = useClipboard(initialShot.prompt); const { query } = useRouter(); const [shot, setShot] = useState(initialShot); + const [isHd, setIsHd] = useState(Boolean(shot.hdOutputUrl)); const { mutate: bookmark, isLoading } = useMutation( `update-shot-${initialShot.id}`, @@ -54,6 +72,19 @@ const ShotCard = ({ } ); + const { mutate: createdHd, isLoading: isCreatingHd } = useMutation( + `create-hd-${initialShot.id}`, + () => + axios.post<{ shot: Shot }>( + `/api/projects/${query.id}/predictions/${initialShot.id}/hd` + ), + { + onSuccess: (response) => { + setShot(response.data.shot); + }, + } + ); + useQuery( `shot-${initialShot.id}`, () => @@ -65,10 +96,33 @@ const ShotCard = ({ { refetchInterval: (data) => (data?.shot.outputUrl ? false : 5000), refetchOnWindowFocus: false, - enabled: !initialShot.outputUrl, + enabled: !initialShot.outputUrl && initialShot.status !== "failed", + initialData: { shot: initialShot }, + onSuccess: (response) => { + setShot(response.shot); + }, + } + ); + + useQuery( + `shot-hd-${initialShot.id}`, + () => + axios + .get<{ shot: Shot }>( + `/api/projects/${query.id}/predictions/${initialShot.id}/hd` + ) + .then((res) => res.data), + { + refetchInterval: (data) => + data?.shot.hdStatus !== "PENDING" ? false : 5000, + refetchOnWindowFocus: false, + enabled: shot.hdStatus === "PENDING", initialData: { shot: initialShot }, onSuccess: (response) => { setShot(response.shot); + if (response.shot.hdOutputUrl) { + setIsHd(true); + } }, } ); @@ -82,7 +136,7 @@ const ShotCard = ({ position="relative" > {shot.outputUrl ? ( - + ) : ( @@ -104,10 +158,7 @@ const ShotCard = ({ )} - - - {formatRelative(new Date(shot.createdAt), new Date())} - + {shot.seed && shot.outputUrl && ( @@ -129,17 +180,41 @@ const ShotCard = ({ )} {shot.outputUrl && ( - } - /> + <> + } + /> + + } + color={isHd ? "red.400" : "gray.600"} + isLoading={shot.hdStatus === "PENDING" || isCreatingHd} + onClick={() => { + if (shot.hdStatus === "NO") { + createdHd(); + } else if ( + shot.hdStatus === "PROCESSED" && + shot.hdOutputUrl + ) { + setIsHd(!isHd); + } + }} + size="sm" + variant="ghost" + aria-label="Make 4K" + fontSize="lg" + /> + + )} + - + + + {formatRelative(new Date(shot.createdAt), new Date())} +