Spaces:
Running
Running
File size: 4,598 Bytes
88d98bf b11f3cf 241e984 88d98bf b11f3cf 88d98bf 241e984 b11f3cf 241e984 b11f3cf 241e984 88d98bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | #!/usr/bin/env python3
"""Start Potato in a Hugging Face Spaces friendly way."""
from __future__ import annotations
import os
import sys
from pathlib import Path
def running_on_hugging_face() -> bool:
"""Return true inside HF Spaces or when explicitly requested."""
return bool(
os.environ.get("SPACE_ID")
or os.environ.get("SPACE_REPO_NAME")
or os.environ.get("POTATO_ENABLE_IFRAME_COOKIES") == "1"
)
def configure_iframe_cookies(app) -> None:
"""Allow Flask session cookies to survive inside the Hugging Face iframe."""
app.config.update(
SESSION_COOKIE_SAMESITE="None",
SESSION_COOKIE_SECURE=True,
SESSION_COOKIE_HTTPONLY=True,
)
def patch_login_template_for_iframe() -> None:
"""Patch Potato's login form so HF iframe login navigates top-level."""
import potato
template_path = Path(potato.__file__).resolve().parent / "templates" / "home.html"
template = template_path.read_text(encoding="utf-8")
marker = "<!-- HF_DIRECT_LOGIN_PATCH -->"
if marker in template:
return
patch = f"""
{marker}
<script>
(function() {{
function isInIframe() {{
try {{
return window.self !== window.top;
}} catch (error) {{
return true;
}}
}}
document.addEventListener('DOMContentLoaded', function() {{
if (!isInIframe()) return;
var form = document.querySelector('form[action="/auth"]');
var usernameInput = document.querySelector('#login-email');
if (!form || !usernameInput) return;
form.setAttribute('method', 'GET');
form.setAttribute('action', '/hf-direct-login');
form.setAttribute('target', '_top');
usernameInput.setAttribute('name', 'username');
var note = document.createElement('div');
note.className = 'potato-alert';
note.style.marginBottom = '0.75rem';
note.style.fontSize = '0.875rem';
note.textContent = 'Login will open the app directly so your browser can keep the session.';
form.parentNode.insertBefore(note, form);
}});
}}());
</script>
"""
template = template.replace("</body>", patch + "\n</body>")
template_path.write_text(template, encoding="utf-8")
def register_direct_login_route(app) -> None:
"""Register a cookie-friendly top-level login route for HF Spaces."""
from flask import redirect, request, session, url_for
from potato.authentication import UserAuthenticator
from potato.flask_server import (
UserPhase,
get_item_state_manager,
get_user_state_manager,
init_user_state,
)
def hf_direct_login():
username = (request.args.get("username") or "").strip()
if not username:
return redirect(url_for("home"))
authenticator = UserAuthenticator.get_instance()
if not authenticator.is_valid_username(username):
authenticator.add_user(username, None)
session.clear()
session["username"] = username
session.permanent = True
user_state_manager = get_user_state_manager()
if not user_state_manager.has_user(username):
init_user_state(username)
user_state = user_state_manager.get_user_state(username)
if user_state and user_state.get_phase() == UserPhase.LOGIN:
user_state_manager.advance_phase(username)
if user_state and not user_state.has_assignments():
get_item_state_manager().assign_instances_to_user(user_state)
return redirect(url_for("annotate"))
app.add_url_rule("/hf-direct-login", "hf_direct_login", hf_direct_login, methods=["GET"])
def main() -> int:
port = os.environ.get("PORT", "7860")
config = os.environ.get("POTATO_CONFIG", "config.yaml")
if running_on_hugging_face():
patch_login_template_for_iframe()
from potato.flask_server import create_app
app = create_app(config)
register_direct_login_route(app)
if running_on_hugging_face():
configure_iframe_cookies(app)
print("Enabled Hugging Face iframe-compatible session cookies.", flush=True)
print(f"Starting Potato on 0.0.0.0:{port}", flush=True)
app.run(host="0.0.0.0", port=int(port), debug=False, use_reloader=False, threaded=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|