strands-diffusers 0.1.0__py3-none-any.whl

strands_diffusers/core/registry.py

@@ -0,0 +1,349 @@
+"""Dynamic pipeline/model/scheduler registry — 100% diffusers coverage, zero hardcoding.
+
+The single source of truth is diffusers' own `_import_structure` — the lazy map of
+every public symbol it exposes (307 pipelines, 87 models, 54 schedulers as of
+0.38). We read it at runtime, so when diffusers adds a new pipeline (e.g. a fresh
+Cosmos world-foundation model), strands-diffusers supports it automatically — no
+code change required.
+
+Same philosophy as `use_aws` (wraps boto3 dynamically), `use_lerobot` (wraps
+lerobot) and `use_transformers` (wraps the transformers task taxonomy): discover,
+don't hardcode.
+
+Diffusers has no single "task taxonomy" like transformers' SUPPORTED_TASKS, so we
+derive structure from three places, all dynamic:
+
+1. `diffusers._import_structure`        → every public class, grouped by submodule.
+2. The `AutoPipelineFor*` mappings      → the canonical task → pipeline-class maps
+   (text2image / image2image / inpainting), diffusers' closest thing to tasks.
+3. Class-name heuristics                → group the long tail of pipelines by the
+   modality their name implies (TextToImage / ImageToVideo / VideoToWorld / ...).
+"""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+import re
+from functools import lru_cache
+from typing import Any, Dict, List, Optional
+
+
+@lru_cache(maxsize=1)
+def _import_structure() -> Dict[str, List[str]]:
+    import diffusers
+
+    ist = getattr(diffusers, "_import_structure", None) or {}
+    # Flatten dotted keys but keep the group name (submodule) for context.
+    return {k: list(v) for k, v in ist.items()}
+
+
+@lru_cache(maxsize=1)
+def all_symbols() -> Dict[str, str]:
+    """Every public diffusers symbol → its kind (pipeline/model/scheduler/other)."""
+    out: Dict[str, str] = {}
+    for syms in _import_structure().values():
+        for s in syms:
+            out[s] = _classify(s)
+    return out
+
+
+def _classify(name: str) -> str:
+    if name.endswith("Pipeline"):
+        return "pipeline"
+    if name.endswith("Scheduler") or name.endswith("SchedulerOutput"):
+        return "scheduler"
+    if name.endswith("Output"):
+        return "output"
+    if "Model" in name or name.startswith("Autoencoder") or name.endswith("Transformer"):
+        return "model"
+    return "other"
+
+
+@lru_cache(maxsize=1)
+def pipelines() -> List[str]:
+    return sorted(n for n, k in all_symbols().items() if k == "pipeline")
+
+
+@lru_cache(maxsize=1)
+def models() -> List[str]:
+    return sorted(n for n, k in all_symbols().items() if k == "model")
+
+
+@lru_cache(maxsize=1)
+def schedulers() -> List[str]:
+    return sorted(n for n, k in all_symbols().items() if k == "scheduler"
+                  and not n.endswith("Output"))
+
+
+# ───────────────────────── AutoPipeline task maps ─────────────────────────
+
+@lru_cache(maxsize=1)
+def auto_pipeline_tasks() -> Dict[str, Dict[str, str]]:
+    """diffusers' canonical task → {model-family: pipeline-class} maps.
+
+    These are the only first-class "tasks" diffusers ships: text2image,
+    image2image, inpainting. Returned as plain class-name strings.
+    """
+    from diffusers.pipelines.auto_pipeline import (
+        AUTO_TEXT2IMAGE_PIPELINES_MAPPING,
+        AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
+        AUTO_INPAINT_PIPELINES_MAPPING,
+    )
+
+    def _names(m):
+        return {family: cls.__name__ for family, cls in m.items()}
+
+    return {
+        "text-to-image": _names(AUTO_TEXT2IMAGE_PIPELINES_MAPPING),
+        "image-to-image": _names(AUTO_IMAGE2IMAGE_PIPELINES_MAPPING),
+        "inpainting": _names(AUTO_INPAINT_PIPELINES_MAPPING),
+    }
+
+
+# ───────────────────────── modality grouping (derived) ─────────────────────────
+
+# Ordered (pattern → modality) rules applied to a pipeline class name. First match
+# wins. Purely name-derived so new pipelines slot in without code changes.
+_MODALITY_RULES = (
+    # Explicit transition names first (most specific).
+    (r"TextToImage|Text2Image", "text-to-image"),
+    (r"TextToVideo|Text2Video|TextToWorld", "text-to-video"),
+    (r"ImageToVideo|Image2Video", "image-to-video"),
+    (r"VideoToVideo|Video2Video", "video-to-video"),
+    (r"VideoToWorld|World", "video-to-world"),
+    (r"ImageToImage|Image2Image", "image-to-image"),
+    (r"Inpaint", "inpainting"),
+    (r"Upscale|SuperResolution", "upscaling"),
+    (r"TextToAudio|Text2Audio|Audio|Music|Speech|TTS|Bark|MusicGen|Stable[Aa]udio", "audio"),
+    # Image families that share a name-stem with a video family (e.g. HunyuanDiT,
+    # HunyuanImage) — classify as image BEFORE the broad video catch-all below.
+    (r"DiT|HunyuanImage|HunyuanDiT", "image"),
+    # Broad video/world model families (named after the architecture, not a task).
+    (r"Video|Animate|Cosmos|Wan|Hunyuan|Mochi|Allegro|LTX|SkyReels|CogVideo|Latte|Genie", "video"),
+    (r"ControlNet|Adapter|IP", "controlled-image"),
+    # Abbreviated transition tokens — MUST precede architecture-family rules so a
+    # family-named video pipeline (Kandinsky5I2V, *T2V) isn't grabbed as image.
+    (r"I2V", "image-to-video"),
+    (r"T2V", "text-to-video"),
+    (r"V2V", "video-to-video"),
+    # Transition variants on family-named pipelines (task encoded as a suffix).
+    (r"Img2Img|Image2Image", "image-to-image"),
+    (r"Fill", "inpainting"),
+    # NOTE: no bare "Edit" rule — editing can be image OR video (e.g. ChronoEdit
+    # is image-to-video). Let such names fall through to their family/Video rule.
+    # Architecture-named image-gen families (task implicit = text-to-image-class).
+    (r"StableDiffusion|Flux|CogView|Bria|Chroma|AuraFlow|Amused|Kandinsky|"
+     r"PixArt|Sana|Lumina|DeepFloyd|Wuerstchen|Kolors|HiDream|Janus|OmniGen|"
+     r"Marigold|VisualCloze", "image"),
+    # Architecture-named audio families.
+    (r"AudioLDM|StableAudio|MusicGen|Musicgen|AceStep|Dance|Spectrogram", "audio"),
+    # 3D / mesh generation (ShapE, etc.) — emits meshes, not images.
+    (r"ShapE|Shap[_-]?E|Mesh|TripoSR|Hunyuan3D|TRELLIS", "3d"),
+    (r"Image", "image"),
+)
+
+
+# Docstring phrasings → modality. Used ONLY as a fallback when name-rules yield
+# "other". Rules stay authoritative (they encode WFM/inpaint precedence the docs
+# lack); the doc parser just rescues the long tail. Ordered: most-specific first.
+_DOC_MODALITY = (
+    ("text-to-3d", "3d"), ("image-to-3d", "3d"), ("text-to-mesh", "3d"),
+    ("image-to-image", "image-to-image"),
+    ("image-to-video", "image-to-video"),
+    ("video-to-video", "video-to-video"),
+    ("text-to-video", "text-to-video"),
+    ("text-to-image", "text-to-image"),
+    ("text-to-audio", "audio"), ("text-to-speech", "audio"),
+    ("super-resolution", "upscaling"),
+    ("inpainting", "inpainting"),
+    ("unconditional image", "image"), ("class-conditional image", "image"),
+    ("image generation", "image"), ("image synthesis", "image"),
+    ("video generation", "video"), ("audio generation", "audio"),
+)
+
+_DOC_PAT = re.compile(
+    r"[Pp]ipeline for ([\w\- ]+?) (?:generation|synthesis|using|based|with)", re.S)
+
+
+@lru_cache(maxsize=512)
+def _modality_from_doc(name: str) -> str:
+    """Best-effort modality from a pipeline class docstring. Cached; never raises."""
+    try:
+        cls = resolve_attr(name)
+        doc = (cls.__doc__ or "").strip().lower()
+    except Exception:
+        return "other"
+    if not doc:
+        return "other"
+    m = _DOC_PAT.search(doc)
+    phrase = m.group(1).strip() if m else doc[:80]
+    for key, mod in _DOC_MODALITY:
+        if key in phrase:
+            return mod
+    return "other"
+
+
+def modality_of(pipeline_name: str, use_doc: bool = False) -> str:
+    """Derive a pipeline's modality. Name-rules are authoritative; if they yield
+    "other" and use_doc=True, fall back to docstring parsing (slower: imports the
+    class). Keep use_doc=False for hot paths / deterministic tests."""
+    for pat, mod in _MODALITY_RULES:
+        if re.search(pat, pipeline_name):
+            return mod
+    if use_doc:
+        return _modality_from_doc(pipeline_name)
+    return "other"
+
+
+def tasks_by_modality(use_doc: bool = True) -> Dict[str, List[str]]:
+    groups: Dict[str, List[str]] = {}
+    for p in pipelines():
+        groups.setdefault(modality_of(p, use_doc=use_doc), []).append(p)
+    for v in groups.values():
+        v.sort()
+    return groups
+
+
+# World-foundation-model / action-capable pipelines (the ones the agent cares
+# about for robotics). Detected by name — Cosmos*, *World*, action-conditioned.
+def world_foundation_models() -> List[str]:
+    return sorted(p for p in pipelines()
+                  if re.search(r"Cosmos|World|Wan|Hunyuan|Genie", p))
+
+
+# ───────────────────────── resolution & introspection ─────────────────────────
+
+def resolve_attr(dotted: str, root_module: str = "diffusers") -> Any:
+    """Resolve a dotted path against diffusers (or a submodule).
+
+    Examples:
+        resolve_attr("StableDiffusionPipeline")
+        resolve_attr("DiffusionPipeline.from_pretrained")
+        resolve_attr("Cosmos3OmniPipeline")          # from-source builds
+        resolve_attr("utils.export_to_video")
+        resolve_attr("schedulers.UniPCMultistepScheduler.from_config")
+    """
+    full = dotted if dotted.startswith(root_module + ".") else f"{root_module}.{dotted}"
+
+    try:
+        return importlib.import_module(full)
+    except ImportError:
+        pass
+
+    # Fast path: attribute(s) on the root module (diffusers uses lazy __getattr__
+    # that raises AttributeError, not ImportError, for non-module attrs).
+    try:
+        root = importlib.import_module(root_module)
+        obj = root
+        for attr in dotted.split("."):
+            obj = getattr(obj, attr)
+        return obj
+    except AttributeError:
+        pass
+
+    segments = full.split(".")
+    for i in range(len(segments), 0, -1):
+        try:
+            mod = importlib.import_module(".".join(segments[:i]))
+        except Exception:
+            continue
+        obj = mod
+        try:
+            for attr in segments[i:]:
+                obj = getattr(obj, attr)
+            return obj
+        except AttributeError:
+            break
+
+    root = importlib.import_module(root_module)
+    obj = root
+    for attr in dotted.split("."):
+        obj = getattr(obj, attr)
+    return obj
+
+
+def pipeline_info(name: str) -> Optional[Dict[str, Any]]:
+    """Modality + __call__ signature for one pipeline class (lazily resolved)."""
+    if name not in all_symbols():
+        # Known from-source WFM/pipeline classes (e.g. Cosmos3OmniPipeline ships in
+        # diffusers>0.38 from source). Degrade gracefully instead of erroring like a
+        # typo would — the tool resolves these dynamically once the install has them.
+        if re.search(r"Pipeline$", name) and re.search(
+                r"Cosmos|World|Wan|Hunyuan|Genie|Omni", name):
+            return {
+                "name": name,
+                "kind": "pipeline",
+                "modality": modality_of(name),
+                "available": False,
+                "note": (f"'{name}' is not in this diffusers build "
+                         "(likely a from-source >0.38 class). Install with: "
+                         "pip install 'git+https://github.com/huggingface/diffusers' "
+                         "— use_diffusers resolves it dynamically once present."),
+            }
+        return None
+    info: Dict[str, Any] = {
+        "name": name,
+        "kind": all_symbols()[name],
+        "modality": modality_of(name),
+    }
+    try:
+        cls = resolve_attr(name)
+        call = getattr(cls, "__call__", None)
+        if call is not None:
+            info["call_params"] = _sig_params(call)
+        fp = getattr(cls, "from_pretrained", None)
+        if fp is not None and getattr(fp, "__doc__", None):
+            info["from_pretrained_doc"] = fp.__doc__[:400]
+        if cls.__doc__:
+            info["doc"] = cls.__doc__[:600]
+    except Exception as e:  # resolution may fail on from-source-only classes
+        info["note"] = f"class not resolvable in this diffusers build: {e}"
+    return info
+
+
+def describe(obj: Any, max_doc: int = 600) -> Dict[str, Any]:
+    info: Dict[str, Any] = {
+        "kind": type(obj).__name__,
+        "name": getattr(obj, "__name__", str(obj)[:80]),
+    }
+    if inspect.isclass(obj):
+        info["methods"] = [
+            m for m in dir(obj)
+            if not m.startswith("_") and callable(getattr(obj, m, None))
+        ][:40]
+        for ctor in ("from_pretrained", "__call__", "__init__"):
+            fn = getattr(obj, ctor, None)
+            if fn is not None:
+                try:
+                    info[f"{ctor}_params"] = _sig_params(fn)
+                    if fn.__doc__:
+                        info[f"{ctor}_doc"] = fn.__doc__[:max_doc]
+                except (ValueError, TypeError):
+                    continue
+    elif callable(obj):
+        try:
+            info["params"] = _sig_params(obj)
+        except (ValueError, TypeError):
+            pass
+        if obj.__doc__:
+            info["doc"] = obj.__doc__[:max_doc]
+    elif inspect.ismodule(obj):
+        info["public"] = [n for n in dir(obj) if not n.startswith("_")][:50]
+    else:
+        info["value"] = str(obj)[:200]
+    return info
+
+
+def _sig_params(fn: Any) -> Dict[str, Dict[str, Any]]:
+    sig = inspect.signature(fn)
+    return {
+        name: {
+            "default": ("REQUIRED" if p.default is inspect.Parameter.empty
+                        else str(p.default)),
+            "annotation": (None if p.annotation is inspect.Parameter.empty
+                           else str(p.annotation)),
+        }
+        for name, p in sig.parameters.items()
+        if name not in ("self", "args", "kwargs")
+    }

strands_diffusers/core/viz.py

@@ -0,0 +1,256 @@
+"""Visualize robot ACTION chunks — turn raw action tensors into something you can SEE.
+
+A world-foundation model (Cosmos) emits an action chunk of shape
+`[num_chunks, T, action_dim]` in normalized action space. Numbers alone are
+opaque, so this renders them three ways:
+
+1. time-series   — every action dimension plotted over the chunk's timesteps
+                   (joint / delta curves), with the gripper channel highlighted.
+2. trajectory    — if the first 3 dims look like an end-effector position/delta,
+                   the cumulative 3D path of the gripper through space.
+3. animation     — an mp4/gif that sweeps a playhead across the time-series so you
+                   can watch the action unfold, optionally side-by-side with the
+                   generated world video frames.
+
+All outputs are written to ARTIFACT_DIR and returned as paths.
+
+Design notes:
+- We DON'T hardcode an embodiment. action_dim varies (7-DoF arm, 10-DoF, etc.).
+  We label dims generically and treat the LAST dim as the gripper by convention
+  (override with gripper_index=None to disable). The first 3 dims are *optionally*
+  interpreted as an end-effector position for the 3D path (interpret_xyz).
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Any, List, Optional
+
+from strands_diffusers.core.io import ARTIFACT_DIR
+
+
+def _as_chunks(action: Any):
+    """Normalize an action payload to a numpy array [num_chunks, T, action_dim]."""
+    import numpy as np
+
+    # Accept: serialized dict {"data": [...]}, raw nested list, tensor, ndarray,
+    # OR a JSON string (LLM tool-calls serialize list inputs to a string).
+    if isinstance(action, str):
+        import json
+        try:
+            action = json.loads(action)
+        except (ValueError, TypeError):
+            raise ValueError(
+                "action string is not valid JSON; pass a nested list "
+                "[num_chunks, T, action_dim], a .json path, or cached:key")
+    if isinstance(action, dict):
+        action = action.get("data", action)
+    a = np.asarray(action, dtype=float)
+    if a.ndim == 1:          # [action_dim] → one chunk, one step
+        a = a[None, None, :]
+    elif a.ndim == 2:        # [T, action_dim] → one chunk
+        a = a[None, :, :]
+    return a
+
+
+def visualize_action(
+    action: Any,
+    save_prefix: str = "action",
+    interpret_xyz: bool = True,
+    gripper_index: Optional[int] = -1,
+    cumulative_xyz: bool = True,
+    world_video: Optional[str] = None,
+    fps: int = 5,
+    dim_labels: Optional[List[str]] = None,
+) -> dict:
+    """Render an action chunk to plots + an animation. Returns artifact paths.
+
+    Args:
+        action: action payload (serialized dict, nested list, ndarray, or tensor),
+                shape [num_chunks, T, action_dim] (lower ranks are promoted).
+        save_prefix: filename prefix for artifacts.
+        interpret_xyz: if action_dim >= 3, draw a 3D path from dims 0-2.
+        gripper_index: dim treated as gripper (highlighted); None to disable.
+        cumulative_xyz: treat xyz dims as deltas and integrate into a path; if
+                        False, plot them as absolute positions.
+        world_video: optional path to the generated world .mp4 to show beside the
+                     action animation (frame-synced as best as frame counts allow).
+        fps: animation frames-per-second.
+        dim_labels: optional human labels per dimension.
+
+    Returns:
+        {"artifacts": [...], "summary": {...}}
+    """
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    chunks = _as_chunks(action)
+    num_chunks, T, dim = chunks.shape
+    # Flatten chunks end-to-end into a single timeline for plotting.
+    flat = chunks.reshape(num_chunks * T, dim)
+    steps = np.arange(flat.shape[0])
+
+    labels = dim_labels or [f"dim{i}" for i in range(dim)]
+    g_idx = (gripper_index % dim) if (gripper_index is not None) else None
+
+    artifacts: List[str] = []
+    from strands_diffusers.core.io import _stamp
+    ts = _stamp()
+
+    # ── 1. time-series ──────────────────────────────────────────────
+    fig, ax = plt.subplots(figsize=(10, 5))
+    for i in range(dim):
+        is_grip = (i == g_idx)
+        ax.plot(steps, flat[:, i],
+                label=("gripper" if is_grip else labels[i]),
+                lw=2.4 if is_grip else 1.3,
+                color="black" if is_grip else None,
+                ls="--" if is_grip else "-",
+                alpha=1.0 if is_grip else 0.85)
+    for c in range(1, num_chunks):
+        ax.axvline(c * T - 0.5, color="gray", ls=":", alpha=0.4)
+    ax.set_xlabel("timestep")
+    ax.set_ylabel("normalized action value")
+    ax.set_title(f"Action chunk — {num_chunks}×{T} steps × {dim} dims")
+    ax.legend(loc="upper right", fontsize=8, ncol=2)
+    ax.grid(alpha=0.25)
+    p_ts = ARTIFACT_DIR / f"{save_prefix}_timeseries_{ts}.png"
+    fig.tight_layout()
+    fig.savefig(p_ts, dpi=110)
+    plt.close(fig)
+    artifacts.append(str(p_ts))
+
+    # ── 2. 3D end-effector path ─────────────────────────────────────
+    path_xyz = None
+    if interpret_xyz and dim >= 3:
+        xyz = flat[:, :3].copy()
+        if cumulative_xyz:
+            xyz = np.cumsum(xyz, axis=0)  # treat as deltas → integrated trajectory
+        path_xyz = xyz
+        fig = plt.figure(figsize=(6, 5.5))
+        ax = fig.add_subplot(111, projection="3d")
+        ax.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], "-o", ms=3, lw=1.5)
+        ax.scatter(*xyz[0], c="green", s=60, label="start")
+        ax.scatter(*xyz[-1], c="red", s=60, label="end")
+        # mark gripper close/open events along the path if we have a gripper channel
+        if g_idx is not None:
+            g = flat[:, g_idx]
+            thr = (g.max() + g.min()) / 2.0
+            closed = g > thr
+            if closed.any():
+                ax.scatter(xyz[closed, 0], xyz[closed, 1], xyz[closed, 2],
+                           c="orange", s=18, alpha=0.7, label="gripper engaged")
+        ax.set_xlabel("x"); ax.set_ylabel("y"); ax.set_zlabel("z")
+        ax.set_title("End-effector path"
+                     + (" (∫ deltas)" if cumulative_xyz else " (absolute)"))
+        ax.legend(fontsize=8)
+        p_traj = ARTIFACT_DIR / f"{save_prefix}_trajectory_{ts}.png"
+        fig.tight_layout()
+        fig.savefig(p_traj, dpi=110)
+        plt.close(fig)
+        artifacts.append(str(p_traj))
+
+    # ── 3. animation (playhead sweep, optional world video beside it) ──
+    anim_path = _animate(flat, labels, g_idx, path_xyz, world_video, fps,
+                         save_prefix, ts)
+    if anim_path:
+        artifacts.append(anim_path)
+
+    summary = {
+        "num_chunks": int(num_chunks),
+        "timesteps_per_chunk": int(T),
+        "action_dim": int(dim),
+        "gripper_index": g_idx,
+        "value_range": [float(flat.min()), float(flat.max())],
+        "has_3d_path": path_xyz is not None,
+    }
+    return {"artifacts": artifacts, "summary": summary}
+
+
+def _animate(flat, labels, g_idx, path_xyz, world_video, fps, prefix, ts):
+    """Build an mp4/gif sweeping a playhead across the action curves."""
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    n = flat.shape[0]
+    dim = flat.shape[1]
+
+    # Optionally load world video frames to show alongside.
+    world_frames = None
+    if world_video:
+        try:
+            import imageio.v3 as iio
+            world_frames = iio.imread(world_video)  # [F,H,W,C]
+        except Exception:
+            world_frames = None
+
+    have_world = world_frames is not None and len(world_frames) > 0
+    have_3d = path_xyz is not None
+
+    ncols = 1 + int(have_world) + int(have_3d)
+    frames_out = []
+
+    for k in range(n):
+        fig = plt.figure(figsize=(5 * ncols, 4.2))
+        col = 1
+
+        # panel A: action curves with playhead
+        axc = fig.add_subplot(1, ncols, col); col += 1
+        for i in range(dim):
+            is_grip = (i == g_idx)
+            axc.plot(flat[:, i], lw=2.2 if is_grip else 1.0,
+                     color="black" if is_grip else None,
+                     ls="--" if is_grip else "-", alpha=0.9 if is_grip else 0.7)
+        axc.axvline(k, color="red", lw=2)
+        axc.set_title("action")
+        axc.set_xlabel("timestep"); axc.grid(alpha=0.2)
+
+        # panel B: world frame (synced by proportion)
+        if have_world:
+            axw = fig.add_subplot(1, ncols, col); col += 1
+            fi = min(int(k / max(n - 1, 1) * (len(world_frames) - 1)),
+                     len(world_frames) - 1)
+            axw.imshow(world_frames[fi])
+            axw.set_title(f"world frame {fi}")
+            axw.axis("off")
+
+        # panel C: 3D path with current point
+        if have_3d:
+            ax3 = fig.add_subplot(1, ncols, col, projection="3d"); col += 1
+            ax3.plot(path_xyz[:, 0], path_xyz[:, 1], path_xyz[:, 2],
+                     "-", lw=1, alpha=0.5)
+            ax3.scatter(*path_xyz[k], c="red", s=50)
+            ax3.set_title("end-effector"); ax3.set_xticks([]); ax3.set_yticks([])
+            ax3.set_zticks([])
+
+        fig.tight_layout()
+        fig.canvas.draw()
+        buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
+        w, h = fig.canvas.get_width_height()
+        frames_out.append(buf.reshape(h, w, 4)[..., :3].copy())
+        plt.close(fig)
+
+    if not frames_out:
+        return None
+
+    out_mp4 = ARTIFACT_DIR / f"{prefix}_animation_{ts}.mp4"
+    try:
+        import imageio.v3 as iio
+        iio.imwrite(str(out_mp4), np.stack(frames_out), fps=fps, codec="libx264")
+        return str(out_mp4)
+    except Exception:
+        pass
+    try:
+        from PIL import Image
+        out_gif = ARTIFACT_DIR / f"{prefix}_animation_{ts}.gif"
+        imgs = [Image.fromarray(f) for f in frames_out]
+        imgs[0].save(str(out_gif), save_all=True, append_images=imgs[1:],
+                     duration=int(1000 / max(fps, 1)), loop=0)
+        return str(out_gif)
+    except Exception:
+        return None

strands_diffusers/tools/__init__.py

@@ -0,0 +1,4 @@
+"""Tools for strands-diffusers."""
+from strands_diffusers.tools.use_diffusers import use_diffusers
+
+__all__ = ["use_diffusers"]