add login

.env

@@ -7,4 +7,17 @@ MICROSERVICE_API_SECRET_TOKEN="<USE YOUR OWN>"
 
 AI_TUBE_API_SECRET_JWT_KEY=""
 AI_TUBE_API_SECRET_JWT_ISSUER=""
-AI_TUBE_API_SECRET_JWT_AUDIENCE=""
+AI_TUBE_API_SECRET_JWT_AUDIENCE=""
+
+# ------------- HUGGING FACE OAUTH -------------
+NEXT_PUBLIC_ENABLE_HUGGING_FACE_OAUTH=""
+NEXT_PUBLIC_ENABLE_HUGGING_FACE_OAUTH_WALL=""
+NEXT_PUBLIC_HUGGING_FACE_OAUTH_CLIENT_ID=""
+
+# this one must be kept secret (and is unused for now)
+HUGGING_FACE_OAUTH_SECRET=""
+
+# ----------- RATE LIMIT -------
+ENABLE_RATE_LIMIT=""
+UPSTASH_REDIS_REST_URL=""
+UPSTASH_REDIS_REST_TOKEN=""

package-lock.json

@@ -10,6 +10,7 @@
       "dependencies": {
         "@aitube/clap": "0.0.17",
         "@aitube/client": "0.0.25",
+        "@huggingface/hub": "^0.15.0",
         "@radix-ui/react-accordion": "^1.1.2",
         "@radix-ui/react-avatar": "^1.0.4",
         "@radix-ui/react-checkbox": "^1.0.4",
@@ -32,6 +33,8 @@
         "@types/react": "18.3.0",
         "@types/react-dom": "18.3.0",
         "@types/uuid": "^9.0.8",
+        "@upstash/ratelimit": "^1.1.3",
+        "@upstash/redis": "^1.31.1",
         "autoprefixer": "10.4.17",
         "class-variance-authority": "^0.7.0",
         "clsx": "^2.1.0",
@@ -226,6 +229,17 @@
       "resolved": "https://registry.npmjs.org/@floating-ui/utils/-/utils-0.2.2.tgz",
       "integrity": "sha512-J4yDIIthosAsRZ5CPYP/jQvUAQtlZTTD/4suA08/FEnlxqW3sKS9iAhgsa9VYLZ6vDHn/ixJgIqRQPotoBjxIw=="
     },
+    "node_modules/@huggingface/hub": {
+      "version": "0.15.0",
+      "resolved": "https://registry.npmjs.org/@huggingface/hub/-/hub-0.15.0.tgz",
+      "integrity": "sha512-8jV+DjC68FXTNFCJeaKIa2e13rvfE4MBcJSlVtNOoA1cflLNmVBbta7iwKnMbUgdjW6DObztBLFneUcvZ3SHkQ==",
+      "dependencies": {
+        "hash-wasm": "^4.9.0"
+      },
+      "engines": {
+        "node": ">=18"
+      }
+    },
     "node_modules/@humanwhocodes/config-array": {
       "version": "0.11.14",
       "resolved": "https://registry.npmjs.org/@humanwhocodes/config-array/-/config-array-0.11.14.tgz",
@@ -2618,6 +2632,33 @@
       "resolved": "https://registry.npmjs.org/@ungap/structured-clone/-/structured-clone-1.2.0.tgz",
       "integrity": "sha512-zuVdFrMJiuCDQUMCzQaD6KL28MjnqqN8XnAqiEq9PNm/hCPTSGfrXCOfwj1ow4LFb/tNymJPwsNbVePc1xFqrQ=="
     },
+    "node_modules/@upstash/core-analytics": {
+      "version": "0.0.8",
+      "resolved": "https://registry.npmjs.org/@upstash/core-analytics/-/core-analytics-0.0.8.tgz",
+      "integrity": "sha512-MCJoF+Y8fkzq4NRLG7kEHjtGyMsZ2DICBdmEdwoK9umoSrfkzgBlYdZiHTIaewyt9PGaMZCHOasz0LAuMpxwxQ==",
+      "dependencies": {
+        "@upstash/redis": "^1.28.3"
+      },
+      "engines": {
+        "node": ">=16.0.0"
+      }
+    },
+    "node_modules/@upstash/ratelimit": {
+      "version": "1.1.3",
+      "resolved": "https://registry.npmjs.org/@upstash/ratelimit/-/ratelimit-1.1.3.tgz",
+      "integrity": "sha512-rl+GMvPdZJ9xPDIvIrqRl/g0nzAEaH75hwR5lXAKW8zPPplD/AeliDCHwuwcFCPIjg49FKyA1oc5H473WkVFrQ==",
+      "dependencies": {
+        "@upstash/core-analytics": "^0.0.8"
+      }
+    },
+    "node_modules/@upstash/redis": {
+      "version": "1.31.1",
+      "resolved": "https://registry.npmjs.org/@upstash/redis/-/redis-1.31.1.tgz",
+      "integrity": "sha512-lAsOo+kYjD5lpP+lH/nxHfzFYeCkWBwwKsyZZmh0AoOumBA9ZpS52Gorm7c2bmNu3UFijpPiLSFdW/nRdjbRpQ==",
+      "dependencies": {
+        "crypto-js": "^4.2.0"
+      }
+    },
     "node_modules/acorn": {
       "version": "8.11.3",
       "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.11.3.tgz",
@@ -3547,6 +3588,11 @@
         "node": ">= 8"
       }
     },
+    "node_modules/crypto-js": {
+      "version": "4.2.0",
+      "resolved": "https://registry.npmjs.org/crypto-js/-/crypto-js-4.2.0.tgz",
+      "integrity": "sha512-KALDyEYgpY+Rlob/iriUtjV6d5Eq+Y191A5g4UqLAi8CyGP9N1+FdVbkc1SxKc2r4YAYqG8JzO2KGL+AizD70Q=="
+    },
     "node_modules/cssesc": {
       "version": "3.0.0",
       "resolved": "https://registry.npmjs.org/cssesc/-/cssesc-3.0.0.tgz",
@@ -4822,6 +4868,11 @@
         "url": "https://github.com/sponsors/ljharb"
       }
     },
+    "node_modules/hash-wasm": {
+      "version": "4.11.0",
+      "resolved": "https://registry.npmjs.org/hash-wasm/-/hash-wasm-4.11.0.tgz",
+      "integrity": "sha512-HVusNXlVqHe0fzIzdQOGolnFN6mX/fqcrSAOcTBXdvzrXVHwTz11vXeKRmkR5gTuwVpvHZEIyKoePDvuAR+XwQ=="
+    },
     "node_modules/hasown": {
       "version": "2.0.2",
       "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.2.tgz",

package.json

@@ -11,6 +11,7 @@
   "dependencies": {
     "@aitube/clap": "0.0.17",
     "@aitube/client": "0.0.25",
+    "@huggingface/hub": "^0.15.0",
     "@radix-ui/react-accordion": "^1.1.2",
     "@radix-ui/react-avatar": "^1.0.4",
     "@radix-ui/react-checkbox": "^1.0.4",
@@ -33,6 +34,8 @@
     "@types/react": "18.3.0",
     "@types/react-dom": "18.3.0",
     "@types/uuid": "^9.0.8",
+    "@upstash/ratelimit": "^1.1.3",
+    "@upstash/redis": "^1.31.1",
     "autoprefixer": "10.4.17",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.0",

src/app/config.ts

@@ -5,4 +5,11 @@ export const defaultPrompt =
   // "beautiful footage of a Caribbean fishing village and bay, sail ships, during golden hour, no captions"
   "videogame gameplay footage, first person, exploring some mysterious ruins, no commentary"
 
-export const localStorageStoryDraftKey = "AI_STORIES_FACTORY_STORY_PROMPT_DRAFT"
+export const localStorageStoryDraftKey = "AI_STORIES_FACTORY_STORY_PROMPT_DRAFT"
+
+export const localStorageHuggingFaceOAuthKey = "AI_STORIES_FACTORY_HUGGING_FACE_OAUTH"
+
+export const oauthClientId = `${process.env.NEXT_PUBLIC_HUGGING_FACE_OAUTH_CLIENT_ID || ""}`
+export const oauthScopes = "openid profile inference-api"
+export const enableHuggingFaceOAuth = `${process.env.NEXT_PUBLIC_ENABLE_HUGGING_FACE_OAUTH || ""}` === "true"
+export const enableHuggingFaceOAuthWall = `${process.env.NEXT_PUBLIC_ENABLE_HUGGING_FACE_OAUTH_WALL || ""}` === "true"

src/app/main.tsx

@@ -20,6 +20,9 @@ import { LoadClapButton } from "@/components/interface/load-clap-button"
 import { SaveClapButton } from "@/components/interface/save-clap-button"
 import { useProcessors } from "@/lib/hooks/useProcessors"
 import { Characters } from "@/components/interface/characters"
+import { useOAuth } from "@/lib/oauth/useOAuth"
+import { useStore } from "./store"
+import { AuthWall } from "@/components/interface/auth-wall"
 
 export function Main() {
   const { storyPromptDraft, setStoryPromptDraft, promptDraftRef } = useStoryPromptDraft()
@@ -27,6 +30,10 @@ export function Main() {
   const { orientation, toggleOrientation } = useOrientation()
   const { handleSubmit } = useProcessors()
   useQueryStringParams()
+    
+  const showAuthWall = useStore(s => s.showAuthWall)
+  const { isLoggedIn, enableOAuthWall } = useOAuth({ debug: true })
+
   return (
     <TooltipProvider>
       <div className={cn(
@@ -275,7 +282,7 @@ export function Main() {
                         {/* END OF ORIENTATION SWITCH */}
                         <Button
                           onClick={handleSubmit}
-                          disabled={!storyPromptDraft || isBusy}
+                          disabled={!storyPromptDraft || isBusy || !isLoggedIn}
                           // variant="ghost"
                           className={cn(
                             `text-base md:text-lg lg:text-xl xl:text-2xl`,
@@ -308,6 +315,7 @@ export function Main() {
           <BottomBar />
         </div>
         <Toaster />
+        <AuthWall show={showAuthWall} />
       </div>
     </TooltipProvider>
   );

src/app/server/aitube/createClap.ts

@@ -1,10 +1,16 @@
 "use server"
 
+
+import { Ratelimit } from "@upstash/ratelimit"
+import { Redis } from "@upstash/redis"
 import { ClapProject, ClapMediaOrientation } from "@aitube/clap"
 import { createClap as apiCreateClap } from "@aitube/client"
 
 import { getToken } from "./getToken"
-import { RESOLUTION_LONG, RESOLUTION_SHORT } from "./config"
+import { RESOLUTION_LONG, RESOLUTION_SHORT } from "../config"
+import { getRateLimit } from "../redis/getRateLimit"
+
+const rateLimit = getRateLimit()
 
 export async function createClap({
   prompt = "",
@@ -15,6 +21,24 @@ export async function createClap({
   orientation?: ClapMediaOrientation
   turbo?: boolean
 }): Promise<ClapProject> {
+
+  // TODO: use 
+
+  /*
+  if (process.env.ENABLE_RATE_LIMIT) {
+    const user = "anon"
+    const userRateLimitResult = await rateLimit.limit(user)
+
+    // result.limit
+    if (!userRateLimitResult.success) {
+      console.log(`blocking user ${user} who requested "${prompt}${prompt.length > 60 ? "..." : ""}"`)
+      throw new Error(`Rate Limit Reached`)
+    } else {
+      console.log(`allowing user ${user}: who requested "${prompt}${prompt.length >630 ? "..." : ""}"`)
+    }
+  }
+  */
+
   const clap: ClapProject = await apiCreateClap({
     prompt: prompt.slice(0, 512),
 

src/app/server/{aitube/config.ts → config.ts}

@@ -8,4 +8,4 @@ export const serverHuggingfaceApiKey = `${process.env.HF_API_TOKEN || ""}`
 export const RESOLUTION_LONG = 896 // 832 // 768
 export const RESOLUTION_SHORT = 512 // 448 // 384
 
-// ValueError: `height` and `width` have to be divisible by 8 but are 512 and 1.
+// ValueError: `height` and `width` have to be divisible by 8 but are 512 and 1.

src/app/server/redis/getRateLimit.ts

@@ -0,0 +1,28 @@
+import { Ratelimit } from "@upstash/ratelimit"
+
+import { redis } from "./redis"
+
+// Create a global ratelimiter for all users, that allows 14 requests per 60 seconds
+// 14 is roughly the number of requests that can be handled by the server
+/*
+const rateLimitGlobal = new Ratelimit({
+  redis,
+  limiter: Ratelimit.slidingWindow(14, "60 s"),
+  analytics: true,
+  timeout: 1000,
+  prefix: "production"
+})
+*/
+
+// Create a new ratelimiter for anonymous users
+export function getRateLimit() {
+  const rateLimit = new Ratelimit({
+    redis,
+    limiter: Ratelimit.slidingWindow(1, "1 m"), // 1 request every minute
+    analytics: true,
+    // timeout: 120000,
+    prefix: "production:anon"
+  })
+
+  return rateLimit
+}

src/app/server/redis/redis.ts

@@ -0,0 +1,6 @@
+import { Redis } from "@upstash/redis"
+
+export const redis = new Redis({
+  url: `${process.env.UPSTASH_REDIS_REST_URL || ""}`,
+  token: `${process.env.UPSTASH_REDIS_REST_TOKEN || ""}`,
+})

src/app/store.ts

@@ -6,7 +6,7 @@ import { create } from "zustand"
 import { GenerationStage, GlobalStatus, TaskStatus } from "@/types"
 import { getVideoOrientation } from "@/lib/utils/getVideoOrientation"
 
-import { RESOLUTION_LONG, RESOLUTION_SHORT } from "./server/aitube/config"
+import { RESOLUTION_LONG, RESOLUTION_SHORT } from "./server/config"
 import { putTextInTextAreaElement } from "@/lib/utils/putTextInTextAreaElement"
 import { defaultPrompt } from "./config"
 
@@ -41,6 +41,8 @@ export const useStore = create<{
   currentVideoOrientation: ClapMediaOrientation
   progress: number
   error: string
+  showAuthWall: boolean
+  setShowAuthWall: (showAuthWall: boolean) => void
   toggleOrientation: () => void
   setOrientation: (orientation: ClapMediaOrientation) => void
   setCurrentVideoOrientation: (currentVideoOrientation: ClapMediaOrientation) => void
@@ -94,6 +96,8 @@ export const useStore = create<{
   currentVideoOrientation: ClapMediaOrientation.PORTRAIT,
   progress: 0,
   error: "",
+  showAuthWall: false,
+  setShowAuthWall: (showAuthWall: boolean) => { set({ showAuthWall }) },
   toggleOrientation: () => {
     const { orientation: previousOrientation, currentVideoOrientation, currentVideo } = get()
     const orientation =

src/components/interface/auth-wall.tsx

@@ -0,0 +1,24 @@
+
+import { Dialog, DialogContent } from "@/components/ui/dialog"
+import { Login } from "./login"
+
+export function AuthWall({ show }: { show: boolean }) {
+  return (
+    <Dialog open={show}>
+      <DialogContent className="sm:max-w-[800px]">
+        <div className="grid gap-4 py-4 text-stone-800 text-center text-xl">
+        <p className="">
+          The AI Stories Factory is an app to generate consistent video stories.
+        </p>
+        <p>
+          By default it uses Hugging Face for story and image generation,<br/>
+          our service is free of charge but we would like you to sign-in 👇
+         </p>
+         <p>
+          <Login />
+         </p>
+        </div>
+      </DialogContent>
+    </Dialog>
+  )
+}

src/components/interface/login/index.tsx

@@ -0,0 +1,8 @@
+"use client"
+
+import dynamic from "next/dynamic";
+
+export const Login = dynamic(() => import("./login"), {
+  // Make sure we turn SSR off
+  ssr: false,
+});

src/components/interface/login/login.tsx

@@ -0,0 +1,20 @@
+"use client"
+
+import { Button } from "@/components/ui/button"
+import { useOAuth } from "@/lib/oauth/useOAuth"
+
+function Login() {
+  const { login } = useOAuth()
+  return <Button
+    variant="ghost"
+    onClick={login}
+    className="
+    text-xs
+    bg-lime-500 dark:bg-lime-500
+    text-stone-950/80 dark:text-stone-950/80
+    hover:bg-lime-400 dark:hover:bg-lime-400
+    hover:text-stone-950/100 dark:hover:text-stone-950/100
+    ">Login with Hugging Face</Button>
+}
+
+export default Login

src/components/interface/video-preview.tsx

@@ -12,8 +12,13 @@ import { cn } from "@/lib/utils/cn"
 
 import { useStore } from "../../app/store"
 import HFLogo from "../../app/hf-logo.svg"
+import { Login } from "./login"
+import { useOAuth } from "@/lib/oauth/useOAuth"
 
 export function VideoPreview() {
+
+  const { isLoggedIn, enableOAuthWall } = useOAuth()
+
   const status = useStore(s => s.status)
   const parseGenerationStatus = useStore(s => s.parseGenerationStatus)
   const storyGenerationStatus = useStore(s => s.storyGenerationStatus)
@@ -38,6 +43,17 @@ export function VideoPreview() {
     isPortrait,
   } = useOrientation()
 
+  const placeholder = <div
+  className="
+  text-base
+  text-center
+text-stone-50/90 dark:text-stone-50/90
+"
+>{
+  error ? <span>{error}</span> : 
+  <span>No video yet</span>
+}</div>
+
   return (
     <div className={cn(`
       -mt-8 md:mt-0
@@ -67,12 +83,25 @@ export function VideoPreview() {
         w-full h-full
         bg-black text-white
         ">
-          {isBusy ? <div className="
+          {
+            !isLoggedIn ? <div className="
+            flex flex-col items-center justify-center
+            space-y-2
+            ">
+              <div className="
+              text-base
+              text-center
+            text-stone-50/90 dark:text-stone-50/90
+            ">Please login to generate videos:</div>
+              <Login />
+              </div>
+            : isBusy ? <div className="
           flex flex-col 
           items-center justify-center
           text-center space-y-1.5">
             <p className="text-2xl font-bold">{progress}%</p> 
-            <p className="text-base text-white/70">{isBusy
+            <p className="text-base text-white/70">{
+            isBusy
               ? (
                 // note: some of those tasks are running in parallel,
                 // and some are super-slow (like music or video)
@@ -91,7 +120,7 @@ export function VideoPreview() {
               )
               : status === "error"
               ? <span>{error || ""}</span>
-              : <span>{error ? error : <span>&nbsp;</span>}</span> // to prevent layout changes
+              : placeholder // to prevent layout changes
             }</p>
             </div>
           : (currentVideo && currentVideo?.length > 128) ? <video
@@ -106,10 +135,7 @@ export function VideoPreview() {
             className="object-cover"
             style={{
             }}
-          /> : <div  className="
-          flex flex-col 
-          items-center justify-center
-          text-lg text-center"></div>}
+          /> : placeholder}
         </div>
 
         <div className={cn(`

src/lib/hooks/useProcessors.ts

@@ -1,10 +1,12 @@
 "use client"
 
-import React, { useTransition } from "react"
+import React, { useState, useTransition } from "react"
 import { ClapProject, ClapSegmentCategory, getClapAssetSourceType, newEntity, updateClap } from "@aitube/clap"
 
 import { logImage } from "@/lib/utils"
 import { useIsBusy, useStoryPromptDraft } from "@/lib/hooks"
+import { isRateLimitError } from "@/lib/utils/isRateLimitError"
+import { useToast } from "@/components/ui/use-toast"
 
 import { createClap } from "@/app/server/aitube/createClap"
 import { editClapEntities } from "@/app/server/aitube/editClapEntities"
@@ -16,8 +18,11 @@ import { editClapVideos } from "@/app/server/aitube/editClapVideos"
 import { exportClapToVideo } from "@/app/server/aitube/exportClapToVideo"
 
 import { useStore } from "../../app/store"
+import { useOAuth } from "../oauth/useOAuth"
 
 export function useProcessors() {
+  const [isLocked, setLocked] = useState(false)
+  
   const { storyPromptDraft, setStoryPromptDraft, promptDraftRef } = useStoryPromptDraft()
 
   const [_isPending, startTransition] = useTransition()
@@ -30,6 +35,7 @@ export function useProcessors() {
   const setMainCharacterImage = useStore(s => s.setMainCharacterImage)
   const setMainCharacterVoice = useStore(s => s.setMainCharacterVoice)
   const setStatus = useStore(s => s.setStatus)
+  const setShowAuthWall = useStore(s => s.setShowAuthWall)
 
   const error = useStore(s => s.error)
   const setError = useStore(s => s.setError)
@@ -46,8 +52,12 @@ export function useProcessors() {
   const setCurrentVideo = useStore(s => s.setCurrentVideo)
   const setProgress = useStore(s => s.setProgress)
 
+  const { isLoggedIn, enableOAuthWall } = useOAuth()
+
   const  { isBusy, busyRef } = useIsBusy()
 
+  const { toast } = useToast()
+
   const generateStory = async (): Promise<ClapProject> => {
 
     let clap: ClapProject | undefined = undefined
@@ -69,11 +79,11 @@ export function useProcessors() {
 
       if (clap.segments.length <= 1) { throw new Error(`failed to generate more than one segments`) }
 
-      console.log(`handleSubmit(): received a clap = `, clap)
+      console.log(`generateStory(): received a clap = `, clap)
 
-      console.log(`handleSubmit():  copying over entities from the previous clap`)
+      console.log(`generateStory():  copying over entities from the previous clap`)
 
-      console.log(`handleSubmit(): later we can add button(s) to clear the project and/or the character(s)`)
+      console.log(`generateStory(): later we can add button(s) to clear the project and/or the character(s)`)
       const { currentClap } = useStore.getState()
 
       clap.entities = Array.isArray(currentClap?.entities) ? currentClap.entities : []
@@ -153,7 +163,7 @@ export function useProcessors() {
 
       if (!clap) { throw new Error(`failed to edit the sound`) }
 
-      console.log(`handleSubmit(): received a clap with sound = `, clap)
+      console.log(`generateSounds(): received a clap with sound = `, clap)
       setSoundGenerationStatus("finished")
       console.log("---------------- GENERATED SOUND ----------------")
       console.table(clap.segments.filter(s => s.category === ClapSegmentCategory.SOUND), [
@@ -180,7 +190,7 @@ export function useProcessors() {
 
       if (!clap) { throw new Error(`failed to edit the music`) }
 
-      console.log(`handleSubmit(): received a clap with music = `, clap)
+      console.log(`generateMusic(): received a clap with music = `, clap)
       setMusicGenerationStatus("finished")
       console.log("---------------- GENERATED MUSIC ----------------")
       console.table(clap.segments.filter(s => s.category === ClapSegmentCategory.MUSIC), [
@@ -210,7 +220,7 @@ export function useProcessors() {
       if (!clap) { throw new Error(`failed to edit the storyboards`) }
 
       // const fusion = 
-      console.log(`handleSubmit(): received storyboards = `, clap)
+      console.log(`generateStoryboards(): received storyboards = `, clap)
 
       setImageGenerationStatus("finished")
       console.log("---------------- GENERATED STORYBOARDS ----------------")
@@ -282,7 +292,7 @@ export function useProcessors() {
 
       if (!clap) { throw new Error(`failed to edit the dialogues`) }
 
-      console.log(`handleSubmit(): received dialogues = `, clap)
+      console.log(`generateDialogues(): received dialogues = `, clap)
       setVoiceGenerationStatus("finished")
       console.log("---------------- GENERATED DIALOGUES ----------------")
       console.table(clap.segments.filter(s => s.category === ClapSegmentCategory.DIALOGUE), [
@@ -322,6 +332,14 @@ export function useProcessors() {
   }
   
   const handleSubmit = async () => {
+
+    if (busyRef.current) { return }
+
+    if (enableOAuthWall && !isLoggedIn) {
+      setShowAuthWall(true)
+      return
+    }
+
     setStatus("generating")
     busyRef.current = true
 
@@ -425,9 +443,22 @@ export function useProcessors() {
         setStatus("finished")
         setError("")
       } catch (err) {
-        console.error(`failed to generate: `, err)
-        setStatus("error")
-        setError(`Error, please contact an admin on Discord (${err})`)
+
+        if (isRateLimitError(err)) {
+          console.error("Critical error: you are doing too many requests!")
+          toast({
+            title: "You can generate only one video per minute 👀",
+            description: "Don't send too many requests at once 🤗",
+          })
+          return
+        } else {
+          console.error(err)
+          toast({
+            title: "We couldn't generate this video 👀",
+            description: "We are currently experiencing a surge in traffic, please try later in the day 🤗",
+          })
+        }
+
       }
     })
   }

src/lib/oauth/getOAuthRedirectUrl.ts

@@ -0,0 +1,11 @@
+export function getOAuthRedirectUrl(): string {
+  if (typeof window === "undefined") {
+    return "http://localhost:3000"
+  }
+
+  return (
+    window.location.hostname === "aistoriesfactory.app" ? "https://aistoriesfactory.app"
+    : window.location.hostname === "jbilcke-hf-ai-stories-factory.hf.space" ? "https://jbilcke-hf-ai-stories-factory.hf.space"
+    : "http://localhost:3000"
+  )
+}

src/lib/oauth/getValidOAuth.ts

@@ -0,0 +1,30 @@
+import { OAuthResult } from "@huggingface/hub"
+
+// return a valid OAuthResult, or else undefined
+export function getValidOAuth(rawInput?: any): OAuthResult | undefined {
+  try {
+    let untypedOAuthResult: any
+    try {
+      untypedOAuthResult = JSON.parse(rawInput)
+      if (!untypedOAuthResult) { throw new Error("no valid serialized oauth result") }
+    } catch (err) {
+      untypedOAuthResult = rawInput
+    }
+
+    const maybeValidOAuth = untypedOAuthResult as OAuthResult
+
+    const accessTokenExpiresAt = new Date(maybeValidOAuth.accessTokenExpiresAt)
+
+    // Get the current date
+    const currentDate = new Date()
+
+    if (accessTokenExpiresAt.getTime() < currentDate.getTime()) {
+      throw new Error("the serialized oauth result has expired")
+    }
+
+    return maybeValidOAuth
+  } catch (err) {
+    // console.error(err)
+    return undefined
+  }
+}

src/lib/oauth/useOAuth.ts

@@ -0,0 +1,146 @@
+"use client"
+
+import { useEffect } from "react"
+import { useSearchParams } from "next/navigation"
+import { OAuthResult, oauthHandleRedirectIfPresent, oauthLoginUrl } from "@huggingface/hub"
+
+import { enableHuggingFaceOAuth, oauthClientId, oauthScopes } from "@/app/config"
+
+import { usePersistedOAuth } from "./usePersistedOAuth"
+import { getValidOAuth } from "./getValidOAuth"
+import { useShouldDisplayLoginWall } from "./useShouldDisplayLoginWall"
+import { getOAuthRedirectUrl } from "./getOAuthRedirectUrl"
+
+export function useOAuth({
+  debug = false
+}: {
+  debug?: boolean
+} = {
+  debug: false
+}): {
+  clientId: string
+  redirectUrl: string
+  scopes: string
+  canLogin: boolean
+  login: () => Promise<void>
+  isLoggedIn: boolean
+  enableOAuth: boolean
+  enableOAuthWall: boolean
+  oauthResult?: OAuthResult
+ } {
+  const [oauthResult, setOAuthResult] = usePersistedOAuth()
+
+  const clientId = oauthClientId
+
+  // const redirectUrl = config.oauthRedirectUrl
+  const redirectUrl = getOAuthRedirectUrl()
+
+  const scopes = oauthScopes
+  const enableOAuth = enableHuggingFaceOAuth
+
+  const searchParams = useSearchParams()
+  const code = searchParams?.get("code") || ""
+  const state = searchParams?.get("state") || ""
+
+  const hasReceivedFreshOAuth = Boolean(code && state)
+
+  // note: being able to log into hugging face using the popup
+  // is different from seeing the "login wall"
+  const canLogin: boolean = Boolean(oauthClientId && enableOAuth)
+  const isLoggedIn = Boolean(oauthResult)
+
+  const enableOAuthWall = useShouldDisplayLoginWall() 
+
+  if (debug) {
+    console.log("useOAuth debug:", {
+      oauthResult,
+      clientId,
+      redirectUrl,
+      scopes,
+      enableOAuth,
+      enableOAuthWall,
+      code,
+      state,
+      hasReceivedFreshOAuth,
+      canLogin,
+      isLoggedIn,
+    })
+
+    /*
+    useOAuth debug: {
+      oauthResult: '',
+      clientId: '........',
+      redirectUrl: 'http://localhost:3000',
+      scopes: 'openid profile inference-api',
+      isOAuthEnabled: true,
+      code: '...........',
+      state: '{"nonce":".........","redirectUri":"http://localhost:3000"}',
+      hasReceivedFreshOAuth: true,
+      canLogin: false,
+      isLoggedIn: false
+    }
+    */
+  }
+
+  useEffect(() => {
+    // no need to perfor the rest if the operation is there is nothing in the url
+    if (hasReceivedFreshOAuth) {
+
+      (async () => {
+        const maybeValidOAuth = await oauthHandleRedirectIfPresent()
+
+        const newOAuth = getValidOAuth(maybeValidOAuth)
+
+        if (!newOAuth) {
+          if (debug) {
+            console.log("useOAuth::useEffect 1: got something in the url but no valid oauth data to show.. something went terribly wrong")
+          }
+        } else {
+          if (debug) {
+            console.log("useOAuth::useEffect 1: correctly received the new oauth result, saving it to local storage:", newOAuth)
+          }
+          setOAuthResult(newOAuth)
+
+          // once set we can (brutally) reload the page
+          window.location.href = `//${window.location.host}${window.location.pathname}`
+        }
+      })()
+    }
+  }, [debug, hasReceivedFreshOAuth])
+
+  // for debugging purpose
+  useEffect(() => {
+    if (!debug) {
+      return
+    }
+    // console.log(`useOAuth::useEffect 2: canLogin? ${canLogin}`)
+    if (!canLogin) {
+      return
+    }
+    // console.log(`useOAuth::useEffect2: isLoggedIn? ${isLoggedIn}`)
+    if (!isLoggedIn) {
+      return
+    }
+    // console.log(`useOAuth::useEffect 2: oauthResult:`, oauthResult)
+  }, [debug, canLogin, isLoggedIn, oauthResult])
+
+  const login = async () => {
+    window.location.href = await oauthLoginUrl({
+      clientId,
+      redirectUrl,
+      scopes,
+    })
+  }
+
+  return {
+    clientId,
+    redirectUrl,
+    scopes,
+    canLogin,
+    login,
+    isLoggedIn,
+    enableOAuth,
+    enableOAuthWall,
+    oauthResult
+  }
+}

src/lib/oauth/usePersistedOAuth.ts

@@ -0,0 +1,23 @@
+"use client"
+
+import { useLocalStorage } from "usehooks-ts"
+import { OAuthResult } from "@huggingface/hub"
+
+import { localStorageHuggingFaceOAuthKey } from "@/app/config"
+
+import { getValidOAuth } from "./getValidOAuth"
+
+export function usePersistedOAuth(): [OAuthResult | undefined, (oauthResult: OAuthResult) => void] {
+  const [serializedHuggingFaceOAuth, setSerializedHuggingFaceOAuth] = useLocalStorage<string>(
+    localStorageHuggingFaceOAuthKey,
+    ""
+  )
+
+  const oauthResult = getValidOAuth(serializedHuggingFaceOAuth)
+
+  const setOAuthResult = (oauthResult: OAuthResult) => {
+    setSerializedHuggingFaceOAuth(JSON.stringify(oauthResult))
+  }
+
+  return [oauthResult, setOAuthResult]
+}

src/lib/oauth/useShouldDisplayLoginWall.ts

@@ -0,0 +1,16 @@
+
+// we don't want to display the login wall to people forking the project,
+
+import { enableHuggingFaceOAuth, enableHuggingFaceOAuthWall, oauthClientId } from "@/app/config"
+
+// or to people who selected no hugging face server at all
+export function useShouldDisplayLoginWall() {
+
+  const shouldDisplayLoginWall = Boolean(
+    oauthClientId &&
+    enableHuggingFaceOAuth &&
+    enableHuggingFaceOAuthWall
+  )
+
+  return shouldDisplayLoginWall
+}

src/lib/utils/isRateLimitError.ts

@@ -0,0 +1,4 @@
+export function isRateLimitError(something: unknown) {
+  // yeah this is a very crude implementation
+  return `${something || ""}`.includes("Rate Limit Reached")
+}