docs(quickstart): load mirrored local weights (no upstream fetch)

quickstart.py

@@ -1,95 +1,60 @@
-"""
-LumynaX Speech Kokoro 82M TTS — LumynaX quickstart.
-
-This script fetches the upstream model from Hugging Face and runs a short
-LumynaX-flavoured prompt. Run it on a host that satisfies the resource budget
-documented in the README (LumynaX Speech Kokoro 82M TTS).
-
-Usage:
-    python quickstart.py                # one-shot demo prompt
-    python quickstart.py --interactive  # REPL
-    python quickstart.py --gguf         # use the GGUF mirror via llama-cpp
-
-LumynaX package repo: https://huggingface.co/AbteeXAILab/lumynax-speech-kokoro-82m-tts
-Upstream weights:     https://huggingface.co/hexgrad/Kokoro-82M
-"""
-from __future__ import annotations
-import argparse, os, sys
-
-LUMYNAX_SYSTEM = (
-    "You are LumynaX, the AbteeX AI Labs assistant from Aotearoa New Zealand. "
-    "Ko te marama te tuapapa - the light is the foundation. "
-    "Answer with care, cite uncertainty, and prefer local-first reasoning. "
-    "Refuse unsafe, unlawful, or sovereignty-violating requests."
-)
-DEMO_PROMPT = "Explain in 3 bullets why local-first AI matters for Aotearoa New Zealand."
-
-def _run_hf(prompt: str, interactive: bool):
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    print("[lumynax] Loading hexgrad/Kokoro-82M. This is a >100B MoE — multi-GPU or accelerate offload recommended.")
-    tok = AutoTokenizer.from_pretrained("hexgrad/Kokoro-82M", trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        "hexgrad/Kokoro-82M", device_map="auto", torch_dtype="auto", trust_remote_code=True
-    )
-    def chat(user):
-        messages = [
-            {"role": "system", "content": LUMYNAX_SYSTEM},
-            {"role": "user",   "content": user},
-        ]
-        text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = tok(text, return_tensors="pt").to(model.device)
-        out = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.4)
-        return tok.decode(out[0, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-    if interactive:
-        print("[lumynax] interactive mode — empty line exits.")
-        while True:
-            try: q = input("you> ").strip()
-            except EOFError: break
-            if not q: break
-            print("lumynax> " + chat(q))
-    else:
-        print(chat(prompt))
-
-
-def _run_gguf(prompt: str, interactive: bool):
-    from llama_cpp import Llama
-    mirror = ""
-    if not mirror:
-        print("[lumynax] No community GGUF mirror registered for this build."); sys.exit(2)
-    print(f"[lumynax] Loading GGUF from {mirror}...")
-    llm = Llama.from_pretrained(
-        repo_id=mirror, filename="*Q4_K_M*.gguf",
-        n_ctx=510,
-        n_gpu_layers=int(os.environ.get("N_GPU_LAYERS", "-1")), verbose=False,
-    )
-    def chat(user):
-        out = llm.create_chat_completion(messages=[
-            {"role": "system", "content": LUMYNAX_SYSTEM},
-            {"role": "user",   "content": user},
-        ], max_tokens=512, temperature=0.4)
-        return out["choices"][0]["message"]["content"]
-    if interactive:
-        while True:
-            try: q = input("you> ").strip()
-            except EOFError: break
-            if not q: break
-            print("lumynax> " + chat(q))
-    else:
-        print(chat(prompt))
-
-
-def main():
-    p = argparse.ArgumentParser()
-    p.add_argument("--interactive", action="store_true")
-    p.add_argument("--prompt", default=DEMO_PROMPT)
-    p.add_argument("--gguf", action="store_true")
-    args = p.parse_args()
-    if args.gguf:
-        _run_gguf(args.prompt, args.interactive)
-    else:
-        _run_hf(args.prompt, args.interactive)
-
-
-if __name__ == "__main__":
-    main()
+"""
+Lumynax Speech Kokoro 82M Tts — LumynaX quickstart (clone & run, multimodal safetensors).
+
+Loads the local safetensors shards in this repo via transformers.
+Requires significant VRAM (160+ GB VRAM).
+
+Usage:
+  python quickstart.py --interactive
+  python quickstart.py --image foo.jpg --prompt "describe this"
+"""
+from __future__ import annotations
+import argparse, os, sys
+from pathlib import Path
+
+LUMYNAX_SYSTEM = "You are LumynaX, the AbteeX AI Labs assistant from Aotearoa New Zealand. Ko te marama te tuapapa. Answer with care; cite uncertainty; refuse unsafe asks."
+DEMO_PROMPT = "Explain in 3 bullets why local-first AI matters for Aotearoa New Zealand."
+HERE = Path(__file__).resolve().parent
+
+def main():
+    import torch
+    from transformers import AutoProcessor, AutoModelForImageTextToText
+    p = argparse.ArgumentParser()
+    p.add_argument("--interactive", action="store_true")
+    p.add_argument("--prompt", default=DEMO_PROMPT)
+    p.add_argument("--image", default=None)
+    args = p.parse_args()
+    if not (HERE / "kokoro-v1_0.pth").exists():
+        print(f"[lumynax] weight index missing in {HERE}", file=sys.stderr)
+        print(f"[lumynax] run: hf download AbteeXAILab/lumynax-speech-kokoro-82m-tts --local-dir <dir> first.", file=sys.stderr)
+        sys.exit(2)
+    print(f"[lumynax] loading from local repo {HERE}")
+    processor = AutoProcessor.from_pretrained(str(HERE), trust_remote_code=True)
+    model = AutoModelForImageTextToText.from_pretrained(
+        str(HERE), device_map="auto", torch_dtype="auto", trust_remote_code=True
+    )
+    def chat(user, img):
+        content = [{"type":"text","text":user}]
+        if img: content.insert(0, {"type":"image","url":img})
+        messages = [
+            {"role":"system","content":[{"type":"text","text":LUMYNAX_SYSTEM}]},
+            {"role":"user","content":content},
+        ]
+        inputs = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=True,
+                                               return_dict=True, return_tensors="pt").to(model.device)
+        out = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.4)
+        return processor.batch_decode(out[:, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)[0]
+    if args.interactive:
+        print("[lumynax] interactive — '/img <path>' attaches, empty line exits.")
+        pending = None
+        while True:
+            try: q = input("you> ").strip()
+            except EOFError: break
+            if not q: break
+            if q.startswith("/img "): pending = q[5:].strip(); print(f"[lumynax] attached: {pending}"); continue
+            print("lumynax> " + chat(q, pending)); pending = None
+    else:
+        print(chat(args.prompt, args.image))
+
+if __name__ == "__main__":
+    main()