Hggd/mage-flow
0
1"""Mage-Flow: Efficient Native-Resolution Foundation Model for Image Generation and Editing.2 3Gradio Space demo with a single unified interface: image presence selects4editing vs. generation, while the model control selects fast vs. quality.5"""6import gc7import os8import threading9 10# Use flash_attention_2 for the HF text encoder (flash_attn is installed via wheel)11os.environ.setdefault("VF_HF_ATTN_IMPL", "flash_attention_2")12 13import spaces # MUST be first (after env setup)14import torch15import gradio as gr16from PIL import Image17 18from mage_flow.pipeline import MageFlowPipeline19 20MODEL_VARIANTS = {21 "turbo": {22 "t2i": "microsoft/Mage-Flow-Turbo", "edit": "microsoft/Mage-Flow-Edit-Turbo",23 "t2i_steps": 4, "edit_steps": 4, "cfg": 1.0,24 },25 "quality": {26 "t2i": "microsoft/Mage-Flow", "edit": "microsoft/Mage-Flow-Edit",27 "t2i_steps": 20, "edit_steps": 30, "cfg": 5.0,28 },29}30 31_pipe_slots = {32 "t2i": {"variant": "turbo", "pipe": MageFlowPipeline.from_pretrained(MODEL_VARIANTS["turbo"]["t2i"], device="cuda")},33 "edit": {"variant": "turbo", "pipe": MageFlowPipeline.from_pretrained(MODEL_VARIANTS["turbo"]["edit"], device="cuda")},34}35_pipe_lock = threading.Lock()36 37 38def _get_pipe(task: str, variant: str):39 """Keep one loaded variant per task, matching the original two-pipeline footprint."""40 with _pipe_lock:41 slot = _pipe_slots.get(task)42 if slot and slot["variant"] == variant:43 return slot["pipe"]44 if slot:45 del _pipe_slots[task]46 del slot47 gc.collect()48 torch.cuda.empty_cache()49 pipe = MageFlowPipeline.from_pretrained(MODEL_VARIANTS[variant][task], device="cuda")50 _pipe_slots[task] = {"variant": variant, "pipe": pipe}51 return pipe52 53 54def _recommended(variant: str, image):55 spec = MODEL_VARIANTS[variant]56 return (spec["edit_steps"] if image is not None else spec["t2i_steps"], spec["cfg"])57 58 59@spaces.GPU(duration=120)60def generate(61 prompt: str,62 image=None,63 negative_prompt: str = " ",64 steps: int = 4,65 cfg: float = 1.0,66 height: int = 1024,67 width: int = 1024,68 max_size: int = 1024,69 seed: int = 42,70 model_variant: str = "turbo",71 progress=gr.Progress(track_tqdm=True),72):73 """Generate or edit an image with Mage-Flow.74 75 If ``image`` is provided, route to the selected edit model; otherwise route76 to the selected text-to-image model.77 78 Args:79 prompt: Text description (generation) or edit instruction (editing).80 image: Optional reference image. When given, routes to the edit model.81 negative_prompt: What to avoid in the result.82 steps: Number of denoising steps (Turbo uses 4).83 cfg: Classifier-free guidance scale (Turbo uses 1.0).84 height: Output image height for text-to-image (multiple of 16).85 width: Output image width for text-to-image (multiple of 16).86 max_size: Longest side of edited output (0 = keep source resolution).87 seed: Random seed for reproducibility.88 """89 if not (prompt or "").strip():90 raise gr.Error("Prompt is empty.")91 92 if image is not None:93 # Route to the edit model when an image is provided.94 pipe_edit = _get_pipe("edit", model_variant)95 if isinstance(image, str):96 image = Image.open(image)97 refs = [image.convert("RGB")]98 99 # Content-safety gate: blocked requests return a blank image.100 verdict = pipe_edit.model.txt_enc.screen_edit(prompt, refs)101 if verdict.violates:102 w, h = refs[0].size103 return Image.new("RGB", (w, h), (255, 255, 255))104 105 out = pipe_edit.edit(106 [prompt],107 [refs],108 neg_prompts=[negative_prompt or " "],109 seeds=[int(seed)],110 steps=int(steps),111 cfg=float(cfg),112 max_size=int(max_size) if max_size else None,113 )[0]114 return out115 116 # No image: route to the text-to-image model.117 # Content-safety gate: blocked requests return a blank image.118 pipe_t2i = _get_pipe("t2i", model_variant)119 verdict = pipe_t2i.model.txt_enc.screen_text(prompt)120 if verdict.violates:121 return Image.new("RGB", (int(width), int(height)), (255, 255, 255))122 123 img = pipe_t2i.generate(124 [prompt],125 neg_prompts=[negative_prompt or " "],126 seeds=[int(seed)],127 steps=int(steps),128 cfg=float(cfg),129 heights=[int(height)],130 widths=[int(width)],131 )[0]132 return img133 134 135ASSETS_DIR = os.path.join(os.path.dirname(__file__), "mage_flow", "assets")136 137CSS = """138#col-container { margin: 0 auto; max-width: 1100px; }139.dark .gradio-container { color: var(--body-text-color); }140"""141 142with gr.Blocks(css=CSS) as demo:143 with gr.Column(elem_id="col-container"):144 gr.Markdown(145 "# Mage-Flow\n"146 "Efficient Native-Resolution Foundation Model for Image Generation and Editing. "147 "Enter a prompt to generate an image, or upload an image to edit it.\n\n"148 "Models: [Mage-Flow](https://huggingface.co/microsoft/Mage-Flow), "149 "[Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo), "150 "[Mage-Flow-Edit](https://huggingface.co/microsoft/Mage-Flow-Edit), "151 "[Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo) | "152 "[Paper](https://huggingface.co/papers/2607.19064) | "153 "[GitHub](https://github.com/microsoft/Mage)"154 )155 156 with gr.Row():157 with gr.Column(scale=1):158 with gr.Row():159 prompt = gr.Textbox(160 label="Prompt",161 show_label=False,162 max_lines=3,163 placeholder="Describe an image to generate, or an edit instruction for an uploaded image",164 container=False,165 scale=4,166 )167 run_btn = gr.Button("Run", variant="primary", scale=1)168 169 model_variant = gr.Radio(170 [("Mage-Flow-Turbo · Fast", "turbo"), ("Mage-Flow · Quality", "quality")],171 value="turbo", label="Model",172 )173 174 with gr.Accordion("Input image (optional — enables editing)", open=True):175 image = gr.Image(176 type="pil",177 label="Input image",178 show_label=False,179 height=300,180 )181 182 with gr.Accordion("Advanced Settings", open=False):183 negative_prompt = gr.Textbox(label="Negative prompt", value=" ", lines=1)184 with gr.Row():185 steps = gr.Slider(1, 50, value=4, step=1, label="Steps")186 cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")187 with gr.Row():188 height = gr.Slider(256, 1536, value=1024, step=16, label="Height (text→image)")189 width = gr.Slider(256, 1536, value=1024, step=16, label="Width (text→image)")190 max_size = gr.Slider(191 0, 1536, value=1024, step=16,192 label="Max output side for editing (0 = keep source size)",193 )194 seed = gr.Number(value=42, precision=0, label="Seed")195 196 with gr.Column(scale=1):197 result = gr.Image(type="pil", label="Output", height=560)198 199 gr.Markdown("### Text → Image examples")200 gr.Examples(201 examples=[202 ["A close-up portrait of an elderly Hausa man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic."],203 ["A serene mountain landscape at sunset, with snow-capped peaks reflecting golden light, photorealistic."],204 ["A cute robot playing a guitar in a neon-lit cyberpunk city, digital art style."],205 ],206 inputs=[prompt],207 outputs=result,208 fn=generate,209 cache_examples=True,210 cache_mode="lazy",211 )212 213 gr.Markdown("### Image editing examples")214 gr.Examples(215 examples=[216 ["change the background to a city street", os.path.join(ASSETS_DIR, "dog.jpg")],217 ["make it look like a painting", os.path.join(ASSETS_DIR, "cuisine.jpg")],218 ["add a hat to the person", os.path.join(ASSETS_DIR, "portrait.jpg")],219 ],220 inputs=[prompt, image],221 outputs=result,222 fn=generate,223 cache_examples=True,224 cache_mode="lazy",225 )226 227 model_variant.change(_recommended, [model_variant, image], [steps, cfg], api_name=False)228 image.change(_recommended, [model_variant, image], [steps, cfg], api_name=False)229 230 inputs = [prompt, image, negative_prompt, steps, cfg, height, width, max_size, seed, model_variant]231 run_btn.click(lambda: None, None, result).then(232 generate, inputs, result, api_name="generate",233 )234 prompt.submit(lambda: None, None, result).then(235 generate, inputs, result, api_name=False,236 )237 238if __name__ == "__main__":239 demo.launch(theme=gr.themes.Citrus(), mcp_server=True, show_error=True)240