compose/provenance.py

from __future__ import annotations

import hashlib
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

BASE_MODEL = "black-forest-labs/FLUX.2-klein-base-4B"


def load_registry(registry_path: str | Path = "registry/loras.json") -> dict[str, Any]:
    with open(registry_path, "r", encoding="utf-8") as f:
        return json.load(f)


def lora_map(registry: dict[str, Any]) -> dict[str, dict[str, Any]]:
    return {entry["id"]: entry for entry in registry.get("loras", [])}


def expand_ancestry(
    lora_id: str,
    weight: float,
    registry_by_id: dict[str, dict[str, Any]],
    *,
    max_depth: int = 8,
    _depth: int = 0,
    _seen: set[str] | None = None,
) -> dict[str, Any]:
    """Expand a LoRA node into a provenance tree.

    Defaults to 8 levels: effectively complete for the hackathon, but bounded.
    Circular references are marked instead of recursing forever.
    """
    seen = set() if _seen is None else set(_seen)
    if lora_id not in registry_by_id:
        raise KeyError(f"Unknown LoRA id: {lora_id}")

    lora = registry_by_id[lora_id]
    node = {
        "lora_id": lora["id"],
        "weight": round(float(weight), 4),
        "weight_pct": round(float(weight) * 100, 2),
        "creator": lora.get("creator"),
        "cultural_source": lora.get("cultural_source"),
        "hf_repo": lora.get("hf_repo"),
        "checkpoint_step": lora.get("checkpoint_step"),
        "type": lora.get("type"),
        "parent_ids": lora.get("parent_ids", []),
        "status": lora.get("status", "unknown"),
        "children": [],
    }

    if lora_id in seen:
        node["cycle_detected"] = True
        return node
    if _depth >= max_depth:
        node["max_depth_reached"] = True
        return node

    seen.add(lora_id)
    parent_ids = lora.get("parent_ids", [])
    if parent_ids:
        parent_weight = float(weight) / len(parent_ids)
        node["children"] = [
            expand_ancestry(
                parent_id,
                parent_weight,
                registry_by_id,
                max_depth=max_depth,
                _depth=_depth + 1,
                _seen=seen,
            )
            for parent_id in parent_ids
            if parent_id in registry_by_id
        ]
    return node


def normalize_weights(weights: list[float]) -> list[float]:
    total = sum(weights)
    if total <= 0:
        raise ValueError("Blend weights must sum to a positive number")
    return [w / total for w in weights]


def file_sha256(path: str | Path | None) -> str | None:
    if path is None:
        return None
    p = Path(path)
    if not p.exists() or not p.is_file():
        return None
    h = hashlib.sha256()
    with p.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return "sha256:" + h.hexdigest()


def build_provenance(
    lora_ids: list[str],
    weights: list[float],
    registry: dict[str, Any],
    prompt: str,
    seed: int,
    output_path: str | None = None,
    *,
    max_depth: int = 8,
) -> dict[str, Any]:
    if len(lora_ids) != len(weights):
        raise ValueError("lora_ids and weights must be the same length")

    weights = normalize_weights(weights)
    by_id = lora_map(registry)
    ancestry = [
        expand_ancestry(lora_id, weight, by_id, max_depth=max_depth)
        for lora_id, weight in zip(lora_ids, weights)
    ]

    return {
        "generation_id": str(uuid.uuid4()),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "prompt": prompt,
        "seed": seed,
        "base_model": BASE_MODEL,
        "ancestry": ancestry,
        "image_hash": file_sha256(output_path),
        "output_path": output_path,
    }


def provenance_sentence(provenance: dict[str, Any]) -> str:
    parts = [
        f"{node['weight_pct']:g}% {node['cultural_source']}"
        for node in provenance.get("ancestry", [])
    ]
    if not parts:
        return "No lineage data available yet."
    return "This image blends " + " and ".join(parts) + "."