Allow splitting flow to two GPUs

mage_flow/README.md

@@ -335,6 +335,7 @@ mage-flow-edit --prompt "Replace the background with a field of sunflowers" "ble
 
 ```bash
 mage-flow-app                     # serve on http://0.0.0.0:7860  (or: python -m mage_flow.app)
+mage-flow-app --device cuda:0 --text-device cuda:1 # split to two GPUs
 ```
 
 A web UI with **Text → Image** and **Image Edit** tabs; models load lazily on first use and are cached. Presets default to the `microsoft/Mage-Flow*` **Hugging Face repos** (downloaded + cached on first use); set `MAGEFLOW_HF_DIR` to load local checkpoint dirs instead.
@@ -346,6 +347,7 @@ A web UI with **Text → Image** and **Image Edit** tabs; models load lazily on
 | `--host`    | `0.0.0.0` | bind address                                                                |
 | `--port`    | `7860`    | port                                                                        |
 | `--device`  | `cuda`    | inference device                                                            |
+| `--text-device`  | `cuda`    | text encoding device                                                            |
 | `--share`   | off         | create a public Gradio share link                                           |
 | `--preload` | *(lazy)*  | comma-separated repo ids / paths to load at startup instead of on first use |
 

mage_flow/app.py

@@ -42,6 +42,7 @@ def _repo(hf_id: str, local_name: str) -> str:
 }
 
 DEVICE = "cuda"
+TEXT_DEVICE: str | None = None
 _CACHE: dict[str, MageFlowPipeline] = {}
 
 
@@ -56,7 +57,8 @@ def _get_pipe(repo: str) -> MageFlowPipeline:
         raise gr.Error("No model specified.")
     if repo not in _CACHE:
         try:
-            _CACHE[repo] = MageFlowPipeline.from_pretrained(repo, device=DEVICE)
+            _CACHE[repo] = MageFlowPipeline.from_pretrained(
+                repo, device=DEVICE, text_device=TEXT_DEVICE)
         except Exception as e:  # noqa: BLE001
             raise gr.Error(f"Failed to load model '{repo}': {type(e).__name__}: {e}")
     return _CACHE[repo]
@@ -178,16 +180,20 @@ def build_ui():
 
 
 def main():
-    global DEVICE
+    global DEVICE, TEXT_DEVICE
     ap = argparse.ArgumentParser()
     ap.add_argument("--device", default="cuda")
+    ap.add_argument("--text-device", default=None,
+                    help="device for the text encoder (default: same as --device; "
+                         "e.g. --device cuda:0 --text-device cuda:1 splits across two GPUs)")
     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
+    TEXT_DEVICE = args.text_device
     if args.preload:
         for repo in args.preload.split(","):
             _get_pipe(repo.strip())

mage_flow/inference.py

@@ -51,6 +51,9 @@ def _add_common_args(p):
     p.add_argument("--static_shift", type=float, default=None,
                    help="override scheduler shift (default: repo scheduler_config.json, 6.0)")
     p.add_argument("--device", default="cuda")
+    p.add_argument("--text-device", default=None,
+                   help="device for the text encoder (default: same as --device; "
+                        "e.g. --device cuda:0 --text-device cuda:1 splits across two GPUs)")
     p.add_argument("--out", default="./outputs")
 
 
@@ -86,7 +89,7 @@ def main():
     args = p.parse_args()
 
     os.makedirs(args.out, exist_ok=True)
-    pipe = MageFlowPipeline.from_pretrained(args.model_path, args.device)
+    pipe = MageFlowPipeline.from_pretrained(args.model_path, args.device, args.text_device)
     n = len(args.prompt)
     imgs = pipe.generate(
         args.prompt,
@@ -133,7 +136,7 @@ def main_edit():
         p.error(f"--ref count ({len(args.ref)}) must match --prompt count ({len(args.prompt)})")
 
     os.makedirs(args.out, exist_ok=True)
-    pipe = MageFlowPipeline.from_pretrained(args.model_path, args.device)
+    pipe = MageFlowPipeline.from_pretrained(args.model_path, args.device, args.text_device)
     n = len(args.prompt)
     # Each --ref token is one prompt's reference(s); commas split multi-image refs.
     ref_images = [[s.strip() for s in r.split(",") if s.strip()] for r in args.ref]

mage_flow/pipeline.py

@@ -228,8 +228,12 @@ def _encode_texts_packed(model, prompts, template, drop_idx, device):
                   return_tensors="pt").input_ids.squeeze(0)
         for p in prompts
     ]
-    input_ids = torch.cat(ids_list).to(device)
-    cu_seqlens = _lens_to_cu([int(t.numel()) for t in ids_list], device)
+    # Text encoder may live on a different device than the transformer/VAE
+    # (two-GPU split) — encode on ITS device; `device` param is where the
+    # caller (_slice_packed) will move the result afterward.
+    txt_device = next(model.txt_enc.parameters()).device
+    input_ids = torch.cat(ids_list).to(txt_device)
+    cu_seqlens = _lens_to_cu([int(t.numel()) for t in ids_list], txt_device)
     res = model.txt_enc(
         input_ids, cu_seqlens, drop_idx_override=drop_idx)
     return res["txt"], res["vec"], res["txt_seq_lens"].tolist()
@@ -397,17 +401,18 @@ def _encode_edits_packed(model, ref_pils_per_sample, instructions, template, dro
     """Encode ALL image-conditioned edit instructions in ONE packed multimodal
     varlen forward (pixel_values/image_grid_thw concatenated across samples,
     cu_seqlens isolates each). Returns (txt_flat [ΣLi, D], vec [N, D], per-sample lens)."""
+    txt_device = next(model.txt_enc.parameters()).device
     processor = model.txt_enc.processor
     ids_list, pv_list, thw_list = [], [], []
     for ref_pils, instr in zip(ref_pils_per_sample, instructions, strict=False):
         formatted = template.format(_edit_prompt_body(instr, len(ref_pils)))
         vl = processor(text=[formatted], images=list(ref_pils), padding=True, return_tensors="pt")
-        vl = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in vl.items()}
+        vl = {k: (v.to(txt_device) if hasattr(v, "to") else v) for k, v in vl.items()}
         ids_list.append(vl["input_ids"].squeeze(0))
         if vl.get("pixel_values") is not None:
             pv_list.append(vl["pixel_values"]); thw_list.append(vl["image_grid_thw"])
-    input_ids = torch.cat(ids_list).to(device)
-    cu = _lens_to_cu([int(t.numel()) for t in ids_list], device)
+    input_ids = torch.cat(ids_list).to(txt_device)
+    cu = _lens_to_cu([int(t.numel()) for t in ids_list], txt_device)
     inputs = {"input_ids": input_ids, "cu_seqlens": cu}
     if pv_list:
         inputs["pixel_values"] = torch.cat(pv_list, dim=0)
@@ -649,15 +654,20 @@ def __init__(self, model, device="cuda"):
         self.device = device
 
     @classmethod
-    def from_pretrained(cls, repo_dir: str, device: str = "cuda"):
+    def from_pretrained(cls, repo_dir: str, device: str = "cuda", text_device: str | None = None):
         """Load a Mage-Flow diffusers-style repo (``model_index.json`` +
         ``transformer/`` ``vae/`` ``scheduler/`` ``text_encoder/``).
 
         ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id
         (e.g. ``"microsoft/Mage-Flow-4B"``), which is downloaded and cached
         automatically on first use.
+
+        ``device`` hosts the transformer + VAE (this is also what ``generate``/
+        ``edit`` use for image tensors, scheduler, etc). Pass ``text_device``
+        (e.g. ``"cuda:1"`` while ``device="cuda:0"``) to put the Qwen3-VL text
+        encoder on a second GPU and roughly halve peak VRAM per device.
         """
-        return cls(load_from_repo(repo_dir, device), device)
+        return cls(load_from_repo(repo_dir, device, text_device), device)
 
     def generate(self, prompts, **kw) -> list[Image.Image]:
         """Packed multi-resolution t2i. ``prompts`` is a list (or a single
@@ -710,13 +720,17 @@ def _resolve_repo_dir(repo_dir: str) -> str:
     return snapshot_download(repo_id=repo_dir)
 
 
-def load_from_repo(repo_dir: str, device: str = "cuda") -> MageFlowModel:
+def load_from_repo(repo_dir: str, device: str = "cuda", text_device: str | None = None) -> MageFlowModel:
     """Load a Mage-Flow diffusers-style repo (model_index.json + transformer/
     vae/ scheduler/). Transformer weights come from the bf16 safetensors;
     VAE + text encoder are built from the sources recorded in model_index.json.
 
     ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id (e.g.
     ``microsoft/Mage-Flow-4B``), which is downloaded/cached automatically.
+
+    ``device`` hosts the transformer + VAE. ``text_device`` (defaults to
+    ``device``) hosts the Qwen3-VL text encoder — pass a different value
+    (e.g. ``"cuda:1"`` vs ``"cuda:0"``) to split the stack across two GPUs.
     """
     from safetensors.torch import load_file
     repo_dir = _resolve_repo_dir(repo_dir)
@@ -750,11 +764,17 @@ def _resolve(p):
     sd = load_file(_safe_subpath(repo_dir, "transformer", "diffusion_pytorch_model.safetensors"),
                    device="cpu")
     model.transformer.load_state_dict(sd, strict=False, assign=True)
-    model.to(device)
+    # Move each component to its target device separately to avoid OOM
+    # from putting everything on one GPU simultaneously.
+    model.transformer.to(device)
     model.transformer.to(torch.bfloat16)
-    model.txt_enc.to(torch.bfloat16)
     if model.vae is not None:
+        model.vae.to(device)
         model.vae.to(torch.bfloat16)
+    # Text encoder goes to text_device if specified, otherwise to device.
+    _txt_dev = text_device if text_device is not None else device
+    model.txt_enc.to(_txt_dev)
+    model.txt_enc.to(torch.bfloat16)
     model.eval()
     # Diffusers FlowMatchEulerDiscreteScheduler (scheduler/scheduler_config.json).
     model.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(