adam/oasis_player.py

"""Small, UI-independent helpers for ADAM's native Oasis Player page."""

from __future__ import annotations

import json
import random
from pathlib import Path


IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
ACTION_MODEL_TYPES = {
    "action_conditioned_rectified_flow_video",
    "action_conditioned_latent_vae_flow_video",
    "action_conditioned_temporal_latent_flow",
    "action_conditioned_temporal_pixel_flow",
}
LATENT_MODEL_TYPES = {
    "action_conditioned_latent_vae_flow_video",
    "action_conditioned_temporal_latent_flow",
}


def is_action_model(path: str | Path) -> bool:
    """Return whether *path* is a complete Oasis action-world-model folder."""
    folder = Path(path)
    try:
        info = json.loads((folder / "action_flow_model_info.json").read_text(encoding="utf-8"))
    except (OSError, ValueError, TypeError):
        return False
    weights = (
        folder / "unet" / "diffusion_pytorch_model.safetensors",
        folder / "unet" / "diffusion_pytorch_model.bin",
    )
    model_type = info.get("model_type")
    has_unet = (folder / "unet" / "config.json").is_file() and any(candidate.is_file() for candidate in weights)
    has_vae = (folder / "vae" / "config.json").is_file() and (folder / "vae" / "pytorch_model.bin").is_file()
    return model_type in ACTION_MODEL_TYPES and has_unet and (model_type not in LATENT_MODEL_TYPES or has_vae)


def model_info(path: str | Path) -> dict:
    """Read compatible model metadata, raising a useful error for the UI."""
    folder = Path(path)
    if not is_action_model(folder):
        raise ValueError("Choose a complete Oasis action model folder.")
    return json.loads((folder / "action_flow_model_info.json").read_text(encoding="utf-8"))


def frame_size(info: dict) -> tuple[int, int]:
    """Return the trained (width, height), including legacy checkpoint metadata."""
    if info.get("width") and info.get("height"):
        return int(info["width"]), int(info["height"])
    value = str(info.get("resolution", "256x144")).lower().replace("×", "x")
    if "x" in value:
        width, height = (int(part.strip()) for part in value.split("x", 1))
        return width, height
    width = int(value)
    return width, round(width * 9 / 16)


def discover_models(oasis_root: str | Path, assets=()) -> list[tuple[str, Path]]:
    """Find local Oasis models without copying, moving, or modifying them."""
    found: dict[Path, str] = {}
    root = Path(oasis_root)
    library = root / "output_action_flow_models"
    if library.is_dir():
        for folder in library.iterdir():
            if folder.is_dir() and is_action_model(folder):
                found[folder.resolve()] = folder.name
    for asset in assets:
        if getattr(asset, "kind", "") != "model" or getattr(asset, "trainer", "") != "oasis":
            continue
        folder = Path(getattr(asset, "path", ""))
        if is_action_model(folder):
            found[folder.resolve()] = getattr(asset, "name", "") or folder.name
    return sorted(((name, path) for path, name in found.items()), key=lambda item: item[0].casefold())


def capture_path(root: str | Path, prefix: str = "oasis_frame") -> Path:
    """Choose an ADAM-owned, collision-resistant PNG destination."""
    from datetime import datetime
    from uuid import uuid4

    folder = Path(root) / "data" / "oasis_captures"
    folder.mkdir(parents=True, exist_ok=True)
    return folder / f"{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}_{uuid4().hex[:6]}.png"


def random_roblox_dataset_frame(oasis_root: str | Path, chooser=None) -> Path:
    """Pick one image from a deeply nested Roblox Dataset without assuming its layout.

    Oasis recordings often live in several named recording folders, each with a
    ``frames`` folder.  Reservoir sampling avoids holding every path in memory.
    ``chooser`` is injectable for deterministic tests.
    """
    root = Path(oasis_root) / "OldDatasets" / "Roblox Dataset"
    if not root.is_dir():
        raise FileNotFoundError("The connected Oasis Trainer has no OldDatasets/Roblox Dataset folder.")
    pick = None
    count = 0
    chooser = chooser or random.randrange
    try:
        for path in root.rglob("*"):
            if not path.is_file() or path.suffix.casefold() not in IMAGE_EXTENSIONS:
                continue
            count += 1
            if chooser(count) == 0:
                pick = path
    except OSError as exc:
        raise OSError(f"ADAM could not read the Roblox Dataset: {exc}") from exc
    if pick is None:
        raise FileNotFoundError("No PNG, JPG, WEBP, or BMP frames were found in the Roblox Dataset.")
    return pick