app.py

import os
import base64
import json
import importlib.machinery
import importlib.util
import zlib
from pathlib import Path
from typing import Any, Dict

_CONFIG_MARKER = b"OMNICFG1"


def _apply_frame_interpolation_settings() -> None:
    os.environ["OMNI_BASE_FPS"] = "12"
    os.environ["OMNI_FRAME_MULTIPLIER"] = "2"
    os.environ["OMNI_CRF"] = "0" 
    os.environ["OMNI_ALLOWED_FPS"] = "16,32,64,128"
    print("[startup] Frame interpolation settings applied: BASE_FPS=12, MULTIPLIER=2, CRF=auto, ALLOWED_FPS=16,32,64,128")


def _collect_config_values(namespace: Dict[str, Any]) -> Dict[str, Any]:
    data: Dict[str, Any] = {}
    for key, value in namespace.items():
        if key.isupper() and key != "APP_RUNTIME_CONFIG":
            data[key] = value
    class_obj = namespace.get("APP_RUNTIME_CONFIG")
    if class_obj is not None:
        for key in dir(class_obj):
            if key.isupper() and key not in data:
                data[key] = getattr(class_obj, key)
    if not data:
        raise RuntimeError("private config source contains no uppercase settings")
    return data


def _load_config_from_source(content: str) -> Dict[str, Any]:
    scope: Dict[str, Any] = {"__builtins__": __builtins__, "__name__": "omni_private_config"}
    exec(content, scope, scope)
    return _collect_config_values(scope)


def _load_config_from_pyc(path: Path) -> Dict[str, Any]:
    loader = importlib.machinery.SourcelessFileLoader("omni_private_config_compat", str(path))
    spec = importlib.util.spec_from_loader("omni_private_config_compat", loader)
    module = importlib.util.module_from_spec(spec) if spec else None
    if not spec or not module or not spec.loader:
        raise RuntimeError(f"failed to load config bytecode module spec from {path}")
    spec.loader.exec_module(module)
    namespace = {key: getattr(module, key) for key in dir(module)}
    return _collect_config_values(namespace)


def _write_packed_private_config(path: Path, data: Dict[str, Any]) -> None:
    payload = json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    packed = _CONFIG_MARKER + zlib.compress(payload, level=9)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(packed)


def _materialize_private_config_from_env() -> None:
    repo_dir = Path(__file__).resolve().parent
    pyc_target = repo_dir / "src" / "config.pyc"
    raw_b64 = (os.getenv("OMNI_PRIVATE_CONFIG_B64") or "").strip()
    raw_py = (os.getenv("OMNI_PRIVATE_CONFIG_PY") or "").strip()

    content = ""
    source = ""
    if raw_b64:
        try:
            content = base64.b64decode(raw_b64.encode("utf-8")).decode("utf-8")
            source = "OMNI_PRIVATE_CONFIG_B64"
        except Exception as exc:
            print(f"[startup] warning: decode OMNI_PRIVATE_CONFIG_B64 failed: {type(exc).__name__}: {exc}")
            return
    elif raw_py:
        content = raw_py
        source = "OMNI_PRIVATE_CONFIG_PY"

    if content:
        try:
            data = _load_config_from_source(content.rstrip() + "\n")
            _write_packed_private_config(pyc_target, data)
            print(f"[startup] private config packed from {source} -> {pyc_target}")
            return
        except Exception as exc:
            print(f"[startup] warning: build private config failed: {type(exc).__name__}: {exc}")

    local_source_candidates = (
        repo_dir / "tools_local" / "config.py",
        repo_dir / "src_codes" / "src" / "config.py",
    )
    for local_src in local_source_candidates:
        if not local_src.exists():
            continue
        try:
            data = _load_config_from_source(local_src.read_text(encoding="utf-8"))
            _write_packed_private_config(pyc_target, data)
            print(f"[startup] private config packed from local source -> {pyc_target}")
            return
        except Exception as exc:
            print(f"[startup] warning: pack local private config failed: {type(exc).__name__}: {exc}")

    if not pyc_target.exists():
        return
    try:
        raw = pyc_target.read_bytes()
        if raw.startswith(_CONFIG_MARKER):
            return
        data = _load_config_from_pyc(pyc_target)
        _write_packed_private_config(pyc_target, data)
        print(f"[startup] migrated private config bytecode to packed format -> {pyc_target}")
    except Exception as exc:
        print(f"[startup] warning: private config migration skipped: {type(exc).__name__}: {exc}")


def _load_runtime_module():
    repo_dir = Path(__file__).resolve().parent
    pyc_path = repo_dir / "src" / "app_lib.pyc"
    if pyc_path.exists():
        got = pyc_path.read_bytes()[:4]
        expected = importlib.util.MAGIC_NUMBER
        if got == expected:
            loader = importlib.machinery.SourcelessFileLoader("app_lib_runtime", str(pyc_path))
            spec = importlib.util.spec_from_loader("app_lib_runtime", loader)
            if spec is not None:
                module = importlib.util.module_from_spec(spec)
                loader.exec_module(module)
                print("[startup] loaded runtime module from src/app_lib.pyc")
                return module
            print("[startup] warning: failed to create module spec for src/app_lib.pyc; fallback to source")
        else:
            py_ver = f"{os.sys.version_info.major}.{os.sys.version_info.minor}"
            print(
                "[startup] warning: incompatible src/app_lib.pyc, fallback to source. "
                f"got={got.hex()} expected={expected.hex()} for Python {py_ver}."
            )

    from src import app_lib as module

    print("[startup] loaded runtime module from src/app_lib.py")
    return module


_materialize_private_config_from_env()
_RUNTIME = _load_runtime_module()
build_demo = _RUNTIME.build_demo
parse_model_names = _RUNTIME.parse_model_names
DEFAULT_OMNI_VIDEOS = getattr(_RUNTIME, "DEFAULT_OMNI_VIDEOS", "")
ensure_models_ready_on_startup = getattr(_RUNTIME, "ensure_models_ready_on_startup", lambda: "startup model hook not available")
kickoff_model_prepare_background = getattr(
    _RUNTIME,
    "kickoff_model_prepare_background",
    None,
)
kickoff_runtime_prepare_background = getattr(
    _RUNTIME,
    "kickoff_runtime_prepare_background",
    None,
)

def _log_model_config_status() -> None:
    raw = (os.environ.get("OMNI_VIDEOS") or DEFAULT_OMNI_VIDEOS or "").strip()
    if not raw:
        print("[startup] OMNI_VIDEOS is empty; running in bootstrap mode.")
        return
    try:
        parse_model_names(raw)
        print("[startup] OMNI_VIDEOS format looks valid.")
    except Exception as exc:
        print(f"[startup] OMNI_VIDEOS format invalid: {type(exc).__name__}: {exc}")


if __name__ == "__main__":
    _apply_frame_interpolation_settings()
    _log_model_config_status()
    print("[startup] runtime preparation:")
    if callable(kickoff_runtime_prepare_background):
        print(kickoff_runtime_prepare_background())
    else:
        print("runtime startup hook not available")
    print("[startup] model preparation:")
    if callable(kickoff_model_prepare_background):
        print(kickoff_model_prepare_background())
    else:
        print(ensure_models_ready_on_startup())
    demo = build_demo()
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")), ssr_mode=False)