File size: 4,661 Bytes
f8c73f9 | 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 | """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
|