adam/model_inspector/ddpm.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from .generic import GenericModelInspector


class DDPMInspector(GenericModelInspector):
    architecture = "DDPM / Diffusers"

    def inspect(self, path: str | Path, *, recorded_architecture: str = "", run_settings: dict[str, Any] | None = None, progress=None, cancelled=None):
        summary = super().inspect(
            path,
            recorded_architecture=recorded_architecture or "ddpm",
            run_settings=run_settings,
            progress=progress,
            cancelled=cancelled,
        )
        if summary.architecture == "Generic / Unknown":
            summary.architecture = self.architecture
            summary.messages.insert(0, "Model recognized as DDPM from ADAM trainer metadata")
        expected = ("unet", "scheduler")
        root = Path(summary.resolved_path)
        existing = {part.name.casefold() for part in (root.iterdir() if root.is_dir() else [])}
        for component in expected:
            if root.is_dir() and component not in existing and not any(component in item.casefold() for item in summary.configs):
                summary.health.append(f"Unusual: expected DDPM component not found: {component}")
        return summary