add max generations functionality and do enforce pw functionality

app.py

@@ -29,9 +29,16 @@ credentials = service_account.Credentials.from_service_account_info(service_acco
 project="pdr-imagen"
 aiplatform.init(project=project, credentials=credentials)
 
+# enforce password is True if DO_ENFORCE_PW is set to "true"
+DO_ENFORCE_PW = os.getenv("DO_ENFORCE_PW")
+
+
+def trigger_max_gens():
+    gr.Warning("🖼️ Max Image Generations Reached! 🖼️")
+
 def generate_image(pw,prompt,model_name):
 
-    if pw != os.getenv("PW"):
+    if pw != os.getenv("PW") and DO_ENFORCE_PW == "true":
         raise gr.Error("Invalid password. Please try again.")
 
     try:
@@ -46,16 +53,59 @@ def generate_image(pw,prompt,model_name):
 
     except Exception as e:
         print(e)
-        raise gr.Error(f"An error occurred while generating the image for: {entry}")
+        raise gr.Error(f"An error occurred while generating the image")
     return image_url
 
-with gr.Blocks() as demo:
+custom_js = """
+function customJS() {
+    //Limit Image Generation
+    const MAX_GENERATIONS = 10;
+    const DO_ENFORCE_MAX_GENERATIONS = true;
+
+    disableGenerateButton = function() {
+        const btn = document.getElementById('btn_generate-images');
+        btn.disabled = true;
+        btn.classList.add('not-visible');
+    }
+
+    triggerMaxGenerationsToast = function() {
+        const trigger_max_gens_btn = document.getElementById('trigger-max-gens-btn');
+        trigger_max_gens_btn.click();
+    }
+
+    setCurrentGenerations = function() {
+        if (!DO_ENFORCE_MAX_GENERATIONS) {
+            return;
+        }
+        const curGenerations = localStorage.getItem('currentGenerations');
+        console.log(`${curGenerations} / ${MAX_GENERATIONS}`)
+        if (curGenerations) {
+            if (curGenerations >= MAX_GENERATIONS) {
+                triggerMaxGenerationsToast();
+                disableGenerateButton();
+            } else {
+                localStorage.setItem('currentGenerations', parseInt(curGenerations) + 1);
+            }
+        } else {
+            localStorage.setItem('currentGenerations', 1);
+        }
+    }
+
+    setCurrentGenerations();
+
+    document.getElementById('btn_generate-images').addEventListener('click', function() {
+        setCurrentGenerations();
+    });
+
+}
+"""
+
+with gr.Blocks(js=custom_js) as demo:
 
     gr.Markdown("# <center>Google Vertex Imagen Generator</center>")
-
     #password
-    pw = gr.Textbox(label="Password", type="password", placeholder="Enter the password to unlock the service")
-    gr.Markdown("Need access? Send a DM to @HeaversMike on Twitter or send me an email / Slack msg.")
+    pw = gr.Textbox(label="Password", type="password", placeholder="Enter the password to unlock the service",visible=False if DO_ENFORCE_PW == "false" else True)
+    gr.Markdown("Need access? Send a DM to @HeaversMike on Twitter or send me an email / Slack msg.",visible=False if DO_ENFORCE_PW == "false" else True)
 
     #instructions
     with gr.Accordion("Instructions & Tips",label="instructions",open=False):
@@ -69,13 +119,19 @@ with gr.Blocks() as demo:
     model = gr.Dropdown(choices=["imagegeneration@002", "imagegeneration@005"], label="Model", value="imagegeneration@005")
 
     with gr.Row():
-        btn = gr.Button("Generate Images")
+        btn = gr.Button("Generate Images", variant="primary", elem_id="btn_generate-images")
 
     #output
     with gr.Accordion("Image Output",label="Image Output",open=True):
         output_image = gr.Image(label="Image")
 
+    with gr.Row():
+        trigger_max_gens_btn = gr.Button(value="Show Max Gens Reached",visible=False,elem_id="trigger-max-gens-btn")
+
     btn.click(fn=generate_image, inputs=[pw,text, model ], outputs=output_image, api_name=False)
     text.submit(fn=generate_image, inputs=[pw,text, model ], outputs=output_image, api_name="generate_image") # Generate an api endpoint in Gradio / HF
 
+    #js-triggered functionality
+    trigger_max_gens_btn.click(trigger_max_gens, None, None)
+
 demo.launch(share=False)