studio/backends/animatediff.py

"""AnimateDiff backend adapter (Apache 2.0).

Wraps `diffusers.AnimateDiffSparseControlNetPipeline` so a `PixelCursor` over a
single Image can be animated through one of the 8 verified MotionLoRA presets
(zoom, pan, tilt, roll). SparseCtrl is the I2V path: frame 0 is pinned to the
cursor's source image, frames 1..N-1 are diffused with motion-LoRA conditioning.

Light surface (works without torch):
  - `register()` wires bind_motion_exemplar + write_motion into the op registry
  - `bind_motion_exemplar(cursor, ["zoom_in"])` packs the resolved MotionLoRA
    id into an IdentityLock — no model load
  - `dry_run(cursor)` returns a diagnostic dict (what would happen, what it costs)

Heavy surface (requires `pip install -e '.[backends]'`):
  - `write_motion(cursor, motion_spec)` lazy-imports torch + diffusers, loads
    the pipeline, runs SparseCtrl-conditioned generation, returns FrameStack

References:
  diffusers AnimateDiffSparseControlNetPipeline docs (Apache 2.0)
  guoyww/animatediff-motion-adapter-v1-5-2 (Apache 2.0)
  guoyww/animatediff-sparsectrl-rgb (Apache 2.0)
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Sequence

import numpy as np

from pixel_cursor import PixelCursor, Image, IdentityLock, FrameStack, ops
from pixel_cursor.artifact import _new_framestack


# Verified by A5 research (background subagent, confidence: high)
MOTION_LORA_MAP: dict[str, str] = {
    "zoom_in": "guoyww/animatediff-motion-lora-zoom-in",
    "zoom_out": "guoyww/animatediff-motion-lora-zoom-out",
    "pan_left": "guoyww/animatediff-motion-lora-pan-left",
    "pan_right": "guoyww/animatediff-motion-lora-pan-right",
    "tilt_up": "guoyww/animatediff-motion-lora-tilt-up",
    "tilt_down": "guoyww/animatediff-motion-lora-tilt-down",
    "roll_anticlockwise": "guoyww/animatediff-motion-lora-rolling-anticlockwise",
    "roll_clockwise": "guoyww/animatediff-motion-lora-rolling-clockwise",
}

DEFAULT_MOTION_ADAPTER = "guoyww/animatediff-motion-adapter-v1-5-2"
DEFAULT_SPARSECTRL = "guoyww/animatediff-sparsectrl-rgb"
DEFAULT_BASE_MODEL = "SG161222/Realistic_Vision_V5.1_noVAE"

INSTALL_HINT = (
    "AnimateDiff inference requires torch + diffusers. Install with:\n"
    "  pip install -e '.[backends]'\n"
    "(or: pip install torch diffusers transformers)"
)


class AnimateDiffAdapter:
    """Adapter exposing AnimateDiff primitives to the PixelCursor op registry."""

    SOURCE_TAG = "animatediff"

    def __init__(
        self,
        motion_adapter_id: str = DEFAULT_MOTION_ADAPTER,
        base_model_id: str = DEFAULT_BASE_MODEL,
        sparsectrl_id: str = DEFAULT_SPARSECTRL,
    ) -> None:
        self.motion_adapter_id = motion_adapter_id
        self.base_model_id = base_model_id
        self.sparsectrl_id = sparsectrl_id

    def register(self) -> None:
        ops.register_backend("bind_motion_exemplar", "animatediff", self.bind_motion_exemplar)
        ops.register_backend("write_motion", "animatediff", self.write_motion)

    def bind_motion_exemplar(
        self, cursor: PixelCursor, exemplar: Sequence[str]
    ) -> PixelCursor:
        """Resolve a preset name (e.g. ["zoom_in"]) to a MotionLoRA id and pack into
        an IdentityLock on the cursor. No model load happens here.

        Accepts a sequence to keep the signature uniform with MotionDirector
        (which takes K reference clips). For AnimateDiff a single preset is expected;
        a list of clip paths is rejected with a helpful error pointing to session 3.
        """
        if len(exemplar) != 1:
            raise ValueError(
                f"animatediff.bind_motion_exemplar expects exactly one preset name; "
                f"got {len(exemplar)} items. For few-shot training from K reference "
                f"clips, use backend='motiondirector' (session 3)."
            )
        preset = exemplar[0]
        if preset.endswith((".mp4", ".webm", ".mkv", ".gif", ".avi", ".mov")):
            raise ValueError(
                f"animatediff backend takes preset names from MOTION_LORA_MAP "
                f"(e.g. 'zoom_in'); got what looks like a video path: {preset!r}. "
                f"For exemplar-based few-shot, use backend='motiondirector'."
            )
        if preset not in MOTION_LORA_MAP:
            valid = sorted(MOTION_LORA_MAP)
            raise ValueError(
                f"Unknown AnimateDiff preset {preset!r}. Valid presets: {valid}"
            )
        lora_id = MOTION_LORA_MAP[preset]
        lock = IdentityLock(
            embedding=np.zeros(0, dtype=np.float32),
            source=f"{self.SOURCE_TAG}:{preset}:{lora_id}",
        )
        return cursor.lock_identity(lock)

    def dry_run(self, cursor: PixelCursor, motion_spec: dict | None = None) -> dict:
        """Return what `write_motion` WOULD do without running it.

        Useful for verifying the preset mapping and resource estimate before
        paying the install + model-download cost. Safe to call without torch.
        """
        if cursor.identity_lock is None or not cursor.identity_lock.source.startswith(
            f"{self.SOURCE_TAG}:"
        ):
            raise ValueError(
                "dry_run requires bind_motion_exemplar(..., backend='animatediff') first"
            )
        if not isinstance(cursor.artifact, Image):
            raise TypeError(
                f"AnimateDiff dry_run expects an Image artifact; got {type(cursor.artifact).__name__}"
            )

        _, preset, lora_id = cursor.identity_lock.source.split(":", 2)
        ms = motion_spec or {}
        return {
            "backend": "animatediff",
            "pipeline_class": "diffusers.AnimateDiffSparseControlNetPipeline",
            "base_model": self.base_model_id,
            "motion_adapter": self.motion_adapter_id,
            "sparsectrl": self.sparsectrl_id,
            "preset": preset,
            "motion_lora_id": lora_id,
            "input_image_shape": tuple(cursor.artifact.pixels.shape),
            "num_frames": ms.get("num_frames", 16),
            "native_fps": 8,
            "target_fps": 24,
            "post_interpolation": "RIFE 3x (8fps -> 24fps), MIT licensed",
            "expected_output_shape": (ms.get("num_frames", 16), *cursor.artifact.pixels.shape),
            "hardware_floor_vram_gb": 13,
            "expected_wall_time_s_4090": 45,
            "prompt": ms.get("prompt", "(none — cursor has no text prompt yet)"),
            "guidance_scale": ms.get("guidance_scale", 7.5),
            "num_inference_steps": ms.get("num_inference_steps", 25),
            "torch_installed": _torch_available(),
            "note": (
                "If torch_installed is False, `write_motion` will raise ImportError. "
                "Install with: pip install -e '.[backends]'"
            ),
        }

    def write_motion(self, cursor: PixelCursor, motion_spec: dict | None = None) -> FrameStack:
        if cursor.identity_lock is None or not cursor.identity_lock.source.startswith(
            f"{self.SOURCE_TAG}:"
        ):
            raise ValueError(
                "write_motion(animatediff) requires bind_motion_exemplar(..., backend='animatediff') first"
            )
        if not isinstance(cursor.artifact, Image):
            raise TypeError(
                f"AnimateDiff write_motion expects an Image artifact; got {type(cursor.artifact).__name__}"
            )

        try:
            import torch  # noqa: F401
            from diffusers import (  # noqa: F401
                AnimateDiffSparseControlNetPipeline,
                MotionAdapter,
                SparseControlNetModel,
                DDIMScheduler,
            )
        except ImportError as e:
            raise ImportError(INSTALL_HINT) from e

        return self._run_inference(cursor, motion_spec or {})

    def _run_inference(self, cursor: PixelCursor, motion_spec: dict) -> FrameStack:
        import torch
        from PIL import Image as PILImage
        from diffusers import (
            AnimateDiffSparseControlNetPipeline,
            MotionAdapter,
            SparseControlNetModel,
            DDIMScheduler,
        )

        _, preset, lora_id = cursor.identity_lock.source.split(":", 2)
        dtype = torch.float16

        motion_adapter = MotionAdapter.from_pretrained(self.motion_adapter_id, torch_dtype=dtype)
        controlnet = SparseControlNetModel.from_pretrained(self.sparsectrl_id, torch_dtype=dtype)
        pipe = AnimateDiffSparseControlNetPipeline.from_pretrained(
            self.base_model_id,
            motion_adapter=motion_adapter,
            controlnet=controlnet,
            torch_dtype=dtype,
        )
        pipe.scheduler = DDIMScheduler.from_config(
            pipe.scheduler.config,
            clip_sample=False,
            timestep_spacing="linspace",
            beta_schedule="linear",
            steps_offset=1,
        )
        pipe.enable_vae_slicing()
        pipe.enable_model_cpu_offload()
        pipe.load_lora_weights(lora_id, adapter_name=preset)
        pipe.set_adapters([preset], adapter_weights=[motion_spec.get("lora_weight", 1.0)])

        conditioning_pil = PILImage.fromarray(cursor.artifact.pixels)
        output = pipe(
            prompt=motion_spec.get("prompt", "high quality, detailed"),
            negative_prompt=motion_spec.get("negative_prompt", "bad quality, blurry"),
            num_frames=motion_spec.get("num_frames", 16),
            guidance_scale=motion_spec.get("guidance_scale", 7.5),
            num_inference_steps=motion_spec.get("num_inference_steps", 25),
            conditioning_frames=[conditioning_pil],
            controlnet_conditioning_scale=motion_spec.get("controlnet_scale", 1.0),
            controlnet_frame_indices=[0],
            generator=torch.Generator("cpu").manual_seed(motion_spec.get("seed", 42)),
        )
        frames_pil = output.frames[0]
        frames_np = np.stack([np.array(f) for f in frames_pil], axis=0).astype(np.uint8)
        return _new_framestack(frames_np, fps=8, name=f"animatediff:{preset}")


@dataclass(frozen=True)
class _ProbeResult:
    torch_installed: bool
    diffusers_installed: bool


def _torch_available() -> bool:
    try:
        import torch  # noqa: F401
        return True
    except ImportError:
        return False


def probe() -> _ProbeResult:
    """Diagnostic: are the heavy deps installed?"""
    try:
        import torch  # noqa: F401
        t = True
    except ImportError:
        t = False
    try:
        import diffusers  # noqa: F401
        d = True
    except ImportError:
        d = False
    return _ProbeResult(torch_installed=t, diffusers_installed=d)