waypoint_ttnn 0.1.1: pinned weights, only the files the server reads, chip selection fix

bfe08bd verified 23 days ago

9.05 kB

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
"""Module-level mesh-device + model singleton, modeled on tt-skyreels'
`skyreels_ttnn/session.py` (itself modeled on tt-animatediff's). Manages the TTNN mesh
device, the three loaded models (`WaypointWorldModel`, `WaypointVAEEncoder`,
`WaypointVAEDecoder`), and the SINGLE active interactive `WaypointGenerator` session, so
the Gradio app doesn't pay weight-load cost more than once per process.
Unlike tt-skyreels (stateless: every call is an independent one-shot generation), this
model is inherently STATEFUL -- an interactive session holds a live KV cache across many
`step()` calls. This module therefore holds at most ONE active `WaypointGenerator` at a
time; starting a new session (a new seed image) discards the previous one's cache state.
A real multi-tenant server would need one generator per session id, not a module global --
out of scope for this first demo app (see PORT_PLAN.md Stage 6's serving-contract note).
"""

import threading
from typing import Optional

_lock = threading.Lock()
_device = None
_world_model = None
_vae_encoder = None
_vae_decoder = None
_hf_config = None
_generator = None

MESH_SHAPE = (1, 1)
L1_SMALL_SIZE = 21760  # conv2d's halo op needs this set; see vae_decoder.py's own note.

# --- Weights: which repo, which exact revision, which files --------------------------
#
#: Upstream weights repo. Never embedded in this package; fetched into the HF cache.
WEIGHTS_REPO_ID = "Overworld/Waypoint-1.5-1B"

#: The upstream commit this port was brought up, benchmarked and hardware-verified
#: against (HF API `sha` on 2026-09-27; upstream lastModified 2026-07-16, i.e. before
#: the 2026-09-09 bring-up started, so this is the revision every number in
#: BRINGUP_LOG.md was measured on). Pinned so a later upstream push can't silently
#: change what a served bundle loads.
PINNED_WEIGHTS_REVISION = "391f92827075edcf4a8b3c8a2ddae010698f8636"

#: tt-model-manager's v6 `run.sh` exports this when the bundle manifest carries a
#: weights revision, so the manifest and the code can't disagree. When it is unset
#: (dev checkout, bare `uvicorn`), the pin above applies. An empty string counts as
#: unset -- `or`, not a `.get()` default -- so `TT_MODEL_WEIGHTS_REVISION=` can't
#: accidentally resolve `main`.
WEIGHTS_REVISION_ENV = "TT_MODEL_WEIGHTS_REVISION"

#: The ONLY files the served path reads -- see `ensure_waypoint_models()` and
#: `_load_config()` below, which open exactly these three paths and nothing else under
#: the snapshot (verified by grepping the served closure -- session.py, server/app.py,
#: tt/*.py -- for every file open / from_pretrained / hub call; none exist outside this
#: module). The upstream repo is ~11.4 GB; this filter skips the root
#: `model.safetensors` (3.72 GB, an alternative packaging of the same model that this
#: port never loads), `assets/` (~250 MB of demo media), and the upstream Python
#: sources (`transformer/model.py`, `vae/ae_model.py`, `modular_*`), which this port
#: replaces rather than imports. `vae/config.json` is deliberately NOT included: the
#: VAE's shape is hardcoded in tt/vae_{encoder,decoder}.py, not read from config.
#: Keep this list and the three `os.path.join(snapshot_dir, ...)` reads below (the config in
#: `_load_config()`, the two weight files in `ensure_waypoint_models()`) in lockstep --
#: `tests/test_weights_pin.py` asserts they match. A repackage should pass the same
#: list to `tt-model package-thin` so the bundle's own pull fetches the same subset.
WEIGHTS_ALLOW_PATTERNS = (
    "transformer/config.json",
    "transformer/diffusion_pytorch_model.safetensors",
    "vae/diffusion_pytorch_model.safetensors",
)


class _Cfg:
    def __init__(self, d):
        self.__dict__.update(d)


def ensure_waypoint_models():
    """Open the mesh device and load all three TTNN models (once per process).
    Returns:
        (device, world_model, vae_encoder, vae_decoder, hf_config)
    """
    global _device, _world_model, _vae_encoder, _vae_decoder, _hf_config
    if _device is not None:
        return _device, _world_model, _vae_encoder, _vae_decoder, _hf_config

    with _lock:
        if _device is not None:
            return _device, _world_model, _vae_encoder, _vae_decoder, _hf_config

        import os

        import ttnn
        from safetensors.torch import load_file

        from waypoint_ttnn.tt.full_model import WaypointWorldModel
        from waypoint_ttnn.tt.vae_decoder import WaypointVAEDecoder
        from waypoint_ttnn.tt.vae_encoder import WaypointVAEEncoder

        device = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(*MESH_SHAPE), l1_small_size=L1_SMALL_SIZE)
        try:
            snapshot_dir = _resolve_snapshot_dir()
            hf_config = _load_config(snapshot_dir)
            # Both paths must be covered by WEIGHTS_ALLOW_PATTERNS, or the filtered
            # snapshot won't contain them.
            tf_weights = os.path.join(snapshot_dir, "transformer", "diffusion_pytorch_model.safetensors")
            vae_weights = os.path.join(snapshot_dir, "vae", "diffusion_pytorch_model.safetensors")
            tf_state_dict = load_file(tf_weights)
            vae_state_dict = load_file(vae_weights)

            world_model = WaypointWorldModel.from_state_dict(tf_state_dict, hf_config=hf_config, mesh_device=device)
            vae_encoder = WaypointVAEEncoder.from_state_dict(vae_state_dict, mesh_device=device)
            vae_decoder = WaypointVAEDecoder.from_state_dict(vae_state_dict, mesh_device=device)
        except Exception:
            _close_device(device)
            raise

        _world_model = world_model
        _vae_encoder = vae_encoder
        _vae_decoder = vae_decoder
        _hf_config = hf_config
        _device = device
        return _device, _world_model, _vae_encoder, _vae_decoder, _hf_config


def start_session():
    """Ends any current interactive session and returns a fresh WaypointGenerator
    (models loaded lazily via ensure_waypoint_models() if not already). Call `.seed()`
    on the result with a real image before the first `.step()`."""
    global _generator
    device, world_model, vae_encoder, vae_decoder, hf_config = ensure_waypoint_models()

    from waypoint_ttnn.tt.generation_loop import WaypointGenerator

    with _lock:
        _generator = WaypointGenerator(world_model, vae_encoder, vae_decoder, hf_config, device)
        return _generator


def current_session() -> Optional["WaypointGenerator"]:  # noqa: F821 -- forward ref, avoids importing ttnn at module load
    return _generator


def _resolve_snapshot_dir() -> str:
    """Downloads (or reuses an already-cached) local snapshot of the weights repo,
    returning its directory. Uses huggingface_hub's own resolution rather than a
    hardcoded `~/.cache/huggingface/...` path -- that path doesn't exist inside the
    served container, which mounts the HF cache at `/hf` (`HF_HOME=/hf`, set by
    tt-model-manager's own container.py) instead. A real bug hit on the first real
    `tt-model serve` attempt: `glob.glob(...)[0]` raised `IndexError` because the
    hardcoded host path was simply absent in the container.
    Fetches a pinned revision (`weights_revision()`), and only the files the served
    path reads (`WEIGHTS_ALLOW_PATTERNS`). Before 0.1.1 this downloaded `main` in full,
    ~4 GB more than needed, and would have silently followed any upstream push."""
    from huggingface_hub import snapshot_download

    return snapshot_download(
        repo_id=WEIGHTS_REPO_ID,
        revision=weights_revision(),
        allow_patterns=list(WEIGHTS_ALLOW_PATTERNS),
    )


def weights_revision() -> str:
    """The weights revision to load: `$TT_MODEL_WEIGHTS_REVISION` if set and non-empty
    (exported by tt-model-manager's v6 run.sh from the bundle manifest), else the
    revision this port was verified on. Pure function of the environment -- no hub call,
    no device -- so it is unit-testable without hardware."""
    import os

    return os.environ.get(WEIGHTS_REVISION_ENV) or PINNED_WEIGHTS_REVISION


def _load_config(snapshot_dir: str):
    """Loads the real transformer config the way every test in this repo does --
    `AutoConfig`-free, since the shared venv's diffusers predates it (see CLAUDE.md)."""
    import json
    import os

    with open(os.path.join(snapshot_dir, "transformer", "config.json")) as f:
        return _Cfg(json.load(f))


def close() -> None:
    """Closes the TTNN mesh device. Rarely needed -- process exit reclaims it."""
    global _device, _world_model, _vae_encoder, _vae_decoder, _hf_config, _generator
    with _lock:
        if _device is None:
            return
        _close_device(_device)
        _device = None
        _world_model = None
        _vae_encoder = None
        _vae_decoder = None
        _hf_config = None
        _generator = None


def _close_device(device) -> None:
    try:
        import ttnn

        ttnn.close_mesh_device(device)
    except Exception:
        pass

	# SPDX-License-Identifier: Apache-2.0
	# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
	"""Module-level mesh-device + model singleton, modeled on tt-skyreels'
	`skyreels_ttnn/session.py` (itself modeled on tt-animatediff's). Manages the TTNN mesh
	device, the three loaded models (`WaypointWorldModel`, `WaypointVAEEncoder`,
	`WaypointVAEDecoder`), and the SINGLE active interactive `WaypointGenerator` session, so
	the Gradio app doesn't pay weight-load cost more than once per process.

	Unlike tt-skyreels (stateless: every call is an independent one-shot generation), this
	model is inherently STATEFUL -- an interactive session holds a live KV cache across many
	`step()` calls. This module therefore holds at most ONE active `WaypointGenerator` at a
	time; starting a new session (a new seed image) discards the previous one's cache state.
	A real multi-tenant server would need one generator per session id, not a module global --
	out of scope for this first demo app (see PORT_PLAN.md Stage 6's serving-contract note).
	"""

	import threading
	from typing import Optional

	_lock = threading.Lock()
	_device = None
	_world_model = None
	_vae_encoder = None
	_vae_decoder = None
	_hf_config = None
	_generator = None

	MESH_SHAPE = (1, 1)
	L1_SMALL_SIZE = 21760 # conv2d's halo op needs this set; see vae_decoder.py's own note.

	# --- Weights: which repo, which exact revision, which files --------------------------
	#
	#: Upstream weights repo. Never embedded in this package; fetched into the HF cache.
	WEIGHTS_REPO_ID = "Overworld/Waypoint-1.5-1B"

	#: The upstream commit this port was brought up, benchmarked and hardware-verified
	#: against (HF API `sha` on 2026-09-27; upstream lastModified 2026-07-16, i.e. before
	#: the 2026-09-09 bring-up started, so this is the revision every number in
	#: BRINGUP_LOG.md was measured on). Pinned so a later upstream push can't silently
	#: change what a served bundle loads.
	PINNED_WEIGHTS_REVISION = "391f92827075edcf4a8b3c8a2ddae010698f8636"

	#: tt-model-manager's v6 `run.sh` exports this when the bundle manifest carries a
	#: weights revision, so the manifest and the code can't disagree. When it is unset
	#: (dev checkout, bare `uvicorn`), the pin above applies. An empty string counts as
	#: unset -- `or`, not a `.get()` default -- so `TT_MODEL_WEIGHTS_REVISION=` can't
	#: accidentally resolve `main`.
	WEIGHTS_REVISION_ENV = "TT_MODEL_WEIGHTS_REVISION"

	#: The ONLY files the served path reads -- see `ensure_waypoint_models()` and
	#: `_load_config()` below, which open exactly these three paths and nothing else under
	#: the snapshot (verified by grepping the served closure -- session.py, server/app.py,
	#: tt/*.py -- for every file open / from_pretrained / hub call; none exist outside this
	#: module). The upstream repo is ~11.4 GB; this filter skips the root
	#: `model.safetensors` (3.72 GB, an alternative packaging of the same model that this
	#: port never loads), `assets/` (~250 MB of demo media), and the upstream Python
	#: sources (`transformer/model.py`, `vae/ae_model.py`, `modular_*`), which this port
	#: replaces rather than imports. `vae/config.json` is deliberately NOT included: the
	#: VAE's shape is hardcoded in tt/vae_{encoder,decoder}.py, not read from config.
	#: Keep this list and the three `os.path.join(snapshot_dir, ...)` reads below (the config in
	#: `_load_config()`, the two weight files in `ensure_waypoint_models()`) in lockstep --
	#: `tests/test_weights_pin.py` asserts they match. A repackage should pass the same
	#: list to `tt-model package-thin` so the bundle's own pull fetches the same subset.
	WEIGHTS_ALLOW_PATTERNS = (
	"transformer/config.json",
	"transformer/diffusion_pytorch_model.safetensors",
	"vae/diffusion_pytorch_model.safetensors",
	)


	class _Cfg:
	def __init__(self, d):
	self.__dict__.update(d)


	def ensure_waypoint_models():
	"""Open the mesh device and load all three TTNN models (once per process).

	Returns:
	(device, world_model, vae_encoder, vae_decoder, hf_config)
	"""
	global _device, _world_model, _vae_encoder, _vae_decoder, _hf_config
	if _device is not None:
	return _device, _world_model, _vae_encoder, _vae_decoder, _hf_config

	with _lock:
	if _device is not None:
	return _device, _world_model, _vae_encoder, _vae_decoder, _hf_config

	import os

	import ttnn
	from safetensors.torch import load_file

	from waypoint_ttnn.tt.full_model import WaypointWorldModel
	from waypoint_ttnn.tt.vae_decoder import WaypointVAEDecoder
	from waypoint_ttnn.tt.vae_encoder import WaypointVAEEncoder

	device = ttnn.open_mesh_device(mesh_shape=ttnn.MeshShape(*MESH_SHAPE), l1_small_size=L1_SMALL_SIZE)
	try:
	snapshot_dir = _resolve_snapshot_dir()
	hf_config = _load_config(snapshot_dir)
	# Both paths must be covered by WEIGHTS_ALLOW_PATTERNS, or the filtered
	# snapshot won't contain them.
	tf_weights = os.path.join(snapshot_dir, "transformer", "diffusion_pytorch_model.safetensors")
	vae_weights = os.path.join(snapshot_dir, "vae", "diffusion_pytorch_model.safetensors")
	tf_state_dict = load_file(tf_weights)
	vae_state_dict = load_file(vae_weights)

	world_model = WaypointWorldModel.from_state_dict(tf_state_dict, hf_config=hf_config, mesh_device=device)
	vae_encoder = WaypointVAEEncoder.from_state_dict(vae_state_dict, mesh_device=device)
	vae_decoder = WaypointVAEDecoder.from_state_dict(vae_state_dict, mesh_device=device)
	except Exception:
	_close_device(device)
	raise

	_world_model = world_model
	_vae_encoder = vae_encoder
	_vae_decoder = vae_decoder
	_hf_config = hf_config
	_device = device
	return _device, _world_model, _vae_encoder, _vae_decoder, _hf_config


	def start_session():
	"""Ends any current interactive session and returns a fresh WaypointGenerator
	(models loaded lazily via ensure_waypoint_models() if not already). Call `.seed()`
	on the result with a real image before the first `.step()`."""
	global _generator
	device, world_model, vae_encoder, vae_decoder, hf_config = ensure_waypoint_models()

	from waypoint_ttnn.tt.generation_loop import WaypointGenerator

	with _lock:
	_generator = WaypointGenerator(world_model, vae_encoder, vae_decoder, hf_config, device)
	return _generator


	def current_session() -> Optional["WaypointGenerator"]: # noqa: F821 -- forward ref, avoids importing ttnn at module load
	return _generator


	def _resolve_snapshot_dir() -> str:
	"""Downloads (or reuses an already-cached) local snapshot of the weights repo,
	returning its directory. Uses huggingface_hub's own resolution rather than a
	hardcoded `~/.cache/huggingface/...` path -- that path doesn't exist inside the
	served container, which mounts the HF cache at `/hf` (`HF_HOME=/hf`, set by
	tt-model-manager's own container.py) instead. A real bug hit on the first real
	`tt-model serve` attempt: `glob.glob(...)[0]` raised `IndexError` because the
	hardcoded host path was simply absent in the container.

	Fetches a pinned revision (`weights_revision()`), and only the files the served
	path reads (`WEIGHTS_ALLOW_PATTERNS`). Before 0.1.1 this downloaded `main` in full,
	~4 GB more than needed, and would have silently followed any upstream push."""
	from huggingface_hub import snapshot_download

	return snapshot_download(
	repo_id=WEIGHTS_REPO_ID,
	revision=weights_revision(),
	allow_patterns=list(WEIGHTS_ALLOW_PATTERNS),
	)


	def weights_revision() -> str:
	"""The weights revision to load: `$TT_MODEL_WEIGHTS_REVISION` if set and non-empty
	(exported by tt-model-manager's v6 run.sh from the bundle manifest), else the
	revision this port was verified on. Pure function of the environment -- no hub call,
	no device -- so it is unit-testable without hardware."""
	import os

	return os.environ.get(WEIGHTS_REVISION_ENV) or PINNED_WEIGHTS_REVISION


	def _load_config(snapshot_dir: str):
	"""Loads the real transformer config the way every test in this repo does --
	`AutoConfig`-free, since the shared venv's diffusers predates it (see CLAUDE.md)."""
	import json
	import os

	with open(os.path.join(snapshot_dir, "transformer", "config.json")) as f:
	return _Cfg(json.load(f))


	def close() -> None:
	"""Closes the TTNN mesh device. Rarely needed -- process exit reclaims it."""
	global _device, _world_model, _vae_encoder, _vae_decoder, _hf_config, _generator
	with _lock:
	if _device is None:
	return
	_close_device(_device)
	_device = None
	_world_model = None
	_vae_encoder = None
	_vae_decoder = None
	_hf_config = None
	_generator = None


	def _close_device(device) -> None:
	try:
	import ttnn

	ttnn.close_mesh_device(device)
	except Exception:
	pass