Create README.md

README.md

@@ -0,0 +1,153 @@
+---
+base_model:
+- HKUSTAudio/Llasa-3B
+---
+
+# Sample Inference Script
+```py
+import random
+import re
+from argparse import ArgumentParser
+
+import torch
+import torchaudio
+from exllamav2 import (
+    ExLlamaV2,
+    ExLlamaV2Cache,
+    ExLlamaV2Config,
+    ExLlamaV2Tokenizer,
+    Timer,
+)
+from exllamav2.generator import (
+    ExLlamaV2DynamicGenerator,
+    ExLlamaV2DynamicJob,
+    ExLlamaV2Sampler,
+)
+from rich import print
+from torchaudio import functional as F
+from xcodec2.modeling_xcodec2 import XCodec2Model
+
+parser = ArgumentParser()
+parser.add_argument("-m", "--model", required=True)
+parser.add_argument("-v", "--vocoder", required=True)
+parser.add_argument("-i", "--input", required=True)
+parser.add_argument("-a", "--audio", default="")
+parser.add_argument("-t", "--transcript", default="")
+parser.add_argument("-o", "--output", default="output.wav")
+parser.add_argument("-d", "--debug", action="store_true")
+parser.add_argument("--max_seq_len", type=int, default=2048)
+parser.add_argument("--sample_rate", type=int, default=16000)
+parser.add_argument("--seed", type=int, default=None)
+parser.add_argument("--temperature", type=float, default=0.8)
+parser.add_argument("--top_p", type=float, default=1.0)
+args = parser.parse_args()
+
+with Timer() as timer:
+    config = ExLlamaV2Config(args.model)
+    config.max_seq_len = args.max_seq_len
+    model = ExLlamaV2(config, lazy_load=True)
+    cache = ExLlamaV2Cache(model, lazy=True)
+    model.load_autosplit(cache, progress=True)
+    tokenizer = ExLlamaV2Tokenizer(config)
+    generator = ExLlamaV2DynamicGenerator(model, cache, tokenizer)
+
+print(f"Loaded model in {timer.interval:.2f} seconds.")
+
+with Timer() as timer:
+    vocoder = XCodec2Model.from_pretrained(args.vocoder)
+    vocoder = vocoder.cuda().eval()
+
+print(f"Loaded vocoder in {timer.interval:.2f} seconds.")
+
+if args.audio and args.transcript:
+    with Timer() as timer:
+        transcript = f"{args.transcript} "
+        audio, sample_rate = torchaudio.load(args.audio)
+        audio = audio.cuda()
+
+        if audio.shape[0] > 1:
+            audio = torch.mean(audio, dim=0, keepdim=True)
+
+        if sample_rate != args.sample_rate:
+            audio = F.resample(audio, sample_rate, args.sample_rate)
+
+    print(f"Loaded audio in {timer.interval:.2f} seconds.")
+
+    with Timer() as timer:
+        audio = vocoder.encode_code(audio)
+        audio = audio[0, 0, :]
+        audio = [f"<|s_{a}|>" for a in audio]
+        audio = "".join(audio)
+
+    print(f"Encoded audio in {timer.interval:.2f} seconds.")
+else:
+    transcript = ""
+    audio = ""
+
+with Timer() as timer:
+    input = (
+        "<|start_header_id|>user<|end_header_id|>\n\n"
+        "Convert the text to speech:"
+        "<|TEXT_UNDERSTANDING_START|>"
+        f"{transcript}{args.input}"
+        "<|TEXT_UNDERSTANDING_END|>"
+        "<|eot_id|>\n"
+        "<|start_header_id|>assistant<|end_header_id|>\n\n"
+        "<|SPEECH_GENERATION_START|>"
+        f"{audio}"
+    )
+
+    input_ids = tokenizer.encode(input, add_bos=True, encode_special_tokens=True)
+
+print(f"Encoded input in {timer.interval:.2f} seconds.")
+
+with Timer() as timer:
+    max_new_tokens = config.max_seq_len - input_ids.shape[-1]
+    gen_settings = ExLlamaV2Sampler.Settings()
+    gen_settings.temperature = args.temperature
+    gen_settings.top_p = args.top_p
+    seed = args.seed if args.seed else random.randint(0, 2**64 - 1)
+    stop_conditions = ["<|SPEECH_GENERATION_END|>"]
+
+    job = ExLlamaV2DynamicJob(
+        input_ids=input_ids,
+        max_new_tokens=max_new_tokens,
+        gen_settings=gen_settings,
+        seed=seed,
+        stop_conditions=stop_conditions,
+        decode_special_tokens=True,
+    )
+
+    generator.enqueue(job)
+    output = []
+
+    while generator.num_remaining_jobs():
+        for result in generator.iterate():
+            if result.get("stage") == "streaming":
+                text = result.get("text", "")
+                output.append(text)
+
+                if args.debug:
+                    print(text, end="", flush=True)
+
+            if result.get("eos"):
+                generator.clear_queue()
+
+                if args.debug:
+                    print()
+
+print(
+    f"Generated {len(output)} tokens with seed {seed} in {timer.interval:.2f} seconds."
+)
+
+with Timer() as timer:
+    output = "".join(output)
+    output = [int(o) for o in re.findall(r"<\|s_(\d+)\|>", output)]
+    output = torch.tensor([[output]]).cuda()
+    output = vocoder.decode_code(output)
+    output = output[0, 0, :]
+    output = output.unsqueeze(0).cpu()
+    torchaudio.save(args.output, output, args.sample_rate)
+
+print(f"Decoded audio in {timer.interval:.2f} seconds.")
+```