"""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())