    @torch.no_grad()
    def get_ref_logps(self, text_ids, mel, sampled_classes):
        """
        Get log probabilities from the reference model for the sampled classes
        """
        B = text_ids.shape[0]
        K = self.num_pre_samples
        with torch.no_grad():
            ref_logits = self.ref_model(text_ids=text_ids, mel=mel)[:, -1, :]
            ref_logits = ref_logits.unsqueeze(1).repeat(1, K, 1).view(B*K, -1)
            ref_log_probs = F.log_softmax(ref_logits, dim=-1)
            ref_logps = torch.gather(
                ref_log_probs,
                dim=-1,
                index=sampled_classes.unsqueeze(-1)
            ).squeeze(-1)
        return ref_logps

    @torch.no_grad()
    def generate_duration_samples(self, batch_inputs):
        """
        Generate multiple duration predictions from the model for each input
        and evaluate them using the inference function and reward model

        Args:
            batch_inputs: Dictionary with text, prompt audio, etc.

        Returns:
            Dictionary with duration samples, rewards, and reference logits
        """

        if self.rank == 0:
            print("Generating duration samples...")

        # all_logits = []
        all_text_ids = []
        all_mels = []
        all_sampled_classes = []
        all_durations = []
        all_rewards = []
        all_gen_logps = []

        all_ctc_loss = []
        all_sv_loss = []

        # Fetch batch inputs
        # prompt_mel = batch_inputs['mel'].permute(0, 2, 1).to(self.device)
        prompt_mel = batch_inputs['mel'].permute(0, 2, 1) # (B, T, 100)
        prompt_text = batch_inputs['text']

        batch_size = prompt_mel.shape[0]

        # Shift text to unpair 'mel' and 'text'; The shifted text will be synthesized
        target_text = batch_inputs['target_text']
        target_text_lengths = torch.LongTensor([len(t) for t in target_text]).to(prompt_mel.device)
        try:
            full_text = [prompt+[' ']+target for prompt, target in zip(prompt_text, target_text)]
        except:
            target_text = [batch_inputs['text'][-1]] + batch_inputs['text'][:-1]
            target_text_lengths = batch_inputs['text_lengths'].clone().roll(1, 0)
            full_text = [prompt+[' ']+target for prompt, target in zip(prompt_text, target_text)]

        # Goes to reward model
        target_text_ids = list_str_to_idx(target_text, self.vocab_char_map).to(self.accelerator.device)    # to device, the dataloader only gives list

        # Goes to duration model and TTS
        full_text_ids = list_str_to_idx(full_text, self.vocab_char_map).to(self.accelerator.device)

        # Deepcopy to separate text_ids for SLP and TTS
        slp_text_ids = full_text_ids.detach().clone()
        slp_text_ids = slp_text_ids.masked_fill(slp_text_ids==-1, self.vocab_size) # (B, L)

        # Pre-compute duration logits
        K = self.num_pre_samples
        B, T, _ = prompt_mel.shape
        _, L = slp_text_ids.shape
        # prompt_mel_k_repeats = prompt_mel.unsqueeze(1).repeat(1, K, 1, 1)  # (B, K, T, 100)
        # slp_text_ids_k_repeats = slp_text_ids.unsqueeze(1).repeat(1, K, 1)  # (B, K, L)

        # Run model once for B inputs
        old_logits = self.model(
            text_ids=slp_text_ids, # (B, L)
            mel=prompt_mel         # (B, T, 100)
        )[:, -1, :]  # (B, n_class)

        # Repeat each result K times along batch dimension
        old_logits = old_logits.unsqueeze(1).repeat(1, K, 1) # (B, K, n_class)
        # logits_nograd = logits_grad.detach().clone().view(B, K, -1) # (B, K, n_class)

        for _full_text_ids, _target_text_ids, _target_text_lengths, \
            _prompt_mel, _old_logits in zip(
                full_text_ids, target_text_ids, target_text_lengths,
                prompt_mel, old_logits
            ):

            duration_sample = F.gumbel_softmax(_old_logits, tau=self.gumbel_tau, hard=True, dim=-1)
            duration2frames = torch.arange(self.n_class).float().to(self.accelerator.device) * self.n_frame_per_class
            est_frames = (duration_sample * duration2frames).sum(-1) # (K, )

            # Compute log probabilities of the samples
            sampled_classes = duration_sample.argmax(dim=-1)
            log_probs = F.log_softmax(_old_logits, dim=-1)
            gen_logps = torch.gather(
                log_probs,
                dim=-1,
                index=sampled_classes.unsqueeze(-1)
            ).squeeze(-1)  # Shape: [K, n_class]

            # Generate speech using the sampled durations
            sampled_rewards = []

            for i in range(K):
                cur_duration = est_frames[i]
                if cur_duration == 0:
                    cur_duration = cur_duration + 50 # prevent 0 duration
                infer_full_text_ids = _full_text_ids.unsqueeze(0)
                infer_prompt_mel = _prompt_mel.unsqueeze(0)
                cur_duration = cur_duration.unsqueeze(0)
                infer_target_text_ids = _target_text_ids.unsqueeze(0)
                infer_target_text_lengths = _target_text_lengths.unsqueeze(0)
                with torch.inference_mode():
                    try:
                        _est_mel = self.inference_fn(
                            full_text_ids=infer_full_text_ids,
                            prompt_mel=infer_prompt_mel,
                            target_duration=cur_duration,
                            teacher_steps=0
                        )
                        _est_mel = _est_mel.permute(0, 2, 1) # (1, T, 100)

                        loss_dict = self.reward_fn(
                            prompt_mel=infer_prompt_mel,
                            est_mel=_est_mel,
                            target_text_id=infer_target_text_ids,
                            target_text_length=infer_target_text_lengths
                        )
                        # #TODO reweight the loss for reward
                        reward_sim = loss_dict['loss_sim'] # 0 to 1
                        reward_ctc = loss_dict['loss_ctc']
                        reward = -(reward_ctc + reward_sim * 3)
                        all_ctc_loss.append(reward_ctc)
                        all_sv_loss.append(reward_sim)
                    except Exception as e:
                        if self.rank == 0:
                            print(f"Error in speech synthesis: {e}")
                        reward = torch.tensor(-1.0).to(cur_duration.device)

                    sampled_rewards.append(reward)
                        # list with length of K
            sampled_rewards = torch.stack(sampled_rewards)  # (K, )
            # Normalize rewards
            if (sampled_rewards.max() - sampled_rewards.min()).item() > 1e-6:
                sampled_rewards = (sampled_rewards - sampled_rewards.mean()) / (sampled_rewards.std() + 1e-8)

            # Store all data
            # all_logits.append(duration_logits)
            # all_text_ids.append(duration_input_expanded["text_ids"])
            # all_mels.append(duration_input_expanded["mel"])
            all_sampled_classes.append(sampled_classes)
            all_durations.append(est_frames)
            all_gen_logps.append(gen_logps)
            all_rewards.extend(sampled_rewards)  # list with length of B*K

        # Concatenate all data
        # logits = torch.cat(all_logits, dim=0)
        # text_ids = torch.cat(all_text_ids, dim=0)
        # mels = torch.cat(all_mels, dim=0)
        sampled_classes = torch.cat(all_sampled_classes, dim=0)
        durations = torch.cat(all_durations, dim=0)
        rewards = torch.stack(all_rewards)    # use stack to keep the same device of elements
        gen_logps = torch.cat(all_gen_logps, dim=0)

        ctc_losses = torch.stack(all_ctc_loss)
        sv_losses = torch.stack(all_sv_loss)

        if self.is_main:
            self.accelerator.log({
                "ctc_loss": ctc_losses.mean().item(),
                "sv_loss": sv_losses.mean().item(),
                "reward": rewards.mean().item(),
                "reward_min": rewards.min().item(),
                "reward_max": rewards.max().item(),
            }, step=self.global_step)

        # # Normalize rewards
        # if (rewards.max() - rewards.min()).item() > 1e-6:
        #     rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-8)

        ref_logps = self.get_ref_logps(slp_text_ids, prompt_mel, sampled_classes)

        # Create batch dict similar to Qwen2.5 implementation
        batch_outputs = {
            # "logits": logits_grad,
            "text_ids": slp_text_ids,
            "prompt_mel": prompt_mel,
            "rewards": rewards,
            "refs": ref_logps,
            "sampled_classes": sampled_classes,
            "durations": durations,
        }

        if self.compute_gen_logps:
            batch_outputs["gen_logps"] = gen_logps

        if self.rank == 0:
            print(f"Generated {len(rewards)} samples with reward min/mean/max: {rewards.min().item():.4f}/{rewards.mean().item():.4f}/{rewards.max().item():.4f}")

        return batch_outputs

    def GRPO_step(self, batch):
        """
        Perform a GRPO update step

        Args:
            batch: Dictionary with inputs, rewards, reference logits, etc.

        Returns:
            Loss value
        """
        # Extract batch data
        # NOTE: why .unsqueeze(1) ???
        rewards = batch['rewards'] #.unsqueeze(1)
        ref_logps = batch['refs']  # (B)
        sampled_classes = batch['sampled_classes'] # (B)
        prompt_mel = batch['prompt_mel']
        text_ids = batch['text_ids']

        # Forward pass to get current model logits


        # Create batch dict similar to Qwen2.5 implementation
        batch_outputs = {
            # "logits": logits_grad,
            "text_ids": slp_text_ids,
            "prompt_mel": prompt_mel,
            "rewards": rewards,
            "refs": ref_logps,
            "sampled_classes": sampled_classes,
            "durations": durations,
        }

        if self.compute_gen_logps:
            batch_outputs["gen_logps"] = gen_logps

        if self.rank == 0:
            print(f"Generated {len(rewards)} samples with reward min/mean/max: {rewards.min().item():.4f}/{rewards.mean().item():.4f}/{rewards.max().item():.4f}")

        return batch_outputs

    def GRPO_step(self, batch):
        """
        Perform a GRPO update step

        Args:
            batch: Dictionary with inputs, rewards, reference logits, etc.

        Returns:
            Loss value
        """
        # Extract batch data
        # NOTE: why .unsqueeze(1) ???
        rewards = batch['rewards'] #.unsqueeze(1)
        ref_logps = batch['refs']  # (B)
        sampled_classes = batch['sampled_classes'] # (B)
        prompt_mel = batch['prompt_mel']
        text_ids = batch['text_ids']

        # Forward pass to get current model logits
        K = self.num_pre_samples
        B, T, _ = prompt_mel.shape
        _, L = text_ids.shape
        cur_logits = self.model(
            text_ids=text_ids, # (B, L)
            mel=prompt_mel         # (B, T, 100)
        )[:, -1, :]
        cur_logits = cur_logits.unsqueeze(1).repeat(1, K, 1).view(B*K, -1)

        # Compute current log probabilities for sampled actions
        log_probs = F.log_softmax(cur_logits, dim=-1)
        cur_logps = torch.gather(
            log_probs,
            dim=-1,
            index=sampled_classes.unsqueeze(-1)
        ).squeeze(-1)  # (B)

        # KL divergence computation (same as in Qwen2.5 code)
        # KL = exp(ref - cur) - (ref - cur) - 1
        kl_div = torch.exp(ref_logps - cur_logps) - (ref_logps - cur_logps) - 1 # (B)

        # Compute probability ratio for PPO
        if "gen_logps" in batch:
            gen_logps = batch['gen_logps']
            ratio = torch.exp(cur_logps - gen_logps)
            clipped_ratio = torch.clamp(ratio, 1 - self.clip_param, 1 + self.clip_param)
            loss = torch.min(ratio * rewards, clipped_ratio * rewards)
        else:
            # Simplification if gen_logps not available
            loss = torch.exp(cur_logps - cur_logps.detach()) * rewards

        # Final GRPO loss with KL regularization
        loss = -(loss - self.beta * kl_div) # (B)
        loss = loss.mean()

        return loss

    def get_batch(self):
        """Get a batch from the queue or return None if empty"""
        if not self.batch_queue:
            return None
        return self.batch_queue.pop(0)

    def generate_mode(self, num_batches=5):
        """
        Generate samples and add them to the batch queue

        Args:
            dataset: Dataset to sample from
            num_batches: Number of batches to generate
        """
        if self.rank == 0:
            print("Entering generate mode...")

        tic = time.time()
        for _ in range(num_batches):


        # KL divergence computation (same as in Qwen2.5 code)
        # KL = exp(ref - cur) - (ref - cur) - 1
        kl_div = torch.exp(ref_logps - cur_logps) - (ref_logps - cur_logps) - 1 # (B)

        # Compute probability ratio for PPO
        if "gen_logps" in batch:
            gen_logps = batch['gen_logps']
            ratio = torch.exp(cur_logps - gen_logps)
            clipped_ratio = torch.clamp(ratio, 1 - self.clip_param, 1 + self.clip_param)
            loss = torch.min(ratio * rewards, clipped_ratio * rewards)
        else:
            # Simplification if gen_logps not available
            loss = torch.exp(cur_logps - cur_logps.detach()) * rewards

        # Final GRPO loss with KL regularization
        loss = -(loss - self.beta * kl_div) # (B)
        loss = loss.mean()

        return loss

    def get_batch(self):
        """Get a batch from the queue or return None if empty"""
        if not self.batch_queue:
            return None
        return self.batch_queue.pop(0)

    def generate_mode(self, num_batches=5):
        """
        Generate samples and add them to the batch queue

        Args:
            dataset: Dataset to sample from
            num_batches: Number of batches to generate
        """
        if self.rank == 0:
            print("Entering generate mode...")

        tic = time.time()
        for _ in range(num_batches):
            try:
                batch_inputs = next(self.train_iterator)
            except StopIteration:
                self.train_iterator = iter(self.train_dataloader)
                batch_inputs = next(self.train_iterator)


            gen_logps = batch['gen_logps']
            ratio = torch.exp(cur_logps - gen_logps)
            clipped_ratio = torch.clamp(ratio, 1 - self.clip_param, 1 + self.clip_param)
            loss = torch.min(ratio * rewards, clipped_ratio * rewards)
        else:
            # Simplification if gen_logps not available
            loss = torch.exp(cur_logps - cur_logps.detach()) * rewards

        # Final GRPO loss with KL regularization
        loss = -(loss - self.beta * kl_div) # (B)
        loss = loss.mean()

        return loss

    def get_batch(self):
        """Get a batch from the queue or return None if empty"""
        if not self.batch_queue:
            return None
        return self.batch_queue.pop(0)

    def generate_mode(self, num_batches=5):
        """
        Generate samples and add them to the batch queue

        Args:
            dataset: Dataset to sample from
            num_batches: Number of batches to generate
        """
        if self.rank == 0:
            print("Entering generate mode...")

        tic = time.time()
        for _ in range(num_batches):
            try:
                batch_inputs = next(self.train_iterator)
            except StopIteration:
                self.train_iterator = iter(self.train_dataloader)
                batch_inputs = next(self.train_iterator)

            # Generate samples and compute rewards
            batch_outputs = self.generate_duration_samples(batch_inputs)
            # Check if batch has sufficient reward diversity
            rewards = batch_outputs["rewards"]
            if (rewards.max() - rewards.min()).item() < 0.01:
                if self.rank == 0:
                    print("Skipping batch with low reward diversity")
                continue
            # Add batch to queue
            self.batch_queue.append(batch_outputs)

        if self.rank == 0:
            print(f"Exiting generate mode: {time.time() - tic:.3f}s")

    def train(self, train_dataset, valid_dataset=None, num_workers=64, resumable_with_seed=666):
        """
        Train the model using GRPO

        Args:
            train_dataset: Training dataset
            valid_dataset: Validation dataset (optional)
            num_workers: Number of workers for data loading
        """

        # Create training dataloader using the appropriate batching strategy
        if self.batch_size_type == "sample":
            self.train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_size=self.batch_size,


train

            try:
                batch_inputs = next(self.train_iterator)
            except StopIteration:
                self.train_iterator = iter(self.train_dataloader)
                batch_inputs = next(self.train_iterator)

            # Generate samples and compute rewards
            batch_outputs = self.generate_duration_samples(batch_inputs)
            # Check if batch has sufficient reward diversity
            rewards = batch_outputs["rewards"]
            if (rewards.max() - rewards.min()).item() < 0.01:
                if self.rank == 0:
                    print("Skipping batch with low reward diversity")
                continue
            # Add batch to queue
            self.batch_queue.append(batch_outputs)

        if self.rank == 0:
            print(f"Exiting generate mode: {time.time() - tic:.3f}s")

    def train(self, train_dataset, valid_dataset=None, num_workers=64, resumable_with_seed=666):
        """
        Train the model using GRPO

        Args:
            train_dataset: Training dataset
            valid_dataset: Validation dataset (optional)
            num_workers: Number of workers for data loading
        """

        # Create training dataloader using the appropriate batching strategy
        if self.batch_size_type == "sample":
            self.train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_size=self.batch_size,
                shuffle=True,
                generator=generator,
            )
            # Create validation dataloader (always sequential, no shuffling)
            self.valid_dataloader = DataLoader(
                valid_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                batch_size=self.batch_size,
                shuffle=False,
            )

            self.train_iterator = iter(self.train_dataloader)
            self.valid_iterator = iter(self.valid_dataloader)

        elif self.batch_size_type == "frame":
            self.accelerator.even_batches = False

            sampler = SequentialSampler(train_dataset)
            batch_sampler = DynamicBatchSampler(
                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
            )
            self.train_dataloader = DataLoader(
                train_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_sampler=batch_sampler,
            )

            sampler = SequentialSampler(valid_dataset)
            batch_sampler = DynamicBatchSampler(
                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
            )
            # Create validation dataloader (always sequential, no shuffling)
            self.valid_dataloader = DataLoader(
                valid_dataset,
                collate_fn=collate_fn,
                num_workers=num_workers,
                pin_memory=True,
                persistent_workers=True,
                batch_sampler=batch_sampler,
            )

            self.train_dataloader, self.valid_dataloader = self.accelerator.prepare(self.train_dataloader, self.valid_dataloader)

            self.train_iterator = iter(self.train_dataloader)
            self.valid_iterator = iter(self.valid_dataloader)

        else:
            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")

        # Setup schedulers
        warmup_steps = self.num_warmup_updates * self.accelerator.num_processes
        total_steps = self.all_steps
        decay_steps = total_steps - warmup_steps

        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)

        self.scheduler = SequentialLR(
            self.optimizer,
            schedulers=[warmup_scheduler, decay_scheduler],
            milestones=[warmup_steps]
        )

        self.scheduler = self.accelerator.prepare(self.scheduler)

        # Load checkpoint if available
        start_step = self.load_checkpoint()
        self.global_step = start_step

        # Generate initial batches
        self.generate_mode()

        # Training loop
        progress = range(1, self.all_steps + 1)

        # Skip steps that are already done
        progress = [step for step in progress if step > start_step]
        if self.is_main:
            progress = tqdm(progress, desc="Training", unit="step")

        for step in progress:
            # Get batch from queue or generate more
            batch = self.get_batch()
            while batch is None:
                self.generate_mode()
                batch = self.get_batch()

            # GRPO update
            with self.accelerator.accumulate(self.model):
                loss = self.GRPO_step(batch)
                # for param in self.model.parameters():
                #     custom_loss = loss + 0 * param.sum()
                self.accelerator.backward(loss)

                if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
                    total_norm = self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
                else:
                    total_norm = torch.norm(
                        torch.stack([
                            torch.norm(p.grad.detach(), 2)
                            for p in self.model.parameters()
                            if p.grad is not None
                        ]),
                        2
                    )

                self.accelerator.log({
                    "grad_norm": total_norm.item()
                }, step=self.global_step)

                self.optimizer.step()
                self.scheduler.step()
                self.optimizer.zero_grad()

            self.global_step += 1

            # Log metrics
            if self.is_main:
                self.accelerator.log({
                    "loss": loss.item(),
                    "lr": self.scheduler.get_last_lr()[0],
                    # "avg_reward": batch["rewards"].mean().item(),
                    # "max_reward": batch["rewards"].max().item(),
                    # "min_reward": batch["rewards"].min().item(),
                }, step=self.global_step)
                progress.set_postfix(
                    loss=f"{loss.item():.4f}",
                    lr=f"{self.scheduler.get_last_lr()[0]:.8f}"
                )

            # Save checkpoint
            if self.global_step % self.save_per_updates == 0:
                self.save_checkpoint(self.global_step)

                # Optional validation logic could be added here

        # Save final checkpoint
        self.save_checkpoint(self.global_step, last=True)
        self.accelerator.end_training()


src.infer

import os
import torch
import torchaudio
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torchdiffeq import odeint
from safetensors.torch import load_file
import IPython.display as ipd

# Import F5-TTS modules
from f5_tts.model import CFM, UNetT, DiT
from f5_tts.model.modules import MelSpec
from f5_tts.model.utils import (
    default, exists, list_str_to_idx, list_str_to_tensor,
    lens_to_mask, mask_from_frac_lengths, get_tokenizer
)
from f5_tts.infer.utils_infer import (
    load_vocoder, preprocess_ref_audio_text, chunk_text,
    convert_char_to_pinyin, transcribe, target_rms,
    target_sample_rate, hop_length, speed
)

# Import custom modules
from unimodel import UniModel
from duration_predictor import SpeechLengthPredictor

class DMOInference:
    """F5-TTS Inference wrapper class for easy text-to-speech generation."""

    def __init__(
        self,
        student_checkpoint_path="",
        duration_predictor_path="",
        device="cuda",
        model_type="F5TTS_Base",  # "F5TTS_Base" or "E2TTS_Base"
        tokenizer="pinyin",
        dataset_name="Emilia_ZH_EN",
    ):
        """
        Initialize F5-TTS inference model.

        Args:
            student_checkpoint_path: Path to student model checkpoint
            duration_predictor_path: Path to duration predictor checkpoint
            device: Device to run inference on
            model_type: Model architecture type
            tokenizer: Tokenizer type ("pinyin", "char", or "custom")
            dataset_name: Dataset name for tokenizer
            cuda_device_id: CUDA device ID to use
        """

        self.device = device
        self.model_type = model_type
        self.tokenizer = tokenizer
        self.dataset_name = dataset_name

        # Model parameters
        self.target_sample_rate = 24000
        self.n_mel_channels = 100
        self.hop_length = 256
        self.real_guidance_scale = 2
        self.fake_guidance_scale = 0
        self.gen_cls_loss = False
        self.num_student_step = 4

        # Initialize components
        self._setup_tokenizer()
        self._setup_models(student_checkpoint_path)
        self._setup_mel_spec()
        self._setup_vocoder()
        self._setup_duration_predictor(duration_predictor_path)

    def _setup_tokenizer(self):
        """Setup tokenizer and vocabulary."""
        if self.tokenizer == "custom":
            tokenizer_path = self.tokenizer_path
        else:
            tokenizer_path = self.dataset_name

        self.vocab_char_map, self.vocab_size = get_tokenizer(tokenizer_path, self.tokenizer)

    def _setup_models(self, student_checkpoint_path):
        """Initialize teacher and student models."""
        # Model configuration
        if self.model_type == "F5TTS_Base":
            model_cls = DiT
            model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
        elif self.model_type == "E2TTS_Base":
            model_cls = UNetT
            model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")

        # Initialize UniModel (student)
        self.model = UniModel(
            model_cls(**model_cfg, text_num_embeds=self.vocab_size, mel_dim=self.n_mel_channels,
                      second_time=self.num_student_step > 1),
            checkpoint_path="",
            vocab_char_map=self.vocab_char_map,
            frac_lengths_mask=(0.5, 0.9),
            real_guidance_scale=self.real_guidance_scale,
            fake_guidance_scale=self.fake_guidance_scale,
            gen_cls_loss=self.gen_cls_loss,
            sway_coeff=0,
        )

        # Load student checkpoint
        checkpoint = torch.load(student_checkpoint_path, map_location='cpu')
        self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)

        # Setup generator and teacher
        self.generator = self.model.feedforward_model.to(self.device)
        self.teacher = self.model.guidance_model.real_unet.to(self.device)

        self.scale = checkpoint['scale']

    def _setup_mel_spec(self):
        """Initialize mel spectrogram module."""
        mel_spec_kwargs = dict(
            target_sample_rate=self.target_sample_rate,
            n_mel_channels=self.n_mel_channels,
            hop_length=self.hop_length,
        )
        self.mel_spec = MelSpec(**mel_spec_kwargs)

    def _setup_vocoder(self):
        """Initialize vocoder."""
        self.vocos = load_vocoder(is_local=False, local_path="")
        self.vocos = self.vocos.to(self.device)

    def _setup_duration_predictor(self, checkpoint_path):
        """Initialize duration predictor."""
        self.wav2mel = MelSpec(
            target_sample_rate=24000,
            n_mel_channels=100,
            hop_length=256,
            win_length=1024,
            n_fft=1024,
            mel_spec_type='vocos'
        ).to(self.device)

        self.SLP = SpeechLengthPredictor(
            vocab_size=2545,
            n_mel=100,
            hidden_dim=512,
            n_text_layer=4,
            n_cross_layer=4,
            n_head=8,
            output_dim=301
        ).to(self.device)

        self.SLP.eval()
        self.SLP.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['model_state_dict'])

    def predict_duration(self, pmt_wav_path, tar_text, pmt_text, dp_softmax_range=0.7, temperature=0):
        """
        Predict duration for target text based on prompt audio.

        Args:
            pmt_wav_path: Path to prompt audio
            tar_text: Target text to generate
            pmt_text: Prompt text
            dp_softmax_range: softmax annliation range from rate-based duration
            temperature: temperature for softmax sampling (if 0, will use argmax)
        Returns:
            Estimated duration in frames
        """
        pmt_wav, sr = torchaudio.load(pmt_wav_path)
        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            pmt_wav = resampler(pmt_wav)
        if pmt_wav.size(0) > 1:
            pmt_wav = pmt_wav[0].unsqueeze(0)
        pmt_wav = pmt_wav.to(self.device)

        pmt_mel = self.wav2mel(pmt_wav).permute(0, 2, 1)
        tar_tokens = self._convert_to_pinyin(list(tar_text))
        pmt_tokens = self._convert_to_pinyin(list(pmt_text))

        # Calculate duration
        ref_text_len = len(pmt_tokens)
        gen_text_len = len(tar_tokens)
        ref_audio_len = pmt_mel.size(1)
        duration = int(ref_audio_len / ref_text_len * gen_text_len / speed)
        duration = duration // 10

        min_duration = max(int(duration * dp_softmax_range), 0)
        max_duration = min(int(duration * (1 + dp_softmax_range)), 301)

        all_tokens = pmt_tokens + [' '] + tar_tokens

        text_ids = list_str_to_idx([all_tokens], self.vocab_char_map).to(self.device)
        text_ids = text_ids.masked_fill(text_ids == -1, self.vocab_size)

        with torch.no_grad():
            predictions = self.SLP(text_ids=text_ids, mel=pmt_mel)
        predictions = predictions[:, -1, :]
        predictions[:, :min_duration] = float('-inf')
        predictions[:, max_duration:] = float('-inf')

        if temperature == 0:
            est_label = predictions.argmax(-1)[..., -1].item() * 10
        else:
            probs = torch.softmax(predictions / temperature, dim=-1)
            sampled_idx = torch.multinomial(probs.squeeze(0), num_samples=1)  # Remove the -1 index
            est_label = sampled_idx.item() * 10

        return est_label

    def _convert_to_pinyin(self, char_list):
        """Convert character list to pinyin."""
        result = []
        for x in convert_char_to_pinyin(char_list):
            result = result + x
        while result[0] == ' ' and len(result) > 1:
            result = result[1:]
        return result

    def generate(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=16,
        teacher_stopping_time=0.07,
        student_start_step=1,
        duration=None,
        dp_softmax_range=0.7,
        temperature=0,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0,
        verbose=False
    ):
        """
        Generate speech from text using teacher-student distillation.

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of teacher guidance steps
            teacher_stopping_time: When to stop teacher sampling
            student_start_step: When to start student sampling
            duration: Total duration (if None, will predict)
            dp_softmax_range: Duration predictor softmax range allowed around rate based duration
            temperature: Temperature for duration predictor sampling (0 means use argmax)
            eta: Stochasticity control (0=DDIM, 1=DDPM)
            cfg_strength: Classifier-free guidance strength
            sway_coefficient: Sway sampling coefficient
            verbose: Output sampling steps

        Returns:
            Generated audio waveform
        """
        if prompt_text is None:
            prompt_text = transcribe(audio_path)

        # Predict duration if not provided
        if duration is None:
            duration = self.predict_duration(audio_path, gen_text, prompt_text, dp_softmax_range, temperature)

        # Preprocess audio and text
        ref_audio, ref_text = preprocess_ref_audio_text(audio_path, prompt_text)
        audio, sr = torchaudio.load(ref_audio)

        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        # Normalize audio
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms

        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            audio = resampler(audio)

        audio = audio.to(self.device)

        # Prepare text
        text_list = [ref_text + gen_text]
        final_text_list = convert_char_to_pinyin(text_list)

        # Calculate durations
        ref_audio_len = audio.shape[-1] // self.hop_length
        if duration is None:
            ref_text_len = len(ref_text.encode("utf-8"))
            gen_text_len = len(gen_text.encode("utf-8"))
            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
        else:
            duration = ref_audio_len + duration

        if verbose:
            print('audio:', audio.shape)
            print('text:', final_text_list)
            print('duration:', duration)
            print('eta (stochasticity):', eta)  # Print eta value for debugging

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration, duration_tensor = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-student sampling
            if teacher_steps > 0 and student_start_step > 0:
                if verbose:
                    print('Start teacher sampling with hybrid DDIM/DDPM (eta={})....'.format(eta))
                x1 = self._teacher_sampling(
                    step_cond, text, cond_mask, max_duration, duration_tensor,  # Use duration_tensor
                    teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_coefficient
                )
            else:
                x1 = step_cond

            if verbose:
                print('Start student sampling...')
            # Student sampling
            x1 = self._student_sampling(x1, cond, text, student_start_step, verbose, sway_coefficient)

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave.cpu().numpy().squeeze()

    def generate_teacher_only(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=32,
        duration=None,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0
    ):
        """
        Generate speech using teacher model only (no student distillation).

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of sampling steps
            duration: Total duration (if None, will predict)
            eta: Stochasticity control (0=DDIM, 1=DDPM)
            cfg_strength: Classifier-free guidance strength
            sway_coefficient: Sway sampling coefficient

        Returns:
            Generated audio waveform
        """
        if prompt_text is None:
            prompt_text = transcribe(audio_path)

        # Predict duration if not provided
        if duration is None:
            duration = self.predict_duration(audio_path, gen_text, prompt_text)

        # Preprocess audio and text
        ref_audio, ref_text = preprocess_ref_audio_text(audio_path, prompt_text)
        audio, sr = torchaudio.load(ref_audio)

        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        # Normalize audio
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms

        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            audio = resampler(audio)

        audio = audio.to(self.device)

        # Prepare text
        text_list = [ref_text + gen_text]
        final_text_list = convert_char_to_pinyin(text_list)

        # Calculate durations
        ref_audio_len = audio.shape[-1] // self.hop_length
        if duration is None:
            ref_text_len = len(ref_text.encode("utf-8"))
            gen_text_len = len(gen_text.encode("utf-8"))
            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
        else:
            duration = ref_audio_len + duration

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-only sampling
            x1 = self._teacher_sampling(
                step_cond, text, cond_mask, max_duration, duration,
                teacher_steps, 1.0, eta, cfg_strength, sway_coefficient  # stopping_time=1.0 for full sampling
            )

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave

    def _prepare_inputs(self, audio, text_list, duration):
        """Prepare inputs for generation."""
        lens = None
        max_duration_limit = 4096

        cond = audio
        text = text_list

        if cond.ndim == 2:
            cond = self.mel_spec(cond)
            cond = cond.permute(0, 2, 1)
            assert cond.shape[-1] == 100

        cond = cond / self.scale
        batch, cond_seq_len, device = *cond.shape[:2], cond.device

        if not exists(lens):
            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)

        # Process text
        if isinstance(text, list):
            if exists(self.vocab_char_map):
                text = list_str_to_idx(text, self.vocab_char_map).to(device)
            else:
                text = list_str_to_tensor(text).to(device)
            assert text.shape[0] == batch

        if exists(text):
            text_lens = (text != -1).sum(dim=-1)
            lens = torch.maximum(text_lens, lens)

        # Process duration
        cond_mask = lens_to_mask(lens)

        if isinstance(duration, int):
            duration = torch.full((batch,), duration, device=device, dtype=torch.long)

        duration = torch.maximum(lens + 1, duration)
        duration = duration.clamp(max=max_duration_limit)
        max_duration = duration.amax()

        # Pad conditioning
        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
        cond_mask = cond_mask.unsqueeze(-1)
        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))

        return cond, text, step_cond, cond_mask, max_duration, duration

    def _teacher_sampling(self, step_cond, text, cond_mask, max_duration, duration,
                         teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_sampling_coef = -1):
        """Perform teacher model sampling."""
        device = step_cond.device

        # Pre-generate noise sequence for stochastic sampling
        noise_seq = None
        if eta > 0:
            noise_seq = [torch.randn(1, max_duration, 100, device=device)
                        for _ in range(teacher_steps)]

        def fn(t, x):
            with torch.inference_mode():
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    if verbose:
                        print(f'current t: {t}')
                    step_frac = 1.0 - t.item()
                    step_idx = min(int(step_frac * len(noise_seq)), len(noise_seq) - 1) if noise_seq else 0

                    # Predict flow
                    pred = self.teacher(
                        x=x, cond=step_cond, text=text, time=t, mask=None,
                        drop_audio_cond=False, drop_text=False
                    )

                    if cfg_strength > 1e-5:
                        null_pred = self.teacher(
                            x=x, cond=step_cond, text=text, time=t, mask=None,
                            drop_audio_cond=True, drop_text=True
                        )
                        pred = pred + (pred - null_pred) * cfg_strength

                    # Add stochasticity if eta > 0
                    if eta > 0 and noise_seq is not None:
                        alpha_t = 1.0 - t.item()
                        sigma_t = t.item()
                        noise_scale = torch.sqrt(torch.tensor(
                            (sigma_t**2) / (alpha_t**2 + sigma_t**2) * eta,
                            device=device
                        ))
                        return pred + noise_scale * noise_seq[step_idx]
                    else:
                        return pred

        # Initialize noise
        y0 = []
        for dur in duration:
            y0.append(torch.randn(dur, 100, device=device, dtype=step_cond.dtype))
        y0 = pad_sequence(y0, padding_value=0, batch_first=True)

        # Setup time steps
        t = torch.linspace(0, 1, teacher_steps + 1, device=device, dtype=step_cond.dtype)
        if sway_sampling_coef is not None:
            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
        t = t[:(t > teacher_stoppin
# ... truncated ...


DMOInference

import IPython.display as ipd

# Import F5-TTS modules
from f5_tts.model import CFM, UNetT, DiT
from f5_tts.model.modules import MelSpec
from f5_tts.model.utils import (
    default, exists, list_str_to_idx, list_str_to_tensor,
    lens_to_mask, mask_from_frac_lengths, get_tokenizer
)
from f5_tts.infer.utils_infer import (
    load_vocoder, preprocess_ref_audio_text, chunk_text,
    convert_char_to_pinyin, transcribe, target_rms,
    target_sample_rate, hop_length, speed
)

# Import custom modules
from unimodel import UniModel
from duration_predictor import SpeechLengthPredictor

class DMOInference:
    """F5-TTS Inference wrapper class for easy text-to-speech generation."""

    def __init__(
        self,
        student_checkpoint_path="",
        duration_predictor_path="",
        device="cuda",
        model_type="F5TTS_Base",  # "F5TTS_Base" or "E2TTS_Base"
        tokenizer="pinyin",
        dataset_name="Emilia_ZH_EN",
    ):
        """
        Initialize F5-TTS inference model.

        Args:
            student_checkpoint_path: Path to student model checkpoint
            duration_predictor_path: Path to duration predictor checkpoint
            device: Device to run inference on
            model_type: Model architecture type
            tokenizer: Tokenizer type ("pinyin", "char", or "custom")
            dataset_name: Dataset name for tokenizer
            cuda_device_id: CUDA device ID to use
        """

        self.device = device
        self.model_type = model_type
        self.tokenizer = tokenizer
        self.dataset_name = dataset_name

        # Model parameters
        self.target_sample_rate = 24000
        self.n_mel_channels = 100
        self.hop_length = 256
        self.real_guidance_scale = 2
        self.fake_guidance_scale = 0
        self.gen_cls_loss = False
        self.num_student_step = 4

        # Initialize components
        self._setup_tokenizer()
        self._setup_models(student_checkpoint_path)
        self._setup_mel_spec()
        self._setup_vocoder()
        self._setup_duration_predictor(duration_predictor_path)

    def _setup_tokenizer(self):
        """Setup tokenizer and vocabulary."""
        if self.tokenizer == "custom":
            tokenizer_path = self.tokenizer_path
        else:
            tokenizer_path = self.dataset_name

        self.vocab_char_map, self.vocab_size = get_tokenizer(tokenizer_path, self.tokenizer)

    def _setup_models(self, student_checkpoint_path):
        """Initialize teacher and student models."""
        # Model configuration
        if self.model_type == "F5TTS_Base":
            model_cls = DiT
            model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
        elif self.model_type == "E2TTS_Base":
            model_cls = UNetT
            model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")

        # Initialize UniModel (student)
        self.model = UniModel(
            model_cls(**model_cfg, text_num_embeds=self.vocab_size, mel_dim=self.n_mel_channels,
                      second_time=self.num_student_step > 1),
            checkpoint_path="",
            vocab_char_map=self.vocab_char_map,
            frac_lengths_mask=(0.5, 0.9),
            real_guidance_scale=self.real_guidance_scale,
            fake_guidance_scale=self.fake_guidance_scale,
            gen_cls_loss=self.gen_cls_loss,
            sway_coeff=0,
        )

        # Load student checkpoint
        checkpoint = torch.load(student_checkpoint_path, map_location='cpu')
        self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)

        # Setup generator and teacher
        self.generator = self.model.feedforward_model.to(self.device)
        self.teacher = self.model.guidance_model.real_unet.to(self.device)

        self.scale = checkpoint['scale']

    def _setup_mel_spec(self):
        """Initialize mel spectrogram module."""
        mel_spec_kwargs = dict(
            target_sample_rate=self.target_sample_rate,
            n_mel_channels=self.n_mel_channels,
            hop_length=self.hop_length,
        )
        self.mel_spec = MelSpec(**mel_spec_kwargs)

    def _setup_vocoder(self):
        """Initialize vocoder."""
        self.vocos = load_vocoder(is_local=False, local_path="")
        self.vocos = self.vocos.to(self.device)

    def _setup_duration_predictor(self, checkpoint_path):
        """Initialize duration predictor."""
        self.wav2mel = MelSpec(
            target_sample_rate=24000,
            n_mel_channels=100,
            hop_length=256,
            win_length=1024,
            n_fft=1024,
            mel_spec_type='vocos'
        ).to(self.device)

        self.SLP = SpeechLengthPredictor(
            vocab_size=2545,
            n_mel=100,
            hidden_dim=512,
            n_text_layer=4,
            n_cross_layer=4,
            n_head=8,
            output_dim=301
        ).to(self.device)

        self.SLP.eval()
        self.SLP.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['model_state_dict'])

    def predict_duration(self, pmt_wav_path, tar_text, pmt_text, dp_softmax_range=0.7, temperature=0):
        """
        Predict duration for target text based on prompt audio.

        Args:
            pmt_wav_path: Path to prompt audio
            tar_text: Target text to generate
            pmt_text: Prompt text
            dp_softmax_range: softmax annliation range from rate-based duration
            temperature: temperature for softmax sampling (if 0, will use argmax)
        Returns:
            Estimated duration in frames
        """
        pmt_wav, sr = torchaudio.load(pmt_wav_path)
        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            pmt_wav = resampler(pmt_wav)
        if pmt_wav.size(0) > 1:
            pmt_wav = pmt_wav[0].unsqueeze(0)
        pmt_wav = pmt_wav.to(self.device)

        pmt_mel = self.wav2mel(pmt_wav).permute(0, 2, 1)
        tar_tokens = self._convert_to_pinyin(list(tar_text))
        pmt_tokens = self._convert_to_pinyin(list(pmt_text))

        # Calculate duration
        ref_text_len = len(pmt_tokens)
        gen_text_len = len(tar_tokens)
        ref_audio_len = pmt_mel.size(1)
        duration = int(ref_audio_len / ref_text_len * gen_text_len / speed)
        duration = duration // 10

        min_duration = max(int(duration * dp_softmax_range), 0)
        max_duration = min(int(duration * (1 + dp_softmax_range)), 301)

        all_tokens = pmt_tokens + [' '] + tar_tokens

        text_ids = list_str_to_idx([all_tokens], self.vocab_char_map).to(self.device)
        text_ids = text_ids.masked_fill(text_ids == -1, self.vocab_size)

        with torch.no_grad():
            predictions = self.SLP(text_ids=text_ids, mel=pmt_mel)
        predictions = predictions[:, -1, :]
        predictions[:, :min_duration] = float('-inf')
        predictions[:, max_duration:] = float('-inf')

        if temperature == 0:
            est_label = predictions.argmax(-1)[..., -1].item() * 10
        else:
            probs = torch.softmax(predictions / temperature, dim=-1)
            sampled_idx = torch.multinomial(probs.squeeze(0), num_samples=1)  # Remove the -1 index
            est_label = sampled_idx.item() * 10

        return est_label

    def _convert_to_pinyin(self, char_list):
        """Convert character list to pinyin."""
        result = []
        for x in convert_char_to_pinyin(char_list):
            result = result + x
        while result[0] == ' ' and len(result) > 1:
            result = result[1:]
        return result

    def generate(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=16,
        teacher_stopping_time=0.07,
        student_start_step=1,
        duration=None,
        dp_softmax_range=0.7,
        temperature=0,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0,
        verbose=False
    ):
        """
        Generate speech from text using teacher-student distillation.

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of teacher guidance steps
            teacher_stopping_time: When to stop teacher sampling
            student_start_step: When to start student sampling
            duration: Total duration (if None, will predict)
            dp_softmax_range: Duration predictor softmax range allowed around rate based duration
            temperature: Temperature for duration predictor sampling (0 means use argmax)
            eta: Stochasticity control (0=DDIM, 1=DDPM)
            cfg_strength: Classifier-free guidance strength
            sway_coefficient: Sway sampling coefficient
            verbose: Output sampling steps

        Returns:
            Generated audio waveform
        """
        if prompt_text is None:
            prompt_text = transcribe(audio_path)

        # Predict duration if not provided
        if duration is None:
            duration = self.predict_duration(audio_path, gen_text, prompt_text, dp_softmax_range, temperature)

        # Preprocess audio and text
        ref_audio, ref_text = preprocess_ref_audio_text(audio_path, prompt_text)
        audio, sr = torchaudio.load(ref_audio)

        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        # Normalize audio
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms

        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            audio = resampler(audio)

        audio = audio.to(self.device)

        # Prepare text
        text_list = [ref_text + gen_text]
        final_text_list = convert_char_to_pinyin(text_list)

        # Calculate durations
        ref_audio_len = audio.shape[-1] // self.hop_length
        if duration is None:
            ref_text_len = len(ref_text.encode("utf-8"))
            gen_text_len = len(gen_text.encode("utf-8"))
            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
        else:
            duration = ref_audio_len + duration

        if verbose:
            print('audio:', audio.shape)
            print('text:', final_text_list)
            print('duration:', duration)
            print('eta (stochasticity):', eta)  # Print eta value for debugging

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration, duration_tensor = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-student sampling
            if teacher_steps > 0 and student_start_step > 0:
                if verbose:
                    print('Start teacher sampling with hybrid DDIM/DDPM (eta={})....'.format(eta))
                x1 = self._teacher_sampling(
                    step_cond, text, cond_mask, max_duration, duration_tensor,  # Use duration_tensor
                    teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_coefficient
                )
            else:
                x1 = step_cond

            if verbose:
                print('Start student sampling...')
            # Student sampling
            x1 = self._student_sampling(x1, cond, text, student_start_step, verbose, sway_coefficient)

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave.cpu().numpy().squeeze()

    def generate_teacher_only(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=32,
        duration=None,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0
    ):
        """
        Generate speech using teacher model only (no student distillation).

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of sampling steps
            duration: Total duration (if None, will predict)
            eta: Stochasticity control (0=DDIM, 1=DDPM)
            cfg_strength: Classifier-free guidance strength
            sway_coefficient: Sway sampling coefficient

        Returns:
            Generated audio waveform
        """
        if prompt_text is None:
            prompt_text = transcribe(audio_path)

        # Predict duration if not provided
        if duration is None:
            duration = self.predict_duration(audio_path, gen_text, prompt_text)

        # Preprocess audio and text
        ref_audio, ref_text = preprocess_ref_audio_text(audio_path, prompt_text)
        audio, sr = torchaudio.load(ref_audio)

        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        # Normalize audio
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms

        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            audio = resampler(audio)

        audio = audio.to(self.device)

        # Prepare text
        text_list = [ref_text + gen_text]
        final_text_list = convert_char_to_pinyin(text_list)

        # Calculate durations
        ref_audio_len = audio.shape[-1] // self.hop_length
        if duration is None:
            ref_text_len = len(ref_text.encode("utf-8"))
            gen_text_len = len(gen_text.encode("utf-8"))
            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
        else:
            duration = ref_audio_len + duration

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-only sampling
            x1 = self._teacher_sampling(
                step_cond, text, cond_mask, max_duration, duration,
                teacher_steps, 1.0, eta, cfg_strength, sway_coefficient  # stopping_time=1.0 for full sampling
            )

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave

    def _prepare_inputs(self, audio, text_list, duration):
        """Prepare inputs for generation."""
        lens = None
        max_duration_limit = 4096

        cond = audio
        text = text_list

        if cond.ndim == 2:
            cond = self.mel_spec(cond)
            cond = cond.permute(0, 2, 1)
            assert cond.shape[-1] == 100

        cond = cond / self.scale
        batch, cond_seq_len, device = *cond.shape[:2], cond.device

        if not exists(lens):
            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)

        # Process text
        if isinstance(text, list):
            if exists(self.vocab_char_map):
                text = list_str_to_idx(text, self.vocab_char_map).to(device)
            else:
                text = list_str_to_tensor(text).to(device)
            assert text.shape[0] == batch

        if exists(text):
            text_lens = (text != -1).sum(dim=-1)
            lens = torch.maximum(text_lens, lens)

        # Process duration
        cond_mask = lens_to_mask(lens)

        if isinstance(duration, int):
            duration = torch.full((batch,), duration, device=device, dtype=torch.long)

        duration = torch.maximum(lens + 1, duration)
        duration = duration.clamp(max=max_duration_limit)
        max_duration = duration.amax()

        # Pad conditioning
        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
        cond_mask = cond_mask.unsqueeze(-1)
        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))

        return cond, text, step_cond, cond_mask, max_duration, duration

    def _teacher_sampling(self, step_cond, text, cond_mask, max_duration, duration,
                         teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_sampling_coef = -1):
        """Perform teacher model sampling."""
        device = step_cond.device

        # Pre-generate noise sequence for stochastic sampling
        noise_seq = None
        if eta > 0:
            noise_seq = [torch.randn(1, max_duration, 100, device=device)
                        for _ in range(teacher_steps)]

        def fn(t, x):
            with torch.inference_mode():
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    if verbose:
                        print(f'current t: {t}')
                    step_frac = 1.0 - t.item()
                    step_idx = min(int(step_frac * len(noise_seq)), len(noise_seq) - 1) if noise_seq else 0

                    # Predict flow
                    pred = self.teacher(
                        x=x, cond=step_cond, text=text, time=t, mask=None,
                        drop_audio_cond=False, drop_text=False
                    )

                    if cfg_strength > 1e-5:
                        null_pred = self.teacher(
                            x=x, cond=step_cond, text=text, time=t, mask=None,
                            drop_audio_cond=True, drop_text=True
                        )
                        pred = pred + (pred - null_pred) * cfg_strength

                    # Add stochasticity if eta > 0
                    if eta > 0 and noise_seq is not None:
                        alpha_t = 1.0 - t.item()
                        sigma_t = t.item()
                        noise_scale = torch.sqrt(torch.tensor(
                            (sigma_t**2) / (alpha_t**2 + sigma_t**2) * eta,
                            device=device
                        ))
                        return pred + noise_scale * noise_seq[step_idx]
                    else:
                        return pred

        # Initialize noise
        y0 = []
        for dur in duration:
            y0.append(torch.randn(dur, 100, device=device, dtype=step_cond.dtype))
        y0 = pad_sequence(y0, padding_value=0, batch_first=True)

        # Setup time steps
        t = torch.linspace(0, 1, teacher_steps + 1, device=device, dtype=step_cond.dtype)
        if sway_sampling_coef is not None:
            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
        t = t[:(t > teacher_stopping_time).float().argmax() + 2]
        t = t[:-1]

        # Solve ODE
        trajectory = odeint(fn, y0, t, method="euler")

        if teacher_stopping_time < 1.0:
# ... truncated ...


from f5_tts.model import CFM, UNetT, DiT
from f5_tts.model.modules import MelSpec
from f5_tts.model.utils import (
    default, exists, list_str_to_idx, list_str_to_tensor,
    lens_to_mask, mask_from_frac_lengths, get_tokenizer
)
from f5_tts.infer.utils_infer import (
    load_vocoder, preprocess_ref_audio_text, chunk_text,
    convert_char_to_pinyin, transcribe, target_rms,
    target_sample_rate, hop_length, speed
)

# Import custom modules
from unimodel import UniModel
from duration_predictor import SpeechLengthPredictor

class DMOInference:
    """F5-TTS Inference wrapper class for easy text-to-speech generation."""

    def __init__(
        self,
        student_checkpoint_path="",
        duration_predictor_path="",
        device="cuda",
        model_type="F5TTS_Base",  # "F5TTS_Base" or "E2TTS_Base"
        tokenizer="pinyin",
        dataset_name="Emilia_ZH_EN",
    ):
        """
        Initialize F5-TTS inference model.

        Args:
            student_checkpoint_path: Path to student model checkpoint
            duration_predictor_path: Path to duration predictor checkpoint
            device: Device to run inference on
            model_type: Model architecture type
            tokenizer: Tokenizer type ("pinyin", "char", or "custom")
            dataset_name: Dataset name for tokenizer
            cuda_device_id: CUDA device ID to use
        """

        self.device = device
        self.model_type = model_type
        self.tokenizer = tokenizer
        self.dataset_name = dataset_name

        # Model parameters
        self.target_sample_rate = 24000
        self.n_mel_channels = 100
        self.hop_length = 256
        self.real_guidance_scale = 2
        self.fake_guidance_scale = 0
        self.gen_cls_loss = False
        self.num_student_step = 4

        # Initialize components
        self._setup_tokenizer()
        self._setup_models(student_checkpoint_path)
        self._setup_mel_spec()
        self._setup_vocoder()
        self._setup_duration_predictor(duration_predictor_path)

    def _setup_tokenizer(self):
        """Setup tokenizer and vocabulary."""
        if self.tokenizer == "custom":
            tokenizer_path = self.tokenizer_path
        else:
            tokenizer_path = self.dataset_name

        self.vocab_char_map, self.vocab_size = get_tokenizer(tokenizer_path, self.tokenizer)

    def _setup_models(self, student_checkpoint_path):
        """Initialize teacher and student models."""
        # Model configuration
        if self.model_type == "F5TTS_Base":
            model_cls = DiT
            model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
        elif self.model_type == "E2TTS_Base":
            model_cls = UNetT
            model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
        else:


        self.model_type = model_type
        self.tokenizer = tokenizer
        self.dataset_name = dataset_name

        # Model parameters
        self.target_sample_rate = 24000
        self.n_mel_channels = 100
        self.hop_length = 256
        self.real_guidance_scale = 2
        self.fake_guidance_scale = 0
        self.gen_cls_loss = False
        self.num_student_step = 4

        # Initialize components
        self._setup_tokenizer()
        self._setup_models(student_checkpoint_path)
        self._setup_mel_spec()
        self._setup_vocoder()
        self._setup_duration_predictor(duration_predictor_path)

    def _setup_tokenizer(self):
        """Setup tokenizer and vocabulary."""
        if self.tokenizer == "custom":
            tokenizer_path = self.tokenizer_path
        else:
            tokenizer_path = self.dataset_name

        self.vocab_char_map, self.vocab_size = get_tokenizer(tokenizer_path, self.tokenizer)

    def _setup_models(self, student_checkpoint_path):
        """Initialize teacher and student models."""
        # Model configuration
        if self.model_type == "F5TTS_Base":
            model_cls = DiT
            model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
        elif self.model_type == "E2TTS_Base":
            model_cls = UNetT
            model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")

        # Initialize UniModel (student)
        self.model = UniModel(
            model_cls(**model_cfg, text_num_embeds=self.vocab_size, mel_dim=self.n_mel_channels,
                      second_time=self.num_student_step > 1),
            checkpoint_path="",
            vocab_char_map=self.vocab_char_map,
            frac_lengths_mask=(0.5, 0.9),


        self.fake_guidance_scale = 0
        self.gen_cls_loss = False
        self.num_student_step = 4

        # Initialize components
        self._setup_tokenizer()
        self._setup_models(student_checkpoint_path)
        self._setup_mel_spec()
        self._setup_vocoder()
        self._setup_duration_predictor(duration_predictor_path)

    def _setup_tokenizer(self):
        """Setup tokenizer and vocabulary."""
        if self.tokenizer == "custom":
            tokenizer_path = self.tokenizer_path
        else:
            tokenizer_path = self.dataset_name

        self.vocab_char_map, self.vocab_size = get_tokenizer(tokenizer_path, self.tokenizer)

    def _setup_models(self, student_checkpoint_path):
        """Initialize teacher and student models."""
        # Model configuration
        if self.model_type == "F5TTS_Base":
            model_cls = DiT
            model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
        elif self.model_type == "E2TTS_Base":
            model_cls = UNetT
            model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
        else:
            raise ValueError(f"Unknown model type: {self.model_type}")

        # Initialize UniModel (student)
        self.model = UniModel(
            model_cls(**model_cfg, text_num_embeds=self.vocab_size, mel_dim=self.n_mel_channels,
                      second_time=self.num_student_step > 1),
            checkpoint_path="",
            vocab_char_map=self.vocab_char_map,
            frac_lengths_mask=(0.5, 0.9),
            real_guidance_scale=self.real_guidance_scale,
            fake_guidance_scale=self.fake_guidance_scale,
            gen_cls_loss=self.gen_cls_loss,
            sway_coeff=0,
        )

        # Load student checkpoint
        checkpoint = torch.load(student_checkpoint_path, map_location='cpu')
        self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)

        # Setup generator and teacher
        self.generator = self.model.feedforward_model.to(self.device)
        self.teacher = self.model.guidance_model.real_unet.to(self.device)

        self.scale = checkpoint['scale']

    def _setup_mel_spec(self):
        """Initialize mel spectrogram module."""
        mel_spec_kwargs = dict(
            target_sample_rate=self.target_sample_rate,
            n_mel_channels=self.n_mel_channels,
            hop_length=self.hop_length,
        )
        self.mel_spec = MelSpec(**mel_spec_kwargs)

    def _setup_vocoder(self):
        """Initialize vocoder."""
        self.vocos = load_vocoder(is_local=False, local_path="")
        self.vocos = self.vocos.to(self.device)

    def _setup_duration_predictor(self, checkpoint_path):
        """Initialize duration predictor."""
        self.wav2mel = MelSpec(
            target_sample_rate=24000,
            n_mel_channels=100,


                      second_time=self.num_student_step > 1),
            checkpoint_path="",
            vocab_char_map=self.vocab_char_map,
            frac_lengths_mask=(0.5, 0.9),
            real_guidance_scale=self.real_guidance_scale,
            fake_guidance_scale=self.fake_guidance_scale,
            gen_cls_loss=self.gen_cls_loss,
            sway_coeff=0,
        )

        # Load student checkpoint
        checkpoint = torch.load(student_checkpoint_path, map_location='cpu')
        self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)

        # Setup generator and teacher
        self.generator = self.model.feedforward_model.to(self.device)
        self.teacher = self.model.guidance_model.real_unet.to(self.device)

        self.scale = checkpoint['scale']

    def _setup_mel_spec(self):
        """Initialize mel spectrogram module."""
        mel_spec_kwargs = dict(
            target_sample_rate=self.target_sample_rate,
            n_mel_channels=self.n_mel_channels,
            hop_length=self.hop_length,
        )
        self.mel_spec = MelSpec(**mel_spec_kwargs)

    def _setup_vocoder(self):
        """Initialize vocoder."""
        self.vocos = load_vocoder(is_local=False, local_path="")
        self.vocos = self.vocos.to(self.device)

    def _setup_duration_predictor(self, checkpoint_path):
        """Initialize duration predictor."""
        self.wav2mel = MelSpec(
            target_sample_rate=24000,
            n_mel_channels=100,
            hop_length=256,
            win_length=1024,
            n_fft=1024,
            mel_spec_type='vocos'
        ).to(self.device)

        self.SLP = SpeechLengthPredictor(
            vocab_size=2545,
            n_mel=100,


        # Load student checkpoint
        checkpoint = torch.load(student_checkpoint_path, map_location='cpu')
        self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)

        # Setup generator and teacher
        self.generator = self.model.feedforward_model.to(self.device)
        self.teacher = self.model.guidance_model.real_unet.to(self.device)

        self.scale = checkpoint['scale']

    def _setup_mel_spec(self):
        """Initialize mel spectrogram module."""
        mel_spec_kwargs = dict(
            target_sample_rate=self.target_sample_rate,
            n_mel_channels=self.n_mel_channels,
            hop_length=self.hop_length,
        )
        self.mel_spec = MelSpec(**mel_spec_kwargs)

    def _setup_vocoder(self):
        """Initialize vocoder."""
        self.vocos = load_vocoder(is_local=False, local_path="")
        self.vocos = self.vocos.to(self.device)

    def _setup_duration_predictor(self, checkpoint_path):
        """Initialize duration predictor."""
        self.wav2mel = MelSpec(
            target_sample_rate=24000,
            n_mel_channels=100,
            hop_length=256,
            win_length=1024,
            n_fft=1024,
            mel_spec_type='vocos'
        ).to(self.device)

        self.SLP = SpeechLengthPredictor(
            vocab_size=2545,
            n_mel=100,
            hidden_dim=512,
            n_text_layer=4,
            n_cross_layer=4,
            n_head=8,
            output_dim=301


        # Setup generator and teacher
        self.generator = self.model.feedforward_model.to(self.device)
        self.teacher = self.model.guidance_model.real_unet.to(self.device)

        self.scale = checkpoint['scale']

    def _setup_mel_spec(self):
        """Initialize mel spectrogram module."""
        mel_spec_kwargs = dict(
            target_sample_rate=self.target_sample_rate,
            n_mel_channels=self.n_mel_channels,
            hop_length=self.hop_length,
        )
        self.mel_spec = MelSpec(**mel_spec_kwargs)

    def _setup_vocoder(self):
        """Initialize vocoder."""
        self.vocos = load_vocoder(is_local=False, local_path="")
        self.vocos = self.vocos.to(self.device)

    def _setup_duration_predictor(self, checkpoint_path):
        """Initialize duration predictor."""
        self.wav2mel = MelSpec(
            target_sample_rate=24000,
            n_mel_channels=100,
            hop_length=256,
            win_length=1024,
            n_fft=1024,
            mel_spec_type='vocos'
        ).to(self.device)

        self.SLP = SpeechLengthPredictor(
            vocab_size=2545,
            n_mel=100,
            hidden_dim=512,
            n_text_layer=4,
            n_cross_layer=4,
            n_head=8,
            output_dim=301
        ).to(self.device)

        self.SLP.eval()
        self.SLP.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['model_state_dict'])

    def predict_duration(self, pmt_wav_path, tar_text, pmt_text, dp_softmax_range=0.7, temperature=0):
        """
        Predict duration for target text based on prompt audio.

        Args:
            pmt_wav_path: Path to prompt audio
            tar_text: Target text to generate
            pmt_text: Prompt text
            dp_softmax_range: softmax annliation range from rate-based duration
            temperature: temperature for softmax sampling (if 0, will use argmax)
        Returns:
            Estimated duration in frames
        """
        pmt_wav, sr = torchaudio.load(pmt_wav_path)
        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            pmt_wav = resampler(pmt_wav)
        if pmt_wav.size(0) > 1:
            pmt_wav = pmt_wav[0].unsqueeze(0)


            n_mel_channels=100,
            hop_length=256,
            win_length=1024,
            n_fft=1024,
            mel_spec_type='vocos'
        ).to(self.device)

        self.SLP = SpeechLengthPredictor(
            vocab_size=2545,
            n_mel=100,
            hidden_dim=512,
            n_text_layer=4,
            n_cross_layer=4,
            n_head=8,
            output_dim=301
        ).to(self.device)

        self.SLP.eval()
        self.SLP.load_state_dict(torch.load(checkpoint_path, map_location='cpu')['model_state_dict'])

    def predict_duration(self, pmt_wav_path, tar_text, pmt_text, dp_softmax_range=0.7, temperature=0):
        """
        Predict duration for target text based on prompt audio.

        Args:
            pmt_wav_path: Path to prompt audio
            tar_text: Target text to generate
            pmt_text: Prompt text
            dp_softmax_range: softmax annliation range from rate-based duration
            temperature: temperature for softmax sampling (if 0, will use argmax)
        Returns:
            Estimated duration in frames
        """
        pmt_wav, sr = torchaudio.load(pmt_wav_path)
        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            pmt_wav = resampler(pmt_wav)
        if pmt_wav.size(0) > 1:
            pmt_wav = pmt_wav[0].unsqueeze(0)
        pmt_wav = pmt_wav.to(self.device)

        pmt_mel = self.wav2mel(pmt_wav).permute(0, 2, 1)
        tar_tokens = self._convert_to_pinyin(list(tar_text))
        pmt_tokens = self._convert_to_pinyin(list(pmt_text))

        # Calculate duration
        ref_text_len = len(pmt_tokens)
        gen_text_len = len(tar_tokens)
        ref_audio_len = pmt_mel.size(1)
        duration = int(ref_audio_len / ref_text_len * gen_text_len / speed)
        duration = duration // 10

        min_duration = max(int(duration * dp_softmax_range), 0)
        max_duration = min(int(duration * (1 + dp_softmax_range)), 301)

        all_tokens = pmt_tokens + [' '] + tar_tokens

        text_ids = list_str_to_idx([all_tokens], self.vocab_char_map).to(self.device)
        text_ids = text_ids.masked_fill(text_ids == -1, self.vocab_size)

        with torch.no_grad():
            predictions = self.SLP(text_ids=text_ids, mel=pmt_mel)
        predictions = predictions[:, -1, :]
        predictions[:, :min_duration] = float('-inf')
        predictions[:, max_duration:] = float('-inf')

        if temperature == 0:
            est_label = predictions.argmax(-1)[..., -1].item() * 10
        else:
            probs = torch.softmax(predictions / temperature, dim=-1)
            sampled_idx = torch.multinomial(probs.squeeze(0), num_samples=1)  # Remove the -1 index
            est_label = sampled_idx.item() * 10

        return est_label

    def _convert_to_pinyin(self, char_list):
        """Convert character list to pinyin."""
        result = []
        for x in convert_char_to_pinyin(char_list):
            result = result + x
        while result[0] == ' ' and len(result) > 1:
            result = result[1:]
        return result

    def generate(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=16,
        teacher_stopping_time=0.07,
        student_start_step=1,
        duration=None,
        dp_softmax_range=0.7,


        all_tokens = pmt_tokens + [' '] + tar_tokens

        text_ids = list_str_to_idx([all_tokens], self.vocab_char_map).to(self.device)
        text_ids = text_ids.masked_fill(text_ids == -1, self.vocab_size)

        with torch.no_grad():
            predictions = self.SLP(text_ids=text_ids, mel=pmt_mel)
        predictions = predictions[:, -1, :]
        predictions[:, :min_duration] = float('-inf')
        predictions[:, max_duration:] = float('-inf')

        if temperature == 0:
            est_label = predictions.argmax(-1)[..., -1].item() * 10
        else:
            probs = torch.softmax(predictions / temperature, dim=-1)
            sampled_idx = torch.multinomial(probs.squeeze(0), num_samples=1)  # Remove the -1 index
            est_label = sampled_idx.item() * 10

        return est_label

    def _convert_to_pinyin(self, char_list):
        """Convert character list to pinyin."""
        result = []
        for x in convert_char_to_pinyin(char_list):
            result = result + x
        while result[0] == ' ' and len(result) > 1:
            result = result[1:]
        return result

    def generate(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=16,
        teacher_stopping_time=0.07,
        student_start_step=1,
        duration=None,
        dp_softmax_range=0.7,
        temperature=0,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0,
        verbose=False
    ):
        """
        Generate speech from text using teacher-student distillation.


generate

        predictions[:, max_duration:] = float('-inf')

        if temperature == 0:
            est_label = predictions.argmax(-1)[..., -1].item() * 10
        else:
            probs = torch.softmax(predictions / temperature, dim=-1)
            sampled_idx = torch.multinomial(probs.squeeze(0), num_samples=1)  # Remove the -1 index
            est_label = sampled_idx.item() * 10

        return est_label

    def _convert_to_pinyin(self, char_list):
        """Convert character list to pinyin."""
        result = []
        for x in convert_char_to_pinyin(char_list):
            result = result + x
        while result[0] == ' ' and len(result) > 1:
            result = result[1:]
        return result

    def generate(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=16,
        teacher_stopping_time=0.07,
        student_start_step=1,
        duration=None,
        dp_softmax_range=0.7,
        temperature=0,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0,
        verbose=False
    ):
        """
        Generate speech from text using teacher-student distillation.

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of teacher guidance steps
            teacher_stopping_time: When to stop teacher sampling
            student_start_step: When to start student sampling
            duration: Total duration (if None, will predict)
            dp_softmax_range: Duration predictor softmax range allowed around rate based duration
            temperature: Temperature for duration predictor sampling (0 means use argmax)
            eta: Stochasticity control (0=DDIM, 1=DDPM)
            cfg_strength: Classifier-free guidance strength
            sway_coefficient: Sway sampling coefficient
            verbose: Output sampling steps

        Returns:
            Generated audio waveform
        """
        if prompt_text is None:
            prompt_text = transcribe(audio_path)

        # Predict duration if not provided
        if duration is None:
            duration = self.predict_duration(audio_path, gen_text, prompt_text, dp_softmax_range, temperature)

        # Preprocess audio and text
        ref_audio, ref_text = preprocess_ref_audio_text(audio_path, prompt_text)
        audio, sr = torchaudio.load(ref_audio)

        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        # Normalize audio
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms

        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            audio = resampler(audio)

        audio = audio.to(self.device)

        # Prepare text
        text_list = [ref_text + gen_text]
        final_text_list = convert_char_to_pinyin(text_list)

        # Calculate durations
        ref_audio_len = audio.shape[-1] // self.hop_length
        if duration is None:
            ref_text_len = len(ref_text.encode("utf-8"))
            gen_text_len = len(gen_text.encode("utf-8"))
            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
        else:
            duration = ref_audio_len + duration

        if verbose:
            print('audio:', audio.shape)
            print('text:', final_text_list)
            print('duration:', duration)
            print('eta (stochasticity):', eta)  # Print eta value for debugging

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration, duration_tensor = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-student sampling
            if teacher_steps > 0 and student_start_step > 0:
                if verbose:
                    print('Start teacher sampling with hybrid DDIM/DDPM (eta={})....'.format(eta))
                x1 = self._teacher_sampling(
                    step_cond, text, cond_mask, max_duration, duration_tensor,  # Use duration_tensor
                    teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_coefficient
                )
            else:
                x1 = step_cond

            if verbose:
                print('Start student sampling...')
            # Student sampling
            x1 = self._student_sampling(x1, cond, text, student_start_step, verbose, sway_coefficient)

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave.cpu().numpy().squeeze()

    def generate_teacher_only(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=32,
        duration=None,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0
    ):
        """
        Generate speech using teacher model only (no student distillation).

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of sampling steps


                if verbose:
                    print('Start teacher sampling with hybrid DDIM/DDPM (eta={})....'.format(eta))
                x1 = self._teacher_sampling(
                    step_cond, text, cond_mask, max_duration, duration_tensor,  # Use duration_tensor
                    teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_coefficient
                )
            else:
                x1 = step_cond

            if verbose:
                print('Start student sampling...')
            # Student sampling
            x1 = self._student_sampling(x1, cond, text, student_start_step, verbose, sway_coefficient)

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave.cpu().numpy().squeeze()

    def generate_teacher_only(
        self,
        gen_text,
        audio_path,
        prompt_text=None,
        teacher_steps=32,
        duration=None,
        eta=1.0,
        cfg_strength=2.0,
        sway_coefficient=-1.0
    ):
        """
        Generate speech using teacher model only (no student distillation).

        Args:
            gen_text: Text to generate
            audio_path: Path to prompt audio
            prompt_text: Prompt text (if None, will use ASR)
            teacher_steps: Number of sampling steps
            duration: Total duration (if None, will predict)
            eta: Stochasticity control (0=DDIM, 1=DDPM)
            cfg_strength: Classifier-free guidance strength
            sway_coefficient: Sway sampling coefficient

        Returns:
            Generated audio waveform
        """
        if prompt_text is None:
            prompt_text = transcribe(audio_path)

        # Predict duration if not provided
        if duration is None:
            duration = self.predict_duration(audio_path, gen_text, prompt_text)

        # Preprocess audio and text
        ref_audio, ref_text = preprocess_ref_audio_text(audio_path, prompt_text)
        audio, sr = torchaudio.load(ref_audio)

        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)

        # Normalize audio
        rms = torch.sqrt(torch.mean(torch.square(audio)))
        if rms < target_rms:
            audio = audio * target_rms / rms

        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            audio = resampler(audio)

        audio = audio.to(self.device)

        # Prepare text
        text_list = [ref_text + gen_text]
        final_text_list = convert_char_to_pinyin(text_list)

        # Calculate durations
        ref_audio_len = audio.shape[-1] // self.hop_length
        if duration is None:
            ref_text_len = len(ref_text.encode("utf-8"))
            gen_text_len = len(gen_text.encode("utf-8"))
            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
        else:
            duration = ref_audio_len + duration

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-only sampling
            x1 = self._teacher_sampling(
                step_cond, text, cond_mask, max_duration, duration,
                teacher_steps, 1.0, eta, cfg_strength, sway_coefficient  # stopping_time=1.0 for full sampling
            )

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave

    def _prepare_inputs(self, audio, text_list, duration):
        """Prepare inputs for generation."""
        lens = None
        max_duration_limit = 4096

        cond = audio
        text = text_list

        if cond.ndim == 2:
            cond = self.mel_spec(cond)
            cond = cond.permute(0, 2, 1)
            assert cond.shape[-1] == 100

        cond = cond / self.scale
        batch, cond_seq_len, device = *cond.shape[:2], cond.device

        if not exists(lens):
            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)


            duration = ref_audio_len + duration

        # Run inference
        with torch.inference_mode():
            cond, text, step_cond, cond_mask, max_duration = self._prepare_inputs(
                audio, final_text_list, duration
            )

            # Teacher-only sampling
            x1 = self._teacher_sampling(
                step_cond, text, cond_mask, max_duration, duration,
                teacher_steps, 1.0, eta, cfg_strength, sway_coefficient  # stopping_time=1.0 for full sampling
            )

            # Decode to audio
            mel = x1.permute(0, 2, 1) * self.scale
            generated_wave = self.vocos.decode(mel[..., cond_mask.sum():])

        return generated_wave

    def _prepare_inputs(self, audio, text_list, duration):
        """Prepare inputs for generation."""
        lens = None
        max_duration_limit = 4096

        cond = audio
        text = text_list

        if cond.ndim == 2:
            cond = self.mel_spec(cond)
            cond = cond.permute(0, 2, 1)
            assert cond.shape[-1] == 100

        cond = cond / self.scale
        batch, cond_seq_len, device = *cond.shape[:2], cond.device

        if not exists(lens):
            lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)

        # Process text
        if isinstance(text, list):
            if exists(self.vocab_char_map):
                text = list_str_to_idx(text, self.vocab_char_map).to(device)
            else:
                text = list_str_to_tensor(text).to(device)
            assert text.shape[0] == batch

        if exists(text):
            text_lens = (text != -1).sum(dim=-1)
            lens = torch.maximum(text_lens, lens)

        # Process duration
        cond_mask = lens_to_mask(lens)

        if isinstance(duration, int):
            duration = torch.full((batch,), duration, device=device, dtype=torch.long)

        duration = torch.maximum(lens + 1, duration)
        duration = duration.clamp(max=max_duration_limit)
        max_duration = duration.amax()

        # Pad conditioning
        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
        cond_mask = cond_mask.unsqueeze(-1)
        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))

        return cond, text, step_cond, cond_mask, max_duration, duration

    def _teacher_sampling(self, step_cond, text, cond_mask, max_duration, duration,
                         teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_sampling_coef = -1):
        """Perform teacher model sampling."""
        device = step_cond.device

        # Pre-generate noise sequence for stochastic sampling
        noise_seq = None
        if eta > 0:
            noise_seq = [torch.randn(1, max_duration, 100, device=device)
                        for _ in range(teacher_steps)]

        def fn(t, x):
            with torch.inference_mode():
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    if verbose:
                        print(f'current t: {t}')
                    step_frac = 1.0 - t.item()
                    step_idx = min(int(step_frac * len(noise_seq)), len(noise_seq) - 1) if noise_seq else 0


            lens = torch.maximum(text_lens, lens)

        # Process duration
        cond_mask = lens_to_mask(lens)

        if isinstance(duration, int):
            duration = torch.full((batch,), duration, device=device, dtype=torch.long)

        duration = torch.maximum(lens + 1, duration)
        duration = duration.clamp(max=max_duration_limit)
        max_duration = duration.amax()

        # Pad conditioning
        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
        cond_mask = cond_mask.unsqueeze(-1)
        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))

        return cond, text, step_cond, cond_mask, max_duration, duration

    def _teacher_sampling(self, step_cond, text, cond_mask, max_duration, duration,
                         teacher_steps, teacher_stopping_time, eta, cfg_strength, verbose, sway_sampling_coef = -1):
        """Perform teacher model sampling."""
        device = step_cond.device

        # Pre-generate noise sequence for stochastic sampling
        noise_seq = None
        if eta > 0:
            noise_seq = [torch.randn(1, max_duration, 100, device=device)
                        for _ in range(teacher_steps)]

        def fn(t, x):
            with torch.inference_mode():
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    if verbose:
                        print(f'current t: {t}')
                    step_frac = 1.0 - t.item()
                    step_idx = min(int(step_frac * len(noise_seq)), len(noise_seq) - 1) if noise_seq else 0

                    # Predict flow
                    pred = self.teacher(
                        x=x, cond=step_cond, text=text, time=t, mask=None,
                        drop_audio_cond=False, drop_text=False
                    )

                    if cfg_strength > 1e-5:
                        null_pred = self.teacher(
                            x=x, cond=step_cond, text=text, time=t, mask=None,
                            drop_audio_cond=True, drop_text=True
                        )
                        pred = pred + (pred - null_pred) * cfg_strength

                    # Add stochasticity if eta > 0
                    if eta > 0 and noise_seq is not None:
                        alpha_t = 1.0 - t.item()
                        sigma_t = t.item()
                        noise_scale = torch.sqrt(torch.tensor(
                            (sigma_t**2) / (alpha_t**2 + sigma_t**2) * eta,
                            device=device
                        ))
                        return pred + noise_scale * noise_seq[step_idx]
                    else:
                        return pred

        # Initialize noise
        y0 = []
        for dur in duration:
            y0.append(torch.randn(dur, 100, device=device, dtype=step_cond.dtype))
        y0 = pad_sequence(y0, padding_value=0, batch_first=True)

        # Setup time steps
        t = torch.linspace(0, 1, teacher_steps + 1, device=device, dtype=step_cond.dtype)
        if sway_sampling_coef is not None:
            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
        t = t[:(t > teacher_stopping_time).float().argmax() + 2]
        t = t[:-1]

        # Solve ODE
        trajectory = odeint(fn, y0, t, method="euler")

        if teacher_stopping_time < 1.0:
            # If early stopping, compute final step
            pred = fn(t[-1], trajectory[-1])
            test_out = trajectory[-1] + (1 - t[-1]) * pred
            return test_out
        else:
            return trajectory[-1]

    def _student_sampling(self, x1, cond, text, student_start_step, verbose, sway_coeff = -1):
        """Perform student model sampling."""
        steps = torch.Tensor([0, 0.25, 0.5, 0.75])
        steps = steps + sway_coeff * (torch.cos(torch.pi / 2 * steps) - 1 + steps)
        steps = steps[student_start_step:]

        for step in steps:
            time = torch.Tensor([step]).to(x1.device)

            x0 = torch.randn_like(x1)
            t = time.unsqueeze(-1).unsqueeze(-1)
            phi = (1 - t) * x0 + t * x1

            if verbose:
                print(f'current step: {step}')
            with torch.no_grad():
                pred = self.generator(
                    x=phi,
                    cond=cond,