pipelines/mr.py

from __future__ import annotations

import time

from .base import GenerationRequest, GenerationResult
from . import _runner


def generate(req: GenerationRequest) -> GenerationResult:
    """MR (rflow-mr) — image-only across multiple contrasts and anatomies."""
    if req.modality_class is None:
        raise ValueError("MR generation requires modality_class (8=mri, 9=t1, 10=t2, 11=flair, ...)")
    t0 = time.time()
    path = _runner.run_image_only(
        version="rflow-mr",
        output_size=req.output_size,
        spacing=req.spacing,
        modality=req.modality_class,
        seed=req.seed,
        num_inference_steps=req.num_steps,
        cfg_guidance_scale=req.cfg_guidance_scale or 15.0,
    )
    return GenerationResult(
        volume_path=str(path),
        runtime_seconds=time.time() - t0,
        seed=req.seed,
        modality=req.modality_class,
    )