anima_mlx/generate.py

"""Command line prompt-to-image MVP for Anima MLX."""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from anima_mlx.models.pipeline import AnimaTinyPipeline
from anima_mlx.runtime.scheduler import FlowSchedulerConfig
from anima_mlx.runtime.tokenizer import AnimaTokenizer
from anima_mlx.runtime.workflow import DEFAULT_WORKFLOW_SPEC
from anima_mlx.utils.compare import to_numpy


ROOT = Path(__file__).resolve().parents[1]


@dataclass(frozen=True)
class GenerationSummary:
    output: str
    image_shape: tuple[int, ...]
    decoded_shape: tuple[int, ...]
    latent_shape: tuple[int, ...]
    steps: int
    seed: int
    cfg: float
    sampler: str
    dit_load_mode: str
    dit_eval_interval: int
    peak_memory_gb: float | None
    timings: dict[str, float]


def default_text_encoder_path() -> Path:
    converted = ROOT / "split_files/text_encoders/qwen_3_06b_base-mlx.safetensors"
    if converted.exists():
        return converted
    converted = default_mlx_weights_dir() / "text_encoder.safetensors"
    if converted.exists():
        return converted
    return ROOT / "split_files/text_encoders/qwen_3_06b_base-mlx.safetensors"


def default_diffusion_path() -> Path:
    converted = ROOT / "split_files/diffusion_models/anima-base-v1.0-mlx.safetensors"
    if converted.exists():
        return converted
    converted = default_mlx_weights_dir() / "diffusion_core.safetensors"
    if converted.exists():
        return converted
    converted = default_mlx_weights_dir() / "diffusion.safetensors"
    if converted.exists():
        return converted
    return ROOT / "split_files/diffusion_models/anima-base-v1.0-mlx.safetensors"


def default_vae_path() -> Path:
    converted = ROOT / "split_files/vae/qwen_image_vae-mlx.safetensors"
    if converted.exists():
        return converted
    converted = default_mlx_weights_dir() / "vae.safetensors"
    if converted.exists():
        return converted
    return ROOT / "split_files/vae/qwen_image_vae-mlx.safetensors"


def default_mlx_weights_dir() -> Path:
    return ROOT / "mlx_weights"


def default_comfy_root() -> Path:
    return ROOT


def postprocess_decoded_to_uint8(decoded: Any) -> Any:
    import numpy as np

    array = to_numpy(decoded)
    if array.ndim == 5:
        image = array[0, :, 0, :, :].transpose(1, 2, 0)
    elif array.ndim == 4:
        image = array[0].transpose(1, 2, 0)
    elif array.ndim == 3:
        image = array.transpose(1, 2, 0) if array.shape[0] == 3 else array
    else:
        raise ValueError(f"decoded tensor must have 3, 4, or 5 dimensions, got {array.shape}")
    image = np.clip((image + 1.0) / 2.0, 0.0, 1.0)
    return (image * 255.0 + 0.5).astype(np.uint8)


def save_png(image: Any, output: str | Path) -> Path:
    from PIL import Image

    output_path = Path(output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(image).save(output_path)
    return output_path


def generate_image(args: argparse.Namespace) -> GenerationSummary:
    import mlx.core as mx

    timings: dict[str, float] = {}
    set_mlx_limits(args.memory_limit_gb, args.cache_limit_gb)
    _require_file(args.text_encoder, "text encoder weights")
    _require_file(args.diffusion_model, "diffusion weights")
    _require_file(args.vae, "VAE weights")
    _require_dir(args.comfy_root, "ComfyUI tokenizer root")
    total_start = time.perf_counter()
    stage_start = total_start
    tokenizer = AnimaTokenizer.from_comfy_root(args.comfy_root)
    tokenized = tokenizer.tokenize_pair(args.prompt, args.negative_prompt)
    timings["tokenize"] = time.perf_counter() - stage_start
    stage_start = time.perf_counter()
    pipeline = AnimaTinyPipeline.from_safetensors(
        args.text_encoder,
        args.diffusion_model,
        vae_path=args.vae,
        dtype=args.dtype,
        dit_load_mode=args.dit_load_mode,
        dit_eval_interval=args.dit_eval_interval,
    )
    timings["load_pipeline"] = time.perf_counter() - stage_start
    scheduler_config = FlowSchedulerConfig(shift=args.flow_shift, multiplier=args.flow_multiplier)
    stage_start = time.perf_counter()
    result = pipeline.generate_from_tokens(
        tokenized,
        height=args.height,
        width=args.width,
        frames=args.frames,
        batch_size=args.batch_size,
        seed=args.seed,
        steps=args.steps,
        cfg=args.cfg,
        scheduler_config=scheduler_config,
        sampler_name=args.sampler,
        dtype=args.dtype,
        vae_decode_mode=args.vae_decode_mode,
        vae_tile_size=args.vae_tile_size,
        vae_overlap=args.vae_overlap,
    )
    timings["generate_tensor"] = time.perf_counter() - stage_start
    if result.timings is not None:
        timings.update({f"pipeline_{key}": value for key, value in result.timings.items()})
    stage_start = time.perf_counter()
    mx.eval(result.decoded)
    timings["eval_decoded"] = time.perf_counter() - stage_start
    peak_memory_gb = get_peak_memory_gb()
    stage_start = time.perf_counter()
    image = postprocess_decoded_to_uint8(result.decoded)
    timings["postprocess"] = time.perf_counter() - stage_start
    stage_start = time.perf_counter()
    output = save_png(image, args.output)
    timings["save_png"] = time.perf_counter() - stage_start
    mx.clear_cache()
    timings["total"] = time.perf_counter() - total_start
    return GenerationSummary(
        output=output.as_posix(),
        image_shape=tuple(image.shape),
        decoded_shape=tuple(result.decoded.shape),
        latent_shape=tuple(result.latent.shape),
        steps=args.steps,
        seed=args.seed,
        cfg=args.cfg,
        sampler=args.sampler,
        dit_load_mode=args.dit_load_mode,
        dit_eval_interval=args.dit_eval_interval,
        peak_memory_gb=peak_memory_gb,
        timings=timings,
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Generate a single image with the Anima MLX MVP pipeline.")
    parser.add_argument("--prompt", default=DEFAULT_WORKFLOW_SPEC.positive_prompt)
    parser.add_argument("--negative-prompt", default=DEFAULT_WORKFLOW_SPEC.negative_prompt)
    parser.add_argument("--steps", type=int, default=1)
    parser.add_argument("--seed", type=int, default=DEFAULT_WORKFLOW_SPEC.seed)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--frames", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--cfg", type=float, default=DEFAULT_WORKFLOW_SPEC.cfg)
    parser.add_argument("--sampler", choices=("euler", "er_sde"), default="euler")
    parser.add_argument("--scheduler", choices=("simple",), default=DEFAULT_WORKFLOW_SPEC.scheduler)
    parser.add_argument("--flow-shift", type=float, default=DEFAULT_WORKFLOW_SPEC.flow_shift)
    parser.add_argument("--flow-multiplier", type=float, default=DEFAULT_WORKFLOW_SPEC.flow_multiplier)
    parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16")
    parser.add_argument("--dit-load-mode", choices=("lazy", "preload"), default="lazy")
    parser.add_argument("--dit-eval-interval", type=int, default=0)
    parser.add_argument("--vae-decode-mode", choices=("auto", "full", "tiled"), default="full")
    parser.add_argument("--vae-tile-size", type=int, default=64)
    parser.add_argument("--vae-overlap", type=int, default=16)
    parser.add_argument("--memory-limit-gb", type=float, default=12.0)
    parser.add_argument("--cache-limit-gb", type=float, default=1.0)
    parser.add_argument("--text-encoder", type=Path, default=default_text_encoder_path())
    parser.add_argument("--diffusion-model", type=Path, default=default_diffusion_path())
    parser.add_argument("--vae", type=Path, default=default_vae_path())
    parser.add_argument("--comfy-root", type=Path, default=default_comfy_root())
    parser.add_argument("--output", type=Path, default=ROOT / "arona_mlx.png")
    return parser


def _require_file(path: Path, label: str) -> None:
    if not path.exists():
        raise FileNotFoundError(f"{label} not found: {path}")


def _require_dir(path: Path, label: str) -> None:
    if not path.exists() or not path.is_dir():
        raise FileNotFoundError(f"{label} not found: {path}")


def set_mlx_limits(memory_limit_gb: float, cache_limit_gb: float) -> None:
    import mlx.core as mx

    mx.set_memory_limit(int(memory_limit_gb * 1024**3))
    mx.set_cache_limit(int(cache_limit_gb * 1024**3))
    mx.reset_peak_memory()


def get_peak_memory_gb() -> float | None:
    try:
        import mlx.core as mx

        return mx.get_peak_memory() / 1024**3
    except Exception:
        return None


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        summary = generate_image(args)
    except ModuleNotFoundError as exc:
        print(f"Missing optional dependency: {exc.name}. Install the project with the golden/mlx extras.", file=sys.stderr)
        return 2
    except FileNotFoundError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    print(json.dumps(asdict(summary), ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())