File size: 10,669 Bytes
1e3ce4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
"""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)