tt-animatediff / tt_model_server.py

animatediff-ttnn 0.11.1: pinned weights, restore Gradio app.py, runner renamed to tt_model_server.py

5b3ff58 verified about 1 month ago

17.4 kB

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
"""ASGI app serving tt-animatediff, for tt-model-manager's ``tt-dit-server`` kind.
A diffusion transformer has no tokens, no KV cache and no continuous batching, so vLLM has
nothing to do for it. ``tt-dit-server`` therefore installs a small HTTP stack instead of an
engine and launches this module with uvicorn:
    python -m uvicorn --host 0.0.0.0 --port 8000 --lifespan on \\
        animatediff_ttnn.server.app:app
TWO CONTRACTS THIS FILE HAS TO HONOUR
-------------------------------------
**1. Readiness is the lifespan.** tt-model-manager decides the server is up when uvicorn
logs ``Application startup complete``, which it prints *after* ASGI lifespan startup
returns. So the device open and the model warm belong in the lifespan and nowhere else: do
them lazily on the first request and the supervisor calls the server ready while it is
still loading, then times out the first request it sends.
**2. Importing this module must not touch hardware.** ``verify_lines`` imports the ASGI
attribute at image-build time, on a machine with no card, purely to prove the code
allowlist actually shipped the server. Every ttnn/tt-metal import is therefore inside a
function, never at module scope. That is also what lets the tests below run on CPU.
ENVIRONMENT, set by the launcher
--------------------------------
``MESH_DEVICE``            the SKU, e.g. ``P300x2`` — informational here
``ANIMATEDIFF_MESH_SHAPE`` the resolved shape, e.g. ``1x4`` (the manifest points
                           ``runtime.mesh_shape_env`` at this name; the kind's default is
                           FLUX.2's, which means nothing to this model)
``HF_MODEL``               the weights repo id, reported by ``/v1/models``
ONE REQUEST AT A TIME
---------------------
There is no continuous batching to hide behind: the pipeline owns the device, and two
concurrent denoise loops would interleave on it. Requests are serialised on a lock and the
blocking work runs in a worker thread so the event loop can still answer ``/health`` while
a generation is in flight. This is a deliberate cap, not an oversight — a queue that
admits work the device cannot do concurrently is a queue that reports success and returns
corrupted frames.
"""
from __future__ import annotations

import asyncio
import base64
import io
import os
import time
from contextlib import asynccontextmanager
from typing import Any, List, Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

#: The env var carrying the resolved mesh shape. Named for THIS model: the kind's default
#: is ``FLUX2_MESH_SHAPE``, which is the default only because FLUX.2 is what it was built
#: for. A manifest serving this app sets ``runtime.mesh_shape_env`` to the name below.
MESH_SHAPE_ENV = "ANIMATEDIFF_MESH_SHAPE"

#: Frames per second baked into the returned GIF.
DEFAULT_FPS = 8


#: The only image size this server can produce. The TTNN UNet is compiled at a fixed
#: 64x64 latent, and 64 * 8 = 512. Named rather than repeated so the request model, the
#: error message and any future widening all move together.
SERVABLE_SIZE = 512


class VideoGenerationRequest(BaseModel):
    """OpenAI-shaped video request.
    Field names follow ``/v1/videos/generations`` as tt-inference-server's tt-media-server
    already serves it, so the same model can sit behind either surface without a client
    change. Everything beyond ``prompt`` is optional with the pipeline's own defaults.
    """

    prompt: str = Field(min_length=1)
    negative_prompt: str = ""
    model: Optional[str] = None
    num_frames: int = Field(default=16, ge=1, le=64)
    num_inference_steps: int = Field(default=25, ge=1, le=100)
    guidance_scale: float = Field(default=7.5, ge=0.0, le=20.0)
    seed: int = 42
    # 0.35, because that is what the rest of the project calibrated: generate_animation(),
    # hf/pipeline.py AND generate_frames_temporal -- the function this server calls -- all
    # use it. This said 0.5, which matched none of them, so an HTTP client omitting the
    # field got quietly different motion from every other entry point into the same code.
    # test_the_server_request_default_matches_the_librarys_calibrated_one asserts it
    # against the library rather than against a literal, so it cannot drift again.
    temporal_alpha: float = Field(default=0.35, ge=0.0, le=1.0)
    # 512 exactly, and not as a conservative default: the TTNN UNet is compiled ONCE at a
    # fixed 64x64 latent (generation_helpers.py: `UNet2D(device, parameters, 2, 64, 64)`),
    # so 512x512 is the only thing this server can produce. These were bounded 64..1024,
    # which is pydantic-valid and kernel-invalid -- the request passed validation, took the
    # device lock, and then raised a TT_FATAL deep in the denoise loop as an uncaught 500,
    # having held the one device to do it. Refuse at the edge instead, with a message that
    # says why. Widening this means compiling the UNet per requested shape, which is a
    # ~2-3 minute cost per shape and a different design.
    height: int = Field(default=SERVABLE_SIZE, ge=SERVABLE_SIZE, le=SERVABLE_SIZE,
                        description=f"only {SERVABLE_SIZE} is supported: the UNet is "
                                    "compiled at a fixed 64x64 latent")
    width: int = Field(default=SERVABLE_SIZE, ge=SERVABLE_SIZE, le=SERVABLE_SIZE,
                       description=f"only {SERVABLE_SIZE} is supported: the UNet is "
                                   "compiled at a fixed 64x64 latent")
    #: Only b64_json is offered. A URL response would need the server to host files it has
    #: no store for, and a bundle that returns links to a path the consumer cannot read is
    #: worse than one that returns bytes.
    response_format: str = "b64_json"


class VideoData(BaseModel):
    b64_json: str


class VideoGenerationResponse(BaseModel):
    created: int
    data: List[VideoData]


def mesh_shape_from_env(env: Optional[dict] = None) -> tuple:
    """Parse ``ANIMATEDIFF_MESH_SHAPE`` (``"RxC"``) into ``(rows, cols)``.
    Defaults to ``(1, 1)`` when unset, which is what a bare local run wants. A malformed
    value RAISES rather than falling back: silently converting against the wrong mesh is
    the failure this variable exists to prevent, and it would show up as bad frames rather
    than as an error.
    """
    raw = (env if env is not None else os.environ).get(MESH_SHAPE_ENV, "").strip()
    if not raw:
        return (1, 1)
    try:
        rows, cols = (int(p) for p in raw.lower().split("x", 1))
    except ValueError as exc:
        raise ValueError(
            f"{MESH_SHAPE_ENV}={raw!r} is not a mesh shape like '1x4'"
        ) from exc
    if rows < 1 or cols < 1:
        raise ValueError(f"{MESH_SHAPE_ENV}={raw!r} must be positive")
    # Only 1x1, and refused here rather than absorbed. _open_device_and_models passes
    # `device_ids=None` for any other shape, and setup_blackhole(None) opens EVERY chip on
    # the box as a 1xN mesh -- so on a shared machine a stray 1x4 takes four chips out of
    # circulation, and then dies in the lifespan anyway, because this model cannot shard:
    # a real serve at P300x2 failed with `Can't convert a tensor distributed on
    # MeshShape([1, 4]) mesh to row-major logical tensor`. The manifest declares P150 for
    # the same reason. Widen this only alongside the mesh frame-sharding work.
    if rows * cols != 1:
        raise ValueError(
            f"{MESH_SHAPE_ENV}={raw!r} requests {rows * cols} chips, but this model cannot "
            f"shard -- only 1x1 is supported until mesh frame-sharding ships "
            f"(docs/superpowers/plans/2026-06-15-mesh-frame-sharding.md). A larger shape "
            f"would claim every chip on the box and still fail in the lifespan."
        )
    return (rows, cols)


def _open_device_and_models(shape: tuple) -> tuple:
    """Claim the mesh and load the TTNN model set. Blocking, hardware, lifespan-only."""
    from animatediff_ttnn.generation_helpers import load_sd14_ttnn
    from animatediff_ttnn.ttnn_pipeline import setup_blackhole

    from animatediff_ttnn.generation_helpers import encode_prompt

    rows, cols = shape
    # setup_blackhole(device_ids=None) opens every available chip as a 1xN mesh, which is
    # what a (1, N) shape means. An explicit id list is passed only for a single chip so a
    # 1x1 run cannot accidentally claim a neighbour someone else is using.
    device = setup_blackhole(device_ids=[0] if rows * cols == 1 else None)
    # Everything after the open is wrapped: a failure here used to leave the chip claimed
    # by a container that never started, and on a restart loop that is one claimed device
    # per attempt while TTNN refuses to reopen any of them. Same leak session.py was fixed
    # for; the server has its own open path, so it needs its own rollback.
    try:
        models = load_sd14_ttnn(device)

    # Warm the TEXT encoder too, and not as an optimisation. load_sd14_ttnn brings up the
    # UNet and the VAE; CLIP's tokenizer and text encoder were left to be downloaded
    # lazily by the first encode_prompt() call -- which happens inside _generate, under
    # the device lock, AFTER /health has already answered 200. So the readiness contract
    # this module is built on ("the device open and the model warm happen in the lifespan
    # and nowhere else") was false for one of the three models, and in a container with
    # only unet/vae cached and no egress, the server reported ready and then failed every
    # request. Downloading here means a missing weight fails startup, which is what a
    # readiness probe is for.
    #
    # NOT a full denoise pass: that would also pay TTNN kernel compilation here, which is
    # minutes on a cold cache and would delay readiness for every boot. The first request
    # is therefore still slower than the rest -- slow, but correct, and measured in
    # docs/measurements/serving-benchmark.json rather than hidden.
        encode_prompt("warmup", "")
    except Exception:
        _close_device(device)
        raise
    return device, models


def _close_device(device) -> None:
    """Release the mesh, swallowing any error.
    Swallowed deliberately: on the failure path this must not mask the exception the
    operator actually needs to read, and a close that fails is not a reason to hold the
    chip. Mirrors animatediff_ttnn.session._close_device.
    """
    try:
        import ttnn
        ttnn.close_mesh_device(device)
    except Exception:
        pass


def _frames_to_gif_b64(frames: List[Any], fps: int = DEFAULT_FPS) -> str:
    """PIL frames -> base64 GIF, in memory. Never touches disk: the container has no
    volume to write to and a served artifact should not depend on one."""
    buf = io.BytesIO()
    frames[0].save(
        buf, format="GIF", save_all=True, append_images=frames[1:],
        duration=1000 // max(fps, 1), loop=0,
    )
    return base64.b64encode(buf.getvalue()).decode("ascii")


def _generate(state: dict, req: VideoGenerationRequest) -> List[Any]:
    """Run the denoise loop. Blocking; called in a worker thread under the device lock."""
    from animatediff_ttnn.generation_helpers import encode_prompt
    from animatediff_ttnn.temporal_attention import generate_frames_temporal

    ttnn_model, ttnn_vae, config, torch_time_proj = state["models"]
    return generate_frames_temporal(
        device=state["device"],
        ttnn_model=ttnn_model,
        ttnn_vae=ttnn_vae,
        config=config,
        torch_time_proj=torch_time_proj,
        text_embeddings=encode_prompt(req.prompt, req.negative_prompt),
        num_frames=req.num_frames,
        num_steps=req.num_inference_steps,
        guidance_scale=req.guidance_scale,
        seed=req.seed,
        temporal_alpha=req.temporal_alpha,
        height=req.height,
        width=req.width,
    )


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Claim the device and warm the pipeline BEFORE the server reports ready.
    Run in a worker thread so a slow warm cannot wedge the event loop, and awaited to
    completion so ``Application startup complete`` means what tt-model-manager reads it to
    mean. A failure here is deliberately fatal: a server that starts without a device
    would answer /health cheerfully and fail every generation.
    """
    shape = mesh_shape_from_env()
    device, models = await asyncio.to_thread(_open_device_and_models, shape)
    app.state.engine = {
        "device": device, "models": models, "mesh_shape": shape,
        "hf_model": os.environ.get("HF_MODEL", "unknown"),
        "mesh_device": os.environ.get("MESH_DEVICE", "unknown"),
    }
    app.state.device_lock = asyncio.Lock()
    try:
        yield
    finally:
        # Let ttnn close the mesh itself. Bypassing its teardown is what triggers an abort
        # in MetalContext::destroy_all_instances.
        try:
            import ttnn

            ttnn.close_mesh_device(device)
        except Exception:  # noqa: BLE001 - shutdown must not mask the original error
            pass


app = FastAPI(title="tt-animatediff", lifespan=lifespan)


def _readiness() -> dict:
    """What both probes report about the model, in one place so they cannot disagree."""
    engine = getattr(app.state, "engine", None)
    return {
        "model_ready": engine is not None,
        "model": (engine or {}).get("hf_model"),
        "mesh_device": (engine or {}).get("mesh_device"),
        "mesh_shape": "x".join(str(n) for n in (engine or {}).get("mesh_shape", ())),
    }


@app.get("/tt-liveness")
async def tt_liveness() -> dict:
    """LIVENESS probe -- "is this process able to respond at all", restart on failure.
    Matches tt-inference-server's tt-media-server
    (``open_ai_api/tt_maintenance_api.py``): 200 with ``{"status": "alive", ...}`` plus the
    model-readiness payload, and a non-200 ONLY on unrecoverable failure. A model that is
    still warming is alive, not broken -- that is exactly the distinction this endpoint
    exists to draw, and reporting it as failure would have an orchestrator restart a server
    that is loading correctly.
    Deliberately lock-free and engine-free, so it keeps answering while a generation holds
    the device. A liveness probe that blocks behind the work it is monitoring is how a
    healthy-but-busy server gets killed.
    """
    try:
        return {"status": "alive", **_readiness()}
    except Exception as exc:  # noqa: BLE001 - the one case that IS unrecoverable
        raise HTTPException(status_code=500, detail=f"Liveness check failed: {exc}")


@app.get("/health")
async def health() -> dict:
    """READINESS probe -- "should traffic be routed here", gate on it.
    Despite the names, `/health` is readiness and `/tt-liveness` is liveness. That is
    tt-media-server's convention, kept for compatibility with the Python server, and the
    body is vLLM's: an empty dict when ready, 503 when not.
    This returned 200 with ``{"status": "starting"}`` until the semantics were checked
    against tt-media-server. A readiness probe answering 200 while the pipeline is still
    loading tells a load balancer to send work to a server that cannot do it -- the precise
    failure the lifespan design exists to prevent, reintroduced one layer up. The
    diagnostic payload now lives on /tt-liveness, which is where tt-media-server puts it.
    """
    if not _readiness()["model_ready"]:
        raise HTTPException(status_code=503, detail="Model not ready")
    return {}


@app.get("/v1/models")
async def models() -> dict:
    engine = getattr(app.state, "engine", None)
    name = (engine or {}).get("hf_model", "unknown")
    return {"object": "list", "data": [{"id": name, "object": "model",
                                        "owned_by": "tenstorrent"}]}


@app.post("/v1/videos/generations", response_model=VideoGenerationResponse)
async def videos_generations(req: VideoGenerationRequest) -> VideoGenerationResponse:
    engine = getattr(app.state, "engine", None)
    if engine is None:
        raise HTTPException(status_code=503, detail="model is still starting")
    if req.response_format != "b64_json":
        raise HTTPException(
            status_code=400,
            detail=f"response_format {req.response_format!r} is not supported; this "
                   "server returns b64_json only",
        )
    # One denoise loop at a time: the pipeline owns the device.
    async with app.state.device_lock:
        frames = await asyncio.to_thread(_generate, engine, req)
    # Off the loop as well, and NOT because it is slow in absolute terms: PIL's GIF
    # palette quantization plus base64 of a multi-MB payload is seconds of pure CPU, and
    # run inline here it blocks every coroutine -- including /health and /tt-liveness,
    # whose whole purpose is to answer while this endpoint is busy. An orchestrator with a
    # 1-2 s liveness timeout would restart a server that is merely encoding. Outside the
    # device lock, because the device is free once the frames are in host memory.
    b64 = await asyncio.to_thread(_frames_to_gif_b64, frames)
    return VideoGenerationResponse(
        created=int(time.time()),
        data=[VideoData(b64_json=b64)],
    )

	# SPDX-License-Identifier: Apache-2.0
	# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
	"""ASGI app serving tt-animatediff, for tt-model-manager's ``tt-dit-server`` kind.

	A diffusion transformer has no tokens, no KV cache and no continuous batching, so vLLM has
	nothing to do for it. ``tt-dit-server`` therefore installs a small HTTP stack instead of an
	engine and launches this module with uvicorn:

	python -m uvicorn --host 0.0.0.0 --port 8000 --lifespan on \\
	animatediff_ttnn.server.app:app

	TWO CONTRACTS THIS FILE HAS TO HONOUR
	-------------------------------------
	1. Readiness is the lifespan. tt-model-manager decides the server is up when uvicorn
	logs ``Application startup complete``, which it prints after ASGI lifespan startup
	returns. So the device open and the model warm belong in the lifespan and nowhere else: do
	them lazily on the first request and the supervisor calls the server ready while it is
	still loading, then times out the first request it sends.

	2. Importing this module must not touch hardware. ``verify_lines`` imports the ASGI
	attribute at image-build time, on a machine with no card, purely to prove the code
	allowlist actually shipped the server. Every ttnn/tt-metal import is therefore inside a
	function, never at module scope. That is also what lets the tests below run on CPU.

	ENVIRONMENT, set by the launcher
	--------------------------------
	``MESH_DEVICE`` the SKU, e.g. ``P300x2`` — informational here
	``ANIMATEDIFF_MESH_SHAPE`` the resolved shape, e.g. ``1x4`` (the manifest points
	``runtime.mesh_shape_env`` at this name; the kind's default is
	FLUX.2's, which means nothing to this model)
	``HF_MODEL`` the weights repo id, reported by ``/v1/models``

	ONE REQUEST AT A TIME
	---------------------
	There is no continuous batching to hide behind: the pipeline owns the device, and two
	concurrent denoise loops would interleave on it. Requests are serialised on a lock and the
	blocking work runs in a worker thread so the event loop can still answer ``/health`` while
	a generation is in flight. This is a deliberate cap, not an oversight — a queue that
	admits work the device cannot do concurrently is a queue that reports success and returns
	corrupted frames.
	"""
	from __future__ import annotations

	import asyncio
	import base64
	import io
	import os
	import time
	from contextlib import asynccontextmanager
	from typing import Any, List, Optional

	from fastapi import FastAPI, HTTPException
	from pydantic import BaseModel, Field

	#: The env var carrying the resolved mesh shape. Named for THIS model: the kind's default
	#: is ``FLUX2_MESH_SHAPE``, which is the default only because FLUX.2 is what it was built
	#: for. A manifest serving this app sets ``runtime.mesh_shape_env`` to the name below.
	MESH_SHAPE_ENV = "ANIMATEDIFF_MESH_SHAPE"

	#: Frames per second baked into the returned GIF.
	DEFAULT_FPS = 8


	#: The only image size this server can produce. The TTNN UNet is compiled at a fixed
	#: 64x64 latent, and 64 * 8 = 512. Named rather than repeated so the request model, the
	#: error message and any future widening all move together.
	SERVABLE_SIZE = 512


	class VideoGenerationRequest(BaseModel):
	"""OpenAI-shaped video request.

	Field names follow ``/v1/videos/generations`` as tt-inference-server's tt-media-server
	already serves it, so the same model can sit behind either surface without a client
	change. Everything beyond ``prompt`` is optional with the pipeline's own defaults.
	"""

	prompt: str = Field(min_length=1)
	negative_prompt: str = ""
	model: Optional[str] = None
	num_frames: int = Field(default=16, ge=1, le=64)
	num_inference_steps: int = Field(default=25, ge=1, le=100)
	guidance_scale: float = Field(default=7.5, ge=0.0, le=20.0)
	seed: int = 42
	# 0.35, because that is what the rest of the project calibrated: generate_animation(),
	# hf/pipeline.py AND generate_frames_temporal -- the function this server calls -- all
	# use it. This said 0.5, which matched none of them, so an HTTP client omitting the
	# field got quietly different motion from every other entry point into the same code.
	# test_the_server_request_default_matches_the_librarys_calibrated_one asserts it
	# against the library rather than against a literal, so it cannot drift again.
	temporal_alpha: float = Field(default=0.35, ge=0.0, le=1.0)
	# 512 exactly, and not as a conservative default: the TTNN UNet is compiled ONCE at a
	# fixed 64x64 latent (generation_helpers.py: `UNet2D(device, parameters, 2, 64, 64)`),
	# so 512x512 is the only thing this server can produce. These were bounded 64..1024,
	# which is pydantic-valid and kernel-invalid -- the request passed validation, took the
	# device lock, and then raised a TT_FATAL deep in the denoise loop as an uncaught 500,
	# having held the one device to do it. Refuse at the edge instead, with a message that
	# says why. Widening this means compiling the UNet per requested shape, which is a
	# ~2-3 minute cost per shape and a different design.
	height: int = Field(default=SERVABLE_SIZE, ge=SERVABLE_SIZE, le=SERVABLE_SIZE,
	description=f"only {SERVABLE_SIZE} is supported: the UNet is "
	"compiled at a fixed 64x64 latent")
	width: int = Field(default=SERVABLE_SIZE, ge=SERVABLE_SIZE, le=SERVABLE_SIZE,
	description=f"only {SERVABLE_SIZE} is supported: the UNet is "
	"compiled at a fixed 64x64 latent")
	#: Only b64_json is offered. A URL response would need the server to host files it has
	#: no store for, and a bundle that returns links to a path the consumer cannot read is
	#: worse than one that returns bytes.
	response_format: str = "b64_json"


	class VideoData(BaseModel):
	b64_json: str


	class VideoGenerationResponse(BaseModel):
	created: int
	data: List[VideoData]


	def mesh_shape_from_env(env: Optional[dict] = None) -> tuple:
	"""Parse ``ANIMATEDIFF_MESH_SHAPE`` (``"RxC"``) into ``(rows, cols)``.

	Defaults to ``(1, 1)`` when unset, which is what a bare local run wants. A malformed
	value RAISES rather than falling back: silently converting against the wrong mesh is
	the failure this variable exists to prevent, and it would show up as bad frames rather
	than as an error.
	"""
	raw = (env if env is not None else os.environ).get(MESH_SHAPE_ENV, "").strip()
	if not raw:
	return (1, 1)
	try:
	rows, cols = (int(p) for p in raw.lower().split("x", 1))
	except ValueError as exc:
	raise ValueError(
	f"{MESH_SHAPE_ENV}={raw!r} is not a mesh shape like '1x4'"
	) from exc
	if rows < 1 or cols < 1:
	raise ValueError(f"{MESH_SHAPE_ENV}={raw!r} must be positive")
	# Only 1x1, and refused here rather than absorbed. _open_device_and_models passes
	# `device_ids=None` for any other shape, and setup_blackhole(None) opens EVERY chip on
	# the box as a 1xN mesh -- so on a shared machine a stray 1x4 takes four chips out of
	# circulation, and then dies in the lifespan anyway, because this model cannot shard:
	# a real serve at P300x2 failed with `Can't convert a tensor distributed on
	# MeshShape([1, 4]) mesh to row-major logical tensor`. The manifest declares P150 for
	# the same reason. Widen this only alongside the mesh frame-sharding work.
	if rows * cols != 1:
	raise ValueError(
	f"{MESH_SHAPE_ENV}={raw!r} requests {rows * cols} chips, but this model cannot "
	f"shard -- only 1x1 is supported until mesh frame-sharding ships "
	f"(docs/superpowers/plans/2026-06-15-mesh-frame-sharding.md). A larger shape "
	f"would claim every chip on the box and still fail in the lifespan."
	)
	return (rows, cols)


	def _open_device_and_models(shape: tuple) -> tuple:
	"""Claim the mesh and load the TTNN model set. Blocking, hardware, lifespan-only."""
	from animatediff_ttnn.generation_helpers import load_sd14_ttnn
	from animatediff_ttnn.ttnn_pipeline import setup_blackhole

	from animatediff_ttnn.generation_helpers import encode_prompt

	rows, cols = shape
	# setup_blackhole(device_ids=None) opens every available chip as a 1xN mesh, which is
	# what a (1, N) shape means. An explicit id list is passed only for a single chip so a
	# 1x1 run cannot accidentally claim a neighbour someone else is using.
	device = setup_blackhole(device_ids=[0] if rows * cols == 1 else None)
	# Everything after the open is wrapped: a failure here used to leave the chip claimed
	# by a container that never started, and on a restart loop that is one claimed device
	# per attempt while TTNN refuses to reopen any of them. Same leak session.py was fixed
	# for; the server has its own open path, so it needs its own rollback.
	try:
	models = load_sd14_ttnn(device)

	# Warm the TEXT encoder too, and not as an optimisation. load_sd14_ttnn brings up the
	# UNet and the VAE; CLIP's tokenizer and text encoder were left to be downloaded
	# lazily by the first encode_prompt() call -- which happens inside _generate, under
	# the device lock, AFTER /health has already answered 200. So the readiness contract
	# this module is built on ("the device open and the model warm happen in the lifespan
	# and nowhere else") was false for one of the three models, and in a container with
	# only unet/vae cached and no egress, the server reported ready and then failed every
	# request. Downloading here means a missing weight fails startup, which is what a
	# readiness probe is for.
	#
	# NOT a full denoise pass: that would also pay TTNN kernel compilation here, which is
	# minutes on a cold cache and would delay readiness for every boot. The first request
	# is therefore still slower than the rest -- slow, but correct, and measured in
	# docs/measurements/serving-benchmark.json rather than hidden.
	encode_prompt("warmup", "")
	except Exception:
	_close_device(device)
	raise
	return device, models


	def _close_device(device) -> None:
	"""Release the mesh, swallowing any error.

	Swallowed deliberately: on the failure path this must not mask the exception the
	operator actually needs to read, and a close that fails is not a reason to hold the
	chip. Mirrors animatediff_ttnn.session._close_device.
	"""
	try:
	import ttnn
	ttnn.close_mesh_device(device)
	except Exception:
	pass


	def _frames_to_gif_b64(frames: List[Any], fps: int = DEFAULT_FPS) -> str:
	"""PIL frames -> base64 GIF, in memory. Never touches disk: the container has no
	volume to write to and a served artifact should not depend on one."""
	buf = io.BytesIO()
	frames[0].save(
	buf, format="GIF", save_all=True, append_images=frames[1:],
	duration=1000 // max(fps, 1), loop=0,
	)
	return base64.b64encode(buf.getvalue()).decode("ascii")


	def _generate(state: dict, req: VideoGenerationRequest) -> List[Any]:
	"""Run the denoise loop. Blocking; called in a worker thread under the device lock."""
	from animatediff_ttnn.generation_helpers import encode_prompt
	from animatediff_ttnn.temporal_attention import generate_frames_temporal

	ttnn_model, ttnn_vae, config, torch_time_proj = state["models"]
	return generate_frames_temporal(
	device=state["device"],
	ttnn_model=ttnn_model,
	ttnn_vae=ttnn_vae,
	config=config,
	torch_time_proj=torch_time_proj,
	text_embeddings=encode_prompt(req.prompt, req.negative_prompt),
	num_frames=req.num_frames,
	num_steps=req.num_inference_steps,
	guidance_scale=req.guidance_scale,
	seed=req.seed,
	temporal_alpha=req.temporal_alpha,
	height=req.height,
	width=req.width,
	)


	@asynccontextmanager
	async def lifespan(app: FastAPI):
	"""Claim the device and warm the pipeline BEFORE the server reports ready.

	Run in a worker thread so a slow warm cannot wedge the event loop, and awaited to
	completion so ``Application startup complete`` means what tt-model-manager reads it to
	mean. A failure here is deliberately fatal: a server that starts without a device
	would answer /health cheerfully and fail every generation.
	"""
	shape = mesh_shape_from_env()
	device, models = await asyncio.to_thread(_open_device_and_models, shape)
	app.state.engine = {
	"device": device, "models": models, "mesh_shape": shape,
	"hf_model": os.environ.get("HF_MODEL", "unknown"),
	"mesh_device": os.environ.get("MESH_DEVICE", "unknown"),
	}
	app.state.device_lock = asyncio.Lock()
	try:
	yield
	finally:
	# Let ttnn close the mesh itself. Bypassing its teardown is what triggers an abort
	# in MetalContext::destroy_all_instances.
	try:
	import ttnn

	ttnn.close_mesh_device(device)
	except Exception: # noqa: BLE001 - shutdown must not mask the original error
	pass


	app = FastAPI(title="tt-animatediff", lifespan=lifespan)


	def _readiness() -> dict:
	"""What both probes report about the model, in one place so they cannot disagree."""
	engine = getattr(app.state, "engine", None)
	return {
	"model_ready": engine is not None,
	"model": (engine or {}).get("hf_model"),
	"mesh_device": (engine or {}).get("mesh_device"),
	"mesh_shape": "x".join(str(n) for n in (engine or {}).get("mesh_shape", ())),
	}


	@app.get("/tt-liveness")
	async def tt_liveness() -> dict:
	"""LIVENESS probe -- "is this process able to respond at all", restart on failure.

	Matches tt-inference-server's tt-media-server
	(``open_ai_api/tt_maintenance_api.py``): 200 with ``{"status": "alive", ...}`` plus the
	model-readiness payload, and a non-200 ONLY on unrecoverable failure. A model that is
	still warming is alive, not broken -- that is exactly the distinction this endpoint
	exists to draw, and reporting it as failure would have an orchestrator restart a server
	that is loading correctly.

	Deliberately lock-free and engine-free, so it keeps answering while a generation holds
	the device. A liveness probe that blocks behind the work it is monitoring is how a
	healthy-but-busy server gets killed.
	"""
	try:
	return {"status": "alive", **_readiness()}
	except Exception as exc: # noqa: BLE001 - the one case that IS unrecoverable
	raise HTTPException(status_code=500, detail=f"Liveness check failed: {exc}")


	@app.get("/health")
	async def health() -> dict:
	"""READINESS probe -- "should traffic be routed here", gate on it.

	Despite the names, `/health` is readiness and `/tt-liveness` is liveness. That is
	tt-media-server's convention, kept for compatibility with the Python server, and the
	body is vLLM's: an empty dict when ready, 503 when not.

	This returned 200 with ``{"status": "starting"}`` until the semantics were checked
	against tt-media-server. A readiness probe answering 200 while the pipeline is still
	loading tells a load balancer to send work to a server that cannot do it -- the precise
	failure the lifespan design exists to prevent, reintroduced one layer up. The
	diagnostic payload now lives on /tt-liveness, which is where tt-media-server puts it.
	"""
	if not _readiness()["model_ready"]:
	raise HTTPException(status_code=503, detail="Model not ready")
	return {}


	@app.get("/v1/models")
	async def models() -> dict:
	engine = getattr(app.state, "engine", None)
	name = (engine or {}).get("hf_model", "unknown")
	return {"object": "list", "data": [{"id": name, "object": "model",
	"owned_by": "tenstorrent"}]}


	@app.post("/v1/videos/generations", response_model=VideoGenerationResponse)
	async def videos_generations(req: VideoGenerationRequest) -> VideoGenerationResponse:
	engine = getattr(app.state, "engine", None)
	if engine is None:
	raise HTTPException(status_code=503, detail="model is still starting")
	if req.response_format != "b64_json":
	raise HTTPException(
	status_code=400,
	detail=f"response_format {req.response_format!r} is not supported; this "
	"server returns b64_json only",
	)
	# One denoise loop at a time: the pipeline owns the device.
	async with app.state.device_lock:
	frames = await asyncio.to_thread(_generate, engine, req)
	# Off the loop as well, and NOT because it is slow in absolute terms: PIL's GIF
	# palette quantization plus base64 of a multi-MB payload is seconds of pure CPU, and
	# run inline here it blocks every coroutine -- including /health and /tt-liveness,
	# whose whole purpose is to answer while this endpoint is busy. An orchestrator with a
	# 1-2 s liveness timeout would restart a server that is merely encoding. Outside the
	# device lock, because the device is free once the frames are in host memory.
	b64 = await asyncio.to_thread(_frames_to_gif_b64, frames)
	return VideoGenerationResponse(
	created=int(time.time()),
	data=[VideoData(b64_json=b64)],
	)