Add vendor/mage_flow/app.py

vendor/mage_flow/app.py

@@ -0,0 +1,199 @@
+"""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()