Upload scripts/run_omnistep_12a3b.py with huggingface_hub

scripts/run_omnistep_12a3b.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python3
+"""
+OmniStep 12A3B — main entry point.
+
+A complete multimodal model: real-time ASR, text reasoning, TTS, image understanding,
+and infinite background music generation. Built as a paper-exact Darwin weight-space
+recombination (arXiv:2605.14386) of Qwen2.5-Omni-3B and ACE-Step v1.5 XL SFT 4B.
+
+Usage:
+    python run_omnistep_6a3b.py text "Explain the Darwin Family paper."
+    python run_omnistep_6a3b.py music "chill lofi beats" --output ~/music/track.wav
+    python run_omnistep_6a3b.py music-loop "chill lofi beats"   # infinite background music
+    python run_omnistep_6a3b.py voice  # streaming voice assistant (ASR + TTS)
+    python run_omnistep_6a3b.py serve  # start the vllm server
+
+The model has ALL modalities. Pick a mode. They all work together.
+"""
+import argparse
+import asyncio
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+MODEL_ID = os.environ.get("OMNISTEP_MODEL", "sovthpaw/omnistep-12a3b")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# ============================================================================
+# Text reasoning — works with vllm OR llama-server
+# ============================================================================
+def cmd_text(prompt: str, max_tokens: int = 200, server_url: str = None):
+    """Quick text reasoning via the running server (vllm or llama-server)."""
+    import urllib.request
+    url = server_url or os.environ.get("OMNISTEP_SERVER", "http://localhost:8080")
+    payload = {
+        "model": "omnistep-12a3b",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "temperature": 0.7,
+    }
+    req = urllib.request.Request(
+        f"{url}/v1/chat/completions",
+        data=json.dumps(payload).encode(),
+        headers={"Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(req) as resp:
+        result = json.loads(resp.read())
+    print(result["choices"][0]["message"]["content"])
+
+# ============================================================================
+# Music generation — uses the ACE DiT (the diffusion part, F16 / unquantized)
+# ============================================================================
+def cmd_music(prompt: str, output: str, duration: int = 60, infer_steps: int = 8):
+    """One-shot music generation via the ACE-Step diffusion decoder.
+
+    Loads the model with transformers (full multimodal safetensors), uses the
+    OmniStep text body to craft the music prompt, and runs the ACE DiT to
+    generate audio. The DiT stays at F16 (unquantized) for max quality.
+    """
+    import torch
+    from transformers import AutoModel, AutoProcessor
+    import soundfile as sf
+
+    print(f"🎵 OmniStep 12A3B — music generation")
+    print(f"   prompt: {prompt}")
+    print(f"   duration: {duration}s, infer_steps: {infer_steps}")
+    print(f"   output: {output}")
+
+    # Load the full multimodal model (with the music head)
+    model = AutoModel.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+        token=HF_TOKEN,
+    )
+    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
+
+    # Use the OmniStep text body to craft a richer prompt
+    crafted = model.craft_music_prompt(prompt, duration=duration)
+    print(f"   crafted prompt: {crafted[:120]}...")
+
+    # Generate via the ACE DiT (unquantized F16 weights)
+    audio = model.generate_music(
+        prompt=crafted,
+        duration=duration,
+        infer_steps=infer_steps,
+    )
+
+    # Save
+    out_path = Path(output).expanduser()
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    sf.write(str(out_path), audio, 48000)
+    print(f"   ✓ wrote {out_path} ({out_path.stat().st_size/1e6:.1f}MB)")
+
+# ============================================================================
+# Infinite background music (the Evolutionary Radio)
+# ============================================================================
+def cmd_music_loop(prompt: str, queue_size: int = 5, duration: int = 60):
+    """Infinite background music generation.
+
+    Runs the 4 concurrent loops of the Evolutionary Radio:
+      1. Playback (mpv)
+      2. Queue fill (generate one track ahead, target 5-track queue)
+      3. GEPA prompt evolution (background, every 50 generations)
+      4. Darwin weight evolution (background, nightly)
+
+    This is the "infinitely generate its own background music" feature.
+    """
+    print(f"🎵 Evolutionary Radio (OmniStep 12A3B)")
+    print(f"   prompt: {prompt}")
+    print(f"   queue: {queue_size} tracks, duration: {duration}s each")
+    print(f"   Press Ctrl+C to stop.")
+    print()
+
+    # Lazy imports — only needed for music-loop
+    try:
+        from omnistep_radio import EvolutionaryRadio
+    except ImportError:
+        print("  Note: omnistep_radio.py not in the same directory.")
+        print("  Run from the OmniStep 12A3B repo root, or set PYTHONPATH.")
+        sys.exit(1)
+
+    radio = EvolutionaryRadio(prompt=prompt, queue_size=queue_size, duration=duration)
+    try:
+        asyncio.run(radio.run())
+    except KeyboardInterrupt:
+        print("\n  ✓ stopped.")
+
+# ============================================================================
+# Streaming voice (ASR + reasoning + TTS) — needs vllm
+# ============================================================================
+def cmd_voice(server_url: str = None):
+    """Streaming voice assistant: ASR (audio in) → text reasoning → TTS (audio out)."""
+    print("🎤 OmniStep 12A3B — streaming voice assistant")
+    print()
+    print("This mode uses vllm with the full safetensors for streaming ASR + TTS.")
+    print("If vllm is not running, start it first:")
+    print(f"    vllm serve {MODEL_ID} --max-model-len 32768 --gpu-memory-utilization 0.85")
+    print()
+    print("Then run this command. It will:")
+    print("  1. Stream audio in (your microphone)")
+    print("  2. Real-time ASR via the Whisper audio encoder")
+    print("  3. Text reasoning via the OmniStep text body")
+    print("  4. Streaming TTS via the Talker + token2wav speech-out heads")
+    print("  5. (Optional) background music via the ACE DiT, mixed in")
+    print()
+    print("Requires: sounddevice, numpy, vllm-client, the model running via vllm.")
+    print()
+
+    try:
+        from omnistep_voice import VoiceAssistant
+    except ImportError:
+        print("  Note: omnistep_voice.py not in the same directory.")
+        print("  Run from the OmniStep 12A3B repo root, or set PYTHONPATH.")
+        sys.exit(1)
+
+    va = VoiceAssistant(server_url=server_url or os.environ.get("OMNISTEP_SERVER", "http://localhost:8080"))
+    try:
+        va.run()
+    except KeyboardInterrupt:
+        print("\n  ✓ stopped.")
+
+# ============================================================================
+# Start the vllm server (the easiest way to run the model)
+# ============================================================================
+def cmd_serve(gpu_memory: float = 0.85, max_len: int = 32768, port: int = 8080):
+    """Start the vllm server with the full multimodal OmniStep 12A3B model."""
+    import subprocess
+    print(f"🚀 Starting vllm server for {MODEL_ID}")
+    print(f"   gpu_memory_utilization: {gpu_memory}, max_model_len: {max_len}, port: {port}")
+    print(f"   vllm must be installed (pip install vllm)")
+    cmd = [
+        "vllm", "serve", MODEL_ID,
+        "--max-model-len", str(max_len),
+        "--gpu-memory-utilization", str(gpu_memory),
+        "--port", str(port),
+        "--trust-remote-code",
+    ]
+    print(f"   $ {' '.join(cmd)}")
+    subprocess.run(cmd)
+
+# ============================================================================
+# Main
+# ============================================================================
+def main():
+    p = argparse.ArgumentParser(
+        description="OmniStep 12A3B — fast 4o-style streaming voice assistant with infinite background music",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Text reasoning (need a server running: vllm or llama-server)
+  python run_omnistep_6a3b.py text "Explain the Darwin Family paper."
+
+  # One-shot music generation (uses the ACE DiT, F16 unquantized)
+  python run_omnistep_6a3b.py music "chill lofi beats" --output ~/music/track.wav
+
+  # Infinite background music
+  python run_omnistep_6a3b.py music-loop "chill lofi beats"
+
+  # Streaming voice assistant (ASR + reasoning + TTS)
+  python run_omnistep_6a3b.py voice
+
+  # Start the vllm server (the easiest deployment)
+  python run_omnistep_6a3b.py serve
+
+The model has ALL modalities. Pick a mode. They all work together.
+        """,
+    )
+    sub = p.add_subparsers(dest="mode", required=True)
+
+    # text
+    pt = sub.add_parser("text", help="Text reasoning (via running vllm or llama-server)")
+    pt.add_argument("prompt")
+    pt.add_argument("--max-tokens", type=int, default=200)
+    pt.add_argument("--server", default=None, help="vllm/llama-server URL (default: $OMNISTEP_SERVER or http://localhost:8080)")
+
+    # music
+    pm = sub.add_parser("music", help="One-shot music generation via the ACE DiT (F16 unquantized)")
+    pm.add_argument("prompt")
+    pm.add_argument("--output", default="~/music/omnistep_track.wav")
+    pm.add_argument("--duration", type=int, default=60)
+    pm.add_argument("--infer-steps", type=int, default=8)
+
+    # music-loop (infinite background music)
+    pl = sub.add_parser("music-loop", help="Infinite background music (Evolutionary Radio)")
+    pl.add_argument("prompt")
+    pl.add_argument("--queue-size", type=int, default=5)
+    pl.add_argument("--duration", type=int, default=60)
+
+    # voice
+    pv = sub.add_parser("voice", help="Streaming voice assistant (ASR + TTS)")
+    pv.add_argument("--server", default=None)
+
+    # serve
+    ps = sub.add_parser("serve", help="Start the vllm server")
+    ps.add_argument("--gpu-memory", type=float, default=0.85)
+    ps.add_argument("--max-len", type=int, default=32768)
+    ps.add_argument("--port", type=int, default=8080)
+
+    args = p.parse_args()
+
+    if args.mode == "text":
+        cmd_text(args.prompt, args.max_tokens, args.server)
+    elif args.mode == "music":
+        cmd_music(args.prompt, args.output, args.duration, args.infer_steps)
+    elif args.mode == "music-loop":
+        cmd_music_loop(args.prompt, args.queue_size, args.duration)
+    elif args.mode == "voice":
+        cmd_voice(args.server)
+    elif args.mode == "serve":
+        cmd_serve(args.gpu_memory, args.max_len, args.port)
+
+if __name__ == "__main__":
+    main()