#!/usr/bin/env python3
"""MODUS any-to-any demo — HuggingFace Space entrypoint (ZeroGPU).

Thin ZeroGPU wrapper over the existing 3-tab demo backend (``demo_modus.py`` +
``demo_my/*``). The model is loaded ONCE at startup on CPU (no GPU, no time
limit); each inference call moves it to the ZeroGPU-provided GPU via the
``@spaces.GPU`` decorator (moving ~30GB bf16 over PCIe is seconds, whereas the
build+weight-load is ~5min and must NOT happen inside the GPU-time window).

Space setup (Settings -> Variables and secrets):
    HF_TOKEN   read token for the gated weights repo (below)
Optional env (have sane defaults):
    MODUS_WEIGHTS_REPO   gated HF model repo with the bf16 weights + config + VAE
"""
import os
import sys

try:
    import torch
    print(f"[app] torch={torch.__version__} python={sys.version.split()[0]}", flush=True)
except Exception as _e:
    print(f"[app] torch import failed: {_e}", flush=True)

# Teach PIL to decode in-the-wild formats (AVIF, iPhone HEIC). gradio opens each
# uploaded file with PIL.Image.open before it reaches our code, so registering
# these openers here (import-time) is what lets those uploads work at all.
try:
    import pillow_avif  # noqa: F401  (registers the AVIF opener on import)
except Exception as _e:
    print(f"[app] pillow-avif-plugin unavailable: {_e}", flush=True)
try:
    from pillow_heif import register_heif_opener
    register_heif_opener()
except Exception as _e:
    print(f"[app] pillow-heif unavailable: {_e}", flush=True)

# ── Install the MODUS backend from the (private) GitHub repo at startup ──────────
# Single source of truth: this Space only carries app.py + fourm/ + test_images/;
# the demo/inference code (demo_modus, any2any, modeling, data, core, conf, ...)
# is pip-installed from EPFL-VILAB/Modus. Needs a GH_TOKEN Space secret with read
# access. --no-deps so it does NOT touch the ZeroGPU torch stack.
import subprocess  # noqa: E402
_GH = os.environ.get("GH_TOKEN")
if _GH:
    print("[app] installing modus backend from EPFL-VILAB/Modus ...", flush=True)
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-deps", "--quiet",
         "--force-reinstall", "--no-cache-dir",  # always pull the latest repo HEAD
         f"git+https://x-access-token:{_GH}@github.com/EPFL-VILAB/Modus.git"],
        check=True,
    )
    print("[app] modus backend installed.", flush=True)
else:
    print("[app] GH_TOKEN not set; expecting the modus backend to be present.", flush=True)

# ── Env MUST be set before importing the demo backend (it reads these at import) ─
os.environ.setdefault("MODUS_NO_MEAN_RESIZING", "1")  # avoid gradio-BLAS deadlock
os.environ.setdefault("MODUS_FORCE_SDPA_ATTN", "1")   # no flash-attn on the Space
os.environ.setdefault("MODUS_TORCHVISION_FREE", "1")  # ZeroGPU torch has no torchvision
os.environ.setdefault("MODUS_DEMO_MODALITY_CONFIG",
                      "conf/modalities/instruction_16mod_stage2.yaml")
os.environ.setdefault("MODUS_DEMO_MODEL_NAME", "bagel_from_json")
os.environ.setdefault("MODUS_DEMO_USE_EMA", "0")

WEIGHTS_REPO = os.environ.get("MODUS_WEIGHTS_REPO", "mqye/modus-16mod-stage3")

# ── Pull the gated weights once (model.safetensors + ae + config + tokenizer) ────
from huggingface_hub import snapshot_download  # noqa: E402

_weights_dir = snapshot_download(
    repo_id=WEIGHTS_REPO,
    repo_type="model",
    token=os.environ.get("HF_TOKEN"),
)
# The snapshot dir holds BOTH the checkpoint (model.safetensors) and the base
# config/tokenizer/VAE, so it serves as CHECKPOINT_PATH and MODEL_PATH at once.
os.environ["MODUS_DEMO_CHECKPOINT"] = _weights_dir
os.environ["BAGEL_MODEL_PATH"] = _weights_dir
print(f"[app] weights ready at {_weights_dir}", flush=True)

import spaces  # noqa: E402  (ZeroGPU)

# diffusers 0.20 (imported by the fourm VQVAE feature tokenizers) does
# `from huggingface_hub import cached_download`, which was removed in hub>=0.26
# (the Space has 0.36). Shim it to hf_hub_download so `import diffusers` succeeds
# and the dino/clip/imagebind tokenizers can load.
import huggingface_hub as _hh  # noqa: E402
if not hasattr(_hh, "cached_download"):
    _hh.cached_download = _hh.hf_hub_download

# Importing demo_modus builds the UI + backend and reads the env set above.
import demo_modus  # noqa: E402


# ── ZeroGPU: wrap the two inference entry points so each call gets a GPU slice ───
# tab1/tab2 call demo_modus.run_task; tab3 calls demo_modus.run_representation_task.
# The tab handlers resolve these by module-global name at call time, so replacing
# the module attribute makes them use the GPU-wrapped versions.
# duration = max GPU-seconds ZeroGPU reserves per call. The Chained tab runs TWO
# generations inside a single run_task call (~60s+), so 60 is too tight ("GPU task
# aborted"); 120 covers chained + cold-start model materialisation. (With PRO the
# reservation size is a non-issue quota-wise.)
_GPU_DURATION = int(os.environ.get("MODUS_GPU_DURATION", "120"))
demo_modus.run_task = spaces.GPU(duration=_GPU_DURATION)(demo_modus.run_task)
# Tab3 runs THREE generations (vit/vae/both) per click. Wrapping run_representation_task
# would take three separate GPU acquisitions in one handler → the 3rd hits "Expired
# ZeroGPU proxy token". Instead wrap the whole tab3 HANDLER so all three run inside a
# single GPU session (one token, one reservation big enough for 3 gens).
demo_modus.tab3_generate = spaces.GPU(duration=180)(demo_modus.tab3_generate)

# Load the model ONCE at startup. The heavy CPU work (build arch + read 30GB
# weights, ~5min) runs here, OUTSIDE any @spaces.GPU window; the model's .cuda()
# moves are deferred by `spaces` and materialise on the first GPU call.
try:
    demo_modus.HOLDER.ensure_loaded()
    print("[app] model loaded (CPU) at startup", flush=True)
except Exception as e:  # surface in UI, retry lazily on first request
    demo_modus.HOLDER.load_error = str(e)
    print(f"[app] startup model load failed: {e}", flush=True)

# The demo backend is pip-installed (site-packages), so demo_modus.REPO_ROOT points
# into site-packages and its `test_images/` lookup finds nothing. The example images
# live in THIS Space repo (CWD /home/user/app/test_images), so point the example
# gallery there before build_ui() reads it.
_EX_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_images")
def _space_example_images():
    if not os.path.isdir(_EX_DIR):
        return []
    return [os.path.join(_EX_DIR, f) for f in sorted(os.listdir(_EX_DIR))
            if f.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))
            and "_seg." not in f.lower()]  # exclude precomputed seg previews
demo_modus._example_images = _space_example_images
print(f"[app] {len(_space_example_images())} example images from {_EX_DIR}", flush=True)

# Work around a gradio 4.44.1 bug: get_api_info() (called when rendering the main
# "/" route) crashes with `TypeError: argument of type 'bool' is not iterable` when
# a component's JSON schema contains a bool (additionalProperties: true/false).
# Patch gradio_client's schema helpers to tolerate bool schemas.
try:
    import gradio_client.utils as _gcu

    _orig_j2p = _gcu._json_schema_to_python_type

    def _safe_j2p(schema, defs=None):
        if isinstance(schema, bool):
            return "Any"
        return _orig_j2p(schema, defs)

    _gcu._json_schema_to_python_type = _safe_j2p

    _orig_get_type = _gcu.get_type

    def _safe_get_type(schema):
        if isinstance(schema, bool):
            return "Any"
        return _orig_get_type(schema)

    _gcu.get_type = _safe_get_type
except Exception as _e:
    print(f"[app] gradio_client schema patch skipped: {_e}", flush=True)

demo = demo_modus.build_ui()
demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_api=False)