import numpy as np import torch from math import inf from logitbias_json import json2logitbias from sample import ( BiasFormer, SearchException, SyntaxChecker, ) from dataclasses import dataclass from typing import Optional, List, Any from logitbias import eval as logitbias_eval, VariableStub from logitbias_parser import parse as logitbias_parse from threading import Thread from queue import Queue, Empty from contextlib import contextmanager from loader import load_by_path from lib.perf_recorder import BasePerfRecorder, PerfRecorder import time # TODO placement of description anchor token for inference! class RedundancyChecker: def __init__( self, vocab, pitch_penalty: float, duration_penalty: float, own_symbols: np.ndarray, clip_refs: np.ndarray, ctx_refs: np.ndarray, ): self.vocab = vocab self.pitch_penalty = pitch_penalty self.duration_penalty = duration_penalty self.own_symbols = own_symbols self.clip_refs = clip_refs self.ctx_refs = ctx_refs self.reset() def before_sample(self, pos, logits): if self.all_distinct: return # evaluate distinctness over [self.pos, pos) already-sampled symbols for i in range(self.pos, pos): sym = self.own_symbols[i].item() assert sym != self.vocab.pad.index if self.first_pitch_pos == -1 and self.vocab.is_pitch(sym): self.first_pitch_pos = i if self.first_pitch_pos != -1: assert self.first_pitch_pos <= i nd_syms = self.clip_refs[self.clip_refs_nd, i - self.first_pitch_pos] nd_pred = nd_syms == sym self.clip_refs_nd = self.clip_refs_nd[nd_pred] nd_syms = self.ctx_refs[self.ctx_refs_nd, i] nd_pred = nd_syms == sym self.ctx_refs_nd = self.ctx_refs_nd[nd_pred] self.pos = pos if self.ctx_refs_nd.shape[0] == 0 and self.clip_refs_nd.shape[0] == 0: self.all_distinct = True return # penalize logits matching next symbol in non-distinct clips + ctxs nd_syms = np.unique( np.concatenate( [ self.clip_refs[ self.clip_refs_nd, ( (pos - self.first_pitch_pos) if self.first_pitch_pos != -1 else 0 ), ].reshape(-1), self.ctx_refs[self.ctx_refs_nd, pos].reshape(-1), ] ) ) pitch_syms = nd_syms[ np.logical_or( np.logical_and( nd_syms >= self.vocab.pitch_least_index, nd_syms < self.vocab.pitch_greatest_index, ), np.logical_and( nd_syms >= self.vocab.polyphony_least_index, nd_syms < self.vocab.polyphony_greatest_index, ), ) ] duration_syms = nd_syms[ np.logical_and( nd_syms >= self.vocab.durations_least_index, nd_syms < self.vocab.durations_greatest_index, ) ] logits[pitch_syms] -= self.pitch_penalty logits[duration_syms] -= self.duration_penalty def reset(self): self.first_pitch_pos = -1 self.pos = 0 self.all_distinct = False self.clip_refs_nd = np.arange(0, self.clip_refs.shape[0], dtype=np.int32) self.ctx_refs_nd = np.arange(0, self.ctx_refs.shape[0], dtype=np.int32) def maybe_coerce(typ, x): if x is None: return None return typ(x) @dataclass class GenerationRequest: temperature: float stop_cond: str bias: Any = None min_polyphony: Optional[int] = None max_polyphony: Optional[int] = None dists: bool = False rest_temperature: float = None note_length_temperature: float = None cfg_scale: float = 1.5 def __post_init__(self): if self.rest_temperature is None: self.rest_temperature = self.temperature if self.note_length_temperature is None: self.note_length_temperature = self.temperature @classmethod def from_json(cls, obj): bias = obj.get("bias", None) if bias is not None: bias = json2logitbias(bias) if "stopCondition" in obj: if "count" in obj: raise RuntimeError("Cannot specify both stopCondition and count") stop_cond = str(obj["stopCondition"]) else: stop_cond = f"steps == {int(obj.get('count', 5))}" # for reverse compatibility timing_temperature = obj.get("timingTemperature", None) return cls( temperature=float(obj.get("temperature", 0.9)), stop_cond=stop_cond, bias=bias, min_polyphony=maybe_coerce(int, obj.get("minPolyphony", None)), max_polyphony=maybe_coerce(int, obj.get("maxPolyphony", None)), dists=maybe_coerce(bool, obj.get("dists", None)), rest_temperature=maybe_coerce( float, obj.get("restTemperature", timing_temperature) ), note_length_temperature=maybe_coerce( float, obj.get("noteLengthTemperature", timing_temperature) ), cfg_scale=float(obj.get("cfgScale", 1.5)), ) @dataclass class GenerationRequestCtx: stop_cond: Any bias_former: Optional[BiasFormer] logit_processors: List[Any] request: GenerationRequest idx: int cfg_idx: int redundancy_checker: Optional[RedundancyChecker] stop_pos: int = -1 step_count_pos: int = 0 in_chord: bool = False step_count: int = 0 backtrack_count: int = 0 backtrack_from_pos: int = -1 pending_search_error: bool = False backtrack_total_count: int = 0 @dataclass class SearchCtx: request_ctxts: List[GenerationRequestCtx] encoder_input_ids: Optional[torch.Tensor] encoder_attention_mask: Optional[torch.Tensor] batch_size: int start: torch.Tensor logit_masks: np.ndarray ys: np.ndarray dists: np.ndarray pos: int recorder: BasePerfRecorder start_time: float deadline: float logit_processor_time: float = 0.0 sampling_time: float = 0.0 redundancy_checker_time: float = 0.0 stop_cond_time: float = 0.0 logits0: Optional[torch.Tensor] = None logits: Any = None @contextmanager def perf_stopwatch(recorder, name): start = time.time() yield recorder.record(name, 1000.0 * (time.time() - start), "ms") class ComposerAPI: @classmethod def from_path(cls, path): vocab, model, critic = load_by_path(path) return cls(vocab, model, critic) def __init__(self, vocab, model, critic=None, backtrack_limit=10): super().__init__() self.vocab = vocab self.model = model self.critic = critic self.backtrack_limit = backtrack_limit self.max_seq_len = model.params.max_inference_seq_len # self.prompt_prefix = "tags: [analyzed aligned] " # TODO cleanup - also this must match data_aug.py self.prompt_prefix = "" self.inference_batch_size = model.params.inference_batch_size # TODO accompany becomes list of midi def generate(self, *args, recorder=PerfRecorder(), **kwargs): with torch.no_grad(), perf_stopwatch(recorder, "request"): return self.generate_impl(*args, recorder=recorder, **kwargs) def setup_request_context( self, idx: int, count: int, request: GenerationRequest, redundant_pitch_penalty: float, redundant_duration_penalty: float, redundant_clips: np.ndarray, ys: np.ndarray, start_len: int, ): stop_cond = logitbias_parse(request.stop_cond) stop_cond = logitbias_eval( stop_cond, {k: VariableStub() for k in ["steps", "error", "maxProbability", "beat"]}, ) logit_processors = [] logit_processors.append( SyntaxChecker( self.vocab, request.min_polyphony, request.max_polyphony, ) ) bias_former = None if request.bias is not None: bias_former = BiasFormer(self.vocab, request.bias) logit_processors.append(bias_former) if redundant_pitch_penalty > 0.0 or redundant_duration_penalty > 0.0: redundancy_checker = RedundancyChecker( self.vocab, redundant_pitch_penalty, redundant_duration_penalty, ys[idx, start_len:, 0], redundant_clips, ys[:idx, start_len:, 0], ) else: redundancy_checker = None return GenerationRequestCtx( stop_cond, bias_former, logit_processors, request, idx, count + idx, redundancy_checker, ) @contextmanager def forward_pass_thread(self, recorder: BasePerfRecorder): threaded_forward_q_in: Queue = Queue(1) threaded_forward_q_out: Queue = Queue() logits_bottleneck_count = 0 model_bottleneck_count = 0 model_forward_time = 0.0 def threaded_forward_entry(): nonlocal model_forward_time while True: req = threaded_forward_q_in.get() if req is None: break try: start = time.time() res = self.model.forward(**req) if callable(res): res = res() res = res.cpu().detach() model_forward_time += time.time() - start threaded_forward_q_out.put(res) except Exception as exn: # pylint: disable=broad-except threaded_forward_q_out.put(exn) threaded_forward_thd = Thread(target=threaded_forward_entry, daemon=True) threaded_forward_thd.start() def threaded_forward(**kwargs): threaded_forward_q_in.put(kwargs) def thunk(): nonlocal logits_bottleneck_count nonlocal model_bottleneck_count try: res = threaded_forward_q_out.get(block=False) logits_bottleneck_count += 1 except Empty: res = threaded_forward_q_out.get() model_bottleneck_count += 1 if isinstance(res, Exception): raise res return res return thunk try: yield threaded_forward finally: threaded_forward_q_in.put(None) threaded_forward_thd.join() recorder.record("threaded_forward", 1000.0 * model_forward_time, "msThd") recorder.record("logits_bottleneck", logits_bottleneck_count, "count") recorder.record("model_bottleneck", model_bottleneck_count, "count") def search_step(self, threaded_forward, sctx: SearchCtx): logit_processor_start = time.time() running_ctxts = [c for c in sctx.request_ctxts if c.stop_pos == -1] if len(running_ctxts) == 0: return False if sctx.pos >= sctx.ys.shape[1]: for c in running_ctxts: sctx.recorder.log( f"search reached max length {sctx.pos} for context {c.idx}" ) c.stop_pos = sctx.pos return False if time.time() >= sctx.deadline: sctx.recorder.log(f"search reached deadline at pos {sctx.pos}") for c in running_ctxts: c.stop_pos = sctx.pos return False # Evaluate logit processors for non-done contexts to get masks for ctx in running_ctxts: sctx.logit_masks[ctx.idx].fill(0.0) for logit_processor in ctx.logit_processors: logit_processor(sctx.ys[ctx.idx, : sctx.pos], sctx.logit_masks[ctx.idx]) if ctx.backtrack_from_pos != -1 and sctx.pos > ctx.backtrack_from_pos: sctx.recorder.log( f"backtrack: progress for context {ctx.idx} after {ctx.backtrack_count} attempts" ) ctx.backtrack_from_pos = -1 sctx.logit_processor_time += time.time() - logit_processor_start if callable(sctx.logits): sctx.logits = sctx.logits()[:, -1] sampling_start = time.time() # CFG if self.model.encoder is not None: for ctx in running_ctxts: logits_uncond = sctx.logits[ctx.cfg_idx] sctx.logits[ctx.idx] = logits_uncond + ctx.request.cfg_scale * ( sctx.logits[ctx.idx] - logits_uncond ) # TODO this only needs to be done for running contexts sctx.logits += torch.from_numpy(sctx.logit_masks) probs = torch.softmax(sctx.logits, dim=-1).numpy() sctx.dists[:, sctx.pos] = probs # Sample next token for non-done contexts backtracking = False backtracking_to_start = False for ctx in running_ctxts: if torch.all(sctx.logits[ctx.idx] < -1000.0): # backtrack if ctx.backtrack_from_pos == -1: sctx.recorder.log( f"backtrack: starting from {sctx.pos} for context {ctx.idx}" ) ctx.backtrack_from_pos = sctx.pos ctx.backtrack_count = 0 if ctx.backtrack_count > self.backtrack_limit: ctx.pending_search_error = True continue # skip sampling else: sctx.pos = max( sctx.start.shape[0], ctx.backtrack_from_pos - 4 * int(pow(1.5, ctx.backtrack_count)), ) backtracking = True if sctx.pos == sctx.start.shape[0]: backtracking_to_start = True ctx.backtrack_total_count += 1 ctx.backtrack_count += 1 for ctx2 in running_ctxts: ctx2.step_count_pos = 0 if ctx2.bias_former is not None: ctx2.bias_former.count = 0 if ctx2.redundancy_checker is not None: ctx2.redundancy_checker.reset() sctx.ys[ctx2.idx, sctx.pos + 1 :, 0] = 0 break if ctx.redundancy_checker is not None: start = time.time() ctx.redundancy_checker.before_sample( sctx.pos - sctx.start.shape[0], sctx.logits[ctx.idx] ) sctx.redundancy_checker_time += time.time() - start temperature = ctx.request.temperature if self.vocab.is_rest(sctx.ys[ctx.idx, sctx.pos - 1, 0]): temperature = ctx.request.rest_temperature if self.vocab.is_pitch(sctx.ys[ctx.idx, sctx.pos - 1, 0]): temperature = ctx.request.note_length_temperature temperature = float(temperature) if temperature == 0.0: next_token = torch.argmax(sctx.logits[ctx.idx]) else: next_token = torch.multinomial( torch.softmax(sctx.logits[ctx.idx] / temperature, dim=-1), num_samples=1, ) self.vocab.append_symbols( sctx.ys[ctx.idx : ctx.idx + 1], next_token.view(1, 1), strict=False, seq_len=sctx.pos, ) if self.model.encoder is not None: sctx.ys[ctx.cfg_idx, sctx.pos] = sctx.ys[ctx.idx, sctx.pos] sctx.sampling_time += time.time() - sampling_start if backtracking_to_start: sctx.logits = sctx.logits0 return True # Kick off model forward pass -> logits sctx.logits = threaded_forward( x=torch.from_numpy(sctx.ys[:, sctx.pos : sctx.pos + 1]), start_pos=sctx.pos, encoder_input_ids=None, encoder_attention_mask=None, ) stop_cond_start = time.time() # Evaluate stop conditions for non-done contexts if not backtracking: for ctx in running_ctxts: if ctx.step_count_pos == 0: ctx.step_count = 0 ctx.in_chord = False for j in range(sctx.start.shape[0] + ctx.step_count_pos, sctx.pos + 1): sym = sctx.ys[ctx.idx, j, 0].item() if self.vocab.is_pitch(sym) and not ctx.in_chord: ctx.in_chord = True if self.vocab.is_rest(sym) and ctx.in_chord: ctx.in_chord = False ctx.step_count += 1 ctx.step_count_pos = max(sctx.pos + 1 - sctx.start.shape[0], 0) if ctx.bias_former is not None: ctx.bias_former.count = ctx.step_count if self.vocab.embedding_time_wrapped(sctx.ys[ctx.idx, sctx.pos]): sctx.recorder.log(f"context {ctx.idx} wrapped around") ctx.pending_search_error = True max_probability = np.max(probs[ctx.idx]) env = { "steps": ctx.step_count, "error": ctx.pending_search_error, "maxProbability": max_probability.item(), "beat": self.vocab.get_tensor_subbeat( sctx.ys[ctx.idx, sctx.pos - 1], next=True ).item() / self.vocab.quantize_divisions, "generationTime": time.time() - sctx.start_time, } if logitbias_eval(ctx.stop_cond, env): sctx.recorder.log( f"stop: context {ctx.idx} at pos {sctx.pos} with env {env}" ) ctx.stop_pos = sctx.pos elif ctx.pending_search_error: raise SearchException( f"search failed for context {ctx.idx} at pos {sctx.pos}" ) sctx.stop_cond_time += time.time() - stop_cond_start sctx.pos += 1 return True def generate_ctx(self, sctx: SearchCtx): with self.forward_pass_thread(sctx.recorder) as threaded_forward: # common forward pass with perf_stopwatch(sctx.recorder, "common_forward"): rearranger = [0] * len(sctx.request_ctxts) if self.model.encoder is not None: rearranger += [1] * len(sctx.request_ctxts) rearranger = torch.tensor(rearranger, dtype=torch.long) self.model.hint_inference_batch_size(sctx.batch_size) sctx.logits0 = ( self.model.forward( x=torch.from_numpy(sctx.ys[0, : sctx.start.shape[0]]) .unsqueeze(0) .repeat(2 if self.model.encoder is not None else 1, 1, 1), start_pos=0, encoder_input_ids=sctx.encoder_input_ids, encoder_attention_mask=sctx.encoder_attention_mask, )[:, -1] .cpu() .detach()[rearranger] ) self.model.rearrange(rearranger) sctx.logits = sctx.logits0 with perf_stopwatch(sctx.recorder, "search"): while self.search_step(threaded_forward, sctx): pass def generate_impl( self, prefix, generation_requests, recorder, accompany=None, text_prompt=None, deadline=inf, redundant_pitch_penalty=0.0, redundant_duration_penalty=0.0, redundant_clips=[], ): start_time = time.time() assert redundant_pitch_penalty >= 0.0 and redundant_duration_penalty >= 0.0 if self.model.encoder is None: assert text_prompt is None encoder_input_ids, encoder_attention_mask = None, None else: with perf_stopwatch(recorder, "tokenize"): encoder_input_ids, encoder_attention_mask = self.model.encoder.tokenize( [ self.prompt_prefix + (text_prompt if text_prompt is not None else ""), "", # for CFG ] ) with perf_stopwatch(recorder, "midi_to_tensor"): start = self.vocab.midi_to_tensor( prefix, accompany=accompany, rest_end=True ) with perf_stopwatch(recorder, "redundant_clips_to_tensor"): redundant_clip_tensor = np.zeros( (len(redundant_clips), self.max_seq_len), dtype=np.int32 ) for i, clip in enumerate(redundant_clips): clip_tensor = self.vocab.midi_to_tensor(clip) for j in range(clip_tensor.shape[0]): if self.vocab.is_pitch(clip_tensor[j, 0].item()): insert_len = min(self.max_seq_len, clip_tensor.shape[0] - j) redundant_clip_tensor[i, :insert_len] = clip_tensor[ j : insert_len + j, 0 ] break recorder.record("start", start.shape[0], "count") if start.shape[0] >= self.max_seq_len - 4: raise RuntimeError("input too long for model") start_midi, _ = self.vocab.tensor_to_midi(start, strict=False) ctxts = [] batch_size_per_request = 1 if self.model.encoder is None else 2 if ( batch_size_per_request * len(generation_requests) > self.inference_batch_size ): raise RuntimeError("too many generation requests") if self.critic is None: runs_per_request = 1 else: runs_per_request = self.inference_batch_size // ( batch_size_per_request * len(generation_requests) ) num_ctxs = len(generation_requests) * runs_per_request batch_size = num_ctxs * batch_size_per_request with perf_stopwatch(recorder, "alloc"): # allocate working tensors ys = np.zeros( (batch_size, self.max_seq_len, 9), dtype=np.int32, ) dists = np.zeros( (batch_size, self.max_seq_len, self.vocab.N), dtype=np.float32, ) # faster for complex indexing logit_masks = np.zeros((batch_size, self.vocab.N), dtype=np.float32) ys[:, : start.shape[0]] = start.unsqueeze(0) for i, request in enumerate(generation_requests): with perf_stopwatch(recorder, f"setup_context[{i}]"): for j in range(runs_per_request): ctxts.append( self.setup_request_context( i * runs_per_request + j, num_ctxs, request, redundant_pitch_penalty, redundant_duration_penalty, redundant_clip_tensor, ys, start.shape[0], ) ) search_ctx = SearchCtx( batch_size=batch_size, request_ctxts=ctxts, encoder_input_ids=encoder_input_ids, encoder_attention_mask=encoder_attention_mask, start=start, logit_masks=logit_masks, ys=ys, dists=dists, pos=start.shape[0], recorder=recorder, start_time=start_time, deadline=deadline, ) self.generate_ctx(search_ctx) recorder.record( "logit_processors", 1000.0 * search_ctx.logit_processor_time, "msThd" ) recorder.record("stop_cond", 1000.0 * search_ctx.stop_cond_time, "msThd") recorder.record("sampling", 1000.0 * search_ctx.sampling_time, "ms") recorder.record( "redundancy_checker", 1000.0 * search_ctx.redundancy_checker_time, "ms" ) recorder.record("steps", search_ctx.pos - start.shape[0], "count") if self.critic is None: scores = torch.tensor([0.0] * len(ctxts)) else: with perf_stopwatch(search_ctx.recorder, "embeddings"): max_stop_pos = max(ctx.stop_pos for ctx in ctxts) embeddings = ( self.model.forward( x=torch.from_numpy(search_ctx.ys[:, : max_stop_pos + 1]), start_pos=0, encoder_input_ids=search_ctx.encoder_input_ids, encoder_attention_mask=search_ctx.encoder_attention_mask, return_embeddings=True, ) .cpu() .detach() ) scores = self.critic(embeddings).cpu().detach().view(-1) with perf_stopwatch(recorder, "tensor_to_midi"): res = [] for i, req in enumerate(generation_requests): group = [c for c in ctxts if id(c.request) == id(req)] recorder.record( f"backtracks[{i}]", sum(c.backtrack_total_count for c in group), "count", ) best_ctx = max(group, key=lambda c: scores[c.idx]) midi, _ = self.vocab.tensor_to_midi( ys[best_ctx.idx, : best_ctx.stop_pos + 1], strict=False, dists=( dists[best_ctx.idx, start.shape[0] : best_ctx.stop_pos + 1] if req.dists else None ), ) res.append(midi[0 if start_midi is None else len(start_midi) :]) return res