Update app.py

app.py

@@ -25,6 +25,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,')
 
@@ -104,7 +106,11 @@ def resize_image(img, target_height):
     return resized_img
 
 # @spaces.GPU(enable_queue=True)
-def inference(input_image_b64, prompt, a_prompt, n_prompt, denoise_steps, upscale, alpha, cfg, seed):
+def inference(secret_token, input_image_b64, prompt, a_prompt, n_prompt, denoise_steps, upscale, alpha, cfg, seed):
+    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.')
+
     input_image = readb64(input_image_b64)
     
     input_image = resize_image(input_image, 512)
@@ -172,6 +178,7 @@ with gr.Blocks() as demo:
     with gr.Column():
         with gr.Row():
             with gr.Column():
+                secret_token = gr.Textbox()
                 input_image_b64 = gr.Textbox()
                 prompt_in = gr.Textbox(label="Prompt", value="Frog")
                 with gr.Accordion(label="Advanced settings", open=False):
@@ -189,6 +196,7 @@ with gr.Blocks() as demo:
     submit_btn.click(
         fn = inference,
         inputs = [
+            secret_token,
             input_image_b64, prompt_in,
             added_prompt, neg_prompt,
             denoise_steps,