add pw

app.py

@@ -29,7 +29,11 @@ credentials = service_account.Credentials.from_service_account_info(service_acco
 project="pdr-imagen"
 aiplatform.init(project=project, credentials=credentials)
 
-def generate_image(prompt,model_name):
+def generate_image(pw,prompt,model_name):
+
+    if pw != os.getenv("PW"):
+        raise gr.Error("Invalid password. Please try again.")
+
     try:
         model = ImageGenerationModel.from_pretrained(model_name)
         response = model.generate_images(
@@ -49,6 +53,10 @@ with gr.Blocks() 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.")
+
     #instructions
     with gr.Accordion("Instructions & Tips",label="instructions",open=False):
         with gr.Row():
@@ -67,7 +75,7 @@ with gr.Blocks() as demo:
     with gr.Accordion("Image Output",label="Image Output",open=True):
         output_image = gr.Image(label="Image")
 
-    btn.click(fn=generate_image, inputs=[text, model ], outputs=output_image, api_name=False)
-    text.submit(fn=generate_image, inputs=[text, model ], outputs=output_image, api_name="generate_image") # Generate an api endpoint in Gradio / HF
+    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
 
 demo.launch(share=False)