Update app.py

app.py

@@ -14,6 +14,8 @@ from io import BytesIO
 import base64
 import re
 
+SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
+
 # Regex pattern to match data URI scheme
 data_uri_pattern = re.compile(r'data:image/(png|jpeg|jpg|webp);base64,')
 
@@ -49,8 +51,12 @@ def resize_image(image):
     return image
 
 
-def process(image_base64):
-    orig_image = readb64(image_base64)
+def process(secret_token, base64_in):
+    if secret_token != SECRET_TOKEN:
+        raise gr.Error(
+            f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
+
+    orig_image = readb64(base64_in)
     
     # prepare input
     w,h = orig_im_size = orig_image.size
@@ -77,9 +83,30 @@ def process(image_base64):
     new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
     new_im.paste(orig_image, mask=pil_im)
 
-    return writeb64(new_im)
+    base64_out = writeb64(new_im)
+
+    return base64_out
+
 
-demo = gr.Interface(fn=process, inputs="text", outputs="text")
+with gr.Blocks() as demo:
+    secret_token = gr.Text(
+        label='Secret Token',
+        max_lines=1,
+        placeholder='Enter your secret token')
+    gr.HTML("""
+        <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
+        <div style="text-align: center; color: black;">
+        <p style="color: black;">This space is a REST API to programmatically remove the background of an image.</p>
+        <p style="color: black;">Interested in using it? Please use the <a href="https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4" target="_blank">original space</a>, thank you!</p>
+        </div>
+        </div>""")
+    base64_in = gr.Textbox(label="Base64 Input")
+    base64_out = gr.Textbox(label="Base64 Output")
+       
+    generate_btn.click(
+        fn=process,
+        inputs=[secret_token, base64_in],
+        outputs=base64_out,
+        api_name="run")
 
-if __name__ == "__main__":
-    demo.launch(share=False)
+demo.queue(max_size=20).launch()