o oiH @sdZddlZddlmZddlmZmZmZddlm Z m Z ddl m Z ddl mZddlmZddlmZdd lmZmZmZmZdd lmZ dd ejdedededejf ddZGdddejZGdddeZ GdddeeZ!dS)a Custom LFM2 implementation with KaniTTS-2 frame-level position encoding. Key Innovation: - Frame-level position encoding: All 4 tokens within an audio frame share the same position ID This reduces RoPE distance between tokens across frames, improving long-form generation. Compatible with Flash Attention 2 for 10-20x training speedup. FIXED: Proper frame-level position tracking during generation with KV-cache. N)OptionalUnionTuple)CausalLMOutputWithPastBaseModelOutputWithPast)TransformersKwargs)Unpack)Cache)GenerationMixin) Lfm2ModelLfm2ForCausalLMLfm2PreTrainedModelLfm2HybridConvCache) Lfm2Config? input_idsaudio_tokens_starttokens_per_frame audio_stepreturncCs|j\}}|j}||k}|}tj|d|tjd} tj| |gddjddddddf} tj| |gddjddddddf} | |} | | |} | S)a Vectorized computation of frame-level position IDs (10-50x faster than Python loops). Key insight: Use cumulative counts to determine positions. - Text tokens: sequential positions (step 1.0) - Audio tokens: frame-level positions (step audio_step per frame) Args: input_ids: Input token IDs [batch_size, seq_len] audio_tokens_start: Token ID where audio tokens begin (typically 64410) tokens_per_frame: Number of tokens per audio frame (typically 4) audio_step: Position step size per audio frame (default 1.0). Set to < 1.0 (e.g., 0.5) to compress audio position space. Returns: position_ids: Position IDs [batch_size, seq_len]. if audio_step is float, returns FloatTensor. Example: >>> input_ids = torch.tensor([[100, 200, 64410, 68442, 72474, 76506, 300]]) >>> # Tokens: [text, text, aud0, aud1, aud2, aud3, text] >>> pos = compute_frame_level_positions(input_ids, 64410, 4, audio_step=0.5) >>> pos tensor([[0., 1., 2., 2., 2., 2., 3.]]) # Text at 0, 1. Audio frame at 2. Next text at 3 (1+1+1?) # Note: Text logic accumulates 1 per text token. # Audio logic accumulates audio_step per frame. devicedtypedimN)shapertorchzeroslongcatcumsum)rrrr batch_sizeseq_lenris_audio text_maskr text_countaudio_token_countaudio_frame_count position_idsr,B/home/ubuntu/.local/lib/python3.10/site-packages/kani_tts/model.pycompute_frame_level_positions s #.. r.csHeZdZdZ   d fdd ZeddZed d Zd d ZZ S)LearnableRotaryEmbeddingu Learnable RoPE with layer-wise frequency scaling. Each layer has a learnable alpha parameter that scales the RoPE frequencies: θᵢ^(l) = α^(l) · base^(-2i/d) where α^(l) is constrained to [alpha_min, alpha_max] via sigmoid reparameterization: α^(l) = alpha_min + (alpha_max - alpha_min) · sigmoid(w^(l)) This allows the model to learn optimal positional encoding frequencies per layer. 皙?@Nc st||_||_||_||_|j|j}|j}|j } ||_ ||_ | |_ d|t j d|dt jdj|t jd|} |jd| ddtt d |_dS) Nrrrr inv_freq_baseF) persistentg)super__init__ layer_idxtotal_attention_layers alpha_min alpha_max hidden_sizenum_attention_heads rope_thetamax_position_embeddingsrbaserarangeint64tofloatregister_buffernn Parametertensor alpha_weight) selfconfigr8r9r:r;rrr@r?r4 __class__r,r-r7ls  &z!LearnableRotaryEmbedding.__init__cCs|j|j|jt|jS)z Compute constrained alpha via sigmoid reparameterization. Returns: Scalar alpha value in range [alpha_min, alpha_max] )r:r;rsigmoidrIrJr,r,r-alphaszLearnableRotaryEmbedding.alphacCs |j|jS)u Compute scaled inverse frequencies: α^(l) · θᵢ Returns: Tensor of shape [d/2] with scaled frequencies )r4rPrOr,r,r-inv_freqs z!LearnableRotaryEmbedding.inv_freqc Cs|jddddf|jddd}|dddddf}|jj}t|tr0|dkr0|nd}tj |dd%|| dd }tj ||fdd }| }| } Wdn1sawY|j|jd | j|jd fS) aX Compute rotary position embeddings for the input. Args: x: Input tensor of shape [batch_size, num_heads, seq_len, head_dim] position_ids: Position indices of shape [batch_size, seq_len] Returns: Tuple of (cos, sin) embeddings, each of shape [batch_size, seq_len, head_dim] NrrrmpscpuF) device_typeenabledr2rr3)rQrDexpandrrtype isinstancestrrautocast transposer"cossinrCr) rJxr+inv_freq_expandedposition_ids_expandedrTfreqsembr\r]r,r,r-forwards(  z LearnableRotaryEmbedding.forward)r0r1N) __name__ __module__ __qualname____doc__r7propertyrPrQrc __classcell__r,r,rLr-r/_s!   r/cseZdZdZ      dded ed ed ed ed ededeffdd Z        d de e j de e j de e j de e de e jde e jde ede e j deedeffdd ZZS)!Lfm2ForKaniModelz Custom LFM2 model with KaniTTS-2 frame-level position encoding. This version only overrides position ID computation - everything else uses the standard Lfm2Model implementation. rrFr0r1rKrrruse_learnable_roper:r;speaker_emb_dimc  st|||_||_||_||_||_||_||_t j ||j dd|_ |rg} t |drBt|jD] \} } | dkr@| | q3ntt|j} t| } t |_t|jD]%} | | vrvt|| | ||t |drk|jndd} |j| qW|jdqWtdtd|td |td |d |j td td | dtd|d|ddSd|_tdtd|td |td|td |d |j td tddS)NFbias layer_typesfull_attentionr)rKr8r9r:r;ru!✅ Lfm2ForKaniModel initialized:z - Audio tokens start: z - Tokens per frame: z - Speaker embedding: z -> z4 - Using frame-level position encoding (KaniTTS-2)z - Learnable RoPE ENABLED for z attention layersz - Alpha range: [z, ]z - Audio step: z, - Learnable RoPE DISABLED (standard RoPE))r6r7rrrrlr:r;rmrFLinearr<speaker_emb_projectionhasattr enumeraterpappendlistrangenum_hidden_layerslen ModuleListlearnable_rope_layersr/rprint)rJrKrrrrlr:r;rmattention_layer_indicesidx layer_typer9learnable_roperLr,r-r7s^     zLfm2ForKaniModel.__init__Nrattention_maskr+past_key_values inputs_embeds speaker_emb use_cachecache_positionkwargsrc  s|dur|durt||j|j|jd}|js'tjd |||||||d| Sddlm} m } |dur;|dur;t d|durD| |}|rZ|durZ|j d} | |j | |j|jd}|durx|durf|nd} |j d}tj| | ||jd }|dur|d}| |j |||||d }|}d}t|jd|j jD],\}}|j|dur|j|||}n |dur|||}||f|||||d | }q||}t||d S)a^ Forward pass with custom frame-level position IDs and speaker embeddings. Speaker embeddings are inserted at position 1 (after the first token). All subsequent positions are shifted by +1 to maintain sequential ordering. If learnable RoPE is disabled: Delegates to parent class after computing frame-level position IDs. If learnable RoPE is enabled: Overrides position embedding computation to use per-layer learnable RoPE. Args: speaker_emb: Speaker embeddings [batch_size, speaker_emb_dim] (e.g., [1, 128]) Nrrrr)rrr+rrrrr)rcreate_causal_maskz;You must specify at least one of input_ids or inputs_embeds)rKmax_batch_sizerrrr)rK input_embedsrrrr+)rr+rrposition_embeddings)last_hidden_staterr,)r.rrrrlr6rc&transformers.models.lfm2.modeling_lfm2rr ValueError embed_tokensrrKrrget_seq_lengthrrA unsqueezervlayersrzr}pos_embembedding_normr)rJrrr+rrrrrrrrr$past_seen_tokens seq_length causal_mask hidden_statesrr8 decoder_layerrLr,r-rc s          zLfm2ForKaniModel.forwardrrFr0r1rk)NNNNNNNN)rdrerfrgrintrDboolr7rr LongTensorTensorr FloatTensorrrrrcrir,r,rLr-rjsl  I   rjcseZdZdZdgZedefddZ    d?d ed e de de dede de de ffdd Z d@de e fddZ de fddZ          dAde ejde ejd e ejd!e ed"e ejd#e ejd$e ejd%e ed&e ejd'ee ejfd(eedefd)d*Z      +dBd,d-Zfd.d/Ze       dCd0ed e de de dede de de fd1d2Zd3d4Zd5d6Zd7d8Zd9d:Z d;d<Z!d=d>Z"Z#S)DKaniTTS2ForCausalLMa Flash Attention compatible LFM2 for causal language modeling with KaniTTS-2 frame-level positions. Features: - Frame-level position encoding for audio tokens (KaniTTS-2 innovation) - Optional learnable RoPE with per-layer frequency scaling (alpha parameters) - Proper position tracking during generation with KV-cache - Flash Attention 2 optimized - Standard causal masking - Compatible with existing KaniTTS inference pipeline - Includes GenerationMixin for text generation capabilities lm_head.weightrcCsdS)NFr,)clsr,r,r-_supports_default_dynamic_cachesz3KaniTTS2ForCausalLM._supports_default_dynamic_cacherrFr0r1rkrKrrrrlr:r;rmc st|t||||||||d|_|j|_tj|j|jdd|_||_ ||_ ||_ ||_ ||_ ||_||_d|_d|_t|drE|jnd|_d|_|dS)N)rrlr:r;rmFrngeneration_configr)r6r7rjmodel vocab_sizerFrsr<lm_headrrrrlr:r;rm_generation_state_current_speaker_embrurmain_input_name post_init) rJrKrrrrlr:r;rmrLr,r-r7s2   zKaniTTS2ForCausalLM.__init__Nstarting_frame_positioncCs d|dur t|ndd|_dS)a Reset generation state before starting new generation. This tracks: - The position where the first audio frame should start - How many audio tokens we've generated - What the current frame position should be Args: starting_frame_position: The position ID where the first audio frame begins. If None, will be determined when first audio token is generated. rN)audio_tokens_generatedcurrent_frame_position)rDr)rJrr,r,r-_reset_generation_states z+KaniTTS2ForCausalLM._reset_generation_state new_token_idcCs2|jdurdS||jkr|jdd7<dSdS)z Update generation state after generating a token. Args: new_token_id: The token ID that was just generated Nrr)rr)rJrr,r,r-_update_generation_states  z,KaniTTS2ForCausalLM._update_generation_staterrrr+rrrlabelsrrlogits_to_keeprc  Ks|jd|||||||| d| } | j} t| trt| dn| }|| dd|ddf}d}|durC|jd|||jjd| }t ||| j | j | j dS)z Forward pass - delegates to flash-compatible model with frame-level position encoding. Args: speaker_emb: Speaker embeddings [batch_size, speaker_emb_dim] (e.g., [1, 128]) )rrr+rrrrrN)logitsrr)lossrrr attentionsr,) rrrXrslicer loss_functionrKrrrrr)rJrrr+rrrrrrrroutputsr slice_indicesrrr,r,r-rcs2  zKaniTTS2ForCausalLM.forwardTcKsR|dur|jdur|j|}|j|j} | d} tj|ddddddf| |ddddddfgdd}|duretj|ddddftj|jdd|j |j d|ddddfgdd}|durtj|dd|ddd|dddgdd}d}|dur t |t t fr|} | } nt|dkr|ddjdnd} } |dur|dur|jd|jdkr|dd|jd|  df}n;| |jdkr|dd| df}n)| |jdkr|ddddf}n|dur | |jdkr |dd| df}|dur=|dur|nd} |dur#|jdn|jd} |dur0|j n|j } tj| | | | d}|dur|dur|jdur|durU|j n|j } |d }||jkri|}n5|jd durz|}||jd <|jd |j}|dkr|jd dkr|jd |j7<|jd }t |trtj|gg| tjd}n tj|gg| tjd}||n<|dur|jdur|jd}|j } tj|| dd}n|durt||j|j|jd }|dur|r|jdd |||||d}|dur|dur|jdus||d<|dur'|dur'||d<|S)z Prepare inputs for generation with proper frame-level position encoding. CRITICAL FIX: Maintains frame-level positions during generation with KV-cache. Nrrrrr2rr)rrrrr)r)rr+rrrrr)rrrrtrrr"onesrrrrXr rrr{rAritemrrrrDrHr!rr.r)rJrrrrrr+rrspeaker_emb_projected cache_length past_lengthrr current_tokenposfirst_frame_postoken_in_framer% model_inputsr,r,r-prepare_inputs_for_generations    &"              z1KaniTTS2ForCausalLM.prepare_inputs_for_generationcsL|dd}d|_||_ztj|i|}Wd|_d|_|Sd|_d|_w)a6 Override generate to reset state before generation. This ensures frame-level position tracking starts fresh for each generation call. Also handles speaker embeddings if provided. Args: speaker_emb: Optional speaker embedding [batch_size, speaker_emb_dim] rN)poprrr6generate)rJargsrrresultrLr,r-rs zKaniTTS2ForCausalLM.generatepretrained_model_name_or_pathc  Osdd| D} ddlm} | j|fi| } |dur*t| dd}|dur*td|dur4t| dd }|dur>t| d d }|durHt| d d }|durRt| dd}|dur\t| dd}|durft| dd}|| |||||||d}|rddlm}ddlm }ddl }|j |r|j |d}n||dd}||}|j|d d\}}d|vrd|vr|jjj|j_dd|D}|rtdt|t|dkr|D] }td |q|rtd!t|dd"lm}z ||}||_Wnty||_Ynw| d#d$}|d$krtjrd%nd&}||}n-tj|fi| }|jj|jd d|j|jt|d'rE|j|_||j }td(||S))a Load a pretrained LFM2 model with KaniTTS-2 flash-compatible implementation. Args: pretrained_model_name_or_path: HuggingFace model ID or local path audio_tokens_start: Token ID where audio tokens begin (typically 64410). If None, reads from model config. tokens_per_frame: Number of tokens per audio frame (default: 4). If None, reads from model config. audio_step: Step size per audio frame (default: 1.0). Use 0.5 for new models. If None, reads from model config. use_learnable_rope: Enable learnable RoPE with per-layer alpha (default: False). If None, reads from model config. alpha_min: Minimum alpha value for learnable RoPE (default: 0.1). If None, reads from model config. alpha_max: Maximum alpha value for learnable RoPE (default: 2.0). If None, reads from model config. speaker_emb_dim: Dimension of speaker embeddings (default: 128). If None, reads from model config. cSsi|] \}}|dvr||qS))rlr:r;rmr,).0kvr,r,r- s z7KaniTTS2ForCausalLM.from_pretrained..r) AutoConfigNrzaudio_tokens_start not provided and not found in model config. Please specify audio_tokens_start explicitly or add it to the model's config.jsonrrrrrlFr:r0r;r1rmrk)rKrrrrlr:r;rm) load_file)hf_hub_downloadzmodel.safetensors)repo_idfilename)strictrzmodel.embed_tokens.weightcSsg|]}|dkr|qS)rr,)rrr,r,r- Lsz7KaniTTS2ForCausalLM.from_pretrained..u: ⚠️ Missing keys (will use random initialization): z - u& ⚠️ Unexpected keys (ignored): )GenerationConfig device_mapautocudarSru✅ Model loaded from )!items transformersrfrom_pretrainedgetattrrsafetensors.torchrhuggingface_hubrospathisdirjoinload_state_dictrrweightrr~r{rr Exceptiongetrr is_availablerCr state_dictrur)rrrrrrlr:r;rm model_argsr base_kwargsrrKrrrrsafetensors_pathr missing_keysunexpected_keyskeyrrrr base_modelr,r,r-rs#                     z#KaniTTS2ForCausalLM.from_pretrainedcCs|jjSRequired by GenerationMixin.rrrOr,r,r-get_input_embeddingsysz(KaniTTS2ForCausalLM.get_input_embeddingscCs ||j_dSrNr)rJvaluer,r,r-set_input_embeddings}s z(KaniTTS2ForCausalLM.set_input_embeddingscC|jSrrrOr,r,r-get_output_embeddingsz)KaniTTS2ForCausalLM.get_output_embeddingscC ||_dSrr)rJnew_embeddingsr,r,r-set_output_embeddings z)KaniTTS2ForCausalLM.set_output_embeddingscCrrr)rJdecoderr,r,r- set_decoderrzKaniTTS2ForCausalLM.set_decodercCrrrrOr,r,r- get_decoderrzKaniTTS2ForCausalLM.get_decoderr)N) NNNNNNNNNr)NNNNNT)NNNNNNN)$rdrerfrg_tied_weights_keys classmethodrrrrrDr7rrrrrrr rrrrrrcrrrYrrrrrrrrir,r,rLr-rs   0     4  * r)rr)"rgrtorch.nnrFtypingrrrtransformers.modeling_outputsrrtransformers.utilsrtransformers.processing_utilsrtransformers.cache_utilsr transformers.generation.utilsr rr r r r+transformers.models.lfm2.configuration_lfm2rrrrDr.Moduler/rjrr,r,r,r-s8       ?^I