Allow password query API access

app.py

@@ -267,14 +267,20 @@ def gpu_duration(prompt, width, height, steps, p=None):
     return min(120, max(10, int(work / 4) + 6))
 
 
+def valid_pass(p):
+    return secrets.compare_digest(str(p), PASSWORD)
+
+
 @spaces.GPU(duration=gpu_duration)
-def generate(prompt, width=1152, height=896, steps=16):
+def generate(prompt, width=1152, height=896, steps=16, p=""):
+    if not valid_pass(p):
+        raise gr.Error("Invalid password")
     return generate_core(prompt, width, height, steps)
 
 
 @spaces.GPU(duration=gpu_duration)
 def api_generate(prompt, width=1152, height=896, steps=16, p=""):
-    if not secrets.compare_digest(str(p), PASSWORD):
+    if not valid_pass(p):
         raise gr.Error("Invalid password")
     return generate_core(prompt, width, height, steps)
 
@@ -282,7 +288,7 @@ def api_generate(prompt, width=1152, height=896, steps=16, p=""):
 def queued_generate(prompt, width, height, steps):
     from gradio_client import Client
 
-    client = Client(LOCAL_URL, auth=("user", PASSWORD), verbose=False)
+    client = Client(LOCAL_URL, verbose=False)
     result = client.predict(prompt, width, height, steps, PASSWORD, api_name="/generate")
     return Image.open(result).copy()
 
@@ -339,7 +345,7 @@ api.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], all
 
 
 def require_pass(p: str = Query(...)):
-    if not secrets.compare_digest(p, PASSWORD):
+    if not valid_pass(p):
         raise HTTPException(401, "invalid password")
 
 
@@ -409,23 +415,37 @@ def direct(body: DirectRequest):
     return FileResponse(path, media_type="image/png", filename=filename)
 
 
+def login(p):
+    if not valid_pass(p):
+        raise gr.Error("Invalid password")
+    return gr.Column(visible=False), gr.Column(visible=True)
+
+
 with gr.Blocks(title="Image API") as demo:
-    gr.HTML("<model-note>ComfyUI-compatible API and direct image generation</model-note>")
-    prompt_input = gr.Textbox(label="Prompt", lines=5)
-    with gr.Row():
-        width_input = gr.Number(1152, label="Width", precision=0)
-        height_input = gr.Number(896, label="Height", precision=0)
-        steps_input = gr.Slider(1, 100, 16, step=1, label="Steps")
-    button = gr.Button("Generate", variant="primary")
+    with gr.Column() as login_panel:
+        pass_input = gr.Textbox(label="Password", type="password")
+        login_button = gr.Button("Login", variant="primary")
+    with gr.Column(visible=False) as generate_panel:
+        gr.HTML("<model-note>ComfyUI-compatible API and direct image generation</model-note>")
+        prompt_input = gr.Textbox(label="Prompt", lines=5)
+        with gr.Row():
+            width_input = gr.Number(1152, label="Width", precision=0)
+            height_input = gr.Number(896, label="Height", precision=0)
+            steps_input = gr.Slider(1, 100, 16, step=1, label="Steps")
+        button = gr.Button("Generate", variant="primary")
+        output = gr.Image(label="Image", format="png")
     api_button = gr.Button(visible=False)
-    pass_input = gr.Textbox(visible=False)
-    output = gr.Image(label="Image", format="png")
+    api_pass_input = gr.Textbox(visible=False)
+    login_button.click(
+        login, pass_input, [login_panel, generate_panel], queue=False,
+        api_visibility="private",
+    )
     button.click(
-        generate, [prompt_input, width_input, height_input, steps_input], output,
+        generate, [prompt_input, width_input, height_input, steps_input, pass_input], output,
         api_visibility="private",
     )
     api_button.click(
-        api_generate, [prompt_input, width_input, height_input, steps_input, pass_input], output,
+        api_generate, [prompt_input, width_input, height_input, steps_input, api_pass_input], output,
         api_name="generate",
     )
 
@@ -440,7 +460,5 @@ if __name__ == "__main__":
         server_port=PORT,
         css_paths="style.css",
         head='<link rel="icon" href="data:;base64,=">',
-        auth=("user", PASSWORD),
-        auth_message="Username: user",
         _app=api,
     )