[Model] Add Tencent AuK speech generation and editing (encoder + diffusion pipeline) (#7385)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Sy03 <1370724210@qq.com>

tools/auk_parity_reference.py

@@ -0,0 +1,227 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Generate the upstream AuK reference artifacts consumed by the parity tests.
+
+This runs the *upstream* implementation, so it must be executed in an
+environment where the upstream package is installed
+(``pip install -e /path/to/AuK``; it pins its own torch/transformers line and
+does not need vLLM-Omni). It writes:
+
+    <out>/parity_ref/<variant>/<case>.pt, <case>.wav, meta.json
+        Upstream generations for four cookbook cases (zero-shot TTS, instruct
+        TTS, content edit, denoise) with the reference latents, the generated
+        latents and the waveform. Point ``AUK_PARITY_REF`` at
+        ``<out>/parity_ref/<variant>`` for tests/diffusion/models/auk/test_pipeline_auk.py.
+    <out>/fusion_ref/zs_tts_audio.pt, instruct_text_only.pt
+        The fp32 layer-fused text condition from the HF thinker for the
+        zero-shot and text-only prompts (``AUK_FUSION_REF``).
+
+By default the VAE encode uses the posterior mean, matching the port's default
+and making the artifacts deterministic; ``--vae-sample`` reproduces the
+upstream stochastic draw (seeded through ``--global-seed``). The ODE noise is
+drawn from ``--seed`` on both sides.
+
+Usage:
+    python tools/auk_parity_reference.py \\
+        --auk-repo /path/to/AuK --ckpt-dir /path/to/AuK \\
+        --qwen-dir /path/to/Qwen2.5-Omni-3B --out /path/to/auk-parity \\
+        [--variant base|flash] [--vae-sample] [--fusion-device cpu|cuda] [--skip-fusion]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+WEIGHTS = {"base": "auk_base.safetensors", "flash": "auk_flash.safetensors"}
+
+# Cookbook cases; audio paths are relative to <auk-repo>/assets/demo-input-audio.
+CASES: dict[str, dict[str, Any]] = {
+    "zs_tts_en": {
+        "text": (
+            "Say the following with the same voice: 'Ladies and gentlemen, it's an honor to have the "
+            "opportunity to address such a distinguished audience'"
+        ),
+        "audio": "zero-shot-tts/ref.wav",
+        "gen_seconds": 6.0,
+    },
+    "instruct_en": {
+        "text": (
+            'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". '
+            'The content to speak is: "Welcome back, how was your day at work?".'
+        ),
+        "audio": None,
+        "gen_seconds": 3.5,
+    },
+    "content_edit": {
+        "text": "Replace 'but accepting what we cannot have' with 'and living well with dreams unmet'.",
+        "audio": "content-edit/content.wav",
+        "gen_seconds": 7.0,
+    },
+    "denoise": {
+        "text": "Keep pure speech voice, remove noise and reverberation.",
+        "audio": "vocal-extraction/vocal-1-input.wav",
+        "gen_seconds": None,
+    },
+}
+# The two prompts the fusion reference covers (they must match the cases above).
+FUSION_CASES = {"zs_tts_audio": "zs_tts_en", "instruct_text_only": "instruct_en"}
+NO_PROMPT_AUDIO = "|<no_prompt_audio>|"
+
+
+def _messages(case: dict[str, Any], assets: Path) -> list[dict[str, Any]]:
+    content: list[dict[str, Any]] = [{"type": "text", "text": case["text"]}]
+    if case["audio"]:
+        content.append({"type": "audio", "audio": str(assets / case["audio"])})
+    return [{"role": "user", "content": content}]
+
+
+def _mean_encode(vae):
+    """Posterior-mean version of the upstream ``encoding_and_normalization``."""
+
+    def encode(sample, sample_lengths=None):
+        latent_stats = vae.audio_encoder(sample)
+        if sample_lengths is None:
+            sample_lengths = torch.LongTensor([sample.size(-1)] * sample.size(0)).to(sample.device)
+        latent_lens = sample_lengths // vae.hop_size
+        mean, _log_std = latent_stats.chunk(2, 1)
+        latents = mean.transpose(1, 2).float()
+        latents = (latents - vae.global_mean.float()) / torch.sqrt(vae.global_log_std.float())
+        return latents, torch.clamp(latent_lens, max=latents.size(1))
+
+    return encode
+
+
+def generate_references(args: argparse.Namespace) -> None:
+    from auk.infer.infer_auk import AukInfer, save_audio
+
+    ckpt = Path(args.ckpt_dir)
+    out = Path(args.out) / "parity_ref" / args.variant
+    out.mkdir(parents=True, exist_ok=True)
+    assets = Path(args.auk_repo) / "assets" / "demo-input-audio"
+
+    engine = AukInfer(str(ckpt / "config.yaml"), str(ckpt / WEIGHTS[args.variant]), qwen_path=args.qwen_dir)
+    encode = engine.vae_model.encoding_and_normalization if args.vae_sample else _mean_encode(engine.vae_model)
+    recorded: dict[str, torch.Tensor] = {}
+
+    def encode_hook(sample, sample_lengths=None):
+        latents, lens = encode(sample, sample_lengths)
+        recorded["ref_latents"] = latents.detach().float().cpu()
+        recorded["ref_latent_lens"] = lens.detach().cpu()
+        return latents, lens
+
+    engine.vae_model.encoding_and_normalization = encode_hook
+    upstream_sample = engine.model.sample
+
+    def sample_hook(*a, **k):
+        latents, trajectory = upstream_sample(*a, **k)
+        recorded["generated"] = latents.detach().float().cpu()
+        return latents, trajectory
+
+    engine.model.sample = sample_hook
+
+    meta: dict[str, Any] = {}
+    for name, case in CASES.items():
+        messages = _messages(case, assets)
+        recorded.clear()
+        torch.manual_seed(args.global_seed)
+        started = time.perf_counter()
+        audio, sr = engine.generate(messages, gen_seconds=case["gen_seconds"], seed=args.seed)
+        wall = time.perf_counter() - started
+        save_audio(audio, sr, str(out / f"{name}.wav"))
+        torch.save(
+            {
+                "messages": messages,
+                "gen_seconds": case["gen_seconds"],
+                "seed": args.seed,
+                "global_seed": args.global_seed,
+                "vae_sample": bool(args.vae_sample),
+                "audio": audio.float().cpu(),
+                "sr": sr,
+                **recorded,
+            },
+            out / f"{name}.pt",
+        )
+        ref_frames = int(recorded["ref_latent_lens"][0]) if "ref_latent_lens" in recorded else 0
+        meta[name] = {
+            "wall_s": wall,
+            "out_s": audio.shape[-1] / sr,
+            "sr": sr,
+            "ref_frames": ref_frames,
+            "gen_frames": int(recorded["generated"].shape[1]) - ref_frames,
+        }
+        print(name, meta[name], flush=True)
+    (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n")
+
+
+def generate_fusion(args: argparse.Namespace) -> None:
+    from auk.model.cfm_edit import CFMEdit
+    from safetensors import safe_open
+    from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
+
+    out = Path(args.out) / "fusion_ref"
+    out.mkdir(parents=True, exist_ok=True)
+    assets = Path(args.auk_repo) / "assets" / "demo-input-audio"
+    with safe_open(str(Path(args.ckpt_dir) / WEIGHTS[args.variant]), "pt") as f:
+        layer_weights = f.get_tensor("layer_weights").float()
+        layer_scale = f.get_tensor("layer_scale").float()
+
+    processor = Qwen2_5OmniProcessor.from_pretrained(args.qwen_dir)
+    thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(args.qwen_dir, torch_dtype=torch.float32)
+    thinker.visual = None
+    thinker = thinker.to(args.fusion_device).eval()
+
+    for out_name, case_name in FUSION_CASES.items():
+        case = dict(CASES[case_name])
+        if not case["audio"] and not case["text"].endswith(NO_PROMPT_AUDIO):
+            case["text"] = case["text"] + NO_PROMPT_AUDIO
+        cond = CFMEdit.build_cond_inputs([_messages(case, assets)], processor).to(args.fusion_device)
+        with torch.no_grad():
+            hidden_states = thinker(**cond, output_hidden_states=True).hidden_states
+        width = hidden_states[0].shape[-1]
+        stacked = torch.stack([F.layer_norm(h.float(), [width]) for h in hidden_states[1:]], 0)
+        weights = F.softmax(layer_weights.to(stacked.device), 0)
+        fused = (stacked * weights[:, None, None, None]).sum(0) * layer_scale.to(stacked.device)
+        torch.save(
+            {
+                "input_ids": cond["input_ids"].cpu(),
+                "attention_mask": cond["attention_mask"].cpu(),
+                "fused": fused[0].float().cpu(),
+                "final_normed": hidden_states[-1][0].float().cpu(),
+                "text": case["text"],
+            },
+            out / f"{out_name}.pt",
+        )
+        print(out_name, {"tokens": int(cond["input_ids"].shape[1]), "fused_norm": float(fused.norm())}, flush=True)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--auk-repo", required=True, help="checkout of the upstream AuK repository (for demo clips)")
+    parser.add_argument("--ckpt-dir", required=True, help="tencent/AuK or tencent/AuK-Flash snapshot")
+    parser.add_argument("--qwen-dir", required=True, help="Qwen/Qwen2.5-Omni-3B snapshot")
+    parser.add_argument("--out", required=True, help="output root; writes parity_ref/<variant> and fusion_ref")
+    parser.add_argument("--variant", choices=sorted(WEIGHTS), default="base")
+    parser.add_argument("--seed", type=int, default=0, help="ODE noise seed passed to the upstream sampler")
+    parser.add_argument("--global-seed", type=int, default=1234, help="global RNG seed set before each generate()")
+    parser.add_argument("--vae-sample", action="store_true", help="draw the VAE posterior sample as upstream does")
+    parser.add_argument("--fusion-device", default="cpu", help="device for the fp32 HF thinker forward")
+    parser.add_argument("--skip-fusion", action="store_true", help="only write the generation references")
+    args = parser.parse_args()
+
+    generate_references(args)
+    if not args.skip_fusion:
+        generate_fusion(args)
+    print(f"reference artifacts written under {args.out}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tools/prepare_auk_checkpoint.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Assemble an AuK checkpoint directory that ``--model <path>`` can load.
+
+The released AuK repos ship ``config.yaml`` plus loose safetensors and depend
+on a separate Qwen2.5-Omni-3B snapshot for the frozen text encoder. This tool
+writes the ``config.json`` the vLLM-Omni pipeline expects and links every
+artifact both stages need into one flat directory:
+
+    <out>/
+      config.json          the Qwen2.5-Omni-3B config with model_type "auk",
+                           AuK's architecture, and the dit / vae / variant /
+                           flash_t_grid / defaults sections added at the top
+                           level (see transformers_utils/configs/auk.py)
+      auk.safetensors      -> auk_base.safetensors or auk_flash.safetensors
+      vae.safetensors      -> the AuK VAE weights
+      thinker/             -> the Qwen2.5-Omni-3B snapshot, for reference only
+      model-0000*.safetensors, model.safetensors.index.json,
+      tokenizer*, vocab.json, merges.txt, added_tokens.json,
+      special_tokens_map.json, chat_template.json,
+      preprocessor_config.json, generation_config.json, spk_dict.pt
+                           -> the same snapshot's entries
+
+The Qwen entries are linked at the top level rather than left under
+``thinker/`` because both stages load from this one directory: the encoder
+reads the thinker shards, the tokenizer and the audio processor from here,
+and takes the layer-fusion tensors out of ``auk.safetensors``. The weight
+index keeps ``auk.safetensors`` and ``vae.safetensors`` out of the thinker's
+own weight load, and the DiT and VAE weights are read by model code directly.
+
+Usage:
+    python tools/prepare_auk_checkpoint.py \
+        --auk-dir  /path/to/AuK \
+        --qwen-dir /path/to/Qwen2.5-Omni-3B \
+        --out      /path/to/auk-omni-base
+
+Re-running against an existing directory is safe: config.json is rewritten
+and the links are repointed. Pass ``--copy`` on filesystems without symlinks.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+# The distilled student's 4-step grid. Duplicated from
+# vllm_omni/transformers_utils/configs/auk.py so this tool stays runnable
+# without importing the package (and its torch dependency).
+FLASH_T_GRID = [
+    0.0,
+    0.07612049579620361,
+    0.2928932309150696,
+    0.6173166036605835,
+    1.0,
+]
+
+# Names the assembled directory owns; everything else in it is a Qwen entry.
+WEIGHTS_LINK = "auk.safetensors"
+VAE_LINK = "vae.safetensors"
+THINKER_LINK = "thinker"
+CONFIG_NAME = "config.json"
+
+# Snapshot entries that are not reused: the config is rewritten, and the repo
+# paperwork would misdescribe the assembled directory.
+_SKIP_QWEN_ENTRIES = frozenset({CONFIG_NAME, "README.md", "LICENSE"})
+
+DIT_KEYS = ("dim", "heads", "ff_mult", "text_hidden_dim", "num_layers", "num_single_layers")
+VAE_KEYS = ("latent_dim", "downsample_rate", "target_sample_rate", "model_init_kwargs")
+
+_INTERPOLATION = re.compile(r"^\$\{([A-Za-z0-9_.]+)\}$")
+
+
+def _strip_comment(line: str) -> str:
+    """Drop a trailing ``#`` comment that is not inside a quoted scalar."""
+    out: list[str] = []
+    quote: str | None = None
+    for ch in line:
+        if quote is not None:
+            if ch == quote:
+                quote = None
+            out.append(ch)
+            continue
+        if ch in "\"'":
+            quote = ch
+            out.append(ch)
+            continue
+        if ch == "#":
+            break
+        out.append(ch)
+    return "".join(out).rstrip()
+
+
+def _parse_scalar(text: str) -> Any:
+    """Parse one YAML scalar or inline sequence."""
+    text = text.strip()
+    if not text:
+        return None
+    lowered = text.lower()
+    if lowered in ("true", "yes"):
+        return True
+    if lowered in ("false", "no"):
+        return False
+    if lowered in ("null", "~"):
+        return None
+    if text[0] in "[{":
+        normalized = re.sub(r"\b(True|true|yes)\b", "True", text)
+        normalized = re.sub(r"\b(False|false|no)\b", "False", normalized)
+        normalized = re.sub(r"\b(null|~)\b", "None", normalized)
+        try:
+            return ast.literal_eval(normalized)
+        except (SyntaxError, ValueError):
+            return text
+    try:
+        return ast.literal_eval(text)
+    except (SyntaxError, ValueError):
+        return text.strip("\"'")
+
+
+def _parse_yaml_minimal(text: str) -> dict[str, Any]:
+    """Parse the small subset of YAML the AuK configs use.
+
+    Nested mappings by indentation, scalars, and inline sequences. Used only
+    when PyYAML is unavailable.
+    """
+    root: dict[str, Any] = {}
+    stack: list[tuple[int, dict[str, Any]]] = [(-1, root)]
+    for raw in text.splitlines():
+        line = _strip_comment(raw)
+        if not line.strip():
+            continue
+        indent = len(line) - len(line.lstrip(" "))
+        body = line.strip()
+        if ":" not in body:
+            raise ValueError(f"unsupported YAML line (no key): {raw!r}")
+        key, _, rest = body.partition(":")
+        key = key.strip()
+        while stack and indent <= stack[-1][0]:
+            stack.pop()
+        if not stack:
+            raise ValueError(f"unsupported YAML indentation at: {raw!r}")
+        parent = stack[-1][1]
+        if rest.strip():
+            parent[key] = _parse_scalar(rest)
+        else:
+            child: dict[str, Any] = {}
+            parent[key] = child
+            stack.append((indent, child))
+    return root
+
+
+def _resolve_interpolations(node: Any, root: dict[str, Any]) -> Any:
+    """Replace ``${a.b.c}`` scalars with the value at that dotted path."""
+    if isinstance(node, dict):
+        return {k: _resolve_interpolations(v, root) for k, v in node.items()}
+    if isinstance(node, list):
+        return [_resolve_interpolations(v, root) for v in node]
+    if isinstance(node, str):
+        match = _INTERPOLATION.match(node.strip())
+        if match is None:
+            return node
+        target: Any = root
+        for part in match.group(1).split("."):
+            if not isinstance(target, dict) or part not in target:
+                raise ValueError(f"cannot resolve interpolation ${{{match.group(1)}}}")
+            target = target[part]
+        return target
+    return node
+
+
+def load_auk_yaml(path: Path) -> dict[str, Any]:
+    """Load the AuK ``config.yaml`` with interpolations resolved."""
+    text = path.read_text()
+    try:
+        import yaml
+
+        parsed = yaml.safe_load(text)
+    except ImportError:
+        parsed = _parse_yaml_minimal(text)
+    if not isinstance(parsed, dict) or "model" not in parsed:
+        raise ValueError(f"{path} has no top-level 'model' section")
+    return _resolve_interpolations(parsed, parsed)
+
+
+def detect_variant(auk_dir: Path, model_section: dict[str, Any]) -> str:
+    """Infer base vs flash from the weight filename, then the model name."""
+    if (auk_dir / "auk_flash.safetensors").exists():
+        return "flash"
+    if (auk_dir / "auk_base.safetensors").exists():
+        return "base"
+    name = str(model_section.get("name", "")).lower()
+    return "flash" if "flash" in name else "base"
+
+
+def weights_name(variant: str) -> str:
+    return "auk_flash.safetensors" if variant == "flash" else "auk_base.safetensors"
+
+
+def build_config(
+    auk_yaml: dict[str, Any],
+    qwen_config: dict[str, Any],
+    variant: str,
+) -> dict[str, Any]:
+    """Assemble the ``config.json`` body for the output directory.
+
+    The Qwen config is kept whole, including ``talker_config`` and
+    ``token2wav_config``: ``AuKConfig`` subclasses ``Qwen2_5OmniConfig`` and
+    the encoder stage reuses the in-tree thinker processor, which asserts that
+    type. AuK's own sections are added at the top level.
+    """
+    model = auk_yaml["model"]
+    arch = model.get("arch", {})
+    vae = model.get("vae", {})
+
+    missing = [key for key in DIT_KEYS if key not in arch]
+    if missing:
+        raise ValueError(f"AuK config.yaml arch section is missing {missing}")
+    missing = [key for key in VAE_KEYS if key not in vae]
+    if missing:
+        raise ValueError(f"AuK config.yaml vae section is missing {missing}")
+    if not isinstance(qwen_config.get("thinker_config"), dict):
+        raise ValueError("Qwen snapshot config.json has no 'thinker_config' mapping")
+
+    dit = {key: arch[key] for key in DIT_KEYS}
+    # dim_head is implicit upstream (dim // heads) but the port takes it
+    # explicitly, so it is recorded rather than recomputed at load time.
+    dit["dim_head"] = int(arch["dim"]) // int(arch["heads"])
+    # Carried from the released config: the DiT builds an attention mask over
+    # the padded [audio | text] sequence and parity depends on it.
+    dit["attn_mask_enabled"] = bool(arch.get("attn_mask_enabled", True))
+
+    config = dict(qwen_config)
+    config["model_type"] = "auk"
+    config["architectures"] = ["AuKForConditionalGeneration"]
+    config["dit"] = dit
+    config["vae"] = {key: vae[key] for key in VAE_KEYS}
+    config["variant"] = variant
+    config["flash_t_grid"] = FLASH_T_GRID
+    config["defaults"] = {"nfe": 32, "cfg": 2.0, "sway": -1.0}
+    return config
+
+
+def _replace_path(target: Path) -> None:
+    """Remove an existing file, directory, or link at ``target``."""
+    if target.is_symlink() or target.is_file():
+        target.unlink()
+    elif target.is_dir():
+        shutil.rmtree(target)
+
+
+def link_or_copy(source: Path, target: Path, *, copy: bool) -> None:
+    """Point ``target`` at ``source``, replacing whatever is already there."""
+    if not source.exists():
+        raise FileNotFoundError(f"missing source artifact: {source}")
+    if target.is_symlink() and not copy and target.readlink() == source.resolve():
+        return
+    _replace_path(target)
+    if copy:
+        if source.is_dir():
+            shutil.copytree(source, target)
+        else:
+            shutil.copy2(source, target)
+        return
+    target.symlink_to(source.resolve(), target_is_directory=source.is_dir())
+
+
+MANIFEST_NAME = ".auk-assembled.json"
+
+
+def _reject_overlap(out: Path, *inputs: Path) -> None:
+    """Refuse an output directory that is, contains, or lives inside an input."""
+    out_r = out.resolve()
+    for src in inputs:
+        src_r = src.resolve()
+        if out_r == src_r or out_r in src_r.parents or src_r in out_r.parents:
+            raise ValueError(f"--out {out} overlaps input directory {src}; choose a separate directory")
+
+
+def assemble(auk_dir: Path, qwen_dir: Path, out: Path, variant: str, *, copy: bool) -> None:
+    """Link the AuK weights and the Qwen snapshot entries into ``out``."""
+    _reject_overlap(out, auk_dir, qwen_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    manifest_path = out / MANIFEST_NAME
+    previous: set[str] = set()
+    if manifest_path.is_file():
+        previous = set(json.loads(manifest_path.read_text()).get("entries", []))
+    weights = auk_dir / weights_name(variant)
+
+    entries: dict[str, Path] = {
+        WEIGHTS_LINK: weights,
+        VAE_LINK: auk_dir / VAE_LINK,
+        THINKER_LINK: qwen_dir,
+    }
+    # The assembled config replaces the Qwen one; every other entry is reused.
+    for entry in qwen_dir.iterdir():
+        if entry.name.startswith(".") or entry.name in _SKIP_QWEN_ENTRIES:
+            continue
+        entries[entry.name] = entry
+
+    for name, source in entries.items():
+        link_or_copy(source, out / name, copy=copy)
+
+    # Drop only entries this tool created in an earlier run and no longer needs;
+    # anything else in the directory belongs to the user.
+    for name in previous - set(entries) - {CONFIG_NAME}:
+        stale = out / name
+        if stale.exists() or stale.is_symlink():
+            _replace_path(stale)
+    manifest_path.write_text(json.dumps({"entries": sorted(entries)}, indent=1) + "\n")
+
+
+def print_tree(out: Path) -> None:
+    """Print the assembled directory, one entry per line."""
+    print(f"\n{out}/")
+    for entry in sorted(out.iterdir()):
+        if entry.is_symlink():
+            print(f"  {entry.name} -> {entry.readlink()}")
+        elif entry.is_dir():
+            print(f"  {entry.name}/")
+        else:
+            print(f"  {entry.name}  ({entry.stat().st_size / 1e9:.2f} GB)")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--auk-dir", required=True, help="HF snapshot of the AuK or AuK-Flash repo")
+    parser.add_argument("--qwen-dir", required=True, help="HF snapshot of Qwen2.5-Omni-3B")
+    parser.add_argument("--out", required=True, help="Directory to assemble (created if absent)")
+    parser.add_argument(
+        "--variant",
+        choices=("base", "flash"),
+        help="Override the variant; detected from the AuK snapshot by default",
+    )
+    parser.add_argument(
+        "--copy",
+        action="store_true",
+        help="Copy the weights and the snapshot entries instead of symlinking",
+    )
+    args = parser.parse_args()
+
+    auk_dir = Path(args.auk_dir).expanduser().resolve()
+    qwen_dir = Path(args.qwen_dir).expanduser().resolve()
+    out = Path(args.out).expanduser().resolve()
+    _reject_overlap(out, auk_dir, qwen_dir)
+
+    auk_yaml = load_auk_yaml(auk_dir / "config.yaml")
+    variant = args.variant or detect_variant(auk_dir, auk_yaml["model"])
+    qwen_config = json.loads((qwen_dir / CONFIG_NAME).read_text())
+
+    config = build_config(auk_yaml, qwen_config, variant)
+    assemble(auk_dir, qwen_dir, out, variant, copy=args.copy)
+    (out / CONFIG_NAME).write_text(json.dumps(config, indent=2) + "\n")
+
+    print(f"variant: {variant}  (weights {weights_name(variant)})")
+    print_tree(out)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())