compose/merge.py

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

from compose.provenance import build_provenance, load_registry, lora_map, normalize_weights

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

_PIPE = None
_LOADED_ADAPTERS: set[str] = set()


def _device():
    import torch

    if torch.cuda.is_available():
        return "cuda", torch.bfloat16
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return "mps", torch.float16
    return "cpu", torch.float32


def _load_pipe():
    global _PIPE
    if _PIPE is not None:
        return _PIPE

    import torch
    from diffusers import Flux2KleinPipeline

    device, dtype = _device()
    token = os.environ.get("HF_TOKEN")
    pipe = Flux2KleinPipeline.from_pretrained(
        BASE_INFERENCE_MODEL,
        torch_dtype=dtype,
        token=token,
    )
    pipe = pipe.to(device)
    _PIPE = pipe
    return pipe


def _load_adapters(pipe, lora_ids: list[str], registry_by_id: dict[str, dict[str, Any]]) -> None:
    for lora_id in lora_ids:
        if lora_id in _LOADED_ADAPTERS:
            continue
        repo = registry_by_id[lora_id]["hf_repo"]
        pipe.load_lora_weights(repo, adapter_name=lora_id, token=os.environ.get("HF_TOKEN"))
        _LOADED_ADAPTERS.add(lora_id)


def compose(
    lora_ids: list[str],
    weights: list[float],
    prompt: str,
    registry_path: str = "registry/loras.json",
    seed: int = 42,
    output_dir: str = "outputs",
    mode: str = "live",
    width: int = 768,
    height: int = 768,
    num_inference_steps: int = 12,
    guidance_scale: float = 2.0,
) -> dict[str, Any]:
    """Generate with FLUX.2 klein + LoRA adapters and return image/provenance.

    `mode="live"` is the real path. `mode="mock"` was removed from the UI, but
    kept as an explicit developer escape hatch only through NotImplementedError
    so accidental demo mocks fail loudly.
    """
    if mode != "live":
        raise NotImplementedError("Mock generation is disabled. Use mode='live'.")

    import torch

    if len(lora_ids) != len(weights):
        raise ValueError("lora_ids and weights must be the same length")
    if not prompt.strip():
        raise ValueError("Prompt is required")

    registry = load_registry(registry_path)
    by_id = lora_map(registry)
    missing = [lid for lid in lora_ids if lid not in by_id]
    if missing:
        raise KeyError(f"Unknown LoRA ids: {missing}")

    weights = normalize_weights(weights)
    pipe = _load_pipe()
    _load_adapters(pipe, lora_ids, by_id)
    pipe.set_adapters(lora_ids, adapter_weights=weights)

    triggers = " ".join(by_id[lid]["trigger"] for lid in lora_ids)
    full_prompt = f"{triggers}. {prompt.strip()}"

    device, _ = _device()
    generator = torch.Generator(device=device if device == "cuda" else "cpu").manual_seed(int(seed))
    result = pipe(
        prompt=full_prompt,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        height=int(height),
        width=int(width),
        generator=generator,
    )
    image = result.images[0]

    Path(output_dir).mkdir(parents=True, exist_ok=True)
    generation_id = f"{seed}_{'_'.join(lora_ids)}"
    output_path = Path(output_dir) / f"{generation_id}.png"
    image.save(output_path)

    provenance = build_provenance(lora_ids, weights, registry, prompt, seed, str(output_path))
    provenance["full_prompt"] = full_prompt
    provenance["inference_model"] = BASE_INFERENCE_MODEL
    provenance["mode"] = "live"
    return {"image": image, "provenance": provenance, "output_path": str(output_path)}