mage_flow/app.py

"""Gradio app for MageFlow — text-to-image and instruction-based image editing.

    python app.py                       # serve on 0.0.0.0:7860
    python app.py --share --port 7861

Each tab has a model preset dropdown (base / rl / turbo) plus a free-form
"Custom model" box for any Hugging Face repo id or local path. Models load
lazily on first use and are cached. Notes:
  - By default the presets point at the `microsoft/Mage-Flow*` Hugging Face
    repos (downloaded + cached on first use). Set ``MAGEFLOW_HF_DIR`` to load
    local checkpoint dirs instead.
  - Turbo checkpoints are few-step: use steps=4, cfg=1.
"""
from __future__ import annotations

import argparse
import os

import gradio as gr
from PIL import Image

from mage_flow.pipeline import MageFlowPipeline

# Default to Hugging Face repo ids; if MAGEFLOW_HF_DIR is set, use local
# checkpoint dirs under it instead (local dir names match the HF repo basename).
HF_DIR = os.environ.get("MAGEFLOW_HF_DIR")


def _repo(hf_id: str, local_name: str) -> str:
    return f"{HF_DIR}/{local_name}" if HF_DIR else hf_id


T2I_MODELS = {
    "base":  _repo("microsoft/Mage-Flow-Base",  "Mage-Flow-Base"),
    "rl":    _repo("microsoft/Mage-Flow",       "Mage-Flow"),
    "turbo": _repo("microsoft/Mage-Flow-Turbo", "Mage-Flow-Turbo"),
}
EDIT_MODELS = {
    "base":  _repo("microsoft/Mage-Flow-Edit-Base",  "Mage-Flow-Edit-Base"),
    "rl":    _repo("microsoft/Mage-Flow-Edit",       "Mage-Flow-Edit"),
    "turbo": _repo("microsoft/Mage-Flow-Edit-Turbo", "Mage-Flow-Edit-Turbo"),
}

DEVICE = "cuda"
_CACHE: dict[str, MageFlowPipeline] = {}


def _get_pipe(repo: str) -> MageFlowPipeline:
    """Load (and cache) a pipeline from a local dir OR a Hugging Face repo id.

    ``MageFlowPipeline.from_pretrained`` resolves a repo id via
    ``snapshot_download`` automatically, so both are accepted here.
    """
    repo = (repo or "").strip()
    if not repo:
        raise gr.Error("No model specified.")
    if repo not in _CACHE:
        try:
            _CACHE[repo] = MageFlowPipeline.from_pretrained(repo, device=DEVICE)
        except Exception as e:  # noqa: BLE001
            raise gr.Error(f"Failed to load model '{repo}': {type(e).__name__}: {e}")
    return _CACHE[repo]


def _resolve(preset_map, model_key, custom_model):
    """Custom repo id / path (if given) overrides the preset dropdown."""
    return (custom_model or "").strip() or preset_map[model_key]


def run_t2i(model_key, custom_model, prompt, neg_prompt, steps, cfg, height, width, seed,
            progress=gr.Progress(track_tqdm=False)):
    if not (prompt or "").strip():
        raise gr.Error("Prompt is empty.")
    repo = _resolve(T2I_MODELS, model_key, custom_model)
    progress(0.1, desc=f"loading {repo} …")
    pipe = _get_pipe(repo)
    progress(0.4, desc="generating …")
    img = pipe.generate(
        [prompt], neg_prompts=[neg_prompt or " "], seeds=[int(seed)],
        steps=int(steps), cfg=float(cfg),
        heights=[int(height)], widths=[int(width)],
    )[0]
    return img


def run_edit(model_key, custom_model, prompt, neg_prompt, ref_img, extra_files, steps, cfg, max_size, seed,
             progress=gr.Progress(track_tqdm=False)):
    if not (prompt or "").strip():
        raise gr.Error("Edit instruction is empty.")
    refs = []
    if ref_img is not None:
        refs.append(ref_img if isinstance(ref_img, Image.Image) else Image.open(ref_img))
    for f in (extra_files or []):
        refs.append(Image.open(f).convert("RGB"))
    if not refs:
        raise gr.Error("Upload at least one reference image.")
    refs = [r.convert("RGB") for r in refs]
    repo = _resolve(EDIT_MODELS, model_key, custom_model)
    progress(0.1, desc=f"loading {repo} …")
    pipe = _get_pipe(repo)
    progress(0.4, desc="editing …")
    out = pipe.edit(
        [prompt], [refs], neg_prompts=[neg_prompt or " "], seeds=[int(seed)],
        steps=int(steps), cfg=float(cfg),
        max_size=int(max_size) if max_size else None,
    )[0]
    return out


_NOTE = (
    "Pick a **preset** (base / rl / turbo) or type a **custom model** — any "
    "Hugging Face repo id (e.g. `microsoft/Mage-Flow-Turbo`) or local path; it "
    "is downloaded and cached on first use. **Turbo** models are few-step: set "
    "**steps=4, cfg=1**."
)

_CUSTOM_PH_T2I = "microsoft/Mage-Flow  (repo id or local path — overrides preset)"
_CUSTOM_PH_EDIT = "microsoft/Mage-Flow-Edit  (repo id or local path — overrides preset)"


def build_ui():
    with gr.Blocks(title="MageFlow") as demo:
        gr.Markdown("# MageFlow\nText-to-image generation and instruction-based image editing.")
        gr.Markdown(_NOTE)

        with gr.Tab("Text → Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    t_model = gr.Dropdown(list(T2I_MODELS), value="base", label="Model preset")
                    t_custom = gr.Textbox(label="Custom model (optional)", placeholder=_CUSTOM_PH_T2I, lines=1)
                    t_prompt = gr.Textbox(label="Prompt", lines=3,
                                          value="A close-up portrait of an elderly African man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic.")
                    t_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
                    with gr.Row():
                        t_steps = gr.Slider(1, 50, value=30, step=1, label="Steps")
                        t_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG")
                    with gr.Row():
                        t_h = gr.Slider(256, 1536, value=1024, step=16, label="Height")
                        t_w = gr.Slider(256, 1536, value=1024, step=16, label="Width")
                    t_seed = gr.Number(value=42, precision=0, label="Seed")
                    t_btn = gr.Button("Generate", variant="primary")
                with gr.Column(scale=1):
                    t_out = gr.Image(type="pil", label="Output", height=560)
            # Clear the previous output first so the stale image isn't shown as
            # the result while the new one is still transferring (esp. over a
            # gradio share tunnel, where the image download can lag a few seconds).
            t_btn.click(lambda: None, None, t_out).then(
                        run_t2i,
                        [t_model, t_custom, t_prompt, t_neg, t_steps, t_cfg, t_h, t_w, t_seed],
                        t_out)

        with gr.Tab("Image Edit"):
            with gr.Row():
                with gr.Column(scale=1):
                    e_model = gr.Dropdown(list(EDIT_MODELS), value="base", label="Model preset")
                    e_custom = gr.Textbox(label="Custom model (optional)", placeholder=_CUSTOM_PH_EDIT, lines=1)
                    e_prompt = gr.Textbox(label="Edit instruction", lines=2,
                                          value="change the background to a city street")
                    e_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
                    e_ref = gr.Image(type="pil", label="Reference image", height=280,
                                     value=os.path.join(os.path.dirname(__file__), "assets", "dog.jpg"))
                    e_extra = gr.File(file_count="multiple", type="filepath",
                                      label="Extra references (optional, multi-image edit)")
                    with gr.Row():
                        e_steps = gr.Slider(1, 50, value=30, step=1, label="Steps")
                        e_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG")
                    e_max = gr.Slider(0, 1536, value=1024, step=16,
                                      label="Max output side (0 = keep source size)")
                    e_seed = gr.Number(value=42, precision=0, label="Seed")
                    e_btn = gr.Button("Edit", variant="primary")
                with gr.Column(scale=1):
                    e_out = gr.Image(type="pil", label="Output", height=560)
            e_btn.click(lambda: None, None, e_out).then(
                        run_edit,
                        [e_model, e_custom, e_prompt, e_neg, e_ref, e_extra, e_steps, e_cfg, e_max, e_seed],
                        e_out)
    return demo


def main():
    global DEVICE
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=7860)
    ap.add_argument("--share", action="store_true")
    ap.add_argument("--preload", default=None,
                    help="comma-separated repo ids / paths to load at startup (else lazy)")
    args = ap.parse_args()
    DEVICE = args.device
    if args.preload:
        for repo in args.preload.split(","):
            _get_pipe(repo.strip())
    build_ui().queue().launch(server_name=args.host, server_port=args.port,
                              share=args.share, theme=gr.themes.Soft())


if __name__ == "__main__":
    main()