#!/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 = "" if marker in template: return patch = f""" {marker} """ template = template.replace("", patch + "\n") 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())