Upload t2i.py with huggingface_hub

t2i.py

@@ -0,0 +1,174 @@
+"""Text-to-Image — hybrid MLX + PyTorch.
+
+Phase 1 (MLX): block-diffusion VQ token generation with CFG.
+Phase 2 (PyTorch): SigVQ + ZImageTransformer2DModel + VAE → pixel image.
+
+The MLX backbone is released before the PyTorch decoder loads to fit in 64 GB
+unified memory (decoder is ~12 GB, backbone is ~32 GB).
+"""
+import argparse
+import gc
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import mlx.core as mx
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer
+
+REPO_ROOT = Path(__file__).resolve().parent.parent / "llada2-uni-repo"
+sys.path.insert(0, str(REPO_ROOT))
+
+# Stub out flash_attn (not available on Apple Silicon). The decoder has a
+# dispatch_attention_fn fallback via diffusers that we use instead.
+import types as _types, importlib.machinery as _im
+if "flash_attn" not in sys.modules:
+    _stub = _types.ModuleType("flash_attn")
+    _stub.__spec__ = _im.ModuleSpec(name="flash_attn", loader=None)
+    _stub.__version__ = "0.0.0-stub"
+    _stub.flash_attn_func = lambda *a, **k: (_ for _ in ()).throw(
+        RuntimeError("flash_attn unavailable"))
+    sys.modules["flash_attn"] = _stub
+
+from llada2.model import LLaDA2Config, LLaDA2Model
+from llada2.weights import load_weights_into_model
+from llada2.generate_image import generate_image_tokens, extract_vq_tokens
+
+
+def build_t2i_prompt(tokenizer, prompt_text: str, image_h: int, image_w: int):
+    """Return (cond_ids, uncond_ids) — prompt id lists for CFG."""
+    sys_tmpl = "You are a text-to-image generation assistant."
+    # _build_chat equivalent
+    sys_ids = tokenizer(f"<role>SYSTEM</role> {sys_tmpl} <role>HUMAN</role>").input_ids
+    asst_ids = tokenizer("<role>ASSISTANT</role>").input_ids
+
+    soi = tokenizer("<|image|>").input_ids
+    boi = tokenizer("<boi>").input_ids
+    h_tok = tokenizer(f"<|reserved_token_{image_h}|>").input_ids
+    w_tok = tokenizer(f"<|reserved_token_{image_w}|>").input_ids
+    img_header = soi + h_tok + w_tok + boi
+
+    cond_ids = sys_ids + tokenizer(prompt_text).input_ids + asst_ids + img_header
+    uncond_ids = sys_ids + tokenizer("<uncondition>").input_ids + asst_ids + img_header
+    return cond_ids, uncond_ids
+
+
+def decode_to_pixels(token_ids: list[int], h: int, w: int, model_path: Path,
+                     decoder_steps: int, resolution_multiplier: int,
+                     decode_mode: str = "decoder-turbo"):
+    """Call the official decoder to render pixels."""
+    import torch
+    from decoder import decode_vq_tokens
+    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+    return decode_vq_tokens(
+        token_ids, h, w, str(model_path), device,
+        resolution_multiplier=resolution_multiplier,
+        num_steps=decoder_steps, decode_mode=decode_mode,
+    )
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--prompt", required=True, type=str)
+    ap.add_argument("--image-h", default=512, type=int)
+    ap.add_argument("--image-w", default=512, type=int)
+    ap.add_argument("--steps", default=16, type=int)
+    ap.add_argument("--block-length", default=32, type=int)
+    ap.add_argument("--cfg-scale", default=4.0, type=float)
+    ap.add_argument("--decoder-steps", default=50, type=int)
+    ap.add_argument("--decode-mode", default="normal",
+                    choices=["decoder-turbo", "normal"],
+                    help="'normal' = full 50-step decoder (cleaner, ~8 min), "
+                         "'decoder-turbo' = 8-step distilled (faster but brittle ≈ striping)")
+    ap.add_argument("--resolution-multiplier", default=2, type=int)
+    ap.add_argument("--output", default="t2i_output.png", type=str)
+    ap.add_argument("--repo-id", default="inclusionAI/LLaDA2.0-Uni", type=str)
+    ap.add_argument("--save-vq", default=None, type=str, help="Save intermediate VQ tokens to .json")
+    ap.add_argument("--load-vq", default=None, type=str, help="Skip phase 1, load VQ tokens from .json")
+    args = ap.parse_args()
+
+    print("[t2i] fetching model files…")
+    snap = Path(snapshot_download(
+        args.repo_id,
+        allow_patterns=[
+            "model-*.safetensors", "model.safetensors.index.json",
+            "config.json", "tokenizer*", "special_tokens_map.json",
+            "decoder-turbo/*", "decoder/*", "image_tokenizer/*", "vae/*",
+        ],
+    ))
+
+    # Generate image: LLaDA2 divides H and W by 2 internally before computing grid.
+    # Net result: grid = (image_h // 2 // 16) x (image_w // 2 // 16)
+    grid_h = args.image_h // 2 // 16
+    grid_w = args.image_w // 2 // 16
+    gen_length = grid_h * grid_w
+
+    if args.load_vq:
+        with open(args.load_vq) as f:
+            cached = json.load(f)
+        vq_tokens = cached["token_ids"]
+        grid_h, grid_w = cached["h"], cached["w"]
+        print(f"[t2i] loaded {len(vq_tokens)} VQ tokens from {args.load_vq}")
+    else:
+        # ---------- Phase 1: MLX VQ-token generation ----------
+        tokenizer = AutoTokenizer.from_pretrained(str(snap), trust_remote_code=True)
+        config = LLaDA2Config.from_hf(json.loads((snap / "config.json").read_text()))
+
+        cond_ids, uncond_ids = build_t2i_prompt(tokenizer, args.prompt, grid_h, grid_w)
+        print(f"[t2i] prompt tokens: {len(cond_ids)} | grid: {grid_h}x{grid_w} ({gen_length} VQ tokens)")
+
+        print("[t2i] building model + loading backbone…")
+        model = LLaDA2Model(config)
+        t0 = time.time()
+        load_weights_into_model(model, snap, dtype=mx.bfloat16, verbose=False)
+        print(f"[t2i] backbone loaded in {time.time()-t0:.1f}s")
+
+        prompt_ids = mx.array([cond_ids], dtype=mx.int32)
+        uc_ids = mx.array([uncond_ids], dtype=mx.int32)
+
+        t0 = time.time()
+        out = generate_image_tokens(
+            model, prompt_ids, uc_ids,
+            gen_length=gen_length,
+            block_length=args.block_length,
+            steps_per_block=args.steps,
+            cfg_scale=args.cfg_scale,
+            mask_token_id=config.mask_token_id,
+            image_token_offset=config.image_token_offset,
+            vocab_size=config.vocab_size,
+        )
+        mx.eval(out)
+        vq_tokens = (out[0, len(cond_ids):len(cond_ids) + gen_length] - config.image_token_offset).tolist()
+        print(f"[t2i] VQ generation in {time.time()-t0:.1f}s, {len(vq_tokens)} tokens, "
+              f"range [{min(vq_tokens)}, {max(vq_tokens)}]")
+
+        if args.save_vq:
+            with open(args.save_vq, "w") as f:
+                json.dump({"token_ids": vq_tokens, "h": grid_h, "w": grid_w,
+                           "prompt": args.prompt}, f)
+            print(f"[t2i] saved VQ tokens → {args.save_vq}")
+
+        # ---------- Free MLX backbone before PyTorch decoder loads ----------
+        del model, out
+        gc.collect()
+        mx.clear_cache()
+
+    # ---------- Phase 2: PyTorch decode → pixels ----------
+    print(f"[t2i] decoding VQ tokens → pixels ({args.decoder_steps} steps)…")
+    t0 = time.time()
+    img = decode_to_pixels(
+        vq_tokens, grid_h, grid_w, snap,
+        decoder_steps=args.decoder_steps,
+        resolution_multiplier=args.resolution_multiplier,
+        decode_mode=args.decode_mode,
+    )
+    print(f"[t2i] decoded in {time.time()-t0:.1f}s")
+
+    img.save(args.output)
+    print(f"[t2i] wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()